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

gobby/llm/codex_executor.py

@@ -1,30 +1,19 @@
 """
-Codex (OpenAI) implementation of AgentExecutor.
+Codex (OpenAI) implementation of AgentExecutor for CLI/subscription mode only.
 
-Supports two authentication modes with different capabilities:
+This executor spawns `codex exec --json` CLI and parses JSONL events.
+It uses Codex's built-in tools (bash, file operations, etc.) - NO custom tool
+injection is supported.
 
-1. api_key mode (OPENAI_API_KEY):
-   - Uses OpenAI API with function calling
-   - Full tool injection support
-   - Requires OPENAI_API_KEY environment variable
-
-2. subscription mode (ChatGPT Plus/Pro/Team/Enterprise):
-   - Spawns `codex exec --json` CLI and parses JSONL events
-   - Uses Codex's built-in tools (bash, file operations, etc.)
-   - NO custom tool injection - tools parameter is IGNORED
-   - Good for delegating complete autonomous tasks
-
-IMPORTANT: These modes have fundamentally different capabilities.
-Use api_key mode if you need custom MCP tool injection.
-Use subscription mode for delegating complete tasks to Codex.
+Note: api_key mode is now routed through LiteLLMExecutor for unified cost tracking.
+Use the resolver.create_executor() function which handles routing automatically.
 """
 
 import asyncio
 import json
 import logging
-import os
 import shutil
