diff --git a/.licenserc.yaml b/.licenserc.yaml index e43aaf7..1f34816 100644 --- a/.licenserc.yaml +++ b/.licenserc.yaml @@ -15,4 +15,4 @@ header: - 'docs/**/*' - 'LICENSE' - comment: on-failure \ No newline at end of file + comment: on-failure diff --git a/.vscode/settings.json b/.vscode/settings.json index 083f60b..1182e98 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -6,4 +6,4 @@ "packageManager": "ms-python.python:conda" } ] -} \ No newline at end of file +} diff --git a/LICENSE b/LICENSE index 895c63f..ede97f7 100644 --- a/LICENSE +++ b/LICENSE @@ -26,4 +26,4 @@ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md index ab8f479..4e586d6 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ Agent Codemode generates **programmatic tools** from two sources: 1. **MCP Servers** - Connect to any MCP server and generate typed Python bindings for its tools -2. **Skills** - Reusable code patterns that compose multiple tools into higher-level operations +1. **Skills** - Reusable code patterns that compose multiple tools into higher-level operations These programmatic tools can be: @@ -59,14 +59,21 @@ Same task, same MCP server — Code Mode uses significantly fewer tokens by comp ## Configuration Highlights -| Option | Description | -|--------|-------------| -| `allow_direct_tool_calls` | When `False` (default), `call_tool` is hidden; all execution flows through `execute_code` | -| `max_tool_calls` | Safety cap limiting tool invocations per `execute_code` run | -| `sandbox_variant` | Sandbox type for code execution (default: `"eval"`) | -| `workspace_path` | Working directory for sandbox execution | -| `generated_path` | Path where tool bindings are generated | -| `skills_path` | Path for saved skills | +| Option | Description | +| ------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `allow_direct_tool_calls` | When `False` (default), `call_tool` is hidden; all execution flows through `execute_code` | +| `max_tool_calls` | Safety cap limiting tool invocations per `execute_code` run | +| `sandbox_variant` | Sandbox type for code execution (default: `"eval"`). One of `eval`, `monty`, `docker`, `jupyter`, `colab`, `kaggle`, `modal`, `datalayer` | +| `sandbox_gpu` | Optional GPU flavor / accelerator for supported variants (Modal/Datalayer examples: `T4`, `A10G`, `A100`, `H100`; Kaggle batch examples: `NvidiaTeslaT4`, `NvidiaTeslaP100`, `T4`, `P100`) | +| `workspace_path` | Working directory for sandbox execution | +| `generated_path` | Path where tool bindings are generated | +| `skills_path` | Path for saved skills | + +When using cloud notebook variants via code-sandboxes: + +- `colab` reuses an already-running Colab kernel (from explicit runtime values or a channels URL). +- `kaggle` defaults to batch execution via Kaggle's API when no runtime URL/channels are provided; it can also create/attach interactive kernels when runtime values are provided. +- `kaggle` accelerator selection is available in batch mode via `gpu` / `accelerator` (for example `T4`, `P100`, `NvidiaTeslaT4`). ### Tool Discovery Options @@ -108,14 +115,14 @@ await registry.discover_all() async with CodeModeExecutor(registry) as executor: result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file - + # Read multiple files content1 = await read_file({"path": "/tmp/file1.txt"}) content2 = await read_file({"path": "/tmp/file2.txt"}) - + # Process and combine combined = content1 + "\\n---\\n" + content2 - + # Write result await write_file({"path": "/tmp/combined.txt", "content": combined}) """) @@ -150,10 +157,10 @@ async with CodeModeExecutor(registry) as executor: execution = await executor.execute(""" import asyncio from generated.mcp.filesystem import ls, read_file - + # List all files files = await ls({"path": "/data"}) - + # Read all files in parallel contents = await asyncio.gather(*[ read_file({"path": f}) for f in files @@ -186,25 +193,25 @@ The primary pattern is skills as Python files in a `skills/` directory: async def batch_process(input_dir: str, output_dir: str) -> dict: """Process all files in a directory. - + Args: input_dir: Input directory path. output_dir: Output directory path. - + Returns: Processing statistics. """ from generated.mcp.filesystem import list_directory, read_file, write_file - + entries = await list_directory({"path": input_dir}) processed = 0 - + for entry in entries.get("entries", []): content = await read_file({"path": f"{input_dir}/{entry}"}) # Process content... await write_file({"path": f"{output_dir}/{entry}", "content": content.upper()}) processed += 1 - + return {"processed": processed} ``` @@ -368,17 +375,17 @@ Then configure your MCP client to run the launcher script. **Tools exposed by the MCP server:** -| Tool | Description | -|------|-------------| -| `search_tools` | Progressive tool discovery | -| `list_servers` | List connected MCP servers | -| `list_tool_names` | Fast listing of tool names | -| `get_tool_details` | Get full tool schema | -| `execute_code` | Run code that composes tools | -| `call_tool` | Direct tool invocation | -| `save_skill` / `run_skill` | Skill management | -| `list_skills` / `delete_skill` | Skill management | -| `add_mcp_server` | Dynamically add servers | +| Tool | Description | +| ------------------------------ | ---------------------------- | +| `search_tools` | Progressive tool discovery | +| `list_servers` | List connected MCP servers | +| `list_tool_names` | Fast listing of tool names | +| `get_tool_details` | Get full tool schema | +| `execute_code` | Run code that composes tools | +| `call_tool` | Direct tool invocation | +| `save_skill` / `run_skill` | Skill management | +| `list_skills` / `delete_skill` | Skill management | +| `add_mcp_server` | Dynamically add servers | ## Recommended System Prompt @@ -394,7 +401,7 @@ You are an AI assistant with access to MCP tools via Code Mode. - **execute_code** - Execute Python code in a sandboxed environment ## Execution Model -ALL tool execution must go through execute_code. Write Python code that imports +ALL tool execution must go through execute_code. Write Python code that imports and uses the generated tool bindings: from generated.mcp.filesystem import read_file @@ -405,9 +412,8 @@ print(content) ## Workflow 1. Discover tools using search_tools or list_tool_names -2. Write Python code that imports tools from generated.mcp. -3. Execute using execute_code - +1. Write Python code that imports tools from generated.mcp. +1. Execute using execute_code See the [Getting Started guide](https://agent-codemode.datalayer.tech/getting-started) for a complete system prompt example. diff --git a/agent_codemode/__init__.py b/agent_codemode/__init__.py index 0eac064..79965cd 100644 --- a/agent_codemode/__init__.py +++ b/agent_codemode/__init__.py @@ -28,9 +28,28 @@ ''') """ +# Import skills functionality from agent_skills +from agent_skills import ( # type: ignore[import-untyped] + RateLimiter, + Skill, + SkillDirectory, + SkillFile, + SkillsManager, + parallel, + retry, + run_with_timeout, + setup_skills_directory, + wait_for, +) + from .composition.executor import CodeModeExecutor from .discovery.codegen import PythonCodeGenerator from .discovery.registry import ToolRegistry +from .proxy.mcp_client import MCPClient +from .proxy.meta_tools import MetaToolProvider +from .server import configure as configure_server +from .server import mcp as codemode_server +from .toolset import PYDANTIC_AI_AVAILABLE, CodemodeToolset from .types import ( CodeModeConfig, MCPServerConfig, @@ -40,58 +59,39 @@ ToolDefinition, ToolParameter, ) -from .proxy.mcp_client import MCPClient -from .proxy.meta_tools import MetaToolProvider - -# Import skills functionality from agent_skills -from agent_skills import ( - Skill, - SkillDirectory, - SkillFile, - SkillsManager, - setup_skills_directory, - wait_for, - retry, - run_with_timeout, - parallel, - RateLimiter, -) - -from .server import mcp as codemode_server, configure as configure_server -from .toolset import CodemodeToolset, PYDANTIC_AI_AVAILABLE __all__ = [ - # Core components - "ToolRegistry", + "PYDANTIC_AI_AVAILABLE", + "CodeModeConfig", "CodeModeExecutor", - "PythonCodeGenerator", + # Pydantic AI Toolset + "CodemodeToolset", # Proxy "MCPClient", + "MCPServerConfig", "MetaToolProvider", + "PythonCodeGenerator", + "RateLimiter", + "SearchResult", + "ServerInfo", # Skills (from agent_skills) "Skill", - "SkillsManager", "SkillDirectory", "SkillFile", - "setup_skills_directory", - # Helpers (from agent_skills) - "wait_for", - "retry", - "run_with_timeout", - "parallel", - "RateLimiter", - # MCP Server - "codemode_server", - "configure_server", - # Pydantic AI Toolset - "CodemodeToolset", - "PYDANTIC_AI_AVAILABLE", + "SkillsManager", + "ToolCallResult", # Models "ToolDefinition", "ToolParameter", - "ToolCallResult", - "MCPServerConfig", - "CodeModeConfig", - "SearchResult", - "ServerInfo", + # Core components + "ToolRegistry", + # MCP Server + "codemode_server", + "configure_server", + "parallel", + "retry", + "run_with_timeout", + "setup_skills_directory", + # Helpers (from agent_skills) + "wait_for", ] diff --git a/agent_codemode/__version__.py b/agent_codemode/__version__.py index bc08b8a..3e40be8 100644 --- a/agent_codemode/__version__.py +++ b/agent_codemode/__version__.py @@ -4,4 +4,4 @@ """Agent Codemode.""" -__version__ = "0.1.3" +__version__ = "0.1.5" diff --git a/agent_codemode/composition/executor.py b/agent_codemode/composition/executor.py index 9f177a8..73f8368 100644 --- a/agent_codemode/composition/executor.py +++ b/agent_codemode/composition/executor.py @@ -13,12 +13,12 @@ into the execution environment. When identities are set in the request context (via agent_runtimes.context.identities), they are automatically made available as environment variables in the sandbox: - + - GITHUB_TOKEN for GitHub OAuth - GITLAB_TOKEN for GitLab OAuth - GOOGLE_ACCESS_TOKEN for Google OAuth - AZURE_ACCESS_TOKEN for Microsoft OAuth - + This allows code executed via execute_code to access authenticated APIs without explicitly passing credentials. """ @@ -26,12 +26,13 @@ import logging import time from pathlib import Path +from types import TracebackType from typing import Any, Optional -from code_sandboxes import Sandbox, ExecutionResult, SandboxConfig +from code_sandboxes import ExecutionResult, Sandbox, SandboxConfig # type: ignore[import-untyped] -from ..discovery.registry import ToolRegistry from ..discovery.codegen import PythonCodeGenerator +from ..discovery.registry import ToolRegistry from ..types import CodeModeConfig, ToolCallResult logger = logging.getLogger(__name__) @@ -39,17 +40,22 @@ def _get_identity_env() -> dict[str, str]: """Get identity environment variables from request context. - + This function attempts to import the identity context from agent_runtimes. If not available (standalone codemode usage), returns empty dict. - + Returns: Dictionary of environment variable names to token values. """ try: - from agent_runtimes.context.identities import get_identity_env + import importlib + + get_identity_env = importlib.import_module( + "agent_runtimes.context.identities" + ).get_identity_env + return get_identity_env() - except ImportError: + except Exception: return {} @@ -134,11 +140,16 @@ def set_skills_metadata(self, metadata: list[dict[str, Any]]) -> None: def _is_local_eval_sandbox(self) -> bool: """Check if the sandbox is a eval type (has in-memory namespaces). - + This checks the actual sandbox instance, not the config, to handle cases where an external sandbox is passed that differs from config. """ - return self._sandbox is not None and hasattr(self._sandbox, '_namespaces') + return self._sandbox is not None and hasattr(self._sandbox, "_namespaces") + + def _require_sandbox(self) -> Sandbox: + if self._sandbox is None: + raise RuntimeError("Sandbox is not initialized") + return self._sandbox @property def sandbox(self) -> Optional[Sandbox]: @@ -152,8 +163,12 @@ async def setup(self) -> None: prepares the sandbox environment. """ import sys as _sys - print(f"[EXECUTOR.setup] Starting setup, sandbox_variant={self.config.sandbox_variant}", file=_sys.stderr) - + + print( + f"[EXECUTOR.setup] Starting setup, sandbox_variant={self.config.sandbox_variant}", + file=_sys.stderr, + ) + # Generate code bindings on the host filesystem. Skip when running in # sandbox-only mode (no generated modules exposed). if self.config.setup_generated_modules: @@ -163,19 +178,21 @@ async def setup(self) -> None: # Create sandbox if not provided if self._sandbox is None: import os + # Pass the complete environment to the sandbox env_vars = dict(os.environ) sandbox_config = SandboxConfig( - timeout=self.config.sandbox_variant == "datalayer-runtime" and 300 or 30, + timeout=300 if self.config.sandbox_variant == "datalayer" else 30, working_dir=self.config.workspace_path, env_vars=env_vars, + gpu=self.config.sandbox_gpu, ) sandbox_kwargs: dict[str, Any] = {} if self.config.sandbox_image: sandbox_kwargs["image"] = self.config.sandbox_image self._sandbox = Sandbox.create( - variant=self.config.sandbox_variant, # type: ignore + variant=self.config.sandbox_variant, config=sandbox_config, **sandbox_kwargs, ) @@ -205,7 +222,7 @@ async def _setup_sandbox_environment(self) -> None: # agent on the same host. if not self.config.setup_generated_modules: generated_parent = str(Path(self.config.generated_path).resolve().parent) - purge_code = f''' + purge_code = f""" import sys _generated_parent = {generated_parent!r} @@ -232,7 +249,7 @@ def find_spec(self, name, path=None, target=None): type(f).__name__ == "_BlockGeneratedFinder" for f in sys.meta_path ): sys.meta_path.insert(0, _BlockGeneratedFinder()) -''' +""" self._sandbox.run_code(purge_code) # Register the tool caller so ``call_tool`` still works inside # ``execute_code`` if the agent invokes it directly (raw MCP). @@ -252,7 +269,7 @@ def find_spec(self, name, path=None, target=None): sandbox_generated_path = str(generated_path.parent) # Add generated path and skills path to sys.path and clear any stale module cache - setup_code = f''' + setup_code = f""" import sys generated_path = {sandbox_generated_path!r} skills_path = {str(skills_path)!r} @@ -286,28 +303,35 @@ def find_spec(self, name, path=None, target=None): # Add skills path to sys.path (for skills imports) if skills_path not in sys.path: sys.path.insert(0, str(skills_path)) -''' +""" self._sandbox.run_code(setup_code) # Register tool caller with the sandbox import sys as _sys - print(f"[SETUP ENV DEBUG] About to call register_tool_caller, sandbox={self._sandbox} id={id(self._sandbox)}", file=_sys.stderr) + + print( + f"[SETUP ENV DEBUG] About to call register_tool_caller, sandbox={self._sandbox} id={id(self._sandbox)}", + file=_sys.stderr, + ) self._sandbox.register_tool_caller(self.call_tool) - print(f"[SETUP ENV DEBUG] register_tool_caller called", file=_sys.stderr) - + print("[SETUP ENV DEBUG] register_tool_caller called", file=_sys.stderr) + # Verify __call_tool__ was set - verify_code = ''' + verify_code = """ import sys try: print(f"[VERIFY] __call_tool__ = {__call_tool__}", file=sys.stderr) except NameError: print("[VERIFY] __call_tool__ NOT SET after register_tool_caller!", file=sys.stderr) -''' +""" self._sandbox.run_code(verify_code) # For Jupyter/remote sandboxes, set up in-sandbox registry for tool calling # Use actual sandbox type detection, not config - print(f"[SETUP ENV] is_local_eval={is_local_eval}, config.mcp_proxy_url={self.config.mcp_proxy_url}", file=_sys.stderr) + print( + f"[SETUP ENV] is_local_eval={is_local_eval}, config.mcp_proxy_url={self.config.mcp_proxy_url}", + file=_sys.stderr, + ) if not is_local_eval: # ======================================================================= # Two-Container Codemode Architecture @@ -337,11 +361,11 @@ def find_spec(self, name, path=None, target=None): # - Requires stdio MCP processes accessible from the sandbox (not always possible) # # ======================================================================= - + if self.config.mcp_proxy_url: # HTTP Proxy Mode: Use HTTP calls to agent-runtimes proxy endpoint proxy_url = self.config.mcp_proxy_url.rstrip("/") - + in_sandbox_http_caller_setup = f''' # ======================================================================= # HTTP Proxy Tool Caller for Two-Container Codemode @@ -371,29 +395,29 @@ def find_spec(self, name, path=None, target=None): async def __call_tool__(tool_name: str, arguments: dict) -> dict: """Call a tool via HTTP proxy to agent-runtimes. - + This function routes tool calls through the MCP proxy endpoint, which forwards them to the appropriate stdio MCP server. - + Args: tool_name: Full tool name in format "server__toolname" (e.g., "github__star_repo") arguments: Tool arguments dictionary - + Returns: Tool result dictionary """ # Parse server name from tool_name (format: server__toolname) if "__" not in tool_name: return {{"isError": True, "content": [{{"type": "text", "text": f"Invalid tool name format: {{tool_name}}. Expected server__toolname"}}]}} - + server_name, original_tool_name = tool_name.split("__", 1) - + # Build the proxy URL # Format: /api/v1/mcp/proxy/{{server_name}}/tools/{{tool_name}} url = f"{{__MCP_PROXY_URL__}}/{{server_name}}/tools/{{original_tool_name}}" - + print(f"[HTTP Proxy] Calling tool: {{tool_name}} -> {{url}}", file=sys.stderr) - + try: async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client: response = await client.post( @@ -401,21 +425,21 @@ async def __call_tool__(tool_name: str, arguments: dict) -> dict: json={{"arguments": arguments}}, headers={{"Content-Type": "application/json"}}, ) - + if response.status_code == 404: return {{ "isError": True, "content": [{{"type": "text", "text": f"MCP server '{{server_name}}' not found or tool '{{original_tool_name}}' not available"}}] }} - + if response.status_code != 200: return {{ "isError": True, "content": [{{"type": "text", "text": f"HTTP {{response.status_code}}: {{response.text}}"}}] }} - + result = response.json() - + # Convert proxy response to MCP format if result.get("success", False): return {{ @@ -427,7 +451,7 @@ async def __call_tool__(tool_name: str, arguments: dict) -> dict: "isError": True, "content": [{{"type": "text", "text": result.get("error", "Unknown error")}}] }} - + except httpx.ConnectError as e: return {{ "isError": True, @@ -455,7 +479,7 @@ async def __call_tool__(tool_name: str, arguments: dict) -> dict: } for config in self.registry._servers.values() ] - + in_sandbox_registry_setup = f''' try: from agent_codemode.proxy.mcp_client import MCPClient @@ -479,7 +503,7 @@ def __init__(self, server_configs): env=config["env"], ) self._clients[config["name"]] = client - + async def call_tool(self, tool_name, arguments): """Call a tool with arguments.""" # Parse server name from tool_name (format: server__toolname) @@ -487,18 +511,18 @@ async def call_tool(self, tool_name, arguments): server_name, original_name = tool_name.split("__", 1) else: return {{"error": f"Invalid tool name format: {{tool_name}}"}} - + client = self._clients.get(server_name) if not client: return {{"error": f"Server not available: {{server_name}}"}} - + return await client.call_tool(original_name, arguments) class _SandboxExecutor: """In-sandbox executor for tool calls.""" def __init__(self, registry): self._registry = registry - + async def call_tool(self, tool_name, arguments): return await self._registry.call_tool(tool_name, arguments) @@ -511,40 +535,42 @@ async def __call_tool__(tool_name, arguments): self._sandbox.run_code(in_sandbox_registry_setup) # Set up the generated client to use __call_tool__ - caller_setup_code = ''' + caller_setup_code = """ try: from generated.client import set_tool_caller set_tool_caller(__call_tool__) except (ImportError, NameError) as e: import sys print(f"[SETUP] caller_setup_code error: {type(e).__name__}: {e}", file=sys.stderr) -''' +""" self._sandbox.run_code(caller_setup_code) async def _generate_tools_in_sandbox(self) -> None: """Generate tool bindings directly in the remote sandbox. - + For Jupyter/remote sandboxes, we can't easily upload files, so instead we send the code generation logic to be executed in the sandbox. This way the generated modules exist in the sandbox's filesystem. """ if self._sandbox is None: return - + # Get tool definitions for code generation tools_dict = {tool.name: tool for tool in self.registry.list_tools()} - + # Build serializable tool data tools_data = [] for name, tool in tools_dict.items(): - tools_data.append({ - "name": tool.name, - "description": tool.description, - "input_schema": tool.input_schema, - "output_schema": tool.output_schema, - "server_name": tool.server_name, - }) - + tools_data.append( + { + "name": tool.name, + "description": tool.description, + "input_schema": tool.input_schema, + "output_schema": tool.output_schema, + "server_name": tool.server_name, + } + ) + # Code to generate bindings in the sandbox generation_code = f''' import os @@ -611,7 +637,7 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: if _tool_caller is None: raise RuntimeError("No tool caller configured.") result = await _tool_caller(tool_name, arguments) - + if not isinstance(result, (dict, object)) or result is None: return result @@ -626,38 +652,38 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: content_list = result.get("content") elif hasattr(result, "content"): content_list = result.content - + if not isinstance(content_list, list): return result text_content = "" has_text = False - + for part in content_list: part_type = None part_text = None - + if isinstance(part, dict): part_type = part.get("type") part_text = part.get("text") elif hasattr(part, "type") and hasattr(part, "text"): part_type = part.type part_text = part.text - + if part_type == "text" and part_text is not None: text_content += part_text has_text = True - + if is_error and has_text: raise RuntimeError(text_content) - + if has_text: try: import json return json.loads(text_content) except Exception: return text_content - + return result """ (__generated_path__ / "client.py").write_text(__client_code__) @@ -674,22 +700,22 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: for server_name, tools_list in __server_tools__.items(): server_dir = __mcp_path__ / server_name server_dir.mkdir(parents=True, exist_ok=True) - + imports = [] exports = [] - + for tool in tools_list: tool_name = tool["name"] if tool_name.startswith(f"{{server_name}}__"): short_name = tool_name[len(server_name) + 2:] else: short_name = tool_name - + func_name = _sanitize_name(short_name) input_type = _schema_to_type_hint(tool.get("input_schema", {{}})) output_type = _schema_to_type_hint(tool.get("output_schema")) if tool.get("output_schema") else "Any" description = tool.get("description", f"Call {{tool_name}} tool.") - + # Generate tool file tool_code = f"""# Auto-generated tool binding for {{tool_name}} from typing import Any, Optional @@ -704,10 +730,10 @@ async def {{func_name}}(arguments: Optional[{{input_type}}] = None, **kwargs: An return await call_tool("{{tool_name}}", arguments) """ (server_dir / f"{{func_name}}.py").write_text(tool_code) - + imports.append(f"from .{{func_name}} import {{func_name}}") exports.append(f' "{{func_name}}",') - + # Generate server index server_index = f"""# Auto-generated server module for {{server_name}} {{chr(10).join(imports)}} @@ -1005,9 +1031,7 @@ async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any: if self._in_execute and self.config.max_tool_calls is not None: if self._tool_calls_in_run >= self.config.max_tool_calls: - raise RuntimeError( - f"Tool call limit exceeded ({self.config.max_tool_calls})." - ) + raise RuntimeError(f"Tool call limit exceeded ({self.config.max_tool_calls}).") self._tool_calls_in_run += 1 try: @@ -1051,7 +1075,7 @@ async def execute( The code can import from the generated modules and call tools using async/await syntax. - + Identity tokens from the request context are automatically injected as environment variables (e.g., GITHUB_TOKEN, GITLAB_TOKEN). @@ -1074,7 +1098,7 @@ async def execute( try: # Get identity environment variables from request context identity_env = _get_identity_env() - + # Get the generated path for sys.path setup # For remote sandboxes, use /tmp so 'from generated.mcp...' works (files at /tmp/generated/) # For eval, use parent of generated_path so 'from generated.mcp...' works @@ -1084,7 +1108,7 @@ async def execute( generated_path = "/tmp" else: generated_path = str(Path(self.config.generated_path).resolve().parent) - + # Build identity injection code if we have tokens identity_injection = "" if identity_env: @@ -1094,9 +1118,9 @@ async def execute( import os os.environ.update({identity_env!r}) """ - + # Set up the environment before running user code - setup_code = f'''{identity_injection} + setup_code = f"""{identity_injection} import sys # Ensure generated path is first on sys.path and purge stale generated modules @@ -1128,7 +1152,7 @@ async def execute( __generated_spec__.loader.exec_module(__generated_module__) except Exception: pass -''' +""" # Branch based on actual sandbox type (already computed above) if is_local_eval: # For eval, we can access _namespaces directly @@ -1138,6 +1162,7 @@ async def execute( return await self._execute_jupyter(code, setup_code, timeout) finally: self._in_execute = False + async def _execute_local_eval( self, code: str, @@ -1145,33 +1170,41 @@ async def _execute_local_eval( timeout: Optional[float] = None, ) -> ExecutionResult: """Execute code in eval sandbox with direct namespace access.""" - import sys import io import time - from contextlib import redirect_stdout, redirect_stderr - from code_sandboxes.models import ExecutionResult, Logs, OutputMessage - + from contextlib import redirect_stderr, redirect_stdout + + from code_sandboxes.models import ( # type: ignore[import-untyped] + ExecutionResult, + Logs, + OutputMessage, + ) + + sandbox = self._require_sandbox() + # Get the namespace directly - namespace = self._sandbox._namespaces[self._sandbox._default_context.id] - + namespace = sandbox._namespaces[sandbox._default_context.id] + # Execute setup_code directly in namespace (avoids async wrapper issues) exec(setup_code, namespace, namespace) - + # Configure the generated.client tool caller if available - if '__call_tool__' in namespace: + if "__call_tool__" in namespace: try: - from generated.client import set_tool_caller - set_tool_caller(namespace['__call_tool__']) + import importlib + + set_tool_caller = importlib.import_module("generated.client").set_tool_caller + set_tool_caller(namespace["__call_tool__"]) except ImportError: pass - + # For async code, we need to handle it specially to avoid event loop conflicts if "await " in code or "async " in code: # Wrap user code in async function def _indent_code(value: str, spaces: int) -> str: indent = " " * spaces return "\n".join(indent + line for line in value.split("\n")) - + async_wrapper = f""" async def __user_code__(): {_indent_code(code, 4)} @@ -1179,11 +1212,11 @@ async def __user_code__(): """ # Execute the wrapper in namespace exec(async_wrapper, namespace, namespace) - + # Capture stdout/stderr stdout_buffer = io.StringIO() stderr_buffer = io.StringIO() - + exit_code = None # Call the async function directly (we're already in async context) @@ -1199,31 +1232,45 @@ async def __user_code__(): else: exit_code = 0 locals_value = {} - + # Update namespace with returned locals if isinstance(locals_value, dict): for key, value in locals_value.items(): - if key in ("__builtins__", "__name__", "__doc__", "__package__", - "__loader__", "__spec__", "__annotations__", "__cached__", - "__file__"): + if key in ( + "__builtins__", + "__name__", + "__doc__", + "__package__", + "__loader__", + "__spec__", + "__annotations__", + "__cached__", + "__file__", + ): continue namespace[key] = value - + stdout_lines = stdout_buffer.getvalue().splitlines() stderr_lines = stderr_buffer.getvalue().splitlines() timestamp = time.time() - + result = ExecutionResult( execution_ok=True, code_error=None, exit_code=exit_code, results=[], logs=Logs( - stdout=[OutputMessage(line=line, timestamp=timestamp, error=False) for line in stdout_lines], - stderr=[OutputMessage(line=line, timestamp=timestamp, error=True) for line in stderr_lines], + stdout=[ + OutputMessage(line=line, timestamp=timestamp, error=False) + for line in stdout_lines + ], + stderr=[ + OutputMessage(line=line, timestamp=timestamp, error=True) + for line in stderr_lines + ], ), - execution_count=self._sandbox._execution_count[self._sandbox._default_context.id], - context_id=self._sandbox._default_context.id, + execution_count=sandbox._execution_count[sandbox._default_context.id], + context_id=sandbox._default_context.id, ) else: # For sync code, run in a worker thread so FastAPI's event loop @@ -1231,11 +1278,11 @@ async def __user_code__(): import asyncio result = await asyncio.to_thread( - self._sandbox.run_code, + sandbox.run_code, code, timeout=timeout, ) - + return result async def _execute_jupyter( @@ -1245,13 +1292,13 @@ async def _execute_jupyter( timeout: Optional[float] = None, ) -> ExecutionResult: """Execute code in Jupyter/remote sandbox using run_code(). - + IMPORTANT: The sandbox.run_code() is synchronous and blocks waiting for the kernel to complete. When the kernel code calls back to the agent-runtimes server (e.g., via MCP proxy for tool calls), we need the event loop to be free to handle those requests. Therefore, we run the blocking code in a thread pool using asyncio.to_thread(). - + This prevents the deadlock: 1. FastAPI endpoint -> Jupyter sandbox (waiting) 2. Jupyter sandbox -> MCP proxy HTTP request @@ -1259,27 +1306,84 @@ async def _execute_jupyter( 4. With to_thread: MCP proxy handles request, kernel continues """ import asyncio - + + sandbox = self._require_sandbox() + # Run setup code in thread pool to avoid blocking event loop - await asyncio.to_thread(self._sandbox.run_code, setup_code, timeout=timeout) - + await asyncio.to_thread(sandbox.run_code, setup_code, timeout=timeout) + # Re-register the tool caller since the module cache was cleared # The __call_tool__ function was defined during initial setup and persists # in the kernel's global namespace, but we need to re-wire it to the # freshly-loaded generated.client module - tool_caller_rewire = ''' + tool_caller_rewire = """ try: from generated.client import set_tool_caller set_tool_caller(__call_tool__) except (ImportError, NameError) as e: import sys print(f"[EXECUTE] Failed to rewire tool caller: {type(e).__name__}: {e}", file=sys.stderr) -''' - await asyncio.to_thread(self._sandbox.run_code, tool_caller_rewire, timeout=timeout) - +""" + await asyncio.to_thread(sandbox.run_code, tool_caller_rewire, timeout=timeout) + + if hasattr(sandbox, "run_code_streaming"): + from code_sandboxes.models import ( + CodeError, + Logs, + OutputMessage, + Result, + ) + + def _collect_streaming_result() -> ExecutionResult: + stdout: list[OutputMessage] = [] + stderr: list[OutputMessage] = [] + results: list[Result] = [] + code_error: CodeError | None = None + + try: + for event in sandbox.run_code_streaming(code, timeout=timeout): + if hasattr(event, "line"): + message = OutputMessage( + line=str(getattr(event, "line", "") or ""), + timestamp=float(getattr(event, "timestamp", 0.0) or 0.0), + error=bool(getattr(event, "error", False)), + ) + if message.error: + stderr.append(message) + else: + stdout.append(message) + elif hasattr(event, "data"): + results.append( + Result( + data=getattr(event, "data", {}) or {}, + is_main_result=bool(getattr(event, "is_main_result", False)), + extra=getattr(event, "extra", {}) or {}, + ) + ) + elif hasattr(event, "name") and hasattr(event, "value"): + code_error = CodeError( + name=str(getattr(event, "name", "Error") or "Error"), + value=str(getattr(event, "value", "") or ""), + traceback=str(getattr(event, "traceback", "") or ""), + ) + + return ExecutionResult( + logs=Logs(stdout=stdout, stderr=stderr), + results=results, + code_error=code_error, + execution_ok=True, + ) + except Exception as exc: + return ExecutionResult( + execution_ok=False, + execution_error=str(exc), + ) + + return await asyncio.to_thread(_collect_streaming_result) + # Run user code in thread pool - this is where tool calls happen # and the kernel may call back to the MCP proxy - return await asyncio.to_thread(self._sandbox.run_code, code, timeout=timeout) + return await asyncio.to_thread(sandbox.run_code, code, timeout=timeout) def _indent_code(self, code: str, spaces: int) -> str: """Indent code by a number of spaces. @@ -1309,7 +1413,7 @@ async def execute_skill( Returns: Execution result. """ - from agent_skills import SkillsManager + from agent_skills import SkillsManager # type: ignore[import-untyped] manager = SkillsManager(self.config.skills_path) skill = manager.get(skill_name) @@ -1345,7 +1449,12 @@ async def __aenter__(self) -> "CodeModeExecutor": await self.setup() return self - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: """Async context manager exit.""" await self.cleanup() diff --git a/agent_codemode/discovery/__init__.py b/agent_codemode/discovery/__init__.py index 2a6b0b8..7f92259 100644 --- a/agent_codemode/discovery/__init__.py +++ b/agent_codemode/discovery/__init__.py @@ -4,7 +4,7 @@ """Tool discovery and registration.""" -from .registry import ToolRegistry from .codegen import PythonCodeGenerator +from .registry import ToolRegistry -__all__ = ["ToolRegistry", "PythonCodeGenerator"] +__all__ = ["PythonCodeGenerator", "ToolRegistry"] diff --git a/agent_codemode/discovery/codegen.py b/agent_codemode/discovery/codegen.py index 95df5e1..7e43819 100644 --- a/agent_codemode/discovery/codegen.py +++ b/agent_codemode/discovery/codegen.py @@ -10,6 +10,7 @@ import logging from pathlib import Path + logger = logging.getLogger(__name__) from typing import Any @@ -97,7 +98,7 @@ def _generate_client_module(self) -> None: def set_tool_caller(caller) -> None: """Set the global tool caller function. - + Args: caller: Async function that takes (tool_name, arguments) and returns result. """ @@ -107,14 +108,14 @@ def set_tool_caller(caller) -> None: async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: """Call an MCP tool. - + Args: tool_name: Full tool name (server__toolname format). arguments: Tool arguments. - + Returns: Tool result. - + Raises: RuntimeError: If no tool caller is configured. """ @@ -124,7 +125,7 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: "Use set_tool_caller() or run through CodeModeExecutor." ) result = await _tool_caller(tool_name, arguments) - + # helper to check if something is a list if not isinstance(result, (dict, object)) or result is None: return result @@ -142,33 +143,33 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: content_list = result.get("content") elif hasattr(result, "content"): content_list = result.content - + if not isinstance(content_list, list): return result # Concatenate text parts text_content = "" has_text = False - + for part in content_list: part_type = None part_text = None - + if isinstance(part, dict): part_type = part.get("type") part_text = part.get("text") elif hasattr(part, "type") and hasattr(part, "text"): part_type = part.type part_text = part.text - + if part_type == "text" and part_text is not None: text_content += part_text has_text = True - + # If this was an error response, raise an exception if is_error and has_text: raise RuntimeError(text_content) - + if has_text: # Try to parse as JSON first, as many tools return JSON string try: @@ -176,15 +177,13 @@ async def call_tool(tool_name: str, arguments: dict[str, Any]) -> Any: return json.loads(text_content) except Exception: return text_content - + return result ''' client_path = self.output_path / "client.py" client_path.write_text(client_code) - def _generate_server_module( - self, server_name: str, tools: list[ToolDefinition] - ) -> None: + def _generate_server_module(self, server_name: str, tools: list[ToolDefinition]) -> None: """Generate a module for a server's tools. Args: @@ -201,9 +200,7 @@ def _generate_server_module( # Generate server index self._generate_server_index(server_dir, server_name, tools) - def _generate_tool_file( - self, server_dir: Path, server_name: str, tool: ToolDefinition - ) -> None: + def _generate_tool_file(self, server_dir: Path, server_name: str, tool: ToolDefinition) -> None: """Generate a file for a single tool. Args: @@ -213,7 +210,7 @@ def _generate_tool_file( """ # Extract tool name without server prefix if tool.name.startswith(f"{server_name}__"): - short_name = tool.name[len(server_name) + 2:] + short_name = tool.name[len(server_name) + 2 :] else: short_name = tool.name @@ -235,18 +232,20 @@ def _generate_tool_file( # Extract the main parameter names from the schema schema_props = tool.input_schema.get("properties", {}) if tool.input_schema else {} param_names = list(schema_props.keys()) - + # Build parameter list for flexible functions param_list = [] for prop_name, prop_def in schema_props.items(): param_type = self._schema_property_to_type_hint(prop_def) param_list.append(f" {prop_name}: Optional[{param_type}] = None") - - param_signature = ",\n".join(param_list) + ",\n **kwargs: Any" if param_list else "**kwargs: Any" - + + param_signature = ( + ",\n".join(param_list) + ",\n **kwargs: Any" if param_list else "**kwargs: Any" + ) + # Generate flexible call logic call_logic = self._generate_flexible_call_logic(param_names) - + code = f'''# Auto-generated tool binding for {tool.name} # Copyright (c) 2025-2026 Datalayer, Inc. # BSD 3-Clause License @@ -311,7 +310,7 @@ def _generate_server_index( for tool in tools: if tool.name.startswith(f"{server_name}__"): - short_name = tool.name[len(server_name) + 2:] + short_name = tool.name[len(server_name) + 2 :] else: short_name = tool.name func_name = self._sanitize_name(short_name) @@ -412,6 +411,7 @@ def generate_skill_bindings( # Embed the skill catalog as a constant so list_skills is self-contained import json as _json + skill_catalog_json = _json.dumps(skills, ensure_ascii=False, indent=2) # --- list_skills --- @@ -582,7 +582,10 @@ async def run_skill( ''' (skills_dir / "__init__.py").write_text(init_code) - logger.info("Generated skill bindings: list_skills, load_skill, read_skill_resource, run_skill") + logger.info( + "Generated skill bindings: list_skills, load_skill, read_skill_resource, run_skill" + ) + def _sanitize_name(self, name: str) -> str: """Sanitize a name to be a valid Python identifier. @@ -606,6 +609,7 @@ def _sanitize_name(self, name: str) -> str: # Handle Python keywords import keyword + if keyword.iskeyword(result): result = result + "_" @@ -653,9 +657,7 @@ def _has_doc_sections(text: str) -> bool: """ import re - return bool( - re.search(r"^\s*(Args|Returns|Raises|Yields)\s*:", text, re.MULTILINE) - ) + return bool(re.search(r"^\s*(Args|Returns|Raises|Yields)\s*:", text, re.MULTILINE)) def _generate_docstring(self, tool: ToolDefinition) -> str: """Generate a docstring for a tool function. @@ -711,10 +713,12 @@ def _generate_docstring(self, tool: ToolDefinition) -> str: lines.append(" Tool execution result.") return "\n ".join(lines) - - def _add_schema_description(self, lines: list[str], schema: dict[str, Any], indent: str) -> None: + + def _add_schema_description( + self, lines: list[str], schema: dict[str, Any], indent: str + ) -> None: """Add schema description to docstring lines. - + Args: lines: List of docstring lines to append to schema: JSON Schema to describe @@ -722,7 +726,7 @@ def _add_schema_description(self, lines: list[str], schema: dict[str, Any], inde """ if not schema: return - + schema_type = schema.get("type") if schema_type == "object" and "properties" in schema: lines.append(f"{indent}Object with properties:") @@ -746,46 +750,47 @@ def _add_schema_description(self, lines: list[str], schema: dict[str, Any], inde def _needs_flexible_parameters(self, schema: dict[str, Any]) -> bool: """Check if a tool schema needs flexible parameter handling. - + Returns True if the schema has properties that suggest it needs individual parameter handling (like tool_name + arguments pattern). - + Args: schema: Tool input schema - + Returns: True if flexible parameters are needed """ if not schema or "properties" not in schema: return False - + props = schema["properties"] - + # Check for common patterns that need flexible handling: - # 1. Has both "tool_name" and "arguments" + # 1. Has both "tool_name" and "arguments" # 2. Has "function_name" or "method_name" with "arguments"/"parameters" # 3. Any combination that suggests nested tool calling - - has_tool_identifier = any(key in props for key in [ - "tool_name", "function_name", "method_name", "action", "command" - ]) - has_arguments = any(key in props for key in [ - "arguments", "parameters", "args", "params", "input", "data" - ]) - + + has_tool_identifier = any( + key in props + for key in ["tool_name", "function_name", "method_name", "action", "command"] + ) + has_arguments = any( + key in props for key in ["arguments", "parameters", "args", "params", "input", "data"] + ) + return has_tool_identifier and has_arguments - + def _schema_property_to_type_hint(self, prop_def: dict[str, Any]) -> str: """Convert a single JSON Schema property to Python type hint. - + Args: prop_def: Property definition from JSON Schema - + Returns: Python type hint string """ prop_type = prop_def.get("type", "any") - + if prop_type == "string": return "str" elif prop_type == "number": @@ -800,19 +805,19 @@ def _schema_property_to_type_hint(self, prop_def: dict[str, Any]) -> str: return "dict[str, Any]" else: return "Any" - + def _generate_flexible_call_logic(self, param_names: list[str]) -> str: """Generate the call logic for flexible parameter functions. - + Args: param_names: List of parameter names from schema - + Returns: Indented Python code for parameter handling """ lines = [ " # Support both calling styles:", - " # 1. Natural keyword arguments: func(param1=val1, param2=val2)", + " # 1. Natural keyword arguments: func(param1=val1, param2=val2)", " # 2. Single arguments dict: func(arguments={'param1': val1, 'param2': val2})", " ", " # Check if any schema parameters were provided directly", @@ -823,7 +828,7 @@ def _generate_flexible_call_logic(self, param_names: list[str]) -> str: " call_args = direct_params", " call_args.update(kwargs)", " elif 'arguments' in kwargs and isinstance(kwargs['arguments'], dict):", - " # Style 2: Single arguments dict provided", + " # Style 2: Single arguments dict provided", " call_args = kwargs['arguments'].copy()", " # Add any other kwargs that aren't 'arguments'", " call_args.update({k: v for k, v in kwargs.items() if k != 'arguments'})", @@ -831,5 +836,5 @@ def _generate_flexible_call_logic(self, param_names: list[str]) -> str: " # Fallback: use kwargs as arguments", " call_args = kwargs", ] - + return "\n".join(lines) diff --git a/agent_codemode/discovery/registry.py b/agent_codemode/discovery/registry.py index 2b2d339..c575d74 100644 --- a/agent_codemode/discovery/registry.py +++ b/agent_codemode/discovery/registry.py @@ -11,8 +11,8 @@ import logging from typing import Optional -from ..types import MCPServerConfig, SearchResult, ServerInfo, ToolDefinition from ..proxy.mcp_client import MCPClient +from ..types import MCPServerConfig, SearchResult, ServerInfo, ToolDefinition logger = logging.getLogger(__name__) @@ -317,9 +317,7 @@ async def search_tools( query=query, ) - async def call_tool( - self, tool_name: str, arguments: dict - ) -> dict: + async def call_tool(self, tool_name: str, arguments: dict) -> dict: """Call a tool with arguments. Args: diff --git a/agent_codemode/proxy/mcp_client.py b/agent_codemode/proxy/mcp_client.py index 04f1923..036aa4c 100644 --- a/agent_codemode/proxy/mcp_client.py +++ b/agent_codemode/proxy/mcp_client.py @@ -6,8 +6,8 @@ import asyncio import json -from typing import Any, Optional from contextlib import suppress +from typing import Any, Optional import httpx @@ -49,10 +49,10 @@ def __init__( self._http_client: Optional[httpx.AsyncClient] = None self._stdio_process: Optional[asyncio.subprocess.Process] = None self._request_id = 0 - self._stdio_session = None - self._stdio_ctx = None - self._http_session = None - self._http_ctx = None + self._stdio_session: Any | None = None + self._stdio_ctx: Any | None = None + self._http_session: Any | None = None + self._http_ctx: Any | None = None @property def is_http(self) -> bool: @@ -82,6 +82,7 @@ async def _start_stdio_process(self) -> asyncio.subprocess.Process: """Start the stdio process.""" if self._stdio_process is None or self._stdio_process.returncode is not None: import os + env = {**os.environ, **self.env} self._stdio_process = await asyncio.create_subprocess_exec( self.command, @@ -99,6 +100,7 @@ async def _get_stdio_session(self): return self._stdio_session import os as _os + from mcp.client.session import ClientSession from mcp.client.stdio import StdioServerParameters, stdio_client @@ -107,12 +109,14 @@ async def _get_stdio_session(self): args=self.args, env={**_os.environ, **self.env} if self.env else None, ) - self._stdio_ctx = stdio_client(params) - read_stream, write_stream = await self._stdio_ctx.__aenter__() - self._stdio_session = ClientSession(read_stream, write_stream) - await self._stdio_session.__aenter__() - await self._stdio_session.initialize() - return self._stdio_session + stdio_ctx = stdio_client(params) + self._stdio_ctx = stdio_ctx + read_stream, write_stream = await stdio_ctx.__aenter__() + stdio_session = ClientSession(read_stream, write_stream) + self._stdio_session = stdio_session + await stdio_session.__aenter__() + await stdio_session.initialize() + return stdio_session async def _get_http_session(self): """Create or return an MCP StreamableHTTP session.""" @@ -122,12 +126,14 @@ async def _get_http_session(self): from mcp.client.session import ClientSession from mcp.client.streamable_http import streamablehttp_client - self._http_ctx = streamablehttp_client(self.url) - read_stream, write_stream, _get_session_id = await self._http_ctx.__aenter__() - self._http_session = ClientSession(read_stream, write_stream) - await self._http_session.__aenter__() - await self._http_session.initialize() - return self._http_session + http_ctx = streamablehttp_client(self.url) + self._http_ctx = http_ctx + read_stream, write_stream, _get_session_id = await http_ctx.__aenter__() + http_session = ClientSession(read_stream, write_stream) + self._http_session = http_session + await http_session.__aenter__() + await http_session.initialize() + return http_session async def _send_jsonrpc(self, method: str, params: dict) -> Any: """Send a JSON-RPC request. @@ -196,7 +202,7 @@ async def list_tools(self) -> list[dict[str, Any]]: result = await self._send_jsonrpc("tools/list", {}) return result.get("tools", []) - except Exception as e: + except Exception: # Fallback: try REST endpoint if self.is_http: try: diff --git a/agent_codemode/proxy/meta_tools.py b/agent_codemode/proxy/meta_tools.py index 580502e..6c3beac 100644 --- a/agent_codemode/proxy/meta_tools.py +++ b/agent_codemode/proxy/meta_tools.py @@ -19,11 +19,13 @@ from __future__ import annotations import logging +from collections.abc import Awaitable +from inspect import isawaitable from typing import TYPE_CHECKING, Any, Callable, Optional if TYPE_CHECKING: - from ..discovery.registry import ToolRegistry from ..composition.executor import CodeModeExecutor + from ..discovery.registry import ToolRegistry from ..types import ToolDefinition @@ -31,7 +33,7 @@ # Type for AI tool selector function -AIToolSelector = Callable[[str, list[dict[str, Any]]], list[str]] +AIToolSelector = Callable[[str, list[dict[str, Any]]], list[str] | Awaitable[list[str]]] class MetaToolProvider: @@ -55,10 +57,10 @@ class MetaToolProvider: # Fast listing names = provider.list_tool_names(keywords=["file", "read"]) - + # AI-powered search tools = await provider.search_tools("read CSV files and analyze data") - + # Execute code result = await provider.execute_code(''' from generated.mcp.filesystem import read_file @@ -107,15 +109,13 @@ async def search_tools( """ # Get all tools all_tools = self.registry.list_tools(server=server, include_deferred=include_deferred) - + # Use AI selector if available if self._ai_selector and query: - tool_list = [ - {"name": t.name, "description": t.description} - for t in all_tools - ] + tool_list = [{"name": t.name, "description": t.description} for t in all_tools] try: - selected_names = await self._ai_selector(query, tool_list) + selected = self._ai_selector(query, tool_list) + selected_names = await selected if isawaitable(selected) else selected all_tools = [t for t in all_tools if t.name in selected_names] except Exception as e: logger.debug( @@ -125,7 +125,7 @@ async def search_tools( all_tools = self._keyword_filter(all_tools, query) elif query: all_tools = self._keyword_filter(all_tools, query) - + # Apply limit tools = all_tools[:limit] @@ -144,10 +144,8 @@ async def search_tools( for tool in tools ], } - - def _keyword_filter( - self, tools: list[ToolDefinition], query: str - ) -> list[ToolDefinition]: + + def _keyword_filter(self, tools: list[ToolDefinition], query: str) -> list[ToolDefinition]: """Filter tools by keywords in query. Args: @@ -159,7 +157,7 @@ def _keyword_filter( """ query_lower = query.lower() keywords = query_lower.split() - + scored_tools = [] for tool in tools: tool_text = f"{tool.name} {tool.description or ''}".lower() @@ -167,7 +165,7 @@ def _keyword_filter( score = sum(1 for kw in keywords if kw in tool_text) if score > 0: scored_tools.append((score, tool)) - + # Sort by score descending scored_tools.sort(key=lambda x: x[0], reverse=True) return [t for _, t in scored_tools] @@ -317,9 +315,7 @@ def get_meta_tools(self) -> list[dict[str, Any]]: }, { "name": "list_tool_names", - "description": ( - "List available tool names. Fast way to see what tools exist." - ), + "description": ("List available tool names. Fast way to see what tools exist."), "inputSchema": { "type": "object", "properties": { @@ -342,9 +338,7 @@ def get_meta_tools(self) -> list[dict[str, Any]]: }, { "name": "get_tool_definition", - "description": ( - "Get the full definition and schema of a specific tool." - ), + "description": ("Get the full definition and schema of a specific tool."), "inputSchema": { "type": "object", "properties": { @@ -360,39 +354,39 @@ def get_meta_tools(self) -> list[dict[str, Any]]: # Add execute_code if executor is available if self.executor is not None: - tools.append({ - "name": "execute_code", - "description": ( - "Execute Python code that composes MCP tools. " - "The code can import from generated tool bindings and " - "call multiple tools efficiently. Use this for complex " - "operations that require loops, conditionals, or state." - ), - "inputSchema": { - "type": "object", - "properties": { - "code": { - "type": "string", - "description": ( - "Python code to execute. Can use async/await " - "and import from generated.mcp.*" - ), - }, - "timeout": { - "type": "number", - "description": "Execution timeout in seconds (default: 60)", - "default": 60, + tools.append( + { + "name": "execute_code", + "description": ( + "Execute Python code that composes MCP tools. " + "The code can import from generated tool bindings and " + "call multiple tools efficiently. Use this for complex " + "operations that require loops, conditionals, or state." + ), + "inputSchema": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": ( + "Python code to execute. Can use async/await " + "and import from generated.mcp.*" + ), + }, + "timeout": { + "type": "number", + "description": "Execution timeout in seconds (default: 60)", + "default": 60, + }, }, + "required": ["code"], }, - "required": ["code"], - }, - }) + } + ) return tools - async def handle_tool_call( - self, tool_name: str, arguments: dict[str, Any] - ) -> dict[str, Any]: + async def handle_tool_call(self, tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]: """Handle a meta-tool call. Args: diff --git a/agent_codemode/server.py b/agent_codemode/server.py index 31cf131..156b194 100644 --- a/agent_codemode/server.py +++ b/agent_codemode/server.py @@ -21,7 +21,7 @@ import json import logging -from typing import Any, Optional +from typing import Any, Optional, cast import anyio import mcp.types as types @@ -47,15 +47,15 @@ def configure( registry: Optional[ToolRegistry] = None, ) -> None: """Configure the Codemode MCP server. - + Args: config: Configuration for the server. registry: Optional pre-configured tool registry. """ global _registry, _executor, _config - + _config = config or CodeModeConfig() - + if registry is not None: logger.debug(f"Using provided registry with {len(registry._servers)} servers") logger.debug(f"Server names: {list(registry._servers.keys())}") @@ -63,9 +63,9 @@ def configure( else: logger.debug("Creating new empty registry") _registry = ToolRegistry() - + _executor = CodeModeExecutor(_registry, _config) - + logger.info(f"Codemode MCP server configured with {len(_registry._servers)} servers") @@ -77,6 +77,8 @@ def get_registry() -> ToolRegistry: configure() else: logger.debug(f"Returning existing registry with {len(_registry._servers)} servers") + if _registry is None: + raise RuntimeError("Tool registry is not configured") return _registry @@ -85,6 +87,8 @@ def get_executor() -> CodeModeExecutor: global _executor if _executor is None: configure() + if _executor is None: + raise RuntimeError("Code executor is not configured") return _executor @@ -92,100 +96,114 @@ def get_executor() -> CodeModeExecutor: # Tool Definitions # ============================================================================= + def _build_tools() -> list[types.Tool]: """Build MCP tool definitions from shared schemas.""" from .tool_definitions import TOOL_SCHEMAS - + tools = [] - + # Core codemode tools - for name in ["search_tools", "list_tool_names", "list_servers", "get_tool_details", - "execute_code", "call_tool"]: + for name in [ + "search_tools", + "list_tool_names", + "list_servers", + "get_tool_details", + "execute_code", + "call_tool", + ]: schema = TOOL_SCHEMAS[name] - tools.append(types.Tool( - name=name, - description=schema["description"], - inputSchema=schema["parameters"], - )) - + description = cast(str | None, schema.get("description")) + parameters = cast(dict[str, Any], schema.get("parameters", {})) + tools.append( + types.Tool( + name=name, + description=description, + inputSchema=parameters, + ) + ) + # Skill management tools - tools.extend([ - types.Tool( - name="save_skill", - description="Save a reusable skill (code-based tool composition).", - inputSchema={ - "type": "object", - "required": ["name", "code", "description"], - "properties": { - "name": {"type": "string", "description": "Unique skill name"}, - "code": {"type": "string", "description": "Python code"}, - "description": {"type": "string", "description": "What the skill does"}, - "tags": {"type": "array", "items": {"type": "string"}}, - "parameters": {"type": "object"}, + tools.extend( + [ + types.Tool( + name="save_skill", + description="Save a reusable skill (code-based tool composition).", + inputSchema={ + "type": "object", + "required": ["name", "code", "description"], + "properties": { + "name": {"type": "string", "description": "Unique skill name"}, + "code": {"type": "string", "description": "Python code"}, + "description": {"type": "string", "description": "What the skill does"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "parameters": {"type": "object"}, + }, }, - }, - ), - types.Tool( - name="run_skill", - description="Execute a saved skill.", - inputSchema={ - "type": "object", - "required": ["name"], - "properties": { - "name": {"type": "string", "description": "Skill name"}, - "arguments": {"type": "object"}, + ), + types.Tool( + name="run_skill", + description="Execute a saved skill.", + inputSchema={ + "type": "object", + "required": ["name"], + "properties": { + "name": {"type": "string", "description": "Skill name"}, + "arguments": {"type": "object"}, + }, }, - }, - ), - types.Tool( - name="list_skills", - description="List available skills.", - inputSchema={ - "type": "object", - "properties": { - "tags": {"type": "array", "items": {"type": "string"}}, + ), + types.Tool( + name="list_skills", + description="List available skills.", + inputSchema={ + "type": "object", + "properties": { + "tags": {"type": "array", "items": {"type": "string"}}, + }, }, - }, - ), - types.Tool( - name="delete_skill", - description="Delete a saved skill.", - inputSchema={ - "type": "object", - "required": ["name"], - "properties": { - "name": {"type": "string"}, + ), + types.Tool( + name="delete_skill", + description="Delete a saved skill.", + inputSchema={ + "type": "object", + "required": ["name"], + "properties": { + "name": {"type": "string"}, + }, }, - }, - ), - types.Tool( - name="get_execution_history", - description="Get recent tool execution history.", - inputSchema={ - "type": "object", - "properties": { - "limit": {"type": "integer", "default": 10}, + ), + types.Tool( + name="get_execution_history", + description="Get recent tool execution history.", + inputSchema={ + "type": "object", + "properties": { + "limit": {"type": "integer", "default": 10}, + }, }, - }, - ), - types.Tool( - name="add_mcp_server", - description="Add a new MCP server to discover tools from.", - inputSchema={ - "type": "object", - "required": ["name"], - "properties": { - "name": {"type": "string"}, - "url": {"type": "string"}, - "command": {"type": "string"}, - "args": {"type": "array", "items": {"type": "string"}}, + ), + types.Tool( + name="add_mcp_server", + description="Add a new MCP server to discover tools from.", + inputSchema={ + "type": "object", + "required": ["name"], + "properties": { + "name": {"type": "string"}, + "url": {"type": "string"}, + "command": {"type": "string"}, + "args": {"type": "array", "items": {"type": "string"}}, + }, }, - }, - ), - ]) - + ), + ] + ) + return tools + TOOLS = _build_tools() @@ -193,6 +211,7 @@ def _build_tools() -> list[types.Tool]: # Tool Handlers # ============================================================================= + async def handle_search_tools(arguments: dict[str, Any]) -> dict[str, Any]: """Search for available tools matching a query.""" query = arguments["query"] @@ -205,12 +224,12 @@ async def handle_search_tools(arguments: dict[str, Any]) -> dict[str, Any]: result = await registry.search_tools( query, server=server, limit=limit, include_deferred=include_deferred ) - + # Filter by category if specified tools = result.tools if category: tools = [t for t in tools if category.lower() in (t.description or "").lower()] - + return { "tools": [ { @@ -233,7 +252,7 @@ async def handle_list_servers(arguments: dict[str, Any]) -> dict[str, Any]: """List all connected MCP servers.""" registry = get_registry() servers = await registry.list_servers() - + return { "servers": [ { @@ -262,7 +281,7 @@ async def handle_list_tool_names(arguments: dict[str, Any]) -> dict[str, Any]: include_deferred=include_deferred, ) total = len(registry.list_tools(server=server, include_deferred=include_deferred)) - + return { "tool_names": names, "returned": len(names), @@ -276,10 +295,10 @@ async def handle_get_tool_details(arguments: dict[str, Any]) -> dict[str, Any]: tool_name = arguments["tool_name"] registry = get_registry() tool = registry.get_tool(tool_name) - + if tool is None: return {"error": f"Tool not found: {tool_name}"} - + return { "name": tool.name, "description": tool.description, @@ -298,16 +317,16 @@ async def handle_execute_code(arguments: dict[str, Any]) -> dict[str, Any]: context = arguments.get("context") executor = get_executor() - + # Ensure executor is set up if not executor._setup_done: await executor.setup() - + # Inject context variables if provided if context and executor._sandbox: for name, value in context.items(): executor._sandbox.set_variable(name, value) - + try: execution = await executor.execute(code, timeout=timeout) error_message = ( @@ -319,7 +338,7 @@ async def handle_execute_code(arguments: dict[str, Any]) -> dict[str, Any]: if execution.code_error else None ) - + return { "success": execution.success, "execution_ok": execution.execution_ok, @@ -346,7 +365,7 @@ async def handle_call_tool(arguments: dict[str, Any]) -> dict[str, Any]: tool_arguments = arguments["arguments"] executor = get_executor() - + try: result = await executor.call_tool(tool_name, tool_arguments) return { @@ -365,8 +384,8 @@ async def handle_call_tool(arguments: dict[str, Any]) -> dict[str, Any]: async def handle_save_skill(arguments: dict[str, Any]) -> dict[str, Any]: """Save a reusable skill (code-based tool composition).""" - from agent_skills import SkillsManager - + from agent_skills import SkillsManager # type: ignore[import-untyped] + name = arguments["name"] code = arguments["code"] description = arguments["description"] @@ -374,9 +393,11 @@ async def handle_save_skill(arguments: dict[str, Any]) -> dict[str, Any]: config = _config or CodeModeConfig() manager = SkillsManager(config.skills_path) - + try: - manager.create(name=name, description=description, content=code, python_code=code, tags=tags) + manager.create( + name=name, description=description, content=code, python_code=code, tags=tags + ) return { "success": True, "skill_id": name, @@ -396,7 +417,7 @@ async def handle_run_skill(arguments: dict[str, Any]) -> dict[str, Any]: skill_arguments = arguments.get("arguments") executor = get_executor() - + try: execution = await executor.execute_skill(name, skill_arguments) error_message = ( @@ -408,7 +429,7 @@ async def handle_run_skill(arguments: dict[str, Any]) -> dict[str, Any]: if execution.code_error else None ) - + return { "success": execution.success, "execution_ok": execution.execution_ok, @@ -434,14 +455,14 @@ async def handle_run_skill(arguments: dict[str, Any]) -> dict[str, Any]: async def handle_list_skills(arguments: dict[str, Any]) -> dict[str, Any]: """List available skills.""" from agent_skills import SkillsManager - + tags = arguments.get("tags") config = _config or CodeModeConfig() manager = SkillsManager(config.skills_path) - + skills = manager.list(tags=tags) - + return { "skills": [ { @@ -458,15 +479,15 @@ async def handle_list_skills(arguments: dict[str, Any]) -> dict[str, Any]: async def handle_delete_skill(arguments: dict[str, Any]) -> dict[str, Any]: """Delete a saved skill.""" from agent_skills import SkillsManager - + name = arguments["name"] config = _config or CodeModeConfig() manager = SkillsManager(config.skills_path) - + skill = manager.get(name) success = manager.delete(skill.skill_id) if skill and skill.skill_id else False - + return { "success": success, "error": None if success else f"Skill not found: {name}", @@ -479,7 +500,7 @@ async def handle_get_execution_history(arguments: dict[str, Any]) -> dict[str, A executor = get_executor() history = executor.tool_call_history[-limit:] - + return { "history": [ { @@ -496,15 +517,15 @@ async def handle_get_execution_history(arguments: dict[str, Any]) -> dict[str, A async def handle_add_mcp_server(arguments: dict[str, Any]) -> dict[str, Any]: """Add a new MCP server to discover tools from.""" - from .models import MCPServerConfig - + from .models import MCPServerConfig # type: ignore[import-untyped] + name = arguments["name"] url = arguments.get("url") command = arguments.get("command") args = arguments.get("args", []) registry = get_registry() - + if url: config = MCPServerConfig( name=name, @@ -522,13 +543,13 @@ async def handle_add_mcp_server(arguments: dict[str, Any]) -> dict[str, Any]: "tools_discovered": 0, "error": "Either url or command must be provided", } - + try: registry.add_server(config) await registry.discover_server(name) - + tools = registry.list_tools(server=name) - + return { "success": True, "tools_discovered": len(tools), @@ -566,6 +587,7 @@ async def handle_add_mcp_server(arguments: dict[str, Any]) -> dict[str, Any]: # MCP Server Handlers # ============================================================================= + @mcp.list_tools() async def list_tools() -> list[types.Tool]: """Return the list of available tools.""" @@ -581,7 +603,7 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[types.TextCont handler = TOOL_HANDLERS.get(name) if handler is None: raise ValueError(f"Unknown tool: {name}") - + result = await handler(arguments) json_str = json.dumps(result, indent=2) return [types.TextContent(type="text", text=json_str)] @@ -591,9 +613,10 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[types.TextCont # Server Entry Points # ============================================================================= + def run(transport: str = "stdio", host: str = "127.0.0.1", port: int = 8000) -> None: """Run the MCP server. - + Args: transport: Transport type - "stdio" or "streamable-http" (default: "stdio"). host: Host for HTTP transport (default: "127.0.0.1"). @@ -603,17 +626,18 @@ def run(transport: str = "stdio", host: str = "127.0.0.1", port: int = 8000) -> # Ensure registry is initialized (will use existing if already configured) if _registry is None: configure() - + if transport == "streamable-http": + import uvicorn from mcp.server.streamable_http import StreamableHTTPServerTransport from starlette.applications import Starlette from starlette.routing import Route - import uvicorn async def handle_mcp(request): - async with StreamableHTTPServerTransport( + transport_ctx: Any = StreamableHTTPServerTransport( "/mcp", request.scope, request.receive, request._send - ) as transport: + ) + async with transport_ctx as transport: await mcp.run( transport.read_stream, transport.write_stream, @@ -641,22 +665,20 @@ async def arun(): logger.info(f"Upfront discovery complete: {len(registry.list_tools())} tools") except Exception as e: logger.warning(f"Upfront discovery failed: {e}") - + async with stdio_server() as streams: - await mcp.run( - streams[0], streams[1], mcp.create_initialization_options() - ) + await mcp.run(streams[0], streams[1], mcp.create_initialization_options()) anyio.run(arun) if __name__ == "__main__": import sys - + transport = "stdio" host = "127.0.0.1" port = 8000 - + # Simple CLI argument parsing args = sys.argv[1:] for i, arg in enumerate(args): @@ -666,5 +688,5 @@ async def arun(): host = args[i + 1] elif arg == "--port" and i + 1 < len(args): port = int(args[i + 1]) - + run(transport=transport, host=host, port=port) diff --git a/agent_codemode/tool_definitions.py b/agent_codemode/tool_definitions.py index edac47a..f0d700b 100644 --- a/agent_codemode/tool_definitions.py +++ b/agent_codemode/tool_definitions.py @@ -141,13 +141,13 @@ def get_tool_schema(tool_name: str) -> dict[str, Any]: """Get the schema for a specific tool. - + Args: tool_name: Name of the tool - + Returns: Dictionary with 'description' and 'parameters' keys - + Raises: KeyError: If tool not found """ diff --git a/agent_codemode/toolset.py b/agent_codemode/toolset.py index d407eef..9dd1690 100644 --- a/agent_codemode/toolset.py +++ b/agent_codemode/toolset.py @@ -16,15 +16,15 @@ Example: from pydantic_ai import Agent from agent_codemode import CodemodeToolset, ToolRegistry - + # Set up registry registry = ToolRegistry() registry.add_server(MCPServerConfig(name="bash", url="...")) await registry.discover_all() - + # Create toolset toolset = CodemodeToolset(registry=registry) - + # Use with agent agent = Agent( model='openai:gpt-4o', @@ -36,14 +36,15 @@ import logging import time +from collections.abc import Awaitable from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional +from typing import TYPE_CHECKING, Any, Callable, Optional # Lazy imports to avoid circular dependencies # ToolRegistry and CodeModeExecutor are imported at runtime in methods if TYPE_CHECKING: - from .discovery.registry import ToolRegistry from .composition.executor import CodeModeExecutor + from .discovery.registry import ToolRegistry from .types import CodeModeConfig @@ -52,52 +53,50 @@ # Check if pydantic-ai is available try: + from pydantic_ai._run_context import RunContext + from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets import AbstractToolset from pydantic_ai.toolsets.abstract import ToolsetTool - from pydantic_ai.tools import ToolDefinition - from pydantic_ai._run_context import RunContext from pydantic_core import SchemaValidator, core_schema - + PYDANTIC_AI_AVAILABLE = True except ImportError: PYDANTIC_AI_AVAILABLE = False - AbstractToolset = object if PYDANTIC_AI_AVAILABLE: - # Schema validator for any args CODEMODE_ARGS_VALIDATOR = SchemaValidator(schema=core_schema.any_schema()) - + @dataclass class CodemodeToolset(AbstractToolset): """Codemode toolset for pydantic-ai with method-based tool execution. - + This provides the same tools as the MCP server but via direct method calls, which is more efficient for in-process agent usage. - + Provides: - search_tools: Find relevant tools by query - get_tool_details: Get full tool definition - list_servers: List connected MCP servers - execute_code: Run Python code that composes tools - call_tool: Call a single tool directly - + Example: from agent_codemode import CodemodeToolset, ToolRegistry from pydantic_ai import Agent - + registry = ToolRegistry() # ... configure registry ... - + toolset = CodemodeToolset(registry=registry) - + agent = Agent( model='openai:gpt-4o', toolsets=[toolset], ) """ - + registry: ToolRegistry | None = None config: CodeModeConfig = field(default_factory=CodeModeConfig) sandbox: Any | None = None # Optional pre-configured sandbox (e.g., EvalSandbox) @@ -108,16 +107,22 @@ class CodemodeToolset(AbstractToolset): # Used by agent-runtimes to push websocket status updates without polling. status_change_callback: Callable[[bool], Awaitable[None]] | None = None _id: str | None = None - + # Internal state _executor: CodeModeExecutor | None = field(default=None, repr=False) _initialized: bool = field(default=False, repr=False) _codemode_call_count: int = field(default=0, repr=False) _runtime_is_executing: bool = field(default=False, repr=False) - _post_init_callbacks: list[Callable[["CodemodeToolset"], None]] = field( - default_factory=list, repr=False, + _post_init_callbacks: list[Callable[[CodemodeToolset], None]] = field( + default_factory=list, + repr=False, ) + def _get_registry(self) -> ToolRegistry: + if self.registry is None: + raise RuntimeError("Tool registry is not initialized") + return self.registry + @property def runtime_is_executing(self) -> bool: """Execution state tracked by toolset lifecycle hooks.""" @@ -135,19 +140,20 @@ async def _set_runtime_executing(self, value: bool) -> None: "Codemode status_change_callback failed: %s", exc, ) - + def __post_init__(self): if self.registry is None: # Import at runtime to avoid circular dependency from .discovery.registry import ToolRegistry + self.registry = ToolRegistry() # Default the direct-call policy from config if not provided if self.allow_direct_tool_calls is None: self.allow_direct_tool_calls = self.config.allow_direct_tool_calls - + def add_post_init_callback( self, - callback: Callable[["CodemodeToolset"], None], + callback: Callable[[CodemodeToolset], None], ) -> None: """Register a callback to run after executor initialisation. @@ -156,31 +162,33 @@ def add_post_init_callback( wire skill bindings because the executor already exists. """ self._post_init_callbacks.append(callback) - + @property def id(self) -> str | None: return self._id - + @property def label(self) -> str: return "Codemode Toolset" - + async def _ensure_initialized(self) -> None: """Initialize the executor if not already done.""" if self._initialized: return - + if self._executor is None: # Import at runtime to avoid circular dependency from .composition.executor import CodeModeExecutor + # Ensure tools are discovered before generating bindings - if self.registry is not None and not self.registry.list_tools(): + registry = self._get_registry() + if not registry.list_tools(): logger.info("Codemode registry empty; discovering tools...") - await self.registry.discover_all() - tool_count = len(self.registry.list_tools()) if self.registry is not None else 0 + await registry.discover_all() + tool_count = len(registry.list_tools()) logger.info("Codemode registry tool count: %s", tool_count) self._executor = CodeModeExecutor( - registry=self.registry, + registry=registry, config=self.config, sandbox=self.sandbox, ) @@ -189,7 +197,7 @@ async def _ensure_initialized(self) -> None: "Codemode executor setup complete (generated_path=%s)", self.config.generated_path, ) - + self._initialized = True # Run any registered post-init callbacks (e.g. skill wiring) @@ -202,7 +210,7 @@ async def _ensure_initialized(self) -> None: async def start(self) -> None: """Start the toolset and executor.""" await self._ensure_initialized() - + async def cleanup(self) -> None: """Clean up resources.""" if self._executor: @@ -213,12 +221,17 @@ async def cleanup(self) -> None: async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]: """Get the tools provided by this toolset.""" from .tool_definitions import get_tool_schema - + tools = {} - + if self.allow_discovery_tools: # Discovery tools - for tool_name in ["list_tool_names", "search_tools", "get_tool_details", "list_servers"]: + for tool_name in [ + "list_tool_names", + "search_tools", + "get_tool_details", + "list_servers", + ]: schema = get_tool_schema(tool_name) tools[tool_name] = ToolsetTool( toolset=self, @@ -230,7 +243,7 @@ async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]: max_retries=0, args_validator=CODEMODE_ARGS_VALIDATOR, ) - + # execute_code - always available schema = get_tool_schema("execute_code") tools["execute_code"] = ToolsetTool( @@ -243,7 +256,7 @@ async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]: max_retries=0, args_validator=CODEMODE_ARGS_VALIDATOR, ) - + # call_tool (optional) if self.allow_direct_tool_calls: schema = get_tool_schema("call_tool") @@ -257,9 +270,9 @@ async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]: max_retries=1, args_validator=CODEMODE_ARGS_VALIDATOR, ) - + return tools - + async def call_tool( self, name: str, @@ -270,7 +283,7 @@ async def call_tool( """Call a tool by name.""" await self._ensure_initialized() self._codemode_call_count += 1 - + if name == "list_tool_names": return await self._list_tool_names( server=tool_args.get("server"), @@ -303,7 +316,7 @@ async def call_tool( ) else: raise ValueError(f"Unknown tool: {name}") - + async def _list_tool_names( self, server: Optional[str] = None, @@ -312,7 +325,8 @@ async def _list_tool_names( include_deferred: bool = False, ) -> dict[str, Any]: """List all tool names quickly without descriptions.""" - tools = self.registry.list_tools(server=server, include_deferred=include_deferred) + registry = self._get_registry() + tools = registry.list_tools(server=server, include_deferred=include_deferred) total_available = len(tools) if keywords: lowered = [kw.lower() for kw in keywords] @@ -325,7 +339,7 @@ async def _list_tool_names( total_available = len(filtered) if limit: tools = tools[:limit] - + # Group by server for better organization with import hints by_server: dict[str, list[str]] = {} import_hints: dict[str, str] = {} @@ -336,11 +350,13 @@ async def _list_tool_names( # Convert tool name to function name (replace dashes with underscores) func_name = t.name.split("__")[-1].replace("-", "_") by_server[server_name].append(func_name) - + # Generate import hints for each server for server_name, funcs in by_server.items(): - import_hints[server_name] = f"from generated.mcp.{server_name} import {', '.join(funcs)}" - + import_hints[server_name] = ( + f"from generated.mcp.{server_name} import {', '.join(funcs)}" + ) + # Add skills import hint if skills are available executor = getattr(self, "_executor", None) if executor and getattr(executor, "_skills_metadata", None): @@ -359,7 +375,7 @@ async def _list_tool_names( "include_deferred": include_deferred, "usage": "Use import_hints to get the correct import statement for execute_code", } - + async def _search_tools( self, query: str, @@ -368,7 +384,8 @@ async def _search_tools( include_deferred: bool = True, ) -> dict[str, Any]: """Search for tools matching a query.""" - result = await self.registry.search_tools( + registry = self._get_registry() + result = await registry.search_tools( query, server=server, limit=limit, include_deferred=include_deferred ) tools = result.tools @@ -383,7 +400,7 @@ async def _search_tools( if isawaitable(reranked): tools = await reranked else: - tools = reranked # type: ignore[assignment] + tools = reranked after = [t.name for t in tools] logger.debug( "Applied tool reranker: before=%s, after=%s", @@ -395,7 +412,7 @@ async def _search_tools( "Tool reranker failed; falling back to registry order: %s", e, ) - + return { "tools": [ { @@ -412,14 +429,15 @@ async def _search_tools( "total": result.total, "has_more": result.total > limit, } - + async def _get_tool_details(self, tool_name: str) -> dict[str, Any]: """Get detailed information about a tool.""" - tool = self.registry.get_tool(tool_name) - + registry = self._get_registry() + tool = registry.get_tool(tool_name) + if tool is None: return {"error": f"Tool not found: {tool_name}"} - + return { "name": tool.name, "description": tool.description, @@ -429,21 +447,22 @@ async def _get_tool_details(self, tool_name: str) -> dict[str, Any]: "input_examples": tool.input_examples, "defer_loading": tool.defer_loading, } - + async def _list_servers(self) -> dict[str, Any]: """List all connected MCP servers with import paths.""" - servers = await self.registry.list_servers() - + registry = self._get_registry() + servers = await registry.list_servers() + # Get tools for each server to provide function names server_tools: dict[str, list[str]] = {} - for tool in self.registry.list_tools(include_deferred=True): + for tool in registry.list_tools(include_deferred=True): sname = tool.server_name or "unknown" if sname not in server_tools: server_tools[sname] = [] # Convert tool name to function name (replace dashes with underscores) func_name = tool.name.split("__")[-1].replace("-", "_") server_tools[sname].append(func_name) - + server_entries = [ { "name": s.name, @@ -467,21 +486,23 @@ async def _list_servers(self) -> dict[str, Any]: skill_scripts_summary.append( f"{skill_meta['name']}/{script['name']}: {script.get('description', '')}" ) - server_entries.append({ - "name": "skills", - "description": "Agent skills \u2013 reusable task scripts", - "tool_count": len(skill_funcs), - "import_path": f"from generated.skills import {', '.join(skill_funcs)}", - "functions": skill_funcs, - "available_scripts": skill_scripts_summary, - }) + server_entries.append( + { + "name": "skills", + "description": "Agent skills \u2013 reusable task scripts", + "tool_count": len(skill_funcs), + "import_path": f"from generated.skills import {', '.join(skill_funcs)}", + "functions": skill_funcs, + "available_scripts": skill_scripts_summary, + } + ) return { "servers": server_entries, "total": len(server_entries), "usage_hint": "Use the import_path to import tools in execute_code", } - + async def _execute_code( self, code: str, @@ -490,14 +511,18 @@ async def _execute_code( """Execute Python code that composes tools.""" if self._executor is None: return {"success": False, "error": "Executor not initialized"} - + try: await self._set_runtime_executing(True) start_time = time.monotonic() # Log full code for debugging (truncate only for single-line display) - code_lines = (code or "").strip().split('\n') + code_lines = (code or "").strip().split("\n") if len(code_lines) > 1: - logger.info("Codemode execute_code: calling executor.execute() with %d lines:\n%s", len(code_lines), code) + logger.info( + "Codemode execute_code: calling executor.execute() with %d lines:\n%s", + len(code_lines), + code, + ) else: logger.info("Codemode execute_code: calling executor.execute() code=%r", code) execution = await self._executor.execute(code, timeout=timeout) @@ -518,7 +543,7 @@ async def _execute_code( else: logger.info(log_message) elapsed = time.monotonic() - start_time - + error_message = ( execution.execution_error if not execution.execution_ok @@ -562,7 +587,7 @@ async def _execute_code( } finally: await self._set_runtime_executing(False) - + async def _call_tool( self, tool_name: str, @@ -570,7 +595,8 @@ async def _call_tool( ) -> dict[str, Any]: """Call a single tool directly.""" try: - result = await self.registry.call_tool(tool_name, arguments) + registry = self._get_registry() + result = await registry.call_tool(tool_name, arguments) return { "success": True, "result": result, @@ -596,7 +622,7 @@ def get_call_counts(self) -> dict[str, int]: # Fallback when pydantic-ai is not available class CodemodeToolset: # type: ignore """Placeholder when pydantic-ai is not installed.""" - + def __init__(self, *args, **kwargs): raise ImportError( "pydantic-ai is required for CodemodeToolset. " diff --git a/agent_codemode/types.py b/agent_codemode/types.py index c6cbc56..d6d0fea 100644 --- a/agent_codemode/types.py +++ b/agent_codemode/types.py @@ -62,7 +62,7 @@ class ToolDefinition(BaseModel): server_url: str = "" @model_validator(mode="after") - def _unwrap_fastmcp_output_schema(self) -> "ToolDefinition": + def _unwrap_fastmcp_output_schema(self) -> ToolDefinition: """Unwrap FastMCP's artificial ``x-fastmcp-wrap-result`` wrapper. FastMCP wraps simple return types (e.g. ``str``) in an object schema @@ -72,11 +72,7 @@ def _unwrap_fastmcp_output_schema(self) -> "ToolDefinition": inner type). Unwrap it so every consumer sees the real schema. """ schema = self.output_schema - if ( - schema - and schema.get("x-fastmcp-wrap-result") - and schema.get("type") == "object" - ): + if schema and schema.get("x-fastmcp-wrap-result") and schema.get("type") == "object": props = schema.get("properties", {}) if "result" in props: # Replace with the inner type schema @@ -198,8 +194,14 @@ class CodeModeConfig(BaseModel): workspace_path: Path for workspace files. skills_path: Path for saved skills. generated_path: Path for generated code bindings. - sandbox_variant: Which sandbox to use for execution. + sandbox_variant: Which sandbox to use for execution. One of the + code-sandboxes variants: ``eval`` (in-process, default), + ``monty`` (secure in-process interpreter), ``docker``, + ``jupyter``, ``colab``, ``kaggle``, ``modal``, or ``datalayer``. sandbox_image: Optional sandbox image (for Docker-based sandboxes). + sandbox_gpu: Optional GPU flavor / accelerator for supported sandboxes + (for example Modal/Datalayer: ``T4``, ``A100``; Kaggle batch: + ``NvidiaTeslaT4`` or alias ``T4``). allow_direct_tool_calls: Whether to expose call_tool in the toolset. max_tool_calls: Optional safety cap for tool calls per execute() run. skills_directories: Directories to scan for skill definitions. @@ -211,9 +213,9 @@ class CodeModeConfig(BaseModel): HTTP to this URL instead of trying to use stdio. This enables the two-container architecture where MCP servers run in the agent-runtimes container and code executes in a Jupyter container. - + Example: "http://agent-runtimes:8765/api/v1/mcp/proxy" - + For local development with local Jupyter: "http://localhost:8765/api/v1/mcp/proxy" """ @@ -225,6 +227,7 @@ class CodeModeConfig(BaseModel): generated_path: str = "./generated" sandbox_variant: str = "eval" sandbox_image: str | None = None + sandbox_gpu: str | None = None allow_direct_tool_calls: bool = False max_tool_calls: int | None = None skills_directories: list[str] = Field(default_factory=list) diff --git a/docs/Makefile b/docs/Makefile index 3370dc1..3a992cc 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -28,7 +28,7 @@ env-rm: env: -conda env create -f environment.yml - @echo + @echo @echo -------------------------------------------------- @echo ✨ Datalayer environment is created. @echo -------------------------------------------------- diff --git a/docs/README.md b/docs/README.md index 005d7e8..694c0a2 100644 --- a/docs/README.md +++ b/docs/README.md @@ -18,7 +18,6 @@ make pydoc make typedoc ``` - ```bash # Local Development: This command starts a local development server and opens up a browser window. # Most changes are reflected live without having to restart the server. @@ -27,7 +26,7 @@ make start ``` ```bash -# Build: This command generates static content into the `build` directory +# Build: This command generates static content into the `build` directory # and can be served using any static contents hosting service. make build ``` diff --git a/docs/docs/concept/index.mdx b/docs/docs/concept/index.mdx index 89a4195..19f8092 100644 --- a/docs/docs/concept/index.mdx +++ b/docs/docs/concept/index.mdx @@ -78,10 +78,10 @@ Agent Codemode generates Python code that wraps MCP tools: # Generated: generated/mcp/filesystem.py async def read_file(arguments: dict) -> str: """Read file contents. - + Args: path: File path to read - + Returns: File contents as string """ diff --git a/docs/docs/getting-started/index.mdx b/docs/docs/getting-started/index.mdx index 7607cb5..0fda714 100644 --- a/docs/docs/getting-started/index.mdx +++ b/docs/docs/getting-started/index.mdx @@ -81,16 +81,16 @@ Now you can run code that composes MCP tools: ```python result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file - + # Read a file content = await read_file({"path": "/tmp/input.txt"}) - + # Process it processed = content.upper() - + # Write the result await write_file({"path": "/tmp/output.txt", "content": processed}) - + print(f"Processed {len(content)} characters") """) @@ -116,30 +116,30 @@ async def main(): args=["-y", "@anthropic/mcp-server-filesystem", "/workspace"] )) await registry.discover_all() - + # Set up executor executor = CodeModeExecutor(registry) await executor.setup() - + # Execute code result = await executor.execute(""" from generated.mcp.filesystem import list_directory, read_file import asyncio - + # List files entries = await list_directory({"path": "/"}) print(f"Found {len(entries['entries'])} files") - + # Read all text files in parallel txt_files = [e for e in entries['entries'] if e.endswith('.txt')] contents = await asyncio.gather(*[ read_file({"path": f"/{f}"}) for f in txt_files ]) - + total_chars = sum(len(c) for c in contents) print(f"Total characters: {total_chars}") """) - + if result.success: print("Execution successful!") print(result.output) diff --git a/docs/docs/index.mdx b/docs/docs/index.mdx index ba21632..47f87e2 100644 --- a/docs/docs/index.mdx +++ b/docs/docs/index.mdx @@ -200,7 +200,7 @@ Use reusable code patterns that compose multiple tools: # skills/batch_process.py async def batch_process(input_dir: str, output_dir: str) -> dict: from generated.mcp.filesystem import list_directory, read_file, write_file - + entries = await list_directory({"path": input_dir}) for entry in entries["entries"]: content = await read_file({"path": f"{input_dir}/{entry}"}) @@ -286,7 +286,7 @@ await executor.setup() result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file - + content = await read_file({"path": "/tmp/input.txt"}) await write_file({"path": "/tmp/output.txt", "content": content.upper()}) """) diff --git a/docs/docs/integrations/index.mdx b/docs/docs/integrations/index.mdx index 16b92f9..6e52ca5 100644 --- a/docs/docs/integrations/index.mdx +++ b/docs/docs/integrations/index.mdx @@ -48,8 +48,7 @@ from agent_codemode import CodemodeToolset, CodeModeConfig config = CodeModeConfig( workspace_path="/workspace", generated_path="/tmp/generated", - sandbox_variant="local", - default_timeout=60.0, + sandbox_variant="eval", ) toolset = CodemodeToolset( @@ -82,7 +81,7 @@ Execute Python code that uses MCP tools: # Agent writes and executes code result = execute_code(code=""" from generated.mcp.filesystem import read_file, list_directory - + entries = await list_directory({"path": "/workspace"}) for entry in entries["entries"]: if entry.endswith(".py"): @@ -106,7 +105,7 @@ from agent_codemode import CodemodeToolset agent = Agent( "anthropic:claude-sonnet-4-20250514", - system_prompt="""You are a code-first assistant. + system_prompt="""You are a code-first assistant. When given a task: 1. Use search_tools to find relevant MCP tools @@ -262,7 +261,7 @@ Use Code Mode for complex, multi-step operations: # Good for Code Mode: Complex multi-tool workflow "Read all Python files, analyze their imports, generate dependency graph" -# Use direct tools: Simple single operations +# Use direct tools: Simple single operations "Read the file /config.json" ``` @@ -277,7 +276,7 @@ agent = Agent( You have two modes: - Code Mode: For complex operations involving multiple files or steps - Direct Tools: For simple single operations - + Use search_tools first to understand available capabilities. When writing code, always include error handling. """ @@ -310,8 +309,10 @@ Configure sandbox security: from agent_codemode import CodeModeConfig config = CodeModeConfig( - sandbox_variant="datalayer-runtime", # Isolated cloud sandbox - default_timeout=30.0, # Limit execution time + sandbox_variant="datalayer", # Isolated cloud sandbox workspace_path="/safe/workspace", # Restricted workspace ) + +# Set timeout per execution call +result = await executor.execute(code, timeout=30.0) ``` diff --git a/docs/docs/mcp-server/index.mdx b/docs/docs/mcp-server/index.mdx index fcccbad..c098f4a 100644 --- a/docs/docs/mcp-server/index.mdx +++ b/docs/docs/mcp-server/index.mdx @@ -136,7 +136,7 @@ result = await client.call_tool("save_skill", { "code": ''' async def batch_uppercase(input_dir, output_dir): from generated.mcp.filesystem import list_directory, read_file, write_file - + entries = await list_directory({"path": input_dir}) for entry in entries["entries"]: content = await read_file({"path": f"{input_dir}/{entry}"}) @@ -256,8 +256,8 @@ Build an agent that uses Codemode to intelligently compose tools: ``` User: "Analyze all Python files in /project and create a summary report" -Agent: -1. Calls search_tools("python file analysis") +Agent: +1. Calls search_tools("python file analysis") 2. Discovers filesystem and analysis tools 3. Writes code that reads files, analyzes them, generates report 4. Calls execute_code with the generated code diff --git a/docs/docs/mcp-servers/index.mdx b/docs/docs/mcp-servers/index.mdx index 27e71ae..ab71640 100644 --- a/docs/docs/mcp-servers/index.mdx +++ b/docs/docs/mcp-servers/index.mdx @@ -106,10 +106,10 @@ Agent Codemode generates Python bindings in the `generated/mcp/` directory: # Generated: generated/mcp/filesystem.py async def read_file(arguments: dict) -> str: """Read the complete contents of a file. - + Args: path: Path to the file to read - + Returns: File contents as string """ @@ -117,11 +117,11 @@ async def read_file(arguments: dict) -> str: async def write_file(arguments: dict) -> dict: """Write content to a file. - + Args: path: Path to the file to write content: Content to write - + Returns: Success status """ @@ -282,17 +282,17 @@ async with CodeModeExecutor(registry) as executor: result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file from generated.mcp.github import create_issue - + # Read local file content = await read_file({"path": "/workspace/bug-report.txt"}) - + # Create GitHub issue with the content issue = await create_issue({ "repo": "myorg/myrepo", "title": "Bug Report", "body": content }) - + print(f"Created issue: {issue['url']}") """) ``` diff --git a/docs/docs/programmatic-tools/index.mdx b/docs/docs/programmatic-tools/index.mdx index 3dff9a1..a1082f8 100644 --- a/docs/docs/programmatic-tools/index.mdx +++ b/docs/docs/programmatic-tools/index.mdx @@ -34,7 +34,7 @@ await executor.setup() # Execute code result = await executor.execute(""" from generated.mcp.filesystem import read_file - + content = await read_file({"path": "/tmp/data.txt"}) print(f"Read {len(content)} characters") """) @@ -110,7 +110,7 @@ Call tools one after another: ```python result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file - + # Read, transform, write content = await read_file({"path": "/input.txt"}) transformed = content.upper() @@ -152,14 +152,14 @@ Use `asyncio.gather` for concurrent tool calls: result = await executor.execute(""" import asyncio from generated.mcp.filesystem import read_file - + files = ["a.txt", "b.txt", "c.txt", "d.txt"] - + # Read all files in parallel contents = await asyncio.gather(*[ read_file({"path": f"/data/{f}"}) for f in files ]) - + total = sum(len(c) for c in contents) print(f"Total: {total} characters from {len(files)} files") """) @@ -173,7 +173,7 @@ Use if/else for decision making: result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file from generated.mcp.web import fetch_url - + # Try local file first, fall back to web try: content = await read_file({"path": "/cache/data.json"}) @@ -192,9 +192,9 @@ Process collections of items: ```python result = await executor.execute(""" from generated.mcp.filesystem import list_directory, read_file - + entries = await list_directory({"path": "/documents"}) - + for entry in entries["entries"]: if entry.endswith(".md"): content = await read_file({"path": f"/documents/{entry}"}) @@ -210,17 +210,17 @@ Use try/except for robust execution: ```python result = await executor.execute(""" from generated.mcp.filesystem import read_file, write_file - + files = ["important.txt", "optional.txt", "extra.txt"] results = [] - + for f in files: try: content = await read_file({"path": f"/data/{f}"}) results.append({"file": f, "status": "success", "size": len(content)}) except Exception as e: results.append({"file": f, "status": "failed", "error": str(e)}) - + successful = [r for r in results if r["status"] == "success"] print(f"Processed {len(successful)}/{len(files)} files") """) @@ -243,19 +243,27 @@ This means you get the safety and isolation of MCP servers while writing natural from agent_codemode import CodeModeConfig, CodeModeExecutor config = CodeModeConfig( - sandbox_variant="local", # "local" or "datalayer-runtime" + sandbox_variant="eval", # eval | monty | docker | jupyter | colab | kaggle | modal | datalayer workspace_path="/workspace", # Working directory generated_path="/tmp/generated", # Where to put generated bindings - default_timeout=30.0, # Default execution timeout ) executor = CodeModeExecutor(registry, config=config) + +# Set timeout per execution +result = await executor.execute(code, timeout=30.0) ``` ### Sandbox Variants -- **local**: Runs in a subprocess on the same machine -- **datalayer-runtime**: Runs in a cloud-based Datalayer Runtime +- **eval**: In-process Python execution on the same machine (default) +- **monty**: Secure in-process interpreter (pydantic-monty) +- **docker**: Docker container with an isolated kernel +- **jupyter**: Jupyter Server with a persistent kernel +- **colab**: Google Colab runtime (existing running kernel reuse only) +- **kaggle**: Kaggle runtime (API-token kernel creation or existing kernel session) +- **modal**: Modal cloud containers (optional GPU) +- **datalayer**: Cloud-based Datalayer Runtime ### Available in Sandbox diff --git a/docs/docs/skills/index.mdx b/docs/docs/skills/index.mdx index ec235c7..0920c32 100644 --- a/docs/docs/skills/index.mdx +++ b/docs/docs/skills/index.mdx @@ -102,25 +102,25 @@ The primary pattern is organizing skills as Python files in a `skills/` director async def batch_process(input_dir: str, output_dir: str) -> dict: """Process all files in a directory. - + Args: input_dir: Input directory path. output_dir: Output directory path. - + Returns: Processing statistics. """ from generated.mcp.filesystem import list_directory, read_file, write_file - + entries = await list_directory({"path": input_dir}) processed = 0 - + for entry in entries.get("entries", []): content = await read_file({"path": f"{input_dir}/{entry}"}) # Process content... await write_file({"path": f"{output_dir}/{entry}", "content": content.upper()}) processed += 1 - + return {"processed": processed} ``` @@ -142,12 +142,12 @@ import json async def analyze_file(file_path: str) -> dict: """Analyze a file and return statistics.""" from generated.mcp.filesystem import read_file - + content = await read_file({"path": file_path}) lines = content.split('\n') word_count = len(content.split()) char_count = len(content) - + return { "lines": len(lines), "words": word_count, @@ -160,7 +160,7 @@ if __name__ == "__main__": if len(sys.argv) < 2: print("Usage: python analyze_file.py ") sys.exit(1) - + file_path = sys.argv[1] result = asyncio.run(analyze_file(file_path)) print(json.dumps(result, indent=2)) @@ -197,22 +197,22 @@ async def wait_for( """Wait for an async condition to become true.""" import time start_time = time.time() - + while True: result = await condition() if result: return - + if time.time() - start_time > timeout_seconds: raise TimeoutError(f"Timeout waiting for condition after {timeout_seconds}s") - + await asyncio.sleep(interval_seconds) async def retry(fn, max_attempts: int = 3, delay_seconds: float = 1.0): """Retry a function until it succeeds or max attempts reached.""" last_error = None - + for attempt in range(1, max_attempts + 1): try: return await fn() @@ -220,7 +220,7 @@ async def retry(fn, max_attempts: int = 3, delay_seconds: float = 1.0): last_error = e if attempt < max_attempts: await asyncio.sleep(delay_seconds) - + raise RuntimeError(f"Failed after {max_attempts} attempts: {last_error}") ``` @@ -335,13 +335,13 @@ Skills can import and use other skills to build higher-level operations: async def analyze_and_report(data_dir: str) -> dict: from skills.batch_process import batch_process from skills.generate_report import generate_report - + # First process the files process_result = await batch_process(data_dir, f"{data_dir}/processed") - + # Then generate a report report = await generate_report(f"{data_dir}/processed") - + return {"processed": process_result["processed"], "report": report} ``` diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index 425bf0e..afa9761 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -201,7 +201,7 @@ module.exports = { docs: { routeBasePath: '/', sidebarPath: require.resolve('./sidebars.js'), - docItemComponent: '@theme/CustomDocItem', + docItemComponent: '@theme/CustomDocItem', editUrl: 'https://github.com/datalayer/agent-codemode/edit/main/', }, theme: { diff --git a/docs/static/img/datalayer/logo.svg b/docs/static/img/datalayer/logo.svg old mode 100755 new mode 100644 index a9527b0..f74aefa --- a/docs/static/img/datalayer/logo.svg +++ b/docs/static/img/datalayer/logo.svg @@ -155,4 +155,4 @@ height="14.174" width="17.01" y="48.188" - x="-5255.4331" /> \ No newline at end of file + x="-5255.4331" /> diff --git a/docs/static/img/datalayer/personas/business.svg b/docs/static/img/datalayer/personas/business.svg old mode 100755 new mode 100644 diff --git a/docs/static/img/datalayer/personas/data-scientist.svg b/docs/static/img/datalayer/personas/data-scientist.svg old mode 100755 new mode 100644 diff --git a/docs/static/img/datalayer/personas/devops.svg b/docs/static/img/datalayer/personas/devops.svg old mode 100755 new mode 100644 diff --git a/docs/static/img/datalayer/project/project-line.svg b/docs/static/img/datalayer/project/project-line.svg old mode 100755 new mode 100644 diff --git a/docs/static/img/open-source.svg b/docs/static/img/open-source.svg index f4b623e..4b1f782 100644 --- a/docs/static/img/open-source.svg +++ b/docs/static/img/open-source.svg @@ -1 +1 @@ -Pixel Icons \ No newline at end of file +Pixel Icons diff --git a/examples/patterns/codemode_example.py b/examples/patterns/codemode_example.py index 990c984..6f5e851 100644 --- a/examples/patterns/codemode_example.py +++ b/examples/patterns/codemode_example.py @@ -23,80 +23,80 @@ import asyncio import logging -from pathlib import Path logger = logging.getLogger(__name__) async def example_tool_discovery(): """Example 1: Progressive Tool Discovery. - + Instead of loading all tools upfront, use the Tool Search Tool to discover relevant tools based on the task at hand. """ - from agent_codemode import ToolRegistry, MCPServerConfig - + from agent_codemode import MCPServerConfig, ToolRegistry + logger.debug("=" * 60) logger.debug("Example 1: Progressive Tool Discovery") logger.debug("=" * 60) - + # Create a registry and add MCP servers registry = ToolRegistry() - + # Add an example server (filesystem operations) # In production, this would be a real MCP server - registry.add_server(MCPServerConfig( - name="filesystem", - command="npx", - args=["-y", "@anthropic/mcp-server-filesystem", "/tmp"], - )) - + registry.add_server( + MCPServerConfig( + name="filesystem", + command="npx", + args=["-y", "@anthropic/mcp-server-filesystem", "/tmp"], + ) + ) + # Discover all tools from configured servers logger.debug("\nDiscovering tools from MCP servers...") await registry.discover_all() - + # List all available tools all_tools = registry.list_tools() logger.debug("Discovered %s tools", len(all_tools)) - + # Search for specific tools (progressive discovery) logger.debug("\nSearching for 'file' operations...") result = await registry.search_tools("file operations", limit=5) - + for tool in result.tools: logger.debug(" - %s: %s", tool.name, tool.description) - + return registry async def example_code_execution(): """Example 2: Code-Based Tool Composition. - + Execute Python code that calls multiple tools. The code runs in an isolated sandbox with generated Python bindings for all tools. """ - from agent_codemode import ToolRegistry, CodeModeExecutor, CodeModeConfig - + from agent_codemode import CodeModeConfig, CodeModeExecutor, ToolRegistry + logger.debug("\n" + "=" * 60) logger.debug("Example 2: Code-Based Tool Composition") logger.debug("=" * 60) - + # Set up the registry registry = ToolRegistry() - + # Configure the executor config = CodeModeConfig( sandbox_variant="eval", # For development generated_path="./generated", skills_path="./skills", ) - + # Use the executor as an async context manager async with CodeModeExecutor(registry, config) as executor: - # Example: Execute code that would call tools # (This is a simplified example - real code would import generated bindings) - code = ''' + code = """ # This code runs in an isolated sandbox import os @@ -114,11 +114,11 @@ async def example_code_execution(): # Return the result result = f"Processed {data['files_processed']} files, total size: {data['total_size']} bytes" print(result) -''' - +""" + logger.debug("\nExecuting code in sandbox...") execution = await executor.execute(code) - + if execution.error: logger.debug("Error: %s", execution.error) else: @@ -126,21 +126,21 @@ async def example_code_execution(): "Output:\n%s", execution.logs.stdout if execution.logs else "No output", ) - + # Show tool call history logger.debug("\nTool calls made: %s", len(executor.tool_call_history)) async def _server(): """Example 4: Running the Codemode MCP Server. - + Shows how to configure and run the Codemode MCP server which exposes code execution capabilities to AI agents. """ logger.debug("\n" + "=" * 60) logger.debug("Example 4: Codemode MCP Server") logger.debug("=" * 60) - + logger.debug(""" The Codemode MCP Server exposes these tools to AI agents: @@ -154,10 +154,10 @@ async def _server(): from agent_codemode import codemode_server, configure_server from agent_codemode import MCPServerConfig - + # Configure with MCP servers configure_server() - + # Run the MCP server (uses FastMCP under the hood) codemode_server.run() @@ -172,12 +172,12 @@ async def main(): logger.debug("\n" + "=" * 60) logger.debug("Agent Codemode Examples") logger.debug("=" * 60) - + # Run examples await example_tool_discovery() await example_code_execution() await _server() - + logger.debug("\n" + "=" * 60) logger.debug("Examples Complete!") logger.debug("=" * 60) diff --git a/examples/patterns/codemode_patterns_example.py b/examples/patterns/codemode_patterns_example.py index 29b5f42..1c830a3 100644 --- a/examples/patterns/codemode_patterns_example.py +++ b/examples/patterns/codemode_patterns_example.py @@ -37,20 +37,20 @@ import asyncio import logging - # ============================================================================= # Example 1: Progressive Tool Discovery (Meta-Tool Pattern) # ============================================================================= + async def example_meta_tools(): """Demonstrate the meta-tool proxy pattern. - + The agent uses 4 meta-tools: 1. list_tool_names - Fast listing of tool names 2. search_tools - AI-powered tool discovery 3. get_tool_definition - Get full schema for a tool 4. execute_code - Run Python code in sandbox - + All actual tool execution goes through execute_code, which runs Python code using the generated tool bindings. """ @@ -58,27 +58,27 @@ async def example_meta_tools(): from agent_codemode.proxy.meta_tools import MetaToolProvider logger = logging.getLogger(__name__) - + logger.debug("=" * 70) logger.debug("Example 1: Meta-Tool Proxy Pattern") logger.debug("=" * 70) - + # Create registry with mock tools for demonstration registry = ToolRegistry() - + # In production, you'd add real MCP servers: # registry.add_server(MCPServerConfig(name="filesystem", ...)) # await registry.discover_all() - + # Create the meta-tool provider provider = MetaToolProvider(registry) - + # Get the meta-tool schemas (these are what the agent sees) meta_tools = provider.get_meta_tools() logger.debug("\nMeta-tools available to agent:") for tool in meta_tools: logger.debug(" - %s: %s...", tool["name"], tool["description"][:60]) - + # Example: Fast tool name listing logger.debug("\n1. list_tool_names (fast, no full schemas):") result = provider.list_tool_names(keywords=["file", "read"], limit=10) @@ -87,7 +87,7 @@ async def example_meta_tools(): result["total"], result["returned"], ) - + # Example: AI-powered search (if AI selector configured) logger.debug("\n2. search_tools (with full schemas):") result = await provider.search_tools("read CSV files and analyze data", limit=5) @@ -98,7 +98,7 @@ async def example_meta_tools(): tool["name"], tool.get("description", "")[:50], ) - + # Example: Get specific tool definition logger.debug("\n3. get_tool_definition:") # result = await provider.get_tool_definition("filesystem__read_file") @@ -109,21 +109,22 @@ async def example_meta_tools(): # Example 2: Code Execution (Tools as Code Pattern) # ============================================================================= + async def example_code_execution(): """Demonstrate code-based tool composition. - + The agent writes Python code that: 1. Imports from generated tool bindings 2. Calls multiple tools with regular Python 3. Uses loops, conditionals, error handling 4. Returns results - + This avoids LLM inference for each tool call! """ logger.debug("\n" + "=" * 70) logger.debug("Example 2: Code-Based Tool Composition") logger.debug("=" * 70) - + # Example code the agent would write and execute example_code = ''' # The agent writes code like this: @@ -131,44 +132,44 @@ async def example_code_execution(): async def process_files(): """Process multiple files efficiently.""" - + # List files in a directory entries = await list_directory({"path": "/tmp/data"}) - + results = [] for entry in entries.get("entries", []): if entry.endswith(".csv"): # Read each CSV file content = await read_file({"path": f"/tmp/data/{entry}"}) - + # Process it (in code, not LLM!) lines = content.split("\\n") row_count = len(lines) - + # Save result results.append({ "file": entry, "rows": row_count, }) - + # Write summary import json await write_file({ "path": "/tmp/summary.json", "content": json.dumps(results, indent=2) }) - + return results # Execute await process_files() ''' - + logger.debug("\nExample code the agent would write:") logger.debug("-" * 60) logger.debug(example_code) logger.debug("-" * 60) - + logger.debug("\nBenefits of Code Mode:") logger.debug(" ✓ One LLM call generates code that does many tool calls") logger.debug(" ✓ No LLM inference between each tool call") diff --git a/examples/simple/agent_cli.py b/examples/simple/agent_cli.py index c719f7d..686bfbe 100644 --- a/examples/simple/agent_cli.py +++ b/examples/simple/agent_cli.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- """Pydantic AI Agent CLI with Agent Codemode (STDIO). This agent connects to a local MCP stdio server and provides an @@ -10,16 +9,17 @@ from __future__ import annotations import asyncio +import inspect import io -import sys import logging +import sys from pathlib import Path from typing import Optional -import inspect try: from pydantic_ai import Agent from pydantic_ai.mcp import MCPServerStdio + HAS_PYDANTIC_AI = True except ImportError: HAS_PYDANTIC_AI = False @@ -27,6 +27,7 @@ try: from rich.console import Console from rich.table import Table + HAS_RICH = True except ImportError: HAS_RICH = False @@ -34,6 +35,7 @@ try: from agent_skills import AgentSkillsToolset, SandboxExecutor from code_sandboxes.eval_sandbox import EvalSandbox + HAS_AGENT_SKILLS = True except ImportError: HAS_AGENT_SKILLS = False @@ -57,9 +59,7 @@ def _build_prompt_examples(codemode: bool) -> str: "without returning the full content each time." ) else: - base.append( - "(Standard) Use the MCP tools directly for each step." - ) + base.append("(Standard) Use the MCP tools directly for each step.") return "\n".join(f" - {item}" for item in base) @@ -176,21 +176,28 @@ def _format_usage(usage: object, keys: Optional[list[str]] = None) -> str: return "N/A" parts = [] - preferred = [k for k in (keys or [ - "input_tokens", - "output_tokens", - "total_tokens", - "requests", - "cached_tokens", - "billable_tokens", - "cache_write_tokens", - "cache_read_tokens", - "input_audio_tokens", - "cache_audio_read_tokens", - "output_audio_tokens", - "tool_calls", - "details", - ]) if k != "details"] + preferred = [ + k + for k in ( + keys + or [ + "input_tokens", + "output_tokens", + "total_tokens", + "requests", + "cached_tokens", + "billable_tokens", + "cache_write_tokens", + "cache_read_tokens", + "input_audio_tokens", + "cache_audio_read_tokens", + "output_audio_tokens", + "tool_calls", + "details", + ] + ) + if k != "details" + ] for key in preferred: if key in data: parts.append(f"{key}={data[key]}") @@ -205,7 +212,7 @@ def _usage_to_table( prompt_usage: object, session_usage: object, keys: Optional[list[str]] = None, -) -> "Table": +) -> Table: prompt_data = _usage_to_dict(prompt_usage) session_data = _usage_to_dict(session_usage) table = Table(title="Token usage") @@ -245,7 +252,6 @@ def _format_value(value: object) -> str: _format_value(session_data.get(key, "-")), ) - if table.row_count == 0: table.add_row("usage", str(prompt_data or "-"), str(session_data or "-")) @@ -323,7 +329,7 @@ def create_agent(model: str, codemode: bool) -> tuple[Agent, object | None, obje mcp_server_path = _resolve_mcp_server_path() if codemode: - from agent_codemode import CodemodeToolset, ToolRegistry, MCPServerConfig, CodeModeConfig + from agent_codemode import CodeModeConfig, CodemodeToolset, MCPServerConfig, ToolRegistry registry = ToolRegistry() # Server name becomes the tool prefix (e.g., example_mcp__read_text_file) @@ -366,7 +372,9 @@ def create_agent(model: str, codemode: bool) -> tuple[Agent, object | None, obje executor=SandboxExecutor(shared_sandbox), ) toolsets.append(skills_toolset) - logger.info("Added AgentSkillsToolset with skills from %s", (repo_root / "skills").resolve()) + logger.info( + "Added AgentSkillsToolset with skills from %s", (repo_root / "skills").resolve() + ) else: logger.debug("agent_skills not available, skipping AgentSkillsToolset") else: @@ -410,22 +418,19 @@ def main() -> None: # Suppress verbose MCP server logs logging.getLogger("mcp.server").setLevel(logging.WARNING) - + import argparse + parser = argparse.ArgumentParser(description="MCP Agent CLI with Agent Codemode") parser.add_argument( "--model", type=str, default="anthropic:claude-sonnet-4-0", - help="Model to use (default: anthropic:claude-sonnet-4-0)" - ) - parser.add_argument( - "--codemode", - action="store_true", - help="Enable Agent Codemode mode" + help="Model to use (default: anthropic:claude-sonnet-4-0)", ) + parser.add_argument("--codemode", action="store_true", help="Enable Agent Codemode mode") args = parser.parse_args() - + model = args.model codemode = args.codemode @@ -522,15 +527,15 @@ async def _run_cli() -> None: node_type = type(node).__name__ # Print all node types for debugging logger.debug(" [iter %s] %s", iteration_count, node_type) - - if node_type == 'CallToolsNode': - mr = getattr(node, 'model_response', None) - if mr and hasattr(mr, 'parts'): + + if node_type == "CallToolsNode": + mr = getattr(node, "model_response", None) + if mr and hasattr(mr, "parts"): for p in mr.parts: - if hasattr(p, 'tool_name'): - args = getattr(p, 'args', {}) - if isinstance(args, dict) and 'code' in args: - code_preview = args['code'][:100].replace('\n', '\\n') + if hasattr(p, "tool_name"): + args = getattr(p, "args", {}) + if isinstance(args, dict) and "code" in args: + code_preview = args["code"][:100].replace("\n", "\\n") logger.debug( " -> %s(code=%s...)", p.tool_name, @@ -538,9 +543,9 @@ async def _run_cli() -> None: ) else: logger.debug(" -> %s(%s)", p.tool_name, args) - elif node_type == 'HandleResponseNode': + elif node_type == "HandleResponseNode": # Tool results might be here - data = getattr(node, 'data', None) + data = getattr(node, "data", None) if data: logger.debug(" -> data: %s", str(data)[:200]) run_result = run.result @@ -576,15 +581,22 @@ async def _run_cli() -> None: counts = codemode_toolset.get_call_counts() # type: ignore[assignment] if counts: prompt_usage_payload["codemode_tool_calls"] = ( - counts.get("codemode_tool_calls", 0) - previous_counts["codemode_tool_calls"] + counts.get("codemode_tool_calls", 0) + - previous_counts["codemode_tool_calls"] ) prompt_usage_payload["mcp_tool_calls"] = ( counts.get("mcp_tool_calls", 0) - previous_counts["mcp_tool_calls"] ) - previous_counts["codemode_tool_calls"] = counts.get("codemode_tool_calls", 0) + previous_counts["codemode_tool_calls"] = counts.get( + "codemode_tool_calls", 0 + ) previous_counts["mcp_tool_calls"] = counts.get("mcp_tool_calls", 0) - session_usage["codemode_tool_calls"] += float(prompt_usage_payload["codemode_tool_calls"]) - session_usage["mcp_tool_calls"] += float(prompt_usage_payload["mcp_tool_calls"]) + session_usage["codemode_tool_calls"] += float( + prompt_usage_payload["codemode_tool_calls"] + ) + session_usage["mcp_tool_calls"] += float( + prompt_usage_payload["mcp_tool_calls"] + ) # Track skills tool calls separately if codemode and skills_toolset is not None: @@ -603,14 +615,19 @@ async def _run_cli() -> None: prompt_usage_payload["skills_calls"] = "N/A" if not codemode: - if "tool_calls" in prompt_usage_payload and "mcp_tool_calls" not in prompt_usage_payload: + if ( + "tool_calls" in prompt_usage_payload + and "mcp_tool_calls" not in prompt_usage_payload + ): prompt_usage_payload["mcp_tool_calls"] = prompt_usage_payload["tool_calls"] if "tool_calls" in prompt_usage_payload: prompt_usage_payload.pop("tool_calls", None) prompt_usage_payload.setdefault("mcp_tool_calls", 0) prompt_usage_payload["codemode_tool_calls"] = "N/A" prompt_usage_payload["skills_calls"] = "N/A" - if "tool_calls" in usage_data and isinstance(usage_data["tool_calls"], (int, float)): + if "tool_calls" in usage_data and isinstance( + usage_data["tool_calls"], (int, float) + ): session_usage["mcp_tool_calls"] += float(usage_data["tool_calls"]) elif "tool_calls" in usage_data: try: @@ -650,7 +667,9 @@ async def _run_cli() -> None: elif skills_toolset is None: session_usage_payload["skills_calls"] = "N/A" console.print( - _usage_to_table(prompt_usage_payload, session_usage_payload, prompt_usage_keys) + _usage_to_table( + prompt_usage_payload, session_usage_payload, prompt_usage_keys + ) ) console.print() else: diff --git a/examples/simple/example_mcp_server.py b/examples/simple/example_mcp_server.py index 94722c1..5249a58 100644 --- a/examples/simple/example_mcp_server.py +++ b/examples/simple/example_mcp_server.py @@ -13,15 +13,15 @@ from pathlib import Path from typing import Optional -from typing_extensions import TypedDict - from mcp.server.fastmcp import FastMCP +from typing_extensions import TypedDict mcp = FastMCP("example-mcp-server") # Base directory for file operations - defaults to /tmp if CWD is not writable _BASE_DIR: Path | None = None + def _get_base_dir() -> Path: """Get the base directory for file operations.""" global _BASE_DIR @@ -35,28 +35,94 @@ def _get_base_dir() -> Path: _BASE_DIR.mkdir(parents=True, exist_ok=True) return _BASE_DIR + _WORDS = [ - "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", - "india", "juliet", "kilo", "lima", "mike", "november", "oscar", "papa", - "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey", - "xray", "yankee", "zulu", "apricot", "banana", "cherry", "date", "elderberry", - "fig", "grape", "honeydew", "kiwi", "lemon", "mango", "nectarine", "orange", - "papaya", "quince", "raspberry", "strawberry", "tangerine", "ugli", "vanilla", - "watermelon", "xigua", "yam", "zucchini", "azure", "binary", "cache", - "docker", "elastic", "feature", "gateway", "hash", "index", "json", "kafka", - "lambda", "micro", "node", "object", "python", "query", "router", "schema", - "token", "vector", "worker", "yaml", + "alpha", + "bravo", + "charlie", + "delta", + "echo", + "foxtrot", + "golf", + "hotel", + "india", + "juliet", + "kilo", + "lima", + "mike", + "november", + "oscar", + "papa", + "quebec", + "romeo", + "sierra", + "tango", + "uniform", + "victor", + "whiskey", + "xray", + "yankee", + "zulu", + "apricot", + "banana", + "cherry", + "date", + "elderberry", + "fig", + "grape", + "honeydew", + "kiwi", + "lemon", + "mango", + "nectarine", + "orange", + "papaya", + "quince", + "raspberry", + "strawberry", + "tangerine", + "ugli", + "vanilla", + "watermelon", + "xigua", + "yam", + "zucchini", + "azure", + "binary", + "cache", + "docker", + "elastic", + "feature", + "gateway", + "hash", + "index", + "json", + "kafka", + "lambda", + "micro", + "node", + "object", + "python", + "query", + "router", + "schema", + "token", + "vector", + "worker", + "yaml", ] class GenerateRandomTextResult(TypedDict): """Result from generate_random_text.""" + text: str word_count: int class WriteTextFileResult(TypedDict): """Result from write_text_file.""" + path: str bytes: int words: int @@ -64,6 +130,7 @@ class WriteTextFileResult(TypedDict): class ReadTextFileResult(TypedDict, total=False): """Result from read_text_file.""" + path: str bytes: int words: int @@ -72,6 +139,7 @@ class ReadTextFileResult(TypedDict, total=False): class ReadTextFileManyResult(TypedDict): """Result from read_text_file_many.""" + path: str reads: int last: ReadTextFileResult @@ -79,7 +147,7 @@ class ReadTextFileManyResult(TypedDict): def _normalize_path(path: str) -> Path: """Normalize a file path. - + Absolute paths are used as-is. Relative paths are resolved relative to the base directory (CWD or /tmp/mcp_files if CWD is not writable). """ @@ -90,7 +158,9 @@ def _normalize_path(path: str) -> Path: @mcp.tool() -def generate_random_text(word_count: int = 1000, seed: Optional[int] = None) -> GenerateRandomTextResult: +def generate_random_text( + word_count: int = 1000, seed: Optional[int] = None +) -> GenerateRandomTextResult: """Generate pseudo-random text. Args: @@ -146,7 +216,7 @@ def read_text_file( target = _normalize_path(path) content = target.read_text(encoding="utf-8") if max_chars is not None: - content = content[: max_chars] + content = content[:max_chars] response = { "path": str(target), "bytes": len(content.encode("utf-8")), diff --git a/examples/skills/README.md b/examples/skills/README.md index 57379a2..fb51d14 100644 --- a/examples/skills/README.md +++ b/examples/skills/README.md @@ -35,8 +35,8 @@ make agent-standard # Standard MCP mode The `skills/` folder contains skill definitions that the agent can discover and use: -| Skill | Description | -|-------|-------------| +| Skill | Description | +| ----- | --------------------------------------------------------------------------------- | | `pdf` | PDF manipulation toolkit - extract text/tables, merge/split documents, fill forms | ### Example prompts @@ -66,7 +66,7 @@ If you don't see it, run a prompt that triggers tool discovery (e.g., `/list_too ## How Skills Work 1. **Discovery**: The agent discovers skills from the `skills/` folder -2. **Understanding**: Skills contain SKILL.md with instructions and scripts -3. **Execution**: The agent uses codemode to execute skill scripts in a sandbox +1. **Understanding**: Skills contain SKILL.md with instructions and scripts +1. **Execution**: The agent uses codemode to execute skill scripts in a sandbox For more on skills, see the [agent-skills](https://github.com/datalayer/agent-skills) repository. diff --git a/examples/skills/agent_cli.py b/examples/skills/agent_cli.py index 874dbe8..3aaa98d 100644 --- a/examples/skills/agent_cli.py +++ b/examples/skills/agent_cli.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- """Skills Agent CLI with Agent Codemode (STDIO). This agent connects to a local MCP stdio server and provides an @@ -13,16 +12,17 @@ from __future__ import annotations import asyncio +import inspect import io -import sys import logging +import sys from pathlib import Path from typing import Optional -import inspect try: from pydantic_ai import Agent from pydantic_ai.mcp import MCPServerStdio + HAS_PYDANTIC_AI = True except ImportError: HAS_PYDANTIC_AI = False @@ -30,6 +30,7 @@ try: from rich.console import Console from rich.table import Table + HAS_RICH = True except ImportError: HAS_RICH = False @@ -53,9 +54,7 @@ def _build_prompt_examples(codemode: bool) -> str: "without returning the full content each time." ) else: - base.append( - "(Standard) Use the MCP tools directly for each step." - ) + base.append("(Standard) Use the MCP tools directly for each step.") return "\n".join(f" - {item}" for item in base) @@ -177,21 +176,28 @@ def _format_usage(usage: object, keys: Optional[list[str]] = None) -> str: return "N/A" parts = [] - preferred = [k for k in (keys or [ - "input_tokens", - "output_tokens", - "total_tokens", - "requests", - "cached_tokens", - "billable_tokens", - "cache_write_tokens", - "cache_read_tokens", - "input_audio_tokens", - "cache_audio_read_tokens", - "output_audio_tokens", - "tool_calls", - "details", - ]) if k != "details"] + preferred = [ + k + for k in ( + keys + or [ + "input_tokens", + "output_tokens", + "total_tokens", + "requests", + "cached_tokens", + "billable_tokens", + "cache_write_tokens", + "cache_read_tokens", + "input_audio_tokens", + "cache_audio_read_tokens", + "output_audio_tokens", + "tool_calls", + "details", + ] + ) + if k != "details" + ] for key in preferred: if key in data: parts.append(f"{key}={data[key]}") @@ -206,7 +212,7 @@ def _usage_to_table( prompt_usage: object, session_usage: object, keys: Optional[list[str]] = None, -) -> "Table": +) -> Table: prompt_data = _usage_to_dict(prompt_usage) session_data = _usage_to_dict(session_usage) table = Table(title="Token usage") @@ -246,7 +252,6 @@ def _format_value(value: object) -> str: _format_value(session_data.get(key, "-")), ) - if table.row_count == 0: table.add_row("usage", str(prompt_data or "-"), str(session_data or "-")) @@ -327,9 +332,10 @@ def create_agent(model: str, codemode: bool) -> tuple[Agent, object | None, obje skills_toolset = None if codemode: - from agent_codemode import CodemodeToolset, ToolRegistry, MCPServerConfig, CodeModeConfig from agent_skills import AgentSkillsToolset + from agent_codemode import CodeModeConfig, CodemodeToolset, MCPServerConfig, ToolRegistry + registry = ToolRegistry() # Server name becomes the tool prefix (e.g., example_mcp__read_text_file) # and matches the generated bindings path: generated.mcp.example_mcp @@ -354,12 +360,12 @@ def create_agent(model: str, codemode: bool) -> tuple[Agent, object | None, obje config=config, allow_discovery_tools=True, # Enable discovery tools (search_tools, get_tool_details, list_tool_names, list_servers) ) - + # Add skills toolset for skill discovery and execution skills_toolset = AgentSkillsToolset( directories=[str(skills_dir.resolve())], ) - + toolsets = [codemode_toolset, skills_toolset] toolset = codemode_toolset else: @@ -402,22 +408,21 @@ def main() -> None: # Suppress verbose MCP server logs logging.getLogger("mcp.server").setLevel(logging.WARNING) - + import argparse + parser = argparse.ArgumentParser(description="Skills Agent CLI with Agent Codemode") parser.add_argument( "--model", type=str, default="anthropic:claude-sonnet-4-0", - help="Model to use (default: anthropic:claude-sonnet-4-0)" + help="Model to use (default: anthropic:claude-sonnet-4-0)", ) parser.add_argument( - "--standard", - action="store_true", - help="Disable codemode (use standard MCP mode)" + "--standard", action="store_true", help="Disable codemode (use standard MCP mode)" ) args = parser.parse_args() - + model = args.model codemode = not args.standard # Codemode is the default @@ -514,15 +519,15 @@ async def _run_cli() -> None: node_type = type(node).__name__ # Print all node types for debugging logger.debug(" [iter %s] %s", iteration_count, node_type) - - if node_type == 'CallToolsNode': - mr = getattr(node, 'model_response', None) - if mr and hasattr(mr, 'parts'): + + if node_type == "CallToolsNode": + mr = getattr(node, "model_response", None) + if mr and hasattr(mr, "parts"): for p in mr.parts: - if hasattr(p, 'tool_name'): - args = getattr(p, 'args', {}) - if isinstance(args, dict) and 'code' in args: - code_preview = args['code'][:100].replace('\n', '\\n') + if hasattr(p, "tool_name"): + args = getattr(p, "args", {}) + if isinstance(args, dict) and "code" in args: + code_preview = args["code"][:100].replace("\n", "\\n") logger.debug( " -> %s(code=%s...)", p.tool_name, @@ -530,9 +535,9 @@ async def _run_cli() -> None: ) else: logger.debug(" -> %s(%s)", p.tool_name, args) - elif node_type == 'HandleResponseNode': + elif node_type == "HandleResponseNode": # Tool results might be here - data = getattr(node, 'data', None) + data = getattr(node, "data", None) if data: logger.debug(" -> data: %s", str(data)[:200]) run_result = run.result @@ -568,15 +573,22 @@ async def _run_cli() -> None: counts = codemode_toolset.get_call_counts() # type: ignore[assignment] if counts: prompt_usage_payload["codemode_tool_calls"] = ( - counts.get("codemode_tool_calls", 0) - previous_counts["codemode_tool_calls"] + counts.get("codemode_tool_calls", 0) + - previous_counts["codemode_tool_calls"] ) prompt_usage_payload["mcp_tool_calls"] = ( counts.get("mcp_tool_calls", 0) - previous_counts["mcp_tool_calls"] ) - previous_counts["codemode_tool_calls"] = counts.get("codemode_tool_calls", 0) + previous_counts["codemode_tool_calls"] = counts.get( + "codemode_tool_calls", 0 + ) previous_counts["mcp_tool_calls"] = counts.get("mcp_tool_calls", 0) - session_usage["codemode_tool_calls"] += float(prompt_usage_payload["codemode_tool_calls"]) - session_usage["mcp_tool_calls"] += float(prompt_usage_payload["mcp_tool_calls"]) + session_usage["codemode_tool_calls"] += float( + prompt_usage_payload["codemode_tool_calls"] + ) + session_usage["mcp_tool_calls"] += float( + prompt_usage_payload["mcp_tool_calls"] + ) # Track skills tool calls if skills_toolset is not None: @@ -593,13 +605,18 @@ async def _run_cli() -> None: prompt_usage_payload["skills_calls"] = "N/A" if not codemode: - if "tool_calls" in prompt_usage_payload and "mcp_tool_calls" not in prompt_usage_payload: + if ( + "tool_calls" in prompt_usage_payload + and "mcp_tool_calls" not in prompt_usage_payload + ): prompt_usage_payload["mcp_tool_calls"] = prompt_usage_payload["tool_calls"] if "tool_calls" in prompt_usage_payload: prompt_usage_payload.pop("tool_calls", None) prompt_usage_payload.setdefault("mcp_tool_calls", 0) prompt_usage_payload["codemode_tool_calls"] = "N/A" - if "tool_calls" in usage_data and isinstance(usage_data["tool_calls"], (int, float)): + if "tool_calls" in usage_data and isinstance( + usage_data["tool_calls"], (int, float) + ): session_usage["mcp_tool_calls"] += float(usage_data["tool_calls"]) elif "tool_calls" in usage_data: try: @@ -638,7 +655,9 @@ async def _run_cli() -> None: if skills_toolset is None: session_usage_payload["skills_calls"] = "N/A" console.print( - _usage_to_table(prompt_usage_payload, session_usage_payload, prompt_usage_keys) + _usage_to_table( + prompt_usage_payload, session_usage_payload, prompt_usage_keys + ) ) console.print() else: diff --git a/examples/skills/example_mcp_server.py b/examples/skills/example_mcp_server.py index 94722c1..5249a58 100644 --- a/examples/skills/example_mcp_server.py +++ b/examples/skills/example_mcp_server.py @@ -13,15 +13,15 @@ from pathlib import Path from typing import Optional -from typing_extensions import TypedDict - from mcp.server.fastmcp import FastMCP +from typing_extensions import TypedDict mcp = FastMCP("example-mcp-server") # Base directory for file operations - defaults to /tmp if CWD is not writable _BASE_DIR: Path | None = None + def _get_base_dir() -> Path: """Get the base directory for file operations.""" global _BASE_DIR @@ -35,28 +35,94 @@ def _get_base_dir() -> Path: _BASE_DIR.mkdir(parents=True, exist_ok=True) return _BASE_DIR + _WORDS = [ - "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", - "india", "juliet", "kilo", "lima", "mike", "november", "oscar", "papa", - "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey", - "xray", "yankee", "zulu", "apricot", "banana", "cherry", "date", "elderberry", - "fig", "grape", "honeydew", "kiwi", "lemon", "mango", "nectarine", "orange", - "papaya", "quince", "raspberry", "strawberry", "tangerine", "ugli", "vanilla", - "watermelon", "xigua", "yam", "zucchini", "azure", "binary", "cache", - "docker", "elastic", "feature", "gateway", "hash", "index", "json", "kafka", - "lambda", "micro", "node", "object", "python", "query", "router", "schema", - "token", "vector", "worker", "yaml", + "alpha", + "bravo", + "charlie", + "delta", + "echo", + "foxtrot", + "golf", + "hotel", + "india", + "juliet", + "kilo", + "lima", + "mike", + "november", + "oscar", + "papa", + "quebec", + "romeo", + "sierra", + "tango", + "uniform", + "victor", + "whiskey", + "xray", + "yankee", + "zulu", + "apricot", + "banana", + "cherry", + "date", + "elderberry", + "fig", + "grape", + "honeydew", + "kiwi", + "lemon", + "mango", + "nectarine", + "orange", + "papaya", + "quince", + "raspberry", + "strawberry", + "tangerine", + "ugli", + "vanilla", + "watermelon", + "xigua", + "yam", + "zucchini", + "azure", + "binary", + "cache", + "docker", + "elastic", + "feature", + "gateway", + "hash", + "index", + "json", + "kafka", + "lambda", + "micro", + "node", + "object", + "python", + "query", + "router", + "schema", + "token", + "vector", + "worker", + "yaml", ] class GenerateRandomTextResult(TypedDict): """Result from generate_random_text.""" + text: str word_count: int class WriteTextFileResult(TypedDict): """Result from write_text_file.""" + path: str bytes: int words: int @@ -64,6 +130,7 @@ class WriteTextFileResult(TypedDict): class ReadTextFileResult(TypedDict, total=False): """Result from read_text_file.""" + path: str bytes: int words: int @@ -72,6 +139,7 @@ class ReadTextFileResult(TypedDict, total=False): class ReadTextFileManyResult(TypedDict): """Result from read_text_file_many.""" + path: str reads: int last: ReadTextFileResult @@ -79,7 +147,7 @@ class ReadTextFileManyResult(TypedDict): def _normalize_path(path: str) -> Path: """Normalize a file path. - + Absolute paths are used as-is. Relative paths are resolved relative to the base directory (CWD or /tmp/mcp_files if CWD is not writable). """ @@ -90,7 +158,9 @@ def _normalize_path(path: str) -> Path: @mcp.tool() -def generate_random_text(word_count: int = 1000, seed: Optional[int] = None) -> GenerateRandomTextResult: +def generate_random_text( + word_count: int = 1000, seed: Optional[int] = None +) -> GenerateRandomTextResult: """Generate pseudo-random text. Args: @@ -146,7 +216,7 @@ def read_text_file( target = _normalize_path(path) content = target.read_text(encoding="utf-8") if max_chars is not None: - content = content[: max_chars] + content = content[:max_chars] response = { "path": str(target), "bytes": len(content.encode("utf-8")), diff --git a/examples/skills/skills/pdf/SKILL.md b/examples/skills/skills/pdf/SKILL.md index f6a22dd..decb4a8 100644 --- a/examples/skills/skills/pdf/SKILL.md +++ b/examples/skills/skills/pdf/SKILL.md @@ -30,6 +30,7 @@ for page in reader.pages: ### pypdf - Basic Operations #### Merge PDFs + ```python from pypdf import PdfWriter, PdfReader @@ -44,6 +45,7 @@ with open("merged.pdf", "wb") as output: ``` #### Split PDF + ```python reader = PdfReader("input.pdf") for i, page in enumerate(reader.pages): @@ -54,6 +56,7 @@ for i, page in enumerate(reader.pages): ``` #### Extract Metadata + ```python reader = PdfReader("document.pdf") meta = reader.metadata @@ -64,6 +67,7 @@ print(f"Creator: {meta.creator}") ``` #### Rotate Pages + ```python reader = PdfReader("input.pdf") writer = PdfWriter() @@ -79,6 +83,7 @@ with open("rotated.pdf", "wb") as output: ### pdfplumber - Text and Table Extraction #### Extract Text with Layout + ```python import pdfplumber @@ -89,6 +94,7 @@ with pdfplumber.open("document.pdf") as pdf: ``` #### Extract Tables + ```python with pdfplumber.open("document.pdf") as pdf: for i, page in enumerate(pdf.pages): @@ -100,6 +106,7 @@ with pdfplumber.open("document.pdf") as pdf: ``` #### Advanced Table Extraction + ```python import pandas as pd @@ -121,6 +128,7 @@ if all_tables: ### reportlab - Create PDFs #### Basic PDF Creation + ```python from reportlab.lib.pagesizes import letter from reportlab.pdfgen import canvas @@ -140,6 +148,7 @@ c.save() ``` #### Create PDF with Multiple Pages + ```python from reportlab.lib.pagesizes import letter from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak @@ -169,6 +178,7 @@ doc.build(story) ## Command-Line Tools ### pdftotext (poppler-utils) + ```bash # Extract text pdftotext input.pdf output.txt @@ -181,6 +191,7 @@ pdftotext -f 1 -l 5 input.pdf output.txt # Pages 1-5 ``` ### qpdf + ```bash # Merge PDFs qpdf --empty --pages file1.pdf file2.pdf -- merged.pdf @@ -197,6 +208,7 @@ qpdf --password=mypassword --decrypt encrypted.pdf decrypted.pdf ``` ### pdftk (if available) + ```bash # Merge pdftk file1.pdf file2.pdf cat output merged.pdf @@ -211,6 +223,7 @@ pdftk input.pdf rotate 1east output rotated.pdf ## Common Tasks ### Extract Text from Scanned PDFs + ```python # Requires: pip install pytesseract pdf2image import pytesseract @@ -230,6 +243,7 @@ print(text) ``` ### Add Watermark + ```python from pypdf import PdfReader, PdfWriter @@ -249,6 +263,7 @@ with open("watermarked.pdf", "wb") as output: ``` ### Extract Images + ```bash # Using pdfimages (poppler-utils) pdfimages -j input.pdf output_prefix @@ -257,6 +272,7 @@ pdfimages -j input.pdf output_prefix ``` ### Password Protection + ```python from pypdf import PdfReader, PdfWriter @@ -275,16 +291,16 @@ with open("encrypted.pdf", "wb") as output: ## Quick Reference -| Task | Best Tool | Command/Code | -|------|-----------|--------------| -| Merge PDFs | pypdf | `writer.add_page(page)` | -| Split PDFs | pypdf | One page per file | -| Extract text | pdfplumber | `page.extract_text()` | -| Extract tables | pdfplumber | `page.extract_tables()` | -| Create PDFs | reportlab | Canvas or Platypus | -| Command line merge | qpdf | `qpdf --empty --pages ...` | -| OCR scanned PDFs | pytesseract | Convert to image first | -| Fill PDF forms | pdf-lib or pypdf (see forms.md) | See forms.md | +| Task | Best Tool | Command/Code | +| ------------------ | ------------------------------- | -------------------------- | +| Merge PDFs | pypdf | `writer.add_page(page)` | +| Split PDFs | pypdf | One page per file | +| Extract text | pdfplumber | `page.extract_text()` | +| Extract tables | pdfplumber | `page.extract_tables()` | +| Create PDFs | reportlab | Canvas or Platypus | +| Command line merge | qpdf | `qpdf --empty --pages ...` | +| OCR scanned PDFs | pytesseract | Convert to image first | +| Fill PDF forms | pdf-lib or pypdf (see forms.md) | See forms.md | ## Next Steps diff --git a/examples/skills/skills/pdf/forms.md b/examples/skills/skills/pdf/forms.md index 4e23450..a6dbfba 100644 --- a/examples/skills/skills/pdf/forms.md +++ b/examples/skills/skills/pdf/forms.md @@ -1,11 +1,14 @@ **CRITICAL: You MUST complete these steps in order. Do not skip ahead to writing code.** If you need to fill out a PDF form, first check to see if the PDF has fillable form fields. Run this script from this file's directory: - `python scripts/check_fillable_fields `, and depending on the result go to either the "Fillable fields" or "Non-fillable fields" and follow those instructions. +`python scripts/check_fillable_fields `, and depending on the result go to either the "Fillable fields" or "Non-fillable fields" and follow those instructions. # Fillable fields + If the PDF has fillable form fields: + - Run this script from this file's directory: `python scripts/extract_form_field_info.py `. It will create a JSON file with a list of fields in this format: + ``` [ { @@ -50,10 +53,12 @@ If the PDF has fillable form fields: } ] ``` + - Convert the PDF to PNGs (one image for each page) with this script (run from this file's directory): -`python scripts/convert_pdf_to_images.py ` -Then analyze the images to determine the purpose of each form field (make sure to convert the bounding box PDF coordinates to image coordinates). + `python scripts/convert_pdf_to_images.py ` + Then analyze the images to determine the purpose of each form field (make sure to convert the bounding box PDF coordinates to image coordinates). - Create a `field_values.json` file in this format with the values to be entered for each field: + ``` [ { @@ -71,64 +76,81 @@ Then analyze the images to determine the purpose of each form field (make sure t // more fields ] ``` + - Run the `fill_fillable_fields.py` script from this file's directory to create a filled-in PDF: -`python scripts/fill_fillable_fields.py ` -This script will verify that the field IDs and values you provide are valid; if it prints error messages, correct the appropriate fields and try again. + `python scripts/fill_fillable_fields.py ` + This script will verify that the field IDs and values you provide are valid; if it prints error messages, correct the appropriate fields and try again. # Non-fillable fields + If the PDF doesn't have fillable form fields, you'll need to visually determine where the data should be added and create text annotations. Follow the below steps *exactly*. You MUST perform all of these steps to ensure that the the form is accurately completed. Details for each step are below. + - Convert the PDF to PNG images and determine field bounding boxes. - Create a JSON file with field information and validation images showing the bounding boxes. - Validate the the bounding boxes. - Use the bounding boxes to fill in the form. ## Step 1: Visual Analysis (REQUIRED) + - Convert the PDF to PNG images. Run this script from this file's directory: -`python scripts/convert_pdf_to_images.py ` -The script will create a PNG image for each page in the PDF. + `python scripts/convert_pdf_to_images.py ` + The script will create a PNG image for each page in the PDF. - Carefully examine each PNG image and identify all form fields and areas where the user should enter data. For each form field where the user should enter text, determine bounding boxes for both the form field label, and the area where the user should enter text. The label and entry bounding boxes MUST NOT INTERSECT; the text entry box should only include the area where data should be entered. Usually this area will be immediately to the side, above, or below its label. Entry bounding boxes must be tall and wide enough to contain their text. These are some examples of form structures that you might see: *Label inside box* + ``` ┌────────────────────────┐ │ Name: │ └────────────────────────┘ ``` + The input area should be to the right of the "Name" label and extend to the edge of the box. *Label before line* + ``` Email: _______________________ ``` + The input area should be above the line and include its entire width. *Label under line* + ``` _________________________ Name ``` + The input area should be above the line and include the entire width of the line. This is common for signature and date fields. *Label above line* + ``` Please enter any special requests: ________________________________________________ ``` + The input area should extend from the bottom of the label to the line, and should include the entire width of the line. *Checkboxes* + ``` Are you a US citizen? Yes □ No □ ``` + For checkboxes: + - Look for small square boxes (□) - these are the actual checkboxes to target. They may be to the left or right of their labels. - Distinguish between label text ("Yes", "No") and the clickable checkbox squares. - The entry bounding box should cover ONLY the small square, not the text label. ### Step 2: Create fields.json and validation images (REQUIRED) + - Create a file named `fields.json` with information for the form fields and bounding boxes in this format: + ``` { "pages": [ @@ -177,29 +199,37 @@ For checkboxes: ``` Create validation images by running this script from this file's directory for each page: -`python scripts/create_validation_image.py +\`python scripts/create_validation_image.py \ \ \ \ The validation images will have red rectangles where text should be entered, and blue rectangles covering label text. ### Step 3: Validate Bounding Boxes (REQUIRED) + #### Automated intersection check + - Verify that none of bounding boxes intersect and that the entry bounding boxes are tall enough by checking the fields.json file with the `check_bounding_boxes.py` script (run from this file's directory): -`python scripts/check_bounding_boxes.py ` + `python scripts/check_bounding_boxes.py ` If there are errors, reanalyze the relevant fields, adjust the bounding boxes, and iterate until there are no remaining errors. Remember: label (blue) bounding boxes should contain text labels, entry (red) boxes should not. #### Manual image inspection + **CRITICAL: Do not proceed without visually inspecting validation images** + - Red rectangles must ONLY cover input areas + - Red rectangles MUST NOT contain any text + - Blue rectangles should contain label text + - For checkboxes: + - Red rectangle MUST be centered on the checkbox square - Blue rectangle should cover the text label for the checkbox - If any rectangles look wrong, fix fields.json, regenerate the validation images, and verify again. Repeat this process until the bounding boxes are fully accurate. - ### Step 4: Add annotations to the PDF + Run this script from this file's directory to create a filled-out PDF using the information in fields.json: -`python scripts/fill_pdf_form_with_annotations.py +\`python scripts/fill_pdf_form_with_annotations.py \ \ \ diff --git a/examples/skills/skills/pdf/reference.md b/examples/skills/skills/pdf/reference.md index 41400bf..3f21575 100644 --- a/examples/skills/skills/pdf/reference.md +++ b/examples/skills/skills/pdf/reference.md @@ -5,9 +5,11 @@ This document contains advanced PDF processing features, detailed examples, and ## pypdfium2 Library (Apache/BSD License) ### Overview + pypdfium2 is a Python binding for PDFium (Chromium's PDF library). It's excellent for fast PDF rendering, image generation, and serves as a PyMuPDF replacement. ### Render PDF to Images + ```python import pypdfium2 as pdfium from PIL import Image @@ -34,6 +36,7 @@ for i, page in enumerate(pdf): ``` ### Extract Text with pypdfium2 + ```python import pypdfium2 as pdfium @@ -50,6 +53,7 @@ for i, page in enumerate(pdf): pdf-lib is a powerful JavaScript library for creating and modifying PDF documents in any JavaScript environment. #### Load and Manipulate Existing PDF + ```javascript import { PDFDocument } from 'pdf-lib'; import fs from 'fs'; @@ -78,6 +82,7 @@ async function manipulatePDF() { ``` #### Create Complex PDFs from Scratch + ```javascript import { PDFDocument, rgb, StandardFonts } from 'pdf-lib'; import fs from 'fs'; @@ -139,6 +144,7 @@ async function createPDF() { ``` #### Advanced Merge and Split Operations + ```javascript import { PDFDocument } from 'pdf-lib'; import fs from 'fs'; @@ -172,6 +178,7 @@ async function mergePDFs() { PDF.js is Mozilla's JavaScript library for rendering PDFs in the browser. #### Basic PDF Loading and Rendering + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -206,6 +213,7 @@ async function renderPDF() { ``` #### Extract Text with Coordinates + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -242,6 +250,7 @@ async function extractText() { ``` #### Extract Annotations and Forms + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -267,6 +276,7 @@ async function extractAnnotations() { ### poppler-utils Advanced Features #### Extract Text with Bounding Box Coordinates + ```bash # Extract text with bounding box coordinates (essential for structured data) pdftotext -bbox-layout document.pdf output.xml @@ -275,6 +285,7 @@ pdftotext -bbox-layout document.pdf output.xml ``` #### Advanced Image Conversion + ```bash # Convert to PNG images with specific resolution pdftoppm -png -r 300 document.pdf output_prefix @@ -287,6 +298,7 @@ pdftoppm -jpeg -jpegopt quality=85 -r 200 document.pdf jpeg_output ``` #### Extract Embedded Images + ```bash # Extract all embedded images with metadata pdfimages -j -p document.pdf page_images @@ -301,6 +313,7 @@ pdfimages -all document.pdf images/img ### qpdf Advanced Features #### Complex Page Manipulation + ```bash # Split PDF into groups of pages qpdf --split-pages=3 input.pdf output_group_%02d.pdf @@ -313,6 +326,7 @@ qpdf --empty --pages doc1.pdf 1-3 doc2.pdf 5-7 doc3.pdf 2,4 -- combined.pdf ``` #### PDF Optimization and Repair + ```bash # Optimize PDF for web (linearize for streaming) qpdf --linearize input.pdf optimized.pdf @@ -329,6 +343,7 @@ qpdf --show-all-pages input.pdf > structure.txt ``` #### Advanced Encryption + ```bash # Add password protection with specific permissions qpdf --encrypt user_pass owner_pass 256 --print=none --modify=none -- input.pdf encrypted.pdf @@ -345,29 +360,31 @@ qpdf --password=secret123 --decrypt encrypted.pdf decrypted.pdf ### pdfplumber Advanced Features #### Extract Text with Precise Coordinates + ```python import pdfplumber with pdfplumber.open("document.pdf") as pdf: page = pdf.pages[0] - + # Extract all text with coordinates chars = page.chars for char in chars[:10]: # First 10 characters print(f"Char: '{char['text']}' at x:{char['x0']:.1f} y:{char['y0']:.1f}") - + # Extract text by bounding box (left, top, right, bottom) bbox_text = page.within_bbox((100, 100, 400, 200)).extract_text() ``` #### Advanced Table Extraction with Custom Settings + ```python import pdfplumber import pandas as pd with pdfplumber.open("complex_table.pdf") as pdf: page = pdf.pages[0] - + # Extract tables with custom settings for complex layouts table_settings = { "vertical_strategy": "lines", @@ -376,7 +393,7 @@ with pdfplumber.open("complex_table.pdf") as pdf: "intersection_tolerance": 15 } tables = page.extract_tables(table_settings) - + # Visual debugging for table extraction img = page.to_image(resolution=150) img.save("debug_layout.png") @@ -385,6 +402,7 @@ with pdfplumber.open("complex_table.pdf") as pdf: ### reportlab Advanced Features #### Create Professional Reports with Tables + ```python from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph from reportlab.lib.styles import getSampleStyleSheet @@ -428,12 +446,14 @@ doc.build(elements) ### Extract Figures/Images from PDF #### Method 1: Using pdfimages (fastest) + ```bash # Extract all images with original quality pdfimages -all document.pdf images/img ``` #### Method 2: Using pypdfium2 + Image Processing + ```python import pypdfium2 as pdfium from PIL import Image @@ -441,26 +461,27 @@ import numpy as np def extract_figures(pdf_path, output_dir): pdf = pdfium.PdfDocument(pdf_path) - + for page_num, page in enumerate(pdf): # Render high-resolution page bitmap = page.render(scale=3.0) img = bitmap.to_pil() - + # Convert to numpy for processing img_array = np.array(img) - + # Simple figure detection (non-white regions) mask = np.any(img_array != [255, 255, 255], axis=2) - + # Find contours and extract bounding boxes # (This is simplified - real implementation would need more sophisticated detection) - + # Save detected figures # ... implementation depends on specific needs ``` ### Batch PDF Processing with Error Handling + ```python import os import glob @@ -472,7 +493,7 @@ logger = logging.getLogger(__name__) def batch_process_pdfs(input_dir, operation='merge'): pdf_files = glob.glob(os.path.join(input_dir, "*.pdf")) - + if operation == 'merge': writer = PdfWriter() for pdf_file in pdf_files: @@ -484,10 +505,10 @@ def batch_process_pdfs(input_dir, operation='merge'): except Exception as e: logger.error(f"Failed to process {pdf_file}: {e}") continue - + with open("batch_merged.pdf", "wb") as output: writer.write(output) - + elif operation == 'extract_text': for pdf_file in pdf_files: try: @@ -495,18 +516,19 @@ def batch_process_pdfs(input_dir, operation='merge'): text = "" for page in reader.pages: text += page.extract_text() - + output_file = pdf_file.replace('.pdf', '.txt') with open(output_file, 'w', encoding='utf-8') as f: f.write(text) logger.info(f"Extracted text from: {pdf_file}") - + except Exception as e: logger.error(f"Failed to extract text from {pdf_file}: {e}") continue ``` ### Advanced PDF Cropping + ```python from pypdf import PdfWriter, PdfReader @@ -528,37 +550,42 @@ with open("cropped.pdf", "wb") as output: ## Performance Optimization Tips ### 1. For Large PDFs + - Use streaming approaches instead of loading entire PDF in memory - Use `qpdf --split-pages` for splitting large files - Process pages individually with pypdfium2 ### 2. For Text Extraction + - `pdftotext -bbox-layout` is fastest for plain text extraction - Use pdfplumber for structured data and tables - Avoid `pypdf.extract_text()` for very large documents ### 3. For Image Extraction + - `pdfimages` is much faster than rendering pages - Use low resolution for previews, high resolution for final output ### 4. For Form Filling + - pdf-lib maintains form structure better than most alternatives - Pre-validate form fields before processing ### 5. Memory Management + ```python # Process PDFs in chunks def process_large_pdf(pdf_path, chunk_size=10): reader = PdfReader(pdf_path) total_pages = len(reader.pages) - + for start_idx in range(0, total_pages, chunk_size): end_idx = min(start_idx + chunk_size, total_pages) writer = PdfWriter() - + for i in range(start_idx, end_idx): writer.add_page(reader.pages[i]) - + # Process chunk with open(f"chunk_{start_idx//chunk_size}.pdf", "wb") as output: writer.write(output) @@ -567,6 +594,7 @@ def process_large_pdf(pdf_path, chunk_size=10): ## Troubleshooting Common Issues ### Encrypted PDFs + ```python # Handle password-protected PDFs from pypdf import PdfReader @@ -580,6 +608,7 @@ except Exception as e: ``` ### Corrupted PDFs + ```bash # Use qpdf to repair qpdf --check corrupted.pdf @@ -587,6 +616,7 @@ qpdf --replace-input corrupted.pdf ``` ### Text Extraction Issues + ```python # Fallback to OCR for scanned PDFs import pytesseract @@ -609,4 +639,4 @@ def extract_text_with_ocr(pdf_path): - **poppler-utils**: GPL-2 License - **qpdf**: Apache License - **pdf-lib**: MIT License -- **pdfjs-dist**: Apache License \ No newline at end of file +- **pdfjs-dist**: Apache License diff --git a/examples/skills/skills/pdf/scripts/check_bounding_boxes.py b/examples/skills/skills/pdf/scripts/check_bounding_boxes.py index 7443660..f74a14d 100644 --- a/examples/skills/skills/pdf/scripts/check_bounding_boxes.py +++ b/examples/skills/skills/pdf/scripts/check_bounding_boxes.py @@ -1,7 +1,6 @@ -from dataclasses import dataclass import json import sys - +from dataclasses import dataclass # Script to check that the `fields.json` file that Claude creates when analyzing PDFs # does not have overlapping bounding boxes. See forms.md. @@ -35,12 +34,18 @@ def rects_intersect(r1, r2): # This is O(N^2); we can optimize if it becomes a problem. for j in range(i + 1, len(rects_and_fields)): rj = rects_and_fields[j] - if ri.field["page_number"] == rj.field["page_number"] and rects_intersect(ri.rect, rj.rect): + if ri.field["page_number"] == rj.field["page_number"] and rects_intersect( + ri.rect, rj.rect + ): has_error = True if ri.field is rj.field: - messages.append(f"FAILURE: intersection between label and entry bounding boxes for `{ri.field['description']}` ({ri.rect}, {rj.rect})") + messages.append( + f"FAILURE: intersection between label and entry bounding boxes for `{ri.field['description']}` ({ri.rect}, {rj.rect})" + ) else: - messages.append(f"FAILURE: intersection between {ri.rect_type} bounding box for `{ri.field['description']}` ({ri.rect}) and {rj.rect_type} bounding box for `{rj.field['description']}` ({rj.rect})") + messages.append( + f"FAILURE: intersection between {ri.rect_type} bounding box for `{ri.field['description']}` ({ri.rect}) and {rj.rect_type} bounding box for `{rj.field['description']}` ({rj.rect})" + ) if len(messages) >= 20: messages.append("Aborting further checks; fix bounding boxes and try again") return messages @@ -50,7 +55,9 @@ def rects_intersect(r1, r2): entry_height = ri.rect[3] - ri.rect[1] if entry_height < font_size: has_error = True - messages.append(f"FAILURE: entry bounding box height ({entry_height}) for `{ri.field['description']}` is too short for the text content (font size: {font_size}). Increase the box height or decrease the font size.") + messages.append( + f"FAILURE: entry bounding box height ({entry_height}) for `{ri.field['description']}` is too short for the text content (font size: {font_size}). Increase the box height or decrease the font size." + ) if len(messages) >= 20: messages.append("Aborting further checks; fix bounding boxes and try again") return messages @@ -59,6 +66,7 @@ def rects_intersect(r1, r2): messages.append("SUCCESS: All bounding boxes are valid") return messages + if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: check_bounding_boxes.py [fields.json]") diff --git a/examples/skills/skills/pdf/scripts/check_bounding_boxes_test.py b/examples/skills/skills/pdf/scripts/check_bounding_boxes_test.py index 1dbb463..cb53d92 100644 --- a/examples/skills/skills/pdf/scripts/check_bounding_boxes_test.py +++ b/examples/skills/skills/pdf/scripts/check_bounding_boxes_test.py @@ -1,16 +1,16 @@ -import unittest -import json import io +import json +import unittest + from check_bounding_boxes import get_bounding_box_messages # Currently this is not run automatically in CI; it's just for documentation and manual checking. class TestGetBoundingBoxMessages(unittest.TestCase): - def create_json_stream(self, data): """Helper to create a JSON stream from data""" return io.StringIO(json.dumps(data)) - + def test_no_intersections(self): """Test case with no bounding box intersections""" data = { @@ -19,22 +19,22 @@ def test_no_intersections(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 1, "label_bounding_box": [10, 40, 50, 60], - "entry_bounding_box": [60, 40, 150, 60] - } + "entry_bounding_box": [60, 40, 150, 60], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_label_entry_intersection_same_field(self): """Test intersection between label and entry of the same field""" data = { @@ -43,16 +43,16 @@ def test_label_entry_intersection_same_field(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 60, 30], - "entry_bounding_box": [50, 10, 150, 30] # Overlaps with label + "entry_bounding_box": [50, 10, 150, 30], # Overlaps with label } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "intersection" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_intersection_between_different_fields(self): """Test intersection between bounding boxes of different fields""" data = { @@ -61,22 +61,22 @@ def test_intersection_between_different_fields(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 1, "label_bounding_box": [40, 20, 80, 40], # Overlaps with Name's boxes - "entry_bounding_box": [160, 10, 250, 30] - } + "entry_bounding_box": [160, 10, 250, 30], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "intersection" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_different_pages_no_intersection(self): """Test that boxes on different pages don't count as intersecting""" data = { @@ -85,22 +85,22 @@ def test_different_pages_no_intersection(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 2, "label_bounding_box": [10, 10, 50, 30], # Same coordinates but different page - "entry_bounding_box": [60, 10, 150, 30] - } + "entry_bounding_box": [60, 10, 150, 30], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_entry_height_too_small(self): """Test that entry box height is checked against font size""" data = { @@ -112,16 +112,16 @@ def test_entry_height_too_small(self): "entry_bounding_box": [60, 10, 150, 20], # Height is 10 "entry_text": { "font_size": 14 # Font size larger than height - } + }, } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "height" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_entry_height_adequate(self): """Test that adequate entry box height passes""" data = { @@ -133,16 +133,16 @@ def test_entry_height_adequate(self): "entry_bounding_box": [60, 10, 150, 30], # Height is 20 "entry_text": { "font_size": 14 # Font size smaller than height - } + }, } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_default_font_size(self): """Test that default font size is used when not specified""" data = { @@ -152,16 +152,16 @@ def test_default_font_size(self): "page_number": 1, "label_bounding_box": [10, 10, 50, 30], "entry_bounding_box": [60, 10, 150, 20], # Height is 10 - "entry_text": {} # No font_size specified, should use default 14 + "entry_text": {}, # No font_size specified, should use default 14 } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "height" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_no_entry_text(self): """Test that missing entry_text doesn't cause height check""" data = { @@ -170,30 +170,32 @@ def test_no_entry_text(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 20] # Small height but no entry_text + "entry_bounding_box": [60, 10, 150, 20], # Small height but no entry_text } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_multiple_errors_limit(self): """Test that error messages are limited to prevent excessive output""" fields = [] # Create many overlapping fields for i in range(25): - fields.append({ - "description": f"Field{i}", - "page_number": 1, - "label_bounding_box": [10, 10, 50, 30], # All overlap - "entry_bounding_box": [20, 15, 60, 35] # All overlap - }) - + fields.append( + { + "description": f"Field{i}", + "page_number": 1, + "label_bounding_box": [10, 10, 50, 30], # All overlap + "entry_bounding_box": [20, 15, 60, 35], # All overlap + } + ) + data = {"form_fields": fields} - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) # Should abort after ~20 messages @@ -202,7 +204,7 @@ def test_multiple_errors_limit(self): failure_count = sum(1 for msg in messages if "FAILURE" in msg) self.assertGreater(failure_count, 0) self.assertLess(len(messages), 30) # Should be limited - + def test_edge_touching_boxes(self): """Test that boxes touching at edges don't count as intersecting""" data = { @@ -211,16 +213,16 @@ def test_edge_touching_boxes(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [50, 10, 150, 30] # Touches at x=50 + "entry_bounding_box": [50, 10, 150, 30], # Touches at x=50 } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() diff --git a/examples/skills/skills/pdf/scripts/check_fillable_fields.py b/examples/skills/skills/pdf/scripts/check_fillable_fields.py index dc43d18..ccfec8a 100644 --- a/examples/skills/skills/pdf/scripts/check_fillable_fields.py +++ b/examples/skills/skills/pdf/scripts/check_fillable_fields.py @@ -1,12 +1,14 @@ import sys -from pypdf import PdfReader +from pypdf import PdfReader # Script for Claude to run to determine whether a PDF has fillable form fields. See forms.md. reader = PdfReader(sys.argv[1]) -if (reader.get_fields()): +if reader.get_fields(): print("This PDF has fillable form fields") else: - print("This PDF does not have fillable form fields; you will need to visually determine where to enter data") + print( + "This PDF does not have fillable form fields; you will need to visually determine where to enter data" + ) diff --git a/examples/skills/skills/pdf/scripts/convert_pdf_to_images.py b/examples/skills/skills/pdf/scripts/convert_pdf_to_images.py index f8a4ec5..0c2381b 100644 --- a/examples/skills/skills/pdf/scripts/convert_pdf_to_images.py +++ b/examples/skills/skills/pdf/scripts/convert_pdf_to_images.py @@ -3,7 +3,6 @@ from pdf2image import convert_from_path - # Converts each page of a PDF to a PNG image. @@ -18,7 +17,7 @@ def convert(pdf_path, output_dir, max_dim=1000): new_width = int(width * scale_factor) new_height = int(height * scale_factor) image = image.resize((new_width, new_height)) - + image_path = os.path.join(output_dir, f"page_{i+1}.png") image.save(image_path) print(f"Saved page {i+1} as {image_path} (size: {image.size})") diff --git a/examples/skills/skills/pdf/scripts/create_validation_image.py b/examples/skills/skills/pdf/scripts/create_validation_image.py index 4913f8f..9b78ffc 100644 --- a/examples/skills/skills/pdf/scripts/create_validation_image.py +++ b/examples/skills/skills/pdf/scripts/create_validation_image.py @@ -3,36 +3,37 @@ from PIL import Image, ImageDraw - # Creates "validation" images with rectangles for the bounding box information that # Claude creates when determining where to add text annotations in PDFs. See forms.md. def create_validation_image(page_number, fields_json_path, input_path, output_path): # Input file should be in the `fields.json` format described in forms.md. - with open(fields_json_path, 'r') as f: + with open(fields_json_path) as f: data = json.load(f) img = Image.open(input_path) draw = ImageDraw.Draw(img) num_boxes = 0 - + for field in data["form_fields"]: if field["page_number"] == page_number: - entry_box = field['entry_bounding_box'] - label_box = field['label_bounding_box'] + entry_box = field["entry_bounding_box"] + label_box = field["label_bounding_box"] # Draw red rectangle over entry bounding box and blue rectangle over the label. - draw.rectangle(entry_box, outline='red', width=2) - draw.rectangle(label_box, outline='blue', width=2) + draw.rectangle(entry_box, outline="red", width=2) + draw.rectangle(label_box, outline="blue", width=2) num_boxes += 2 - + img.save(output_path) print(f"Created validation image at {output_path} with {num_boxes} bounding boxes") if __name__ == "__main__": if len(sys.argv) != 5: - print("Usage: create_validation_image.py [page number] [fields.json file] [input image path] [output image path]") + print( + "Usage: create_validation_image.py [page number] [fields.json file] [input image path] [output image path]" + ) sys.exit(1) page_number = int(sys.argv[1]) fields_json_path = sys.argv[2] diff --git a/examples/skills/skills/pdf/scripts/extract_form_field_info.py b/examples/skills/skills/pdf/scripts/extract_form_field_info.py index f42a2df..c1b2e03 100644 --- a/examples/skills/skills/pdf/scripts/extract_form_field_info.py +++ b/examples/skills/skills/pdf/scripts/extract_form_field_info.py @@ -3,7 +3,6 @@ from pypdf import PdfReader - # Extracts data for the fillable form fields in a PDF and outputs JSON that # Claude uses to fill the fields. See forms.md. @@ -12,16 +11,16 @@ def get_full_annotation_field_id(annotation): components = [] while annotation: - field_name = annotation.get('/T') + field_name = annotation.get("/T") if field_name: components.append(field_name) - annotation = annotation.get('/Parent') + annotation = annotation.get("/Parent") return ".".join(reversed(components)) if components else None def make_field_dict(field, field_id): field_dict = {"field_id": field_id} - ft = field.get('/FT') + ft = field.get("/FT") if ft == "/Tx": field_dict["type"] = "text" elif ft == "/Btn": @@ -35,16 +34,21 @@ def make_field_dict(field, field_id): field_dict["checked_value"] = states[0] if states[0] != "/Off" else states[1] field_dict["unchecked_value"] = "/Off" else: - print(f"Unexpected state values for checkbox `${field_id}`. Its checked and unchecked values may not be correct; if you're trying to check it, visually verify the results.") + print( + f"Unexpected state values for checkbox `${field_id}`. Its checked and unchecked values may not be correct; if you're trying to check it, visually verify the results." + ) field_dict["checked_value"] = states[0] field_dict["unchecked_value"] = states[1] elif ft == "/Ch": field_dict["type"] = "choice" states = field.get("/_States_", []) - field_dict["choice_options"] = [{ - "value": state[0], - "text": state[1], - } for state in states] + field_dict["choice_options"] = [ + { + "value": state[0], + "text": state[1], + } + for state in states + ] else: field_dict["type"] = f"unknown ({ft})" return field_dict @@ -82,12 +86,12 @@ def get_field_info(reader: PdfReader): radio_fields_by_id = {} for page_index, page in enumerate(reader.pages): - annotations = page.get('/Annots', []) + annotations = page.get("/Annots", []) for ann in annotations: field_id = get_full_annotation_field_id(ann) if field_id in field_info_by_id: field_info_by_id[field_id]["page"] = page_index + 1 - field_info_by_id[field_id]["rect"] = ann.get('/Rect') + field_info_by_id[field_id]["rect"] = ann.get("/Rect") elif field_id in possible_radio_names: try: # ann['/AP']['/N'] should have two items. One of them is '/Off', @@ -108,10 +112,12 @@ def get_field_info(reader: PdfReader): # radio buttons correctly. (It does if you remove the leading slash # from the value, but that causes them not to appear correctly in # Chrome/Firefox/Acrobat/etc). - radio_fields_by_id[field_id]["radio_options"].append({ - "value": on_values[0], - "rect": rect, - }) + radio_fields_by_id[field_id]["radio_options"].append( + { + "value": on_values[0], + "rect": rect, + } + ) # Some PDFs have form field definitions without corresponding annotations, # so we can't tell where they are. Ignore these fields for now. @@ -120,7 +126,9 @@ def get_field_info(reader: PdfReader): if "page" in field_info: fields_with_location.append(field_info) else: - print(f"Unable to determine location for field id: {field_info.get('field_id')}, ignoring") + print( + f"Unable to determine location for field id: {field_info.get('field_id')}, ignoring" + ) # Sort by page number, then Y position (flipped in PDF coordinate system), then X. def sort_key(f): @@ -130,7 +138,7 @@ def sort_key(f): rect = f.get("rect") or [0, 0, 0, 0] adjusted_position = [-rect[1], rect[0]] return [f.get("page"), adjusted_position] - + sorted_fields = fields_with_location + list(radio_fields_by_id.values()) sorted_fields.sort(key=sort_key) diff --git a/examples/skills/skills/pdf/scripts/fill_fillable_fields.py b/examples/skills/skills/pdf/scripts/fill_fillable_fields.py index ac35753..829e759 100644 --- a/examples/skills/skills/pdf/scripts/fill_fillable_fields.py +++ b/examples/skills/skills/pdf/scripts/fill_fillable_fields.py @@ -1,10 +1,8 @@ import json import sys -from pypdf import PdfReader, PdfWriter - from extract_form_field_info import get_field_info - +from pypdf import PdfReader, PdfWriter # Fills fillable form fields in a PDF. See forms.md. @@ -21,7 +19,7 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: if page not in fields_by_page: fields_by_page[page] = {} fields_by_page[page][field_id] = field["value"] - + reader = PdfReader(input_pdf_path) has_error = False @@ -34,7 +32,9 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: print(f"ERROR: `{field['field_id']}` is not a valid field ID") elif field["page"] != existing_field["page"]: has_error = True - print(f"ERROR: Incorrect page number for `{field['field_id']}` (got {field['page']}, expected {existing_field['page']})") + print( + f"ERROR: Incorrect page number for `{field['field_id']}` (got {field['page']}, expected {existing_field['page']})" + ) else: if "value" in field: err = validation_error_for_field_value(existing_field, field["value"]) @@ -46,12 +46,14 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: writer = PdfWriter(clone_from=reader) for page, field_values in fields_by_page.items(): - writer.update_page_form_field_values(writer.pages[page - 1], field_values, auto_regenerate=False) + writer.update_page_form_field_values( + writer.pages[page - 1], field_values, auto_regenerate=False + ) # This seems to be necessary for many PDF viewers to format the form values correctly. # It may cause the viewer to show a "save changes" dialog even if the user doesn't make any changes. writer.set_need_appearances_writer(True) - + with open(output_pdf_path, "wb") as f: writer.write(f) @@ -67,7 +69,7 @@ def validation_error_for_field_value(field_info, field_value): elif field_type == "radio_group": option_values = [opt["value"] for opt in field_info["radio_options"]] if field_value not in option_values: - return f'ERROR: Invalid value "{field_value}" for radio group field "{field_id}". Valid values are: {option_values}' + return f'ERROR: Invalid value "{field_value}" for radio group field "{field_id}". Valid values are: {option_values}' elif field_type == "choice": choice_values = [opt["value"] for opt in field_info["choice_options"]] if field_value not in choice_values: @@ -88,15 +90,17 @@ def validation_error_for_field_value(field_info, field_value): # We call the original method and adjust the return value only if the argument to `get_inherited` # is `FA.Opt` and if the return value is a list of two-element lists. def monkeypatch_pydpf_method(): - from pypdf.generic import DictionaryObject from pypdf.constants import FieldDictionaryAttributes + from pypdf.generic import DictionaryObject original_get_inherited = DictionaryObject.get_inherited - def patched_get_inherited(self, key: str, default = None): + def patched_get_inherited(self, key: str, default=None): result = original_get_inherited(self, key, default) if key == FieldDictionaryAttributes.Opt: - if isinstance(result, list) and all(isinstance(v, list) and len(v) == 2 for v in result): + if isinstance(result, list) and all( + isinstance(v, list) and len(v) == 2 for v in result + ): result = [r[0] for r in result] return result diff --git a/examples/skills/skills/pdf/scripts/fill_pdf_form_with_annotations.py b/examples/skills/skills/pdf/scripts/fill_pdf_form_with_annotations.py index f980531..9cf8d8f 100644 --- a/examples/skills/skills/pdf/scripts/fill_pdf_form_with_annotations.py +++ b/examples/skills/skills/pdf/scripts/fill_pdf_form_with_annotations.py @@ -4,7 +4,6 @@ from pypdf import PdfReader, PdfWriter from pypdf.annotations import FreeText - # Fills a PDF by adding text annotations defined in `fields.json`. See forms.md. @@ -14,54 +13,52 @@ def transform_coordinates(bbox, image_width, image_height, pdf_width, pdf_height # PDF coordinates: origin at bottom-left, y increases upward x_scale = pdf_width / image_width y_scale = pdf_height / image_height - + left = bbox[0] * x_scale right = bbox[2] * x_scale - + # Flip Y coordinates for PDF top = pdf_height - (bbox[1] * y_scale) bottom = pdf_height - (bbox[3] * y_scale) - + return left, bottom, right, top def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): """Fill the PDF form with data from fields.json""" - + # `fields.json` format described in forms.md. - with open(fields_json_path, "r") as f: + with open(fields_json_path) as f: fields_data = json.load(f) - + # Open the PDF reader = PdfReader(input_pdf_path) writer = PdfWriter() - + # Copy all pages to writer writer.append(reader) - + # Get PDF dimensions for each page pdf_dimensions = {} for i, page in enumerate(reader.pages): mediabox = page.mediabox pdf_dimensions[i + 1] = [mediabox.width, mediabox.height] - + # Process each form field annotations = [] for field in fields_data["form_fields"]: page_num = field["page_number"] - + # Get page dimensions and transform coordinates. page_info = next(p for p in fields_data["pages"] if p["page_number"] == page_num) image_width = page_info["image_width"] image_height = page_info["image_height"] pdf_width, pdf_height = pdf_dimensions[page_num] - + transformed_entry_box = transform_coordinates( - field["entry_bounding_box"], - image_width, image_height, - pdf_width, pdf_height + field["entry_bounding_box"], image_width, image_height, pdf_width, pdf_height ) - + # Skip empty fields if "entry_text" not in field or "text" not in field["entry_text"]: continue @@ -69,7 +66,7 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): text = entry_text["text"] if not text: continue - + font_name = entry_text.get("font", "Arial") font_size = str(entry_text.get("font_size", 14)) + "pt" font_color = entry_text.get("font_color", "000000") @@ -88,11 +85,11 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): annotations.append(annotation) # page_number is 0-based for pypdf writer.add_annotation(page_number=page_num - 1, annotation=annotation) - + # Save the filled PDF with open(output_pdf_path, "wb") as output: writer.write(output) - + print(f"Successfully filled PDF form and saved to {output_pdf_path}") print(f"Added {len(annotations)} text annotations") @@ -104,5 +101,5 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): input_pdf = sys.argv[1] fields_json = sys.argv[2] output_pdf = sys.argv[3] - - fill_pdf_form(input_pdf, fields_json, output_pdf) \ No newline at end of file + + fill_pdf_form(input_pdf, fields_json, output_pdf) diff --git a/pyproject.toml b/pyproject.toml index 8ee6dac..849f52a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,13 @@ pydantic-ai = [ "pydantic-graph>=1.94.0", "pydantic-ai-slim>=1.94.0", ] -test = ["ipykernel", "jupyter_server>=1.6,<3", "pytest>=7.0", "pytest-asyncio>=0.21"] +test = [ + "ipykernel", + "jupyter_server>=1.6,<3", + "pytest>=7.0", + "pytest-asyncio>=0.21", + "pytest-cov>=4.1", +] lint = ["mdformat>0.7", "mdformat-gfm>=0.3.5", "ruff"] typing = ["mypy>=0.990"] @@ -103,4 +109,20 @@ ignore = [ [tool.ruff.lint.per-file-ignores] # S101 Use of `assert` detected -"agent_codemode/tests/*" = ["S101"] \ No newline at end of file +"tests/*" = ["S101"] +"tests/test_stdout_capture.py" = ["S101", "T201", "E501"] +"tests/test_direct_execution.py" = ["T201"] +"tests/test_executor_async.py" = ["T201"] +"tests/test_skill_bindings.py" = ["S102", "F841"] +"tests/test_agent_minimal.py" = ["T201"] +"tests/test_codemode.py" = ["E731", "E501"] +"skills/**/scripts/*.py" = ["E501", "C901", "T201", "RUF013", "E741"] +"examples/skills/skills/**/scripts/*.py" = ["E501", "C901", "T201", "RUF013", "E741"] +"examples/**/*.py" = ["T201", "S108", "S311", "UP007", "E501", "C901", "E741", "RUF013", "S110", "C408", "F841", "F821"] +"agent_codemode/toolset.py" = ["UP007", "E501", "C901"] +"agent_codemode/types.py" = ["UP007"] +"agent_codemode/proxy/mcp_client.py" = ["S110"] +"agent_codemode/proxy/meta_tools.py" = ["UP007"] +"agent_codemode/tool_definitions.py" = ["E501"] +"agent_codemode/discovery/codegen.py" = ["RUF001", "RUF002", "E501", "B007", "E402"] +"agent_codemode/composition/executor.py" = ["E501", "S108", "C901", "S102", "T201", "B007"] diff --git a/skills/crawl/SKILL.md b/skills/crawl/SKILL.md index def8345..fb2ae01 100644 --- a/skills/crawl/SKILL.md +++ b/skills/crawl/SKILL.md @@ -41,6 +41,7 @@ print(text) ### httpx - HTTP Client #### Basic GET Request + ```python import httpx @@ -51,6 +52,7 @@ print(response.text) ``` #### GET with Headers and Parameters + ```python headers = { "User-Agent": "Mozilla/5.0 (compatible; DataBot/1.0)", @@ -67,6 +69,7 @@ response = httpx.get( ``` #### Handle Redirects and Errors + ```python import httpx @@ -85,6 +88,7 @@ except httpx.RequestError as e: ``` #### Async Requests + ```python import httpx import asyncio @@ -103,6 +107,7 @@ results = asyncio.run(fetch_pages(urls)) ### BeautifulSoup - HTML Parsing #### Parse HTML and Extract Elements + ```python from bs4 import BeautifulSoup @@ -130,6 +135,7 @@ links = [(a.text, a.get("href")) for a in soup.find_all("a")] ``` #### Extract Text Content + ```python from bs4 import BeautifulSoup @@ -145,6 +151,7 @@ if main_content: ``` #### Extract Links with Absolute URLs + ```python from urllib.parse import urljoin from bs4 import BeautifulSoup @@ -163,6 +170,7 @@ for a in soup.find_all("a", href=True): ``` #### Extract Tables + ```python from bs4 import BeautifulSoup @@ -178,6 +186,7 @@ for table in soup.find_all("table"): ``` #### CSS Selectors + ```python from bs4 import BeautifulSoup @@ -193,6 +202,7 @@ first_para = soup.select_one("p") ## Complete Crawling Example ### Single Page Crawler + ```python import httpx from bs4 import BeautifulSoup @@ -246,6 +256,7 @@ print(f"Links found: {len(page.links)}") ``` ### Multi-Page Crawler with Depth Control + ```python import httpx from bs4 import BeautifulSoup @@ -320,6 +331,7 @@ for url, data in pages.items(): ## Handling Special Cases ### JavaScript-Rendered Pages + For pages that require JavaScript execution, use Playwright: ```python @@ -337,6 +349,7 @@ def crawl_js_page(url: str) -> str: ``` ### Respecting robots.txt + ```python from urllib.robotparser import RobotFileParser from urllib.parse import urljoin @@ -357,6 +370,7 @@ def can_crawl(url: str, user_agent: str = "*") -> bool: ``` ### Rate Limiting + ```python import time import httpx @@ -385,14 +399,14 @@ class RateLimitedClient: This skill includes ready-to-use scripts in the `scripts/` folder: -| Script | Description | -|--------|-------------| -| `fetch_page.py` | Fetch a single webpage and extract text content | -| `extract_links.py` | Extract all links from a webpage with filtering | -| `extract_tables.py` | Extract HTML tables as JSON or CSV | -| `crawl_site.py` | Multi-page crawler with depth control | -| `check_robots.py` | Check if URL is allowed by robots.txt | -| `fetch_js_page.py` | Fetch JavaScript-rendered pages using Playwright | +| Script | Description | +| ------------------- | ------------------------------------------------ | +| `fetch_page.py` | Fetch a single webpage and extract text content | +| `extract_links.py` | Extract all links from a webpage with filtering | +| `extract_tables.py` | Extract HTML tables as JSON or CSV | +| `crawl_site.py` | Multi-page crawler with depth control | +| `check_robots.py` | Check if URL is allowed by robots.txt | +| `fetch_js_page.py` | Fetch JavaScript-rendered pages using Playwright | ### Usage Examples diff --git a/skills/crawl/scripts/check_robots.py b/skills/crawl/scripts/check_robots.py index fbe239e..82ce999 100644 --- a/skills/crawl/scripts/check_robots.py +++ b/skills/crawl/scripts/check_robots.py @@ -59,9 +59,7 @@ def check_robots(url: str, user_agent: str = "*") -> dict: def main(): - parser = argparse.ArgumentParser( - description="Check if a URL is allowed by robots.txt." - ) + parser = argparse.ArgumentParser(description="Check if a URL is allowed by robots.txt.") parser.add_argument("url", help="The URL to check") parser.add_argument( "--user-agent", diff --git a/skills/crawl/scripts/crawl_site.py b/skills/crawl/scripts/crawl_site.py index bdc318c..91e34b6 100644 --- a/skills/crawl/scripts/crawl_site.py +++ b/skills/crawl/scripts/crawl_site.py @@ -52,9 +52,7 @@ def crawl_site( "User-Agent": "Mozilla/5.0 (compatible; DataBot/1.0)", } - with httpx.Client( - headers=headers, timeout=timeout, follow_redirects=True - ) as client: + with httpx.Client(headers=headers, timeout=timeout, follow_redirects=True) as client: while queue and len(visited) < max_pages: url, depth = queue.popleft() @@ -88,11 +86,7 @@ def crawl_site( element.decompose() # Extract title - title = ( - soup.title.string.strip() - if soup.title and soup.title.string - else "" - ) + title = soup.title.string.strip() if soup.title and soup.title.string else "" # Extract text (limit to first 5000 chars for efficiency) text = soup.get_text(separator="\n", strip=True)[:5000] @@ -110,9 +104,7 @@ def crawl_site( href = a.get("href", "").strip() # Skip empty, anchors, javascript, etc. - if not href or href.startswith( - ("#", "javascript:", "mailto:", "tel:") - ): + if not href or href.startswith(("#", "javascript:", "mailto:", "tel:")): continue next_url = urljoin(url, href) diff --git a/skills/crawl/scripts/extract_links.py b/skills/crawl/scripts/extract_links.py index 6e765ed..eee6479 100644 --- a/skills/crawl/scripts/extract_links.py +++ b/skills/crawl/scripts/extract_links.py @@ -16,7 +16,7 @@ import argparse import json import sys -from urllib.parse import urljoin, urlparse +from urllib.parse import urljoin import httpx from bs4 import BeautifulSoup @@ -101,9 +101,7 @@ def main(): default=True, help="Convert relative URLs to absolute (default: True)", ) - parser.add_argument( - "--filter", "-f", help="Only include links containing this pattern" - ) + parser.add_argument("--filter", "-f", help="Only include links containing this pattern") parser.add_argument( "--timeout", "-t", type=float, default=30.0, help="Request timeout in seconds" ) diff --git a/skills/crawl/scripts/fetch_js_page.py b/skills/crawl/scripts/fetch_js_page.py index d8c4272..268be8a 100644 --- a/skills/crawl/scripts/fetch_js_page.py +++ b/skills/crawl/scripts/fetch_js_page.py @@ -104,9 +104,7 @@ def main(): default=0, help="Additional wait time in milliseconds after page load", ) - parser.add_argument( - "--selector", "-s", help="Wait for a specific CSS selector to appear" - ) + parser.add_argument("--selector", "-s", help="Wait for a specific CSS selector to appear") parser.add_argument( "--timeout", "-t", @@ -114,9 +112,7 @@ def main(): default=30000, help="Page load timeout in milliseconds (default: 30000)", ) - parser.add_argument( - "--html", action="store_true", help="Output raw HTML instead of text" - ) + parser.add_argument("--html", action="store_true", help="Output raw HTML instead of text") args = parser.parse_args() diff --git a/skills/crawl/scripts/fetch_page.py b/skills/crawl/scripts/fetch_page.py index abcbb55..ba2836b 100644 --- a/skills/crawl/scripts/fetch_page.py +++ b/skills/crawl/scripts/fetch_page.py @@ -72,9 +72,7 @@ def fetch_page(url: str, timeout: float = 30.0) -> dict: def main(): - parser = argparse.ArgumentParser( - description="Fetch a webpage and extract its text content." - ) + parser = argparse.ArgumentParser(description="Fetch a webpage and extract its text content.") parser.add_argument("url", help="The URL of the webpage to fetch") parser.add_argument( "--output", "-o", help="Output file path (prints to stdout if not specified)" diff --git a/skills/github/SKILL.md b/skills/github/SKILL.md index 2142d6a..d14145d 100644 --- a/skills/github/SKILL.md +++ b/skills/github/SKILL.md @@ -24,17 +24,19 @@ This skill provides tools to interact with the GitHub API using Python. It enabl This skill uses GitHub OAuth authentication through the Datalayer identity system. When you connect your GitHub account via the **"Connect GitHub"** button in the agent configuration form, the OAuth token is automatically available to this skill. **How it works:** + 1. Click the **GitHub** connect button in the Agent Configuration form -2. Complete the OAuth authorization flow -3. Your GitHub access token is securely stored and automatically provided to skill scripts +1. Complete the OAuth authorization flow +1. Your GitHub access token is securely stored and automatically provided to skill scripts The token is retrieved from the identity store and made available as the `GITHUB_TOKEN` environment variable when scripts are executed. #### Required OAuth Scopes When connecting your GitHub account, ensure the following scopes are granted: + - `repo` - Full control of private repositories (required for private repos) -- `read:user` - Read user profile data +- `read:user` - Read user profile data - `read:org` - Read organization data (if listing org repos) ### Manual Token (Fallback) @@ -46,19 +48,20 @@ export GITHUB_TOKEN="ghp_your_token_here" ``` To create a Personal Access Token: + 1. Go to [GitHub Settings > Developer settings > Personal access tokens](https://github.com/settings/tokens) -2. Click **"Generate new token"** (classic) -3. Select the scopes listed above -4. Generate and copy the token +1. Click **"Generate new token"** (classic) +1. Select the scopes listed above +1. Generate and copy the token ### Token Permissions -| Scope | Description | Required For | -|-------|-------------|--------------| -| `repo` | Full control of private repositories | Private repos, creating issues/PRs | -| `public_repo` | Access public repositories only | Public repos only (limited) | -| `read:user` | Read user profile data | User information | -| `read:org` | Read org membership | Organization repos | +| Scope | Description | Required For | +| ------------- | ------------------------------------ | ---------------------------------- | +| `repo` | Full control of private repositories | Private repos, creating issues/PRs | +| `public_repo` | Access public repositories only | Public repos only (limited) | +| `read:user` | Read user profile data | User information | +| `read:org` | Read org membership | Organization repos | > **Note:** The OAuth flow requests `repo`, `read:user`, and `read:org` scopes by default. @@ -93,6 +96,7 @@ for repo in repos: ### httpx - GitHub API Client #### List All User Repositories + ```python import os import httpx @@ -114,19 +118,19 @@ def list_repos( per_page: int = 100, ) -> list[dict]: """List repositories for the authenticated user. - + Args: visibility: Filter by visibility - 'all', 'public', or 'private' sort: Sort by 'created', 'updated', 'pushed', or 'full_name' per_page: Number of results per page (max 100) - + Returns: List of repository dictionaries. """ headers = get_github_headers() repos = [] page = 1 - + while True: response = httpx.get( "https://api.github.com/user/repos", @@ -141,25 +145,26 @@ def list_repos( ) response.raise_for_status() page_repos = response.json() - + if not page_repos: break - + repos.extend(page_repos) page += 1 - + return repos ``` #### Get Repository Details + ```python def get_repo(owner: str, repo: str) -> dict: """Get details for a specific repository. - + Args: owner: Repository owner (username or organization) repo: Repository name - + Returns: Repository details dictionary. """ @@ -174,21 +179,22 @@ def get_repo(owner: str, repo: str) -> dict: ``` #### List Organization Repositories + ```python def list_org_repos(org: str, type: str = "all") -> list[dict]: """List repositories for an organization. - + Args: org: Organization name type: Filter by type - 'all', 'public', 'private', 'forks', 'sources', 'member' - + Returns: List of repository dictionaries. """ headers = get_github_headers() repos = [] page = 1 - + while True: response = httpx.get( f"https://api.github.com/orgs/{org}/repos", @@ -198,28 +204,29 @@ def list_org_repos(org: str, type: str = "all") -> list[dict]: ) response.raise_for_status() page_repos = response.json() - + if not page_repos: break - + repos.extend(page_repos) page += 1 - + return repos ``` ### Working with Issues #### List Repository Issues + ```python def list_issues(owner: str, repo: str, state: str = "open") -> list[dict]: """List issues for a repository. - + Args: owner: Repository owner repo: Repository name state: Issue state - 'open', 'closed', or 'all' - + Returns: List of issue dictionaries. """ @@ -235,6 +242,7 @@ def list_issues(owner: str, repo: str, state: str = "open") -> list[dict]: ``` #### Create an Issue + ```python def create_issue( owner: str, @@ -244,14 +252,14 @@ def create_issue( labels: list[str] = None, ) -> dict: """Create a new issue. - + Args: owner: Repository owner repo: Repository name title: Issue title body: Issue body (markdown) labels: List of label names - + Returns: Created issue dictionary. """ @@ -259,7 +267,7 @@ def create_issue( data = {"title": title, "body": body} if labels: data["labels"] = labels - + response = httpx.post( f"https://api.github.com/repos/{owner}/{repo}/issues", headers=headers, @@ -273,6 +281,7 @@ def create_issue( ### Working with Pull Requests #### List Pull Requests + ```python def list_pull_requests( owner: str, @@ -280,12 +289,12 @@ def list_pull_requests( state: str = "open", ) -> list[dict]: """List pull requests for a repository. - + Args: owner: Repository owner repo: Repository name state: PR state - 'open', 'closed', or 'all' - + Returns: List of pull request dictionaries. """ @@ -329,6 +338,7 @@ def safe_github_request(url: str, headers: dict) -> dict | None: ## Rate Limiting GitHub API has rate limits: + - **Authenticated requests**: 5,000 requests per hour - **Unauthenticated requests**: 60 requests per hour @@ -350,6 +360,7 @@ def check_rate_limit() -> dict: ## Common Patterns ### Filter Repositories by Language + ```python def list_repos_by_language(language: str) -> list[dict]: """List user repositories filtered by primary language.""" @@ -358,14 +369,15 @@ def list_repos_by_language(language: str) -> list[dict]: ``` ### Get Recent Activity + ```python def get_recent_repos(days: int = 30) -> list[dict]: """Get repositories updated in the last N days.""" from datetime import datetime, timedelta - + cutoff = datetime.now() - timedelta(days=days) repos = list_repos(sort="updated") - + recent = [] for repo in repos: updated = datetime.fromisoformat(repo["updated_at"].replace("Z", "+00:00")) @@ -378,26 +390,29 @@ def get_recent_repos(days: int = 30) -> list[dict]: This skill includes the following executable scripts in the `scripts/` directory: -| Script | Description | Usage | -|--------|-------------|-------| -| `list_repos.py` | List all repositories for the authenticated user | `python list_repos.py [--visibility all\|public\|private] [--format table\|json]` | -| `get_repo.py` | Get details for a specific repository | `python get_repo.py /` | -| `list_issues.py` | List issues for a repository | `python list_issues.py / [--state open\|closed\|all]` | -| `list_prs.py` | List pull requests for a repository | `python list_prs.py / [--state open\|closed\|all]` | -| `search_repos.py` | Search GitHub repositories | `python search_repos.py [--language python]` | +| Script | Description | Usage | +| ----------------- | ------------------------------------------------ | --------------------------------------------------------------------------------- | +| `list_repos.py` | List all repositories for the authenticated user | `python list_repos.py [--visibility all\|public\|private] [--format table\|json]` | +| `get_repo.py` | Get details for a specific repository | `python get_repo.py /` | +| `list_issues.py` | List issues for a repository | `python list_issues.py / [--state open\|closed\|all]` | +| `list_prs.py` | List pull requests for a repository | `python list_prs.py / [--state open\|closed\|all]` | +| `search_repos.py` | Search GitHub repositories | `python search_repos.py [--language python]` | ## Troubleshooting ### "Bad credentials" Error + - Ensure `GITHUB_TOKEN` is set correctly - Check if the token has expired - Verify the token has required scopes ### "Not Found" Error for Private Repos + - The token needs `repo` scope for private repository access - Verify you have access to the repository ### Rate Limit Exceeded + - Wait for the rate limit to reset (check `X-RateLimit-Reset` header) - Use conditional requests with `If-None-Match` header - Consider using GraphQL API for complex queries (higher limits) diff --git a/skills/github/scripts/get_repo.py b/skills/github/scripts/get_repo.py index 72397f0..370d46b 100644 --- a/skills/github/scripts/get_repo.py +++ b/skills/github/scripts/get_repo.py @@ -37,11 +37,11 @@ def get_github_headers() -> dict: def get_repo(owner: str, repo: str) -> dict: """Get details for a specific repository. - + Args: owner: Repository owner (username or organization) repo: Repository name - + Returns: Repository details dictionary. """ @@ -51,7 +51,7 @@ def get_repo(owner: str, repo: str) -> dict: headers=headers, timeout=30.0, ) - + if response.status_code == 401: print("Error: Invalid or expired GITHUB_TOKEN", file=sys.stderr) sys.exit(1) @@ -62,7 +62,7 @@ def get_repo(owner: str, repo: str) -> dict: print(f"Error: Repository '{owner}/{repo}' not found", file=sys.stderr) print("(Check if you have access and the token has 'repo' scope)", file=sys.stderr) sys.exit(1) - + response.raise_for_status() return response.json() @@ -70,52 +70,52 @@ def get_repo(owner: str, repo: str) -> dict: def format_repo_details(repo: dict) -> str: """Format repository details as readable text.""" lines = [] - + visibility = "🔒 Private" if repo["private"] else "🌐 Public" - + lines.append(f"Repository: {repo['full_name']} ({visibility})") lines.append("=" * 60) - + if repo.get("description"): lines.append(f"Description: {repo['description']}") - + lines.append(f"URL: {repo['html_url']}") lines.append(f"Clone: {repo['clone_url']}") - + if repo.get("homepage"): lines.append(f"Homepage: {repo['homepage']}") - + lines.append("") lines.append("Statistics:") lines.append(f" ⭐ Stars: {repo.get('stargazers_count', 0)}") lines.append(f" 🍴 Forks: {repo.get('forks_count', 0)}") lines.append(f" 👀 Watchers: {repo.get('watchers_count', 0)}") lines.append(f" 🐛 Open Issues: {repo.get('open_issues_count', 0)}") - + lines.append("") lines.append("Details:") lines.append(f" Language: {repo.get('language') or 'Not specified'}") lines.append(f" Default Branch: {repo.get('default_branch', 'main')}") lines.append(f" License: {repo.get('license', {}).get('name') or 'Not specified'}") - + lines.append("") lines.append("Dates:") lines.append(f" Created: {repo.get('created_at', '')[:10]}") lines.append(f" Updated: {repo.get('updated_at', '')[:10]}") lines.append(f" Pushed: {repo.get('pushed_at', '')[:10]}") - + # Topics/tags if repo.get("topics"): lines.append("") lines.append(f"Topics: {', '.join(repo['topics'])}") - + return "\n".join(lines) class _HelpOnErrorParser(argparse.ArgumentParser): """ArgumentParser that prints full help on invalid arguments.""" - def error(self, message: str) -> None: # noqa: D401 + def error(self, message: str) -> None: params_help = ( "\nValid parameters for get_repo:\n" " repo (positional) owner/repo e.g. 'datalayer/jupyter-ui'\n" @@ -129,9 +129,7 @@ def error(self, message: str) -> None: # noqa: D401 def main(): - parser = _HelpOnErrorParser( - description="Get details for a specific GitHub repository." - ) + parser = _HelpOnErrorParser(description="Get details for a specific GitHub repository.") parser.add_argument( "repo", help="Repository in format 'owner/repo' (e.g., 'datalayer/jupyter-ui')", @@ -142,19 +140,19 @@ def main(): default="table", help="Output format (default: table)", ) - + args = parser.parse_args() - + # Parse owner/repo if "/" not in args.repo: print("Error: Repository must be in format 'owner/repo'", file=sys.stderr) sys.exit(1) - + owner, repo = args.repo.split("/", 1) - + try: repo_data = get_repo(owner, repo) - + if args.format == "json": # Output simplified JSON simplified = { @@ -178,7 +176,7 @@ def main(): print(json.dumps(simplified, indent=2)) else: print(format_repo_details(repo_data)) - + except httpx.HTTPStatusError as e: print(f"HTTP error: {e.response.status_code}", file=sys.stderr) sys.exit(1) diff --git a/skills/github/scripts/list_issues.py b/skills/github/scripts/list_issues.py index d79c755..d1aed85 100644 --- a/skills/github/scripts/list_issues.py +++ b/skills/github/scripts/list_issues.py @@ -44,20 +44,20 @@ def list_issues( per_page: int = 100, ) -> list[dict]: """List issues for a repository. - + Args: owner: Repository owner repo: Repository name state: Issue state - 'open', 'closed', or 'all' per_page: Number of results per page - + Returns: List of issue dictionaries (excluding pull requests). """ headers = get_github_headers() issues = [] page = 1 - + while True: response = httpx.get( f"https://api.github.com/repos/{owner}/{repo}/issues", @@ -69,25 +69,25 @@ def list_issues( }, timeout=30.0, ) - + if response.status_code == 401: print("Error: Invalid or expired GITHUB_TOKEN", file=sys.stderr) sys.exit(1) elif response.status_code == 404: print(f"Error: Repository '{owner}/{repo}' not found", file=sys.stderr) sys.exit(1) - + response.raise_for_status() page_issues = response.json() - + if not page_issues: break - + # Filter out pull requests (they appear in issues API too) real_issues = [i for i in page_issues if "pull_request" not in i] issues.extend(real_issues) page += 1 - + return issues @@ -95,30 +95,30 @@ def format_table(issues: list[dict]) -> str: """Format issues as a table.""" if not issues: return "No issues found." - + lines = [] lines.append(f"{'#':<7} {'State':<8} {'Title':<55} {'Author':<15} {'Created':<12}") lines.append("-" * 100) - + for issue in issues: number = f"#{issue['number']}" state = "🟢 Open" if issue["state"] == "open" else "🔴 Closed" title = issue["title"][:53] author = (issue.get("user", {}).get("login") or "-")[:13] created = issue.get("created_at", "")[:10] - + lines.append(f"{number:<7} {state:<8} {title:<55} {author:<15} {created:<12}") - + lines.append("-" * 100) lines.append(f"Total: {len(issues)} issues") - + return "\n".join(lines) class _HelpOnErrorParser(argparse.ArgumentParser): """ArgumentParser that prints full help on invalid arguments.""" - def error(self, message: str) -> None: # noqa: D401 + def error(self, message: str) -> None: params_help = ( "\nValid parameters for list_issues:\n" " repo (positional) owner/repo e.g. 'datalayer/agent-runtimes'\n" @@ -134,9 +134,7 @@ def error(self, message: str) -> None: # noqa: D401 def main(): - parser = _HelpOnErrorParser( - description="List issues for a GitHub repository." - ) + parser = _HelpOnErrorParser(description="List issues for a GitHub repository.") parser.add_argument( "repo", help="Repository in format 'owner/repo'", @@ -159,22 +157,22 @@ def main(): default=50, help="Maximum number of issues to display (default: 50)", ) - + args = parser.parse_args() - + # Parse owner/repo if "/" not in args.repo: print("Error: Repository must be in format 'owner/repo'", file=sys.stderr) sys.exit(1) - + owner, repo = args.repo.split("/", 1) - + try: issues = list_issues(owner, repo, state=args.state) - + if args.limit: - issues = issues[:args.limit] - + issues = issues[: args.limit] + if args.format == "json": simplified = [ { @@ -193,7 +191,7 @@ def main(): print(json.dumps(simplified, indent=2)) else: print(format_table(issues)) - + except httpx.HTTPStatusError as e: print(f"HTTP error: {e.response.status_code}", file=sys.stderr) sys.exit(1) diff --git a/skills/github/scripts/list_prs.py b/skills/github/scripts/list_prs.py index a2e71dd..c8dc90d 100644 --- a/skills/github/scripts/list_prs.py +++ b/skills/github/scripts/list_prs.py @@ -44,20 +44,20 @@ def list_pull_requests( per_page: int = 100, ) -> list[dict]: """List pull requests for a repository. - + Args: owner: Repository owner repo: Repository name state: PR state - 'open', 'closed', or 'all' per_page: Number of results per page - + Returns: List of pull request dictionaries. """ headers = get_github_headers() prs = [] page = 1 - + while True: response = httpx.get( f"https://api.github.com/repos/{owner}/{repo}/pulls", @@ -69,23 +69,23 @@ def list_pull_requests( }, timeout=30.0, ) - + if response.status_code == 401: print("Error: Invalid or expired GITHUB_TOKEN", file=sys.stderr) sys.exit(1) elif response.status_code == 404: print(f"Error: Repository '{owner}/{repo}' not found", file=sys.stderr) sys.exit(1) - + response.raise_for_status() page_prs = response.json() - + if not page_prs: break - + prs.extend(page_prs) page += 1 - + return prs @@ -93,11 +93,11 @@ def format_table(prs: list[dict]) -> str: """Format pull requests as a table.""" if not prs: return "No pull requests found." - + lines = [] lines.append(f"{'#':<7} {'State':<10} {'Title':<50} {'Author':<15} {'Branch':<20}") lines.append("-" * 105) - + for pr in prs: number = f"#{pr['number']}" if pr["state"] == "open": @@ -113,19 +113,19 @@ def format_table(prs: list[dict]) -> str: title = pr["title"][:48] author = (pr.get("user", {}).get("login") or "-")[:13] branch = pr.get("head", {}).get("ref", "-")[:18] - + lines.append(f"{number:<7} {state:<10} {title:<50} {author:<15} {branch:<20}") - + lines.append("-" * 105) lines.append(f"Total: {len(prs)} pull requests") - + return "\n".join(lines) class _HelpOnErrorParser(argparse.ArgumentParser): """ArgumentParser that prints full help on invalid arguments.""" - def error(self, message: str) -> None: # noqa: D401 + def error(self, message: str) -> None: params_help = ( "\nValid parameters for list_prs:\n" " repo (positional) owner/repo e.g. 'datalayer/agent-runtimes'\n" @@ -141,9 +141,7 @@ def error(self, message: str) -> None: # noqa: D401 def main(): - parser = _HelpOnErrorParser( - description="List pull requests for a GitHub repository." - ) + parser = _HelpOnErrorParser(description="List pull requests for a GitHub repository.") parser.add_argument( "repo", help="Repository in format 'owner/repo'", @@ -166,22 +164,22 @@ def main(): default=50, help="Maximum number of PRs to display (default: 50)", ) - + args = parser.parse_args() - + # Parse owner/repo if "/" not in args.repo: print("Error: Repository must be in format 'owner/repo'", file=sys.stderr) sys.exit(1) - + owner, repo = args.repo.split("/", 1) - + try: prs = list_pull_requests(owner, repo, state=args.state) - + if args.limit: - prs = prs[:args.limit] - + prs = prs[: args.limit] + if args.format == "json": simplified = [ { @@ -202,7 +200,7 @@ def main(): print(json.dumps(simplified, indent=2)) else: print(format_table(prs)) - + except httpx.HTTPStatusError as e: print(f"HTTP error: {e.response.status_code}", file=sys.stderr) sys.exit(1) diff --git a/skills/github/scripts/list_repos.py b/skills/github/scripts/list_repos.py index 55bafd6..a490aa5 100644 --- a/skills/github/scripts/list_repos.py +++ b/skills/github/scripts/list_repos.py @@ -47,19 +47,19 @@ def list_repos( per_page: int = 100, ) -> list[dict]: """List repositories for the authenticated user. - + Args: visibility: Filter by visibility - 'all', 'public', or 'private' sort: Sort by 'created', 'updated', 'pushed', or 'full_name' per_page: Number of results per page (max 100) - + Returns: List of repository dictionaries. """ headers = get_github_headers() repos = [] page = 1 - + while True: response = httpx.get( "https://api.github.com/user/repos", @@ -72,23 +72,23 @@ def list_repos( }, timeout=30.0, ) - + if response.status_code == 401: print("Error: Invalid or expired GITHUB_TOKEN", file=sys.stderr) sys.exit(1) elif response.status_code == 403: print("Error: Rate limit exceeded or insufficient permissions", file=sys.stderr) sys.exit(1) - + response.raise_for_status() page_repos = response.json() - + if not page_repos: break - + repos.extend(page_repos) page += 1 - + return repos @@ -96,30 +96,30 @@ def format_table(repos: list[dict]) -> str: """Format repositories as a table.""" if not repos: return "No repositories found." - + lines = [] lines.append(f"{'Visibility':<10} {'Name':<50} {'Language':<15} {'Stars':<8} {'Updated':<12}") lines.append("-" * 95) - + for repo in repos: visibility = "🔒 Private" if repo["private"] else "🌐 Public" name = repo["full_name"][:48] language = (repo.get("language") or "-")[:13] stars = str(repo.get("stargazers_count", 0)) updated = repo.get("updated_at", "")[:10] - + lines.append(f"{visibility:<10} {name:<50} {language:<15} {stars:<8} {updated:<12}") - + lines.append("-" * 95) lines.append(f"Total: {len(repos)} repositories") - + return "\n".join(lines) class _HelpOnErrorParser(argparse.ArgumentParser): """ArgumentParser that prints full help on invalid arguments.""" - def error(self, message: str) -> None: # noqa: D401 + def error(self, message: str) -> None: params_help = ( "\nValid parameters for list_repos:\n" " --visibility all | public | private (default: all)\n" @@ -163,15 +163,15 @@ def main(): default=None, help="Maximum number of repos to display", ) - + args = parser.parse_args() - + try: repos = list_repos(visibility=args.visibility, sort=args.sort) - + if args.limit: - repos = repos[:args.limit] - + repos = repos[: args.limit] + if args.format == "json": # Output simplified JSON simplified = [ @@ -190,7 +190,7 @@ def main(): print(json.dumps(simplified, indent=2)) else: print(format_table(repos)) - + except httpx.HTTPStatusError as e: print(f"HTTP error: {e.response.status_code}", file=sys.stderr) sys.exit(1) diff --git a/skills/github/scripts/search_repos.py b/skills/github/scripts/search_repos.py index 1ffa0f8..e8dfd81 100644 --- a/skills/github/scripts/search_repos.py +++ b/skills/github/scripts/search_repos.py @@ -49,7 +49,7 @@ def search_repos( per_page: int = 30, ) -> list[dict]: """Search GitHub repositories. - + Args: query: Search query language: Filter by programming language @@ -57,12 +57,12 @@ def search_repos( org: Filter by organization sort: Sort by 'stars', 'forks', or 'updated' per_page: Number of results per page - + Returns: List of repository dictionaries. """ headers = get_github_headers() - + # Build search query with qualifiers q_parts = [query] if language: @@ -71,21 +71,21 @@ def search_repos( q_parts.append(f"user:{user}") if org: q_parts.append(f"org:{org}") - + full_query = " ".join(q_parts) - + params = {"q": full_query, "per_page": per_page} if sort and sort != "best-match": params["sort"] = sort params["order"] = "desc" - + response = httpx.get( "https://api.github.com/search/repositories", headers=headers, params=params, timeout=30.0, ) - + if response.status_code == 401: print("Error: Invalid or expired GITHUB_TOKEN", file=sys.stderr) sys.exit(1) @@ -95,10 +95,10 @@ def search_repos( elif response.status_code == 422: print("Error: Invalid search query", file=sys.stderr) sys.exit(1) - + response.raise_for_status() data = response.json() - + return data.get("items", []), data.get("total_count", 0) @@ -106,29 +106,29 @@ def format_table(repos: list[dict], total: int) -> str: """Format search results as a table.""" if not repos: return "No repositories found." - + lines = [] lines.append(f"Found {total} repositories (showing {len(repos)})") lines.append("") lines.append(f"{'Name':<45} {'Language':<12} {'⭐ Stars':<10} {'🍴 Forks':<10} {'Updated':<12}") lines.append("-" * 95) - + for repo in repos: name = repo["full_name"][:43] language = (repo.get("language") or "-")[:10] stars = str(repo.get("stargazers_count", 0)) forks = str(repo.get("forks_count", 0)) updated = repo.get("updated_at", "")[:10] - + lines.append(f"{name:<45} {language:<12} {stars:<10} {forks:<10} {updated:<12}") - + return "\n".join(lines) class _HelpOnErrorParser(argparse.ArgumentParser): """ArgumentParser that prints full help on invalid arguments.""" - def error(self, message: str) -> None: # noqa: D401 + def error(self, message: str) -> None: params_help = ( "\nValid parameters for search_repos:\n" " query (positional) search text e.g. 'jupyter notebook'\n" @@ -148,9 +148,7 @@ def error(self, message: str) -> None: # noqa: D401 def main(): - parser = _HelpOnErrorParser( - description="Search GitHub repositories." - ) + parser = _HelpOnErrorParser(description="Search GitHub repositories.") parser.add_argument( "query", help="Search query", @@ -185,9 +183,9 @@ def main(): default=20, help="Maximum number of results (default: 20)", ) - + args = parser.parse_args() - + try: repos, total = search_repos( query=args.query, @@ -197,7 +195,7 @@ def main(): sort=args.sort, per_page=args.limit, ) - + if args.format == "json": simplified = [ { @@ -216,7 +214,7 @@ def main(): print(json.dumps(output, indent=2)) else: print(format_table(repos, total)) - + except httpx.HTTPStatusError as e: print(f"HTTP error: {e.response.status_code}", file=sys.stderr) sys.exit(1) diff --git a/skills/pdf/SKILL.md b/skills/pdf/SKILL.md index f6a22dd..decb4a8 100644 --- a/skills/pdf/SKILL.md +++ b/skills/pdf/SKILL.md @@ -30,6 +30,7 @@ for page in reader.pages: ### pypdf - Basic Operations #### Merge PDFs + ```python from pypdf import PdfWriter, PdfReader @@ -44,6 +45,7 @@ with open("merged.pdf", "wb") as output: ``` #### Split PDF + ```python reader = PdfReader("input.pdf") for i, page in enumerate(reader.pages): @@ -54,6 +56,7 @@ for i, page in enumerate(reader.pages): ``` #### Extract Metadata + ```python reader = PdfReader("document.pdf") meta = reader.metadata @@ -64,6 +67,7 @@ print(f"Creator: {meta.creator}") ``` #### Rotate Pages + ```python reader = PdfReader("input.pdf") writer = PdfWriter() @@ -79,6 +83,7 @@ with open("rotated.pdf", "wb") as output: ### pdfplumber - Text and Table Extraction #### Extract Text with Layout + ```python import pdfplumber @@ -89,6 +94,7 @@ with pdfplumber.open("document.pdf") as pdf: ``` #### Extract Tables + ```python with pdfplumber.open("document.pdf") as pdf: for i, page in enumerate(pdf.pages): @@ -100,6 +106,7 @@ with pdfplumber.open("document.pdf") as pdf: ``` #### Advanced Table Extraction + ```python import pandas as pd @@ -121,6 +128,7 @@ if all_tables: ### reportlab - Create PDFs #### Basic PDF Creation + ```python from reportlab.lib.pagesizes import letter from reportlab.pdfgen import canvas @@ -140,6 +148,7 @@ c.save() ``` #### Create PDF with Multiple Pages + ```python from reportlab.lib.pagesizes import letter from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak @@ -169,6 +178,7 @@ doc.build(story) ## Command-Line Tools ### pdftotext (poppler-utils) + ```bash # Extract text pdftotext input.pdf output.txt @@ -181,6 +191,7 @@ pdftotext -f 1 -l 5 input.pdf output.txt # Pages 1-5 ``` ### qpdf + ```bash # Merge PDFs qpdf --empty --pages file1.pdf file2.pdf -- merged.pdf @@ -197,6 +208,7 @@ qpdf --password=mypassword --decrypt encrypted.pdf decrypted.pdf ``` ### pdftk (if available) + ```bash # Merge pdftk file1.pdf file2.pdf cat output merged.pdf @@ -211,6 +223,7 @@ pdftk input.pdf rotate 1east output rotated.pdf ## Common Tasks ### Extract Text from Scanned PDFs + ```python # Requires: pip install pytesseract pdf2image import pytesseract @@ -230,6 +243,7 @@ print(text) ``` ### Add Watermark + ```python from pypdf import PdfReader, PdfWriter @@ -249,6 +263,7 @@ with open("watermarked.pdf", "wb") as output: ``` ### Extract Images + ```bash # Using pdfimages (poppler-utils) pdfimages -j input.pdf output_prefix @@ -257,6 +272,7 @@ pdfimages -j input.pdf output_prefix ``` ### Password Protection + ```python from pypdf import PdfReader, PdfWriter @@ -275,16 +291,16 @@ with open("encrypted.pdf", "wb") as output: ## Quick Reference -| Task | Best Tool | Command/Code | -|------|-----------|--------------| -| Merge PDFs | pypdf | `writer.add_page(page)` | -| Split PDFs | pypdf | One page per file | -| Extract text | pdfplumber | `page.extract_text()` | -| Extract tables | pdfplumber | `page.extract_tables()` | -| Create PDFs | reportlab | Canvas or Platypus | -| Command line merge | qpdf | `qpdf --empty --pages ...` | -| OCR scanned PDFs | pytesseract | Convert to image first | -| Fill PDF forms | pdf-lib or pypdf (see forms.md) | See forms.md | +| Task | Best Tool | Command/Code | +| ------------------ | ------------------------------- | -------------------------- | +| Merge PDFs | pypdf | `writer.add_page(page)` | +| Split PDFs | pypdf | One page per file | +| Extract text | pdfplumber | `page.extract_text()` | +| Extract tables | pdfplumber | `page.extract_tables()` | +| Create PDFs | reportlab | Canvas or Platypus | +| Command line merge | qpdf | `qpdf --empty --pages ...` | +| OCR scanned PDFs | pytesseract | Convert to image first | +| Fill PDF forms | pdf-lib or pypdf (see forms.md) | See forms.md | ## Next Steps diff --git a/skills/pdf/forms.md b/skills/pdf/forms.md index 4e23450..a6dbfba 100644 --- a/skills/pdf/forms.md +++ b/skills/pdf/forms.md @@ -1,11 +1,14 @@ **CRITICAL: You MUST complete these steps in order. Do not skip ahead to writing code.** If you need to fill out a PDF form, first check to see if the PDF has fillable form fields. Run this script from this file's directory: - `python scripts/check_fillable_fields `, and depending on the result go to either the "Fillable fields" or "Non-fillable fields" and follow those instructions. +`python scripts/check_fillable_fields `, and depending on the result go to either the "Fillable fields" or "Non-fillable fields" and follow those instructions. # Fillable fields + If the PDF has fillable form fields: + - Run this script from this file's directory: `python scripts/extract_form_field_info.py `. It will create a JSON file with a list of fields in this format: + ``` [ { @@ -50,10 +53,12 @@ If the PDF has fillable form fields: } ] ``` + - Convert the PDF to PNGs (one image for each page) with this script (run from this file's directory): -`python scripts/convert_pdf_to_images.py ` -Then analyze the images to determine the purpose of each form field (make sure to convert the bounding box PDF coordinates to image coordinates). + `python scripts/convert_pdf_to_images.py ` + Then analyze the images to determine the purpose of each form field (make sure to convert the bounding box PDF coordinates to image coordinates). - Create a `field_values.json` file in this format with the values to be entered for each field: + ``` [ { @@ -71,64 +76,81 @@ Then analyze the images to determine the purpose of each form field (make sure t // more fields ] ``` + - Run the `fill_fillable_fields.py` script from this file's directory to create a filled-in PDF: -`python scripts/fill_fillable_fields.py ` -This script will verify that the field IDs and values you provide are valid; if it prints error messages, correct the appropriate fields and try again. + `python scripts/fill_fillable_fields.py ` + This script will verify that the field IDs and values you provide are valid; if it prints error messages, correct the appropriate fields and try again. # Non-fillable fields + If the PDF doesn't have fillable form fields, you'll need to visually determine where the data should be added and create text annotations. Follow the below steps *exactly*. You MUST perform all of these steps to ensure that the the form is accurately completed. Details for each step are below. + - Convert the PDF to PNG images and determine field bounding boxes. - Create a JSON file with field information and validation images showing the bounding boxes. - Validate the the bounding boxes. - Use the bounding boxes to fill in the form. ## Step 1: Visual Analysis (REQUIRED) + - Convert the PDF to PNG images. Run this script from this file's directory: -`python scripts/convert_pdf_to_images.py ` -The script will create a PNG image for each page in the PDF. + `python scripts/convert_pdf_to_images.py ` + The script will create a PNG image for each page in the PDF. - Carefully examine each PNG image and identify all form fields and areas where the user should enter data. For each form field where the user should enter text, determine bounding boxes for both the form field label, and the area where the user should enter text. The label and entry bounding boxes MUST NOT INTERSECT; the text entry box should only include the area where data should be entered. Usually this area will be immediately to the side, above, or below its label. Entry bounding boxes must be tall and wide enough to contain their text. These are some examples of form structures that you might see: *Label inside box* + ``` ┌────────────────────────┐ │ Name: │ └────────────────────────┘ ``` + The input area should be to the right of the "Name" label and extend to the edge of the box. *Label before line* + ``` Email: _______________________ ``` + The input area should be above the line and include its entire width. *Label under line* + ``` _________________________ Name ``` + The input area should be above the line and include the entire width of the line. This is common for signature and date fields. *Label above line* + ``` Please enter any special requests: ________________________________________________ ``` + The input area should extend from the bottom of the label to the line, and should include the entire width of the line. *Checkboxes* + ``` Are you a US citizen? Yes □ No □ ``` + For checkboxes: + - Look for small square boxes (□) - these are the actual checkboxes to target. They may be to the left or right of their labels. - Distinguish between label text ("Yes", "No") and the clickable checkbox squares. - The entry bounding box should cover ONLY the small square, not the text label. ### Step 2: Create fields.json and validation images (REQUIRED) + - Create a file named `fields.json` with information for the form fields and bounding boxes in this format: + ``` { "pages": [ @@ -177,29 +199,37 @@ For checkboxes: ``` Create validation images by running this script from this file's directory for each page: -`python scripts/create_validation_image.py +\`python scripts/create_validation_image.py \ \ \ \ The validation images will have red rectangles where text should be entered, and blue rectangles covering label text. ### Step 3: Validate Bounding Boxes (REQUIRED) + #### Automated intersection check + - Verify that none of bounding boxes intersect and that the entry bounding boxes are tall enough by checking the fields.json file with the `check_bounding_boxes.py` script (run from this file's directory): -`python scripts/check_bounding_boxes.py ` + `python scripts/check_bounding_boxes.py ` If there are errors, reanalyze the relevant fields, adjust the bounding boxes, and iterate until there are no remaining errors. Remember: label (blue) bounding boxes should contain text labels, entry (red) boxes should not. #### Manual image inspection + **CRITICAL: Do not proceed without visually inspecting validation images** + - Red rectangles must ONLY cover input areas + - Red rectangles MUST NOT contain any text + - Blue rectangles should contain label text + - For checkboxes: + - Red rectangle MUST be centered on the checkbox square - Blue rectangle should cover the text label for the checkbox - If any rectangles look wrong, fix fields.json, regenerate the validation images, and verify again. Repeat this process until the bounding boxes are fully accurate. - ### Step 4: Add annotations to the PDF + Run this script from this file's directory to create a filled-out PDF using the information in fields.json: -`python scripts/fill_pdf_form_with_annotations.py +\`python scripts/fill_pdf_form_with_annotations.py \ \ \ diff --git a/skills/pdf/reference.md b/skills/pdf/reference.md index 41400bf..3f21575 100644 --- a/skills/pdf/reference.md +++ b/skills/pdf/reference.md @@ -5,9 +5,11 @@ This document contains advanced PDF processing features, detailed examples, and ## pypdfium2 Library (Apache/BSD License) ### Overview + pypdfium2 is a Python binding for PDFium (Chromium's PDF library). It's excellent for fast PDF rendering, image generation, and serves as a PyMuPDF replacement. ### Render PDF to Images + ```python import pypdfium2 as pdfium from PIL import Image @@ -34,6 +36,7 @@ for i, page in enumerate(pdf): ``` ### Extract Text with pypdfium2 + ```python import pypdfium2 as pdfium @@ -50,6 +53,7 @@ for i, page in enumerate(pdf): pdf-lib is a powerful JavaScript library for creating and modifying PDF documents in any JavaScript environment. #### Load and Manipulate Existing PDF + ```javascript import { PDFDocument } from 'pdf-lib'; import fs from 'fs'; @@ -78,6 +82,7 @@ async function manipulatePDF() { ``` #### Create Complex PDFs from Scratch + ```javascript import { PDFDocument, rgb, StandardFonts } from 'pdf-lib'; import fs from 'fs'; @@ -139,6 +144,7 @@ async function createPDF() { ``` #### Advanced Merge and Split Operations + ```javascript import { PDFDocument } from 'pdf-lib'; import fs from 'fs'; @@ -172,6 +178,7 @@ async function mergePDFs() { PDF.js is Mozilla's JavaScript library for rendering PDFs in the browser. #### Basic PDF Loading and Rendering + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -206,6 +213,7 @@ async function renderPDF() { ``` #### Extract Text with Coordinates + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -242,6 +250,7 @@ async function extractText() { ``` #### Extract Annotations and Forms + ```javascript import * as pdfjsLib from 'pdfjs-dist'; @@ -267,6 +276,7 @@ async function extractAnnotations() { ### poppler-utils Advanced Features #### Extract Text with Bounding Box Coordinates + ```bash # Extract text with bounding box coordinates (essential for structured data) pdftotext -bbox-layout document.pdf output.xml @@ -275,6 +285,7 @@ pdftotext -bbox-layout document.pdf output.xml ``` #### Advanced Image Conversion + ```bash # Convert to PNG images with specific resolution pdftoppm -png -r 300 document.pdf output_prefix @@ -287,6 +298,7 @@ pdftoppm -jpeg -jpegopt quality=85 -r 200 document.pdf jpeg_output ``` #### Extract Embedded Images + ```bash # Extract all embedded images with metadata pdfimages -j -p document.pdf page_images @@ -301,6 +313,7 @@ pdfimages -all document.pdf images/img ### qpdf Advanced Features #### Complex Page Manipulation + ```bash # Split PDF into groups of pages qpdf --split-pages=3 input.pdf output_group_%02d.pdf @@ -313,6 +326,7 @@ qpdf --empty --pages doc1.pdf 1-3 doc2.pdf 5-7 doc3.pdf 2,4 -- combined.pdf ``` #### PDF Optimization and Repair + ```bash # Optimize PDF for web (linearize for streaming) qpdf --linearize input.pdf optimized.pdf @@ -329,6 +343,7 @@ qpdf --show-all-pages input.pdf > structure.txt ``` #### Advanced Encryption + ```bash # Add password protection with specific permissions qpdf --encrypt user_pass owner_pass 256 --print=none --modify=none -- input.pdf encrypted.pdf @@ -345,29 +360,31 @@ qpdf --password=secret123 --decrypt encrypted.pdf decrypted.pdf ### pdfplumber Advanced Features #### Extract Text with Precise Coordinates + ```python import pdfplumber with pdfplumber.open("document.pdf") as pdf: page = pdf.pages[0] - + # Extract all text with coordinates chars = page.chars for char in chars[:10]: # First 10 characters print(f"Char: '{char['text']}' at x:{char['x0']:.1f} y:{char['y0']:.1f}") - + # Extract text by bounding box (left, top, right, bottom) bbox_text = page.within_bbox((100, 100, 400, 200)).extract_text() ``` #### Advanced Table Extraction with Custom Settings + ```python import pdfplumber import pandas as pd with pdfplumber.open("complex_table.pdf") as pdf: page = pdf.pages[0] - + # Extract tables with custom settings for complex layouts table_settings = { "vertical_strategy": "lines", @@ -376,7 +393,7 @@ with pdfplumber.open("complex_table.pdf") as pdf: "intersection_tolerance": 15 } tables = page.extract_tables(table_settings) - + # Visual debugging for table extraction img = page.to_image(resolution=150) img.save("debug_layout.png") @@ -385,6 +402,7 @@ with pdfplumber.open("complex_table.pdf") as pdf: ### reportlab Advanced Features #### Create Professional Reports with Tables + ```python from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph from reportlab.lib.styles import getSampleStyleSheet @@ -428,12 +446,14 @@ doc.build(elements) ### Extract Figures/Images from PDF #### Method 1: Using pdfimages (fastest) + ```bash # Extract all images with original quality pdfimages -all document.pdf images/img ``` #### Method 2: Using pypdfium2 + Image Processing + ```python import pypdfium2 as pdfium from PIL import Image @@ -441,26 +461,27 @@ import numpy as np def extract_figures(pdf_path, output_dir): pdf = pdfium.PdfDocument(pdf_path) - + for page_num, page in enumerate(pdf): # Render high-resolution page bitmap = page.render(scale=3.0) img = bitmap.to_pil() - + # Convert to numpy for processing img_array = np.array(img) - + # Simple figure detection (non-white regions) mask = np.any(img_array != [255, 255, 255], axis=2) - + # Find contours and extract bounding boxes # (This is simplified - real implementation would need more sophisticated detection) - + # Save detected figures # ... implementation depends on specific needs ``` ### Batch PDF Processing with Error Handling + ```python import os import glob @@ -472,7 +493,7 @@ logger = logging.getLogger(__name__) def batch_process_pdfs(input_dir, operation='merge'): pdf_files = glob.glob(os.path.join(input_dir, "*.pdf")) - + if operation == 'merge': writer = PdfWriter() for pdf_file in pdf_files: @@ -484,10 +505,10 @@ def batch_process_pdfs(input_dir, operation='merge'): except Exception as e: logger.error(f"Failed to process {pdf_file}: {e}") continue - + with open("batch_merged.pdf", "wb") as output: writer.write(output) - + elif operation == 'extract_text': for pdf_file in pdf_files: try: @@ -495,18 +516,19 @@ def batch_process_pdfs(input_dir, operation='merge'): text = "" for page in reader.pages: text += page.extract_text() - + output_file = pdf_file.replace('.pdf', '.txt') with open(output_file, 'w', encoding='utf-8') as f: f.write(text) logger.info(f"Extracted text from: {pdf_file}") - + except Exception as e: logger.error(f"Failed to extract text from {pdf_file}: {e}") continue ``` ### Advanced PDF Cropping + ```python from pypdf import PdfWriter, PdfReader @@ -528,37 +550,42 @@ with open("cropped.pdf", "wb") as output: ## Performance Optimization Tips ### 1. For Large PDFs + - Use streaming approaches instead of loading entire PDF in memory - Use `qpdf --split-pages` for splitting large files - Process pages individually with pypdfium2 ### 2. For Text Extraction + - `pdftotext -bbox-layout` is fastest for plain text extraction - Use pdfplumber for structured data and tables - Avoid `pypdf.extract_text()` for very large documents ### 3. For Image Extraction + - `pdfimages` is much faster than rendering pages - Use low resolution for previews, high resolution for final output ### 4. For Form Filling + - pdf-lib maintains form structure better than most alternatives - Pre-validate form fields before processing ### 5. Memory Management + ```python # Process PDFs in chunks def process_large_pdf(pdf_path, chunk_size=10): reader = PdfReader(pdf_path) total_pages = len(reader.pages) - + for start_idx in range(0, total_pages, chunk_size): end_idx = min(start_idx + chunk_size, total_pages) writer = PdfWriter() - + for i in range(start_idx, end_idx): writer.add_page(reader.pages[i]) - + # Process chunk with open(f"chunk_{start_idx//chunk_size}.pdf", "wb") as output: writer.write(output) @@ -567,6 +594,7 @@ def process_large_pdf(pdf_path, chunk_size=10): ## Troubleshooting Common Issues ### Encrypted PDFs + ```python # Handle password-protected PDFs from pypdf import PdfReader @@ -580,6 +608,7 @@ except Exception as e: ``` ### Corrupted PDFs + ```bash # Use qpdf to repair qpdf --check corrupted.pdf @@ -587,6 +616,7 @@ qpdf --replace-input corrupted.pdf ``` ### Text Extraction Issues + ```python # Fallback to OCR for scanned PDFs import pytesseract @@ -609,4 +639,4 @@ def extract_text_with_ocr(pdf_path): - **poppler-utils**: GPL-2 License - **qpdf**: Apache License - **pdf-lib**: MIT License -- **pdfjs-dist**: Apache License \ No newline at end of file +- **pdfjs-dist**: Apache License diff --git a/skills/pdf/scripts/check_bounding_boxes.py b/skills/pdf/scripts/check_bounding_boxes.py index b11ca9c..a8a9a4b 100644 --- a/skills/pdf/scripts/check_bounding_boxes.py +++ b/skills/pdf/scripts/check_bounding_boxes.py @@ -1,10 +1,9 @@ # Copyright (c) 2025-2026 Datalayer, Inc. # Distributed under the terms of the Modified BSD License. -from dataclasses import dataclass import json import sys - +from dataclasses import dataclass # Script to check that the `fields.json` file that Claude creates when analyzing PDFs # does not have overlapping bounding boxes. See forms.md. @@ -38,12 +37,18 @@ def rects_intersect(r1, r2): # This is O(N^2); we can optimize if it becomes a problem. for j in range(i + 1, len(rects_and_fields)): rj = rects_and_fields[j] - if ri.field["page_number"] == rj.field["page_number"] and rects_intersect(ri.rect, rj.rect): + if ri.field["page_number"] == rj.field["page_number"] and rects_intersect( + ri.rect, rj.rect + ): has_error = True if ri.field is rj.field: - messages.append(f"FAILURE: intersection between label and entry bounding boxes for `{ri.field['description']}` ({ri.rect}, {rj.rect})") + messages.append( + f"FAILURE: intersection between label and entry bounding boxes for `{ri.field['description']}` ({ri.rect}, {rj.rect})" + ) else: - messages.append(f"FAILURE: intersection between {ri.rect_type} bounding box for `{ri.field['description']}` ({ri.rect}) and {rj.rect_type} bounding box for `{rj.field['description']}` ({rj.rect})") + messages.append( + f"FAILURE: intersection between {ri.rect_type} bounding box for `{ri.field['description']}` ({ri.rect}) and {rj.rect_type} bounding box for `{rj.field['description']}` ({rj.rect})" + ) if len(messages) >= 20: messages.append("Aborting further checks; fix bounding boxes and try again") return messages @@ -53,7 +58,9 @@ def rects_intersect(r1, r2): entry_height = ri.rect[3] - ri.rect[1] if entry_height < font_size: has_error = True - messages.append(f"FAILURE: entry bounding box height ({entry_height}) for `{ri.field['description']}` is too short for the text content (font size: {font_size}). Increase the box height or decrease the font size.") + messages.append( + f"FAILURE: entry bounding box height ({entry_height}) for `{ri.field['description']}` is too short for the text content (font size: {font_size}). Increase the box height or decrease the font size." + ) if len(messages) >= 20: messages.append("Aborting further checks; fix bounding boxes and try again") return messages @@ -62,6 +69,7 @@ def rects_intersect(r1, r2): messages.append("SUCCESS: All bounding boxes are valid") return messages + if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: check_bounding_boxes.py [fields.json]") diff --git a/skills/pdf/scripts/check_bounding_boxes_test.py b/skills/pdf/scripts/check_bounding_boxes_test.py index 23571b8..0493bd6 100644 --- a/skills/pdf/scripts/check_bounding_boxes_test.py +++ b/skills/pdf/scripts/check_bounding_boxes_test.py @@ -1,19 +1,19 @@ # Copyright (c) 2025-2026 Datalayer, Inc. # Distributed under the terms of the Modified BSD License. -import unittest -import json import io +import json +import unittest + from check_bounding_boxes import get_bounding_box_messages # Currently this is not run automatically in CI; it's just for documentation and manual checking. class TestGetBoundingBoxMessages(unittest.TestCase): - def create_json_stream(self, data): """Helper to create a JSON stream from data""" return io.StringIO(json.dumps(data)) - + def test_no_intersections(self): """Test case with no bounding box intersections""" data = { @@ -22,22 +22,22 @@ def test_no_intersections(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 1, "label_bounding_box": [10, 40, 50, 60], - "entry_bounding_box": [60, 40, 150, 60] - } + "entry_bounding_box": [60, 40, 150, 60], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_label_entry_intersection_same_field(self): """Test intersection between label and entry of the same field""" data = { @@ -46,16 +46,16 @@ def test_label_entry_intersection_same_field(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 60, 30], - "entry_bounding_box": [50, 10, 150, 30] # Overlaps with label + "entry_bounding_box": [50, 10, 150, 30], # Overlaps with label } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "intersection" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_intersection_between_different_fields(self): """Test intersection between bounding boxes of different fields""" data = { @@ -64,22 +64,22 @@ def test_intersection_between_different_fields(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 1, "label_bounding_box": [40, 20, 80, 40], # Overlaps with Name's boxes - "entry_bounding_box": [160, 10, 250, 30] - } + "entry_bounding_box": [160, 10, 250, 30], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "intersection" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_different_pages_no_intersection(self): """Test that boxes on different pages don't count as intersecting""" data = { @@ -88,22 +88,22 @@ def test_different_pages_no_intersection(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 30] + "entry_bounding_box": [60, 10, 150, 30], }, { "description": "Email", "page_number": 2, "label_bounding_box": [10, 10, 50, 30], # Same coordinates but different page - "entry_bounding_box": [60, 10, 150, 30] - } + "entry_bounding_box": [60, 10, 150, 30], + }, ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_entry_height_too_small(self): """Test that entry box height is checked against font size""" data = { @@ -115,16 +115,16 @@ def test_entry_height_too_small(self): "entry_bounding_box": [60, 10, 150, 20], # Height is 10 "entry_text": { "font_size": 14 # Font size larger than height - } + }, } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "height" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_entry_height_adequate(self): """Test that adequate entry box height passes""" data = { @@ -136,16 +136,16 @@ def test_entry_height_adequate(self): "entry_bounding_box": [60, 10, 150, 30], # Height is 20 "entry_text": { "font_size": 14 # Font size smaller than height - } + }, } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_default_font_size(self): """Test that default font size is used when not specified""" data = { @@ -155,16 +155,16 @@ def test_default_font_size(self): "page_number": 1, "label_bounding_box": [10, 10, 50, 30], "entry_bounding_box": [60, 10, 150, 20], # Height is 10 - "entry_text": {} # No font_size specified, should use default 14 + "entry_text": {}, # No font_size specified, should use default 14 } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("FAILURE" in msg and "height" in msg for msg in messages)) self.assertFalse(any("SUCCESS" in msg for msg in messages)) - + def test_no_entry_text(self): """Test that missing entry_text doesn't cause height check""" data = { @@ -173,30 +173,32 @@ def test_no_entry_text(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [60, 10, 150, 20] # Small height but no entry_text + "entry_bounding_box": [60, 10, 150, 20], # Small height but no entry_text } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - + def test_multiple_errors_limit(self): """Test that error messages are limited to prevent excessive output""" fields = [] # Create many overlapping fields for i in range(25): - fields.append({ - "description": f"Field{i}", - "page_number": 1, - "label_bounding_box": [10, 10, 50, 30], # All overlap - "entry_bounding_box": [20, 15, 60, 35] # All overlap - }) - + fields.append( + { + "description": f"Field{i}", + "page_number": 1, + "label_bounding_box": [10, 10, 50, 30], # All overlap + "entry_bounding_box": [20, 15, 60, 35], # All overlap + } + ) + data = {"form_fields": fields} - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) # Should abort after ~20 messages @@ -205,7 +207,7 @@ def test_multiple_errors_limit(self): failure_count = sum(1 for msg in messages if "FAILURE" in msg) self.assertGreater(failure_count, 0) self.assertLess(len(messages), 30) # Should be limited - + def test_edge_touching_boxes(self): """Test that boxes touching at edges don't count as intersecting""" data = { @@ -214,16 +216,16 @@ def test_edge_touching_boxes(self): "description": "Name", "page_number": 1, "label_bounding_box": [10, 10, 50, 30], - "entry_bounding_box": [50, 10, 150, 30] # Touches at x=50 + "entry_bounding_box": [50, 10, 150, 30], # Touches at x=50 } ] } - + stream = self.create_json_stream(data) messages = get_bounding_box_messages(stream) self.assertTrue(any("SUCCESS" in msg for msg in messages)) self.assertFalse(any("FAILURE" in msg for msg in messages)) - -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() diff --git a/skills/pdf/scripts/check_fillable_fields.py b/skills/pdf/scripts/check_fillable_fields.py index da4f838..dba675a 100644 --- a/skills/pdf/scripts/check_fillable_fields.py +++ b/skills/pdf/scripts/check_fillable_fields.py @@ -2,14 +2,16 @@ # Distributed under the terms of the Modified BSD License. import sys -from pypdf import PdfReader +from pypdf import PdfReader # Script for Claude to run to determine whether a PDF has fillable form fields. See forms.md. reader = PdfReader(sys.argv[1]) -if (reader.get_fields()): +if reader.get_fields(): print("This PDF has fillable form fields") else: - print("This PDF does not have fillable form fields; you will need to visually determine where to enter data") + print( + "This PDF does not have fillable form fields; you will need to visually determine where to enter data" + ) diff --git a/skills/pdf/scripts/convert_pdf_to_images.py b/skills/pdf/scripts/convert_pdf_to_images.py index 9fbd097..363458a 100644 --- a/skills/pdf/scripts/convert_pdf_to_images.py +++ b/skills/pdf/scripts/convert_pdf_to_images.py @@ -6,7 +6,6 @@ from pdf2image import convert_from_path - # Converts each page of a PDF to a PNG image. @@ -21,7 +20,7 @@ def convert(pdf_path, output_dir, max_dim=1000): new_width = int(width * scale_factor) new_height = int(height * scale_factor) image = image.resize((new_width, new_height)) - + image_path = os.path.join(output_dir, f"page_{i+1}.png") image.save(image_path) print(f"Saved page {i+1} as {image_path} (size: {image.size})") diff --git a/skills/pdf/scripts/create_validation_image.py b/skills/pdf/scripts/create_validation_image.py index d346c4a..861c57a 100644 --- a/skills/pdf/scripts/create_validation_image.py +++ b/skills/pdf/scripts/create_validation_image.py @@ -6,36 +6,37 @@ from PIL import Image, ImageDraw - # Creates "validation" images with rectangles for the bounding box information that # Claude creates when determining where to add text annotations in PDFs. See forms.md. def create_validation_image(page_number, fields_json_path, input_path, output_path): # Input file should be in the `fields.json` format described in forms.md. - with open(fields_json_path, 'r') as f: + with open(fields_json_path) as f: data = json.load(f) img = Image.open(input_path) draw = ImageDraw.Draw(img) num_boxes = 0 - + for field in data["form_fields"]: if field["page_number"] == page_number: - entry_box = field['entry_bounding_box'] - label_box = field['label_bounding_box'] + entry_box = field["entry_bounding_box"] + label_box = field["label_bounding_box"] # Draw red rectangle over entry bounding box and blue rectangle over the label. - draw.rectangle(entry_box, outline='red', width=2) - draw.rectangle(label_box, outline='blue', width=2) + draw.rectangle(entry_box, outline="red", width=2) + draw.rectangle(label_box, outline="blue", width=2) num_boxes += 2 - + img.save(output_path) print(f"Created validation image at {output_path} with {num_boxes} bounding boxes") if __name__ == "__main__": if len(sys.argv) != 5: - print("Usage: create_validation_image.py [page number] [fields.json file] [input image path] [output image path]") + print( + "Usage: create_validation_image.py [page number] [fields.json file] [input image path] [output image path]" + ) sys.exit(1) page_number = int(sys.argv[1]) fields_json_path = sys.argv[2] diff --git a/skills/pdf/scripts/extract_form_field_info.py b/skills/pdf/scripts/extract_form_field_info.py index 6e26453..d9fc446 100644 --- a/skills/pdf/scripts/extract_form_field_info.py +++ b/skills/pdf/scripts/extract_form_field_info.py @@ -6,7 +6,6 @@ from pypdf import PdfReader - # Extracts data for the fillable form fields in a PDF and outputs JSON that # Claude uses to fill the fields. See forms.md. @@ -15,16 +14,16 @@ def get_full_annotation_field_id(annotation): components = [] while annotation: - field_name = annotation.get('/T') + field_name = annotation.get("/T") if field_name: components.append(field_name) - annotation = annotation.get('/Parent') + annotation = annotation.get("/Parent") return ".".join(reversed(components)) if components else None def make_field_dict(field, field_id): field_dict = {"field_id": field_id} - ft = field.get('/FT') + ft = field.get("/FT") if ft == "/Tx": field_dict["type"] = "text" elif ft == "/Btn": @@ -38,16 +37,21 @@ def make_field_dict(field, field_id): field_dict["checked_value"] = states[0] if states[0] != "/Off" else states[1] field_dict["unchecked_value"] = "/Off" else: - print(f"Unexpected state values for checkbox `${field_id}`. Its checked and unchecked values may not be correct; if you're trying to check it, visually verify the results.") + print( + f"Unexpected state values for checkbox `${field_id}`. Its checked and unchecked values may not be correct; if you're trying to check it, visually verify the results." + ) field_dict["checked_value"] = states[0] field_dict["unchecked_value"] = states[1] elif ft == "/Ch": field_dict["type"] = "choice" states = field.get("/_States_", []) - field_dict["choice_options"] = [{ - "value": state[0], - "text": state[1], - } for state in states] + field_dict["choice_options"] = [ + { + "value": state[0], + "text": state[1], + } + for state in states + ] else: field_dict["type"] = f"unknown ({ft})" return field_dict @@ -85,12 +89,12 @@ def get_field_info(reader: PdfReader): radio_fields_by_id = {} for page_index, page in enumerate(reader.pages): - annotations = page.get('/Annots', []) + annotations = page.get("/Annots", []) for ann in annotations: field_id = get_full_annotation_field_id(ann) if field_id in field_info_by_id: field_info_by_id[field_id]["page"] = page_index + 1 - field_info_by_id[field_id]["rect"] = ann.get('/Rect') + field_info_by_id[field_id]["rect"] = ann.get("/Rect") elif field_id in possible_radio_names: try: # ann['/AP']['/N'] should have two items. One of them is '/Off', @@ -111,10 +115,12 @@ def get_field_info(reader: PdfReader): # radio buttons correctly. (It does if you remove the leading slash # from the value, but that causes them not to appear correctly in # Chrome/Firefox/Acrobat/etc). - radio_fields_by_id[field_id]["radio_options"].append({ - "value": on_values[0], - "rect": rect, - }) + radio_fields_by_id[field_id]["radio_options"].append( + { + "value": on_values[0], + "rect": rect, + } + ) # Some PDFs have form field definitions without corresponding annotations, # so we can't tell where they are. Ignore these fields for now. @@ -123,7 +129,9 @@ def get_field_info(reader: PdfReader): if "page" in field_info: fields_with_location.append(field_info) else: - print(f"Unable to determine location for field id: {field_info.get('field_id')}, ignoring") + print( + f"Unable to determine location for field id: {field_info.get('field_id')}, ignoring" + ) # Sort by page number, then Y position (flipped in PDF coordinate system), then X. def sort_key(f): @@ -133,7 +141,7 @@ def sort_key(f): rect = f.get("rect") or [0, 0, 0, 0] adjusted_position = [-rect[1], rect[0]] return [f.get("page"), adjusted_position] - + sorted_fields = fields_with_location + list(radio_fields_by_id.values()) sorted_fields.sort(key=sort_key) diff --git a/skills/pdf/scripts/fill_fillable_fields.py b/skills/pdf/scripts/fill_fillable_fields.py index 3b2edcb..5cce18a 100644 --- a/skills/pdf/scripts/fill_fillable_fields.py +++ b/skills/pdf/scripts/fill_fillable_fields.py @@ -4,10 +4,8 @@ import json import sys -from pypdf import PdfReader, PdfWriter - from extract_form_field_info import get_field_info - +from pypdf import PdfReader, PdfWriter # Fills fillable form fields in a PDF. See forms.md. @@ -24,7 +22,7 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: if page not in fields_by_page: fields_by_page[page] = {} fields_by_page[page][field_id] = field["value"] - + reader = PdfReader(input_pdf_path) has_error = False @@ -37,7 +35,9 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: print(f"ERROR: `{field['field_id']}` is not a valid field ID") elif field["page"] != existing_field["page"]: has_error = True - print(f"ERROR: Incorrect page number for `{field['field_id']}` (got {field['page']}, expected {existing_field['page']})") + print( + f"ERROR: Incorrect page number for `{field['field_id']}` (got {field['page']}, expected {existing_field['page']})" + ) else: if "value" in field: err = validation_error_for_field_value(existing_field, field["value"]) @@ -49,12 +49,14 @@ def fill_pdf_fields(input_pdf_path: str, fields_json_path: str, output_pdf_path: writer = PdfWriter(clone_from=reader) for page, field_values in fields_by_page.items(): - writer.update_page_form_field_values(writer.pages[page - 1], field_values, auto_regenerate=False) + writer.update_page_form_field_values( + writer.pages[page - 1], field_values, auto_regenerate=False + ) # This seems to be necessary for many PDF viewers to format the form values correctly. # It may cause the viewer to show a "save changes" dialog even if the user doesn't make any changes. writer.set_need_appearances_writer(True) - + with open(output_pdf_path, "wb") as f: writer.write(f) @@ -70,7 +72,7 @@ def validation_error_for_field_value(field_info, field_value): elif field_type == "radio_group": option_values = [opt["value"] for opt in field_info["radio_options"]] if field_value not in option_values: - return f'ERROR: Invalid value "{field_value}" for radio group field "{field_id}". Valid values are: {option_values}' + return f'ERROR: Invalid value "{field_value}" for radio group field "{field_id}". Valid values are: {option_values}' elif field_type == "choice": choice_values = [opt["value"] for opt in field_info["choice_options"]] if field_value not in choice_values: @@ -91,15 +93,17 @@ def validation_error_for_field_value(field_info, field_value): # We call the original method and adjust the return value only if the argument to `get_inherited` # is `FA.Opt` and if the return value is a list of two-element lists. def monkeypatch_pydpf_method(): - from pypdf.generic import DictionaryObject from pypdf.constants import FieldDictionaryAttributes + from pypdf.generic import DictionaryObject original_get_inherited = DictionaryObject.get_inherited - def patched_get_inherited(self, key: str, default = None): + def patched_get_inherited(self, key: str, default=None): result = original_get_inherited(self, key, default) if key == FieldDictionaryAttributes.Opt: - if isinstance(result, list) and all(isinstance(v, list) and len(v) == 2 for v in result): + if isinstance(result, list) and all( + isinstance(v, list) and len(v) == 2 for v in result + ): result = [r[0] for r in result] return result diff --git a/skills/pdf/scripts/fill_pdf_form_with_annotations.py b/skills/pdf/scripts/fill_pdf_form_with_annotations.py index bdbd01a..5e211cb 100644 --- a/skills/pdf/scripts/fill_pdf_form_with_annotations.py +++ b/skills/pdf/scripts/fill_pdf_form_with_annotations.py @@ -7,7 +7,6 @@ from pypdf import PdfReader, PdfWriter from pypdf.annotations import FreeText - # Fills a PDF by adding text annotations defined in `fields.json`. See forms.md. @@ -17,54 +16,52 @@ def transform_coordinates(bbox, image_width, image_height, pdf_width, pdf_height # PDF coordinates: origin at bottom-left, y increases upward x_scale = pdf_width / image_width y_scale = pdf_height / image_height - + left = bbox[0] * x_scale right = bbox[2] * x_scale - + # Flip Y coordinates for PDF top = pdf_height - (bbox[1] * y_scale) bottom = pdf_height - (bbox[3] * y_scale) - + return left, bottom, right, top def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): """Fill the PDF form with data from fields.json""" - + # `fields.json` format described in forms.md. - with open(fields_json_path, "r") as f: + with open(fields_json_path) as f: fields_data = json.load(f) - + # Open the PDF reader = PdfReader(input_pdf_path) writer = PdfWriter() - + # Copy all pages to writer writer.append(reader) - + # Get PDF dimensions for each page pdf_dimensions = {} for i, page in enumerate(reader.pages): mediabox = page.mediabox pdf_dimensions[i + 1] = [mediabox.width, mediabox.height] - + # Process each form field annotations = [] for field in fields_data["form_fields"]: page_num = field["page_number"] - + # Get page dimensions and transform coordinates. page_info = next(p for p in fields_data["pages"] if p["page_number"] == page_num) image_width = page_info["image_width"] image_height = page_info["image_height"] pdf_width, pdf_height = pdf_dimensions[page_num] - + transformed_entry_box = transform_coordinates( - field["entry_bounding_box"], - image_width, image_height, - pdf_width, pdf_height + field["entry_bounding_box"], image_width, image_height, pdf_width, pdf_height ) - + # Skip empty fields if "entry_text" not in field or "text" not in field["entry_text"]: continue @@ -72,7 +69,7 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): text = entry_text["text"] if not text: continue - + font_name = entry_text.get("font", "Arial") font_size = str(entry_text.get("font_size", 14)) + "pt" font_color = entry_text.get("font_color", "000000") @@ -91,11 +88,11 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): annotations.append(annotation) # page_number is 0-based for pypdf writer.add_annotation(page_number=page_num - 1, annotation=annotation) - + # Save the filled PDF with open(output_pdf_path, "wb") as output: writer.write(output) - + print(f"Successfully filled PDF form and saved to {output_pdf_path}") print(f"Added {len(annotations)} text annotations") @@ -107,5 +104,5 @@ def fill_pdf_form(input_pdf_path, fields_json_path, output_pdf_path): input_pdf = sys.argv[1] fields_json = sys.argv[2] output_pdf = sys.argv[3] - - fill_pdf_form(input_pdf, fields_json, output_pdf) \ No newline at end of file + + fill_pdf_form(input_pdf, fields_json, output_pdf) diff --git a/tests/conftest.py b/tests/conftest.py index e68b06f..a6ce79e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,10 +4,10 @@ """Pytest configuration and fixtures for agent-codemode tests.""" -import pytest -import tempfile from pathlib import Path +import pytest + @pytest.fixture def skills_dir(tmp_path: Path) -> Path: diff --git a/tests/test_agent_minimal.py b/tests/test_agent_minimal.py index f725a25..261b316 100644 --- a/tests/test_agent_minimal.py +++ b/tests/test_agent_minimal.py @@ -2,39 +2,40 @@ """Minimal test for agent with fixed sandbox.""" import asyncio -from pathlib import Path # Add parent directory to path for imports import sys +from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) from agent_codemode import CodeModeExecutor, ToolRegistry -from agent_codemode.types import MCPServerConfig, CodeModeConfig +from agent_codemode.types import CodeModeConfig, MCPServerConfig + async def main(): print("=== Testing CodeModeExecutor with fixed sandbox ===\n") - + # Set up registry with example MCP server registry = ToolRegistry() server_config = MCPServerConfig( name="example_mcp", command="mcp-run", args=["./example_server.py"], - cwd=str(Path(__file__).parent / "examples" / "agent") + cwd=str(Path(__file__).parent / "examples" / "agent"), ) registry.add_server(server_config) await registry.discover_all() - + print(f"Registered tools: {[tool.name for tool in registry.list_tools()]}\n") - + # Create executor config = CodeModeConfig( - workspace_path=Path(__file__).parent / "examples" / "agent", - sandbox_variant="eval" + workspace_path=Path(__file__).parent / "examples" / "agent", sandbox_variant="eval" ) executor = CodeModeExecutor(registry=registry, config=config) await executor.setup() - + # Test code execution code = """ from generated.mcp.example_mcp import generate_random_text @@ -46,20 +47,21 @@ async def main(): except Exception as e: print(f"Error: {e}") """ - + print("=== Executing code ===") print(code) print("\n=== Output ===") - + execution = await executor.execute(code) - + print(f"Stdout: {execution.stdout}") print(f"Stderr: {execution.stderr}") print(f"Error: {execution.error}") print(f"Success: {not execution.error}") - + await executor.cleanup() print("\n=== Test completed ===") + if __name__ == "__main__": asyncio.run(main()) diff --git a/tests/test_codemode.py b/tests/test_codemode.py index bff9e5e..693d048 100644 --- a/tests/test_codemode.py +++ b/tests/test_codemode.py @@ -5,34 +5,31 @@ """Unit tests for agent-codemode package.""" import asyncio -import tempfile from pathlib import Path from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest - -from agent_codemode import CodemodeToolset, CodeModeConfig -from agent_codemode.types import ToolDefinition, SearchResult - +from agent_skills.files import ( + SkillDirectory, + SkillFile, +) from agent_skills.helpers import ( - wait_for, + RateLimiter, + parallel, retry, run_with_timeout, - parallel, - RateLimiter, -) -from agent_skills.files import ( - SkillFile, - SkillDirectory, - setup_skills_directory, + wait_for, ) +from agent_codemode import CodeModeConfig, CodemodeToolset +from agent_codemode.types import SearchResult, ToolDefinition # ============================================================================= # wait_for Tests # ============================================================================= + class TestWaitFor: """Tests for wait_for helper.""" @@ -102,12 +99,14 @@ def condition(): # retry Tests # ============================================================================= + class TestRetry: """Tests for retry helper.""" @pytest.mark.asyncio async def test_retry_success_first_try(self): """Test retry when function succeeds on first try.""" + async def succeed(): return "success" @@ -137,6 +136,7 @@ async def fail_then_succeed(): @pytest.mark.asyncio async def test_retry_all_failures(self): """Test retry when function always fails.""" + async def always_fail(): raise ValueError("Always fails") @@ -221,12 +221,14 @@ async def raise_type_error(): # run_with_timeout Tests # ============================================================================= + class TestRunWithTimeout: """Tests for run_with_timeout helper.""" @pytest.mark.asyncio async def test_run_with_timeout_completes(self): """Test run_with_timeout when function completes in time.""" + async def fast_func(): await asyncio.sleep(0.05) return "done" @@ -237,6 +239,7 @@ async def fast_func(): @pytest.mark.asyncio async def test_run_with_timeout_exceeds(self): """Test run_with_timeout when function exceeds timeout.""" + async def slow_func(): await asyncio.sleep(10.0) return "done" @@ -249,12 +252,14 @@ async def slow_func(): # parallel Tests # ============================================================================= + class TestParallel: """Tests for parallel helper.""" @pytest.mark.asyncio async def test_parallel_basic(self): """Test running coroutines in parallel.""" + async def task(n): await asyncio.sleep(0.01) return n * 2 @@ -276,6 +281,7 @@ async def test_parallel_empty(self): @pytest.mark.asyncio async def test_parallel_preserves_order(self): """Test that parallel preserves result order.""" + async def task(n, delay): await asyncio.sleep(delay) return n @@ -294,6 +300,7 @@ async def task(n, delay): # RateLimiter Tests # ============================================================================= + class TestRateLimiter: """Tests for RateLimiter class.""" @@ -331,6 +338,7 @@ async def test_rate_limiter_throttles(self): # SkillFile Tests (agent_codemode version) # ============================================================================= + class TestMcpCodemodeSkillFile: """Tests for SkillFile in agent_codemode.""" @@ -354,6 +362,7 @@ async def process(): # SkillDirectory Tests (agent_codemode version) # ============================================================================= + class TestMcpCodemodeSkillDirectory: """Tests for SkillDirectory in agent_codemode.""" @@ -375,8 +384,8 @@ def test_search(self, tmp_path: Path): """Test searching skills.""" skills = SkillDirectory(str(tmp_path)) - skills.create(name="csv_reader", code='async def csv_reader(): pass') - skills.create(name="json_writer", code='async def json_writer(): pass') + skills.create(name="csv_reader", code="async def csv_reader(): pass") + skills.create(name="json_writer", code="async def json_writer(): pass") results = skills.search("csv") assert any(s.name == "csv_reader" for s in results) @@ -386,6 +395,7 @@ def test_search(self, tmp_path: Path): # Integration Tests # ============================================================================= + class TestMcpCodemodeIntegration: """Integration tests for agent-codemode.""" @@ -397,21 +407,21 @@ async def test_skill_with_helpers(self, tmp_path: Path): # Create a skill that uses wait_for pattern internally skills.create( name="polling_skill", - code=''' + code=""" async def polling_skill(): counter = [0] - + def check(): counter[0] += 1 return counter[0] >= 3 - + # Simple polling without importing wait_for import asyncio while not check(): await asyncio.sleep(0.01) - + return {"checks": counter[0]} -''', +""", ) skill = skills.get("polling_skill") @@ -427,17 +437,17 @@ async def test_retry_in_skill(self, tmp_path: Path): skills.create( name="retry_skill", - code=''' + code=""" async def retry_skill(): attempts = 0 - + async def flaky_operation(): nonlocal attempts attempts += 1 if attempts < 3: raise ValueError("Temporary failure") return "success" - + # Simple retry without importing retry helper for i in range(5): try: @@ -446,9 +456,9 @@ async def flaky_operation(): except ValueError: import asyncio await asyncio.sleep(0.01) - + raise RuntimeError("All retries failed") -''', +""", ) skill = skills.get("retry_skill") @@ -595,5 +605,8 @@ async def test_get_tool_details_includes_examples(): result = await toolset._get_tool_details("a__one") - assert result["output_schema"] == {"type": "object", "properties": {"value": {"type": "string"}}} + assert result["output_schema"] == { + "type": "object", + "properties": {"value": {"type": "string"}}, + } assert result["input_examples"] == [{"value": "example"}] diff --git a/tests/test_direct_execution.py b/tests/test_direct_execution.py index b4962e1..549fd61 100644 --- a/tests/test_direct_execution.py +++ b/tests/test_direct_execution.py @@ -8,12 +8,13 @@ sys.path.insert(0, str(Path(__file__).parent)) from agent_codemode import CodeModeExecutor, ToolRegistry -from agent_codemode.types import MCPServerConfig, CodeModeConfig +from agent_codemode.types import CodeModeConfig, MCPServerConfig + async def main(): print("=== Setting up registry ===") registry = ToolRegistry() - + # Add example MCP server server_path = Path(__file__).parent / "examples" / "agent" / "example_mcp_server.py" server_config = MCPServerConfig( @@ -23,19 +24,18 @@ async def main(): ) registry.add_server(server_config) await registry.discover_all() - + tools = registry.list_tools() print(f"Found {len(tools)} tools: {[t.name for t in tools]}") - + print("\n=== Setting up executor ===") config = CodeModeConfig( - workspace_path=Path(__file__).parent / "examples" / "agent", - sandbox_variant="eval" + workspace_path=Path(__file__).parent / "examples" / "agent", sandbox_variant="eval" ) executor = CodeModeExecutor(registry=registry, config=config) await executor.setup() print("Executor setup complete") - + print("\n=== Executing code ===") code = """ from generated.mcp.example_mcp import generate_random_text @@ -46,17 +46,18 @@ async def main(): except Exception as e: print(f"ERROR: {e}") """ - + print(f"Code to execute:\n{code}\n") execution = await executor.execute(code, timeout=10.0) - - print(f"\n=== Results ===") + + print("\n=== Results ===") print(f"Error: {execution.error}") print(f"Stdout:\n{execution.stdout}") print(f"Stderr:\n{execution.stderr}") - + await executor.cleanup() print("\n=== Done ===") + if __name__ == "__main__": asyncio.run(main()) diff --git a/tests/test_executor_async.py b/tests/test_executor_async.py index 07dc292..5c6135b 100644 --- a/tests/test_executor_async.py +++ b/tests/test_executor_async.py @@ -2,23 +2,24 @@ """Test executor with async code and external functions.""" import asyncio -from pathlib import Path + from code_sandboxes import Sandbox + async def main(): # Create a sandbox sandbox = Sandbox.create(variant="eval") sandbox.start() - + # Set up an executor mock class MockExecutor: async def call_tool(self, name, args): print(f"[MockExecutor] Calling tool: {name} with {args}") await asyncio.sleep(0.01) return {"status": "success", "data": f"Result for {name}"} - + sandbox.set_variable("__executor__", MockExecutor()) - + # First: Define __call_tool__ print("\n=== Step 1: Define __call_tool__ ===") result1 = sandbox.run_code(""" @@ -30,7 +31,7 @@ async def __call_tool__(tool_name, arguments): """) print(f"Result 1 - Error: {result1.code_error}") print(f"Result 1 - Stdout: {result1.stdout}") - + # Second: Use __call_tool__ print("\n=== Step 2: Use __call_tool__ with await ===") result2 = sandbox.run_code(""" @@ -41,9 +42,10 @@ async def __call_tool__(tool_name, arguments): print(f"Result 2 - Error: {result2.code_error}") print(f"Result 2 - Stdout: {result2.stdout}") print(f"Result 2 - Results: {result2.results}") - + sandbox.stop() print("\n=== Test completed successfully! ===") + if __name__ == "__main__": asyncio.run(main()) diff --git a/tests/test_executor_streaming.py b/tests/test_executor_streaming.py new file mode 100644 index 0000000..1acdfc0 --- /dev/null +++ b/tests/test_executor_streaming.py @@ -0,0 +1,74 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Tests for streaming execution behavior in CodeModeExecutor.""" + +from __future__ import annotations + +import pytest +from code_sandboxes import ExecutionResult, Logs, OutputMessage +from code_sandboxes.models import Result + +from agent_codemode.composition.executor import CodeModeExecutor +from agent_codemode.discovery.registry import ToolRegistry + + +class _StreamingSandbox: + def __init__(self) -> None: + self.run_code_calls = 0 + self.streaming_called = False + + def run_code(self, code: str, timeout=None) -> ExecutionResult: + _ = (code, timeout) + self.run_code_calls += 1 + return ExecutionResult(logs=Logs()) + + def run_code_streaming(self, code: str, timeout=None): + _ = (code, timeout) + self.streaming_called = True + yield OutputMessage(line="status: RUNNING", timestamp=0.0, error=False) + yield OutputMessage(line="hello", timestamp=0.0, error=False) + yield Result(data={"text/plain": "42"}, is_main_result=True, extra={}) + + +class _NonStreamingSandbox: + def __init__(self) -> None: + self.run_code_calls = 0 + + def run_code(self, code: str, timeout=None) -> ExecutionResult: + _ = (code, timeout) + self.run_code_calls += 1 + if self.run_code_calls >= 3: + return ExecutionResult( + logs=Logs(stdout=[OutputMessage(line="fallback", timestamp=0.0, error=False)]), + ) + return ExecutionResult(logs=Logs()) + + +@pytest.mark.asyncio +async def test_execute_uses_streaming_when_supported(): + executor = CodeModeExecutor(registry=ToolRegistry()) + sandbox = _StreamingSandbox() + executor._sandbox = sandbox + executor._setup_done = True + + result = await executor.execute("print('hi')") + + assert sandbox.streaming_called is True + assert "status: RUNNING" in result.logs.stdout_text + assert "hello" in result.logs.stdout_text + assert result.results and result.results[0].data["text/plain"] == "42" + + +@pytest.mark.asyncio +async def test_execute_falls_back_to_run_code_without_streaming(): + executor = CodeModeExecutor(registry=ToolRegistry()) + sandbox = _NonStreamingSandbox() + executor._sandbox = sandbox + executor._setup_done = True + + result = await executor.execute("print('hi')") + + assert sandbox.run_code_calls >= 3 + assert result.logs.stdout_text == "fallback" diff --git a/tests/test_skill_bindings.py b/tests/test_skill_bindings.py index 33fa71a..425ae8b 100644 --- a/tests/test_skill_bindings.py +++ b/tests/test_skill_bindings.py @@ -10,14 +10,10 @@ 3. CodemodeToolset post-init callbacks fire after lazy initialisation """ -import asyncio -import importlib -import json import sys -import textwrap from pathlib import Path from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest @@ -25,7 +21,6 @@ from agent_codemode.discovery.registry import ToolRegistry from agent_codemode.types import CodeModeConfig - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -64,6 +59,7 @@ def generated_dir(tmp_path: Path) -> Path: # 1. Codegen: generate_skill_bindings produces correct file hierarchy # --------------------------------------------------------------------------- + class TestGenerateSkillBindings: """Test that PythonCodeGenerator.generate_skill_bindings creates the expected files under servers/skills/.""" @@ -115,6 +111,7 @@ def test_init_exports_all_functions(self, generated_dir: Path): # 2. Executor: skills__* routing through set_skill_tool_caller # --------------------------------------------------------------------------- + class TestExecutorSkillRouting: """Test that CodeModeExecutor routes skills__* calls to the caller.""" @@ -134,9 +131,7 @@ async def test_skills_prefix_routes_to_caller(self, tmp_path: Path): mock_caller = AsyncMock(return_value={"result": "ok"}) executor.set_skill_tool_caller(mock_caller) - result = await executor.call_tool( - "skills__list_skills", {} - ) + result = await executor.call_tool("skills__list_skills", {}) mock_caller.assert_awaited_once_with("skills__list_skills", {}) assert result == {"result": "ok"} @@ -184,13 +179,14 @@ async def test_skills_metadata_stored(self, tmp_path: Path): # 3. CodemodeToolset post-init callbacks # --------------------------------------------------------------------------- + class TestCodemodeToolsetPostInit: """Test the post_init_callback mechanism on CodemodeToolset.""" @pytest.mark.asyncio async def test_callback_fires_on_ensure_initialized(self, tmp_path: Path): """Post-init callback should fire after _ensure_initialized.""" - from agent_codemode.toolset import CodemodeToolset, PYDANTIC_AI_AVAILABLE + from agent_codemode.toolset import PYDANTIC_AI_AVAILABLE, CodemodeToolset if not PYDANTIC_AI_AVAILABLE: pytest.skip("pydantic-ai not installed") @@ -216,7 +212,7 @@ async def test_callback_fires_on_ensure_initialized(self, tmp_path: Path): @pytest.mark.asyncio async def test_callback_not_called_twice(self, tmp_path: Path): """Callbacks only fire once even if start() is called twice.""" - from agent_codemode.toolset import CodemodeToolset, PYDANTIC_AI_AVAILABLE + from agent_codemode.toolset import PYDANTIC_AI_AVAILABLE, CodemodeToolset if not PYDANTIC_AI_AVAILABLE: pytest.skip("pydantic-ai not installed") @@ -241,7 +237,7 @@ async def test_callback_not_called_twice(self, tmp_path: Path): @pytest.mark.asyncio async def test_callback_error_does_not_prevent_init(self, tmp_path: Path): """A failing callback should log an error but not crash init.""" - from agent_codemode.toolset import CodemodeToolset, PYDANTIC_AI_AVAILABLE + from agent_codemode.toolset import PYDANTIC_AI_AVAILABLE, CodemodeToolset if not PYDANTIC_AI_AVAILABLE: pytest.skip("pydantic-ai not installed") @@ -272,12 +268,14 @@ def bad_callback(ts): # 3b. generate_skills_in_sandbox embeds catalog # --------------------------------------------------------------------------- + class TestGenerateSkillsInSandbox: """Test that generate_skills_in_sandbox() produces correct bindings for Jupyter/remote sandboxes.""" class MockRemoteSandbox: """Sandbox mock that captures run_code calls (no _namespaces).""" + def __init__(self): self.code_calls: list[str] = [] self._started = True @@ -286,6 +284,7 @@ def run_code(self, code: str, **kwargs): self.code_calls.append(code) exec(compile(code, "", "exec")) from code_sandboxes.models import ExecutionResult + return ExecutionResult() def start(self): @@ -361,6 +360,7 @@ def test_other_bindings_use_direct_execution(self, tmp_path: Path): # 4. Full codegen + import round-trip (no sandbox needed) # --------------------------------------------------------------------------- + class TestCodegenImportRoundTrip: """Verify that generated skill binding files can be imported and that calling the functions triggers call_tool with the correct diff --git a/tests/test_stdout_capture.py b/tests/test_stdout_capture.py index b13da4a..36e19a0 100644 --- a/tests/test_stdout_capture.py +++ b/tests/test_stdout_capture.py @@ -1,6 +1,7 @@ """Test script to verify stdout capture in async code.""" import asyncio + from agent_codemode.composition.executor import CodeModeExecutor from agent_codemode.discovery.registry import ToolRegistry @@ -10,7 +11,7 @@ async def test_async_stdout(): registry = ToolRegistry() executor = CodeModeExecutor(registry=registry) await executor.setup() - + # Test code that prints to stdout in async context code = """ import random @@ -23,9 +24,9 @@ async def generate_text(): result = await generate_text() """ - + result = await executor.execute(code) - + print(f"Execution ok: {result.execution_ok}") print(f"Execution error: {result.execution_error}") print(f"Code error: {result.code_error}") @@ -33,10 +34,12 @@ async def generate_text(): print(f"Stderr text: '{result.logs.stderr_text}'") print(f"Stdout messages: {result.logs.stdout}") print(f"Stderr messages: {result.logs.stderr}") - + # Verify stdout was captured assert result.success, f"Execution failed: {result.code_error}" - assert "Random text:" in result.logs.stdout_text, f"stdout not captured: {result.logs.stdout_text}" + assert ( + "Random text:" in result.logs.stdout_text + ), f"stdout not captured: {result.logs.stdout_text}" print("\n✅ SUCCESS: stdout was captured correctly!")