stravinsky 0.2.40__py3-none-any.whl → 0.3.4__py3-none-any.whl

mcp_bridge/tools/model_invoke.py

@@ -8,22 +8,86 @@ API requests to external model providers.
 import logging
 import os
 import time
+import uuid
+import base64
+import json as json_module
 
 logger = logging.getLogger(__name__)
 
+
+def _summarize_prompt(prompt: str, max_length: int = 120) -> str:
+    """
+    Generate a short summary of the prompt for logging.
+
+    Args:
+        prompt: The full prompt text
+        max_length: Maximum characters to include in summary
+
+    Returns:
+        Truncated prompt suitable for logging (single line, max_length chars)
+    """
+    if not prompt:
+        return "(empty prompt)"
+
+    # Normalize whitespace: collapse newlines and multiple spaces
+    clean = " ".join(prompt.split())
+
+    if len(clean) <= max_length:
+        return clean
+
+    return clean[:max_length] + "..."
+
+
+# Cache for Codex instructions (fetched from GitHub)
+_CODEX_INSTRUCTIONS_CACHE = {}
+_CODEX_INSTRUCTIONS_RELEASE_TAG = "rust-v0.77.0"  # Update as needed
+
+
+async def _fetch_codex_instructions(model: str = "gpt-5.2-codex") -> str:
+    """
+    Fetch official Codex instructions from GitHub.
+    Caches results to avoid repeated fetches.
+    """
+    import httpx
+
+    if model in _CODEX_INSTRUCTIONS_CACHE:
+        return _CODEX_INSTRUCTIONS_CACHE[model]
+
+    # Map model to prompt file
+    prompt_file_map = {
+        "gpt-5.2-codex": "gpt-5.2-codex_prompt.md",
+        "gpt-5.1-codex": "gpt_5_codex_prompt.md",
+        "gpt-5.1-codex-max": "gpt_5_codex_max_prompt.md",
+    }
+
+    prompt_file = prompt_file_map.get(model, "gpt-5.2-codex_prompt.md")
+    url = f"https://raw.githubusercontent.com/openai/codex/{_CODEX_INSTRUCTIONS_RELEASE_TAG}/codex-rs/core/{prompt_file}"
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, timeout=30.0)
+            response.raise_for_status()
+            instructions = response.text
+            _CODEX_INSTRUCTIONS_CACHE[model] = instructions
+            return instructions
+    except Exception as e:
+        logger.error(f"Failed to fetch Codex instructions: {e}")
+        # Return basic fallback instructions
+        return "You are Codex, based on GPT-5. You are running as a coding agent."
+
+
 # Model name mapping: user-friendly names -> Antigravity API model IDs
 # Per API spec: https://github.com/NoeFabris/opencode-antigravity-auth/blob/main/docs/ANTIGRAVITY_API_SPEC.md
