ccproxy-api 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

ccproxy/claude_sdk/client.py

@@ -4,25 +4,36 @@ from collections.abc import AsyncIterator
 from typing import Any
 
 import structlog
+from pydantic import BaseModel
 
 from ccproxy.core.async_utils import patched_typing
 from ccproxy.core.errors import ClaudeProxyError, ServiceUnavailableError
+from ccproxy.models import claude_sdk as sdk_models
 from ccproxy.observability import timed_operation
 
 
 with patched_typing():
     from claude_code_sdk import (
-        AssistantMessage,
+        AssistantMessage as SDKAssistantMessage,
+    )
+    from claude_code_sdk import (
         ClaudeCodeOptions,
         CLIConnectionError,
         CLIJSONDecodeError,
         CLINotFoundError,
         ProcessError,
-        ResultMessage,
-        SystemMessage,
-        UserMessage,
         query,
     )
+    from claude_code_sdk import (
+        ResultMessage as SDKResultMessage,
+    )
+    from claude_code_sdk import (
+        SystemMessage as SDKSystemMessage,
+    )
+    from claude_code_sdk import (
+        UserMessage as SDKUserMessage,
+    )
+
 
 logger = structlog.get_logger(__name__)
 
@@ -53,9 +64,14 @@ class ClaudeSDKClient:
 
     async def query_completion(
         self, prompt: str, options: ClaudeCodeOptions, request_id: str | None = None
-    ) -> AsyncIterator[UserMessage | AssistantMessage | SystemMessage | ResultMessage]:
+    ) -> AsyncIterator[
+        sdk_models.UserMessage
+        | sdk_models.AssistantMessage
+        | sdk_models.SystemMessage
+        | sdk_models.ResultMessage
+    ]:
         """
-        Execute a query using the Claude Code SDK.
+        Execute a query using the Claude Code SDK and yields strongly-typed Pydantic models.
 
         Args:
             prompt: The prompt string to send to Claude
@@ -63,7 +79,7 @@ class ClaudeSDKClient:
             request_id: Optional request ID for correlation
 
         Yields:
-            Messages from the Claude Code SDK
+            Strongly-typed Pydantic messages from ccproxy.claude_sdk.models
 
         Raises:
             ClaudeSDKError: If the query fails
@@ -75,7 +91,74 @@ class ClaudeSDKClient:
                 message_count = 0
                 async for message in query(prompt=prompt, options=options):
                     message_count += 1
-                    yield message
+
+                    logger.debug(
+                        "claude_sdk_raw_message_received",
+                        message_type=type(message).__name__,
+                        message_count=message_count,
+                        request_id=request_id,
+                        has_content=hasattr(message, "content")
+                        and bool(getattr(message, "content", None)),
+                        content_preview=str(message)[:150],
+                    )
+
+                    model_class: type[BaseModel] | None = None
+                    if isinstance(message, SDKUserMessage):
+                        model_class = sdk_models.UserMessage
+                    elif isinstance(message, SDKAssistantMessage):
+                        model_class = sdk_models.AssistantMessage
+                    elif isinstance(message, SDKSystemMessage):
+                        model_class = sdk_models.SystemMessage
+                    elif isinstance(message, SDKResultMessage):
+                        model_class = sdk_models.ResultMessage
+
+                    # Convert Claude SDK message to our Pydantic model
+                    try:
+                        if hasattr(message, "__dict__"):
+                            converted_message = model_class.model_validate(
+                                vars(message)
+                            )
+                        else:
+                            # For dataclass objects, use dataclass.asdict equivalent
+                            message_dict = {}
+                            if hasattr(message, "__dataclass_fields__"):
+                                message_dict = {
+                                    field: getattr(message, field)
+                                    for field in message.__dataclass_fields__
+                                }
+                            else:
+                                # Try to extract common attributes
+                                for attr in [
+                                    "content",
+                                    "subtype",
+                                    "data",
+                                    "session_id",
+                                    "stop_reason",
+                                    "usage",
+                                    "total_cost_usd",
+                                ]:
+                                    if hasattr(message, attr):
+                                        message_dict[attr] = getattr(message, attr)
+
+                            converted_message = model_class.model_validate(message_dict)
+
+                        logger.debug(
+                            "claude_sdk_message_converted_successfully",
+                            original_type=type(message).__name__,
+                            converted_type=type(converted_message).__name__,
+                            message_count=message_count,
+                            request_id=request_id,
+                        )
+                        yield converted_message
+                    except Exception as e:
+                        logger.warning(
+                            "claude_sdk_message_conversion_failed",
+                            message_type=type(message).__name__,
+                            model_class=model_class.__name__,
+                            error=str(e),
+                        )
+                        # Skip invalid messages rather than crashing
+                        continue
 
                 # Store final metrics
                 op["message_count"] = message_count