93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
from pydantic import BaseModel, Field
|
|
from pydantic_ai import Agent, RunContext
|
|
from pydantic_ai.models.openai import OpenAIChatModel
|
|
from pydantic_ai.providers.ollama import OllamaProvider
|
|
from pydantic_ai.providers.openai import OpenAIProvider
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
|
|
|
|
|
|
class SearchResult(BaseModel):
|
|
content: str = Field(description="The document text content")
|
|
score: float = Field(description="Relevance score (higher is more relevant)")
|
|
document_uri: str = Field(
|
|
description="Source title (if available) or URI/path of the document"
|
|
)
|
|
|
|
|
|
class Dependencies(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
client: HaikuRAG
|
|
|
|
|
|
class QuestionAnswerAgent:
|
|
def __init__(
|
|
self,
|
|
client: HaikuRAG,
|
|
provider: str,
|
|
model: str,
|
|
use_citations: bool = False,
|
|
q: float = 0.0,
|
|
system_prompt: str | None = None,
|
|
):
|
|
self._client = client
|
|
|
|
if system_prompt is None:
|
|
system_prompt = (
|
|
QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
|
|
)
|
|
model_obj = self._get_model(provider, model)
|
|
|
|
self._agent = Agent(
|
|
model=model_obj,
|
|
deps_type=Dependencies,
|
|
system_prompt=system_prompt,
|
|
retries=3,
|
|
)
|
|
|
|
@self._agent.tool
|
|
async def search_documents(
|
|
ctx: RunContext[Dependencies],
|
|
query: str,
|
|
limit: int = 3,
|
|
) -> list[SearchResult]:
|
|
"""Search the knowledge base for relevant documents."""
|
|
search_results = await ctx.deps.client.search(query, limit=limit)
|
|
expanded_results = await ctx.deps.client.expand_context(search_results)
|
|
|
|
return [
|
|
SearchResult(
|
|
content=chunk.content,
|
|
score=score,
|
|
document_uri=(chunk.document_title or chunk.document_uri or ""),
|
|
)
|
|
for chunk, score in expanded_results
|
|
]
|
|
|
|
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.providers.ollama.base_url}/v1"
|
|
),
|
|
)
|
|
elif provider == "vllm":
|
|
return OpenAIChatModel(
|
|
model_name=model,
|
|
provider=OpenAIProvider(
|
|
base_url=f"{Config.providers.vllm.qa_base_url}/v1", api_key="none"
|
|
),
|
|
)
|
|
else:
|
|
# For all other providers, use the provider:model format
|
|
return f"{provider}:{model}"
|
|
|
|
async def answer(self, question: str) -> str:
|
|
"""Answer a question using the RAG system."""
|
|
deps = Dependencies(client=self._client)
|
|
result = await self._agent.run(question, deps=deps)
|
|
return result.output
|