Add more functionality to signature description parsing

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