gobby 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

gobby/llm/gemini.py

@@ -1,29 +1,31 @@
 """
-Gemini implementation of LLMProvider.
+Gemini implementation of LLMProvider using LiteLLM.
 
-Supports two authentication modes:
-- api_key: Use GEMINI_API_KEY environment variable (BYOK)
-- adc: Use Google Application Default Credentials (subscription-based via gcloud auth)
+Routes all calls through LiteLLM for unified cost tracking:
+- api_key mode: Uses gemini/model-name prefix
+- adc mode: Uses vertex_ai/model-name prefix (requires VERTEXAI_PROJECT, VERTEXAI_LOCATION)
+
+This provider replaces direct google-generativeai SDK usage with LiteLLM routing.
 """
 
 import json
 import logging
-import os
 from typing import Any, Literal
 
 from gobby.config.app import DaemonConfig
 from gobby.llm.base import AuthMode, LLMProvider
+from gobby.llm.litellm_executor import get_litellm_model, setup_provider_env
 
 logger = logging.getLogger(__name__)
 
 
 class GeminiProvider(LLMProvider):
     """
-    Gemini implementation of LLMProvider using google-generativeai package.
+    Gemini implementation of LLMProvider using LiteLLM for unified cost tracking.
 
-    Supports two authentication modes:
-    - api_key: Use GEMINI_API_KEY environment variable (BYOK)
-    - adc: Use Google Application Default Credentials (run `gcloud auth application-default login`)
+    All calls are routed through LiteLLM:
+    - api_key mode: Uses gemini/model-name prefix (requires GEMINI_API_KEY)
+    - adc mode: Uses vertex_ai/model-name prefix (requires VERTEXAI_PROJECT, VERTEXAI_LOCATION)
     """
 
     def __init__(
@@ -32,7 +34,7 @@ class GeminiProvider(LLMProvider):
         auth_mode: Literal["api_key", "adc"] | None = None,
     ):
         """
-        Initialize GeminiProvider.
+        Initialize GeminiProvider with LiteLLM routing.
 
         Args:
             config: Client configuration.
@@ -41,8 +43,7 @@ class GeminiProvider(LLMProvider):
         """
         self.config = config
         self.logger = logger
-        self.model_summary = None
-        self.model_title = None
+        self._litellm = None
 
         # Determine auth mode from config or parameter
         self._auth_mode: AuthMode = "api_key"  # Default
@@ -51,53 +52,25 @@ class GeminiProvider(LLMProvider):
         elif config.llm_providers and config.llm_providers.gemini:
             self._auth_mode = config.llm_providers.gemini.auth_mode
 
+        # Set up environment for provider/auth_mode
+        setup_provider_env("gemini", self._auth_mode)  # type: ignore[arg-type]
+
         try:
-            import google.generativeai as genai
-
-            # Initialize based on auth mode
-            if self._auth_mode == "adc":
-                # Use Application Default Credentials
-                # User must run: gcloud auth application-default login
-                try:
-                    import google.auth
-
-                    credentials, project = google.auth.default()
-                    genai.configure(credentials=credentials)
-                    self.genai = genai
-                    self.logger.debug("Gemini initialized with ADC credentials")
-                except Exception as e:
-                    self.logger.error(
-                        f"Failed to initialize Gemini with ADC: {e}. "
-                        "Run 'gcloud auth application-default login' to authenticate."
-                    )
-                    self.genai = None
-            else:
-                # Use API key from environment
-                api_key = os.environ.get("GEMINI_API_KEY")
-                if api_key:
-                    genai.configure(api_key=api_key)
-                    self.genai = genai
-                    self.logger.debug("Gemini initialized with API key")
-                else:
-                    self.logger.warning("GEMINI_API_KEY not found in environment variables.")
-                    self.genai = None
-
-            # Initialize models if genai is configured
-            if self.genai:
-                summary_model_name = self.config.session_summary.model or "gemini-1.5-pro"
-                title_model_name = self.config.title_synthesis.model or "gemini-1.5-flash"
-
-                self.model_summary = genai.GenerativeModel(summary_model_name)
-                self.model_title = genai.GenerativeModel(title_model_name)
+            import litellm
+
+            self._litellm = litellm
+            self.logger.debug(
+                f"GeminiProvider initialized with LiteLLM (auth_mode={self._auth_mode})"
+            )
 
         except ImportError:
             self.logger.error(
-                "google-generativeai package not found. Please install with `pip install google-generativeai`."
+                "litellm package not found. Please install with `pip install litellm`."
             )
