ccproxy-api 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

ccproxy/claude_sdk/converter.py

@@ -1,24 +1,22 @@
 """Message format converter for Claude SDK interactions."""
 
+import html
 import json
-from typing import Any, cast
-from xml.sax.saxutils import escape
+from collections.abc import Callable
+from typing import Any
 
 import structlog
 
+from ccproxy.config.claude import SDKMessageMode
 from ccproxy.core.async_utils import patched_typing
+from ccproxy.models import claude_sdk as sdk_models
+from ccproxy.models.messages import MessageResponse
 
 
 logger = structlog.get_logger(__name__)
 
 with patched_typing():
-    from claude_code_sdk import (
-        AssistantMessage,
-        ResultMessage,
-        TextBlock,
-        ToolResultBlock,
-        ToolUseBlock,
-    )
+    pass
 
 
 class MessageConverter:
@@ -26,6 +24,126 @@ class MessageConverter:
     Handles conversion between Anthropic API format and Claude SDK format.
     """
 
+    @staticmethod
+    def _format_json_data(
+        data: dict[str, Any],
+        pretty_format: bool = True,
+    ) -> str:
+        """
+        Format JSON data with optional indentation and newlines.
+
+        Args:
+            data: Dictionary to format as JSON
+            pretty_format: Whether to use pretty formatting (indented JSON with spacing)
+
+        Returns:
+            Formatted JSON string
+        """
+
+        if pretty_format:
+            # Pretty format with indentation and proper spacing
+            return json.dumps(data, indent=2, separators=(", ", ": "))
+        else:
+            # Compact format without indentation or spacing
+            return json.dumps(data, separators=(",", ":"))
+
+    @staticmethod
+    def _create_xml_formatted_text(
+        data: dict[str, Any], tag_name: str, pretty_format: bool = True
+    ) -> str:
+        """
+        Create XML-formatted text from data with consistent formatting.
+
+        Args:
+            data: Dictionary data to format as JSON and wrap in XML
+            tag_name: XML tag name to wrap the content
+            pretty_format: Whether to use pretty formatting
+
+        Returns:
+            Formatted XML string
+        """
+        formatted_json = MessageConverter._format_json_data(data, pretty_format)
+        escaped_json = MessageConverter._escape_content_for_xml(
+            formatted_json, pretty_format
+        )
+
+        if pretty_format:
+            return f"<{tag_name}>\n{escaped_json}\n</{tag_name}>\n"
+        else:
+            return f"<{tag_name}>{escaped_json}</{tag_name}>"
+
+    @staticmethod
+    def _create_streaming_chunks_with_content(
+        content_block: dict[str, Any],
+        index: int,
+        text_content: str | None = None,
+    ) -> list[tuple[str, dict[str, Any]]]:
+        """
+        Create streaming chunks with optional text delta content.
+
+        Args:
+            content_block: Content block for content_block_start
+            index: Content block index
+            text_content: Optional text content for content_block_delta
+
+        Returns:
+            List of streaming chunks
+        """
+        chunks = [
+            (
+                "content_block_start",
+                {
+                    "type": "content_block_start",
+                    "index": index,
+                    "content_block": content_block,
+                },
+            )
+        ]
+
+        if text_content is not None:
+            chunks.append(
+                (
+                    "content_block_delta",
+                    {
+                        "type": "content_block_delta",
+                        "index": index,
+                        "delta": {"type": "text_delta", "text": text_content},
+                    },
+                )
+            )
+
+        chunks.append(
+            (
+                "content_block_stop",
+                {
+                    "type": "content_block_stop",
+                    "index": index,
+                },
+            )
+        )
+
+        return chunks
+
+    @staticmethod
+    def _escape_content_for_xml(content: str, pretty_format: bool = True) -> str:
+        """
+        Escape content for inclusion in XML tags.
+
+        Args:
+            content: Content to escape
+            pretty_format: Whether to use pretty formatting (no escaping) or compact (escaped)
+
+        Returns:
+            Escaped or unescaped content based on formatting mode
+        """
+        if pretty_format:
+            # Pretty format: no escaping, content as-is
+            return content
+        else:
+            # Compact format: escape special XML characters
+
+            return html.escape(content)
+
     @staticmethod
     def format_messages_to_prompt(messages: list[dict[str, Any]]) -> str:
         """
