code-puppy 0.0.287__py3-none-any.whl → 0.0.323__py3-none-any.whl

code_puppy/gemini_code_assist.py

@@ -0,0 +1,385 @@
+"""Gemini Code Assist Model for pydantic_ai.
+
+This module provides a custom Model implementation that uses Google's
+Code Assist API (cloudcode-pa.googleapis.com) instead of the standard
+Generative Language API. The Code Assist API supports OAuth authentication
+and has a different request/response format.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+
+import httpx
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_ai.models import Model, ModelRequestParameters
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.usage import RequestUsage
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiCodeAssistModel(Model):
+    """Model implementation for Google's Code Assist API.
+
+    This uses the cloudcode-pa.googleapis.com endpoint which accepts OAuth
+    tokens and has a wrapped request/response format.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        access_token: str,
+        project_id: str,
+        api_base_url: str = "https://cloudcode-pa.googleapis.com",
+        api_version: str = "v1internal",
+    ):
+        self._model_name = model_name
+        self.access_token = access_token
+        self.project_id = project_id
+        self.api_base_url = api_base_url
+        self.api_version = api_version
+
+    def model_name(self) -> str:
+        """Return the model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        return "google"
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a non-streaming request to the Code Assist API."""
+        request_body = self._build_request(
+            messages, model_settings, model_request_parameters
+        )
+
+        url = f"{self.api_base_url}/{self.api_version}:generateContent"
+        headers = self._get_headers()
+
+        async with httpx.AsyncClient(timeout=180) as client:
+            response = await client.post(url, json=request_body, headers=headers)
+
+            if response.status_code != 200:
+                error_text = response.text
+                raise RuntimeError(
+                    f"Code Assist API error {response.status_code}: {error_text}"
+                )
+
+            data = response.json()
+
+        return self._parse_response(data)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request to the Code Assist API."""
+        request_body = self._build_request(
+            messages, model_settings, model_request_parameters
+        )
+
+        url = f"{self.api_base_url}/{self.api_version}:streamGenerateContent?alt=sse"
+        headers = self._get_headers()
+
+        async with httpx.AsyncClient(timeout=180) as client:
+            async with client.stream(
+                "POST", url, json=request_body, headers=headers
+            ) as response:
+                if response.status_code != 200:
+                    error_text = await response.aread()
+                    raise RuntimeError(
+                        f"Code Assist API error {response.status_code}: {error_text.decode()}"
+                    )
+
+                yield StreamedResponse(response, self._model_name)
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get HTTP headers for the request."""
+        return {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+
+    def _build_request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> Dict[str, Any]:
+        """Build the Code Assist API request body."""
+        contents = []
+        system_instruction = None
+
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                for part in msg.parts:
+                    if isinstance(part, SystemPromptPart):
+                        # Collect system prompt
+                        if system_instruction is None:
+                            system_instruction = {
+                                "role": "user",
+                                "parts": [{"text": part.content}],
+                            }
+                        else:
+                            system_instruction["parts"].append({"text": part.content})
+                    elif isinstance(part, UserPromptPart):
+                        contents.append(
+                            {
+                                "role": "user",
+                                "parts": [{"text": part.content}],
+                            }
+                        )
+                    elif isinstance(part, ToolReturnPart):
+                        # Serialize content to string if it's not already
+                        content = part.content
+                        if not isinstance(content, (str, int, float, bool, type(None))):
+                            try:
+                                content = json.dumps(content, default=str)
+                            except (TypeError, ValueError):
+                                content = str(content)
+                        contents.append(
+                            {
+                                "role": "user",
+                                "parts": [
+                                    {
+                                        "functionResponse": {
+                                            "name": part.tool_name,
+                                            "response": {"result": content},
+                                        }
+                                    }
+                                ],
+                            }
+                        )
+            elif isinstance(msg, ModelResponse):
+                parts = []
+                first_func_call = True
+                for part in msg.parts:
+                    if isinstance(part, TextPart):
+                        parts.append({"text": part.content})
+                    elif isinstance(part, ToolCallPart):
+                        func_call_part = {
+                            "functionCall": {
+                                "name": part.tool_name,
+                                "args": part.args_as_dict(),
+                            }
+                        }
+                        # Code Assist API requires thoughtSignature on function calls
+                        # Use synthetic signature to skip validation
+                        if first_func_call:
+                            func_call_part["thoughtSignature"] = (
+                                "skip_thought_signature_validator"
+                            )
+                            first_func_call = False
+                        parts.append(func_call_part)
+                if parts:
+                    contents.append({"role": "model", "parts": parts})
+
+        # Build the inner request (Vertex-style format)
+        inner_request: Dict[str, Any] = {
+            "contents": contents,
+        }
+
+        if system_instruction:
+            inner_request["systemInstruction"] = system_instruction
+
+        # Add tools if available
+        if model_request_parameters.function_tools:
+            inner_request["tools"] = [
+                self._build_tools(model_request_parameters.function_tools)
+            ]
+
+        # Add generation config
+        generation_config = self._build_generation_config(model_settings)
+        if generation_config:
+            inner_request["generationConfig"] = generation_config
+
+        # Wrap in Code Assist format
+        return {
+            "model": self._model_name,
+            "project": self.project_id,
+            "user_prompt_id": str(uuid.uuid4()),
+            "request": inner_request,
+        }
+
+    def _build_tools(self, tools: list[ToolDefinition]) -> Dict[str, Any]:
+        """Build tool definitions for the API."""
+        function_declarations = []
+
+        for tool in tools:
+            func_decl: Dict[str, Any] = {
+                "name": tool.name,
+                "description": tool.description or "",
+            }
+
+            if tool.parameters_json_schema:
+                func_decl["parametersJsonSchema"] = tool.parameters_json_schema
+
+            function_declarations.append(func_decl)
+
+        return {"functionDeclarations": function_declarations}
+
+    def _build_generation_config(
+        self, model_settings: ModelSettings | None
+    ) -> Optional[Dict[str, Any]]:
+        """Build generation config from model settings."""
+        if not model_settings:
+            return None
+
+        config: Dict[str, Any] = {}
+
+        if (
+            hasattr(model_settings, "temperature")
+            and model_settings.temperature is not None
+        ):
+            config["temperature"] = model_settings.temperature
+
+        if hasattr(model_settings, "top_p") and model_settings.top_p is not None:
+            config["topP"] = model_settings.top_p
+
+        if (
+            hasattr(model_settings, "max_tokens")
+            and model_settings.max_tokens is not None
+        ):
+            config["maxOutputTokens"] = model_settings.max_tokens
+
+        return config if config else None
+
+    def _parse_response(self, data: Dict[str, Any]) -> ModelResponse:
+        """Parse the Code Assist API response."""
+        # Unwrap the Code Assist response format
+        inner_response = data.get("response", data)
+
+        candidates = inner_response.get("candidates", [])
+        if not candidates:
+            raise RuntimeError("No candidates in response")
+
+        candidate = candidates[0]
+        content = candidate.get("content", {})
+        parts = content.get("parts", [])
+
+        response_parts: list[ModelResponsePart] = []
+
+        for part in parts:
+            if "text" in part:
+                response_parts.append(TextPart(content=part["text"]))
+            elif "functionCall" in part:
+                func_call = part["functionCall"]
+                response_parts.append(
+                    ToolCallPart(
+                        tool_name=func_call["name"],
+                        args=func_call.get("args", {}),
+                        tool_call_id=str(uuid.uuid4()),
+                    )
+                )
+
+        # Extract usage metadata
+        usage_meta = inner_response.get("usageMetadata", {})
+        usage = RequestUsage(
+            input_tokens=usage_meta.get("promptTokenCount", 0),
+            output_tokens=usage_meta.get("candidatesTokenCount", 0),
+        )
+
+        return ModelResponse(
+            parts=response_parts, model_name=self._model_name, usage=usage
+        )
+
+
+class StreamedResponse:
+    """Handler for streaming responses from Code Assist API."""
+
+    def __init__(self, response: httpx.Response, model_name: str):
+        self._response = response
+        self._model_name = model_name
+        self._usage: Optional[RequestUsage] = None
+        self._timestamp = datetime.now(timezone.utc)
+
+    def __aiter__(self) -> AsyncIterator[str]:
+        return self._iter_chunks()
+
+    async def _iter_chunks(self) -> AsyncIterator[str]:
+        """Iterate over SSE chunks from the response."""
+        async for line in self._response.aiter_lines():
+            line = line.strip()
+
+            if line.startswith("data: "):
+                data_str = line[6:]
+                if data_str == "[DONE]":
+                    break
+
+                try:
+                    data = json.loads(data_str)
+                    # Unwrap Code Assist format
+                    inner = data.get("response", data)
+
+                    # Extract usage if available
+                    if "usageMetadata" in inner:
+                        meta = inner["usageMetadata"]
+                        self._usage = RequestUsage(
+                            input_tokens=meta.get("promptTokenCount", 0),
+                            output_tokens=meta.get("candidatesTokenCount", 0),
+                        )
+
+                    # Extract text from candidates
+                    for candidate in inner.get("candidates", []):
+                        content = candidate.get("content", {})
+                        for part in content.get("parts", []):
+                            if "text" in part:
+                                yield part["text"]
+
+                except json.JSONDecodeError:
+                    logger.warning("Failed to parse SSE data: %s", data_str)
+                    continue
+
+    async def get_response_parts(self) -> list[ModelResponsePart]:
+        """Get all response parts after streaming is complete."""
+        text_content = ""
+        tool_calls = []
+
+        async for chunk in self:
+            text_content += chunk
+
+        parts: list[ModelResponsePart] = []
+        if text_content:
+            parts.append(TextPart(content=text_content))
+        parts.extend(tool_calls)
+
+        return parts
+
+    def usage(self) -> RequestUsage:
+        """Get usage statistics."""
+        return self._usage or RequestUsage()
+
+    def model_name(self) -> str:
+        """Get the model name."""
+        return self._model_name
+
+    def timestamp(self) -> datetime:
+        """Get the response timestamp."""
+        return self._timestamp

code_puppy/keymap.py

@@ -0,0 +1,126 @@
+"""Keymap configuration for code-puppy.
+
+This module handles configurable keyboard shortcuts, starting with the
+cancel_agent_key feature that allows users to override Ctrl+C with a
+different key for cancelling agent tasks.
+"""
+
+# Character codes for Ctrl+letter combinations (Ctrl+A = 0x01, Ctrl+Z = 0x1A)
+KEY_CODES: dict[str, str] = {
+    "ctrl+a": "\x01",
+    "ctrl+b": "\x02",
+    "ctrl+c": "\x03",
+    "ctrl+d": "\x04",
+    "ctrl+e": "\x05",
+    "ctrl+f": "\x06",
+    "ctrl+g": "\x07",
+    "ctrl+h": "\x08",
+    "ctrl+i": "\x09",
+    "ctrl+j": "\x0a",
+    "ctrl+k": "\x0b",
+    "ctrl+l": "\x0c",
+    "ctrl+m": "\x0d",
+    "ctrl+n": "\x0e",
+    "ctrl+o": "\x0f",
+    "ctrl+p": "\x10",
+    "ctrl+q": "\x11",
+    "ctrl+r": "\x12",
+    "ctrl+s": "\x13",
+    "ctrl+t": "\x14",
+    "ctrl+u": "\x15",
+    "ctrl+v": "\x16",
+    "ctrl+w": "\x17",
+    "ctrl+x": "\x18",
+    "ctrl+y": "\x19",
+    "ctrl+z": "\x1a",
+    "escape": "\x1b",
+}
+
+# Valid keys for cancel_agent_key configuration
+# NOTE: "escape" is excluded because it conflicts with ANSI escape sequences
+# (arrow keys, F-keys, etc. all start with \x1b)
+VALID_CANCEL_KEYS: set[str] = {
+    "ctrl+c",
+    "ctrl+k",
+    "ctrl+q",
+}
+
+DEFAULT_CANCEL_AGENT_KEY: str = "ctrl+c"
+
+
+class KeymapError(Exception):
+    """Exception raised for keymap configuration errors."""
+
+
+def get_cancel_agent_key() -> str:
+    """Get the configured cancel agent key from config.
+
+    Returns:
+        The key name (e.g., "ctrl+c", "ctrl+k") from config,
+        or the default if not configured.
+    """
+    from code_puppy.config import get_value
+
+    key = get_value("cancel_agent_key")
+    if key is None or key.strip() == "":
+        return DEFAULT_CANCEL_AGENT_KEY
+    return key.strip().lower()
+
+
+def validate_cancel_agent_key() -> None:
+    """Validate the configured cancel agent key.
+
+    Raises:
+        KeymapError: If the configured key is invalid.
+    """
+    key = get_cancel_agent_key()
+    if key not in VALID_CANCEL_KEYS:
+        valid_keys_str = ", ".join(sorted(VALID_CANCEL_KEYS))
+        raise KeymapError(
+            f"Invalid cancel_agent_key '{key}' in puppy.cfg. "
+            f"Valid options are: {valid_keys_str}"
+        )
+
+
+def cancel_agent_uses_signal() -> bool:
+    """Check if the cancel agent key uses SIGINT (Ctrl+C).
+
+    Returns:
+        True if the cancel key is ctrl+c AND we're not on Windows
+        (uses SIGINT handler), False if it uses keyboard listener approach.
+    """
+    import sys
+
+    # On Windows, always use keyboard listener - SIGINT is unreliable
+    if sys.platform == "win32":
+        return False
+
+    return get_cancel_agent_key() == "ctrl+c"
+
+
+def get_cancel_agent_char_code() -> str:
+    """Get the character code for the cancel agent key.
+
+    Returns:
+        The character code (e.g., "\x0b" for ctrl+k).
+
+    Raises:
+        KeymapError: If the key is not found in KEY_CODES.
+    """
+    key = get_cancel_agent_key()
+    if key not in KEY_CODES:
+        raise KeymapError(f"Unknown key '{key}' - no character code mapping found.")
+    return KEY_CODES[key]
+
+
+def get_cancel_agent_display_name() -> str:
+    """Get a human-readable display name for the cancel agent key.
+
+    Returns:
+        A formatted display name like "Ctrl+K".
+    """
+    key = get_cancel_agent_key()
+    if key.startswith("ctrl+"):
+        letter = key.split("+")[1].upper()
+        return f"Ctrl+{letter}"
+    return key.upper()