-            self.genai = None
-        except Exception as e:
-            self.logger.error(f"Failed to initialize Gemini client: {e}")
-            self.genai = None
+
+    def _get_model(self, base_model: str) -> str:
+        """Get the LiteLLM-formatted model name with appropriate prefix."""
+        return get_litellm_model(base_model, "gemini", self._auth_mode)  # type: ignore[arg-type]
 
     @property
     def provider_name(self) -> str:
@@ -113,10 +86,10 @@ class GeminiProvider(LLMProvider):
         self, context: dict[str, Any], prompt_template: str | None = None
     ) -> str:
         """
-        Generate session summary using Gemini.
+        Generate session summary using Gemini via LiteLLM.
         """
-        if not self.genai or not self.model_summary:
-            return "Session summary unavailable (Gemini client not initialized)"
+        if not self._litellm:
+            return "Session summary unavailable (LiteLLM not initialized)"
 
         # Build formatted context for prompt template
         formatted_context = {
@@ -140,20 +113,33 @@ class GeminiProvider(LLMProvider):
         prompt = prompt_template.format(**formatted_context)
 
         try:
-            # Gemini async generation
-            response = await self.model_summary.generate_content_async(prompt)
-            return response.text or ""
+            model_name = self.config.session_summary.model or "gemini-1.5-pro"
+            litellm_model = self._get_model(model_name)
+
+            response = await self._litellm.acompletion(
+                model=litellm_model,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a session summary generator. Create comprehensive, actionable summaries.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=4000,
+                timeout=120,
+            )
+            return response.choices[0].message.content or ""
         except Exception as e:
-            self.logger.error(f"Failed to generate summary with Gemini: {e}")
+            self.logger.error(f"Failed to generate summary with Gemini via LiteLLM: {e}")
             return f"Session summary generation failed: {e}"
 
     async def synthesize_title(
         self, user_prompt: str, prompt_template: str | None = None
     ) -> str | None:
         """
-        Synthesize session title using Gemini.
+        Synthesize session title using Gemini via LiteLLM.
         """
-        if not self.genai or not self.model_title:
+        if not self._litellm:
             return None
 
         # Build prompt - prompt_template is required
@@ -165,10 +151,24 @@ class GeminiProvider(LLMProvider):
         prompt = prompt_template.format(user_prompt=user_prompt)
 
         try:
-            response = await self.model_title.generate_content_async(prompt)
-            return (response.text or "").strip()
+            model_name = self.config.title_synthesis.model or "gemini-1.5-flash"
+            litellm_model = self._get_model(model_name)
+
+            response = await self._litellm.acompletion(
+                model=litellm_model,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a session title generator. Create concise, descriptive titles.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=50,
+                timeout=30,
+            )
+            return (response.choices[0].message.content or "").strip()
         except Exception as e:
-            self.logger.error(f"Failed to synthesize title with Gemini: {e}")
+            self.logger.error(f"Failed to synthesize title with Gemini via LiteLLM: {e}")
             return None
 
     async def generate_text(
@@ -178,28 +178,30 @@ class GeminiProvider(LLMProvider):
         model: str | None = None,
     ) -> str:
         """
-        Generate text using Gemini.
+        Generate text using Gemini via LiteLLM.
         """
-        if not self.genai:
-            return "Generation unavailable (Gemini client not initialized)"
+        if not self._litellm:
+            return "Generation unavailable (LiteLLM not initialized)"
 
         model_name = model or "gemini-1.5-flash"
+        litellm_model = self._get_model(model_name)
 
         try:
-            # Note: Gemini system prompts are configured at model creation,
-            # but simple generation usually just includes it in the prompt or uses default.
-            # For simplicity we'll just generate content.
-            model_instance = self.genai.GenerativeModel(model_name)
-
-            full_prompt = prompt
-            if system_prompt:
-                # Prepend system prompt if provided
-                full_prompt = f"{system_prompt}\n\n{prompt}"
-
-            response = await model_instance.generate_content_async(full_prompt)
-            return response.text or ""
+            response = await self._litellm.acompletion(
+                model=litellm_model,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt or "You are a helpful assistant.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=4000,
+                timeout=120,
+            )
+            return response.choices[0].message.content or ""
         except Exception as e:
-            self.logger.error(f"Failed to generate text with Gemini: {e}")
+            self.logger.error(f"Failed to generate text with Gemini via LiteLLM: {e}")
             return f"Generation failed: {e}"
 
     async def describe_image(
@@ -208,9 +210,7 @@ class GeminiProvider(LLMProvider):
         context: str | None = None,
     ) -> str:
         """
-        Generate a text description of an image using Gemini's vision capabilities.
-
-        Uses Gemini 1.5 Flash for efficient image description.
+        Generate a text description of an image using Gemini's vision via LiteLLM.
 
         Args:
             image_path: Path to the image file
@@ -219,40 +219,64 @@ class GeminiProvider(LLMProvider):
         Returns:
             Text description of the image
         """
+        import base64
+        import mimetypes
         from pathlib import Path
 
-        if not self.genai:
-            return "Image description unavailable (Gemini client not initialized)"
+        if not self._litellm:
+            return "Image description unavailable (LiteLLM not initialized)"
 
         path = Path(image_path)
         if not path.exists():
             return f"Image not found: {image_path}"
 
         try:
-            # Use PIL to load the image - Gemini accepts PIL images directly
-            from PIL import Image
-
-            # Use context manager to ensure image file handle is properly closed
-            with Image.open(path) as image:
-                # Build prompt
-                prompt = (
-                    "Please describe this image in detail, focusing on key visual elements, "
-                    "any text visible, and the overall context or meaning."
-                )
-                if context:
-                    prompt = f"{context}\n\n{prompt}"
-
-                # Use gemini-1.5-flash for efficient vision tasks
-                model = self.genai.GenerativeModel("gemini-1.5-flash")
-
-                # Generate content with image and prompt
-                response = await model.generate_content_async([prompt, image])
+            # Read and encode image
+            image_data = path.read_bytes()
+            image_base64 = base64.standard_b64encode(image_data).decode("utf-8")
+
+            # Determine media type
+            mime_type, _ = mimetypes.guess_type(str(path))
+            if mime_type not in [
+                "image/jpeg",
+                "image/png",
+                "image/webp",
+                "image/heic",
+                "image/heif",
+            ]:
+                mime_type = "image/png"
+
+            # Build prompt
+            prompt = (
+                "Please describe this image in detail, focusing on key visual elements, "
+                "any text visible, and the overall context or meaning."
+            )
+            if context:
+                prompt = f"{context}\n\n{prompt}"
+
+            # Use gemini-1.5-flash via LiteLLM for efficient vision tasks
+            litellm_model = self._get_model("gemini-1.5-flash")
+
+            response = await self._litellm.acompletion(
+                model=litellm_model,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt},
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": f"data:{mime_type};base64,{image_base64}"},
+                            },
+                        ],
+                    }
+                ],
+                max_tokens=1000,
+                timeout=60,
+            )
 
