Base research agent & dependencies for research multi-agent

src/haiku/rag/research/__init__.py

@@ -0,0 +1,6 @@
+"""Multi-agent research workflow for advanced RAG queries."""
+
+from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import ResearchDependencies
+
+__all__ = ["ResearchDependencies", "BaseResearchAgent"]

src/haiku/rag/research/base.py

@@ -0,0 +1,96 @@
+"""Base class for research agents with common patterns."""
+
+from abc import ABC, abstractmethod
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from haiku.rag.config import Config
+from haiku.rag.research.dependencies import ResearchDependencies
+
+T = TypeVar("T")
+
+
+class BaseResearchAgent(ABC):
+    """Base class for all research agents."""
+
+    def __init__(
+        self,
+        provider: str,
+        model: str,
+        output_type: type[T] | None = None,
+    ):
+        self.provider = provider
+        self.model = model
+        self.output_type = output_type or str
+
+        model_obj = self._get_model(provider, model)
+
+        self._agent = Agent(
+            model=model_obj,
+            deps_type=ResearchDependencies,
+            output_type=self.output_type,
+            system_prompt=self.get_system_prompt(),
+        )
+
+        # Register tools
+        self.register_tools()
+
+    def _get_model(self, provider: str, model: str):
+        """Get the appropriate model object for the provider."""
+        if provider == "ollama":
+            return OpenAIChatModel(
+                model_name=model,
+                provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+            )
+        elif provider == "vllm":
+            return OpenAIChatModel(
+                model_name=model,
+                provider=OpenAIProvider(
+                    base_url=f"{Config.VLLM_QA_BASE_URL}/v1", api_key="none"
+                ),
+            )
+        else:
+            # For all other providers, use the provider:model format
+            return f"{provider}:{model}"
+
+    @abstractmethod
+    def get_system_prompt(self) -> str:
+        """Return the system prompt for this agent."""
+        pass
+
+    @abstractmethod
+    def register_tools(self) -> None:
+        """Register agent-specific tools."""
+        pass
+
+    async def run(self, prompt: str, deps: ResearchDependencies, **kwargs) -> Any:
+        """Execute the agent."""
+        return await self._agent.run(prompt, deps=deps, **kwargs)
+
+    @property
+    def agent(self) -> Agent[ResearchDependencies, Any]:
+        """Access the underlying Pydantic AI agent."""
+        return self._agent
+
+
+class SearchResult(BaseModel):
+    """Standard search result format."""
+
+    content: str
+    score: float
+    document_uri: str
+    metadata: dict[str, Any] = {}
+
+
+class ResearchOutput(BaseModel):
+    """Standard research output format."""
+
+    summary: str
+    detailed_findings: list[str]
+    sources: list[str]
+    confidence: float

src/haiku/rag/research/dependencies.py

@@ -0,0 +1,59 @@
+"""Shared dependencies for multi-agent research workflow."""
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from haiku.rag.client import HaikuRAG
+
+if TYPE_CHECKING:
+    from haiku.rag.research.base import SearchResult
+
+
+class ResearchContext(BaseModel):
+    """Context shared across research agents."""
+
+    original_question: str = Field(description="The original research question")
+    sub_questions: list[str] = Field(
+        default_factory=list, description="Decomposed sub-questions"
+    )
+    search_results: list[dict[str, Any]] = Field(
+        default_factory=list, description="Accumulated search results"
+    )
+    insights: list[str] = Field(
+        default_factory=list, description="Key insights discovered"
+    )
+    gaps: list[str] = Field(
+        default_factory=list, description="Identified information gaps"
+    )
+    follow_up_questions: list[str] = Field(
+        default_factory=list, description="Generated follow-up questions"
+    )
+
+    def add_search_result(self, query: str, results: list["SearchResult"]) -> None:
+        """Add search results to context."""
+        self.search_results.append(
+            {
+                "query": query,
+                "results": results,
+            }
+        )
+
+    def add_insight(self, insight: str) -> None:
+        """Add a key insight."""
+        if insight not in self.insights:
+            self.insights.append(insight)
+
+    def add_gap(self, gap: str) -> None:
+        """Identify an information gap."""
+        if gap not in self.gaps:
+            self.gaps.append(gap)
+
+
+class ResearchDependencies(BaseModel):
+    """Dependencies for research agents with multi-agent context."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    context: ResearchContext = Field(description="Shared research context")