massgen 0.0.3__py3-none-any.whl

massgen/backend/gemini.py

@@ -0,0 +1,652 @@
+"""
+Gemini backend implementation using structured output for voting and answer submission.
+
+APPROACH: Uses structured output instead of function declarations to handle the limitation
+where Gemini API cannot combine builtin tools with user-defined function declarations.
+
+KEY FEATURES:
+- ✅ Structured output for vote and new_answer mechanisms
+- ✅ Builtin tools support (code_execution + grounding)
+- ✅ Streaming with proper token usage tracking
+- ✅ Error handling and response parsing
+- ✅ Compatible with MassGen StreamChunk architecture
+
+TECHNICAL SOLUTION:
+- Uses Pydantic models to define structured output schemas
+- Prompts model to use specific JSON format for voting/answering
+- Converts structured responses to standard tool call format
+- Maintains compatibility with existing MassGen workflow
+"""
+
+import os
+import json
+import enum
+from typing import Dict, List, Any, AsyncGenerator, Optional
+from .base import LLMBackend, StreamChunk
+
+try:
+    from pydantic import BaseModel, Field
+except ImportError:
+    BaseModel = None
+    Field = None
+
+
+class VoteOption(enum.Enum):
+    """Vote options for agent selection."""
+
+    AGENT1 = "agent1"
+    AGENT2 = "agent2"
+    AGENT3 = "agent3"
+    AGENT4 = "agent4"
+    AGENT5 = "agent5"
+
+
+class ActionType(enum.Enum):
+    """Action types for structured output."""
+
+    VOTE = "vote"
+    NEW_ANSWER = "new_answer"
+
+
+class VoteAction(BaseModel):
+    """Structured output for voting action."""
+
+    action: ActionType = Field(default=ActionType.VOTE, description="Action type")
+    agent_id: str = Field(
+        description="Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')"
+    )
+    reason: str = Field(description="Brief reason why this agent has the best answer")
+
+
+class NewAnswerAction(BaseModel):
+    """Structured output for new answer action."""
+
+    action: ActionType = Field(default=ActionType.NEW_ANSWER, description="Action type")
+    content: str = Field(
+        description="Your improved answer. If any builtin tools like search or code execution were used, include how they are used here."
+    )
+
+
+class CoordinationResponse(BaseModel):
+    """Structured response for coordination actions."""
+
+    action_type: ActionType = Field(description="Type of action to take")
+    vote_data: Optional[VoteAction] = Field(
+        default=None, description="Vote data if action is vote"
+    )
+    answer_data: Optional[NewAnswerAction] = Field(
+        default=None, description="Answer data if action is new_answer"
+    )
+
+
+class GeminiBackend(LLMBackend):
+    """Google Gemini backend using structured output for coordination."""
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key, **kwargs)
+        self.api_key = (
+            api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
+        )
+        self.search_count = 0
+        self.code_execution_count = 0
+
+        if BaseModel is None:
+            raise ImportError(
+                "pydantic is required for Gemini backend. Install with: pip install pydantic"
+            )
+
+    def detect_coordination_tools(self, tools: List[Dict[str, Any]]) -> bool:
+        """Detect if tools contain vote/new_answer coordination tools."""
+        if not tools:
+            return False
+
+        tool_names = set()
+        for tool in tools:
+            if tool.get("type") == "function":
+                if "function" in tool:
+                    tool_names.add(tool["function"].get("name", ""))
+                elif "name" in tool:
+                    tool_names.add(tool.get("name", ""))
+
+        return "vote" in tool_names and "new_answer" in tool_names
+
+    def build_structured_output_prompt(
+        self, base_content: str, valid_agent_ids: Optional[List[str]] = None
+    ) -> str:
+        """Build prompt that encourages structured output for coordination."""
+        agent_list = ""
+        if valid_agent_ids:
+            agent_list = f"Valid agents: {', '.join(valid_agent_ids)}"
+
+        return f"""{base_content}
+
+IMPORTANT: You must respond with a structured JSON decision at the end of your response.
+
+If you want to VOTE for an existing agent's answer:
+{{
+  "action_type": "vote",
+  "vote_data": {{
+    "action": "vote",
+    "agent_id": "agent1",  // Choose from: {agent_list or 'agent1, agent2, agent3, etc.'}
+    "reason": "Brief reason for your vote"
+  }}
+}}
+
+If you want to provide a NEW ANSWER:
+{{
+  "action_type": "new_answer", 
+  "answer_data": {{
+    "action": "new_answer",
+    "content": "Your complete improved answer here"
+  }}
+}}
+
+Make your decision and include the JSON at the very end of your response."""
+
+    def extract_structured_response(
+        self, response_text: str
+    ) -> Optional[Dict[str, Any]]:
+        """Extract structured JSON response from model output."""
+        try:
+            import re
+
+            # Strategy 0: Look for JSON inside markdown code blocks first
+            markdown_json_pattern = r"```json\s*(\{.*?\})\s*```"
+            markdown_matches = re.findall(
+                markdown_json_pattern, response_text, re.DOTALL
+            )
+
+            for match in reversed(markdown_matches):
+                try:
+                    parsed = json.loads(match.strip())
+                    if isinstance(parsed, dict) and "action_type" in parsed:
+                        return parsed
+                except json.JSONDecodeError:
+                    continue
+
+            # Strategy 1: Look for complete JSON blocks with proper braces
+            json_pattern = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
+            json_matches = re.findall(json_pattern, response_text, re.DOTALL)
+
+            # Try parsing each match (in reverse order - last one first)
+            for match in reversed(json_matches):
+                try:
+                    cleaned_match = match.strip()
+                    parsed = json.loads(cleaned_match)
+                    if isinstance(parsed, dict) and "action_type" in parsed:
+                        return parsed
+                except json.JSONDecodeError:
+                    continue
+
+            # Strategy 2: Look for JSON blocks with nested braces (more complex)
+            brace_count = 0
+            json_start = -1
+
+            for i, char in enumerate(response_text):
+                if char == "{":
+                    if brace_count == 0:
+                        json_start = i
+                    brace_count += 1
+                elif char == "}":
+                    brace_count -= 1
+                    if brace_count == 0 and json_start >= 0:
+                        # Found a complete JSON block
+                        json_block = response_text[json_start : i + 1]
+                        try:
+                            parsed = json.loads(json_block)
+                            if isinstance(parsed, dict) and "action_type" in parsed:
+                                return parsed
+                        except json.JSONDecodeError:
+                            pass
+                        json_start = -1
+
+            # Strategy 3: Line-by-line approach (fallback)
+            lines = response_text.strip().split("\n")
+            json_candidates = []
+
+            for i, line in enumerate(lines):
+                stripped = line.strip()
+                if stripped.startswith("{") and stripped.endswith("}"):
+                    json_candidates.append(stripped)
+                elif stripped.startswith("{"):
+                    # Multi-line JSON - collect until closing brace
+                    json_text = stripped
+                    for j in range(i + 1, len(lines)):
+                        json_text += "\n" + lines[j].strip()
+                        if lines[j].strip().endswith("}"):
+                            json_candidates.append(json_text)
+                            break
+
+            # Try to parse each candidate
+            for candidate in reversed(json_candidates):
+                try:
+                    parsed = json.loads(candidate)
+                    if isinstance(parsed, dict) and "action_type" in parsed:
+                        return parsed
+                except json.JSONDecodeError:
+                    continue
+
+            return None
+
+        except Exception:
+            return None
+
+    def convert_structured_to_tool_calls(
+        self, structured_response: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """Convert structured response to tool call format."""
+        action_type = structured_response.get("action_type")
+
+        if action_type == "vote":
+            vote_data = structured_response.get("vote_data", {})
+            return [
+                {
+                    "id": f"vote_{hash(str(vote_data)) % 10000}",
+                    "type": "function",
+                    "function": {
+                        "name": "vote",
+                        "arguments": {
+                            "agent_id": vote_data.get("agent_id", ""),
+                            "reason": vote_data.get("reason", ""),
+                        },
+                    },
+                }
+            ]
+
+        elif action_type == "new_answer":
+            answer_data = structured_response.get("answer_data", {})
+            return [
+                {
+                    "id": f"new_answer_{hash(str(answer_data)) % 10000}",
+                    "type": "function",
+                    "function": {
+                        "name": "new_answer",
+                        "arguments": {"content": answer_data.get("content", "")},
+                    },
+                }
+            ]
+
+        return []
+
+    async def stream_with_tools(
+        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response using Gemini API with structured output for coordination."""
+        try:
+            from google import genai
+
+            # Extract parameters
+            model_name = kwargs.get("model", "gemini-2.5-flash")
+            temperature = kwargs.get("temperature", 0.1)
+            enable_web_search = kwargs.get("enable_web_search", False)
+            enable_code_execution = kwargs.get("enable_code_execution", False)
+
+            # Check if this is a coordination request
+            is_coordination = self.detect_coordination_tools(tools)
+            valid_agent_ids = None
+
+            if is_coordination:
+                # Extract valid agent IDs from vote tool enum if available
+                for tool in tools:
+                    if tool.get("type") == "function":
+                        func_def = tool.get("function", {})
+                        if func_def.get("name") == "vote":
+                            agent_id_param = (
+                                func_def.get("parameters", {})
+                                .get("properties", {})
+                                .get("agent_id", {})
+                            )
+                            if "enum" in agent_id_param:
+                                valid_agent_ids = agent_id_param["enum"]
+                            break
+
+            # Build content string from messages
+            conversation_content = ""
+            system_message = ""
+
+            for msg in messages:
+                if msg.get("role") == "system":
+                    system_message = msg.get("content", "")
+                elif msg.get("role") == "user":
+                    conversation_content += f"User: {msg.get('content', '')}\n"
+                elif msg.get("role") == "assistant":
+                    conversation_content += f"Assistant: {msg.get('content', '')}\n"
+
+            # For coordination requests, modify the prompt to use structured output
+            if is_coordination:
+                conversation_content = self.build_structured_output_prompt(
+                    conversation_content, valid_agent_ids
+                )
+
+            # Combine system message and conversation
+            full_content = ""
+            if system_message:
+                full_content += f"{system_message}\n\n"
+            full_content += conversation_content
+
+            # Use google-genai package
+            client = genai.Client(api_key=self.api_key)
+
+            # Setup builtin tools
+            builtin_tools = []
+            if enable_web_search:
+                try:
+                    from google.genai import types
+
+                    grounding_tool = types.Tool(google_search=types.GoogleSearch())
+                    builtin_tools.append(grounding_tool)
+                except ImportError:
+                    yield StreamChunk(
+                        type="content",
+                        content="\n⚠️  Web search requires google.genai.types\n",
+                    )
+
+            if enable_code_execution:
+                try:
+                    from google.genai import types
+
+                    code_tool = types.Tool(code_execution=types.ToolCodeExecution())
+                    builtin_tools.append(code_tool)
+                except ImportError:
+                    yield StreamChunk(
+                        type="content",
+                        content="\n⚠️  Code execution requires google.genai.types\n",
+                    )
+
+            config = {
+                "temperature": temperature,
+                "max_output_tokens": kwargs.get("max_tokens", 8192),
+            }
+
+            # Add builtin tools to config
+            if builtin_tools:
+                config["tools"] = builtin_tools
+
+            # For coordination requests, use JSON response format (may conflict with builtin tools)
+            if is_coordination and not builtin_tools:
+                config["response_mime_type"] = "application/json"
+                config["response_schema"] = CoordinationResponse.model_json_schema()
+            elif is_coordination and builtin_tools:
+                # Cannot use structured output with builtin tools - fallback to text parsing
+                pass
+
+            # Use streaming for real-time response
+            full_content_text = ""
+            final_response = None
+
+            for chunk in client.models.generate_content_stream(
+                model=model_name, contents=full_content, config=config
+            ):
+                if hasattr(chunk, "text") and chunk.text:
+                    chunk_text = chunk.text
+                    full_content_text += chunk_text
+                    yield StreamChunk(type="content", content=chunk_text)
+
+                # Keep track of the final response for tool processing
+                if hasattr(chunk, "candidates"):
+                    final_response = chunk
+
+                # Check for tools used in each chunk for real-time detection
+                if builtin_tools and hasattr(chunk, "candidates") and chunk.candidates:
+                    candidate = chunk.candidates[0]
+
+                    # Check for code execution in this chunk
+                    if (
+                        enable_code_execution
+                        and hasattr(candidate, "content")
+                        and hasattr(candidate.content, "parts")
+                    ):
+                        for part in candidate.content.parts:
+                            if (
+                                hasattr(part, "executable_code")
+                                and part.executable_code
+                            ):
+                                code_content = getattr(
+                                    part.executable_code,
+                                    "code",
+                                    str(part.executable_code),
+                                )
+                                yield StreamChunk(
+                                    type="content",
+                                    content=f"\n💻 [Code Executed]\n```python\n{code_content}\n```\n",
+                                )
+                            elif (
+                                hasattr(part, "code_execution_result")
+                                and part.code_execution_result
+                            ):
+                                result_content = getattr(
+                                    part.code_execution_result,
+                                    "output",
+                                    str(part.code_execution_result),
+                                )
+                                yield StreamChunk(
+                                    type="content",
+                                    content=f"📊 [Result] {result_content}\n",
+                                )
+
+            content = full_content_text
+
+            # Process coordination FIRST (before adding tool indicators that might confuse parsing)
+            tool_calls_detected = []
+            if is_coordination and content.strip():
+                # For structured output mode, the entire content is JSON
+                structured_response = None
+                # Try multiple parsing strategies
+                try:
+                    # Strategy 1: Parse entire content as JSON (works for both modes)
+                    structured_response = json.loads(content.strip())
+                except json.JSONDecodeError:
+                    # Strategy 2: Extract JSON from mixed text content (handles markdown-wrapped JSON)
+                    structured_response = self.extract_structured_response(content)
+
+                if (
+                    structured_response
+                    and isinstance(structured_response, dict)
+                    and "action_type" in structured_response
+                ):
+                    # Convert to tool calls
+                    tool_calls = self.convert_structured_to_tool_calls(
+                        structured_response
+                    )
+                    if tool_calls:
+                        tool_calls_detected = tool_calls
+
+            # Process builtin tool results if any tools were used
+            builtin_tool_results = []
+            if (
+                builtin_tools
+                and final_response
+                and hasattr(final_response, "candidates")
+                and final_response.candidates
+            ):
+                # Check for grounding or code execution results
+                candidate = final_response.candidates[0]
+
+                # Check for web search results - only show if actually used
+                if (
+                    hasattr(candidate, "grounding_metadata")
+                    and candidate.grounding_metadata
+                ):
+                    # Check if web search was actually used by looking for queries or chunks
+                    search_actually_used = False
+                    search_queries = []
+
+                    # Look for web search queries
+                    if (
+                        hasattr(candidate.grounding_metadata, "web_search_queries")
+                        and candidate.grounding_metadata.web_search_queries
+                    ):
+                        try:
+                            for (
+                                query
+                            ) in candidate.grounding_metadata.web_search_queries:
+                                if query and query.strip():
+                                    search_queries.append(query.strip())
+                                    search_actually_used = True
+                        except (TypeError, AttributeError):
+                            pass
+
+                    # Look for grounding chunks (indicates actual search results)
+                    if (
+                        hasattr(candidate.grounding_metadata, "grounding_chunks")
+                        and candidate.grounding_metadata.grounding_chunks
+                    ):
+                        try:
+                            if len(candidate.grounding_metadata.grounding_chunks) > 0:
+                                search_actually_used = True
+                        except (TypeError, AttributeError):
+                            pass
+
+                    # Only show indicators if search was actually used
+                    if search_actually_used:
+                        yield StreamChunk(
+                            type="content",
+                            content="🔍 [Builtin Tool: Web Search] Results integrated\n",
+                        )
+
+                        # Show search queries
+                        for query in search_queries:
+                            yield StreamChunk(
+                                type="content", content=f"🔍 [Search Query] '{query}'\n"
+                            )
+
+                        builtin_result = {
+                            "id": f"web_search_{hash(str(candidate.grounding_metadata)) % 10000}",
+                            "tool_type": "google_search_retrieval",
+                            "status": "completed",
+                            "metadata": str(candidate.grounding_metadata),
+                        }
+                        builtin_tool_results.append(builtin_result)
+                        self.search_count += 1
+
+                # Check for code execution in the response parts
+                if (
+                    enable_code_execution
+                    and hasattr(candidate, "content")
+                    and hasattr(candidate.content, "parts")
+                ):
+                    # Look for executable_code and code_execution_result parts
+                    code_parts = []
+                    for part in candidate.content.parts:
+                        if hasattr(part, "executable_code") and part.executable_code:
+                            code_content = getattr(
+                                part.executable_code, "code", str(part.executable_code)
+                            )
+                            code_parts.append(f"Code: {code_content}")
+                        elif (
+                            hasattr(part, "code_execution_result")
+                            and part.code_execution_result
+                        ):
+                            result_content = getattr(
+                                part.code_execution_result,
+                                "output",
+                                str(part.code_execution_result),
+                            )
+                            code_parts.append(f"Result: {result_content}")
+
+                    if code_parts:
+                        # Code execution was actually used
+                        yield StreamChunk(
+                            type="content",
+                            content="💻 [Builtin Tool: Code Execution] Code executed\n",
+                        )
+                        # Also show the actual code and result
+                        for part in code_parts:
+                            if part.startswith("Code: "):
+                                code_content = part[6:]  # Remove "Code: " prefix
+                                yield StreamChunk(
+                                    type="content",
+                                    content=f"💻 [Code Executed]\n```python\n{code_content}\n```\n",
+                                )
+                            elif part.startswith("Result: "):
+                                result_content = part[8:]  # Remove "Result: " prefix
+                                yield StreamChunk(
+                                    type="content",
+                                    content=f"📊 [Result] {result_content}\n",
+                                )
+
+                        builtin_result = {
+                            "id": f"code_execution_{hash(str(code_parts)) % 10000}",
+                            "tool_type": "code_execution",
+                            "status": "completed",
+                            "code_parts": code_parts,
+                            "output": "; ".join(code_parts),
+                        }
+                        builtin_tool_results.append(builtin_result)
+                        self.code_execution_count += 1
+
+            # Yield builtin tool results
+            if builtin_tool_results:
+                yield StreamChunk(
+                    type="builtin_tool_results",
+                    builtin_tool_results=builtin_tool_results,
+                )
+
+            # Yield coordination tool calls if detected
+            if tool_calls_detected:
+                yield StreamChunk(type="tool_calls", tool_calls=tool_calls_detected)
+
+            # Build complete message
+            complete_message = {"role": "assistant", "content": content.strip()}
+            if tool_calls_detected:
+                complete_message["tool_calls"] = tool_calls_detected
+
+            yield StreamChunk(
+                type="complete_message", complete_message=complete_message
+            )
+            yield StreamChunk(type="done")
+
+        except Exception as e:
+            yield StreamChunk(type="error", error=f"Gemini API error: {e}")
+
+    def get_provider_name(self) -> str:
+        """Get the provider name."""
+        return "Gemini"
+
+    def get_supported_builtin_tools(self) -> List[str]:
+        """Get list of builtin tools supported by Gemini."""
+        return ["google_search_retrieval", "code_execution"]
+
+    def estimate_tokens(self, text: str) -> int:
+        """Estimate token count for text (Gemini uses ~4 chars per token)."""
+        return len(text) // 4
+
+    def calculate_cost(
+        self, input_tokens: int, output_tokens: int, model: str
+    ) -> float:
+        """Calculate cost for Gemini token usage (2025 pricing)."""
+        model_lower = model.lower()
+
+        if "gemini-2.5-pro" in model_lower:
+            # Gemini 2.5 Pro pricing
+            input_cost = (input_tokens / 1_000_000) * 1.25
+            output_cost = (output_tokens / 1_000_000) * 5.0
+        elif "gemini-2.5-flash" in model_lower:
+            if "lite" in model_lower:
+                # Gemini 2.5 Flash-Lite pricing
+                input_cost = (input_tokens / 1_000_000) * 0.075
+                output_cost = (output_tokens / 1_000_000) * 0.30
+            else:
+                # Gemini 2.5 Flash pricing
+                input_cost = (input_tokens / 1_000_000) * 0.15
+                output_cost = (output_tokens / 1_000_000) * 0.60
+        else:
+            # Default fallback (assume Flash pricing)
+            input_cost = (input_tokens / 1_000_000) * 0.15
+            output_cost = (output_tokens / 1_000_000) * 0.60
+
+        # Add tool usage costs (estimates)
+        tool_costs = 0.0
+        if self.search_count > 0:
+            tool_costs += self.search_count * 0.01  # Estimated search cost
+
+        if self.code_execution_count > 0:
+            tool_costs += self.code_execution_count * 0.005  # Estimated execution cost
+
+        return input_cost + output_cost + tool_costs
+
+    def reset_tool_usage(self):
+        """Reset tool usage tracking."""
+        self.search_count = 0
+        self.code_execution_count = 0
+        super().reset_token_usage()