-                return response.text or "No description generated"
+            return response.choices[0].message.content or "No description generated"
 
-        except ImportError:
-            self.logger.error("PIL/Pillow not installed. Required for image description.")
-            return "Image description unavailable (PIL not installed)"
         except Exception as e:
-            self.logger.error(f"Failed to describe image with Gemini: {e}")
+            self.logger.error(f"Failed to describe image with Gemini via LiteLLM: {e}")
             return f"Image description failed: {e}"

gobby/llm/litellm_executor.py

@@ -4,17 +4,22 @@ LiteLLM implementation of AgentExecutor.
 Provides a unified interface to 100+ LLM providers using OpenAI-compatible
 function calling API. Supports models from OpenAI, Anthropic, Mistral,
 Cohere, and many others through a single interface.
+
+This executor is the unified path for all api_key and adc authentication modes
+across all providers (Claude, Gemini, Codex/OpenAI). Provider-specific executors
+are only used for subscription/cli modes that require special SDK integrations.
 """
 
 import asyncio
 import json
 import logging
 import os
-from typing import Any
+from typing import Any, Literal
 
 from gobby.llm.executor import (
     AgentExecutor,
     AgentResult,
+    CostInfo,
     ToolCallRecord,
     ToolHandler,
     ToolResult,
@@ -23,6 +28,91 @@ from gobby.llm.executor import (
 
 logger = logging.getLogger(__name__)
 
+# Provider type for routing
+ProviderType = Literal["claude", "gemini", "codex", "openai", "litellm"]
+AuthModeType = Literal["api_key", "adc"]
+
+
+def get_litellm_model(
+    model: str,
+    provider: ProviderType | None = None,
+    auth_mode: AuthModeType | None = None,
+) -> str:
+    """
+    Map provider/model/auth_mode to LiteLLM model string format.
+
+    LiteLLM uses prefixes to route to the correct provider:
+    - anthropic/model-name -> Anthropic API
+    - gemini/model-name -> Google AI Studio (API key)
+    - vertex_ai/model-name -> Google Vertex AI (ADC)
+    - No prefix -> OpenAI (default)
+
+    Args:
+        model: The model name (e.g., "claude-sonnet-4-5", "gemini-2.0-flash")
+        provider: The provider type (claude, gemini, codex, openai)
+        auth_mode: The authentication mode (api_key, adc)
+
+    Returns:
+        LiteLLM-formatted model string with appropriate prefix.
+
+    Examples:
+        >>> get_litellm_model("claude-sonnet-4-5", provider="claude")
+        "anthropic/claude-sonnet-4-5"
+        >>> get_litellm_model("gemini-2.0-flash", provider="gemini", auth_mode="api_key")
+        "gemini/gemini-2.0-flash"
+        >>> get_litellm_model("gemini-2.0-flash", provider="gemini", auth_mode="adc")
+        "vertex_ai/gemini-2.0-flash"
+        >>> get_litellm_model("gpt-4o", provider="codex")
+        "gpt-4o"
+    """
+    # If model already has a prefix, assume it's already formatted
+    if "/" in model:
+        return model
+
+    if provider == "claude":
+        return f"anthropic/{model}"
+    elif provider == "gemini":
+        if auth_mode == "adc":
+            # ADC uses Vertex AI endpoint
+            return f"vertex_ai/{model}"
+        # API key uses Gemini API endpoint
+        return f"gemini/{model}"
+    elif provider in ("codex", "openai"):
+        # OpenAI models don't need a prefix
+        return model
+    else:
+        # Default: return as-is (OpenAI-compatible or already prefixed)
+        return model
+
+
+def setup_provider_env(
+    provider: ProviderType | None = None,
+    auth_mode: AuthModeType | None = None,
+) -> None:
+    """
+    Set up environment variables needed for specific provider/auth_mode combinations.
+
+    For Gemini ADC mode via Vertex AI, this ensures VERTEXAI_PROJECT and
+    VERTEXAI_LOCATION are set from common Google Cloud environment variables.
+
+    Args:
+        provider: The provider type
+        auth_mode: The authentication mode
+    """
+    if provider == "gemini" and auth_mode == "adc":
+        # Vertex AI needs project and location
+        # Check if already set, otherwise try common GCP env vars
+        if "VERTEXAI_PROJECT" not in os.environ:
+            project = os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("GCLOUD_PROJECT")
+            if project:
+                os.environ["VERTEXAI_PROJECT"] = project
+                logger.debug(f"Set VERTEXAI_PROJECT from GCP env: {project}")
+
+        if "VERTEXAI_LOCATION" not in os.environ:
+            location = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
+            os.environ["VERTEXAI_LOCATION"] = location
+            logger.debug(f"Set VERTEXAI_LOCATION: {location}")
+
 
 class LiteLLMExecutor(AgentExecutor):
     """
