danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add citation support to analysis agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-17 14:04:14 +03:00
parent cf89ff55cd
commit 7d98d0ec57
No known key found for this signature in database
10 changed files with 82 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
haiku_rag_slim/haiku/rag/agents/analysis/__init__.py
View file

@ -1,6 +1,10 @@
from haiku.rag.agents.analysis.agent import create_analysis_agent
from haiku.rag.agents.analysis.dependencies import AnalysisContext, AnalysisDeps
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
from haiku.rag.agents.analysis.models import (
AnalysisResult,
CodeExecution,
RawAnalysisResult,
)
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
@ -8,6 +12,7 @@ __all__ = [
"ANALYSIS_SYSTEM_PROMPT",
"AnalysisContext",
"AnalysisDeps",
"RawAnalysisResult",
"AnalysisResult",
"CodeExecution",
"Sandbox",

8
haiku_rag_slim/haiku/rag/agents/analysis/agent.py
View file

@ -1,13 +1,13 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, AnalysisResult]:
def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisResult]:
"""Create an analysis agent with code execution capability.
The analysis agent can write and execute Python code in a sandboxed
@ -22,10 +22,10 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, AnalysisResu
"""
model = get_model(config.analysis.model, config)
agent: Agent[AnalysisDeps, AnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
agent: Agent[AnalysisDeps, RawAnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model,
deps_type=AnalysisDeps,
output_type=AnalysisResult,
output_type=RawAnalysisResult,
instructions=ANALYSIS_SYSTEM_PROMPT,
retries=3,
)

18
haiku_rag_slim/haiku/rag/agents/analysis/models.py
View file

@ -1,5 +1,7 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
class CodeExecution(BaseModel):
"""Result of executing a code block in the analysis sandbox."""
@ -10,8 +12,20 @@ class CodeExecution(BaseModel):
success: bool = Field(description="Whether execution completed without error")
class AnalysisResult(BaseModel):
"""Result from analysis agent execution."""
class RawAnalysisResult(BaseModel):
"""Raw result from the analysis agent (LLM output)."""
answer: str = Field(description="The answer to the user's question")
program: str = Field(description="The final consolidated program")
cited_chunks: list[str] = Field(
default_factory=list,
description="Chunk IDs from search results that informed the answer. Copy full UUIDs from search result chunk_id fields.",
)
class AnalysisResult(BaseModel):
"""Result from analysis execution with resolved citations."""
answer: str
program: str
citations: list[Citation] = Field(default_factory=list)

5
haiku_rag_slim/haiku/rag/agents/analysis/prompts.py
View file

@ -116,12 +116,13 @@ Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available)
Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
{"answer": "Your answer here", "program": "Your final program here", "cited_chunks": ["chunk-id-1", "chunk-id-2"]}
```
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
- `cited_chunks`: List of chunk_id values from search results that informed your answer. Copy the full UUID strings from the `chunk_id` field of search results you used.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format above.
You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

4
haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py
View file

@ -10,6 +10,7 @@ from pydantic_monty import CallbackFile, MemoryFile, OSAccess
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
if TYPE_CHECKING:
from pathlib import PurePosixPath
@ -47,6 +48,7 @@ class Sandbox:
_client: "HaikuRAG"
_config: AppConfig
_context: AnalysisContext
_search_results: "list[SearchResult]"
def __init__(
self,
@ -57,6 +59,7 @@ class Sandbox:
self._client = client
self._config = config
self._context = context
self._search_results = []
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
@ -67,6 +70,7 @@ class Sandbox:
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
results = await client.search(query, limit=limit, filter=context.filter)
expanded = await client.expand_context(results)
self._search_results.extend(expanded)
return [
{
"chunk_id": r.chunk_id,

24
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -30,6 +30,7 @@ from textual.worker import Worker
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.skills.analysis import AnalysisState
from haiku.rag.skills.rag import RAGState, get_agent_preamble
from haiku.skills.agent import (
SkillToolset,
@ -51,6 +52,7 @@ if TYPE_CHECKING:
RAG_STATE_NAMESPACE = "rag"
ANALYSIS_STATE_NAMESPACE = "analysis"
class ChatApp(App):
@ -319,15 +321,15 @@ class ChatApp(App):
chat_input.focus()
async def _show_citations(self, chat_history: "ChatHistory") -> None:
"""Show citations from the RAG state after an agent response."""
"""Show citations from skill states after an agent response."""
if not self._toolset:
return
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if rag_state is None:
return
citations = getattr(rag_state, "citations", [])
citations = []
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
state = self._toolset.get_namespace(namespace)
if state:
citations.extend(getattr(state, "citations", []))
if citations:
# Show only new citations (since last response)
await chat_history.add_citations(citations)
async def action_clear_chat(self) -> None:
@ -420,9 +422,11 @@ class ChatApp(App):
self._document_filter = event.selected
if self._toolset:
doc_filter = build_multi_document_filter(self._document_filter)
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if isinstance(rag_state, RAGState):
rag_state.document_filter = build_multi_document_filter(
self._document_filter
)
self._state = self._toolset.build_state_snapshot()
rag_state.document_filter = doc_filter
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if isinstance(analysis_state, AnalysisState):
analysis_state.document_filter = doc_filter
self._state = self._toolset.build_state_snapshot()

11
haiku_rag_slim/haiku/rag/client.py
View file

@ -1239,10 +1239,19 @@ class HaikuRAG:
context=context,
)
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import resolve_citations
agent = create_analysis_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output
output = result.output
citations = resolve_citations(output.cited_chunks, sandbox._search_results)
return AnalysisResult(
answer=output.answer,
program=output.program,
citations=citations,
)
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.

28
haiku_rag_slim/haiku/rag/skills/_tools.py
View file

@ -210,7 +210,7 @@ async def skill_analyze(
question: str,
document: str | None = None,
document_filter: str | None = None,
) -> tuple[str, str, str | None]:
) -> tuple[str, str, str | None, "list[Citation]"]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
@ -222,7 +222,7 @@ async def skill_analyze(
if result.program:
output += f"\n\nProgram:\n{result.program}"
return output, result.answer, result.program
return output, result.answer, result.program, result.citations
def update_documents_state(
@ -246,6 +246,15 @@ def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> An
return None
def _append_citations(state: Any, citations: "list[Citation]") -> None:
"""Index and append citations to a skill state's citations list."""
next_index = len(state.citations) + 1
for citation in citations:
citation.index = next_index
next_index += 1
state.citations.extend(citations)
def create_skill_extras(
db_path: Path,
config: AppConfig,
@ -406,11 +415,7 @@ def create_skill_tools(
)
if state:
next_index = len(state.citations) + 1
for citation in citations:
citation.index = next_index
next_index += 1
state.citations.extend(citations)
_append_citations(state, citations)
state.qa_history.append(
QAHistoryEntry(
question=question, answer=answer, citations=citations
@ -476,9 +481,11 @@ def create_skill_tools(
question: The question to answer.
document: Optional document ID or title to pre-load for analysis.
"""
from haiku.rag.utils import format_citations
state = _get_state(ctx, state_type)
state_filter = state.document_filter if state else None
output, answer, program = await skill_analyze(
output, answer, program, citations = await skill_analyze(
db_path,
config,
question,
@ -493,6 +500,11 @@ def create_skill_tools(
program=program,
)
)
if citations:
_append_citations(state, citations)
if citations:
output += "\n\n" + format_citations(citations)
return output

4
haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -2,8 +2,9 @@ import os
from functools import cache
from pathlib import Path
from pydantic import BaseModel
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.skills._tools import AnalysisEntry
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
@ -13,6 +14,7 @@ from haiku.skills.parser import parse_skill_md
class AnalysisState(BaseModel):
document_filter: str | None = None
analyses: list[AnalysisEntry] = []
citations: list[Citation] = Field(default_factory=list)
STATE_TYPE = AnalysisState

4
tests/agents/analysis/test_agent.py
View file

@ -5,7 +5,7 @@ from pydantic_ai import Agent
from haiku.rag.agents.analysis.agent import create_analysis_agent
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
from haiku.rag.config import AppConfig, Config
@ -19,7 +19,7 @@ class TestCreateAnalysisAgent:
agent = create_analysis_agent(Config)
assert isinstance(agent, Agent)
assert agent.deps_type is AnalysisDeps
assert agent.output_type is AnalysisResult
assert agent.output_type is RawAnalysisResult
def test_agent_has_execute_code_tool(self):
agent = create_analysis_agent(Config)