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

Remove SearchAgent & friends, consolidate duplicate models

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-06 13:40:23 +01:00
parent 84b320765a
commit df57b0cf43
No known key found for this signature in database
18 changed files with 81 additions and 3491 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
docs/agents.md
View file

@ -56,7 +56,7 @@ The chat agent enables multi-turn conversational RAG. It maintains session state
Key features:
- **Session memory**: Previous Q/A pairs are used as context for follow-up questions
- **Query expansion**: SearchAgent generates multiple query variations for better recall
- **Query expansion**: Search toolset generates multiple query variations for better recall
- **Document filtering**: Natural language document filtering ("search in document X about...")
- **Confidence filtering**: Low-confidence answers are flagged
@ -85,13 +85,14 @@ See [Applications](apps.md#chat-tui) for the full TUI interface guide.
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.agents.chat import create_chat_agent, ChatDeps, ChatSessionState
from haiku.rag.agents.chat import create_chat_agent, ChatDeps
from haiku.rag.tools import ToolContext
async with HaikuRAG(path_to_db) as client:
# Create agent and session
agent = create_chat_agent(config)
session = ChatSessionState()
deps = ChatDeps(client=client, config=config, session_state=session)
# Create agent with composed toolsets
context = ToolContext()
agent = create_chat_agent(config, client, context)
deps = ChatDeps(config=config, tool_context=context)
# First question
result = await agent.run("What is haiku.rag?", deps=deps)
@ -134,14 +135,16 @@ Q/A history is used to:
When using the chat agent with AG-UI streaming, state is emitted under a namespaced key to avoid conflicts with other agents:
```python
from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, ChatSessionState
from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps
from haiku.rag.tools import ToolContext
# AGUI_STATE_KEY = "haiku.rag.chat"
context = ToolContext()
agent = create_chat_agent(config, client, context)
deps = ChatDeps(
client=client,
config=config,
session_state=ChatSessionState(),
tool_context=context,
state_key=AGUI_STATE_KEY, # Enables namespaced state emission
)
```

8
haiku_rag_slim/haiku/rag/agents/__init__.py
View file

@ -1,9 +1,7 @@
from haiku.rag.agents.chat import (
ChatDeps,
ChatSessionState,
QAResponse,
SearchAgent,
SearchDeps,
QAHistoryEntry,
create_chat_agent,
)
from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent
@ -34,9 +32,7 @@ __all__ = [
"IterativePlanResult",
# Chat
"create_chat_agent",
"SearchAgent",
"ChatDeps",
"ChatSessionState",
"QAResponse",
"SearchDeps",
"QAHistoryEntry",
]

11
haiku_rag_slim/haiku/rag/agents/chat/__init__.py
View file

@ -8,30 +8,25 @@ from haiku.rag.agents.chat.context import (
summarize_session,
update_session_context,
)
from haiku.rag.agents.chat.search import SearchAgent
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse,
SearchDeps,
SessionContext,
)
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.document import DocumentInfo, DocumentListResponse
from haiku.rag.tools.qa import QAHistoryEntry
__all__ = [
"AGUI_STATE_KEY",
"create_chat_agent",
"run_chat_agent",
"trigger_background_summarization",
"SearchAgent",
"ChatDeps",
"ChatSessionState",
"DocumentInfo",
"DocumentListResponse",
"QAResponse",
"SearchDeps",
"QAHistoryEntry",
"SessionContext",
"ToolContext",
"summarize_session",

9
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -14,16 +14,13 @@ from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
DocumentInfo,
QAResponse,
SearchDeps,
SessionContext,
emit_state_event,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.document import DocumentListResponse, create_document_toolset
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
@ -275,10 +272,6 @@ __all__ = [
"trigger_background_summarization",
"ChatDeps",
"ChatSessionState",
"DocumentInfo",
"DocumentListResponse",
"QAResponse",
"SearchDeps",
"SessionContext",
"emit_state_event",
"AGUI_STATE_KEY",

21
haiku_rag_slim/haiku/rag/agents/chat/context.py
View file

@ -6,12 +6,12 @@ from typing import TYPE_CHECKING
from pydantic_ai import Agent
from haiku.rag.agents.chat.prompts import SESSION_SUMMARY_PROMPT
from haiku.rag.agents.chat.state import ChatSessionState, QAResponse, SessionContext
from haiku.rag.agents.chat.state import ChatSessionState, SessionContext
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
if TYPE_CHECKING:
from haiku.rag.tools.qa import QASessionState
from haiku.rag.tools.qa import QAHistoryEntry, QASessionState
@dataclass
@ -79,7 +79,7 @@ def get_cached_embedding(session_id: str, question: str) -> list[float] | None:
async def summarize_session(
qa_history: list[QAResponse],
qa_history: list["QAHistoryEntry"],
config: AppConfig,
current_context: str | None = None,
) -> str:
@ -113,7 +113,7 @@ async def summarize_session(
async def update_session_context(
qa_history: list[QAResponse],
qa_history: list["QAHistoryEntry"],
config: AppConfig,
session_state: ChatSessionState,
) -> None:
@ -143,7 +143,7 @@ async def update_session_context(
cache_session_context(session_state.session_id, session_state.session_context)
def _format_qa_history(qa_history: list[QAResponse]) -> str:
def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str:
"""Format qa_history for input to summarization."""
lines: list[str] = []
for i, qa in enumerate(qa_history, 1):
@ -165,16 +165,7 @@ async def _update_context_background(
) -> None:
"""Background task to update session context after an ask."""
try:
# Convert QAHistoryEntry to QAResponse format for update_session_context
qa_history = [
QAResponse(
question=entry.question,
answer=entry.answer,
confidence=entry.confidence,
citations=list(entry.citations),
)
for entry in qa_session_state.qa_history
]
qa_history = list(qa_session_state.qa_history)
session_state = ChatSessionState(
session_id=session_id,

25
haiku_rag_slim/haiku/rag/agents/chat/prompts.py
View file

@ -29,18 +29,6 @@ IMPORTANT - When user mentions a document in search/ask:
Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user."""
SEARCH_SYSTEM_PROMPT = """You are a search query optimizer. You MUST use the run_search tool to execute searches.
For each user request:
1. Use the run_search tool with the original query
2. Use run_search again with 1-2 alternative keyword queries
3. Keep all queries SHORT (2-5 words)
4. After all tool calls complete, respond "Search complete"
You can optionally specify a limit parameter (default 5).
IMPORTANT: You must make actual tool calls. Do not output "run_search(...)" as text."""
SESSION_SUMMARY_PROMPT = """You are a session summarizer. Given a conversation history of Q&A pairs (and optionally existing context), produce a structured summary that captures key information for future context.
If a "Current Context" section is provided at the start of the input, incorporate that context into your summary. This might be initial background context from the user or a previous summary - build upon it rather than discard it.
@ -60,16 +48,3 @@ Rules:
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself."""
DOCUMENT_SUMMARY_PROMPT = """Generate a summary of the document content provided below.
Start with a one-paragraph overview, then list the main topics covered, and highlight any key findings or conclusions.
Guidelines:
- Aim for 1-2 paragraphs for short documents, 3-4 paragraphs for longer ones
- Focus on factual content and key information
- Do not include meta-commentary like "This document discusses..." or "The document covers..."
- Do not speculate beyond what's in the content
Document content:
{content}"""

86
haiku_rag_slim/haiku/rag/agents/chat/search.py
View file

@ -1,86 +0,0 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.chat.prompts import SEARCH_SYSTEM_PROMPT
from haiku.rag.agents.chat.state import SearchDeps
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
class SearchAgent:
"""Agent that generates multiple queries and consolidates results."""
def __init__(self, client: HaikuRAG, config: AppConfig):
self._client = client
self._config = config
model = get_model(config.qa.model, config)
self._agent: Agent[SearchDeps, str] = Agent(
model,
deps_type=SearchDeps,
output_type=str,
instructions=SEARCH_SYSTEM_PROMPT,
retries=3,
)
@self._agent.tool
async def run_search(
ctx: RunContext[SearchDeps],
query: str,
limit: int | None = None,
) -> str:
"""Run a single search query against the knowledge base.
Args:
query: The search query
limit: Number of results to fetch (default: 5)
"""
effective_limit = limit or 5
results = await ctx.deps.client.search(
query, limit=effective_limit, filter=ctx.deps.filter
)
results = await ctx.deps.client.expand_context(results)
ctx.deps.search_results.extend(results)
if not results:
return f"No results for: {query}"
return f"Found {len(results)} results for: {query}"
async def search(
self,
query: str,
context: str | None = None,
filter: str | None = None,
limit: int | None = None,
) -> list[SearchResult]:
"""Execute search with query expansion and deduplication.
Args:
query: The user's search request
context: Optional conversation context
filter: Optional SQL WHERE clause to filter documents
limit: Maximum number of results to return (default: config limit)
Returns:
Deduplicated list of SearchResult sorted by score
"""
prompt = query
if context:
prompt = f"Context: {context}\n\nSearch request: {query}"
deps = SearchDeps(client=self._client, config=self._config, filter=filter)
await self._agent.run(prompt, deps=deps)
# Deduplicate by chunk_id, keeping highest score
seen: dict[str, SearchResult] = {}
for result in deps.search_results:
chunk_id = result.chunk_id or ""
if chunk_id not in seen or result.score > seen[chunk_id].score:
seen[chunk_id] = result
# Sort by score descending and apply limit
effective_limit = limit or self._config.search.limit
return sorted(seen.values(), key=lambda r: r.score, reverse=True)[
:effective_limit
]

122
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

@ -1,65 +1,18 @@
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
from typing import TYPE_CHECKING
import jsonpatch
from ag_ui.core import EventType, StateDeltaEvent
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation, SearchAnswer
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.agents.research.models import Citation
MAX_QA_HISTORY = 50
if TYPE_CHECKING:
from haiku.rag.tools.qa import QAHistoryEntry
AGUI_STATE_KEY = "haiku.rag.chat"
class QAResponse(BaseModel):
"""A Q&A pair from conversation history with citations."""
question: str
answer: str
confidence: float = 0.9
citations: list[Citation] = []
question_embedding: list[float] | None = Field(default=None, exclude=True)
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
def to_search_answer(self) -> SearchAnswer:
"""Convert to SearchAnswer for research graph context."""
return SearchAnswer(
query=self.question,
answer=self.answer,
confidence=self.confidence,
cited_chunks=[c.chunk_id for c in self.citations],
citations=self.citations,
)
class DocumentInfo(BaseModel):
"""Document info for list_documents response."""
title: str
uri: str
created: str
class DocumentListResponse(BaseModel):
"""Response from list_documents tool."""
documents: list[DocumentInfo]
page: int
total_pages: int
total_documents: int
class SessionContext(BaseModel):
"""Compressed summary of conversation history for research graph."""
@ -77,7 +30,7 @@ class ChatSessionState(BaseModel):
session_id: str = ""
initial_context: str | None = None
citations: list[Citation] = []
qa_history: list[QAResponse] = []
qa_history: list["QAHistoryEntry"] = []
session_context: SessionContext | None = None
document_filter: list[str] = []
citation_registry: dict[str, int] = {}
@ -97,69 +50,14 @@ class ChatSessionState(BaseModel):
return new_index
@dataclass
class ChatDeps:
"""Dependencies for chat agent.
def _rebuild_models(qa_history_entry_cls: type) -> None:
"""Resolve ChatSessionState forward reference to QAHistoryEntry.
Implements StateHandler protocol for AG-UI state management.
Must be called after QAHistoryEntry is defined, passing the class.
"""
client: HaikuRAG
config: AppConfig
search_results: list[SearchResult] | None = None
session_state: ChatSessionState = field(
default_factory=lambda: ChatSessionState(session_id="")
ChatSessionState.model_rebuild(
_types_namespace={"QAHistoryEntry": qa_history_entry_cls}
)
state_key: str | None = None
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol."""
snapshot = self.session_state.model_dump()
if self.state_key:
return {self.state_key: snapshot}
return snapshot
@state.setter
def state(self, value: dict[str, Any] | None) -> None:
"""Set state from AG-UI protocol."""
if value is None:
return
# Extract from namespaced key if present
state_data: dict[str, Any] = value
if self.state_key and self.state_key in value:
nested = value[self.state_key]
if isinstance(nested, dict):
state_data = nested
# Update session_state from incoming state
if "qa_history" in state_data:
self.session_state.qa_history = [
QAResponse(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
if "citations" in state_data:
self.session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
if state_data.get("session_id"):
self.session_state.session_id = state_data["session_id"]
if "document_filter" in state_data:
self.session_state.document_filter = state_data.get("document_filter", [])
if "citation_registry" in state_data:
self.session_state.citation_registry = state_data["citation_registry"]
if "initial_context" in state_data:
self.session_state.initial_context = state_data.get("initial_context")
@dataclass
class SearchDeps:
"""Dependencies for search agent."""
client: HaikuRAG
config: AppConfig
filter: str | None = None
search_results: list[SearchResult] = field(default_factory=list)
def emit_state_event(

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

@ -203,10 +203,10 @@ class ChatApp(App):
for c in chat_state["citations"]
]
if "qa_history" in chat_state:
from haiku.rag.agents.chat.state import QAResponse
from haiku.rag.tools.qa import QAHistoryEntry
self.session_state.qa_history = [
QAResponse(**qa) if isinstance(qa, dict) else qa
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
for qa in chat_state["qa_history"]
]
if "session_context" in chat_state:

13
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -49,6 +49,13 @@ class QAHistoryEntry(BaseModel):
citations: list[Citation] = []
question_embedding: list[float] | None = Field(default=None, exclude=True)
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
def to_search_answer(self) -> SearchAnswer:
"""Convert to SearchAnswer for research graph context."""
return SearchAnswer(
@ -60,6 +67,12 @@ class QAHistoryEntry(BaseModel):
)
# Resolve ChatSessionState forward reference to QAHistoryEntry
from haiku.rag.agents.chat.state import _rebuild_models # noqa: E402
_rebuild_models(QAHistoryEntry)
class QAState(BaseModel):
"""State for QA toolset.

174
tests/agents/chat/test_chat_agent.py
View file

@ -7,16 +7,15 @@ from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
QAResponse,
SearchAgent,
QAHistoryEntry,
ToolContext,
create_chat_agent,
)
from haiku.rag.agents.chat.context import get_cached_session_context
from haiku.rag.agents.chat.state import MAX_QA_HISTORY
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.tools.qa import MAX_QA_HISTORY
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
@ -226,7 +225,7 @@ def test_citation():
def test_qa_response():
"""Test QAResponse model."""
"""Test QAHistoryEntry model."""
citation = Citation(
index=1,
document_id="doc-123",
@ -235,7 +234,7 @@ def test_qa_response():
document_title="Test Document",
content="Test content",
)
qa = QAResponse(
qa = QAHistoryEntry(
question="What is this?",
answer="This is a test",
confidence=0.95,
@ -249,7 +248,7 @@ def test_qa_response():
def test_qa_response_sources_with_uri_fallback():
"""Test QAResponse.sources falls back to URI when title is None."""
"""Test QAHistoryEntry.sources falls back to URI when title is None."""
citation = Citation(
index=1,
document_id="doc-123",
@ -258,7 +257,7 @@ def test_qa_response_sources_with_uri_fallback():
document_title=None,
content="Test content",
)
qa = QAResponse(
qa = QAHistoryEntry(
question="What is this?",
answer="This is a test",
citations=[citation],
@ -267,7 +266,7 @@ def test_qa_response_sources_with_uri_fallback():
def test_qa_response_to_search_answer():
"""Test QAResponse.to_search_answer() converts to SearchAnswer for research graph."""
"""Test QAHistoryEntry.to_search_answer() converts to SearchAnswer for research graph."""
citation = Citation(
index=1,
document_id="doc-123",
@ -276,7 +275,7 @@ def test_qa_response_to_search_answer():
document_title="Test Document",
content="Test content",
)
qa = QAResponse(
qa = QAHistoryEntry(
question="What is the answer?",
answer="The answer is 42",
confidence=0.95,
@ -293,14 +292,6 @@ def test_qa_response_to_search_answer():
assert search_answer.citations[0].chunk_id == "chunk-456"
def test_search_agent_initialization(temp_db_path):
"""Test SearchAgent can be initialized."""
client = HaikuRAG(temp_db_path, create=True)
search_agent = SearchAgent(client, Config)
assert search_agent is not None
client.close()
# DocLayNet content for testing
DOCLAYNET_CLASS_LABELS = """
DocLayNet Dataset - Class Labels
@ -467,108 +458,6 @@ async def test_chat_agent_get_document_not_found(allow_model_requests, temp_db_p
assert result.output is not None
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_agent_with_context(allow_model_requests, temp_db_path):
"""Test SearchAgent's search method with context."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
search_agent = SearchAgent(client, Config)
# Search with context
results = await search_agent.search(
query="What are the class labels?",
context="We're discussing document layout analysis",
)
assert isinstance(results, list)
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_agent_with_filter(allow_model_requests, temp_db_path):
"""Test SearchAgent's search method with document filter."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
search_agent = SearchAgent(client, Config)
# Search with filter - only the labels document
results = await search_agent.search(
query="What information is available?",
filter="uri LIKE '%labels%'",
)
assert isinstance(results, list)
# Results should only come from the labels document
for r in results:
assert "labels" in (r.document_uri or "")
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_agent_deduplication(allow_model_requests, temp_db_path):
"""Test SearchAgent deduplicates results by chunk_id."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
search_agent = SearchAgent(client, Config)
# Search - the search agent will likely run multiple queries
# that could return the same chunk, which should be deduplicated
results = await search_agent.search(
query="Tell me about class labels and their counts",
)
assert isinstance(results, list)
# Verify no duplicate chunk_ids
chunk_ids = [r.chunk_id for r in results if r.chunk_id]
assert len(chunk_ids) == len(set(chunk_ids)), "Found duplicate chunk_ids"
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_agent_no_results(allow_model_requests, temp_db_path):
"""Test SearchAgent handles no results gracefully."""
async with HaikuRAG(temp_db_path, create=True) as client:
search_agent = SearchAgent(client, Config)
# Search in empty database
results = await search_agent.search(
query="Find information about nonexistent topic xyz123",
)
assert isinstance(results, list)
assert len(results) == 0
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path):
@ -728,7 +617,7 @@ def test_fifo_limit_enforcement():
"""
# Create a session state with MAX_QA_HISTORY + 1 entries
qa_history = [
QAResponse(
QAHistoryEntry(
question=f"Question {i}",
answer=f"Answer {i}",
confidence=0.9,
@ -819,43 +708,6 @@ async def test_chat_agent_search_with_session_filter(
)
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_agent_with_session_filter(allow_model_requests, temp_db_path):
"""Test SearchAgent respects session document filter."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add two distinct documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
from haiku.rag.tools.filters import build_multi_document_filter
search_agent = SearchAgent(client, Config)
# Build filter for only the labels document
doc_filter = build_multi_document_filter(["DocLayNet Class Labels"])
results = await search_agent.search(
query="What information is available?",
filter=doc_filter,
)
assert isinstance(results, list)
# All results should be from the labels document
for r in results:
assert "labels" in (r.document_uri or "").lower() or "Labels" in (
r.document_title or ""
)
def test_ask_tool_citation_registry_logic():
"""Test the citation index assignment logic used by the ask tool.
@ -1033,9 +885,9 @@ def test_prior_answer_matching_below_threshold():
def test_qa_response_embedding_cache():
"""Test that QAResponse stores and retrieves question_embedding correctly."""
"""Test that QAHistoryEntry stores and retrieves question_embedding correctly."""
embedding = [0.1, 0.2, 0.3, 0.4]
qa = QAResponse(
qa = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.9,
@ -1049,8 +901,8 @@ def test_qa_response_embedding_cache():
def test_qa_response_embedding_default_none():
"""Test that QAResponse.question_embedding defaults to None."""
qa = QAResponse(
"""Test that QAHistoryEntry.question_embedding defaults to None."""
qa = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.9,

24
tests/agents/chat/test_context.py
View file

@ -3,12 +3,10 @@ from pathlib import Path
import pytest
from haiku.rag.agents.chat.state import (
QAResponse,
SessionContext,
)
from haiku.rag.agents.chat.state import SessionContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.config import Config
from haiku.rag.tools.qa import QAHistoryEntry
@pytest.fixture(scope="module")
@ -82,7 +80,7 @@ class TestSummarizeSession:
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens for authentication.",
confidence=0.95,
@ -115,7 +113,7 @@ class TestSummarizeSession:
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens for authentication.",
confidence=0.95,
@ -130,7 +128,7 @@ class TestSummarizeSession:
)
],
),
QAResponse(
QAHistoryEntry(
question="What is the rate limit?",
answer="Rate limiting is set to 100 requests per minute.",
confidence=0.9,
@ -145,7 +143,7 @@ class TestSummarizeSession:
)
],
),
QAResponse(
QAHistoryEntry(
question="How do I refresh tokens?",
answer="Use the /refresh endpoint with your refresh token.",
confidence=0.85,
@ -178,7 +176,7 @@ class TestSummarizeSession:
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAResponse(
QAHistoryEntry(
question="What's the rate limit?",
answer="100 requests per minute.",
confidence=0.9,
@ -218,7 +216,7 @@ class TestUpdateSessionContext:
session_state = ChatSessionState(session_id="test-session")
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens.",
confidence=0.95,
@ -339,7 +337,7 @@ class TestSessionContextCache:
session_state = ChatSessionState(session_id="cache-test-session")
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is Python?",
answer="A programming language.",
confidence=0.95,
@ -409,7 +407,7 @@ class TestSessionContextCache:
)
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is JWT?",
answer="JSON Web Token for authentication.",
confidence=0.95,
@ -461,7 +459,7 @@ class TestSessionContextCache:
)
qa_history = [
QAResponse(
QAHistoryEntry(
question="What is JWT?",
answer="JSON Web Token.",
confidence=0.95,

437
tests/agents/chat/test_state.py
View file

@ -1,9 +1,7 @@
from ag_ui.core import StateDeltaEvent
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
ChatSessionState,
QAResponse,
SessionContext,
)
from haiku.rag.tools.filters import (
@ -11,6 +9,7 @@ from haiku.rag.tools.filters import (
build_multi_document_filter,
combine_filters,
)
from haiku.rag.tools.qa import QAHistoryEntry
def test_build_document_filter_simple():
@ -89,304 +88,11 @@ def test_combine_filters_both():
def test_max_qa_history_constant():
"""Test MAX_QA_HISTORY constant value."""
from haiku.rag.tools.qa import MAX_QA_HISTORY
assert MAX_QA_HISTORY == 50
def test_chat_deps_state_getter_returns_namespaced_state():
"""Test ChatDeps.state getter returns state under namespaced key."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(
session_id="test-123",
qa_history=[
QAResponse(question="Q1", answer="A1", confidence=0.9),
],
)
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
state = deps.state
assert state is not None
assert AGUI_STATE_KEY in state
assert state[AGUI_STATE_KEY]["session_id"] == "test-123"
assert len(state[AGUI_STATE_KEY]["qa_history"]) == 1
assert state[AGUI_STATE_KEY]["qa_history"][0]["question"] == "Q1"
def test_chat_deps_state_getter_without_namespace():
"""Test ChatDeps.state getter returns flat state when no state_key."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="test-123")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=None,
)
state = deps.state
assert state is not None
assert "session_id" in state
assert state["session_id"] == "test-123"
def test_chat_deps_state_getter_returns_default_state():
"""Test ChatDeps.state getter returns default state when not explicitly set."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps
mock_client = MagicMock()
mock_config = MagicMock()
deps = ChatDeps(
client=mock_client,
config=mock_config,
)
state = deps.state
assert state is not None
assert "session_id" in state
assert state["qa_history"] == []
assert state["citations"] == []
def test_chat_deps_state_setter_updates_from_namespaced_state():
"""Test ChatDeps.state setter updates session_state from namespaced incoming state."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="initial")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Simulate incoming AG-UI state with namespaced key
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "updated-123",
"qa_history": [
{"question": "Q1", "answer": "A1", "confidence": 0.9, "citations": []}
],
"citations": [],
}
}
deps.state = incoming_state
assert deps.session_state is not None
assert deps.session_state.session_id == "updated-123"
assert len(deps.session_state.qa_history) == 1
assert deps.session_state.qa_history[0].question == "Q1"
def test_chat_deps_state_setter_handles_none():
"""Test ChatDeps.state setter handles None gracefully."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="original")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
)
# Setting None should not raise and should not change state
deps.state = None
assert deps.session_state is not None
assert deps.session_state.session_id == "original"
def test_chat_deps_state_setter_updates_default_state():
"""Test ChatDeps.state setter updates the default session_state."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps
mock_client = MagicMock()
mock_config = MagicMock()
deps = ChatDeps(
client=mock_client,
config=mock_config,
)
original_session_id = deps.session_state.session_id
# Update with incoming state
deps.state = {"session_id": "updated-123", "qa_history": [], "citations": []}
assert deps.session_state.session_id == "updated-123"
assert deps.session_state.session_id != original_session_id
def test_chat_deps_state_setter_with_citation_dicts():
"""Test ChatDeps.state setter converts citation dicts to Citation."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="test")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "test",
"qa_history": [],
"citations": [
{
"index": 1,
"document_id": "doc-1",
"chunk_id": "chunk-1",
"document_uri": "test.md",
"document_title": "Test Doc",
"page_numbers": [1, 2],
"headings": ["Intro"],
"content": "Test content",
}
],
}
}
deps.state = incoming_state
assert deps.session_state is not None
assert len(deps.session_state.citations) == 1
citation = deps.session_state.citations[0]
assert citation.document_id == "doc-1"
assert citation.chunk_id == "chunk-1"
assert citation.page_numbers == [1, 2]
def test_chat_deps_state_getter_includes_session_context():
"""Test ChatDeps.state getter includes session_context when present."""
from datetime import datetime
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
SessionContext,
)
mock_client = MagicMock()
mock_config = MagicMock()
now = datetime.now()
session_state = ChatSessionState(
session_id="test-123",
session_context=SessionContext(
summary="User discussed authentication.",
last_updated=now,
),
)
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
state = deps.state
assert state is not None
assert AGUI_STATE_KEY in state
assert state[AGUI_STATE_KEY]["session_context"] is not None
assert (
state[AGUI_STATE_KEY]["session_context"]["summary"]
== "User discussed authentication."
)
def test_chat_deps_state_setter_ignores_session_context():
"""Test ChatDeps.state setter ignores session_context from client.
The agent owns session_context via server-side cache, so client-provided
session_context should be ignored to prevent stale state overwriting.
"""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
SessionContext,
)
mock_client = MagicMock()
mock_config = MagicMock()
# Start with a session_context (e.g., from cache)
session_state = ChatSessionState(
session_id="test",
session_context=SessionContext(summary="Server-side context"),
)
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Client sends different session_context (stale)
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "test",
"qa_history": [],
"citations": [],
"session_context": {
"summary": "Client-provided stale context",
"last_updated": "2025-01-15T10:30:00",
},
}
}
deps.state = incoming_state
# session_context should NOT be overwritten
assert deps.session_state is not None
assert deps.session_state.session_context is not None
assert deps.session_state.session_context.summary == "Server-side context"
def test_citation_registry_index_assignment():
"""Test get_or_assign_index basic index assignment behavior.
@ -395,8 +101,6 @@ def test_citation_registry_index_assignment():
- Second unique chunk gets index 2
- Same chunk_id always returns same index
"""
from haiku.rag.agents.chat.state import ChatSessionState
session_state = ChatSessionState(session_id="test")
# First chunk gets index 1
@ -414,8 +118,6 @@ def test_citation_registry_index_assignment():
def test_citation_registry_stability():
"""Test citation indices are stable across multiple calls in any order."""
from haiku.rag.agents.chat.state import ChatSessionState
session_state = ChatSessionState(session_id="test")
# First round assigns indices 1, 2, 3
@ -435,8 +137,6 @@ def test_citation_registry_stability():
def test_citation_registry_serialization_roundtrip():
"""Test citation_registry serializes and deserializes correctly for AG-UI state."""
from haiku.rag.agents.chat.state import ChatSessionState
# Create state and assign indices
original = ChatSessionState(session_id="test")
original.get_or_assign_index("chunk-a")
@ -457,127 +157,6 @@ def test_citation_registry_serialization_roundtrip():
assert restored.get_or_assign_index("chunk-c") == 3
def test_chat_deps_state_getter_includes_citation_registry():
"""Test ChatDeps.state getter includes citation_registry."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="test")
session_state.get_or_assign_index("chunk-a")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
state = deps.state
assert state is not None
assert AGUI_STATE_KEY in state
assert state[AGUI_STATE_KEY]["citation_registry"] == {"chunk-a": 1}
def test_chat_deps_state_setter_restores_citation_registry():
"""Test ChatDeps.state setter restores citation_registry from incoming state."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="test")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Simulate incoming AG-UI state with citation_registry
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "test",
"qa_history": [],
"citations": [],
"citation_registry": {"chunk-x": 1, "chunk-y": 2},
}
}
deps.state = incoming_state
assert deps.session_state is not None
# Registry should be restored
assert deps.session_state.get_or_assign_index("chunk-x") == 1
assert deps.session_state.get_or_assign_index("chunk-y") == 2
# New chunk gets next index
assert deps.session_state.get_or_assign_index("chunk-z") == 3
def test_chat_deps_state_setter_restores_document_filter():
"""Test ChatDeps.state setter restores document_filter from incoming state."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(session_id="test")
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "test",
"qa_history": [],
"citations": [],
"document_filter": ["doc1.pdf", "doc2.pdf"],
}
}
deps.state = incoming_state
assert deps.session_state is not None
assert deps.session_state.document_filter == ["doc1.pdf", "doc2.pdf"]
def test_chat_deps_state_getter_includes_document_filter():
"""Test ChatDeps.state getter includes document_filter."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
mock_client = MagicMock()
mock_config = MagicMock()
session_state = ChatSessionState(
session_id="test-123",
document_filter=["doc1.pdf", "doc2.pdf"],
)
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
state = deps.state
assert state is not None
assert AGUI_STATE_KEY in state
assert state[AGUI_STATE_KEY]["document_filter"] == ["doc1.pdf", "doc2.pdf"]
def test_chat_session_state_defaults_to_empty_session_id():
"""New ChatSessionState should default to empty session_id.
@ -685,7 +264,7 @@ def test_emit_state_event_returns_delta_with_changes():
current_state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
new_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
citations=[],
)
@ -709,7 +288,7 @@ def test_emit_state_event_delta_with_state_key():
current_state = ChatSessionState(session_id="test-123", qa_history=[])
new_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
)
event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY)
@ -728,14 +307,14 @@ def test_emit_state_event_delta_produces_valid_patch():
current_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
citations=[],
)
new_state = ChatSessionState(
session_id="test-123",
qa_history=[
QAResponse(question="Q1", answer="A1", confidence=0.9),
QAResponse(question="Q2", answer="A2", confidence=0.8),
QAHistoryEntry(question="Q1", answer="A1", confidence=0.9),
QAHistoryEntry(question="Q2", answer="A2", confidence=0.8),
],
citations=[],
)

601
tests/cassettes/test_chat_agent/test_search_agent_deduplication.yaml
View file

File diff suppressed because one or more lines are too long

386
tests/cassettes/test_chat_agent/test_search_agent_no_results.yaml
View file

File diff suppressed because one or more lines are too long

492
tests/cassettes/test_chat_agent/test_search_agent_with_context.yaml
View file

File diff suppressed because one or more lines are too long

488
tests/cassettes/test_chat_agent/test_search_agent_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

650
tests/cassettes/test_chat_agent/test_search_agent_with_session_filter.yaml
View file

File diff suppressed because one or more lines are too long