From e1c81644c3833243d77871c2eeecaea47b9a4aef Mon Sep 17 00:00:00 2001 From: Farouk Adeleke Date: Fri, 5 Dec 2025 23:40:50 -0500 Subject: [PATCH] Add more functionality to signature description parsing --- README.md | 195 ++++++++++------ auto_classes.json | 4 + claude_dspy/__init__.py | 24 ++ claude_dspy/agent.py | 479 ++++++++++++++++++++++++++++++++++++++++ claude_dspy/trace.py | 52 +++++ claude_dspy/utils.py | 77 +++++++ config.json | 13 +- main.py | 46 ++++ pyproject.toml | 13 ++ 9 files changed, 827 insertions(+), 76 deletions(-) create mode 100644 auto_classes.json create mode 100644 claude_dspy/__init__.py create mode 100644 claude_dspy/agent.py create mode 100644 claude_dspy/trace.py create mode 100644 claude_dspy/utils.py create mode 100644 main.py create mode 100644 pyproject.toml diff --git a/README.md b/README.md index fbe67f3..9d11317 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ A DSPy module that wraps the Claude Code Python SDK with a signature-driven inte - **Smart schema handling** - Automatically handles str vs Pydantic outputs - **Rich outputs** - Get typed results + execution trace + token usage - **Multi-turn conversations** - Context preserved across calls -- **Output field descriptions** - Automatically enhance prompts +- **Enhanced prompts** - Automatically includes signature docstrings + InputField/OutputField descriptions for better context - **Async support** - Both sync and async execution modes - **Modaic Hub Integration** - Push and pull agents from Modaic Hub @@ -75,7 +75,7 @@ agent = AutoProgram.from_precompiled( "farouk1/claude-code", config={ "signature": "message:str -> answer:str", - "model": "sonnet", + "model": "claude-opus-4-5-20251101", "permission_mode": "acceptEdits", "allowed_tools": ["Read", "Write", "Bash"], } @@ -92,14 +92,15 @@ For local development and creating your own agents: from claude_dspy import ClaudeCode, ClaudeCodeConfig # Create config -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() + +# Create agent +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory="." ) -# Create agent -agent = ClaudeCode(config) - # Use it result = agent(message="What files are in this directory?") print(result.answer) # String response @@ -119,13 +120,14 @@ class BugReport(BaseModel): affected_files: list[str] # Create config with Pydantic output -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() + +agent = ClaudeCode( + config, signature="message:str -> report:BugReport", working_directory="." ) -agent = ClaudeCode(config) - result = agent(message="Analyze the bug in error.log") print(result.report.severity) # Typed access! print(result.report.affected_files) @@ -137,15 +139,15 @@ print(result.report.affected_files) from claude_dspy import ClaudeCode, ClaudeCodeConfig # Create your agent -config = ClaudeCodeConfig( +config = ClaudeCodeConfig(model="claude-opus-4-5-20251101") + +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", - model="claude-opus-4-5-20251101", permission_mode="acceptEdits", ) -agent = ClaudeCode(config) - # Test it locally result = agent(message="Test my agent") print(result.answer) @@ -164,9 +166,25 @@ Configuration object for ClaudeCode agents. class ClaudeCodeConfig: def __init__( self, + model: str = "claude-opus-4-5-20251101", # Default model + ) +``` + +**Parameters:** + +- **`model`** - Claude model to use (default: `"claude-opus-4-5-20251101"`) + +### ClaudeCode + +Main agent class. + +```python +class ClaudeCode(PrecompiledProgram): + def __init__( + self, + config: ClaudeCodeConfig, signature: str | type[Signature], # Required working_directory: str = ".", # Default: "." - model: str = "claude-opus-4-5-20251101", # Default model permission_mode: str | None = None, # Optional allowed_tools: list[str] | None = None, # Optional disallowed_tools: list[str] | None = None, # Optional @@ -178,9 +196,9 @@ class ClaudeCodeConfig: **Parameters:** +- **`config`** - ClaudeCodeConfig instance with model configuration - **`signature`** (required) - DSPy signature defining input/output fields (must have exactly 1 input and 1 output) - **`working_directory`** - Directory where Claude will execute commands (default: `"."`) -- **`model`** - Claude model to use (default: `"claude-opus-4-5-20251101"`) - **`permission_mode`** - Permission mode: `"default"`, `"acceptEdits"`, `"plan"`, `"bypassPermissions"` - **`allowed_tools`** - List of allowed tool names (e.g., `["Read", "Write", "Bash"]`) - **`disallowed_tools`** - List of disallowed tool names @@ -188,19 +206,6 @@ class ClaudeCodeConfig: - **`system_prompt`** - Custom system prompt or preset config - **`api_key`** - Anthropic API key (falls back to `ANTHROPIC_API_KEY` env var) -### ClaudeCode - -Main agent class. - -```python -class ClaudeCode(PrecompiledProgram): - def __init__(self, config: ClaudeCodeConfig) -``` - -**Parameters:** - -- **`config`** - ClaudeCodeConfig instance with agent configuration - #### Methods ##### `__call__(**kwargs) -> Prediction` (or `forward`) @@ -218,11 +223,12 @@ Execute the agent with an input message. **Example:** ```python -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory="." ) -agent = ClaudeCode(config) result = agent(message="Hello") print(result.answer) # Access typed output @@ -249,11 +255,12 @@ Async version of `__call__()` for use in async contexts. **Example:** ```python async def main(): - config = ClaudeCodeConfig( + config = ClaudeCodeConfig() + agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory="." ) - agent = ClaudeCode(config) result = await agent.aforward(message="Hello") print(result.answer) ``` @@ -270,11 +277,12 @@ Get the session ID for this agent instance. **Example:** ```python -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory="." ) -agent = ClaudeCode(config) print(agent.session_id) # None @@ -300,11 +308,12 @@ Each agent instance maintains a stateful session: ```python from claude_dspy import ClaudeCode, ClaudeCodeConfig -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", ) -agent = ClaudeCode(config) # Turn 1 result1 = agent(message="What's the main bug?") @@ -329,45 +338,57 @@ Want a new conversation? Create a new agent: ```python from claude_dspy import ClaudeCode, ClaudeCodeConfig -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() + +# Agent 1 - Task A +agent1 = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", ) - -# Agent 1 - Task A -agent1 = ClaudeCode(config) result1 = agent1(message="Analyze bug in module A") # Agent 2 - Task B (no context from Agent 1) -agent2 = ClaudeCode(config) +agent2 = ClaudeCode( + config, + signature="message:str -> answer:str", + working_directory=".", +) result2 = agent2(message="Analyze bug in module B") ``` -### Pattern 3: Output Field Descriptions +### Pattern 3: Field Descriptions for Enhanced Context -Enhance prompts with field descriptions: +Enhance prompts with signature docstrings and field descriptions - all automatically included in the prompt: ```python import dspy from claude_dspy import ClaudeCode, ClaudeCodeConfig class MySignature(dspy.Signature): - """Analyze code architecture.""" + """Analyze code architecture.""" # Used as task description - message: str = dspy.InputField() + message: str = dspy.InputField( + desc="Request to process" # Provides input context + ) analysis: str = dspy.OutputField( desc="A detailed markdown report with sections: " - "1) Architecture overview, 2) Key components, 3) Dependencies" + "1) Architecture overview, 2) Key components, 3) Dependencies" # Guides output format ) -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature=MySignature, working_directory=".", ) -agent = ClaudeCode(config) result = agent(message="Analyze this codebase") -# The description is automatically appended to the prompt +# The prompt sent to Claude will include: +# 1. Task: "Analyze code architecture." (from docstring) +# 2. Input context: "Request to process" (from InputField desc) +# 3. Your message: "Analyze this codebase" +# 4. Output guidance: "Please produce the following output: A detailed markdown report..." (from OutputField desc) ``` ### Pattern 4: Inspecting Execution Trace @@ -377,11 +398,12 @@ Access detailed execution information: ```python from claude_dspy import ClaudeCode, ClaudeCodeConfig, ToolUseItem, ToolResultItem -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", ) -agent = ClaudeCode(config) result = agent(message="Fix the bug") @@ -418,30 +440,33 @@ Control what the agent can do: from claude_dspy import ClaudeCode, ClaudeCodeConfig # Read-only (safest) -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", permission_mode="default", allowed_tools=["Read", "Glob", "Grep"], ) -agent = ClaudeCode(config) # Auto-accept file edits -config = ClaudeCodeConfig( +config = ClaudeCodeConfig(model="claude-opus-4-5-20251101") +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", permission_mode="acceptEdits", allowed_tools=["Read", "Write", "Edit"], ) -agent = ClaudeCode(config) # Sandbox mode for command execution -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", sandbox={"enabled": True}, ) -agent = ClaudeCode(config) ``` ## Advanced Examples @@ -458,14 +483,14 @@ class CodeReview(BaseModel): severity: str = Field(description="critical, high, medium, or low") recommendations: list[str] = Field(description="Actionable recommendations") -config = ClaudeCodeConfig( +config = ClaudeCodeConfig(model="claude-opus-4-5-20251101") +agent = ClaudeCode( + config, signature="message:str -> review:CodeReview", working_directory="/path/to/project", - model="sonnet", permission_mode="default", allowed_tools=["Read", "Glob", "Grep"], ) -agent = ClaudeCode(config) result = agent(message="Review the changes in src/main.py") @@ -479,13 +504,14 @@ for issue in result.review.issues: ```python from claude_dspy import ClaudeCode, ClaudeCodeConfig -config = ClaudeCodeConfig( +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, signature="message:str -> response:str", working_directory=".", permission_mode="acceptEdits", allowed_tools=["Read", "Write", "Bash"], ) -agent = ClaudeCode(config) # Turn 1: Find the bug result1 = agent(message="Find the bug in src/calculator.py") @@ -511,11 +537,12 @@ import asyncio from claude_dspy import ClaudeCode, ClaudeCodeConfig async def main(): - config = ClaudeCodeConfig( + config = ClaudeCodeConfig() + agent = ClaudeCode( + config, signature="message:str -> answer:str", working_directory=".", ) - agent = ClaudeCode(config) # Use aforward in async context result = await agent.aforward(message="Analyze this code") @@ -582,6 +609,46 @@ sig = dspy.Signature('message:str -> report:BugReport') # Response parsed with BugReport.model_validate_json() ``` +### Prompt Building + +ClaudeCode automatically builds rich prompts from your signature to provide maximum context to Claude: + +```python +class MySignature(dspy.Signature): + """Analyze code quality.""" # 1. Task description + + message: str = dspy.InputField( + desc="Path to file or module" # 2. Input context + ) + report: str = dspy.OutputField( + desc="Markdown report with issues and recommendations" # 3. Output guidance + ) + +config = ClaudeCodeConfig() +agent = ClaudeCode( + config, + signature=MySignature, + working_directory="." +) +result = agent(message="Analyze src/main.py") # 4. Your actual input +``` + +**The final prompt sent to Claude:** +``` +Task: Analyze code quality. + +Input context: Path to file or module + +Analyze src/main.py + +Please produce the following output: Markdown report with issues and recommendations +``` + +This automatic context enhancement helps Claude better understand: +- **What** the overall task is (docstring) +- **What** the input represents (InputField desc) +- **What** format the output should have (OutputField desc) + ## Troubleshooting ### Error: "ClaudeCode requires exactly 1 input field" diff --git a/auto_classes.json b/auto_classes.json new file mode 100644 index 0000000..3e02c0b --- /dev/null +++ b/auto_classes.json @@ -0,0 +1,4 @@ +{ + "AutoConfig": "claude_dspy.agent.ClaudeCodeConfig", + "AutoProgram": "claude_dspy.agent.ClaudeCode" +} \ No newline at end of file diff --git a/claude_dspy/__init__.py b/claude_dspy/__init__.py new file mode 100644 index 0000000..212f164 --- /dev/null +++ b/claude_dspy/__init__.py @@ -0,0 +1,24 @@ +from claude_dspy.agent import ClaudeCode, ClaudeCodeConfig +from claude_dspy.trace import ( + TraceItem, + AgentMessageItem, + ThinkingItem, + ToolUseItem, + ToolResultItem, + ErrorItem, +) +from claude_dspy.utils import Usage + +__version__ = "0.1.0" + +__all__ = [ + "ClaudeCode", + "ClaudeCodeConfig", + "TraceItem", + "AgentMessageItem", + "ThinkingItem", + "ToolUseItem", + "ToolResultItem", + "ErrorItem", + "Usage", +] diff --git a/claude_dspy/agent.py b/claude_dspy/agent.py new file mode 100644 index 0000000..5695381 --- /dev/null +++ b/claude_dspy/agent.py @@ -0,0 +1,479 @@ +import asyncio +import json +import os +from pathlib import Path +from typing import Any, Optional + +from modaic import PrecompiledProgram, PrecompiledConfig +from pydantic import BaseModel + +import dspy +from dspy.primitives.prediction import Prediction + +from claude_agent_sdk import ( + ClaudeSDKClient, + ClaudeAgentOptions, + AssistantMessage, + ResultMessage, + SystemMessage, + TextBlock, + ThinkingBlock, + ToolUseBlock, + ToolResultBlock, +) + +from claude_dspy.trace import ( + TraceItem, + AgentMessageItem, + ThinkingItem, + ToolUseItem, + ToolResultItem, + ErrorItem, +) +from claude_dspy.utils import ( + Usage, + is_pydantic_model, + get_json_schema, + parse_json_response, + extract_text_from_response, +) + + +class ClaudeCodeConfig(PrecompiledConfig): + """Configuration for ClaudeCode agent.""" + model: str = "claude-opus-4-5-20251101" + +class ClaudeCodeKwargs(BaseModel): + model_config = {"arbitrary_types_allowed": True} + + signature: Any # str | dspy.Signature (validated manually) + api_key: str | None = None + working_directory: str = "." + permission_mode: str | None = None + allowed_tools: list[str] | None = None + disallowed_tools: list[str] | None = None + sandbox: dict[str, Any] | None = None + system_prompt: str | dict[str, Any] | None = None + +class ClaudeCode(PrecompiledProgram): + """DSPy module that wraps Claude Code SDK. + + Each agent instance maintains a stateful conversation session. + Perfect for multi-turn agentic workflows with context preservation. + + Example: + >>> config = ClaudeCodeConfig() + >>> agent = ClaudeCode( + ... config, + ... signature='message:str -> answer:str', + ... working_directory="." + ... ) + >>> result = agent(message="What files are in this directory?") + >>> print(result.answer) # Typed access + >>> print(result.trace) # Execution trace + >>> print(result.usage) # Token usage + """ + + config: ClaudeCodeConfig + + def __init__( + self, + config: ClaudeCodeConfig, + **kwargs: dict, + ): + super().__init__(config=config) + + args = ClaudeCodeKwargs(**kwargs) + + signature = args.signature + api_key = args.api_key + working_directory = args.working_directory + permission_mode = args.permission_mode + allowed_tools = args.allowed_tools + disallowed_tools = args.disallowed_tools + sandbox = args.sandbox + system_prompt = args.system_prompt + + # parse and validate signature + if isinstance(signature, str): + self.signature = dspy.Signature(signature) + else: + self.signature = signature + + # validate signature has exactly 1 input and 1 output + input_fields = list(self.signature.input_fields.keys()) + output_fields = list(self.signature.output_fields.keys()) + + if len(input_fields) != 1: + raise ValueError( + f"ClaudeCode requires exactly 1 input field, got {len(input_fields)}. " + f"Found: {input_fields}" + ) + + if len(output_fields) != 1: + raise ValueError( + f"ClaudeCode requires exactly 1 output field, got {len(output_fields)}. " + f"Found: {output_fields}" + ) + + self.input_field_name = input_fields[0] + self.output_field_name = output_fields[0] + self.input_field = self.signature.input_fields[self.input_field_name] + self.output_field = self.signature.output_fields[self.output_field_name] + + # store config values + self.working_directory = Path(working_directory).resolve() + self.model = config.model + self.permission_mode = permission_mode + self.allowed_tools = allowed_tools + self.disallowed_tools = disallowed_tools + self.sandbox = sandbox + self.system_prompt = system_prompt + self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY") + # No extra options since all kwargs are parsed by ClaudeCodeKwargs + self.extra_options = {} + + # determine output format upfront + self.output_format = self._get_output_format() + + # session state + self._client: Optional[ClaudeSDKClient] = None + self._session_id: Optional[str] = None + self._is_connected = False + + @property + def session_id(self) -> Optional[str]: + """Get the session ID for this agent instance. + + Returns None until first forward() call. + """ + return self._session_id + + def _create_client(self) -> ClaudeSDKClient: + """Create ClaudeSDKClient with configured options.""" + options = ClaudeAgentOptions( + cwd=str(self.working_directory), + model=self.model, + permission_mode=self.permission_mode, + allowed_tools=self.allowed_tools or [], + disallowed_tools=self.disallowed_tools or [], + sandbox=self.sandbox, + system_prompt=self.system_prompt, + output_format=self.output_format, # include output format + **self.extra_options, + ) + + # set API key if provided + if self.api_key: + os.environ["ANTHROPIC_API_KEY"] = self.api_key + + return ClaudeSDKClient(options=options) + + def _build_prompt(self, input_value: str) -> str: + """Build prompt from signature docstring, field descriptions, and input value.""" + prompt_parts = [] + + # add signature docstring if present + if self.signature.__doc__: + doc = self.signature.__doc__.strip() + if doc: + prompt_parts.append(f"Task: {doc}") + + # add input field description if present + # DSPy fields store desc in json_schema_extra + input_desc = None + if ( + hasattr(self.input_field, "json_schema_extra") + and self.input_field.json_schema_extra + ): + input_desc = self.input_field.json_schema_extra.get("desc") + + if input_desc: + prompt_parts.append(f"Input context: {input_desc}") + + # add the actual input value + prompt_parts.append(input_value) + + # add output field description if present + output_desc = None + if ( + hasattr(self.output_field, "json_schema_extra") + and self.output_field.json_schema_extra + ): + output_desc = self.output_field.json_schema_extra.get("desc") + + if output_desc: + prompt_parts.append(f"\nPlease produce the following output: {output_desc}") + + # for Pydantic outputs, add explicit JSON instructions + if self.output_format: + schema = self.output_format["schema"] + prompt_parts.append( + f"\nYou MUST respond with ONLY valid JSON matching this schema:\n" + f"{json.dumps(schema, indent=2)}\n\n" + f"Do not include any explanatory text, markdown formatting, or code blocks. " + f"Return ONLY the raw JSON object." + ) + + return "\n\n".join(prompt_parts) + + def _get_output_format(self) -> Optional[dict[str, Any]]: + """Get output format configuration for structured outputs.""" + output_type = self.output_field.annotation + + if is_pydantic_model(output_type): + schema = get_json_schema(output_type) + return { + "type": "json_schema", + "schema": schema, + } + + return None + + async def _run_async(self, prompt: str) -> tuple[str, list[TraceItem], Usage]: + """Run the agent asynchronously and collect results.""" + # create client if needed + if self._client is None: + self._client = self._create_client() + + # connect if not already connected + if not self._is_connected: + await self._client.connect() + self._is_connected = True + + # send query (output_format already configured in options) + await self._client.query(prompt) + + # collect messages and build trace + trace: list[TraceItem] = [] + usage = Usage() + response_text = "" + + async for message in self._client.receive_response(): + # handle assistant messages + if isinstance(message, AssistantMessage): + for block in message.content: + if isinstance(block, TextBlock): + response_text += block.text + trace.append( + AgentMessageItem(text=block.text, model=message.model) + ) + elif isinstance(block, ThinkingBlock): + trace.append( + ThinkingItem(text=block.thinking, model=message.model) + ) + elif isinstance(block, ToolUseBlock): + # Handle StructuredOutput tool (contains JSON response) + if block.name == "StructuredOutput": + # The JSON is directly in the tool input (already a dict) + response_text = json.dumps(block.input) + + trace.append( + ToolUseItem( + tool_name=block.name, + tool_input=block.input, + tool_use_id=block.id, + ) + ) + elif isinstance(block, ToolResultBlock): + content_str = "" + if isinstance(block.content, str): + content_str = block.content + elif isinstance(block.content, list): + # Extract text from content blocks + for item in block.content: + if ( + isinstance(item, dict) + and item.get("type") == "text" + ): + content_str += item.get("text", "") + + trace.append( + ToolResultItem( + tool_name="", # Tool name not in ToolResultBlock + tool_use_id=block.tool_use_id, + content=content_str, + is_error=block.is_error or False, + ) + ) + + # handle result messages (final message with usage info) + elif isinstance(message, ResultMessage): + # store session ID + if hasattr(message, "session_id"): + self._session_id = message.session_id + + # extract usage + if hasattr(message, "usage") and message.usage: + usage_data = message.usage + usage = Usage( + input_tokens=usage_data.get("input_tokens", 0), + cached_input_tokens=usage_data.get( + "cache_read_input_tokens", 0 + ), + output_tokens=usage_data.get("output_tokens", 0), + ) + + # check for errors + if hasattr(message, "is_error") and message.is_error: + error_msg = ( + message.result + if hasattr(message, "result") + else "Unknown error" + ) + trace.append( + ErrorItem(message=error_msg, error_type="execution_error") + ) + raise RuntimeError(f"Agent execution failed: {error_msg}") + + # extract result if present (for structured outputs from result field) + # Note: structured outputs may come from StructuredOutput tool instead + if hasattr(message, "result") and message.result: + response_text = message.result + + # handle system messages + elif isinstance(message, SystemMessage): + # log system messages to trace but don't error + if hasattr(message, "data") and message.data: + data_str = str(message.data) + trace.append( + AgentMessageItem(text=f"[System: {data_str}]", model="system") + ) + + return response_text, trace, usage + + def forward(self, **kwargs: Any) -> Prediction: + """Execute the agent with an input message. + + Args: + **kwargs: Must contain the input field specified in signature + + Returns: + Prediction with: + - Typed output field (named according to signature) + - trace: list[TraceItem] - Execution trace + - usage: Usage - Token usage statistics + + Example: + >>> result = agent(message="Hello") + >>> print(result.answer) # Access typed output + >>> print(result.trace) # List of execution items + >>> print(result.usage) # Token usage stats + """ + # extract input value + if self.input_field_name not in kwargs: + raise ValueError( + f"Missing required input field: {self.input_field_name}. " + f"Received: {list(kwargs.keys())}" + ) + + input_value = kwargs[self.input_field_name] + + # build prompt + prompt = self._build_prompt(input_value) + + # run async execution in event loop + try: + loop = asyncio.get_event_loop() + if loop.is_running(): + # If already in async context, create new loop + import nest_asyncio + + nest_asyncio.apply() + response_text, trace, usage = loop.run_until_complete( + self._run_async(prompt) + ) + else: + response_text, trace, usage = loop.run_until_complete( + self._run_async(prompt) + ) + except RuntimeError: + # no event loop, create one + response_text, trace, usage = asyncio.run(self._run_async(prompt)) + + # parse response based on output type + output_type = self.output_field.annotation + if is_pydantic_model(output_type): + try: + parsed_output = parse_json_response(response_text, output_type) + except Exception as e: + raise ValueError( + f"Failed to parse Claude response as {output_type.__name__}: {e}\n" + f"Response: {response_text}" + ) + else: + # string output - extract text + parsed_output = extract_text_from_response(response_text) + + # return prediction with typed output, trace, and usage + return Prediction( + **{ + self.output_field_name: parsed_output, + "trace": trace, + "usage": usage, + } + ) + + async def aforward(self, **kwargs: Any) -> Prediction: + """Async version of forward(). + + Use this when already in an async context to avoid event loop issues. + + Args: + **kwargs: Must contain the input field specified in signature + + Returns: + Prediction with typed output, trace, and usage + """ + # extract input value + if self.input_field_name not in kwargs: + raise ValueError( + f"Missing required input field: {self.input_field_name}. " + f"Received: {list(kwargs.keys())}" + ) + + input_value = kwargs[self.input_field_name] + + # build prompt + prompt = self._build_prompt(input_value) + + # run async execution + response_text, trace, usage = await self._run_async(prompt) + + # parse response based on output type + output_type = self.output_field.annotation + if is_pydantic_model(output_type): + try: + parsed_output = parse_json_response(response_text, output_type) + except Exception as e: + raise ValueError( + f"Failed to parse Claude response as {output_type.__name__}: {e}\n" + f"Response: {response_text}" + ) + else: + # string output - extract text + parsed_output = extract_text_from_response(response_text) + + # return prediction with typed output, trace, and usage + return Prediction( + **{ + self.output_field_name: parsed_output, + "trace": trace, + "usage": usage, + } + ) + + async def disconnect(self) -> None: + """Disconnect from Claude Code and clean up resources.""" + if self._client and self._is_connected: + await self._client.disconnect() + self._is_connected = False + + def __del__(self): + """Cleanup on deletion.""" + if self._client and self._is_connected: + try: + asyncio.run(self.disconnect()) + except Exception: + # best effort cleanup + pass diff --git a/claude_dspy/trace.py b/claude_dspy/trace.py new file mode 100644 index 0000000..6587b0c --- /dev/null +++ b/claude_dspy/trace.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from typing import Any + + +@dataclass +class TraceItem: + """Base class for trace items.""" + + pass + + +@dataclass +class AgentMessageItem(TraceItem): + """Agent's text response.""" + + text: str + model: str + + +@dataclass +class ThinkingItem(TraceItem): + """Agent's internal reasoning (extended thinking).""" + + text: str + model: str + + +@dataclass +class ToolUseItem(TraceItem): + """Tool invocation request.""" + + tool_name: str + tool_input: dict[str, Any] + tool_use_id: str + + +@dataclass +class ToolResultItem(TraceItem): + """Tool execution result.""" + + tool_name: str + tool_use_id: str + content: str | list[dict[str, Any]] | None = None + is_error: bool = False + + +@dataclass +class ErrorItem(TraceItem): + """Error that occurred during execution.""" + + message: str + error_type: str | None = None diff --git a/claude_dspy/utils.py b/claude_dspy/utils.py new file mode 100644 index 0000000..cd38ffe --- /dev/null +++ b/claude_dspy/utils.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from typing import Any, get_origin, get_args + +from pydantic import BaseModel + + +@dataclass +class Usage: + """Token usage statistics.""" + + input_tokens: int = 0 + cached_input_tokens: int = 0 + output_tokens: int = 0 + + @property + def total_tokens(self) -> int: + """Total tokens used (input + output).""" + return self.input_tokens + self.output_tokens + + def __repr__(self) -> str: + return ( + f"Usage(input={self.input_tokens}, " + f"cached={self.cached_input_tokens}, " + f"output={self.output_tokens}, " + f"total={self.total_tokens})" + ) + + +def is_pydantic_model(type_hint: Any) -> bool: + """Check if a type hint is a Pydantic model.""" + try: + return isinstance(type_hint, type) and issubclass(type_hint, BaseModel) + except TypeError: + return False + + +def get_json_schema(pydantic_model: type[BaseModel]) -> dict[str, Any]: + """Generate JSON schema from Pydantic model. + + Sets additionalProperties to false to match Codex behavior. + """ + schema = pydantic_model.model_json_schema() + + # Recursively set additionalProperties: false for all objects + def set_additional_properties(obj: dict[str, Any]) -> None: + if isinstance(obj, dict): + if obj.get("type") == "object": + obj["additionalProperties"] = False + for value in obj.values(): + if isinstance(value, dict): + set_additional_properties(value) + elif isinstance(value, list): + for item in value: + if isinstance(item, dict): + set_additional_properties(item) + + set_additional_properties(schema) + return schema + + +def parse_json_response(response: str, pydantic_model: type[BaseModel]) -> BaseModel: + """Parse JSON response into Pydantic model. + + Raises: + json.JSONDecodeError: If response is not valid JSON + pydantic.ValidationError: If JSON doesn't match model schema + """ + return pydantic_model.model_validate_json(response) + + +def extract_text_from_response(response: str) -> str: + """Extract plain text from response. + + For string outputs, we just return the text as-is. + Claude Code may wrap responses in markdown or other formatting. + """ + return response.strip() diff --git a/config.json b/config.json index af38f78..d071d3a 100644 --- a/config.json +++ b/config.json @@ -1,14 +1,3 @@ { - "model": "claude-opus-4-5-20251101", - "signature": "message:str -> output:Output", - "working_directory": ".", - "permission_mode": "acceptEdits", - "allowed_tools": [ - "Read", - "Glob" - ], - "disallowed_tools": null, - "sandbox": null, - "system_prompt": null, - "api_key": null + "model": "claude-opus-4-5-20251101" } \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..ec22564 --- /dev/null +++ b/main.py @@ -0,0 +1,46 @@ +from claude_dspy import ClaudeCode, ClaudeCodeConfig +from pydantic import BaseModel +import dspy + + +class Output(BaseModel): + files: list[str] + + +class ClaudeCodeSignature(dspy.Signature): + message: str = dspy.InputField(desc="Request to process") + output: list[str] = dspy.OutputField(desc="List of files modified or created") + + +def main(): + # create config + config = ClaudeCodeConfig() + + # create agent + cc = ClaudeCode( + config, + signature=ClaudeCodeSignature, + working_directory=".", + permission_mode="acceptEdits", + allowed_tools=["Read", "Glob", "Write"], + ) + + # use it + print("Running ClaudeCode...") + result = cc( + message="Create a new file called helloworld.txt with the alphabet backwards" + ) + + print(f"Success: {result.output}") + print(result.usage) + + print(f"Session ID: {cc.session_id}") + cc.push_to_hub( + "farouk1/claude-code", + with_code=True, + commit_message="Add more functionality to signature description parsing", + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7ae1b5b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "claude-code" +version = "0.1.0" +description = "Claude Code SDK wrapped in a DSPy module" +readme = "README.md" +requires-python = ">=3.10" +dependencies = ["claude-agent-sdk>=0.1.12", "dspy>=3.0.4", "modaic>=0.7.0", "nest-asyncio>=1.6.0", "pydantic>=2.0.0"] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.21.0", +]