agentshim 0.1.0__py3-none-any.whl

agentshim/codex/events.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, cast
+
+from ..utils import truncate_content, truncate_tool_params
+
+
+class CodexEvent(ABC):
+    """Base class for Codex (``codex exec --json``) stream events."""
+
+    @abstractmethod
+    def render(self, log_prefix: str) -> str | None:
+        """Render the event as a string for terminal output."""
+
+    @staticmethod
+    def from_dict(data: dict[str, Any]) -> CodexEvent | None:
+        """Factory method to create events from JSON data."""
+        event_type = data.get("type")
+
+        if event_type == "thread.started":
+            thread_id_raw = data.get("thread_id")
+            thread_id = thread_id_raw if isinstance(thread_id_raw, str) else None
+            return ThreadStartedEvent(thread_id=thread_id)
+
+        if event_type == "turn.started":
+            return LifecycleEvent(event_type)
+
+        if event_type == "turn.completed":
+            usage_raw = data.get("usage")
+            usage = cast("dict[str, Any]", usage_raw) if isinstance(usage_raw, dict) else {}
+            return TurnCompletedEvent(
+                input_tokens=int(usage.get("input_tokens") or 0),
+                cached_input_tokens=int(usage.get("cached_input_tokens") or 0),
+                output_tokens=int(usage.get("output_tokens") or 0),
+            )
+
+        if event_type in ("item.started", "item.completed"):
+            item_raw = data.get("item")
+            if not isinstance(item_raw, dict):
+                return None
+            item = cast("dict[str, Any]", item_raw)
+            item_type_raw = item.get("type")
+            item_type = item_type_raw if isinstance(item_type_raw, str) else None
+            item_id_raw = item.get("id")
+            item_id = item_id_raw if isinstance(item_id_raw, str) else None
+            completed = event_type == "item.completed"
+
+            if item_type == "agent_message":
+                if not completed:
+                    return None
+                text_raw = item.get("text", "")
+                return TextEvent(text=text_raw if isinstance(text_raw, str) else "")
+
+            if item_type == "command_execution":
+                exit_code_raw = item.get("exit_code")
+                exit_code = exit_code_raw if isinstance(exit_code_raw, int) else None
+                status_raw = item.get("status")
+                status = status_raw if isinstance(status_raw, str) else None
+                if completed:
+                    return ToolResultEvent(
+                        tool_id=item_id,
+                        output=item.get("aggregated_output", ""),
+                        exit_code=exit_code,
+                        status=status,
+                    )
+                command_raw = item.get("command", "")
+                command = command_raw if isinstance(command_raw, str) else ""
+                return ToolUseEvent(
+                    tool_id=item_id,
+                    tool_name="shell",
+                    parameters={"command": command},
+                )
+
+            status_raw = item.get("status")
+            status = status_raw if isinstance(status_raw, str) else None
+            if completed:
+                return ToolResultEvent(
+                    tool_id=item_id,
+                    output=_summarize_item(item),
+                    exit_code=None,
+                    status=status,
+                )
+            return ToolUseEvent(
+                tool_id=item_id,
+                tool_name=item_type or "item",
+                parameters=_item_parameters(item),
+            )
+
+        if event_type == "turn.failed":
+            err_raw = data.get("error")
+            if isinstance(err_raw, dict):
+                err = cast("dict[str, Any]", err_raw)
+                msg_raw = err.get("message", "")
+                message = msg_raw if isinstance(msg_raw, str) else str(msg_raw)
+            else:
+                message = str(err_raw) if err_raw else ""
+            return ErrorEvent(message=message)
+
+        if event_type == "error":
+            message_raw = data.get("message", "")
+            return ErrorEvent(message=message_raw if isinstance(message_raw, str) else str(message_raw))
+
+        return None
+
+
+def _item_parameters(item: dict[str, Any]) -> dict[str, Any]:
+    """Extract a parameter dict from a generic codex item payload."""
+    excluded = {"id", "type", "status"}
+    return {k: v for k, v in item.items() if k not in excluded}
+
+
+def _summarize_item(item: dict[str, Any]) -> str:
+    """Summarize a generic codex item for the tool-result output field."""
+    for key in ("text", "summary", "output", "result"):
+        value = item.get(key)
+        if isinstance(value, str) and value:
+            return value
+    return ""
+
+
+class ThreadStartedEvent(CodexEvent):
+    """Thread-level start event carrying the resumable ``thread_id``."""
+
+    def __init__(self, thread_id: str | None):
+        self.thread_id = thread_id
+
+    def render(self, log_prefix: str) -> str | None:
+        return None
+
+
+class LifecycleEvent(CodexEvent):
+    """Turn lifecycle marker; not rendered."""
+
+    def __init__(self, event_type: str):
+        self.event_type = event_type
+
+    def render(self, log_prefix: str) -> str | None:
+        return None
+
+
+class TextEvent(CodexEvent):
+    """Assistant text content event (``agent_message`` item)."""
+
+    def __init__(self, text: str):
+        self.text = text
+
+    def render(self, log_prefix: str) -> str | None:
+        return self.text
+
+
+class ToolUseEvent(CodexEvent):
+    """Tool call start event (``item.started``)."""
+
+    def __init__(self, tool_name: str, tool_id: str | None, parameters: Any):
+        self.tool_name = tool_name
+        self.tool_id = tool_id
+        self.parameters = parameters
+
+    def render(self, log_prefix: str) -> str:
+        truncated = truncate_tool_params(self.tool_name, self.parameters)
+        return f"{log_prefix} \033[34m[Tool Use] {self.tool_name} {truncated}\033[0m"
+
+
+class ToolResultEvent(CodexEvent):
+    """Tool call completion event (``item.completed``)."""
+
+    def __init__(
+        self,
+        output: Any,
+        tool_id: str | None,
+        exit_code: int | None = None,
+        status: str | None = None,
+    ):
+        if isinstance(output, list):
+            items = cast("list[Any]", output)
+            self.output = "\n".join(str(item) for item in items)
+        else:
+            self.output = str(output) if output else ""
+        self.tool_id = tool_id
+        self.exit_code = exit_code
+        self.status = status
+        self.tool_name_resolved: str = "Tool"
+
+    def render(self, log_prefix: str) -> str:
+        if not self.output:
+            return f"{log_prefix} \033[32m{self.tool_name_resolved} ran successfully\033[0m"
+        truncated = truncate_content(self.output)
+        return f"{log_prefix} \033[32m[Tool Result] {truncated}\033[0m"
+
+
+class TurnCompletedEvent(CodexEvent):
+    """Per-turn usage summary emitted by Codex."""
+
+    def __init__(
+        self,
+        input_tokens: int = 0,
+        cached_input_tokens: int = 0,
+        output_tokens: int = 0,
+    ):
+        self.input_tokens = input_tokens
+        self.cached_input_tokens = cached_input_tokens
+        self.output_tokens = output_tokens
+
+    def render(self, log_prefix: str) -> str | None:
+        return None
+
+
+class ErrorEvent(CodexEvent):
+    """Error event emitted on turn failure or top-level error."""
+
+    def __init__(self, message: str):
+        self.message = message
+
+    def render(self, log_prefix: str) -> str:
+        return f"{log_prefix} \033[31m[Error] {self.message}\033[0m"

