gobby 0.2.5__py3-none-any.whl

gobby/llm/gemini_executor.py

@@ -0,0 +1,339 @@
+"""
+Gemini implementation of AgentExecutor.
+
+Supports two authentication modes:
+- api_key: Use GEMINI_API_KEY environment variable or provided key
+- adc: Use Google Application Default Credentials (gcloud auth)
+"""
+
+import asyncio
+import logging
+import os
+from typing import Any, Literal
+
+from gobby.llm.executor import (
+    AgentExecutor,
+    AgentResult,
+    ToolCallRecord,
+    ToolHandler,
+    ToolResult,
+    ToolSchema,
+)
+
+logger = logging.getLogger(__name__)
+
+# Auth mode type
+GeminiAuthMode = Literal["api_key", "adc"]
+
+
+class GeminiExecutor(AgentExecutor):
+    """
+    Gemini implementation of AgentExecutor.
+
+    Supports two authentication modes:
+    - api_key: Uses GEMINI_API_KEY environment variable or provided key
+    - adc: Uses Google Application Default Credentials (run `gcloud auth application-default login`)
+
+    The executor implements a proper agentic loop:
+    1. Send prompt to Gemini with function declarations
+    2. When Gemini requests a function call, call tool_handler
+    3. Send function result back to Gemini
+    4. Repeat until Gemini stops requesting functions or limits are reached
+
+    Example:
+        >>> executor = GeminiExecutor(auth_mode="api_key", api_key="...")
+        >>> result = await executor.run(
+        ...     prompt="Create a task",
+        ...     tools=[ToolSchema(name="create_task", ...)],
+        ...     tool_handler=my_handler,
+        ... )
+    """
+
+    def __init__(
+        self,
+        auth_mode: GeminiAuthMode = "api_key",
+        api_key: str | None = None,
+        default_model: str = "gemini-2.0-flash",
+    ):
+        """
+        Initialize GeminiExecutor.
+
+        Args:
+            auth_mode: Authentication mode ("api_key" or "adc").
+            api_key: Gemini API key (optional for api_key mode, uses GEMINI_API_KEY env var).
+            default_model: Default model to use if not specified in run().
+        """
+        self.auth_mode = auth_mode
+        self.default_model = default_model
+        self.logger = logger
+        self._genai: Any = None
+
+        try:
+            import google.generativeai as genai
+
+            if auth_mode == "adc":
+                # Use Application Default Credentials
+                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:
+                    raise ValueError(
+                        f"Failed to initialize Gemini with ADC: {e}. "
+                        "Run 'gcloud auth application-default login' to authenticate."
+                    ) from e
+            else:
+                # Use API key from parameter or environment
+                key = api_key or os.environ.get("GEMINI_API_KEY")
+                if not key:
+                    raise ValueError(
+                        "API key required for api_key mode. "
+                        "Provide api_key parameter or set GEMINI_API_KEY env var."
+                    )
+                genai.configure(api_key=key)
+                self._genai = genai
+                self.logger.debug("Gemini initialized with API key")
+
+        except ImportError as e:
+            raise ImportError(
+                "google-generativeai package not found. "
+                "Please install with `pip install google-generativeai`."
+            ) from e
+
+    @property
+    def provider_name(self) -> str:
+        """Return the provider name."""
+        return "gemini"
+
+    def _convert_tools_to_gemini_format(self, tools: list[ToolSchema]) -> list[dict[str, Any]]:
+        """Convert ToolSchema list to Gemini function declarations format."""
+        function_declarations = []
+        for tool in tools:
+            # Build parameter schema
+            params = tool.input_schema.copy()
+            # Ensure type is object
+            if "type" not in params:
+                params["type"] = "object"
+
+            function_declarations.append(
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": params,
+                }
+            )
+        return function_declarations
+
+    async def run(
+        self,
+        prompt: str,
+        tools: list[ToolSchema],
+        tool_handler: ToolHandler,
+        system_prompt: str | None = None,
+        model: str | None = None,
+        max_turns: int = 10,
+        timeout: float = 120.0,
+    ) -> AgentResult:
+        """
+        Execute an agentic loop with function calling.
+
+        Runs Gemini with the given prompt, calling tools via tool_handler
+        until completion, max_turns, or timeout.
+
+        Args:
+            prompt: The user prompt to process.
+            tools: List of available tools with their schemas.
+            tool_handler: Callback to execute tool calls.
+            system_prompt: Optional system prompt.
+            model: Optional model override.
+            max_turns: Maximum turns before stopping (default: 10).
+            timeout: Maximum execution time in seconds (default: 120.0).
+
+        Returns:
+            AgentResult with output, status, and tool call records.
+        """
+        if self._genai is None:
+            return AgentResult(
+                output="",
+                status="error",
+                error="Gemini client not initialized",
+                turns_used=0,
+            )
+
+        tool_calls: list[ToolCallRecord] = []
+        effective_model = model or self.default_model
+
+        # Track turns in outer scope so timeout handler can access the count
+        turns_counter = [0]
+
+        async def _run_loop() -> AgentResult:
+            turns_used = 0
+            final_output = ""
+            genai = self._genai
+            if genai is None:
+                raise RuntimeError("GeminiExecutor genai not initialized")
+
+            # Create the model with tools
+            gemini_tools = self._convert_tools_to_gemini_format(tools)
+
+            # Create Tool instance (SDK expects Tool objects, not plain dicts)
+            tool_instance = None
+            if gemini_tools:
+                tool_instance = genai.protos.Tool(function_declarations=gemini_tools)
+
+            # Create model with system instruction
+            generation_config = {
+                "max_output_tokens": 8192,
+                "temperature": 0.7,
+            }
+
+            model_instance = genai.GenerativeModel(
+                model_name=effective_model,
+                system_instruction=system_prompt or "You are a helpful assistant.",
+                generation_config=generation_config,
+                tools=[tool_instance] if tool_instance else None,
+            )
+
+            # Start chat
+            chat = model_instance.start_chat()
+
+            # Send initial message
+            try:
+                response = await chat.send_message_async(prompt)
+            except Exception as e:
+                self.logger.error(f"Gemini API error: {e}")
+                return AgentResult(
+                    output="",
+                    status="error",
+                    tool_calls=tool_calls,
+                    error=f"Gemini API error: {e}",
+                    turns_used=0,
+                )
+
+            while turns_used < max_turns:
+                turns_used += 1
+                turns_counter[0] = turns_used
+
+                # Extract function calls and text from response
+                function_calls: list[dict[str, Any]] = []
+
+                for candidate in response.candidates:
+                    for part in candidate.content.parts:
+                        # Check for text content
+                        if hasattr(part, "text") and part.text:
+                            final_output = part.text
+
+                        # Check for function call
+                        if hasattr(part, "function_call") and part.function_call:
+                            fc = part.function_call
+                            function_calls.append(
+                                {
+                                    "name": fc.name,
+                                    "args": dict(fc.args) if fc.args else {},
+                                }
+                            )
+
+                # If no function calls, we're done
+                if not function_calls:
+                    return AgentResult(
+                        output=final_output,
+                        status="success",
+                        tool_calls=tool_calls,
+                        turns_used=turns_used,
+                    )
+
+                # Handle function calls
+                function_responses = []
+
+                for fc in function_calls:
+                    tool_name = fc["name"]
+                    arguments = fc["args"]
+
+                    # Record the tool call
+                    record = ToolCallRecord(
+                        tool_name=tool_name,
+                        arguments=arguments,
+                    )
+                    tool_calls.append(record)
+
+                    # Execute via handler
+                    try:
+                        result = await tool_handler(tool_name, arguments)
+                        record.result = result
+
+                        # Format result for Gemini
+                        if result.success:
+                            # Use 'is not None' to preserve legitimate falsy values like 0, False, {}
+                            response_data = (
+                                result.result
+                                if result.result is not None
+                                else {"status": "success"}
+                            )
+                        else:
+                            response_data = {"error": result.error}
+
+                        function_responses.append(
+                            genai.protos.Part(
+                                function_response=genai.protos.FunctionResponse(
+                                    name=tool_name,
+                                    response=(
+                                        response_data
+                                        if isinstance(response_data, dict)
+                                        else {"result": response_data}
+                                    ),
+                                )
+                            )
+                        )
+                    except Exception as e:
+                        self.logger.error(f"Tool handler error for {tool_name}: {e}")
+                        record.result = ToolResult(
+                            tool_name=tool_name,
+                            success=False,
+                            error=str(e),
+                        )
+                        function_responses.append(
+                            genai.protos.Part(
+                                function_response=genai.protos.FunctionResponse(
+                                    name=tool_name,
+                                    response={"error": str(e)},
+                                )
+                            )
+                        )
+
+                # Send function responses back to Gemini
+                try:
+                    response = await chat.send_message_async(function_responses)
+                    # Response will be processed in the next iteration of the while loop
+                    # which extracts function calls and text directly from the response object
+                except Exception as e:
+                    self.logger.error(f"Error sending function response: {e}")
+                    return AgentResult(
+                        output=final_output,
+                        status="error",
+                        tool_calls=tool_calls,
+                        error=f"Error sending function response: {e}",
+                        turns_used=turns_used,
+                    )
+
+            # Max turns reached
+            return AgentResult(
+                output=final_output,
+                status="partial",
+                tool_calls=tool_calls,
+                turns_used=turns_used,
+            )
+
+        # Run with timeout
+        try:
+            return await asyncio.wait_for(_run_loop(), timeout=timeout)
+        except TimeoutError:
+            return AgentResult(
+                output="",
+                status="timeout",
+                tool_calls=tool_calls,
+                error=f"Execution timed out after {timeout}s",
+                turns_used=turns_counter[0],
+            )

