Agent API¶

class F1Agent:
    """AI agent for F1 data analysis using Claude Agent SDK."""

    def __init__(
        self,
        workspace_id: str | None = None,
        workspace_dir: Path | None = None,
        enable_tracing: bool | None = None,
        inject_temporal_context: bool = True,
        sandbox_enabled: bool = True,
    ):
        """Initialize the F1 agent.

        Args:
            workspace_id: Workspace identifier. Auto-generated if None.
            workspace_dir: Explicit workspace path. Derived from workspace_id if None.
            enable_tracing: Enable OpenTelemetry tracing. If None, uses PITLANE_TRACING_ENABLED env var.
            inject_temporal_context: Enable temporal context in system prompt. Default True.
            sandbox_enabled: Enable OS-level bash sandboxing. Default True.
        """
        self.workspace_id = workspace_id or generate_workspace_id()
        self.workspace_dir = workspace_dir or get_workspace_path(self.workspace_id)
        self.inject_temporal_context = inject_temporal_context
        self.sandbox_enabled = sandbox_enabled
        self._agent_session_id: str | None = None  # Captured from Claude SDK

        # Verify workspace exists or create it
        if not self.workspace_dir.exists():
            create_workspace(self.workspace_id)
        elif workspace_exists(self.workspace_id):
            # Update last accessed timestamp
            update_workspace_metadata(self.workspace_id)

        # Keep matplotlib writes within ~/.pitlane/ for sandbox compatibility
        mpl_cache = Path.home() / ".pitlane" / "cache" / "matplotlib"
        mpl_cache.mkdir(parents=True, exist_ok=True)
        os.environ["MPLCONFIGDIR"] = str(mpl_cache)

        # Configure tracing
        if enable_tracing is not None:
            if enable_tracing:
                tracing.enable_tracing()
            else:
                tracing.disable_tracing()

    @property
    def charts_dir(self) -> Path:
        """Get the charts directory for this workspace.

        Backward-compatible property for accessing the charts directory.

        Returns:
            Path to the workspace charts directory.
        """
        return self.workspace_dir / "charts"

    @property
    def agent_session_id(self) -> str | None:
        """Get the Claude Agent SDK session ID for resumption.

        This ID is captured from the SDK's init message during chat()
        and can be used to resume the conversation later.

        Returns:
            The SDK session ID, or None if not yet captured.
        """
        return self._agent_session_id

    def _build_system_prompt(self) -> str | None:
        """Build the system prompt append string with optional temporal context."""
        if not self.inject_temporal_context:
            return None
        try:
            temporal_ctx = get_temporal_context()
            return format_for_system_prompt(temporal_ctx, verbosity="normal")
        except Exception:
            return None

    async def chat(self, message: str, resume_session_id: str | None = None) -> AsyncIterator[str]:
        """Process a chat message and yield response text chunks.

        Args:
            message: The user's question or message.
            resume_session_id: Optional SDK session ID to resume a previous conversation.

        Yields:
            Text chunks from the assistant's response.
        """
        # Set workspace ID as environment variable so skills can access it
        os.environ["PITLANE_WORKSPACE_ID"] = self.workspace_id

        workspace_dir = str(self.workspace_dir)

        skills_dir = str(PACKAGE_DIR / ".claude")

        # PreToolUse is always registered — it's the only mechanism that can block
        # tools listed in allowed_tools (the SDK skips can_use_tool for those).
        hooks: dict[str, list[HookMatcher]] = {
            "PreToolUse": [
                HookMatcher(
                    matcher=None,
                    hooks=[make_pre_tool_use_hook(workspace_dir, self.workspace_id, skills_dir, self.sandbox_enabled)],
                )
            ],
        }
        if tracing.is_tracing_enabled():
            hooks["PostToolUse"] = [HookMatcher(matcher=None, hooks=[tracing.post_tool_use_hook])]

        system_prompt_append = self._build_system_prompt()

        options = ClaudeAgentOptions(
            cwd=str(PACKAGE_DIR),
            setting_sources=["project"],
            allowed_tools=["Skill", "Bash", "Read", "Write", "WebFetch", "WebSearch"],
            # can_use_tool handles tools NOT in allowed_tools; PreToolUse hook above
            # handles the full list (SDK skips can_use_tool for allowed_tools).
            can_use_tool=make_can_use_tool_callback(workspace_dir, self.workspace_id, skills_dir, self.sandbox_enabled),
            hooks=hooks,
            resume=resume_session_id,
            system_prompt={
                "type": "preset",
                "preset": "claude_code",
                "append": system_prompt_append,
            }
            if system_prompt_append
            else None,
            sandbox=SandboxSettings(enabled=True) if self.sandbox_enabled else None,
        )

        async with ClaudeSDKClient(options=options) as client:
            await client.query(message)

            async for msg in client.receive_response():
                # Capture session ID from init message for future resumption
                if isinstance(msg, SystemMessage):
                    session_id = msg.data.get("session_id")
                    if session_id:
                        self._agent_session_id = session_id
                        logger.debug(f"Captured SDK session ID: {session_id}")
                elif isinstance(msg, AssistantMessage):
                    for block in msg.content:
                        if isinstance(block, TextBlock):
                            yield block.text

    async def chat_full(self, message: str, resume_session_id: str | None = None) -> str:
        """Process a chat message and return the full response.

        Args:
            message: The user's question or message.
            resume_session_id: Optional SDK session ID to resume a previous conversation.

        Returns:
            The complete response text.
        """
        parts = []
        async for chunk in self.chat(message, resume_session_id=resume_session_id):
            parts.append(chunk)
        return "\n".join(parts) if parts else ""