webber/webber-api/src/domains/agents/explore/agent.py

"""
Explore Agent implementation using PydanticAI.

Fast codebase exploration with read-only tools.
Uses sanitized Ollama provider for reliable tool calling.
"""
import os
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel

from src.domains.agents.base import BaseAgent, AgentContext, register_agent
from src.domains.agents.explore.prompts import EXPLORE_SYSTEM_PROMPT
from src.ollama.provider import get_ollama_provider
from src.shared.config import get_settings
from src.shared.logging import logged, get_logger, trace_span

logger = get_logger(__name__)


@dataclass
class ExploreContext(AgentContext):
    """
    Context for explore agent tools.

    Passed to all tool functions via RunContext.
    """
    pass


class ExploreAgentImpl(BaseAgent):
    """
    Fast codebase exploration agent.

    Uses glob, grep, read, and bash tools to search and analyze codebases.
    Read-only mode - cannot modify files.
    """

    name = "explore"
    description = "Fast codebase exploration - find files, search content, read code"

    def __init__(self):
        """Initialize the explore agent."""
        self._agent: Agent[ExploreContext, str] | None = None
        self._settings = get_settings()

    def _create_agent(self) -> Agent[ExploreContext, str]:
        """Create the PydanticAI agent with Ollama backend."""
        # Use sanitized Ollama provider to fix content: null issues
        model = OpenAIModel(
            model_name=self._settings.ollama_agent_model,
            provider=get_ollama_provider(),
        )

        agent: Agent[ExploreContext, str] = Agent(
            model=model,
            system_prompt=EXPLORE_SYSTEM_PROMPT,
            deps_type=ExploreContext,
            output_type=str,
            # Mistral Nemo settings:
            # - temperature 0.3 (Nemo needs slightly higher than 0.0)
            # - tool_choice "required" forces tool use
            model_settings={
                "temperature": 0.3,
                "extra_body": {"tool_choice": "required"},
            },
        )

        # Register tools
        self._register_tools(agent)

        return agent

    def _register_tools(self, agent: Agent[ExploreContext, str]) -> None:
        """Register all exploration tools with the agent."""
        from src.domains.agents.explore.tools import register_explore_tools
        register_explore_tools(agent)

    @logged()
    async def run(
        self,
        prompt: str,
        working_dir: str | None = None,
        allowed_paths: list[str] | None = None,
        **kwargs: Any
    ) -> str:
        """
        Run the explore agent with a prompt.

        Args:
            prompt: User query about the codebase
            working_dir: Working directory for exploration
            allowed_paths: Restrict tool access to these paths

        Returns:
            Agent response with findings
        """
        ctx = ExploreContext(
            working_dir=working_dir or os.getcwd(),
            allowed_paths=allowed_paths or self._settings.effective_allowed_paths,
            timeout_seconds=self._settings.tool_timeout_seconds,
        )

        async with trace_span("explore_agent_run"):
            try:
                # Use run() not run_stream() - Ollama has bugs with streaming + tools
                result = await self.agent.run(prompt, deps=ctx)
                return result.output
            except Exception as e:
                logger.exception(f"Explore agent error: {e}")
                raise


    async def run_stream(
        self,
        prompt: str,
        working_dir: str | None = None,
        allowed_paths: list[str] | None = None,
        **kwargs: Any
    ) -> AsyncIterator[str]:
        """
        Run the explore agent with streaming output.

        Yields text chunks as they become available.
        """
        ctx = ExploreContext(
            working_dir=working_dir or os.getcwd(),
            allowed_paths=allowed_paths or self._settings.effective_allowed_paths,
            timeout_seconds=self._settings.tool_timeout_seconds,
        )

        async with trace_span("explore_agent_stream"):
            try:
                async with self.agent.run_stream(prompt, deps=ctx) as result:
                    async for chunk in result.stream_text():
                        yield chunk
            except Exception as e:
                logger.exception(f"Explore agent stream error: {e}")
                raise


# Create and register the singleton instance
explore_agent = ExploreAgentImpl()
register_agent(explore_agent)


async def explore(
    prompt: str,
    working_dir: str | None = None,
    **kwargs: Any
) -> str:
    """Run exploration query."""
    return await explore_agent.run(prompt, working_dir=working_dir, **kwargs)


async def explore_stream(
    prompt: str,
    working_dir: str | None = None,
    **kwargs: Any
) -> AsyncIterator[str]:
    """Run exploration query with streaming."""
    async for chunk in explore_agent.run_stream(prompt, working_dir=working_dir, **kwargs):
        yield chunk