-# VERIFIED GEMINI MODELS (as of 2025-12):
-#   - gemini-3-pro-high, gemini-3-pro-low
-# NOTE: There is NO gemini-3-flash in the API - all flash aliases map to gemini-3-pro-low
+# VERIFIED GEMINI MODELS (as of 2026-01):
+#   - gemini-3-flash, gemini-3-pro-high, gemini-3-pro-low
 # NOTE: Claude models should use Anthropic API directly, NOT Antigravity
 GEMINI_MODEL_MAP = {
     # Antigravity verified Gemini models (pass-through)
     "gemini-3-pro-low": "gemini-3-pro-low",
     "gemini-3-pro-high": "gemini-3-pro-high",
-    # Aliases for convenience (map to closest verified model)
-    "gemini-flash": "gemini-3-pro-low",
-    "gemini-3-flash": "gemini-3-pro-low",  # NOT a real model - redirect to pro-low
+    "gemini-3-flash": "gemini-3-flash",
+    # Aliases for convenience
+    "gemini-flash": "gemini-3-flash",
     "gemini-pro": "gemini-3-pro-low",
     "gemini-3-pro": "gemini-3-pro-low",
     "gemini": "gemini-3-pro-low",  # Default gemini alias
@@ -84,11 +148,11 @@ def _get_session_id(conversation_key: str | None = None) -> str:
     Returns:
         Stable session UUID for this conversation
     """
-    import uuid
+    import uuid as uuid_module  # Local import workaround
 
     key = conversation_key or "default"
     if key not in _SESSION_CACHE:
-        _SESSION_CACHE[key] = str(uuid.uuid4())
+        _SESSION_CACHE[key] = str(uuid_module.uuid4())
     return _SESSION_CACHE[key]
 
 
@@ -241,11 +305,13 @@ async def invoke_gemini(
     temperature: float = 0.7,
     max_tokens: int = 4096,
     thinking_budget: int = 0,
+    image_path: str | None = None,
 ) -> str:
     """
     Invoke a Gemini model with the given prompt.
 
     Uses OAuth authentication with Antigravity credentials.
+    Supports vision API for image/PDF analysis when image_path is provided.
 
     Args:
         token_store: Token store for OAuth credentials
@@ -253,6 +319,8 @@ async def invoke_gemini(
         model: Gemini model to use
         temperature: Sampling temperature (0.0-2.0)
         max_tokens: Maximum tokens in response
+        thinking_budget: Tokens reserved for internal reasoning
+        image_path: Optional path to image/PDF for vision analysis (token optimization)
 
     Returns:
         The model's response text.
@@ -261,6 +329,7 @@ async def invoke_gemini(
         ValueError: If not authenticated with Gemini
         httpx.HTTPStatusError: If API request fails
     """
+    logger.info(f"[DEBUG] invoke_gemini called, uuid module check: {uuid}")
     # Execute pre-model invoke hooks
     params = {
         "prompt": prompt,
@@ -268,6 +337,8 @@ async def invoke_gemini(
         "temperature": temperature,
         "max_tokens": max_tokens,
         "thinking_budget": thinking_budget,
+        "token_store": token_store,  # Pass for hooks that need model access
+        "provider": "gemini",  # Identify which provider is being called
     }
     hook_manager = get_hook_manager()
     params = await hook_manager.execute_pre_model_invoke(params)
@@ -279,6 +350,22 @@ async def invoke_gemini(
     max_tokens = params["max_tokens"]
     thinking_budget = params["thinking_budget"]
 
+    # Extract agent context for logging (may be passed via params or original call)
+    agent_context = params.get("agent_context", {})
+    agent_type = agent_context.get("agent_type", "direct")
+    task_id = agent_context.get("task_id", "")
+    description = agent_context.get("description", "")
+    prompt_summary = _summarize_prompt(prompt)
+
+    # Log with agent context and prompt summary
+    logger.info(f"[{agent_type}] → {model}: {prompt_summary}")
+
+    # USER-VISIBLE NOTIFICATION (stderr) - Shows when Gemini is invoked
+    import sys
+    task_info = f" task={task_id}" if task_id else ""
+    desc_info = f" | {description}" if description else ""
+    print(f"🔮 GEMINI: {model} | agent={agent_type}{task_info}{desc_info}", file=sys.stderr)
+
     access_token = await _ensure_valid_token(token_store, "gemini")
 
     # Resolve user-friendly model name to actual API model ID
@@ -296,8 +383,43 @@ async def invoke_gemini(
 
     # Build inner request payload
     # Per API spec: contents must include role ("user" or "model")
+
+    # Build parts list - text prompt plus optional image
+    parts = [{"text": prompt}]
+
+    # Add image data for vision analysis (token optimization for multimodal)
+    if image_path:
+        import base64
+        from pathlib import Path
+
+        image_file = Path(image_path)
+        if image_file.exists():
+            # Determine MIME type
+            suffix = image_file.suffix.lower()
+            mime_types = {
+                ".png": "image/png",
+                ".jpg": "image/jpeg",
+                ".jpeg": "image/jpeg",
+                ".gif": "image/gif",
+                ".webp": "image/webp",
+                ".pdf": "application/pdf",
+            }
+            mime_type = mime_types.get(suffix, "image/png")
+
+            # Read and base64 encode
+            image_data = base64.b64encode(image_file.read_bytes()).decode("utf-8")
+
+            # Add inline image data for Gemini Vision API
+            parts.append({
+                "inlineData": {
+                    "mimeType": mime_type,
+                    "data": image_data,
+                }
+            })
+            logger.info(f"[multimodal] Added vision data: {image_path} ({mime_type})")
+
     inner_payload = {
-        "contents": [{"role": "user", "parts": [{"text": prompt}]}],
+        "contents": [{"role": "user", "parts": parts}],
         "generationConfig": {
             "temperature": temperature,
             "maxOutputTokens": max_tokens,
@@ -315,11 +437,19 @@ async def invoke_gemini(
         }
 
     # Wrap request body per reference implementation
+    try:
+        import uuid as uuid_module  # Local import workaround for MCP context issue
+
+        request_id = f"invoke-{uuid_module.uuid4()}"
+    except Exception as e:
+        logger.error(f"UUID IMPORT FAILED: {e}")
+        raise RuntimeError(f"CUSTOM ERROR: UUID import failed: {e}")
+
     wrapped_payload = {
         "project": project_id,
         "model": api_model,
         "userAgent": "antigravity",
-        "requestId": f"invoke-{uuid.uuid4()}",
+        "requestId": request_id,
         "request": inner_payload,
     }
 
@@ -385,6 +515,26 @@ async def invoke_gemini(
         break
 
     if response is None:
+        # FALLBACK: Try Claude sonnet-4.5 for agents that support it
+        agent_context = params.get("agent_context", {})
+        agent_type = agent_context.get("agent_type", "unknown")
+
+        if agent_type in ("dewey", "explore", "document_writer", "multimodal"):
+            logger.warning(f"[{agent_type}] Gemini failed, falling back to Claude sonnet-4.5")
+            try:
+                import subprocess
+                fallback_result = subprocess.run(
+                    ["claude", "-p", prompt, "--model", "sonnet", "--output-format", "text"],
+                    capture_output=True,
+                    text=True,
+                    timeout=120,
+                    cwd=os.getcwd(),
+                )
+                if fallback_result.returncode == 0 and fallback_result.stdout.strip():
+                    return fallback_result.stdout.strip()
+            except Exception as fallback_error:
+                logger.error(f"Fallback to Claude also failed: {fallback_error}")
+
         raise ValueError(f"All Antigravity endpoints failed: {last_error}")
 
     response.raise_for_status()
@@ -537,8 +687,6 @@ async def invoke_gemini_agentic(
     Returns:
         Final text response from the model
     """
-    import uuid
-
     access_token = await _ensure_valid_token(token_store, "gemini")
     api_model = resolve_gemini_model(model)
 
@@ -574,11 +722,13 @@ async def invoke_gemini_agentic(
 
         # Wrap request body per reference implementation
         # From request.ts wrapRequestBody()
+        import uuid as uuid_module  # Local import workaround
+
         wrapped_payload = {
             "project": project_id,
             "model": api_model,
             "userAgent": "antigravity",
-            "requestId": f"agent-{uuid.uuid4()}",
+            "requestId": f"agent-{uuid_module.uuid4()}",
             "request": inner_payload,
         }
 
@@ -688,7 +838,7 @@ async def invoke_gemini_agentic(
 async def invoke_openai(
     token_store: TokenStore,
     prompt: str,
-    model: str = "gpt-5.2",
+    model: str = "gpt-5.2-codex",
     temperature: float = 0.7,
     max_tokens: int = 4096,
     thinking_budget: int = 0,
@@ -717,6 +867,8 @@ async def invoke_openai(
         "temperature": temperature,
         "max_tokens": max_tokens,
         "thinking_budget": thinking_budget,
+        "token_store": token_store,  # Pass for hooks that need model access
+        "provider": "openai",  # Identify which provider is being called
     }
     hook_manager = get_hook_manager()
     params = await hook_manager.execute_pre_model_invoke(params)
@@ -728,52 +880,123 @@ async def invoke_openai(
     max_tokens = params["max_tokens"]
     thinking_budget = params["thinking_budget"]
 
+    # Extract agent context for logging (may be passed via params or original call)
+    agent_context = params.get("agent_context", {})
+    agent_type = agent_context.get("agent_type", "direct")
+    task_id = agent_context.get("task_id", "")
+    description = agent_context.get("description", "")
+    prompt_summary = _summarize_prompt(prompt)
+
+    # Log with agent context and prompt summary
+    logger.info(f"[{agent_type}] → {model}: {prompt_summary}")
+
+    # USER-VISIBLE NOTIFICATION (stderr) - Shows when OpenAI is invoked
+    import sys
+    task_info = f" task={task_id}" if task_id else ""
+    desc_info = f" | {description}" if description else ""
+    print(f"🧠 OPENAI: {model} | agent={agent_type}{task_info}{desc_info}", file=sys.stderr)
+
     access_token = await _ensure_valid_token(token_store, "openai")
+    logger.info(f"[invoke_openai] Got access token")
+
+    # ChatGPT Backend API - Uses Codex Responses endpoint
+    # Replicates opencode-openai-codex-auth plugin behavior
+    api_url = "https://chatgpt.com/backend-api/codex/responses"
 
-    # OpenAI Chat Completions API
-    api_url = "https://api.openai.com/v1/chat/completions"
+    # Extract account ID from JWT token
+    logger.info(f"[invoke_openai] Extracting account ID from JWT")
+    try:
+        parts = access_token.split(".")
+        payload_b64 = parts[1]
+        padding = 4 - len(payload_b64) % 4
+        if padding != 4:
+            payload_b64 += "=" * padding
+        jwt_payload = json_module.loads(base64.urlsafe_b64decode(payload_b64))
+        account_id = jwt_payload.get("https://api.openai.com/auth", {}).get("chatgpt_account_id")
+    except Exception as e:
+        logger.error(f"Failed to extract account ID from JWT: {e}")
+        account_id = None
 
+    # Fetch official Codex instructions from GitHub
+    instructions = await _fetch_codex_instructions(model)
+
+    # Headers matching opencode-openai-codex-auth plugin
     headers = {
         "Authorization": f"Bearer {access_token}",
         "Content-Type": "application/json",
+        "Accept": "text/event-stream",  # SSE stream
+        "openai-beta": "responses=experimental",
+        "openai-originator": "codex_cli_rs",
     }
 
+    if account_id:
+        headers["x-openai-account-id"] = account_id
+
+    # Request body matching opencode transformation
     payload = {
         "model": model,
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": temperature,
+        "store": False,  # Required by ChatGPT backend
+        "stream": True,  # Always stream (handler converts to non-stream if needed)
+        "instructions": instructions,
+        "input": [{"role": "user", "content": prompt}],
+        "reasoning": {"effort": "high" if thinking_budget > 0 else "medium", "summary": "auto"},
+        "text": {"verbosity": "medium"},
+        "include": ["reasoning.encrypted_content"],
     }
 
-    # Handle thinking budget for O1/O3 style models (GPT-5.2)
-    if thinking_budget > 0:
-        payload["max_completion_tokens"] = max_tokens + thinking_budget
-        # For O1, temperature must be 1.0 or omitted usually, but we'll try to pass it
-    else:
-        payload["max_tokens"] = max_tokens
-
-    async with httpx.AsyncClient() as client:
-        response = await client.post(
-            api_url,
-            headers=headers,
-            json=payload,
-            timeout=120.0,
-        )
+    # Stream the response and collect text
+    text_chunks = []
 
-        if response.status_code == 401:
-            raise ValueError(
-                "OpenAI authentication failed. Run: python -m mcp_bridge.auth.cli login openai"
-            )
+    logger.info(f"[invoke_openai] Calling {api_url} with model {model}")
+    logger.info(f"[invoke_openai] Payload keys: {list(payload.keys())}")
+    logger.info(f"[invoke_openai] Instructions length: {len(instructions)}")
 
-        response.raise_for_status()
-
-        data = response.json()
+    try:
+        async with httpx.AsyncClient() as client:
+            async with client.stream(
+                "POST", api_url, headers=headers, json=payload, timeout=120.0
+            ) as response:
+                logger.info(f"[invoke_openai] Response status: {response.status_code}")
+                if response.status_code == 401:
+                    raise ValueError(
+                        "OpenAI authentication failed. Run: stravinsky-auth login openai"
+                    )
 
-        # Extract text from response
-        try:
-            choices = data.get("choices", [])
-            if choices:
-                message = choices[0].get("message", {})
-                return message.get("content", "")
+                if response.status_code >= 400:
+                    error_body = await response.aread()
+                    error_text = error_body.decode("utf-8")
+                    logger.error(f"OpenAI API error {response.status_code}: {error_text}")
+                    logger.error(f"Request payload was: {payload}")
+                    logger.error(f"Request headers were: {headers}")
+                    raise ValueError(f"OpenAI API error {response.status_code}: {error_text}")
+
+                # Parse SSE stream for text deltas
+                async for line in response.aiter_lines():
+                    if line.startswith("data: "):
+                        data_json = line[6:]  # Remove "data: " prefix
+                        try:
+                            data = json_module.loads(data_json)
+                            event_type = data.get("type")
+
+                            # Extract text deltas from SSE stream
+                            if event_type == "response.output_text.delta":
+                                delta = data.get("delta", "")
+                                text_chunks.append(delta)
+
+                        except json_module.JSONDecodeError:
+                            pass  # Skip malformed JSON
+                        except Exception as e:
+                            logger.warning(f"Error processing SSE event: {e}")
+
+        # Return collected text
+        result = "".join(text_chunks)
+        if not result:
             return "No response generated"
-        except (KeyError, IndexError) as e:
-            return f"Error parsing response: {e}"
+        return result
+
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error in invoke_openai: {e}")
+        raise ValueError(f"Failed to invoke OpenAI: {e}")