124 lines
4 KiB
Python
124 lines
4 KiB
Python
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel
|
|
from pydantic_ai import Agent, RunContext
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.utils import get_model
|
|
|
|
if TYPE_CHECKING:
|
|
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
|
|
|
|
|
class ChatSessionState(BaseModel):
|
|
"""State shared between frontend and agent via AG-UI."""
|
|
|
|
session_id: str = ""
|
|
|
|
|
|
@dataclass
|
|
class ChatDeps:
|
|
"""Dependencies for chat agent."""
|
|
|
|
client: HaikuRAG
|
|
config: AppConfig
|
|
agui_emitter: "AGUIEmitter | None" = None
|
|
search_results: list[SearchResult] | None = None
|
|
|
|
|
|
CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
|
|
|
|
You have access to a knowledge base of documents. Use your tools to search and answer questions.
|
|
|
|
CRITICAL RULES:
|
|
1. For greetings or casual chat: respond directly WITHOUT using any tools
|
|
2. For substantive questions requiring information: use the search or ask tools
|
|
3. NEVER make up information - always use tools to get facts from the knowledge base
|
|
4. When citing sources, reference the chunk IDs from search results
|
|
|
|
How to decide which tool to use:
|
|
- "search" - When you need to find relevant documents or explore what's in the knowledge base
|
|
- "ask" - When you have a specific question that needs a direct answer with citations
|
|
|
|
Be friendly and conversational. When you use tools, summarize the key findings for the user."""
|
|
|
|
|
|
def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|
"""Create the chat agent with search and ask tools."""
|
|
model = get_model(config.qa.model, config)
|
|
|
|
agent: Agent[ChatDeps, str] = Agent(
|
|
model,
|
|
deps_type=ChatDeps,
|
|
output_type=str,
|
|
instructions=CHAT_SYSTEM_PROMPT,
|
|
)
|
|
|
|
@agent.tool
|
|
async def search(
|
|
ctx: RunContext[ChatDeps],
|
|
query: str,
|
|
limit: int = 5,
|
|
document_filter: str | None = None,
|
|
) -> str:
|
|
"""Search the knowledge base for relevant documents.
|
|
|
|
Use this when you need to find documents or explore the knowledge base.
|
|
Returns relevant chunks with metadata.
|
|
|
|
Args:
|
|
query: The search query
|
|
limit: Maximum number of results (default 5)
|
|
document_filter: Optional SQL WHERE clause to filter documents (e.g. "id IN ('doc1', 'doc2')")
|
|
"""
|
|
if ctx.deps.agui_emitter:
|
|
ctx.deps.agui_emitter.log(f"Searching: {query}")
|
|
|
|
results = await ctx.deps.client.search(
|
|
query, limit=limit, filter=document_filter
|
|
)
|
|
results = await ctx.deps.client.expand_context(results)
|
|
|
|
# Store for potential citation resolution
|
|
ctx.deps.search_results = results
|
|
|
|
if not results:
|
|
return "No results found for your query."
|
|
|
|
# Format results for the agent
|
|
parts = [r.format_for_agent() for r in results]
|
|
return "\n\n".join(parts)
|
|
|
|
@agent.tool
|
|
async def ask(
|
|
ctx: RunContext[ChatDeps],
|
|
question: str,
|
|
document_filter: str | None = None,
|
|
) -> str:
|
|
"""Answer a specific question using the knowledge base.
|
|
|
|
Use this for direct questions that need a focused answer with citations.
|
|
|
|
Args:
|
|
question: The question to answer
|
|
document_filter: Optional SQL WHERE clause to filter documents (e.g. "id IN ('doc1', 'doc2')")
|
|
"""
|
|
if ctx.deps.agui_emitter:
|
|
ctx.deps.agui_emitter.log(f"Answering: {question}")
|
|
|
|
answer, citations = await ctx.deps.client.ask(question, filter=document_filter)
|
|
|
|
# Format answer with citations
|
|
if citations:
|
|
citation_list = "\n".join(
|
|
f" [{i + 1}] {c.document_uri or c.document_title or 'Unknown'} (chunk: {c.chunk_id})"
|
|
for i, c in enumerate(citations)
|
|
)
|
|
return f"{answer}\n\nSources:\n{citation_list}"
|
|
|
|
return answer
|
|
|
|
return agent
|