Revert "Synthesis steps now select only relevant citations instead of including all"

This commit is contained in:
Yiorgis Gozadinos 2026-01-28 16:00:27 +02:00
parent e2158c7dad
commit 57fd08d600
No known key found for this signature in database
13 changed files with 5489 additions and 6048 deletions

View file

@ -21,10 +21,7 @@
- Reduces bandwidth when state grows large (e.g., 50 Q&As with citations)
- First request still sends full snapshot; subsequent requests send only changes
- Backend logging shows incoming/outgoing state events for debugging
- **Selective Citation Filtering**: Synthesis steps now select only relevant citations instead of including all
- LLM receives `<available_citations>` with chunk IDs and content previews
- LLM populates `cited_chunks` with only chunks that directly support the answer
- `ResearchReport` now has `cited_chunks` and `citations` fields; removed `sources_summary`
## [0.27.1] - 2026-01-27

View file

@ -71,7 +71,7 @@ prompts:
Replace the research report synthesis prompt by setting `prompts.synthesis`. This controls how the multi-agent research workflow generates its final report.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `cited_chunks`.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `sources_summary`.
**Example:**
@ -87,13 +87,12 @@ prompts:
- conclusions: 2-4 bullet points
- recommendations: 2-5 actionable recommendations
- limitations: 1-3 limitations or gaps
- cited_chunks: List of chunk IDs that directly support the report
- sources_summary: Brief description of sources used
Guidelines:
- Base all content strictly on collected evidence
- Be specific and objective
- Avoid meta-commentary like "This report covers..."
- Only include chunks in cited_chunks that directly support claims in the report
```
## Picture Description Prompt

View file

@ -26,11 +26,7 @@ Each result includes:
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Citation guidelines:
- In cited_chunks, include ONLY chunk IDs that directly support your answer.
- Do NOT cite chunks that are merely related or that you reviewed but did not use.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT, COMPLETE chunk IDs (full UUIDs).
In your response, include the chunk IDs you used in cited_chunks.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge

View file

@ -32,7 +32,6 @@ from haiku.rag.utils import build_prompt, get_model
def format_context_for_prompt(
context: ResearchContext,
include_pending_questions: bool = True,
include_citations: bool = False,
) -> str:
"""Format the research context as XML for prompts.
@ -40,8 +39,6 @@ def format_context_for_prompt(
context: The research context to format.
include_pending_questions: Whether to include pending sub-questions.
Set to False for synthesis prompts where pending questions aren't relevant.
include_citations: Whether to include available citations for selection.
Set to True for synthesis prompts where the LLM should select relevant citations.
"""
context_data: dict[str, object] = {}
@ -64,26 +61,6 @@ def format_context_for_prompt(
for qa in context.qa_responses
]
if include_citations and context.qa_responses:
seen_chunks: set[str] = set()
available_citations: list[dict[str, str]] = []
for qa in context.qa_responses:
for c in qa.citations:
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
content_preview = (
c.content[:500] + "..." if len(c.content) > 500 else c.content
)
available_citations.append(
{
"chunk_id": c.chunk_id,
"document": c.document_title or c.document_uri,
"content": content_preview,
}
)
if available_citations:
context_data["available_citations"] = available_citations
return format_as_xml(context_data, root_tag="context")
@ -373,10 +350,7 @@ def build_research_graph(
deps_type=ResearchDependencies,
)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False, include_citations=True
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
@ -387,21 +361,7 @@ def build_research_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
report = result.output
citation_lookup: dict[str, Citation] = {}
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
resolved_citations: list[Citation] = []
for chunk_id in report.cited_chunks:
if chunk_id in citation_lookup:
resolved_citations.append(citation_lookup[chunk_id])
report.citations = resolved_citations
return report
return result.output
# Build the graph structure
collect_answers = g.join(
@ -519,19 +479,17 @@ def build_conversational_graph(
state = ctx.state
deps = ctx.deps
# Use RawSearchAnswer so LLM can select which chunks to cite
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[invalid-assignment]
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[invalid-assignment]
model=get_model(config.research.model, config),
output_type=RawSearchAnswer,
output_type=ConversationalAnswer,
instructions=conversational_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False, include_citations=True
state.context, include_pending_questions=False
)
prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
agent_deps = ResearchDependencies(
@ -539,23 +497,20 @@ def build_conversational_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
raw_answer = result.output
citation_lookup: dict[str, Citation] = {}
# Collect unique citations from qa_responses (dedupe by chunk_id)
seen_chunks: set[str] = set()
unique_citations: list[Citation] = []
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
filtered_citations: list[Citation] = []
for chunk_id in raw_answer.cited_chunks:
if chunk_id in citation_lookup:
filtered_citations.append(citation_lookup[chunk_id])
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
unique_citations.append(c)
return ConversationalAnswer(
answer=raw_answer.answer,
citations=filtered_citations,
confidence=raw_answer.confidence,
answer=result.output.answer,
citations=unique_citations,
confidence=result.output.confidence,
)
# Build the graph structure (simplified: plan → search → synthesize)

View file

@ -163,11 +163,6 @@ class ResearchReport(BaseModel):
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
cited_chunks: list[str] = Field(
default_factory=list,
description="Chunk IDs selected by synthesis as directly supporting the report",
)
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
sources_summary: str = Field(
description="Summary of sources used and their reliability"
)

View file

@ -116,7 +116,6 @@ Goals:
2. Present findings clearly and concisely.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
5. Select only the citations that directly support your final answer.
Report guidelines (map to output fields):
- title: concise (5-12 words), informative.
@ -128,17 +127,10 @@ Report guidelines (map to output fields):
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- cited_chunks: list of chunk IDs that DIRECTLY support your report.
- sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific claims in your report.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from the available_citations (full UUIDs).
Style:
- Base all content solely on the collected evidence.
- Be professional, objective, and specific.
@ -149,11 +141,9 @@ CONVERSATIONAL_SYNTHESIS_PROMPT = """Generate a direct, conversational answer
to the question based on the gathered evidence.
Output:
- query: Echo the original question being answered.
- answer: Direct, comprehensive answer with a natural, helpful tone.
Write the actual answer, not a description of what you found.
Use as many sentences as needed to fully address the question.
- cited_chunks: List of chunk IDs that DIRECTLY support your answer.
- confidence: Score from 0.0 to 1.0 indicating answer quality.
Guidelines:
@ -163,11 +153,4 @@ Guidelines:
- Use formatting (bullet points, numbered lists) when it improves clarity.
- Do NOT use meta-commentary like "Based on the research..." or "The evidence shows..."
Instead, directly state the information.
- If the evidence is incomplete, acknowledge limitations briefly.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific statements in your answer.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from available_citations (full UUIDs)."""
- If the evidence is incomplete, acknowledge limitations briefly."""

View file

@ -416,9 +416,10 @@ class HaikuRAGApp:
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"{finding}")
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
if report.sources_summary:
self.console.print()
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
else:
self.console.print("[yellow]No answer generated.[/yellow]")
else:
@ -512,10 +513,10 @@ class HaikuRAGApp:
self.console.print(f"{limitation}")
self.console.print()
# Sources
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
# Sources Summary
if report.sources_summary:
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(

View file

@ -184,46 +184,3 @@ def test_format_context_for_prompt_with_prior_answers():
assert "<prior_answers>" in result
assert "Sub question?" in result
assert "The answer is here." in result
def test_format_context_for_prompt_with_citations():
"""Test format_context_for_prompt includes available_citations when requested."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
from haiku.rag.agents.research.models import Citation, SearchAnswer
context = ResearchContext(original_question="Main question?")
context.add_qa_response(
SearchAnswer(
query="Sub question?",
answer="The answer is here.",
confidence=0.9,
cited_chunks=["chunk-123", "chunk-456"],
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-123",
document_uri="test://doc1",
document_title="Test Document",
content="This is the chunk content.",
),
Citation(
document_id="doc-1",
chunk_id="chunk-456",
document_uri="test://doc1",
document_title="Test Document",
content="More chunk content here.",
),
],
)
)
result_without = format_context_for_prompt(context, include_citations=False)
assert "<available_citations>" not in result_without
result_with = format_context_for_prompt(context, include_citations=True)
assert "<available_citations>" in result_with
assert "chunk-123" in result_with
assert "chunk-456" in result_with
assert "Test Document" in result_with
assert "This is the chunk content." in result_with

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -349,6 +349,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()
@ -386,6 +387,7 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()

View file

@ -277,6 +277,7 @@ async def test_mcp_research_question():
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
recommendations=["Recommendation 1"],
sources_summary="Sources used",
)
with (