@@ -31,6 +121,12 @@ class LiteLLMExecutor(AgentExecutor):
     Uses LiteLLM's unified API to access 100+ LLM providers with OpenAI-compatible
     function calling. Supports models from OpenAI, Anthropic, Mistral, Cohere, etc.
 
+    This is the unified executor for all api_key and adc authentication modes:
+    - Claude (api_key) -> anthropic/model-name
+    - Gemini (api_key) -> gemini/model-name
+    - Gemini (adc) -> vertex_ai/model-name
+    - Codex/OpenAI (api_key) -> model-name (no prefix)
+
     The executor implements a proper agentic loop:
     1. Send prompt to LLM with function/tool schemas
     2. When LLM requests a function call, call tool_handler
@@ -38,12 +134,17 @@ class LiteLLMExecutor(AgentExecutor):
     4. Repeat until LLM stops requesting functions or limits are reached
 
     Example:
-        >>> executor = LiteLLMExecutor(default_model="gpt-4o-mini")
+        >>> executor = LiteLLMExecutor(
+        ...     default_model="claude-sonnet-4-5",
+        ...     provider="claude",
+        ...     auth_mode="api_key",
+        ... )
         >>> result = await executor.run(
         ...     prompt="Create a task",
         ...     tools=[ToolSchema(name="create_task", ...)],
         ...     tool_handler=my_handler,
         ... )