agentshim/codex_events.py

@@ -0,0 +1,23 @@
+"""Compatibility shim for older ``agentshim.codex_events`` imports."""
+
+from .codex.events import (
+    CodexEvent,
+    ErrorEvent,
+    LifecycleEvent,
+    TextEvent,
+    ThreadStartedEvent,
+    ToolResultEvent,
+    ToolUseEvent,
+    TurnCompletedEvent,
+)
+
+__all__ = [
+    "CodexEvent",
+    "ErrorEvent",
+    "LifecycleEvent",
+    "TextEvent",
+    "ThreadStartedEvent",
+    "ToolResultEvent",
+    "ToolUseEvent",
+    "TurnCompletedEvent",
+]

agentshim/events.py

@@ -0,0 +1,24 @@
+from typing import Any, Protocol
+
+
+class AgentEventHandler(Protocol):
+    """Protocol for handling agent events."""
+
+    def on_thinking(self, text: str) -> None:
+        """Handle agent thinking output."""
+        ...
+
+    def on_tool_call(self, tool: str, args: dict[str, Any] | str | None = None) -> None:
+        """Handle tool execution start."""
+        ...
+
+    def on_tool_result(
+        self,
+        tool: str,
+        stdout: str = "",
+        stderr: str = "",
+        exit_code: int | None = None,
+        duration: float | None = None,
+    ) -> None:
+        """Handle tool execution result."""
+        ...

agentshim/gemini/__init__.py

@@ -0,0 +1,3 @@
+from .agent import GeminiCodingAgent, GeminiGenerationSession
+
+__all__ = ["GeminiCodingAgent", "GeminiGenerationSession"]

agentshim/gemini/agent.py