gobby/llm/litellm.py

@@ -0,0 +1,287 @@
+"""
+LiteLLM implementation of LLMProvider.
+
+LiteLLM provides a unified interface to many LLM providers (OpenAI, Anthropic,
+Mistral, Cohere, etc.) through their APIs using BYOK (Bring Your Own Key).
+
+This provider is useful when users want to use their own API keys for
+multiple different providers without needing separate provider implementations.
+"""
+
+import json
+import logging
+from typing import Any
+
+from gobby.config.app import DaemonConfig
+from gobby.llm.base import AuthMode, LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+class LiteLLMProvider(LLMProvider):
+    """
+    LiteLLM implementation of LLMProvider.
+
+    Uses API key-based authentication (BYOK) for multiple providers.
+    API keys are read from:
+    1. llm_providers.api_keys in config (e.g., OPENAI_API_KEY, MISTRAL_API_KEY)
+    2. Environment variables as fallback
+
+    Example config:
+        llm_providers:
+          litellm:
+            models: gpt-4o-mini,mistral-large
+            auth_mode: api_key
+          api_keys:
+            OPENAI_API_KEY: sk-...
+            MISTRAL_API_KEY: ...
+    """
+
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "litellm"
+
+    @property
+    def auth_mode(self) -> AuthMode:
+        """LiteLLM uses API key authentication."""
+        return "api_key"
+
+    def __init__(self, config: DaemonConfig):
+        """
+        Initialize LiteLLMProvider.
+
+        Args:
+            config: Client configuration with optional api_keys in llm_providers.
+        """
+        self.config = config
+        self.logger = logger
+        self._litellm = None
+        self._api_keys: dict[str, str] = {}
+
+        # Load API keys from config
+        if config.llm_providers and config.llm_providers.api_keys:
+            self._api_keys = config.llm_providers.api_keys.copy()
+
+        try:
+            import litellm
+
+            self._litellm = litellm
+
+            # Set API keys in litellm's environment
+            # LiteLLM reads from os.environ, so we set them there
+            import os
+
+            for key, value in self._api_keys.items():
+                if value and key not in os.environ:
+                    os.environ[key] = value
+                    self.logger.debug(f"Set {key} from config")
+
+            self.logger.debug("LiteLLM provider initialized")
+
+        except ImportError:
+            self.logger.error(
+                "litellm package not found. Please install with `pip install litellm`."
+            )
+        except Exception as e:
+            self.logger.error(f"Failed to initialize LiteLLM: {e}")
+
+    def _get_model(self, task: str) -> str:
+        """
+        Get the model to use for a specific task.
+
+        Args:
+            task: Task type ("summary" or "title")
+
+        Returns:
+            Model name string
+        """
+        if task == "summary":
+            if self.config.session_summary:
+                return self.config.session_summary.model or "gpt-4o-mini"
+            return "gpt-4o-mini"
+        elif task == "title":
+            if self.config.title_synthesis:
+                return self.config.title_synthesis.model or "gpt-4o-mini"
+            return "gpt-4o-mini"
+        else:
+            return "gpt-4o-mini"
+
+    async def generate_summary(
+        self, context: dict[str, Any], prompt_template: str | None = None
+    ) -> str:
+        """
+        Generate session summary using LiteLLM.
+        """
+        if not self._litellm:
+            return "Session summary unavailable (LiteLLM not initialized)"
+
+        # Build formatted context for prompt template
+        formatted_context = {
+            "transcript_summary": context.get("transcript_summary", ""),
+            "last_messages": json.dumps(context.get("last_messages", []), indent=2),
+            "git_status": context.get("git_status", ""),
+            "file_changes": context.get("file_changes", ""),
+            **{
+                k: v
+                for k, v in context.items()
+                if k not in ["transcript_summary", "last_messages", "git_status", "file_changes"]
+            },
+        }
+
+        # Build prompt - prompt_template is required
+        if not prompt_template:
+            raise ValueError(
+                "prompt_template is required for generate_summary. "
+                "Configure 'session_summary.prompt' in ~/.gobby/config.yaml"
+            )
+        prompt = prompt_template.format(**formatted_context)
+
+        try:
+            # Use LiteLLM's async completion
+            response = await self._litellm.acompletion(
+                model=self._get_model("summary"),
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a session summary generator. Create comprehensive, actionable summaries.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=4000,
+            )
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            self.logger.error(f"Failed to generate summary with 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 LiteLLM.
+        """
+        if not self._litellm:
+            return None
+
+        # Build prompt - prompt_template is required
+        if not prompt_template:
+            raise ValueError(
+                "prompt_template is required for synthesize_title. "
+                "Configure 'title_synthesis.prompt' in ~/.gobby/config.yaml"
+            )
+        prompt = prompt_template.format(user_prompt=user_prompt)
+
+        try:
+            response = await self._litellm.acompletion(
+                model=self._get_model("title"),
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a session title generator. Create concise, descriptive titles.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=50,
+            )
+            return (response.choices[0].message.content or "").strip()
+        except Exception as e:
+            self.logger.error(f"Failed to synthesize title with LiteLLM: {e}")
+            return None
+
+    async def generate_text(
+        self,
+        prompt: str,
+        system_prompt: str | None = None,
+        model: str | None = None,
+    ) -> str:
+        """
+        Generate text using LiteLLM.
+        """
+        if not self._litellm:
+            return "Generation unavailable (LiteLLM not initialized)"
+
+        try:
+            response = await self._litellm.acompletion(
+                model=model or "gpt-4o-mini",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt or "You are a helpful assistant.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=4000,
+            )
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            self.logger.error(f"Failed to generate text with LiteLLM: {e}")
+            return f"Generation failed: {e}"
+
+    async def describe_image(
+        self,
+        image_path: str,
+        context: str | None = None,
+    ) -> str:
+        """
+        Generate a text description of an image using LiteLLM's vision support.
+
+        Args:
+            image_path: Path to the image file to describe
+            context: Optional context to guide the description
+
+        Returns:
+            Text description of the image
+        """
+        import base64
+        import mimetypes
+        from pathlib import Path
+
+        if not self._litellm:
+            return "Image description unavailable (LiteLLM not initialized)"
+
+        # Validate image exists
+        path = Path(image_path)
+        if not path.exists():
+            return f"Image not found: {image_path}"
+
+        # Read and encode image
+        try:
+            image_data = path.read_bytes()
+            image_base64 = base64.standard_b64encode(image_data).decode("utf-8")
+        except Exception as e:
+            self.logger.error(f"Failed to read image {image_path}: {e}")
+            return f"Failed to read image: {e}"
+
+        # Determine media type
+        mime_type, _ = mimetypes.guess_type(str(path))
+        if mime_type not in ["image/jpeg", "image/png", "image/gif", "image/webp"]:
+            mime_type = "image/png"
+
+        # Build prompt
+        prompt = "Please describe this image in detail, focusing on the key visual elements and any text visible."
+        if context:
+            prompt = f"{context}\n\n{prompt}"
+
+        try:
+            # Use LiteLLM's vision support (works with gpt-4o, claude-3, etc.)
+            response = await self._litellm.acompletion(
+                model="gpt-4o-mini",  # Default to a vision-capable 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,
+            )
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            self.logger.error(f"Failed to describe image with LiteLLM: {e}")
+            return f"Image description failed: {e}"