Simplify citations in qa & graph agents
This commit is contained in:
Yiorgis Gozadinos
parent
16a97cc140
commit
d808c6c425
23 changed files with 287 additions and 328 deletions
|
|
@ -248,7 +248,8 @@ async def run_qa_benchmark(
|
|||
qa = get_qa_agent(rag, system_prompt=system_prompt)
|
||||
|
||||
async def answer_question(question: str) -> str:
|
||||
return await qa.answer(question)
|
||||
answer, _ = await qa.answer(question)
|
||||
return answer
|
||||
|
||||
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from haiku.rag.store.models.document import Document
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.utils import format_bytes
|
||||
from haiku.rag.utils import format_bytes, format_citations
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -291,39 +291,49 @@ class HaikuRAGApp:
|
|||
"""
|
||||
async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
|
||||
try:
|
||||
citations = []
|
||||
if deep:
|
||||
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
|
||||
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
|
||||
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
|
||||
|
||||
graph = build_deep_qa_graph(config=self.config)
|
||||
context = DeepQAContext(
|
||||
original_question=question, use_citations=cite
|
||||
)
|
||||
context = DeepQAContext(original_question=question)
|
||||
state = DeepQAState.from_config(context=context, config=self.config)
|
||||
deps = DeepQADeps(client=self.client)
|
||||
|
||||
if verbose:
|
||||
# Use AG-UI renderer to process and display events
|
||||
from haiku.rag.graph.agui import AGUIConsoleRenderer
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
renderer = AGUIConsoleRenderer(self.console)
|
||||
result_dict = await renderer.render(
|
||||
stream_graph(graph, state, deps)
|
||||
)
|
||||
# Result should be a dict with 'answer' key
|
||||
# Result should be a dict with 'answer' and 'citations' keys
|
||||
answer = result_dict.get("answer", "") if result_dict else ""
|
||||
if cite and result_dict:
|
||||
# Convert dicts to Citation objects
|
||||
raw_citations = result_dict.get("citations", [])
|
||||
citations = [
|
||||
Citation(**c) if isinstance(c, dict) else c
|
||||
for c in raw_citations
|
||||
]
|
||||
else:
|
||||
# Run without rendering events, just get the result
|
||||
result = await graph.run(state=state, deps=deps)
|
||||
answer = result.answer
|
||||
if cite:
|
||||
citations = result.citations
|
||||
else:
|
||||
answer = await self.client.ask(question, cite=cite)
|
||||
answer, citations = await self.client.ask(question)
|
||||
|
||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||
self.console.print()
|
||||
self.console.print("[bold green]Answer:[/bold green]")
|
||||
self.console.print(Markdown(answer))
|
||||
if cite and citations:
|
||||
self.console.print(Markdown(format_citations(citations)))
|
||||
except Exception as e:
|
||||
self.console.print(f"[red]Error: {e}[/red]")
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator
|
|||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
|
@ -22,6 +23,9 @@ from haiku.rag.store.repositories.chunk import ChunkRepository
|
|||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
|
@ -850,23 +854,20 @@ class HaikuRAG:
|
|||
return final_results + passthrough
|
||||
|
||||
async def ask(
|
||||
self, question: str, cite: bool = False, system_prompt: str | None = None
|
||||
) -> str:
|
||||
self, question: str, system_prompt: str | None = None
|
||||
) -> "tuple[str, list[Citation]]":
|
||||
"""Ask a question using the configured QA agent.
|
||||
|
||||
Args:
|
||||
question: The question to ask.
|
||||
cite: Whether to include citations in the response.
|
||||
system_prompt: Optional custom system prompt for the QA agent.
|
||||
|
||||
Returns:
|
||||
The generated answer as a string.
|
||||
Tuple of (answer text, list of resolved citations).
|
||||
"""
|
||||
from haiku.rag.qa import get_qa_agent
|
||||
|
||||
qa_agent = get_qa_agent(
|
||||
self, config=self._config, use_citations=cite, system_prompt=system_prompt
|
||||
)
|
||||
qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
|
||||
return await qa_agent.answer(question)
|
||||
|
||||
async def rebuild_database(
|
||||
|
|
|
|||
|
|
@ -210,8 +210,7 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
|
|||
messages = input_state.get("messages", [])
|
||||
if messages:
|
||||
question = messages[0].get("content", "")
|
||||
use_citations = input_state.get("use_citations", False)
|
||||
context = DeepQAContext(original_question=question, use_citations=use_citations)
|
||||
context = DeepQAContext(original_question=question)
|
||||
return DeepQAState.from_config(context=context, config=config)
|
||||
|
||||
def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,14 @@
|
|||
"""Common models used across different graph implementations."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
from haiku.rag.store.models import BoundingBox
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class ResearchPlan(BaseModel):
|
||||
"""A structured research plan with sub-questions to explore."""
|
||||
|
|
@ -21,18 +28,25 @@ class ResearchPlan(BaseModel):
|
|||
return v
|
||||
|
||||
|
||||
class Citation(BaseModel):
|
||||
"""Resolved citation with full metadata for display/visual grounding."""
|
||||
|
||||
document_uri: str
|
||||
document_title: str | None = None
|
||||
page_numbers: list[int] = Field(default_factory=list)
|
||||
headings: list[str] | None = None
|
||||
content: str
|
||||
bounding_boxes: list[BoundingBox] | None = None
|
||||
|
||||
|
||||
class SearchAnswer(BaseModel):
|
||||
"""Answer from a search operation with sources."""
|
||||
"""Structured answer from a search operation."""
|
||||
|
||||
query: str = Field(..., description="The question that was answered")
|
||||
answer: str = Field(..., description="The comprehensive answer to the question")
|
||||
context: list[str] = Field(
|
||||
answer: str = Field(..., description="The answer to the question")
|
||||
cited_chunks: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Relevant snippets that directly support the answer",
|
||||
)
|
||||
sources: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Source URIs or titles that contributed to this answer",
|
||||
description="IDs of chunks used to form the answer",
|
||||
)
|
||||
confidence: float = Field(
|
||||
default=1.0,
|
||||
|
|
@ -40,3 +54,29 @@ class SearchAnswer(BaseModel):
|
|||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
|
||||
|
||||
def resolve_citations(
|
||||
cited_chunk_ids: list[str],
|
||||
search_results: "list[SearchResult]",
|
||||
) -> list[Citation]:
|
||||
"""Resolve chunk IDs to full Citation objects with metadata."""
|
||||
# Build lookup by chunk_id
|
||||
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
|
||||
|
||||
citations = []
|
||||
for chunk_id in cited_chunk_ids:
|
||||
r = by_id.get(chunk_id)
|
||||
if not r:
|
||||
continue
|
||||
citations.append(
|
||||
Citation(
|
||||
document_uri=r.document_uri or "",
|
||||
document_title=r.document_title,
|
||||
page_numbers=r.page_numbers,
|
||||
headings=r.headings,
|
||||
content=r.content,
|
||||
bounding_boxes=r.bounding_boxes,
|
||||
)
|
||||
)
|
||||
return citations
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ from collections.abc import Awaitable, Callable
|
|||
from typing import Any, Protocol
|
||||
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.format_prompt import format_as_xml
|
||||
from pydantic_ai.output import ToolOutput
|
||||
from pydantic_graph.beta import StepContext
|
||||
|
||||
|
|
@ -16,6 +15,7 @@ from haiku.rag.graph.agui.emitter import AGUIEmitter
|
|||
from haiku.rag.graph.common import get_model
|
||||
from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
|
||||
from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class GraphContext(Protocol):
|
||||
|
|
@ -24,7 +24,9 @@ class GraphContext(Protocol):
|
|||
original_question: str
|
||||
sub_questions: list[str]
|
||||
|
||||
def add_qa_response(self, qa: SearchAnswer) -> None:
|
||||
def add_qa_response(
|
||||
self, qa: SearchAnswer, search_results: list[SearchResult]
|
||||
) -> None:
|
||||
"""Add a QA response to context."""
|
||||
...
|
||||
|
||||
|
|
@ -49,6 +51,7 @@ class GraphAgentDeps(Protocol):
|
|||
|
||||
client: HaikuRAG
|
||||
context: GraphContext
|
||||
search_results: list[SearchResult]
|
||||
|
||||
|
||||
def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||
|
|
@ -228,28 +231,24 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
|||
async def search_and_answer(
|
||||
ctx2: RunContext[AgentDepsT], query: str, limit: int = 5
|
||||
) -> str:
|
||||
"""Search the knowledge base for relevant documents.
|
||||
|
||||
Returns results with chunk IDs and relevance scores.
|
||||
Reference results by their chunk_id in cited_chunks.
|
||||
"""
|
||||
results = await ctx2.deps.client.search(query, limit=limit)
|
||||
results = await ctx2.deps.client.expand_context(results)
|
||||
# Store results for citation resolution
|
||||
ctx2.deps.search_results = results
|
||||
|
||||
entries: list[dict[str, Any]] = []
|
||||
parts = []
|
||||
for r in results:
|
||||
entry: dict[str, Any] = {
|
||||
"text": r.content,
|
||||
"score": r.score,
|
||||
"document_uri": (r.document_uri or ""),
|
||||
}
|
||||
if r.document_title:
|
||||
entry["document_title"] = r.document_title
|
||||
if r.page_numbers:
|
||||
entry["page_numbers"] = r.page_numbers
|
||||
if r.headings:
|
||||
entry["headings"] = r.headings
|
||||
entries.append(entry)
|
||||
parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")
|
||||
|
||||
if not entries:
|
||||
if not parts:
|
||||
return f"No relevant information found in the knowledge base for: {query}"
|
||||
|
||||
return format_as_xml(entries, root_tag="snippets")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
# Tool is registered via decorator above
|
||||
_ = search_and_answer
|
||||
|
|
@ -260,7 +259,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
|||
result = await agent.run(sub_q, deps=agent_deps)
|
||||
answer = result.output
|
||||
if answer:
|
||||
state.context.add_qa_response(answer)
|
||||
state.context.add_qa_response(answer, agent_deps.search_results)
|
||||
# State updated with new answer - emit state update and narrate
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_state(state)
|
||||
|
|
|
|||
|
|
@ -20,32 +20,21 @@ Plan requirements:
|
|||
|
||||
SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
|
||||
|
||||
Tasks:
|
||||
1. Search the knowledge base for relevant evidence.
|
||||
2. Analyze retrieved snippets.
|
||||
3. Provide an answer strictly grounded in that evidence.
|
||||
Process:
|
||||
1. Call search_and_answer with relevant keywords from the question.
|
||||
2. Review the results and their relevance scores.
|
||||
3. If needed, perform follow-up searches with different keywords (max 3 total).
|
||||
4. Provide a concise answer based strictly on the retrieved content.
|
||||
|
||||
Tool usage:
|
||||
- Always call search_and_answer before drafting any answer.
|
||||
- The tool returns snippets with:
|
||||
- `text`: verbatim content
|
||||
- `score`: relevance score
|
||||
- `document_uri`: full path to the source document
|
||||
- `document_title`: title if available
|
||||
- `page_numbers`: list of page numbers where content appears (if available)
|
||||
- `headings`: section heading hierarchy (if available)
|
||||
- You may call the tool multiple times to refine or broaden context, but do not
|
||||
exceed 3 total calls. Favor precision over volume.
|
||||
- Use scores to prioritize evidence, but include only the minimal subset of
|
||||
snippet texts (verbatim) in SearchAnswer.context (typically 1-4).
|
||||
- Set SearchAnswer.sources to include document_uri, page numbers, and headings
|
||||
for each snippet used. Format: "document_uri (p. X, Section: Y)" or just
|
||||
"document_uri" if no page/heading info. One source per context snippet.
|
||||
- If no relevant information is found, clearly say so and return an empty
|
||||
context list and sources list.
|
||||
The search tool returns results like:
|
||||
[chunk_abc123] (score: 0.85) Content text here...
|
||||
[chunk_def456] (score: 0.72) More content...
|
||||
|
||||
Answering rules:
|
||||
- Be direct and specific; avoid meta commentary about the process.
|
||||
- Do not include any claims not supported by the provided snippets.
|
||||
- Prefer concise phrasing; avoid copying long passages.
|
||||
- When evidence is partial, state the limits explicitly in the answer."""
|
||||
In your response, include the chunk IDs you used in cited_chunks.
|
||||
|
||||
Guidelines:
|
||||
- Base answers strictly on retrieved content - do not use external knowledge.
|
||||
- If multiple results are relevant, synthesize them coherently.
|
||||
- If information is insufficient, say so clearly.
|
||||
- Be concise and direct; avoid meta commentary about the process.
|
||||
- Higher scores indicate more relevant results."""
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from pydantic import BaseModel, Field
|
|||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.graph.common.models import SearchAnswer
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class DeepQAContext(BaseModel):
|
||||
|
|
@ -12,11 +13,12 @@ class DeepQAContext(BaseModel):
|
|||
qa_responses: list[SearchAnswer] = Field(
|
||||
default_factory=list, description="QA pairs collected during answering"
|
||||
)
|
||||
use_citations: bool = Field(
|
||||
default=False, description="Whether to include citations in the answer"
|
||||
)
|
||||
|
||||
def add_qa_response(self, qa: SearchAnswer) -> None:
|
||||
def add_qa_response(
|
||||
self, qa: SearchAnswer, search_results: list[SearchResult]
|
||||
) -> None:
|
||||
"""Add a QA response."""
|
||||
del search_results # Not needed with chunk ID-based citations
|
||||
self.qa_responses.append(qa)
|
||||
|
||||
|
||||
|
|
@ -25,3 +27,6 @@ class DeepQADependencies(BaseModel):
|
|||
|
||||
client: HaikuRAG = Field(description="RAG client for document operations")
|
||||
context: DeepQAContext = Field(description="Shared QA context")
|
||||
search_results: list[SearchResult] = Field(
|
||||
default_factory=list, description="Search results for citation resolution"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,16 +6,13 @@ from pydantic_graph.beta.join import reduce_list_append
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.graph.common import get_model
|
||||
from haiku.rag.graph.common.models import SearchAnswer
|
||||
from haiku.rag.graph.common.models import SearchAnswer, resolve_citations
|
||||
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
|
||||
from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies
|
||||
from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation
|
||||
from haiku.rag.graph.deep_qa.prompts import (
|
||||
DECISION_PROMPT,
|
||||
SYNTHESIS_PROMPT,
|
||||
SYNTHESIS_PROMPT_WITH_CITATIONS,
|
||||
)
|
||||
from haiku.rag.graph.deep_qa.prompts import DECISION_PROMPT, SYNTHESIS_PROMPT
|
||||
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
def build_deep_qa_graph(
|
||||
|
|
@ -102,7 +99,7 @@ def build_deep_qa_graph(
|
|||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"sources": qa.sources,
|
||||
"confidence": qa.confidence,
|
||||
}
|
||||
for qa in state.context.qa_responses
|
||||
],
|
||||
|
|
@ -163,16 +160,10 @@ def build_deep_qa_graph(
|
|||
)
|
||||
|
||||
try:
|
||||
prompt_template = (
|
||||
SYNTHESIS_PROMPT_WITH_CITATIONS
|
||||
if state.context.use_citations
|
||||
else SYNTHESIS_PROMPT
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
model=get_model(model_config, config),
|
||||
output_type=DeepQAAnswer,
|
||||
instructions=prompt_template,
|
||||
output_type=SearchAnswer,
|
||||
instructions=SYNTHESIS_PROMPT,
|
||||
retries=3,
|
||||
deps_type=DeepQADependencies,
|
||||
)
|
||||
|
|
@ -183,7 +174,8 @@ def build_deep_qa_graph(
|
|||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"sources": qa.sources,
|
||||
"confidence": qa.confidence,
|
||||
"cited_chunks": qa.cited_chunks,
|
||||
}
|
||||
for qa in state.context.qa_responses
|
||||
],
|
||||
|
|
@ -197,13 +189,22 @@ def build_deep_qa_graph(
|
|||
context=state.context,
|
||||
)
|
||||
result = await agent.run(prompt, deps=agent_deps)
|
||||
llm_answer = result.output
|
||||
|
||||
# Resolve citations by fetching chunks by ID
|
||||
search_results = []
|
||||
for chunk_id in llm_answer.cited_chunks:
|
||||
chunk = await deps.client.chunk_repository.get_by_id(chunk_id)
|
||||
if chunk:
|
||||
search_results.append(SearchResult.from_chunk(chunk, score=1.0))
|
||||
citations = resolve_citations(llm_answer.cited_chunks, search_results)
|
||||
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_activity(
|
||||
"synthesizing", {"message": "Answer complete"}
|
||||
)
|
||||
|
||||
return result.output
|
||||
return DeepQAAnswer(answer=llm_answer.answer, citations=citations)
|
||||
finally:
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.finish_step()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
|
||||
class DeepQAEvaluation(BaseModel):
|
||||
is_sufficient: bool = Field(
|
||||
|
|
@ -13,8 +15,9 @@ class DeepQAEvaluation(BaseModel):
|
|||
|
||||
|
||||
class DeepQAAnswer(BaseModel):
|
||||
"""Final deep QA answer with resolved citations."""
|
||||
|
||||
answer: str = Field(description="The comprehensive answer to the question")
|
||||
sources: list[str] = Field(
|
||||
description="Document titles or URIs used to generate the answer",
|
||||
default_factory=list,
|
||||
citations: list[Citation] = Field(
|
||||
default_factory=list, description="Resolved citations for the answer"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,34 +10,15 @@ Task:
|
|||
|
||||
Output format:
|
||||
- answer: The complete answer to the original question (2-4 paragraphs)
|
||||
- sources: List of document titles/URIs used (extract from the sub-answers)
|
||||
- cited_chunks: List of chunk IDs (from sub-answers) that directly support your answer
|
||||
|
||||
Guidelines:
|
||||
- Start directly with the answer - no preamble like "Based on the research..."
|
||||
- Use a clear, professional tone
|
||||
- Organize information logically
|
||||
- If evidence is incomplete, state limitations clearly
|
||||
- Do not include any claims not supported by the gathered information"""
|
||||
|
||||
SYNTHESIS_PROMPT_WITH_CITATIONS = """You are an expert at synthesizing information into clear, concise answers with proper citations.
|
||||
|
||||
Task:
|
||||
- Combine the gathered information from sub-questions into a single comprehensive answer
|
||||
- Answer the original question directly and completely
|
||||
- Base your answer strictly on the provided evidence
|
||||
- Include inline citations using [Source Title] format
|
||||
|
||||
Output format:
|
||||
- answer: The complete answer with inline citations (2-4 paragraphs)
|
||||
- sources: List of document titles/URIs used (extract from the sub-answers)
|
||||
|
||||
Guidelines:
|
||||
- Start directly with the answer - no preamble like "Based on the research..."
|
||||
- Add citations after each claim: [Source Title]
|
||||
- Use a clear, professional tone
|
||||
- Organize information logically
|
||||
- If evidence is incomplete, state limitations clearly
|
||||
- Do not include any claims not supported by the gathered information"""
|
||||
- Do not include any claims not supported by the gathered information
|
||||
- Each sub-answer includes cited_chunks IDs - include the relevant ones in your response"""
|
||||
|
||||
DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question.
|
||||
|
||||
|
|
|
|||
|
|
@ -14,8 +14,7 @@ def format_context_for_prompt(context: ResearchContext) -> str:
|
|||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"context_snippets": qa.context,
|
||||
"sources": qa.sources, # pyright: ignore[reportAttributeAccessIssue]
|
||||
"confidence": qa.confidence,
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
],
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from haiku.rag.graph.research.models import (
|
|||
InsightAnalysis,
|
||||
InsightRecord,
|
||||
)
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class ResearchContext(BaseModel):
|
||||
|
|
@ -37,8 +38,12 @@ class ResearchContext(BaseModel):
|
|||
self._insights_by_id = {ins.id: ins for ins in self.insights}
|
||||
self._gaps_by_id = {gap.id: gap for gap in self.gaps}
|
||||
|
||||
def add_qa_response(self, qa: SearchAnswer) -> None:
|
||||
"""Add a structured QA response (minimal context already included)."""
|
||||
def add_qa_response(
|
||||
self, qa: SearchAnswer, search_results: list[SearchResult]
|
||||
) -> None:
|
||||
"""Add a structured QA response."""
|
||||
# Research doesn't accumulate search_results for citation resolution
|
||||
del search_results
|
||||
self.qa_responses.append(qa)
|
||||
|
||||
def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
|
||||
|
|
@ -144,6 +149,9 @@ class ResearchDependencies(BaseModel):
|
|||
|
||||
client: HaikuRAG = Field(description="RAG client for document operations")
|
||||
context: ResearchContext = Field(description="Shared research context")
|
||||
search_results: list[SearchResult] = Field(
|
||||
default_factory=list, description="Search results for citation resolution"
|
||||
)
|
||||
|
||||
|
||||
def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from haiku.rag.client import HaikuRAG
|
|||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.graph.research.models import ResearchReport
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
|
||||
class DocumentResult(BaseModel):
|
||||
|
|
@ -176,16 +177,17 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
|
|||
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
|
||||
|
||||
graph = build_deep_qa_graph(config=config)
|
||||
context = DeepQAContext(
|
||||
original_question=question, use_citations=cite
|
||||
)
|
||||
context = DeepQAContext(original_question=question)
|
||||
state = DeepQAState.from_config(context=context, config=config)
|
||||
deps = DeepQADeps(client=rag)
|
||||
|
||||
result = await graph.run(state=state, deps=deps)
|
||||
answer = result.answer
|
||||
citations = result.citations
|
||||
else:
|
||||
answer = await rag.ask(question, cite=cite)
|
||||
answer, citations = await rag.ask(question)
|
||||
if cite and citations:
|
||||
answer += "\n\n" + format_citations(citations)
|
||||
return answer
|
||||
except Exception as e:
|
||||
return f"Error answering question: {e!s}"
|
||||
|
|
|
|||
|
|
@ -6,16 +6,13 @@ from haiku.rag.qa.agent import QuestionAnswerAgent
|
|||
def get_qa_agent(
|
||||
client: HaikuRAG,
|
||||
config: AppConfig = Config,
|
||||
use_citations: bool = False,
|
||||
system_prompt: str | None = None,
|
||||
) -> QuestionAnswerAgent:
|
||||
"""
|
||||
Factory function to get a QA agent based on the configuration.
|
||||
"""Factory function to get a QA agent based on the configuration.
|
||||
|
||||
Args:
|
||||
client: HaikuRAG client instance.
|
||||
config: Configuration to use. Defaults to global Config.
|
||||
use_citations: Whether to include citations in responses.
|
||||
system_prompt: Optional custom system prompt.
|
||||
|
||||
Returns:
|
||||
|
|
@ -24,6 +21,5 @@ def get_qa_agent(
|
|||
return QuestionAnswerAgent(
|
||||
client=client,
|
||||
model_config=config.qa.model,
|
||||
use_citations=use_citations,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,33 +1,24 @@
|
|||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.output import ToolOutput
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
from haiku.rag.graph.common import get_model
|
||||
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
|
||||
|
||||
|
||||
class ToolSearchResult(BaseModel):
|
||||
"""Search result model exposed to the LLM tool."""
|
||||
|
||||
content: str = Field(description="The document text content")
|
||||
score: float = Field(description="Relevance score (higher is more relevant)")
|
||||
document_uri: str = Field(description="The URI/path of the source document")
|
||||
document_title: str | None = Field(
|
||||
default=None, description="The title of the document (if available)"
|
||||
)
|
||||
page_numbers: list[int] = Field(
|
||||
default=[], description="Page numbers where this content appears"
|
||||
)
|
||||
headings: list[str] | None = Field(
|
||||
default=None, description="Section heading hierarchy for this content"
|
||||
)
|
||||
from haiku.rag.graph.common.models import (
|
||||
Citation,
|
||||
SearchAnswer,
|
||||
resolve_citations,
|
||||
)
|
||||
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class Dependencies(BaseModel):
|
||||
model_config = {"arbitrary_types_allowed": True}
|
||||
client: HaikuRAG
|
||||
search_results: list[SearchResult] = []
|
||||
|
||||
|
||||
class QuestionAnswerAgent:
|
||||
|
|
@ -35,22 +26,16 @@ class QuestionAnswerAgent:
|
|||
self,
|
||||
client: HaikuRAG,
|
||||
model_config: ModelConfig,
|
||||
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 = get_model(model_config, Config)
|
||||
|
||||
self._agent = Agent(
|
||||
model=model_obj,
|
||||
deps_type=Dependencies,
|
||||
system_prompt=system_prompt,
|
||||
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
||||
instructions=system_prompt or QA_SYSTEM_PROMPT,
|
||||
retries=3,
|
||||
)
|
||||
|
||||
|
|
@ -59,24 +44,29 @@ class QuestionAnswerAgent:
|
|||
ctx: RunContext[Dependencies],
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
) -> list[ToolSearchResult]:
|
||||
"""Search the knowledge base for relevant documents."""
|
||||
) -> str:
|
||||
"""Search the knowledge base for relevant documents.
|
||||
|
||||
Returns results with chunk IDs and relevance scores.
|
||||
Reference results by their chunk_id in cited_chunks.
|
||||
"""
|
||||
results = await ctx.deps.client.search(query, limit=limit)
|
||||
results = await ctx.deps.client.expand_context(results)
|
||||
return [
|
||||
ToolSearchResult(
|
||||
content=r.content,
|
||||
score=r.score,
|
||||
document_uri=(r.document_uri or ""),
|
||||
document_title=r.document_title,
|
||||
page_numbers=r.page_numbers,
|
||||
headings=r.headings,
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
# Store results for citation resolution
|
||||
ctx.deps.search_results = results
|
||||
# Format with chunk IDs
|
||||
parts = []
|
||||
for r in results:
|
||||
parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")
|
||||
return "\n\n".join(parts) if parts else "No results found."
|
||||
|
||||
async def answer(self, question: str) -> str:
|
||||
"""Answer a question using the RAG system."""
|
||||
async def answer(self, question: str) -> tuple[str, list[Citation]]:
|
||||
"""Answer a question using the RAG system.
|
||||
|
||||
Returns:
|
||||
Tuple of (answer text, list of resolved citations)
|
||||
"""
|
||||
deps = Dependencies(client=self._client)
|
||||
result = await self._agent.run(question, deps=deps)
|
||||
return result.output
|
||||
citations = resolve_citations(result.output.cited_chunks, deps.search_results)
|
||||
return result.output.answer, citations
|
||||
|
|
|
|||
|
|
@ -1,65 +1,21 @@
|
|||
QA_SYSTEM_PROMPT = """
|
||||
You are a knowledgeable assistant that helps users find information from a document knowledge base.
|
||||
QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base.
|
||||
|
||||
Your process:
|
||||
1. When a user asks a question, use the search_documents tool to find relevant information
|
||||
2. Search with specific keywords and phrases from the user's question
|
||||
3. Review the search results and their relevance scores
|
||||
4. If you need additional context, perform follow-up searches with different keywords
|
||||
5. Provide a short and to the point comprehensive answer based only on the retrieved documents
|
||||
Process:
|
||||
1. Call search_documents with relevant keywords from the question
|
||||
2. Review the results and their relevance scores
|
||||
3. If needed, perform follow-up searches with different keywords (max 3 total)
|
||||
4. Provide a concise answer based strictly on the retrieved content
|
||||
|
||||
The search tool returns results like:
|
||||
[chunk_abc123] (score: 0.85) Content text here...
|
||||
[chunk_def456] (score: 0.72) More content...
|
||||
|
||||
In your response, include the chunk IDs you used in cited_chunks.
|
||||
|
||||
Guidelines:
|
||||
- Base your answers strictly on the provided document content
|
||||
- Quote or reference specific information when possible
|
||||
- If multiple documents contain relevant information, synthesize them coherently
|
||||
- Indicate when information is incomplete or when you need to search for additional context
|
||||
- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
|
||||
- For complex questions, consider breaking them down and performing multiple searches
|
||||
- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
|
||||
|
||||
Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
|
||||
/no_think
|
||||
"""
|
||||
|
||||
QA_SYSTEM_PROMPT_WITH_CITATIONS = """
|
||||
You are a knowledgeable assistant that helps users find information from a document knowledge base.
|
||||
|
||||
IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.
|
||||
|
||||
Your process:
|
||||
1. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question
|
||||
2. Review the search results and their relevance scores
|
||||
3. If you need additional context, perform follow-up searches with different keywords
|
||||
4. Provide a short and to the point comprehensive answer based only on the retrieved documents
|
||||
5. Always include citations for the sources used in your answer
|
||||
|
||||
Guidelines:
|
||||
- Base your answers strictly on the provided document content
|
||||
- If multiple documents contain relevant information, synthesize them coherently
|
||||
- Indicate when information is incomplete or when you need to search for additional context
|
||||
- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
|
||||
- For complex questions, consider breaking them down and performing multiple searches
|
||||
- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
|
||||
- ALWAYS include citations at the end of your response using the format below
|
||||
|
||||
Citation Format:
|
||||
After your answer, include a "Citations:" section that lists:
|
||||
- The document URI (from the document_uri field) - always include the full path
|
||||
- The document title if available (from the document_title field)
|
||||
- Page number(s) if available (from the page_numbers field)
|
||||
- Section heading if available (from the headings field)
|
||||
- A VERBATIM excerpt (100-200 characters, copy-paste exact text) from the content field - do NOT summarize, truncate, or paraphrase
|
||||
|
||||
Example response format:
|
||||
[Your answer here]
|
||||
|
||||
Citations:
|
||||
|
||||
- **/docs/user-guide.pdf** - "User Manual" (p. 5, Section: Introduction)
|
||||
The system requires Python 3.10 or higher to run properly...
|
||||
|
||||
- **/reports/quarterly-analysis.md** (pp. 12-13, Section: Results)
|
||||
Revenue increased by 15% compared to the previous quarter...
|
||||
|
||||
Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
|
||||
- Base answers strictly on retrieved content - do not use external knowledge
|
||||
- If multiple results are relevant, synthesize them coherently
|
||||
- If information is insufficient, say: "I cannot find enough information in the knowledge base to answer this question."
|
||||
- Be concise and direct - avoid elaboration unless asked
|
||||
- Higher scores indicate more relevant results
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,10 +4,13 @@ import sys
|
|||
from importlib import metadata
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from packaging.version import Version, parse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
|
||||
def apply_common_settings(
|
||||
settings: Any | None,
|
||||
|
|
@ -271,6 +274,34 @@ def format_bytes(num_bytes: int) -> str:
|
|||
return f"{size:.1f} PB"
|
||||
|
||||
|
||||
def format_citations(citations: "list[Citation]") -> str:
|
||||
"""Format citations as markdown string."""
|
||||
if not citations:
|
||||
return ""
|
||||
lines = ["## Citations\n"]
|
||||
for c in citations:
|
||||
# Build citation header
|
||||
parts = [f"- **{c.document_uri}**"]
|
||||
if c.document_title:
|
||||
parts.append(f' - "{c.document_title}"')
|
||||
location_parts = []
|
||||
if c.page_numbers:
|
||||
if len(c.page_numbers) == 1:
|
||||
location_parts.append(f"p. {c.page_numbers[0]}")
|
||||
else:
|
||||
location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}")
|
||||
if c.headings:
|
||||
location_parts.append(f"Section: {c.headings[-1]}")
|
||||
if location_parts:
|
||||
parts.append(f" ({', '.join(location_parts)})")
|
||||
lines.append("".join(parts))
|
||||
# Add truncated content excerpt
|
||||
excerpt = c.content[:500] + "…" if len(c.content) > 500 else c.content
|
||||
excerpt = excerpt.replace("\r\n", " ").replace("\n", " ").replace("\r", " ")
|
||||
lines.append(f"\n {excerpt}\n")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def get_default_data_dir() -> Path:
|
||||
"""Get the user data directory for the current system platform.
|
||||
|
||||
|
|
|
|||
|
|
@ -25,9 +25,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
|
|||
graph = build_deep_qa_graph()
|
||||
|
||||
state = DeepQAState(
|
||||
context=DeepQAContext(
|
||||
original_question="What is haiku.rag?", use_citations=False
|
||||
),
|
||||
context=DeepQAContext(original_question="What is haiku.rag?"),
|
||||
max_sub_questions=3,
|
||||
)
|
||||
|
||||
|
|
@ -40,42 +38,6 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
|
|||
# TestModel will generate valid structured output based on schemas
|
||||
assert result.answer is not None
|
||||
assert isinstance(result.answer, str)
|
||||
assert isinstance(result.sources, list)
|
||||
|
||||
client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
|
||||
"""Test deep Q&A with citations enabled using TestModel."""
|
||||
|
||||
# Mock get_model to return TestModel
|
||||
def test_model_factory(provider, model, config=None):
|
||||
return TestModel()
|
||||
|
||||
# Patch all locations where get_model is imported
|
||||
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
|
||||
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
|
||||
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
|
||||
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
||||
|
||||
graph = build_deep_qa_graph()
|
||||
|
||||
state = DeepQAState(
|
||||
context=DeepQAContext(original_question="What is Python?", use_citations=True),
|
||||
max_sub_questions=2,
|
||||
)
|
||||
|
||||
# Use real client but with TestModel for LLM calls
|
||||
client = HaikuRAG(temp_db_path, create=True)
|
||||
deps = DeepQADeps(client=client)
|
||||
|
||||
result = await graph.run(state=state, deps=deps)
|
||||
|
||||
# Verify citations flag was used
|
||||
assert state.context.use_citations is True
|
||||
assert result.answer is not None
|
||||
assert isinstance(result.sources, list)
|
||||
|
||||
client.close()
|
||||
|
||||
|
|
@ -87,7 +49,6 @@ async def test_deep_qa_context_operations():
|
|||
assert context.original_question == "Test question?"
|
||||
assert context.sub_questions == []
|
||||
assert context.qa_responses == []
|
||||
assert context.use_citations is False
|
||||
|
||||
context.sub_questions = ["Sub Q1", "Sub Q2"]
|
||||
assert len(context.sub_questions) == 2
|
||||
|
|
@ -95,10 +56,10 @@ async def test_deep_qa_context_operations():
|
|||
qa = SearchAnswer(
|
||||
query="Sub Q1",
|
||||
answer="Answer 1",
|
||||
context=["Context 1"],
|
||||
sources=["source1.md"],
|
||||
cited_chunks=["chunk_1", "chunk_2"],
|
||||
confidence=0.9,
|
||||
)
|
||||
context.add_qa_response(qa)
|
||||
context.add_qa_response(qa, search_results=[])
|
||||
assert len(context.qa_responses) == 1
|
||||
assert context.qa_responses[0].query == "Sub Q1"
|
||||
|
||||
|
|
|
|||
|
|
@ -293,8 +293,9 @@ async def test_serve_all_services(app: HaikuRAGApp, monkeypatch):
|
|||
async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
|
||||
"""Test asking a question without citations."""
|
||||
mock_answer = "Test answer"
|
||||
mock_citations = []
|
||||
mock_client = AsyncMock()
|
||||
mock_client.ask.return_value = mock_answer
|
||||
mock_client.ask.return_value = (mock_answer, mock_citations)
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
|
||||
mock_print = MagicMock()
|
||||
|
|
@ -303,15 +304,25 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
|
|||
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
|
||||
await app.ask("test question")
|
||||
|
||||
mock_client.ask.assert_called_once_with("test question", cite=False)
|
||||
mock_client.ask.assert_called_once_with("test question")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
|
||||
"""Test asking a question with citations."""
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
mock_answer = "Test answer with citations"
|
||||
mock_citations = [
|
||||
Citation(
|
||||
document_uri="test.md",
|
||||
document_title="Test Document",
|
||||
page_numbers=[1],
|
||||
content="Test content",
|
||||
)
|
||||
]
|
||||
mock_client = AsyncMock()
|
||||
mock_client.ask.return_value = mock_answer
|
||||
mock_client.ask.return_value = (mock_answer, mock_citations)
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
|
||||
mock_print = MagicMock()
|
||||
|
|
@ -320,15 +331,18 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
|
|||
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
|
||||
await app.ask("test question", cite=True)
|
||||
|
||||
mock_client.ask.assert_called_once_with("test question", cite=True)
|
||||
mock_client.ask.assert_called_once_with("test question")
|
||||
# Verify print was called (once for answer, once for citations)
|
||||
assert mock_print.call_count >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
|
||||
"""Test asking a question with verbose (should be ignored for non-deep)."""
|
||||
mock_answer = "Test answer"
|
||||
mock_citations = []
|
||||
mock_client = AsyncMock()
|
||||
mock_client.ask.return_value = mock_answer
|
||||
mock_client.ask.return_value = (mock_answer, mock_citations)
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
|
||||
mock_print = MagicMock()
|
||||
|
|
@ -337,7 +351,7 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
|
|||
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
|
||||
await app.ask("test question", verbose=True)
|
||||
|
||||
mock_client.ask.assert_called_once_with("test question", cite=False)
|
||||
mock_client.ask.assert_called_once_with("test question")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -345,7 +359,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
|
|||
"""Test asking a question with deep QA."""
|
||||
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
|
||||
|
||||
mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
|
||||
mock_output = DeepQAAnswer(answer="Deep QA answer")
|
||||
|
||||
mock_graph = AsyncMock()
|
||||
mock_graph.run.return_value = mock_output
|
||||
|
|
@ -365,17 +379,14 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
|
|||
mock_graph.run.assert_called_once()
|
||||
call_kwargs = mock_graph.run.call_args[1]
|
||||
assert call_kwargs["state"].context.original_question == "test question"
|
||||
assert call_kwargs["state"].context.use_citations is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
|
||||
"""Test asking a question with deep QA and citations."""
|
||||
"""Test asking a question with deep QA and citations (cite ignored for deep)."""
|
||||
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
|
||||
|
||||
mock_output = DeepQAAnswer(
|
||||
answer="Deep QA answer with citations [test.md]", sources=["test.md"]
|
||||
)
|
||||
mock_output = DeepQAAnswer(answer="Deep QA answer")
|
||||
|
||||
mock_graph = AsyncMock()
|
||||
mock_graph.run.return_value = mock_output
|
||||
|
|
@ -395,14 +406,13 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
|
|||
mock_graph.run.assert_called_once()
|
||||
call_kwargs = mock_graph.run.call_args[1]
|
||||
assert call_kwargs["state"].context.original_question == "test question"
|
||||
assert call_kwargs["state"].context.use_citations is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
|
||||
"""Test asking a question with deep QA and verbose output."""
|
||||
|
||||
mock_output = {"answer": "Deep QA answer", "sources": ["test.md"]}
|
||||
mock_output = {"answer": "Deep QA answer", "citations": []}
|
||||
|
||||
mock_renderer = AsyncMock()
|
||||
mock_renderer.render.return_value = mock_output
|
||||
|
|
@ -419,9 +429,7 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
|
|||
with patch(
|
||||
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
|
||||
):
|
||||
with patch(
|
||||
"haiku.rag.graph.agui.AGUIConsoleRenderer", return_value=mock_renderer
|
||||
):
|
||||
with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer):
|
||||
await app.ask("test question", deep=True, verbose=True)
|
||||
|
||||
# With verbose, it should use AGUIConsoleRenderer.render, not graph.run
|
||||
|
|
|
|||
|
|
@ -732,8 +732,8 @@ async def test_client_create_document_with_custom_chunks(temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_ask_without_cite(monkeypatch, temp_db_path):
|
||||
"""Test asking questions without citations."""
|
||||
async def test_client_ask(monkeypatch, temp_db_path):
|
||||
"""Test asking questions returns answer and citations."""
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
# Mock get_model to return TestModel
|
||||
|
|
@ -748,35 +748,12 @@ async def test_client_ask_without_cite(monkeypatch, temp_db_path):
|
|||
)
|
||||
|
||||
# Use real QA agent with TestModel
|
||||
answer = await client.ask("What is Python?")
|
||||
|
||||
# TestModel will generate a valid string response
|
||||
assert answer is not None
|
||||
assert isinstance(answer, str)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_ask_with_cite(monkeypatch, temp_db_path):
|
||||
"""Test asking questions with citations."""
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
# Mock get_model to return TestModel
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel()
|
||||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Create a test document
|
||||
await client.create_document(
|
||||
content="Python is a high-level programming language.", uri="test.txt"
|
||||
)
|
||||
|
||||
# Use real QA agent with TestModel
|
||||
answer = await client.ask("What is Python?", cite=True)
|
||||
answer, citations = await client.ask("What is Python?")
|
||||
|
||||
# TestModel will generate a valid string response
|
||||
assert answer is not None
|
||||
assert isinstance(answer, str)
|
||||
assert isinstance(citations, list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ async def test_mcp_ask_question():
|
|||
|
||||
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
|
||||
mock_rag = AsyncMock()
|
||||
mock_rag.ask = AsyncMock(return_value="This is the answer")
|
||||
mock_rag.ask = AsyncMock(return_value=("This is the answer", []))
|
||||
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
|
||||
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
|
|
@ -57,7 +57,7 @@ async def test_mcp_ask_question():
|
|||
)
|
||||
|
||||
assert result == "This is the answer"
|
||||
mock_rag.ask.assert_called_once_with("What is this?", cite=False)
|
||||
mock_rag.ask.assert_called_once_with("What is this?")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -261,14 +261,16 @@ async def test_mcp_ask_question_deep():
|
|||
mock_graph = AsyncMock()
|
||||
mock_result = AsyncMock()
|
||||
mock_result.answer = "Deep answer"
|
||||
mock_result.citations = []
|
||||
mock_graph.run = AsyncMock(return_value=mock_result)
|
||||
mock_graph_builder.return_value = mock_graph
|
||||
|
||||
tools = await mcp.get_tools()
|
||||
ask_tool = next(t for t in tools.values() if t.name == "ask_question")
|
||||
|
||||
# cite=False to avoid citation formatting in output
|
||||
result = await ask_tool.fn( # type: ignore[attr-defined]
|
||||
question="Deep question?", cite=True, deep=True
|
||||
question="Deep question?", cite=False, deep=True
|
||||
)
|
||||
|
||||
assert result == "Deep answer"
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
|
|||
question = doc["question"]
|
||||
expected_answer = doc["answer"]
|
||||
|
||||
answer = await qa.answer(question)
|
||||
answer, _ = await qa.answer(question)
|
||||
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
||||
|
||||
assert is_equivalent, (
|
||||
|
|
@ -55,7 +55,7 @@ async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
|
|||
question = doc["question"]
|
||||
expected_answer = doc["answer"]
|
||||
|
||||
answer = await qa.answer(question)
|
||||
answer, _ = await qa.answer(question)
|
||||
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
||||
|
||||
assert is_equivalent, (
|
||||
|
|
@ -81,7 +81,7 @@ async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path):
|
|||
question = doc["question"]
|
||||
expected_answer = doc["answer"]
|
||||
|
||||
answer = await qa.answer(question)
|
||||
answer, _ = await qa.answer(question)
|
||||
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
||||
|
||||
assert is_equivalent, (
|
||||
|
|
@ -104,7 +104,7 @@ async def test_qa_vllm(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
question = doc["question"]
|
||||
expected_answer = doc["answer"]
|
||||
answer = await qa.answer(question)
|
||||
answer, _ = await qa.answer(question)
|
||||
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
||||
|
||||
assert is_equivalent, (
|
||||
|
|
|
|||
Loading…
Reference in a new issue