@@ -61,62 +179,14 @@ class MessageConverter:
 
         return "\n\n".join(prompt_parts)
 
-    @staticmethod
-    def extract_text_from_content(
-        content: TextBlock | ToolUseBlock | ToolResultBlock,
-    ) -> str:
-        """
-        Extract text content from Claude SDK content blocks.
-
-        Args:
-            content: List of content blocks from Claude SDK
-
-        Returns:
-            Extracted text content
-        """
-        if isinstance(content, TextBlock):
-            return content.text
-        elif isinstance(content, ToolUseBlock):
-            # Return full XML representation of ToolUseBlock
-            tool_id = escape(str(getattr(content, "id", f"tool_{id(content)}")))
-            tool_name = escape(content.name)
-            tool_input = getattr(content, "input", {}) or {}
-            # Convert input dict to JSON string and escape for XML
-            input_json = escape(json.dumps(tool_input, ensure_ascii=False))
-            return f'<tooluseblock id="{tool_id}" name="{tool_name}">{input_json}</tooluseblock>'
-        elif isinstance(content, ToolResultBlock):
-            # Return full XML representation of ToolResultBlock
-            tool_use_id = escape(str(getattr(content, "tool_use_id", "")))
-            result_content = content.content if isinstance(content.content, str) else ""
-            escaped_content = escape(result_content)
-            return f'<toolresultblock tool_use_id="{tool_use_id}">{escaped_content}</toolresultblock>'
-
-    @staticmethod
-    def extract_contents(
-        contents: list[TextBlock | ToolUseBlock | ToolResultBlock],
-    ) -> str:
-        """
-        Extract content from Claude SDK blocks, preserving custom blocks.
-
-        Args:
-            content: List of content blocks from Claude SDK
-
-        Returns:
-            Content with thinking blocks preserved
-        """
-        text_parts = []
-
-        for block in contents:
-            text_parts.append(MessageConverter.extract_text_from_content(block))
-
-        return " ".join(text_parts)
-
     @staticmethod
     def convert_to_anthropic_response(
-        assistant_message: AssistantMessage,
-        result_message: ResultMessage,
+        assistant_message: sdk_models.AssistantMessage,
+        result_message: sdk_models.ResultMessage,
         model: str,
-    ) -> dict[str, Any]:
+        mode: SDKMessageMode = SDKMessageMode.FORWARD,
+        pretty_format: bool = True,
+    ) -> "MessageResponse":
         """
         Convert Claude SDK messages to Anthropic API response format.
 
@@ -124,180 +194,180 @@ class MessageConverter:
             assistant_message: The assistant message from Claude SDK
             result_message: The result message from Claude SDK
             model: The model name used
+            mode: System message handling mode (forward, ignore, formatted)
+            pretty_format: Whether to use pretty formatting (true: indented JSON with newlines, false: compact with escaped content)
 
         Returns:
             Response in Anthropic API format
         """
         # Extract token usage from result message
-        # First try to get usage from the usage field (preferred method)
-        usage = getattr(result_message, "usage", {})
-        if usage:
-            input_tokens = usage.get("input_tokens", 0)
-            output_tokens = usage.get("output_tokens", 0)
-            cache_read_tokens = usage.get("cache_read_input_tokens", 0)
-            cache_write_tokens = usage.get("cache_creation_input_tokens", 0)
-        else:
-            # Fallback to direct attributes
-            input_tokens = getattr(result_message, "input_tokens", 0)
-            output_tokens = getattr(result_message, "output_tokens", 0)
-            cache_read_tokens = getattr(result_message, "cache_read_tokens", 0)
-            cache_write_tokens = getattr(result_message, "cache_write_tokens", 0)
+        usage = result_message.usage_model
 
         # Log token extraction for debugging
