massgen 0.0.3__py3-none-any.whl

massgen/backend/grok.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+"""
+Grok/xAI backend implementation using OpenAI-compatible API.
+Clean implementation with only Grok-specific features.
+
+✅ TESTED: Backend works correctly with architecture
+- ✅ Grok API integration working
+- ✅ Tool message conversion compatible with Chat Completions format
+- ✅ Streaming functionality working correctly  
+- ✅ SingleAgent integration working
+- ✅ Error handling and pricing calculations implemented
+
+TODO for future releases:
+- Test multi-agent orchestrator integration
+- Test web search capabilities with tools
+- Validate advanced Grok-specific features
+"""
+
+import os
+from typing import Dict, List, Any, AsyncGenerator, Optional
+from .chat_completions import ChatCompletionsBackend
+from .base import StreamChunk
+
+
+class GrokBackend(ChatCompletionsBackend):
+    """Grok backend using xAI's OpenAI-compatible API."""
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key, **kwargs)
+        self.api_key = api_key or os.getenv("XAI_API_KEY")
+        self.base_url = "https://api.x.ai/v1"
+
+    async def stream_with_tools(
+        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response using xAI's OpenAI-compatible API."""
+
+        # Convert messages for Grok API compatibility
+        grok_messages = self._convert_messages_for_grok(messages)
+
+        try:
+            import openai
+
+            # Use OpenAI client with xAI base URL
+            client = openai.AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
+
+            # Extract parameters
+            model = kwargs.get("model", "grok-3-mini")
+            max_tokens = kwargs.get("max_tokens", None)
+            temperature = kwargs.get("temperature", None)
+            enable_web_search = kwargs.get("enable_web_search", False)
+
+            # Convert tools to Chat Completions format
+            converted_tools = (
+                self.convert_tools_to_chat_completions_format(tools) if tools else None
+            )
+
+            # Chat Completions API parameters
+            api_params = {
+                "model": model,
+                "messages": grok_messages,
+                "tools": converted_tools,
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "stream": True,
+            }
+
+            # Add Live Search parameters if enabled (Grok-specific)
+            if enable_web_search:
+                search_params_kwargs = {"mode": "auto", "return_citations": True}
+
+                # Allow override of search parameters from backend params
+                max_results = kwargs.get("max_search_results")
+                if max_results is not None:
+                    search_params_kwargs["max_search_results"] = max_results
+
+                search_mode = kwargs.get("search_mode")
+                if search_mode is not None:
+                    search_params_kwargs["mode"] = search_mode
+
+                return_citations = kwargs.get("return_citations")
+                if return_citations is not None:
+                    search_params_kwargs["return_citations"] = return_citations
+
+                # Use extra_body to pass search_parameters to xAI API
+                api_params["extra_body"] = {"search_parameters": search_params_kwargs}
+
+            # Create stream
+            stream = await client.chat.completions.create(**api_params)
+
+            # Use base class streaming handler
+            async for chunk in self.handle_chat_completions_stream(
+                stream, enable_web_search
+            ):
+                yield chunk
+
+        except Exception as e:
+            yield StreamChunk(type="error", error=f"Grok API error: {e}")
+
+    def get_provider_name(self) -> str:
+        """Get the name of this provider."""
+        return "Grok"
+
+    def get_supported_builtin_tools(self) -> List[str]:
+        """Get list of builtin tools supported by Grok."""
+        return ["web_search"]
+
+    def estimate_tokens(self, text: str) -> int:
+        """Estimate token count for text (rough approximation)."""
+        return int(len(text.split()) * 1.3)
+
+    def calculate_cost(
+        self, input_tokens: int, output_tokens: int, model: str
+    ) -> float:
+        """Calculate cost for token usage."""
+        model_lower = model.lower()
+
+        # Handle -mini models with lower costs
+        if "grok-2" in model_lower:
+            if "mini" in model_lower:
+                input_cost = (input_tokens / 1_000_000) * 1.0  # Lower cost for mini
+                output_cost = (output_tokens / 1_000_000) * 5.0
+            else:
+                input_cost = (input_tokens / 1_000_000) * 2.0
+                output_cost = (output_tokens / 1_000_000) * 10.0
+        elif "grok-3" in model_lower:
+            if "mini" in model_lower:
+                input_cost = (input_tokens / 1_000_000) * 2.5  # Lower cost for mini
+                output_cost = (output_tokens / 1_000_000) * 7.5
+            else:
+                input_cost = (input_tokens / 1_000_000) * 5.0
+                output_cost = (output_tokens / 1_000_000) * 15.0
+        elif "grok-4" in model_lower:
+            if "mini" in model_lower:
+                input_cost = (input_tokens / 1_000_000) * 4.0  # Lower cost for mini
+                output_cost = (output_tokens / 1_000_000) * 10.0
+            else:
+                input_cost = (input_tokens / 1_000_000) * 8.0
+                output_cost = (output_tokens / 1_000_000) * 20.0
+        else:
+            # Default fallback (assume grok-3 pricing)
+            input_cost = (input_tokens / 1_000_000) * 5.0
+            output_cost = (output_tokens / 1_000_000) * 15.0
+
+        return input_cost + output_cost
+
+    def _convert_messages_for_grok(
+        self, messages: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """
+        Convert messages for Grok API compatibility.
+
+        Grok expects tool call arguments as JSON strings in conversation history,
+        but returns them as objects in responses.
+        """
+        import json
+
+        converted_messages = []
+
+        for message in messages:
+            # Create a copy to avoid modifying the original
+            converted_msg = dict(message)
+
+            # Convert tool_calls arguments from objects to JSON strings
+            if message.get("role") == "assistant" and "tool_calls" in message:
+                converted_tool_calls = []
+                for tool_call in message["tool_calls"]:
+                    converted_call = dict(tool_call)
+                    if "function" in converted_call:
+                        converted_function = dict(converted_call["function"])
+                        arguments = converted_function.get("arguments")
+
+                        # Convert arguments to JSON string if it's an object
+                        if isinstance(arguments, dict):
+                            converted_function["arguments"] = json.dumps(arguments)
+                        elif arguments is None:
+                            converted_function["arguments"] = "{}"
+                        # If it's already a string, keep it as-is
+
+                        converted_call["function"] = converted_function
+                    converted_tool_calls.append(converted_call)
+                converted_msg["tool_calls"] = converted_tool_calls
+
+            converted_messages.append(converted_msg)
+
+        return converted_messages

massgen/backend/response.py

@@ -0,0 +1,397 @@
+from __future__ import annotations
+
+"""
+Response API backend implementation.
+Standalone implementation optimized for the standard Response API format (originated by OpenAI).
+"""
+
+import os
+from typing import Dict, List, Any, AsyncGenerator, Optional
+from .base import LLMBackend, StreamChunk
+
+
+class ResponseBackend(LLMBackend):
+    """Backend using the standard Response API format."""
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key, **kwargs)
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+
+    def convert_tools_to_response_api_format(
+        self, tools: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Convert tools from Chat Completions format to Response API format if needed.
+
+        Chat Completions format: {"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}
+        Response API format: {"type": "function", "name": ..., "description": ..., "parameters": ...}
+        """
+        if not tools:
+            return tools
+
+        converted_tools = []
+        for tool in tools:
+            if tool.get("type") == "function" and "function" in tool:
+                # Chat Completions format - convert to Response API format
+                func = tool["function"]
+                converted_tools.append(
+                    {
+                        "type": "function",
+                        "name": func["name"],
+                        "description": func["description"],
+                        "parameters": func.get("parameters", {}),
+                    }
+                )
+            else:
+                # Already in Response API format or non-function tool
+                converted_tools.append(tool)
+
+        return converted_tools
+
+    def convert_messages_to_response_api_format(
+        self, messages: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Convert messages from Chat Completions format to Response API format.
+
+        Chat Completions tool message: {"role": "tool", "tool_call_id": "...", "content": "..."}
+        Response API tool message: {"type": "function_call_output", "call_id": "...", "output": "..."}
+
+        Note: Assistant messages with tool_calls should not be in input - they're generated by the backend.
+        """
+        converted_messages = []
+
+        for message in messages:
+            if message.get("role") == "tool":
+                # Convert Chat Completions tool message to Response API format
+                converted_messages.append(
+                    {
+                        "type": "function_call_output",
+                        "call_id": message.get("tool_call_id"),
+                        "output": message.get("content", ""),
+                    }
+                )
+            elif message.get("type") == "function_call_output":
+                # Already in Response API format - keep as-is
+                converted_messages.append(message)
+            elif message.get("role") == "assistant" and "tool_calls" in message:
+                # Assistant message with tool_calls in native Responses API format
+                # Remove tool_calls when sending as input - only results should be sent back
+                cleaned_message = {
+                    k: v for k, v in message.items() if k != "tool_calls"
+                }
+                converted_messages.append(cleaned_message)
+            else:
+                # Keep other message types as-is
+                converted_messages.append(message)
+
+        return converted_messages
+
+    async def stream_with_tools(
+        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response using OpenAI Response API."""
+        try:
+            import openai
+
+            client = openai.AsyncOpenAI(api_key=self.api_key)
+
+            # Extract model and provider tool settings
+            model = kwargs.get("model", "gpt-4o-mini")
+            enable_web_search = kwargs.get("enable_web_search", False)
+            enable_code_interpreter = kwargs.get("enable_code_interpreter", False)
+
+            # Convert messages to Response API format (handles tool messages)
+            converted_messages = self.convert_messages_to_response_api_format(messages)
+
+            # Response API parameters (uses 'input', not 'messages')
+            api_params = {"model": model, "input": converted_messages, "stream": True}
+
+            # Add max_output_tokens if specified (o-series models don't support this)
+            max_tokens = kwargs.get("max_tokens")
+            if max_tokens and not model.startswith("o"):
+                api_params["max_output_tokens"] = max_tokens
+
+            # Add framework tools (convert to Response API format)
+            if tools:
+                converted_tools = self.convert_tools_to_response_api_format(tools)
+                api_params["tools"] = converted_tools
+
+            # Add provider tools (web search, code interpreter) if enabled
+            provider_tools = []
+            if enable_web_search:
+                provider_tools.append({"type": "web_search"})
+
+            if enable_code_interpreter:
+                provider_tools.append(
+                    {"type": "code_interpreter", "container": {"type": "auto"}}
+                )
+
+            if provider_tools:
+                if "tools" not in api_params:
+                    api_params["tools"] = []
+                api_params["tools"].extend(provider_tools)
+
+            stream = await client.responses.create(**api_params)
+
+            content = ""
+
+            async for chunk in stream:
+                # Handle Responses API streaming format
+                if hasattr(chunk, "type"):
+                    if chunk.type == "response.output_text.delta" and hasattr(
+                        chunk, "delta"
+                    ):
+                        content += chunk.delta
+                        yield StreamChunk(type="content", content=chunk.delta)
+                    elif chunk.type == "response.web_search_call.in_progress":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"\n🔍 [Provider Tool: Web Search] Starting search...",
+                        )
+                    elif chunk.type == "response.web_search_call.searching":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"🔍 [Provider Tool: Web Search] Searching...",
+                        )
+                    elif chunk.type == "response.web_search_call.completed":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"✅ [Provider Tool: Web Search] Search completed",
+                        )
+                    elif chunk.type == "response.code_interpreter_call.in_progress":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"\n💻 [Provider Tool: Code Interpreter] Starting execution...",
+                        )
+                    elif chunk.type == "response.code_interpreter_call.executing":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"💻 [Provider Tool: Code Interpreter] Executing...",
+                        )
+                    elif chunk.type == "response.code_interpreter_call.completed":
+                        yield StreamChunk(
+                            type="content",
+                            content=f"✅ [Provider Tool: Code Interpreter] Execution completed",
+                        )
+                    elif chunk.type == "response.output_item.done":
+                        # Get search query or executed code details - show them right after completion
+                        if hasattr(chunk, "item") and chunk.item:
+                            if (
+                                hasattr(chunk.item, "type")
+                                and chunk.item.type == "web_search_call"
+                            ):
+                                if hasattr(chunk.item, "action") and hasattr(
+                                    chunk.item.action, "query"
+                                ):
+                                    search_query = chunk.item.action.query
+                                    if search_query:
+                                        yield StreamChunk(
+                                            type="content",
+                                            content=f"🔍 [Search Query] '{search_query}'",
+                                        )
+                            elif (
+                                hasattr(chunk.item, "type")
+                                and chunk.item.type == "code_interpreter_call"
+                            ):
+                                if hasattr(chunk.item, "code") and chunk.item.code:
+                                    # Format code as a proper code block - don't assume language
+                                    yield StreamChunk(
+                                        type="content",
+                                        content=f"💻 [Code Executed]\n```\n{chunk.item.code}\n```\n",
+                                    )
+
+                                # Also show the execution output if available
+                                if (
+                                    hasattr(chunk.item, "outputs")
+                                    and chunk.item.outputs
+                                ):
+                                    for output in chunk.item.outputs:
+                                        output_text = None
+                                        if hasattr(output, "text") and output.text:
+                                            output_text = output.text
+                                        elif (
+                                            hasattr(output, "content")
+                                            and output.content
+                                        ):
+                                            output_text = output.content
+                                        elif hasattr(output, "data") and output.data:
+                                            output_text = str(output.data)
+                                        elif isinstance(output, str):
+                                            output_text = output
+                                        elif isinstance(output, dict):
+                                            # Handle dict format outputs
+                                            if "text" in output:
+                                                output_text = output["text"]
+                                            elif "content" in output:
+                                                output_text = output["content"]
+                                            elif "data" in output:
+                                                output_text = str(output["data"])
+
+                                        if output_text and output_text.strip():
+                                            yield StreamChunk(
+                                                type="content",
+                                                content=f"📊 [Result] {output_text.strip()}\n",
+                                            )
+                    elif chunk.type == "response.completed":
+                        # Extract and yield tool calls from the complete response
+                        if hasattr(chunk, "response"):
+                            response_dict = self._convert_to_dict(chunk.response)
+
+                            # Extract builtin tool results from output array
+                            builtin_tool_results = []
+                            if (
+                                isinstance(response_dict, dict)
+                                and "output" in response_dict
+                            ):
+                                for item in response_dict["output"]:
+                                    if item.get("type") == "code_interpreter_call":
+                                        # Code execution result
+                                        builtin_tool_results.append(
+                                            {
+                                                "id": item.get("id", ""),
+                                                "tool_type": "code_interpreter",
+                                                "status": item.get("status"),
+                                                "code": item.get("code", ""),
+                                                "outputs": item.get("outputs"),
+                                                "container_id": item.get(
+                                                    "container_id"
+                                                ),
+                                            }
+                                        )
+                                    elif item.get("type") == "web_search_call":
+                                        # Web search result
+                                        builtin_tool_results.append(
+                                            {
+                                                "id": item.get("id", ""),
+                                                "tool_type": "web_search",
+                                                "status": item.get("status"),
+                                                "query": item.get("query", ""),
+                                                "results": item.get("results"),
+                                            }
+                                        )
+
+                            # Yield builtin tool results if any were found
+                            if builtin_tool_results:
+                                yield StreamChunk(
+                                    type="builtin_tool_results",
+                                    builtin_tool_results=builtin_tool_results,
+                                )
+
+                            # Yield the complete response for internal use
+                            yield StreamChunk(
+                                type="complete_response", response=response_dict
+                            )
+                        else:
+                            # Fallback if no response object
+                            complete_message = {
+                                "role": "assistant",
+                                "content": content.strip(),
+                            }
+                            yield StreamChunk(
+                                type="complete_message",
+                                complete_message=complete_message,
+                            )
+
+                        # Signal completion
+                        yield StreamChunk(type="done")
+
+        except Exception as e:
+            yield StreamChunk(type="error", error=str(e))
+
+    def get_provider_name(self) -> str:
+        """Get the provider name."""
+        return "OpenAI"
+
+    def get_supported_builtin_tools(self) -> List[str]:
+        """Get list of builtin tools supported by OpenAI."""
+        return ["web_search", "code_interpreter"]
+
+    def extract_tool_name(self, tool_call: Dict[str, Any]) -> str:
+        """Extract tool name from OpenAI format (handles both Chat Completions and Responses API)."""
+        # Check if it's Chat Completions format
+        if "function" in tool_call:
+            return tool_call.get("function", {}).get("name", "unknown")
+        # Otherwise assume Responses API format
+        return tool_call.get("name", "unknown")
+
+    def extract_tool_arguments(self, tool_call: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract tool arguments from OpenAI format (handles both Chat Completions and Responses API)."""
+        # Check if it's Chat Completions format
+        if "function" in tool_call:
+            return tool_call.get("function", {}).get("arguments", {})
+        # Otherwise assume Responses API format
+        arguments = tool_call.get("arguments", {})
+        if isinstance(arguments, str):
+            try:
+                import json
+
+                return json.loads(arguments)
+            except:
+                return {}
+        return arguments
+
+    def extract_tool_call_id(self, tool_call: Dict[str, Any]) -> str:
+        """Extract tool call ID from OpenAI format (handles both Chat Completions and Responses API)."""
+        # For Responses API, use call_id (for tool results), for Chat Completions use id
+        return tool_call.get("call_id") or tool_call.get("id") or ""
+
+    def create_tool_result_message(
+        self, tool_call: Dict[str, Any], result_content: str
+    ) -> Dict[str, Any]:
+        """Create tool result message for OpenAI Responses API format."""
+        tool_call_id = self.extract_tool_call_id(tool_call)
+        # Use Responses API format directly - no conversion needed
+        return {
+            "type": "function_call_output",
+            "call_id": tool_call_id,
+            "output": result_content,
+        }
+
+    def extract_tool_result_content(self, tool_result_message: Dict[str, Any]) -> str:
+        """Extract content from OpenAI Responses API tool result message."""
+        return tool_result_message.get("output", "")
+
+    def _convert_to_dict(self, obj) -> Dict[str, Any]:
+        """Convert any object to dictionary with multiple fallback methods."""
+        try:
+            if hasattr(obj, "model_dump"):
+                return obj.model_dump()
+            elif hasattr(obj, "dict"):
+                return obj.dict()
+            else:
+                return dict(obj)
+        except:
+            # Final fallback: extract key attributes manually
+            return {
+                key: getattr(obj, key, None)
+                for key in dir(obj)
+                if not key.startswith("_") and not callable(getattr(obj, key, None))
+            }
+
+    def estimate_tokens(self, text: str) -> int:
+        """Estimate token count for text (rough approximation)."""
+        return len(text) // 4
+
+    def calculate_cost(
+        self, input_tokens: int, output_tokens: int, model: str
+    ) -> float:
+        """Calculate cost for OpenAI token usage (2024-2025 pricing)."""
+        model_lower = model.lower()
+
+        if "gpt-4" in model_lower:
+            if "4o-mini" in model_lower:
+                input_cost = input_tokens * 0.00015 / 1000
+                output_cost = output_tokens * 0.0006 / 1000
+            elif "4o" in model_lower:
+                input_cost = input_tokens * 0.005 / 1000
+                output_cost = output_tokens * 0.020 / 1000
+            else:
+                input_cost = input_tokens * 0.03 / 1000
+                output_cost = output_tokens * 0.06 / 1000
+        elif "gpt-3.5" in model_lower:
+            input_cost = input_tokens * 0.0005 / 1000
+            output_cost = output_tokens * 0.0015 / 1000
+        else:
+            input_cost = input_tokens * 0.0005 / 1000
+            output_cost = output_tokens * 0.0015 / 1000
+
+        return input_cost + output_cost