-from typing import Any, Literal
+from typing import Literal
 
 from gobby.llm.executor import (
     AgentExecutor,
@@ -36,39 +25,23 @@ from gobby.llm.executor import (
 
 logger = logging.getLogger(__name__)
 
-# Auth mode type
-CodexAuthMode = Literal["api_key", "subscription"]
+# Auth mode type - subscription/cli only, api_key routes through LiteLLM
+CodexAuthMode = Literal["subscription", "cli"]
 
 
 class CodexExecutor(AgentExecutor):
     """
-    Codex (OpenAI) implementation of AgentExecutor.
-
-    Supports two authentication modes with DIFFERENT CAPABILITIES:
-
-    api_key mode:
-        - Uses OpenAI API function calling (like GPT-4)
-        - Full tool injection support via tools parameter
-        - Requires OPENAI_API_KEY environment variable
-        - Standard agentic loop with custom tools
-
-    subscription mode:
-        - Spawns `codex exec --json` CLI process
-        - Parses JSONL events (thread.started, item.completed, turn.completed)
-        - Uses Codex's built-in tools ONLY (bash, file ops, web search, etc.)
-        - The `tools` parameter is IGNORED in this mode
-        - Cannot inject custom MCP tools
-        - Best for delegating complete autonomous tasks
-
-    Example (api_key mode):
-        >>> executor = CodexExecutor(auth_mode="api_key")
-        >>> result = await executor.run(
-        ...     prompt="Create a task",
-        ...     tools=[ToolSchema(name="create_task", ...)],
-        ...     tool_handler=my_handler,
-        ... )
+    Codex (OpenAI) implementation of AgentExecutor for CLI mode only.
+
+    Spawns `codex exec --json` CLI process and parses JSONL events.
+    Uses Codex's built-in tools ONLY (bash, file ops, web search, etc.).
+    The `tools` parameter is IGNORED - cannot inject custom MCP tools.
+    Best for delegating complete autonomous tasks.
+
+    For api_key mode with custom tool injection, use LiteLLMExecutor with
+    provider="codex" which routes through OpenAI API for unified cost tracking.
 
-    Example (subscription mode):
+    Example:
         >>> executor = CodexExecutor(auth_mode="subscription")
         >>> result = await executor.run(
         ...     prompt="Fix the bug in main.py and run the tests",
@@ -77,83 +50,49 @@ class CodexExecutor(AgentExecutor):
         ... )
     """
 
+    _cli_path: str
+
     def __init__(
         self,
-        auth_mode: CodexAuthMode = "api_key",
-        api_key: str | None = None,
+        auth_mode: CodexAuthMode = "subscription",
         default_model: str = "gpt-4o",
     ):
         """
-        Initialize CodexExecutor.
+        Initialize CodexExecutor for CLI/subscription mode.
 
         Args:
-            auth_mode: Authentication mode.
-                - "api_key": Use OpenAI API with function calling (requires OPENAI_API_KEY)
-                - "subscription": Use Codex CLI with ChatGPT subscription (requires `codex` in PATH)
-            api_key: OpenAI API key (optional for api_key mode, uses OPENAI_API_KEY env var).
-            default_model: Default model for api_key mode (default: gpt-4o).
+            auth_mode: Must be "subscription" or "cli". API key mode is handled by LiteLLMExecutor.
+            default_model: Default model (not used in CLI mode, kept for interface compatibility).
+
+        Raises:
+            ValueError: If auth_mode is not "subscription"/"cli" or Codex CLI not found.
         """
+        if auth_mode not in ("subscription", "cli"):
+            raise ValueError(
+                "CodexExecutor only supports subscription/cli mode. "
+                "For api_key mode with custom tools, use LiteLLMExecutor with provider='codex'."
+            )
+
         self.auth_mode = auth_mode
         self.default_model = default_model
         self.logger = logger
-        self._client: Any = None
-        self._cli_path: str = ""
-
-        if auth_mode == "api_key":
-            # Use provided key or fall back to environment variable
-            key = api_key or os.environ.get("OPENAI_API_KEY")
-            if not key:
-                raise ValueError(
-                    "API key required for api_key mode. "
-                    "Provide api_key parameter or set OPENAI_API_KEY env var."
-                )
-            try:
-                from openai import AsyncOpenAI
-
-                self._client = AsyncOpenAI(api_key=key)
-                self.logger.debug("CodexExecutor initialized with API key")
-            except ImportError as e:
-                raise ImportError(
-                    "openai package not found. Please install with `pip install openai`."
-                ) from e
-
-        elif auth_mode == "subscription":
-            # Verify Codex CLI is available
-            cli_path = shutil.which("codex")
-            if not cli_path:
-                raise ValueError(
-                    "Codex CLI not found in PATH. "
-                    "Install Codex CLI and run `codex login` for subscription mode."
-                )
-            self._cli_path = cli_path
-            self.logger.debug(f"CodexExecutor initialized with CLI at {cli_path}")
-
-        else:
-            raise ValueError(f"Unknown auth_mode: {auth_mode}")
+        self._cli_path = ""
+
+        # Verify Codex CLI is available
+        cli_path = shutil.which("codex")
+        if not cli_path:
+            raise ValueError(
+                "Codex CLI not found in PATH. "
+                "Install Codex CLI and run `codex login` for subscription mode."
+            )
+        self._cli_path = cli_path
+        self.logger.debug(f"CodexExecutor initialized with CLI at {cli_path}")
 
     @property
     def provider_name(self) -> str:
         """Return the provider name."""
         return "codex"
 
-    def _convert_tools_to_openai_format(self, tools: list[ToolSchema]) -> list[dict[str, Any]]:
-        """Convert ToolSchema list to OpenAI function calling format."""
-        openai_tools = []
-        for tool in tools:
-            # Ensure input_schema has "type": "object"
-            params = {"type": "object", **tool.input_schema}
-            openai_tools.append(
-                {
-                    "type": "function",
-                    "function": {
-                        "name": tool.name,
-                        "description": tool.description,
-                        "parameters": params,
-                    },
-                }
-            )
-        return openai_tools
-
     async def run(
         self,
         prompt: str,
@@ -165,200 +104,29 @@ class CodexExecutor(AgentExecutor):
         timeout: float = 120.0,
     ) -> AgentResult:
         """
-        Execute an agentic loop.
+        Execute Codex CLI and parse JSONL events.
+
+        Note: The tools and tool_handler parameters are IGNORED in CLI mode.
+        Codex uses its own built-in tools (bash, file operations, etc.).
 
-        For api_key mode: Uses OpenAI function calling with custom tools.
-        For subscription mode: Spawns Codex CLI (tools parameter is IGNORED).
+        For custom tool injection, use LiteLLMExecutor with provider="codex".
 
         Args:
             prompt: The user prompt to process.
-            tools: List of available tools (IGNORED in subscription mode).
-            tool_handler: Callback for tool calls (NOT CALLED in subscription mode).
-            system_prompt: Optional system prompt (api_key mode only).
-            model: Optional model override (api_key mode only).
-            max_turns: Maximum turns before stopping (api_key mode only).
+            tools: IGNORED - Codex uses its own tools.
+            tool_handler: IGNORED - not called in CLI mode.
+            system_prompt: IGNORED in CLI mode.
+            model: IGNORED in CLI mode.
+            max_turns: IGNORED in CLI mode.
             timeout: Maximum execution time in seconds.
 
         Returns:
             AgentResult with output, status, and tool call records.
         """
-        if self.auth_mode == "api_key":
-            return await self._run_with_api(
-                prompt=prompt,
-                tools=tools,
-                tool_handler=tool_handler,
-                system_prompt=system_prompt,
-                model=model or self.default_model,
-                max_turns=max_turns,
-                timeout=timeout,
-            )
-        else:
-            return await self._run_with_cli(
-                prompt=prompt,
-                timeout=timeout,
-            )
-
-    async def _run_with_api(
-        self,
-        prompt: str,
-        tools: list[ToolSchema],
-        tool_handler: ToolHandler,
-        system_prompt: str | None,
-        model: str,
-        max_turns: int,
-        timeout: float,
-    ) -> AgentResult:
-        """Run using OpenAI API with function calling."""
-        if self._client is None:
-            return AgentResult(
-                output="",
-                status="error",
-                error="OpenAI client not initialized",
-                turns_used=0,
-            )
-
-        tool_calls_list: list[ToolCallRecord] = []
-        openai_tools = self._convert_tools_to_openai_format(tools)
-
-        # Build initial messages
-        messages: list[dict[str, Any]] = []
-        if system_prompt:
-            messages.append({"role": "system", "content": system_prompt})
-        messages.append({"role": "user", "content": prompt})
-
-        # Track turns in outer scope so timeout handler can access the count
-        turns_counter = [0]
-
-        async def _run_loop() -> AgentResult:
-            nonlocal messages
-            turns_used = 0
-            final_output = ""
-            client = self._client
-            if client is None:
-                raise RuntimeError("CodexExecutor client not initialized")
-
-            while turns_used < max_turns:
-                turns_used += 1
-                turns_counter[0] = turns_used
-
-                # Call OpenAI
-                try:
-                    response = await client.chat.completions.create(
-                        model=model,
-                        messages=messages,
-                        tools=openai_tools if openai_tools else None,
-                        max_tokens=8192,
-                    )
-                except Exception as e:
-                    self.logger.error(f"OpenAI API error: {e}")
-                    return AgentResult(
-                        output="",
-                        status="error",
-                        tool_calls=tool_calls_list,
-                        error=f"OpenAI API error: {e}",
-                        turns_used=turns_used,
-                    )
-
-                # Get the assistant's message
-                choice = response.choices[0]
-                message = choice.message
-
-                # Extract text content
-                if message.content:
-                    final_output = message.content
-
-                # Add assistant message to history
-                messages.append(message.model_dump())
-
-                # Check if there are tool calls
-                if not message.tool_calls:
-                    # No tool calls - we're done
-                    return AgentResult(
-                        output=final_output,
-                        status="success",
-                        tool_calls=tool_calls_list,
-                        turns_used=turns_used,
-                    )
-
-                # Handle tool calls
-                for tool_call in message.tool_calls:
-                    tool_name = tool_call.function.name
-                    try:
-                        arguments = json.loads(tool_call.function.arguments)
-                    except json.JSONDecodeError as e:
-                        self.logger.warning(
-                            f"Failed to parse tool call arguments for '{tool_name}' "
-                            f"(id={getattr(tool_call, 'id', 'unknown')}): {e}. "
-                            f"Arguments: {tool_call.function.arguments!r}"
-                        )
-                        arguments = {}
-
-                    # Record the tool call
-                    record = ToolCallRecord(
-                        tool_name=tool_name,
-                        arguments=arguments,
-                    )
-                    tool_calls_list.append(record)
-
-                    # Execute via handler
-                    try:
-                        result = await tool_handler(tool_name, arguments)
-                        record.result = result
-
-                        # Format result for OpenAI
-                        if result.success:
-                            content = json.dumps(result.result) if result.result else "Success"
-                        else:
-                            content = f"Error: {result.error}"
-
-                    except Exception as e:
-                        self.logger.error(f"Tool handler error for {tool_name}: {e}")
-                        from gobby.llm.executor import ToolResult as TR
-
-                        record.result = TR(
-                            tool_name=tool_name,
-                            success=False,
-                            error=str(e),
-                        )
-                        content = f"Error: {e}"
-
-                    # Add tool result to messages
-                    messages.append(
-                        {
-                            "role": "tool",
-                            "tool_call_id": tool_call.id,
-                            "content": content,
-                        }
-                    )
-
-                # Check finish reason
-                if choice.finish_reason == "stop":
-                    return AgentResult(
-                        output=final_output,
-                        status="success",
-                        tool_calls=tool_calls_list,
-                        turns_used=turns_used,
-                    )
-
-            # Max turns reached
-            return AgentResult(
-                output=final_output,
-                status="partial",
-                tool_calls=tool_calls_list,
-                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_list,
-                error=f"Execution timed out after {timeout}s",
-                turns_used=turns_counter[0],
-            )
+        return await self._run_with_cli(
+            prompt=prompt,
+            timeout=timeout,
+        )
 
     async def _run_with_cli(
         self,

gobby/llm/executor.py

@@ -59,6 +59,23 @@ class ToolCallRecord:
     """Result from the tool execution."""
 
 
+@dataclass
+class CostInfo:
+    """Cost information from an LLM call."""
+
+    prompt_tokens: int = 0
+    """Number of tokens in the prompt."""
+
+    completion_tokens: int = 0
+    """Number of tokens in the completion."""
+
+    total_cost: float = 0.0
+    """Total cost in USD for this call."""
+
+    model: str = ""
+    """Model used for this call (LiteLLM format with prefix)."""
+
+
 @dataclass
 class AgentResult:
     """Result from running an agent to completion."""
@@ -93,6 +110,9 @@ class AgentResult:
     child_session_id: str | None = None
     """ID of the child session created for this agent (set by AgentRunner)."""
 
+    cost_info: CostInfo | None = None
+    """Cost tracking information (populated by LiteLLM executor)."""
+
 
 # Type alias for the tool handler callback
 ToolHandler = Callable[[str, dict[str, Any]], Awaitable[ToolResult]]
@@ -310,6 +330,7 @@ class AgentExecutor(ABC):
         if completion_result is not None:
             completion_result.tool_calls = result.tool_calls
             completion_result.turns_used = result.turns_used
+            completion_result.cost_info = result.cost_info
             return completion_result
 
         # Otherwise, return the raw result (might be timeout or natural end)