""" Pydantic AI powered coding runtime with Claude Code style tools. """ from __future__ import annotations import json import os import shlex import subprocess from pathlib import Path from typing import Any, Optional from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext from .config import ( AgentMetadata, AzureBlobResourceConfig, CodingRequestConfig, DatabaseEngine, DatabaseResourceConfig, GitResourceConfig, LiteLLMConfig, ResourceConfig, WorkspaceConfig, ) from .resources import build_authenticated_repo_url, safe_workspace_path, summarize_resources class ToolEvent(BaseModel): tool: str payload: dict[str, Any] = Field(default_factory=dict) class CodingRunContext(BaseModel): workspace: WorkspaceConfig resources: ResourceConfig changed_files: list[str] = Field(default_factory=list) tool_log: list[ToolEvent] = Field(default_factory=list) finish_summary: Optional[str] = None model_config = {"arbitrary_types_allowed": True} class CodingRunResult(BaseModel): response_text: str summary: str workspace_root: str files_changed: list[str] = Field(default_factory=list) tool_log: list[ToolEvent] = Field(default_factory=list) resources_used: dict[str, Any] = Field(default_factory=dict) class CodingA2ARuntime: def __init__( self, llm_config: LiteLLMConfig, metadata: AgentMetadata, ): self.llm_config = llm_config self.metadata = metadata self._apply_llm_env(self.llm_config) self._agent = self._build_agent(self.llm_config) def _build_agent(self, llm_config: LiteLLMConfig) -> Agent: agent: Agent[CodingRunContext] = Agent( llm_config.normalized_model, system_prompt=self.metadata.effective_system_prompt, deps_type=CodingRunContext, ) @agent.tool async def read_file(ctx: RunContext[CodingRunContext], path: str) -> str: target = safe_workspace_path( ctx.deps.workspace.root_dir, path, ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths, ) content = target.read_text(encoding="utf-8", errors="replace") ctx.deps.tool_log.append(ToolEvent(tool="read_file", payload={"path": path, "bytes": len(content)})) return content @agent.tool async def list_files(ctx: RunContext[CodingRunContext], glob_pattern: str = "**/*") -> str: root = Path(ctx.deps.workspace.root_dir).resolve() matched = [ path.relative_to(root).as_posix() for path in sorted(root.glob(glob_pattern)) if path.is_file() ] ctx.deps.tool_log.append(ToolEvent(tool="list_files", payload={"pattern": glob_pattern, "count": len(matched)})) return "\n".join(matched[:500]) if matched else "(no files matched)" @agent.tool async def write_file(ctx: RunContext[CodingRunContext], path: str, content: str) -> str: target = safe_workspace_path( ctx.deps.workspace.root_dir, path, ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths, ) target.parent.mkdir(parents=True, exist_ok=True) target.write_text(content, encoding="utf-8") rel = target.relative_to(Path(ctx.deps.workspace.root_dir).resolve()).as_posix() if rel not in ctx.deps.changed_files: ctx.deps.changed_files.append(rel) ctx.deps.tool_log.append(ToolEvent(tool="write_file", payload={"path": rel, "bytes": len(content.encode())})) return f"written {rel}" @agent.tool async def edit_file(ctx: RunContext[CodingRunContext], path: str, old_text: str, new_text: str) -> str: target = safe_workspace_path( ctx.deps.workspace.root_dir, path, ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths, ) content = target.read_text(encoding="utf-8", errors="replace") if old_text not in content: return "old_text not found" updated = content.replace(old_text, new_text, 1) target.write_text(updated, encoding="utf-8") rel = target.relative_to(Path(ctx.deps.workspace.root_dir).resolve()).as_posix() if rel not in ctx.deps.changed_files: ctx.deps.changed_files.append(rel) ctx.deps.tool_log.append(ToolEvent(tool="edit_file", payload={"path": rel})) return f"edited {rel}" @agent.tool async def run_command(ctx: RunContext[CodingRunContext], command: str) -> str: completed = subprocess.run( ["bash", "-lc", command], cwd=ctx.deps.workspace.root_dir, capture_output=True, text=True, timeout=300, ) ctx.deps.tool_log.append( ToolEvent( tool="run_command", payload={"command": command, "returncode": completed.returncode}, ) ) output = completed.stdout.strip() if completed.stderr.strip(): output = f"{output}\n{completed.stderr.strip()}".strip() return output or f"(exit code {completed.returncode})" @agent.tool async def git_prepare_workspace(ctx: RunContext[CodingRunContext]) -> str: git = ctx.deps.resources.git if not git or not git.repo_url: return "git resource not configured" workspace = Path(ctx.deps.workspace.root_dir) if (workspace / ".git").exists(): return "workspace already contains a git repository" workspace.parent.mkdir(parents=True, exist_ok=True) auth_url = build_authenticated_repo_url( git.repo_url, username=git.username, password=git.password, token=git.token, ) branch = git.default_branch or "main" completed = subprocess.run( ["git", "clone", "--branch", branch, auth_url, workspace.name], cwd=str(workspace.parent), capture_output=True, text=True, timeout=300, ) ctx.deps.tool_log.append( ToolEvent( tool="git_prepare_workspace", payload={"repo_url": git.repo_url, "returncode": completed.returncode}, ) ) if completed.returncode != 0: return (completed.stdout + "\n" + completed.stderr).strip() return f"cloned {git.repo_url} into {workspace}" @agent.tool async def git_status(ctx: RunContext[CodingRunContext]) -> str: return await run_command(ctx, "git status --short") @agent.tool async def git_diff(ctx: RunContext[CodingRunContext], ref: str = "HEAD") -> str: return await run_command(ctx, f"git diff {ref}") @agent.tool async def git_create_branch(ctx: RunContext[CodingRunContext], branch_name: str) -> str: return await run_command(ctx, f"git checkout -B {branch_name}") @agent.tool async def git_commit(ctx: RunContext[CodingRunContext], message: str) -> str: await run_command(ctx, "git add -A") return await run_command(ctx, f"git commit -m {shlex.quote(message)}") @agent.tool async def git_push(ctx: RunContext[CodingRunContext], remote: str = "origin", branch_name: Optional[str] = None) -> str: git = ctx.deps.resources.git if git and git.repo_url: auth_url = build_authenticated_repo_url( git.repo_url, username=git.username, password=git.password, token=git.token, ) subprocess.run( ["git", "remote", "set-url", remote, auth_url], cwd=ctx.deps.workspace.root_dir, capture_output=True, text=True, timeout=60, ) target = branch_name or "HEAD" return await run_command(ctx, f"git push {remote} {target}") @agent.tool async def list_database_tables(ctx: RunContext[CodingRunContext], engine: str = "postgresql") -> str: config = self._select_database_config(ctx.deps.resources, engine) if not config: return f"{engine} resource not configured" result = self._run_database_query(config, self._default_table_query(config.engine)) ctx.deps.tool_log.append(ToolEvent(tool="list_database_tables", payload={"engine": config.engine.value})) return result @agent.tool async def run_database_query(ctx: RunContext[CodingRunContext], engine: str, query: str) -> str: config = self._select_database_config(ctx.deps.resources, engine) if not config: return f"{engine} resource not configured" result = self._run_database_query(config, query) ctx.deps.tool_log.append(ToolEvent(tool="run_database_query", payload={"engine": config.engine.value})) return result @agent.tool async def list_blob_objects(ctx: RunContext[CodingRunContext], limit: int = 50) -> str: config = ctx.deps.resources.azure_blob if not config: return "azure_blob resource not configured" result = self._list_blob_objects(config, limit=limit) ctx.deps.tool_log.append(ToolEvent(tool="list_blob_objects", payload={"count": len(result)})) return json.dumps(result, ensure_ascii=False, indent=2) @agent.tool async def read_blob_text(ctx: RunContext[CodingRunContext], blob_name: str, encoding: str = "utf-8") -> str: config = ctx.deps.resources.azure_blob if not config: return "azure_blob resource not configured" text = self._read_blob_text(config, blob_name=blob_name, encoding=encoding) ctx.deps.tool_log.append(ToolEvent(tool="read_blob_text", payload={"blob_name": blob_name, "bytes": len(text.encode())})) return text @agent.tool async def finish(ctx: RunContext[CodingRunContext], summary: str) -> str: ctx.deps.finish_summary = summary ctx.deps.tool_log.append(ToolEvent(tool="finish", payload={"summary": summary})) return f"done: {summary}" return agent async def run_task( self, prompt: str, request_config: CodingRequestConfig, *, api_key: Optional[str] = None, model: Optional[str] = None, ) -> CodingRunResult: llm_config = LiteLLMConfig( base_url=self.llm_config.base_url, api_key=api_key or request_config.api_key or self.llm_config.api_key, model=model or request_config.model or self.llm_config.model, timeout=self.llm_config.timeout, max_tokens=self.llm_config.max_tokens, ) previous_api_key = os.environ.get("OPENAI_API_KEY") previous_base_url = os.environ.get("OPENAI_BASE_URL") self._apply_llm_env(llm_config) agent = self._build_agent(llm_config) deps = CodingRunContext(workspace=request_config.workspace, resources=request_config.resources) initial_prompt = self._build_initial_prompt(prompt, request_config) try: result = await agent.run(initial_prompt, deps=deps) response_text = getattr(result, "output", None) or getattr(result, "data", None) or str(result) summary = deps.finish_summary or response_text return CodingRunResult( response_text=str(response_text), summary=summary, workspace_root=request_config.workspace.root_dir, files_changed=deps.changed_files, tool_log=deps.tool_log, resources_used=summarize_resources(request_config.resources.model_dump(exclude_none=True)), ) finally: if previous_api_key is None: os.environ.pop("OPENAI_API_KEY", None) else: os.environ["OPENAI_API_KEY"] = previous_api_key if previous_base_url is None: os.environ.pop("OPENAI_BASE_URL", None) else: os.environ["OPENAI_BASE_URL"] = previous_base_url def _apply_llm_env(self, llm_config: LiteLLMConfig) -> None: if llm_config.api_key: os.environ["OPENAI_API_KEY"] = llm_config.api_key os.environ["OPENAI_BASE_URL"] = llm_config.base_url def _build_initial_prompt(self, prompt: str, request_config: CodingRequestConfig) -> str: workspace = request_config.workspace parts = [ f"USER TASK:\n{prompt}", f"WORKSPACE ROOT: {workspace.root_dir}", f"TASK MODE: {request_config.task_mode}", ] if workspace.entry_file: parts.append(f"START BY READING: {workspace.entry_file}") if workspace.context_files: parts.append("ALSO CONSIDER: " + ", ".join(workspace.context_files)) if workspace.allowed_paths: parts.append("YOU MAY ONLY MODIFY: " + ", ".join(workspace.allowed_paths)) if request_config.branch_name: parts.append(f"PREFERRED BRANCH: {request_config.branch_name}") if request_config.commit_message: parts.append(f"SUGGESTED COMMIT MESSAGE: {request_config.commit_message}") parts.append( "Use tools to inspect before editing. Prefer minimal precise changes. " "When the work is complete, call finish(summary)." ) return "\n\n".join(parts) def _select_database_config(self, resources: ResourceConfig, engine: str) -> Optional[DatabaseResourceConfig]: requested = engine.lower() if requested in {"mysql", "mariadb"}: return resources.mysql return resources.postgresql def _default_table_query(self, engine: DatabaseEngine) -> str: if engine == DatabaseEngine.mysql: return "SHOW TABLES" return ( "SELECT table_schema, table_name FROM information_schema.tables " "WHERE table_schema NOT IN ('pg_catalog', 'information_schema') " "ORDER BY table_schema, table_name LIMIT 200" ) def _run_database_query(self, config: DatabaseResourceConfig, query: str) -> str: if config.engine == DatabaseEngine.mysql: import pymysql connection = pymysql.connect( host=config.host, port=config.port or 3306, user=config.username, password=config.password, database=config.database, cursorclass=pymysql.cursors.DictCursor, connect_timeout=10, ) else: import psycopg2 import psycopg2.extras connection = psycopg2.connect( host=config.host, port=config.port or 5432, user=config.username, password=config.password, dbname=config.database, connect_timeout=10, sslmode=config.ssl_mode or "prefer", ) try: with connection.cursor() as cursor: cursor.execute(query) rows = cursor.fetchall() return json.dumps(rows, ensure_ascii=False, default=str, indent=2) finally: connection.close() def _get_blob_client(self, config: AzureBlobResourceConfig): from azure.storage.blob import BlobServiceClient if config.connection_string: return BlobServiceClient.from_connection_string(config.connection_string) if config.account_url and config.sas_token: return BlobServiceClient(account_url=config.account_url, credential=config.sas_token) if config.account_url and config.account_key: return BlobServiceClient(account_url=config.account_url, credential=config.account_key) if config.account_name and config.account_key: account_url = f"https://{config.account_name}.blob.core.windows.net" return BlobServiceClient(account_url=account_url, credential=config.account_key) raise ValueError("azure blob credentials are not configured") def _list_blob_objects(self, config: AzureBlobResourceConfig, limit: int = 50) -> list[dict[str, Any]]: service = self._get_blob_client(config) container = service.get_container_client(config.container_name) items = [] for index, blob in enumerate(container.list_blobs(name_starts_with=config.prefix or None)): if index >= limit: break items.append( { "name": blob.name, "size": blob.size, "content_type": getattr(blob.content_settings, "content_type", None), } ) return items def _read_blob_text(self, config: AzureBlobResourceConfig, blob_name: str, encoding: str = "utf-8") -> str: service = self._get_blob_client(config) blob_client = service.get_blob_client(container=config.container_name, blob=blob_name) return blob_client.download_blob().readall().decode(encoding, errors="replace")