massgen 0.0.3__py3-none-any.whl

massgen/backend/chat_completions.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+"""
+Base class for backends using OpenAI Chat Completions API format.
+Handles common message processing, tool conversion, and streaming patterns.
+"""
+
+import os
+from typing import Dict, List, Any, AsyncGenerator, Optional
+from .base import LLMBackend, StreamChunk
+
+
+class ChatCompletionsBackend(LLMBackend):
+    """Base class for backends using Chat Completions API with shared streaming logic."""
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key, **kwargs)
+
+    def convert_tools_to_chat_completions_format(
+        self, tools: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Convert tools from Response API format to Chat Completions format if needed.
+
+        Response API format: {"type": "function", "name": ..., "description": ..., "parameters": ...}
+        Chat Completions format: {"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}
+        """
+        if not tools:
+            return tools
+
+        converted_tools = []
+        for tool in tools:
+            if tool.get("type") == "function":
+                if "function" in tool:
+                    # Already in Chat Completions format
+                    converted_tools.append(tool)
+                elif "name" in tool and "description" in tool:
+                    # Response API format - convert to Chat Completions format
+                    converted_tools.append(
+                        {
+                            "type": "function",
+                            "function": {
+                                "name": tool["name"],
+                                "description": tool["description"],
+                                "parameters": tool.get("parameters", {}),
+                            },
+                        }
+                    )
+                else:
+                    # Unknown format - keep as-is
+                    converted_tools.append(tool)
+            else:
+                # Non-function tool - keep as-is
+                converted_tools.append(tool)
+
+        return converted_tools
+
+    async def handle_chat_completions_stream(
+        self, stream, enable_web_search: bool = False
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Handle standard Chat Completions API streaming format."""
+        content = ""
+        current_tool_calls = {}
+        search_sources_used = 0
+        citations = []
+
+        async for chunk in stream:
+            try:
+                if hasattr(chunk, "choices") and chunk.choices:
+                    choice = chunk.choices[0]
+
+                    # Handle content delta
+                    if hasattr(choice, "delta") and choice.delta:
+                        if hasattr(choice.delta, "content") and choice.delta.content:
+                            content_chunk = choice.delta.content
+                            content += content_chunk
+                            yield StreamChunk(type="content", content=content_chunk)
+
+                        # Handle tool calls streaming
+                        if (
+                            hasattr(choice.delta, "tool_calls")
+                            and choice.delta.tool_calls
+                        ):
+                            for tool_call_delta in choice.delta.tool_calls:
+                                index = getattr(tool_call_delta, "index", 0)
+
+                                if index not in current_tool_calls:
+                                    current_tool_calls[index] = {
+                                        "id": "",
+                                        "name": "",
+                                        "arguments": "",
+                                    }
+
+                                if (
+                                    hasattr(tool_call_delta, "id")
+                                    and tool_call_delta.id
+                                ):
+                                    current_tool_calls[index]["id"] = tool_call_delta.id
+
+                                if (
+                                    hasattr(tool_call_delta, "function")
+                                    and tool_call_delta.function
+                                ):
+                                    if (
+                                        hasattr(tool_call_delta.function, "name")
+                                        and tool_call_delta.function.name
+                                    ):
+                                        current_tool_calls[index][
+                                            "name"
+                                        ] = tool_call_delta.function.name
+
+                                    if (
+                                        hasattr(tool_call_delta.function, "arguments")
+                                        and tool_call_delta.function.arguments
+                                    ):
+                                        current_tool_calls[index][
+                                            "arguments"
+                                        ] += tool_call_delta.function.arguments
+
+                    # Handle finish reason
+                    if hasattr(choice, "finish_reason") and choice.finish_reason:
+                        if choice.finish_reason == "tool_calls" and current_tool_calls:
+                            # Convert accumulated tool calls to final format
+                            final_tool_calls = []
+                            for index in sorted(current_tool_calls.keys()):
+                                tool_call = current_tool_calls[index]
+
+                                # Parse arguments as JSON
+                                arguments = tool_call["arguments"]
+                                if isinstance(arguments, str):
+                                    try:
+                                        import json
+
+                                        arguments = (
+                                            json.loads(arguments)
+                                            if arguments.strip()
+                                            else {}
+                                        )
+                                    except json.JSONDecodeError:
+                                        arguments = {}
+
+                                final_tool_calls.append(
+                                    {
+                                        "id": tool_call["id"] or f"call_{index}",
+                                        "type": "function",
+                                        "function": {
+                                            "name": tool_call["name"],
+                                            "arguments": arguments,
+                                        },
+                                    }
+                                )
+
+                            yield StreamChunk(
+                                type="tool_calls", tool_calls=final_tool_calls
+                            )
+
+                            # Build and yield complete message
+                            complete_message = {
+                                "role": "assistant",
+                                "content": content.strip(),
+                            }
+                            if final_tool_calls:
+                                complete_message["tool_calls"] = final_tool_calls
+                            yield StreamChunk(
+                                type="complete_message",
+                                complete_message=complete_message,
+                            )
+                        elif choice.finish_reason in ["stop", "length"]:
+                            if search_sources_used > 0:
+                                yield StreamChunk(
+                                    type="content",
+                                    content=f"\n✅ [Live Search Complete] Used {search_sources_used} sources\n",
+                                )
+
+                            # Check for citations before building complete message
+                            if (
+                                hasattr(chunk, "citations")
+                                and chunk.citations
+                                and len(chunk.citations) > 0
+                            ):
+                                if enable_web_search:
+                                    citation_text = "\n📚 **Citations:**\n"
+                                    for i, citation in enumerate(chunk.citations, 1):
+                                        citation_text += f"{i}. {citation}\n"
+                                    yield StreamChunk(
+                                        type="content", content=citation_text
+                                    )
+
+                            # Build and yield complete message (no tool calls)
+                            complete_message = {
+                                "role": "assistant",
+                                "content": content.strip(),
+                            }
+                            yield StreamChunk(
+                                type="complete_message",
+                                complete_message=complete_message,
+                            )
+
+                            yield StreamChunk(type="done")
+                        return
+
+                # Check for usage information (search sources) and citations
+                if hasattr(chunk, "usage") and chunk.usage:
+                    if (
+                        hasattr(chunk.usage, "num_sources_used")
+                        and chunk.usage.num_sources_used
+                    ):
+                        search_sources_used = chunk.usage.num_sources_used
+                        if enable_web_search and search_sources_used > 0:
+                            yield StreamChunk(
+                                type="content",
+                                content=f"\n📊 [Live Search] Using {search_sources_used} sources for real-time data\n",
+                            )
+
+            except Exception as chunk_error:
+                yield StreamChunk(
+                    type="error", error=f"Chunk processing error: {chunk_error}"
+                )
+                continue
+
+        yield StreamChunk(type="done")
+
+    def extract_tool_name(self, tool_call: Dict[str, Any]) -> str:
+        """Extract tool name from Chat Completions format."""
+        return tool_call.get("function", {}).get("name", "unknown")
+
+    def extract_tool_arguments(self, tool_call: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract tool arguments from Chat Completions format."""
+        return tool_call.get("function", {}).get("arguments", {})