@@ -0,0 +1,234 @@
+import json
+import logging
+import subprocess
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import agentshim.trajectory as _trajectory_module
+from agentshim.trajectory import TrajectoryRecorderProtocol
+
+from ..base import register_provider
+from ..cli_agent import CLICodingAgent, CLIGenerationSession
+from ..events import AgentEventHandler
+from ..sandbox import SandboxConfig
+from ..usage import ProviderUsage, TokenUsage
+from .events import GeminiEvent, InitEvent, MessageEvent, ToolResultEvent, ToolUseEvent
+
+_logger = logging.getLogger(__name__)
+
+
+class GeminiGenerationSession(CLIGenerationSession):
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        # Initialize state required for stream processing
+        self.tool_map: dict[str, str] = {}
+        self.tool_start_times: dict[str, float] = {}
+        self.tool_args: dict[str, Any] = {}
+        # Capture call_id and run_id for correlation
+        self.call_id = _trajectory_module.get_current_call_id()
+        self.run_id = _trajectory_module.get_run_id()
+        # Gemini's stream-json does not emit token usage. We count
+        # assistant MessageEvents as a turn proxy; token fields stay 0.
+        self._assistant_message_count: int = 0
+
+    def _write_call_metadata(self):
+        """Write metadata file to help correlate Gemini session with trajectory call."""
+        if self.call_id is None or self.run_id is None:
+            return
+
+        try:
+            gemini_tmp_dir = Path.home() / ".gemini" / "tmp"
+            if not gemini_tmp_dir.exists():
+                return
+
+            if self.cwd:
+                project_name = Path(self.cwd).name
+                project_dir = gemini_tmp_dir / project_name / "chats"
+                if project_dir.exists():
+                    metadata_file = project_dir / f"sds_call_{self.call_id:03d}.json"
+                    metadata = {
+                        "call_id": self.call_id,
+                        "run_id": self.run_id,
+                        "start_time": time.strftime("%Y-%m-%d %H:%M:%S"),
+                        "cwd": self.cwd,
+                    }
+                    with open(metadata_file, "w") as f:
+                        json.dump(metadata, f, indent=2)
+        except Exception:
+            _logger.debug("Failed to write Gemini call metadata", exc_info=True)
+
+    def run(self, prompt: str) -> str:
+        """Execute the generation process, writing call metadata first."""
+        self._write_call_metadata()
+
+        if not self.silent:
+            self._log_raw(f"{self.log_prefix} Input Prompt:\n{prompt}\n")
+
+        return super().run(prompt)
+
+    def _process_stdout(self, line: str) -> None:
+        """Process a line from stdout."""
+        if not line:
+            return
+        try:
+            data = json.loads(line)
+            event = GeminiEvent.from_dict(data)
+            if event:
+                self._handle_event(event)
+        except json.JSONDecodeError:
+            if not self.silent:
+                if self._at_line_start:
+                    self._log_raw(f"{self.log_prefix} ")
+                self._log_raw(line.rstrip() + "\n")
+                self._at_line_start = True
+
+    def _handle_event(self, event: GeminiEvent):
+        """Handle a single parsed Gemini event."""
+        self._update_state(event)
+
+        if not self.silent:
+            self._render_event(event)
+
+    def _update_state(self, event: GeminiEvent):
+        """Update internal state based on the event."""
+        if isinstance(event, InitEvent):
+            if self.session_id is None and event.session_id:
+                self.session_id = event.session_id
+            return
+
+        if isinstance(event, MessageEvent):
+            if event.role == "assistant":
+                self.stdout_lines.append(event.content)
+                self._assistant_message_count += 1
+                self.usage = ProviderUsage(
+                    tokens=TokenUsage(turns=self._assistant_message_count),
+                    provider="gemini",
+                )
+                if self.event_handler:
+                    self.event_handler.on_thinking(event.content)
+
+        elif isinstance(event, ToolUseEvent):
+            if event.tool_id:
+                self.tool_map[event.tool_id] = event.tool_name
+                self.tool_start_times[event.tool_id] = time.time()
+                self.tool_args[event.tool_id] = event.parameters
+                if self.event_handler:
+                    self.event_handler.on_tool_call(event.tool_name, event.parameters)
+
+        elif isinstance(event, ToolResultEvent) and event.tool_id:
+            event.tool_name_resolved = self.tool_map.get(event.tool_id, "Tool")
+
+            start_time = self.tool_start_times.get(event.tool_id)
+            duration = time.time() - start_time if start_time else None
+            args = self.tool_args.get(event.tool_id, {})
+
+            self.recorder.add_tool_call(
+                tool=event.tool_name_resolved,
+                args=args,
+                stdout=event.output,
+                duration=duration,
+            )
+
+            if self.event_handler:
+                self.event_handler.on_tool_result(
+                    tool=event.tool_name_resolved,
+                    stdout=event.output,
+                    duration=duration,
+                )
+
+    def _render_event(self, event: GeminiEvent):
+        """Render the event to stdout."""
+        if isinstance(event, MessageEvent):
+            if event.role == "assistant":
+                self._print_stream_content(event.content)
+            return
+
+        if not self._at_line_start:
+            self._log_raw("\n")
+            self._at_line_start = True
+
+        output = event.render(self.log_prefix)
+        if output:
+            self._log_raw(output + "\n")
+
+
+@register_provider("gemini")
+class GeminiCodingAgent(CLICodingAgent):
+    """Coding agent implementation using the Gemini CLI tool."""
+
+    def __init__(
+        self,
+        model: str | None = None,
+        recorder: TrajectoryRecorderProtocol | None = None,
+        event_handler: AgentEventHandler | None = None,
+        mcp_servers: list[object] | None = None,
+        sandbox: bool | SandboxConfig = False,
+    ):
+        """Initialize the Gemini coding agent.
+
+        Args:
+            model: Optional model name to use.
+            recorder: Trajectory recorder instance.
+            event_handler: Optional event handler for UI updates.
+            mcp_servers: Optional list of MCP server configurations.
+            sandbox: Not supported for Gemini; must be False.
+
+        Raises:
+            ValueError: If mcp_servers is non-empty (not supported).
+            NotImplementedError: If ``sandbox`` is truthy.
+        """
+        if mcp_servers:
+            raise ValueError("GeminiCodingAgent does not support programmatic MCP server configuration via CLI flags")
+        if sandbox:
+            raise NotImplementedError("sandbox is not supported for GeminiCodingAgent")
+        super().__init__("gemini", model, recorder, event_handler)
+
+    @property
+    def gemini_path(self) -> str:
+        """Return path to gemini binary (for backward compatibility)."""
+        return self.binary_path
+
+    @property
+    def _log_prefix(self) -> str:
+        """Return the log prefix for this agent."""
+        return "[Gemini]"
+
+    def _get_command(self, prompt: str, resume_session_id: str | None = None) -> list[str]:
+        cmd = [self.binary_path]
+
+        cmd.extend(["-y"])
+
+        if self.model:
+            cmd.extend(["--model", self.model])
+
+        cmd.extend(["-o", "stream-json"])
+
+        if resume_session_id:
+            cmd.extend(["--resume", resume_session_id])
+
+        return cmd
+
+    def _create_session(
+        self,
+        cmd: list[str],
+        cwd: str | None = None,
+        timeout: int = 300,
+        silent: bool = False,
+        recorder: TrajectoryRecorderProtocol | None = None,
+        on_process_started: Callable[[subprocess.Popen[str]], None] | None = None,
+    ) -> GeminiGenerationSession:
+        return GeminiGenerationSession(
+            binary_name=self.binary_name,
+            env=self.env,
+            log_prefix=self._log_prefix,
+            cmd=cmd,
+            logger=self.logger,
+            cwd=cwd,
+            timeout=timeout,
+            silent=silent,
+            recorder=recorder,
+            event_handler=self.event_handler,
+            on_process_started=on_process_started,
+        )

