diff --git a/haiku_rag_slim/haiku/rag/graph/common/models.py b/haiku_rag_slim/haiku/rag/graph/common/models.py index ba2d1914..e559096e 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ b/haiku_rag_slim/haiku/rag/graph/common/models.py @@ -29,17 +29,17 @@ class ResearchPlan(BaseModel): class Citation(BaseModel): """Resolved citation with full metadata for display/visual grounding.""" - document_id: str = "" - chunk_id: str = "" - document_uri: str = "" + document_id: str + chunk_id: str + document_uri: str document_title: str | None = None page_numbers: list[int] = Field(default_factory=list) headings: list[str] | None = None - content: str = "" + content: str -class SearchAnswer(BaseModel): - """Structured answer from a search operation.""" +class RawSearchAnswer(BaseModel): + """Answer to a search query with chunk references.""" query: str = Field(..., description="The question that was answered") answer: str = Field(..., description="The answer to the question") @@ -53,11 +53,32 @@ class SearchAnswer(BaseModel): ge=0.0, le=1.0, ) + + +class SearchAnswer(RawSearchAnswer): + """Answer to a search query with resolved citations.""" + citations: list[Citation] = Field( default_factory=list, - description="Resolved citations with full metadata (populated after search)", + description="Resolved citations with full metadata", ) + @classmethod + def from_raw( + cls, + raw: RawSearchAnswer, + search_results: "list[SearchResult]", + ) -> "SearchAnswer": + """Create SearchAnswer from RawSearchAnswer with resolved citations.""" + citations = resolve_citations(raw.cited_chunks, search_results) + return cls( + query=raw.query, + answer=raw.answer, + cited_chunks=raw.cited_chunks, + confidence=raw.confidence, + citations=citations, + ) + def resolve_citations( cited_chunk_ids: list[str], diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 13cd7ff2..5781b992 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -13,7 +13,7 @@ from haiku.rag.config import Config from haiku.rag.config.models import AppConfig, ModelConfig 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.models import RawSearchAnswer, ResearchPlan, SearchAnswer from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT from haiku.rag.store.models import SearchResult @@ -24,9 +24,7 @@ class GraphContext(Protocol): original_question: str sub_questions: list[str] - def add_qa_response( - self, qa: SearchAnswer, search_results: list[SearchResult] - ) -> None: + def add_qa_response(self, qa: SearchAnswer) -> None: """Add a QA response to context.""" ... @@ -221,7 +219,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( agent = Agent( model=get_model(model_config, config), - output_type=ToolOutput(SearchAnswer, max_retries=3), + output_type=ToolOutput(RawSearchAnswer, max_retries=3), instructions=SEARCH_AGENT_PROMPT, retries=3, deps_type=deps_type, @@ -257,9 +255,11 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( try: result = await agent.run(sub_q, deps=agent_deps) - answer = result.output - if answer: - state.context.add_qa_response(answer, agent_deps.search_results) + raw_answer = result.output + if raw_answer: + # Convert RawSearchAnswer to SearchAnswer with resolved citations + answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results) + state.context.add_qa_response(answer) # State updated with new answer - emit state update and narrate if deps.agui_emitter: deps.agui_emitter.update_state(state) @@ -279,7 +279,9 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( "confidence": answer.confidence, }, ) - return answer + return answer + # Return empty SearchAnswer if no result + return SearchAnswer(query=sub_q, answer="", confidence=0.0) except Exception as e: if handle_exceptions: # Narrate the error diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py index 85b3fed7..efd6d92f 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py @@ -1,7 +1,7 @@ from pydantic import BaseModel, Field from haiku.rag.client import HaikuRAG -from haiku.rag.graph.common.models import SearchAnswer, resolve_citations +from haiku.rag.graph.common.models import SearchAnswer from haiku.rag.store.models import SearchResult @@ -14,11 +14,8 @@ class DeepQAContext(BaseModel): default_factory=list, description="QA pairs collected during answering" ) - def add_qa_response( - self, qa: SearchAnswer, search_results: list[SearchResult] - ) -> None: - """Add a QA response with resolved citations.""" - qa.citations = resolve_citations(qa.cited_chunks, search_results) + def add_qa_response(self, qa: SearchAnswer) -> None: + """Add a QA response (citations already resolved).""" self.qa_responses.append(qa) diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index aba28df7..542710dd 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -3,7 +3,7 @@ from collections.abc import Iterable from pydantic import BaseModel, Field, PrivateAttr from haiku.rag.client import HaikuRAG -from haiku.rag.graph.common.models import SearchAnswer, resolve_citations +from haiku.rag.graph.common.models import SearchAnswer from haiku.rag.graph.research.models import ( GapRecord, InsightAnalysis, @@ -38,11 +38,8 @@ 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, search_results: list[SearchResult] - ) -> None: - """Add a structured QA response with resolved citations.""" - qa.citations = resolve_citations(qa.cited_chunks, search_results) + def add_qa_response(self, qa: SearchAnswer) -> None: + """Add a structured QA response (citations already resolved).""" self.qa_responses.append(qa) def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index b4d7f8e0..21bdb69e 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -6,11 +6,7 @@ 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.graph.common.models import ( - Citation, - SearchAnswer, - resolve_citations, -) +from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT from haiku.rag.store.models import SearchResult @@ -34,7 +30,7 @@ class QuestionAnswerAgent: self._agent = Agent( model=model_obj, deps_type=Dependencies, - output_type=ToolOutput(SearchAnswer, max_retries=3), + output_type=ToolOutput(RawSearchAnswer, max_retries=3), instructions=system_prompt or QA_SYSTEM_PROMPT, retries=3, ) diff --git a/tests/graph/test_deep_qa.py b/tests/graph/test_deep_qa.py index 05288ec6..69c9c078 100644 --- a/tests/graph/test_deep_qa.py +++ b/tests/graph/test_deep_qa.py @@ -59,7 +59,7 @@ async def test_deep_qa_context_operations(): cited_chunks=["chunk_1", "chunk_2"], confidence=0.9, ) - context.add_qa_response(qa, search_results=[]) + context.add_qa_response(qa) assert len(context.qa_responses) == 1 assert context.qa_responses[0].query == "Sub Q1"