-        from structlog import get_logger
-
-        logger = get_logger(__name__)
-
-        logger.debug(
-            "assistant_message_content",
-            content_blocks=[
-                type(block).__name__ for block in assistant_message.content
-            ],
-            content_count=len(assistant_message.content),
-            first_block_text=(
-                assistant_message.content[0].text[:100]
-                if assistant_message.content
-                and hasattr(assistant_message.content[0], "text")
-                else None
-            ),
-        )
+        # logger.debug(
+        #     "assistant_message_content",
+        #     content_blocks=[block.type for block in assistant_message.content],
+        #     content_count=len(assistant_message.content),
+        # )
 
         logger.debug(
             "token_usage_extracted",
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
-            cache_read_tokens=cache_read_tokens,
-            cache_write_tokens=cache_write_tokens,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            cache_read_tokens=usage.cache_read_input_tokens,
+            cache_write_tokens=usage.cache_creation_input_tokens,
             source="claude_sdk",
         )
 
-        # Calculate total tokens
-        total_tokens = input_tokens + output_tokens
-
         # Build usage information
-        usage_info = {
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_tokens": cache_read_tokens,
-            "cache_write_tokens": cache_write_tokens,
-            "total_tokens": total_tokens,
-        }
+        usage_info = usage.model_dump(mode="json")
 
         # Add cost information if available
-        total_cost_usd = getattr(result_message, "total_cost_usd", None)
-        if total_cost_usd is not None:
-            usage_info["cost_usd"] = total_cost_usd
+        if result_message.total_cost_usd is not None:
+            usage_info["cost_usd"] = result_message.total_cost_usd
 
         # Convert content blocks to Anthropic format, preserving thinking blocks
         content_blocks = []
 
         for block in assistant_message.content:
-            if isinstance(block, TextBlock):
-                # Parse text content for thinking blocks
+            if isinstance(block, sdk_models.TextBlock):
+                # Handle text content directly without thinking block parsing
                 text = block.text
