Add more functionality to signature description parsing

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"

auto_classes.json

@@ -0,0 +1,4 @@
+{
+  "AutoConfig": "claude_dspy.agent.ClaudeCodeConfig",
+  "AutoProgram": "claude_dspy.agent.ClaudeCode"
+}

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",
+]

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

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

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()

config.json

@@ -1,14 +1,3 @@
 {
-  "model": "claude-opus-4-5-20251101",
+  "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
 }

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()

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",
+]