agentshim/gemini/events.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from ..utils import truncate_content, truncate_tool_params
+
+
+class GeminiEvent(ABC):
+    """Base class for Gemini stream events."""
+
+    @abstractmethod
+    def render(self, log_prefix: str) -> str | None:
+        """Render the event as a string for terminal output."""
+
+    @staticmethod
+    def from_dict(data: dict[str, Any]) -> GeminiEvent | None:
+        """Factory method to create events from JSON data."""
+        msg_type = data.get("type")
+
+        if msg_type == "init":
+            return InitEvent(session_id=data.get("session_id"))
+        if msg_type == "message":
+            return MessageEvent(role=data.get("role", ""), content=data.get("content", ""))
+        if msg_type == "tool_use":
+            return ToolUseEvent(
+                tool_name=data.get("tool_name", "Tool"),
+                tool_id=data.get("tool_id"),
+                parameters=data.get("parameters"),
+            )
+        if msg_type == "tool_result":
+            return ToolResultEvent(output=data.get("output", ""), tool_id=data.get("tool_id"))
+        return None
+
+
+class InitEvent(GeminiEvent):
+    """Session initialization event; carries the resumable ``session_id``."""
+
+    def __init__(self, session_id: str | None):
+        self.session_id = session_id
+
+    def render(self, log_prefix: str) -> str | None:
+        return None
+
+
+class MessageEvent(GeminiEvent):
+    def __init__(self, role: str, content: str):
+        self.role = role
+        self.content = content
+
+    def render(self, log_prefix: str) -> str | None:
+        # Message rendering is handled specially due to streaming
+        # This is just a placeholder or could handle non-streaming blocks
+        return self.content
+
+
+class ToolUseEvent(GeminiEvent):
+    def __init__(self, tool_name: str, tool_id: str | None, parameters: Any):
+        self.tool_name = tool_name
+        self.tool_id = tool_id
+        self.parameters = parameters
+
+    def render(self, log_prefix: str) -> str:
+        truncated = truncate_tool_params(self.tool_name, self.parameters)
+        return f"{log_prefix} \033[34m[Tool Use] {self.tool_name} {truncated}\033[0m"
+
+
+class ToolResultEvent(GeminiEvent):
+    def __init__(self, output: str, tool_id: str | None):
+        self.output = output
+        self.tool_id = tool_id
+        self.tool_name_resolved: str = "Tool"  # To be set externally
+
+    def render(self, log_prefix: str) -> str:
+        if not self.output:
+            return f"{log_prefix} \033[32m{self.tool_name_resolved} ran successfully\033[0m"
+        truncated = truncate_content(self.output)
+        return f"{log_prefix} \033[32m[Tool Result] {truncated}\033[0m"