-
-                # Check if the text contains thinking blocks
-                import re
-
-                thinking_pattern = r'<thinking signature="([^"]*)">(.*?)</thinking>'
-
-                # Split the text by thinking blocks
-                last_end = 0
-                for match in re.finditer(thinking_pattern, text, re.DOTALL):
-                    # Add any text before the thinking block
-                    before_text = text[last_end : match.start()].strip()
-                    if before_text:
-                        content_blocks.append({"type": "text", "text": before_text})
-
-                    # Add the thinking block
-                    signature, thinking_text = match.groups()
-                    content_blocks.append(
-                        {
-                            "type": "thinking",
-                            "text": thinking_text,
-                            "signature": signature,
-                        }
+                if mode == SDKMessageMode.FORMATTED:
+                    escaped_text = MessageConverter._escape_content_for_xml(
+                        text, pretty_format
                     )
-
-                    last_end = match.end()
-
-                # Add any remaining text after the last thinking block
-                remaining_text = text[last_end:].strip()
-                if remaining_text:
-                    content_blocks.append({"type": "text", "text": remaining_text})
-
-                # If no thinking blocks were found, add the entire text as a text block
-                if last_end == 0 and text:
+                    formatted_text = (
+                        f"<text>\n{escaped_text}\n</text>\n"
+                        if pretty_format
+                        else f"<text>{escaped_text}</text>"
+                    )
+                    content_blocks.append({"type": "text", "text": formatted_text})
+                else:
                     content_blocks.append({"type": "text", "text": text})
 
-            elif isinstance(block, ToolUseBlock):
-                tool_input = getattr(block, "input", {}) or {}
-                content_blocks.append(
-                    cast(
-                        dict[str, Any],
-                        {
-                            "type": "tool_use",
-                            "id": getattr(block, "id", f"tool_{id(block)}"),
-                            "name": block.name,
-                            "input": tool_input,
-                        },
+            elif isinstance(block, sdk_models.ToolUseBlock):
+                if mode == SDKMessageMode.FORWARD:
+                    content_blocks.append(block.to_sdk_block())
+                elif mode == SDKMessageMode.FORMATTED:
+                    tool_data = block.model_dump(mode="json")
+                    formatted_json = MessageConverter._format_json_data(
+                        tool_data, pretty_format
                     )
-                )
-            elif isinstance(block, ToolResultBlock):
-                content_blocks.append(
-                    {
-                        "type": "tool_result",
-                        "tool_use_id": getattr(block, "tool_use_id", ""),
-                        "content": block.content
-                        if isinstance(block.content, str)
-                        else "",
-                    }
-                )
+                    escaped_json = MessageConverter._escape_content_for_xml(
+                        formatted_json, pretty_format
+                    )
+                    formatted_text = (
+                        f"<tool_use_sdk>\n{escaped_json}\n</tool_use_sdk>\n"
+                        if pretty_format
+                        else f"<tool_use_sdk>{escaped_json}</tool_use_sdk>"
+                    )
+                    content_blocks.append({"type": "text", "text": formatted_text})
+
+            elif isinstance(block, sdk_models.ToolResultBlock):
+                if mode == SDKMessageMode.FORWARD:
+                    content_blocks.append(block.to_sdk_block())
+                elif mode == SDKMessageMode.FORMATTED:
+                    tool_result_data = block.model_dump(mode="json")
+                    formatted_json = MessageConverter._format_json_data(
+                        tool_result_data, pretty_format
+                    )
+                    escaped_json = MessageConverter._escape_content_for_xml(
+                        formatted_json, pretty_format
+                    )
+                    formatted_text = (
+                        f"<tool_result_sdk>\n{escaped_json}\n</tool_result_sdk>\n"
+                        if pretty_format
+                        else f"<tool_result_sdk>{escaped_json}</tool_result_sdk>"
+                    )
+                    content_blocks.append({"type": "text", "text": formatted_text})
 
-        return {
-            "id": f"msg_{result_message.session_id}",
-            "type": "message",
-            "role": "assistant",
-            "content": content_blocks,
-            "model": model,
-            "stop_reason": getattr(result_message, "stop_reason", "end_turn"),
-            "stop_sequence": None,
-            "usage": usage_info,
-        }
+            elif isinstance(block, sdk_models.ThinkingBlock):
+                if mode == SDKMessageMode.FORWARD:
+                    thinking_block = {
+                        "type": "thinking",
+                        "thinking": block.thinking,
+                    }
+                    if block.signature is not None:
+                        thinking_block["signature"] = block.signature
+                    content_blocks.append(thinking_block)
+                elif mode == SDKMessageMode.FORMATTED:
+                    # Format thinking block with signature in XML tag attribute
+                    signature_attr = (
+                        f' signature="{block.signature}"' if block.signature else ""
+                    )
+                    if pretty_format:
+                        escaped_text = MessageConverter._escape_content_for_xml(
+                            block.thinking, pretty_format
+                        )
+                        formatted_text = (
+                            f"<thinking{signature_attr}>\n{escaped_text}\n</thinking>\n"
+                        )
+                    else:
+                        escaped_text = MessageConverter._escape_content_for_xml(
+                            block.thinking, pretty_format
+                        )
+                        formatted_text = (
+                            f"<thinking{signature_attr}>{escaped_text}</thinking>"
+                        )
+                    content_blocks.append({"type": "text", "text": formatted_text})
+
+        return MessageResponse.model_validate(
+            {
+                "id": f"msg_{result_message.session_id}",
+                "type": "message",
+                "role": "assistant",
+                "content": content_blocks,
+                "model": model,
+                "stop_reason": result_message.stop_reason,
+                "stop_sequence": None,
+                "usage": usage_info,
+            }
+        )
 
     @staticmethod
-    def create_streaming_start_chunk(message_id: str, model: str) -> dict[str, Any]:
+    def create_streaming_start_chunks(
+        message_id: str, model: str, input_tokens: int = 0
+    ) -> list[tuple[str, dict[str, Any]]]:
         """
-        Create the initial streaming chunk for Anthropic API format.
+        Create the initial streaming chunks for Anthropic API format.
 
         Args:
             message_id: The message ID
             model: The model name
+            input_tokens: Number of input tokens for the request
 
         Returns:
-            Initial streaming chunk
+            List of tuples (event_type, chunk) for initial streaming chunks
         """
-        return {
-            "id": message_id,
-            "type": "message_start",
-            "message": {
-                "id": message_id,
-                "type": "message",
-                "role": "assistant",
-                "content": [],
-                "model": model,
-                "stop_reason": None,
-                "stop_sequence": None,
-                "usage": {
-                    "input_tokens": 0,
-                    "output_tokens": 0,
-                    "total_tokens": 0,
+        return [
+            # First, send message_start with event type
+            (
+                "message_start",
+                {
+                    "type": "message_start",
+                    "message": {
+                        "id": message_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": model,
+                        "content": [],
+                        "stop_reason": None,
+                        "stop_sequence": None,
+                        "usage": {
+                            "input_tokens": input_tokens,
+                            "cache_creation_input_tokens": 0,
+                            "cache_read_input_tokens": 0,
+                            "output_tokens": 1,
+                            "service_tier": "standard",
+                        },
+                    },
                 },
-            },
-        }
+            ),
+        ]
 
     @staticmethod
-    def create_streaming_delta_chunk(text: str) -> dict[str, Any]:
+    def create_streaming_delta_chunk(text: str) -> tuple[str, dict[str, Any]]:
         """
         Create a streaming delta chunk for Anthropic API format.
 
@@ -305,27 +375,152 @@ class MessageConverter:
             text: The text content to include
 
         Returns:
-            Delta chunk
+            Tuple of (event_type, chunk)
         """
-        return {
-            "type": "content_block_delta",
-            "index": 0,
-            "delta": {"type": "text_delta", "text": text},
-        }
+        return (
+            "content_block_delta",
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": text},
+            },
+        )
 
     @staticmethod
-    def create_streaming_end_chunk(stop_reason: str = "end_turn") -> dict[str, Any]:
+    def create_streaming_end_chunks(
+        stop_reason: str = "end_turn", stop_sequence: str | None = None
+    ) -> list[tuple[str, dict[str, Any]]]:
         """
-        Create the final streaming chunk for Anthropic API format.
+        Create the final streaming chunks for Anthropic API format.
 
         Args:
             stop_reason: The reason for stopping
+            stop_sequence: The stop sequence used (if any)
+
+        Returns:
+            List of tuples (event_type, chunk) for final streaming chunks
+        """
+        return [
+            # Then, send message_delta with stop reason and usage
+            (
+                "message_delta",
+                {
+                    "type": "message_delta",
+                    "delta": {
+                        "stop_reason": stop_reason,
+                        "stop_sequence": stop_sequence,
+                    },
+                    "usage": {"output_tokens": 0},
+                },
+            ),
+            # Finally, send message_stop
+            ("message_stop", {"type": "message_stop"}),
+        ]
+
+    @staticmethod
+    def create_ping_chunk() -> tuple[str, dict[str, Any]]:
+        """
+        Create a ping chunk for keeping the connection alive.
+
+        Returns:
+            Tuple of (event_type, chunk)
+        """
+        return ("ping", {"type": "ping"})
+
+    @staticmethod
+    def _create_sdk_content_block(
+        sdk_object: sdk_models.SystemMessage
+        | sdk_models.ToolUseBlock
+        | sdk_models.ToolResultBlock
+        | sdk_models.ResultMessage,
+        mode: SDKMessageMode = SDKMessageMode.FORWARD,
+        pretty_format: bool = True,
+        xml_tag: str = "sdk_block",
+        forward_converter: Callable[[Any], dict[str, Any]] | None = None,
+    ) -> dict[str, Any] | None:
+        """
+        Generic method to create content blocks for SDK objects in non-streaming responses.
+
+        Args:
+            sdk_object: The SDK object to convert
+            mode: System message handling mode
+            pretty_format: Whether to use pretty formatting
+            xml_tag: XML tag name for FORMATTED mode
+            forward_converter: Optional converter function for FORWARD mode
+
+        Returns:
+            Content block dict for the SDK object, or None if mode is IGNORE
+        """
+        if mode == SDKMessageMode.IGNORE:
+            return None
+        elif mode == SDKMessageMode.FORWARD:
+            if forward_converter:
+                return forward_converter(sdk_object)
+            else:
+                return sdk_object.model_dump(mode="json")
+        elif mode == SDKMessageMode.FORMATTED:
+            object_data = sdk_object.model_dump(mode="json")
+            formatted_json = MessageConverter._format_json_data(
+                object_data, pretty_format
+            )
+            escaped_json = MessageConverter._escape_content_for_xml(
+                formatted_json, pretty_format
+            )
+            formatted_text = (
+                f"<{xml_tag}>\n{escaped_json}\n</{xml_tag}>\n"
+                if pretty_format
+                else f"<{xml_tag}>{escaped_json}</{xml_tag}>"
+            )
+            return {
+                "type": "text",
+                "text": formatted_text,
+            }
+
+    @staticmethod
+    def _create_sdk_content_block_chunks(
+        sdk_object: sdk_models.SystemMessage
+        | sdk_models.ToolUseBlock
+        | sdk_models.ToolResultBlock
+        | sdk_models.ResultMessage,
+        mode: SDKMessageMode = SDKMessageMode.FORWARD,
+        index: int = 0,
+        pretty_format: bool = True,
+        xml_tag: str = "sdk_block",
+        sdk_block_converter: Callable[[Any], dict[str, Any]] | None = None,
+    ) -> list[tuple[str, dict[str, Any]]]:
+        """
+        Generic method to create streaming chunks for SDK content blocks.
+
+        Args:
+            sdk_object: The SDK object (SystemMessage, ToolUseBlock, or ToolResultBlock)
+            mode: System message handling mode
+            index: The content block index
+            pretty_format: Whether to use pretty formatting
+            xml_tag: XML tag name for FORMATTED mode
+            sdk_block_converter: Optional converter function for FORWARD mode
 
         Returns:
-            Final streaming chunk
+            List of tuples (event_type, chunk) for streaming chunks
         """
-        return {
-            "type": "message_delta",
-            "delta": {"stop_reason": stop_reason},
-            "usage": {"output_tokens": 0},
-        }
+        if mode == SDKMessageMode.IGNORE:
+            return []
+        elif mode == SDKMessageMode.FORWARD:
+            content_block = (
+                sdk_block_converter(sdk_object)
+                if sdk_block_converter
+                else sdk_object.model_dump(mode="json")
+            )
+            return MessageConverter._create_streaming_chunks_with_content(
+                content_block=content_block,
+                index=index,
+            )
+        elif mode == SDKMessageMode.FORMATTED:
+            object_data = sdk_object.model_dump(mode="json")
+            formatted_text = MessageConverter._create_xml_formatted_text(
+                object_data, xml_tag, pretty_format
+            )
+            return MessageConverter._create_streaming_chunks_with_content(
+                content_block={"type": "text", "text": ""},
+                index=index,
+                text_content=formatted_text,
+            )