From b568293661f7fde1673dc674ee5a48fe2d662246 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 12 Sep 2025 11:03:27 +0300 Subject: [PATCH] Base research agent & dependencies for research multi-agent --- src/haiku/rag/research/__init__.py | 6 ++ src/haiku/rag/research/base.py | 96 ++++++++++++++++++++++++++ src/haiku/rag/research/dependencies.py | 59 ++++++++++++++++ 3 files changed, 161 insertions(+) create mode 100644 src/haiku/rag/research/__init__.py create mode 100644 src/haiku/rag/research/base.py create mode 100644 src/haiku/rag/research/dependencies.py diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py new file mode 100644 index 00000000..29ac92c0 --- /dev/null +++ b/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"] diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py new file mode 100644 index 00000000..c38961e6 --- /dev/null +++ b/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 diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py new file mode 100644 index 00000000..dc74abee --- /dev/null +++ b/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")