agentshim/gemini_events.py

@@ -0,0 +1,11 @@
+"""Compatibility shim for older ``agentshim.gemini_events`` imports."""
+
+from .gemini.events import GeminiEvent, InitEvent, MessageEvent, ToolResultEvent, ToolUseEvent
+
+__all__ = [
+    "GeminiEvent",
+    "InitEvent",
+    "MessageEvent",
+    "ToolResultEvent",
+    "ToolUseEvent",
+]

agentshim/llm_client.py

@@ -0,0 +1,64 @@
+"""Thin litellm wrapper with automatic token tracking and trajectory recording."""
+
+import os
+from typing import Any
+
+from agentshim.subagent import litellm_call_with_retry
+from agentshim.trajectory import TrajectoryRecorderProtocol
+
+
+class LiteLLMClient:
+    """litellm completion wrapper that accumulates token usage and records to trajectory.
+
+    Every ``complete()`` call:
+    - Delegates to ``litellm_call_with_retry`` for retry-on-network-error logic.
+    - Accumulates prompt/completion/total tokens in ``_token_usage``.
+    - Calls ``recorder.record_token_usage()`` with the running total.
+
+    The caller is responsible for managing conversation history (``messages``).
+    This avoids coupling: pass the full list each call, append results yourself.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        location: str | None = None,
+        recorder: TrajectoryRecorderProtocol | None = None,
+    ):
+        self.model = model
+        self.location = location
+        self.recorder = recorder
+        self._token_usage: dict[str, int] = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+        }
+
+    def complete(self, messages: list[dict[str, Any]], label: str = "llm call") -> str:
+        """Call litellm with *messages*, accumulate tokens, and record to trajectory.
+
+        Args:
+            messages: Full conversation history to send (caller-managed).
+            label: Human-readable label for retry/logging messages.
+
+        Returns:
+            Assistant response text.
+
+        Raises:
+            Exception: Any non-retried litellm error propagates to the caller.
+        """
+        kwargs: dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            "cache": {"no-cache": True},
+        }
+        loc = self.location or os.environ.get("VERTEX_LOCATION")
+        if loc:
+            kwargs["vertex_location"] = loc
+
+        result = litellm_call_with_retry(kwargs, label=label, token_acc=self._token_usage)
+
+        if self.recorder and hasattr(self.recorder, "record_token_usage"):
+            self.recorder.record_token_usage(self._token_usage.copy())  # type: ignore[reportArgumentType]
+
+        return result

agentshim/mcp_config.py

@@ -0,0 +1,25 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class HttpMcpServer(BaseModel):
+    """MCP server accessed over HTTP/SSE."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str = Field(min_length=1)
+    url: str = Field(min_length=1)
+    headers: dict[str, str] = Field(default_factory=dict)
+
+
+class StdioMcpServer(BaseModel):
+    """MCP server launched as a subprocess (stdio transport)."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str = Field(min_length=1)
+    command: str = Field(min_length=1)
+    args: list[str] = Field(default_factory=list)
+    env: dict[str, str] = Field(default_factory=dict)
+
+
+McpServerConfig = HttpMcpServer | StdioMcpServer

agentshim/opencode/__init__.py

@@ -0,0 +1,3 @@
+from .agent import OpencodeCodingAgent, OpencodeGenerationSession
+
+__all__ = ["OpencodeCodingAgent", "OpencodeGenerationSession"]