+        >>> print(result.cost_info)  # Unified cost tracking
     """
 
     def __init__(
@@ -51,20 +152,28 @@ class LiteLLMExecutor(AgentExecutor):
         default_model: str = "gpt-4o-mini",
         api_base: str | None = None,
         api_keys: dict[str, str] | None = None,
+        provider: ProviderType | None = None,
+        auth_mode: AuthModeType | None = None,
     ):
         """
         Initialize LiteLLMExecutor.
 
         Args:
             default_model: Default model to use if not specified in run().
-                          Examples: "gpt-4o-mini", "claude-3-sonnet-20240229",
-                          "mistral/mistral-large-latest"
+                          Examples: "gpt-4o-mini", "claude-sonnet-4-5",
+                          "gemini-2.0-flash"
             api_base: Optional custom API base URL (e.g., OpenRouter endpoint).
             api_keys: Optional dict of API keys to set in environment.
                      Keys should be like "OPENAI_API_KEY", "ANTHROPIC_API_KEY", etc.
+            provider: Provider type for model routing (claude, gemini, codex, openai).
+                     Used to determine the correct LiteLLM model prefix.
+            auth_mode: Authentication mode (api_key, adc).
+                      Used for Gemini to choose between gemini/ and vertex_ai/ prefixes.
         """
         self.default_model = default_model
         self.api_base = api_base
+        self.provider = provider
+        self.auth_mode = auth_mode
         self.logger = logger
         self._litellm: Any = None
 
@@ -80,7 +189,12 @@ class LiteLLMExecutor(AgentExecutor):
                         os.environ[key] = value
                         self.logger.debug(f"Set {key} from config")
 
-            self.logger.debug("LiteLLM executor initialized")
+            # Set up provider-specific environment variables
+            setup_provider_env(provider, auth_mode)
+
+            self.logger.debug(
+                f"LiteLLM executor initialized (provider={provider}, auth_mode={auth_mode})"
+            )
 
         except ImportError as e:
             raise ImportError(
@@ -151,7 +265,13 @@ class LiteLLMExecutor(AgentExecutor):
             )
 
         tool_calls_records: list[ToolCallRecord] = []
-        effective_model = model or self.default_model
+        # Apply model routing based on provider/auth_mode
+        raw_model = model or self.default_model
+        effective_model = get_litellm_model(raw_model, self.provider, self.auth_mode)
+        self.logger.debug(f"Model routing: {raw_model} -> {effective_model}")
+
+        # Track cumulative costs across turns (outer scope for timeout handler)
+        cost_tracker = [CostInfo(model=effective_model)]
 
         # Track turns in outer scope so timeout handler can access the count
         turns_counter = [0]
@@ -197,6 +317,19 @@ class LiteLLMExecutor(AgentExecutor):
                     # Call LiteLLM
                     response = await litellm.acompletion(**completion_kwargs)
 
+                    # Track costs
+                    if hasattr(response, "usage") and response.usage:
+                        cost_tracker[0].prompt_tokens += response.usage.prompt_tokens or 0
+                        cost_tracker[0].completion_tokens += response.usage.completion_tokens or 0
+
+                    # Calculate cost using LiteLLM's cost tracking
+                    try:
+                        turn_cost = litellm.completion_cost(response)
+                        cost_tracker[0].total_cost += turn_cost
+                    except Exception:  # nosec B110 - best effort cost tracking, failure is non-critical
+                        # Cost calculation may fail for some models
+                        pass
+
                 except Exception as e:
                     self.logger.error(f"LiteLLM API error: {e}")
                     return AgentResult(
@@ -205,6 +338,7 @@ class LiteLLMExecutor(AgentExecutor):
                         tool_calls=tool_calls_records,
                         error=f"LiteLLM API error: {e}",
                         turns_used=turns_used,
+                        cost_info=cost_tracker[0],
                     )
 
                 # Process response
@@ -222,6 +356,7 @@ class LiteLLMExecutor(AgentExecutor):
                         status="success",
                         tool_calls=tool_calls_records,
                         turns_used=turns_used,
+                        cost_info=cost_tracker[0],
                     )
 
                 # Add assistant message to history
@@ -288,6 +423,7 @@ class LiteLLMExecutor(AgentExecutor):
                 status="partial",
                 tool_calls=tool_calls_records,
                 turns_used=turns_used,
+                cost_info=cost_tracker[0],
             )
 
         # Run with timeout
@@ -300,4 +436,5 @@ class LiteLLMExecutor(AgentExecutor):
                 tool_calls=tool_calls_records,
                 error=f"Execution timed out after {timeout}s",
                 turns_used=turns_counter[0],
+                cost_info=cost_tracker[0],
             )