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

Remove qa_history ranking in favor of SessionContext, planner can skip searching if context adequate

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-22 10:21:35 +02:00
parent 3b7b7c1bbc
commit 74eac57ffa
No known key found for this signature in database
11 changed files with 39 additions and 2161 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
CHANGELOG.md
View file

@ -16,6 +16,7 @@
- Research graph receives compact context (~1,000-2,000 tokens) instead of raw qa_history (potentially thousands of tokens)
- New `session_context` field on `ChatSessionState` synced via AG-UI state protocol
- Chat TUI: New context modal (`Ctrl+O`) to view current session context
- Planner can now return empty sub_questions when context is sufficient to answer directly
### Changed
@ -28,6 +29,13 @@
- Removed `haiku.rag.embeddings.voyageai` module
- The `voyageai` extra now delegates to `pydantic-ai-slim[voyageai]`
### Removed
- **Q&A History Functions**: Removed unused conversation history utilities
- `rank_qa_history_by_similarity()` - chat agent now uses `SessionContext` instead of ranked Q&A pairs
- `format_conversation_context()` - no longer needed with SessionContext approach
- Associated embedding cache and helper functions also removed
## [0.26.9] - 2026-01-22
### Fixed
@ -35,6 +43,9 @@
- **v0.25.0 Migration Failure**: Fixed "Table 'documents' already exists" error during migration caused by held table references preventing `drop_table()` from succeeding. Added recovery logic to restore documents from staging table if a previous migration attempt failed mid-way.
## [0.26.8] - 2026-01-22
### Added
- **Jina Reranker v3**: Added support for Jina reranking with API mode (`provider: jina`) and local inference (`provider: jina-local`, requires `[jina]` extra)
- **Model Downloads**: `download-models` now pre-downloads HuggingFace models for `sentence-transformers`, `mxbai`, and `jina-local`
- **Reranker Factory**: Removed unreliable `id(config)`-based caching from `get_reranker()`; factory now always instantiates fresh

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

@ -10,7 +10,6 @@ from haiku.rag.agents.chat.state import (
SearchDeps,
SessionContext,
build_document_filter,
format_conversation_context,
)
__all__ = [
@ -24,7 +23,6 @@ __all__ = [
"SearchDeps",
"SessionContext",
"build_document_filter",
"format_conversation_context",
"summarize_session",
"update_session_context",
]

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

@ -14,11 +14,9 @@ from haiku.rag.agents.chat.state import (
CitationInfo,
QAResponse,
build_document_filter,
rank_qa_history_by_similarity,
)
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_conversational_graph
from haiku.rag.agents.research.models import Citation, SearchAnswer
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
@ -181,61 +179,19 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
# Build filter from document_name
doc_filter = build_document_filter(document_name) if document_name else None
# Filter and rank qa_history
ranked_history: list[QAResponse] = []
if ctx.deps.session_state and ctx.deps.session_state.qa_history:
# Step 1: Filter out low-confidence responses
filtered_history = [
qa for qa in ctx.deps.session_state.qa_history if qa.confidence >= 0.3
]
# Step 2: Rank filtered history by similarity to current question
embedder = ctx.deps.client.chunk_repository.embedder
ranked_history = await rank_qa_history_by_similarity(
current_question=question,
qa_history=filtered_history,
embedder=embedder,
top_k=5,
)
# Convert ranked qa_history to SearchAnswers for context seeding
existing_qa: list[SearchAnswer] = []
if ranked_history:
for qa in ranked_history:
citations = [
Citation(
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
for c in qa.citations
]
existing_qa.append(
SearchAnswer(
query=qa.question,
answer=qa.answer,
confidence=qa.confidence,
cited_chunks=[c.chunk_id for c in qa.citations],
citations=citations,
)
)
# Build and run the conversational research graph
graph = build_conversational_graph(config=ctx.deps.config)
# Determine background context:
# 1. Use session_context summary if available (compressed history)
# 2. Fall back to explicit background_context if set
# Determine context strategy:
# 1. If session_context exists, use compressed summary (skip raw qa_history)
# 2. Otherwise, fall back to explicit background_context
background_context: str | None = None
if ctx.deps.session_state:
if (
ctx.deps.session_state.session_context
and ctx.deps.session_state.session_context.summary
):
# Use compressed SessionContext - no need for raw qa_history
background_context = (
ctx.deps.session_state.session_context.render_markdown()
)
@ -244,7 +200,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
context = ResearchContext(
original_question=question,
qa_responses=existing_qa,
background_context=background_context,
)
state = ResearchState(

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

@ -1,31 +1,17 @@
import hashlib
from dataclasses import dataclass, field
from datetime import datetime
from typing import TYPE_CHECKING, Any
from typing import Any
import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel
from pydantic_ai import format_as_xml
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.embeddings import EmbedderWrapper
MAX_QA_HISTORY = 50
AGUI_STATE_KEY = "haiku.rag.chat"
_embedding_cache: dict[str, list[float]] = {}
def _qa_cache_key(question: str, answer: str) -> str:
"""Generate cache key from Q/A content."""
return hashlib.sha256(f"Q: {question}\nA: {answer}".encode()).hexdigest()
class CitationInfo(BaseModel):
"""Citation info for frontend display."""
@ -64,9 +50,7 @@ class SessionContext(BaseModel):
def render_markdown(self) -> str:
"""Render context for injection into research graph."""
if not self.summary:
return ""
return f"## Prior Conversation Context\n\n{self.summary}"
return self.summary
class ChatSessionState(BaseModel):
@ -79,96 +63,6 @@ class ChatSessionState(BaseModel):
session_context: SessionContext | None = None
def format_conversation_context(qa_history: list[QAResponse]) -> str:
"""Format conversation history as XML for inclusion in prompts."""
if not qa_history:
return ""
context_data = {
"previous_qa": [
{
"question": qa.question,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in qa_history
],
}
return format_as_xml(context_data, root_tag="conversation_context")
def _cosine_similarity(a: NDArray[np.float64], b: NDArray[np.float64]) -> float:
"""Compute cosine similarity between two vectors."""
return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
async def rank_qa_history_by_similarity(
current_question: str,
qa_history: list[QAResponse],
embedder: "EmbedderWrapper",
top_k: int = 5,
) -> list[QAResponse]:
"""Rank Q&A history by semantic similarity to current question.
Embeds question+answer pairs and returns the top-K most similar to the
current question. Falls back to returning the last top_k entries if
embedding fails.
Args:
current_question: The current question to compare against.
qa_history: List of previous Q&A pairs.
embedder: Embedder instance to use for embedding.
top_k: Maximum number of entries to return.
Returns:
Top-K Q&A pairs ranked by similarity to current question.
"""
if not qa_history:
return []
if len(qa_history) <= top_k:
return qa_history
# Embed current question
question_embedding = np.array(await embedder.embed_query(current_question))
# Check cache and collect uncached entries
qa_embeddings: list[list[float]] = []
uncached_indices: list[int] = []
uncached_texts: list[str] = []
for i, qa in enumerate(qa_history):
cache_key = _qa_cache_key(qa.question, qa.answer)
if cache_key in _embedding_cache:
qa_embeddings.append(_embedding_cache[cache_key])
else:
qa_embeddings.append([]) # placeholder
uncached_indices.append(i)
uncached_texts.append(f"Q: {qa.question}\nA: {qa.answer}")
# Embed only uncached entries
if uncached_texts:
new_embeddings = await embedder.embed_documents(uncached_texts)
for idx, embedding in zip(uncached_indices, new_embeddings):
qa = qa_history[idx]
cache_key = _qa_cache_key(qa.question, qa.answer)
_embedding_cache[cache_key] = embedding
qa_embeddings[idx] = embedding
# Compute similarities
similarities: list[tuple[int, float]] = []
for i, qa_emb in enumerate(qa_embeddings):
sim = _cosine_similarity(question_embedding, np.array(qa_emb))
similarities.append((i, sim))
# Sort by similarity (descending) and take top-K
similarities.sort(key=lambda x: x[1], reverse=True)
top_indices = sorted([idx for idx, _ in similarities[:top_k]])
# Return in original order
return [qa_history[i] for i in top_indices]
@dataclass
class ChatDeps:
"""Dependencies for chat agent.

2
haiku_rag_slim/haiku/rag/agents/research/models.py
View file

@ -17,8 +17,6 @@ class ResearchPlan(BaseModel):
@field_validator("sub_questions")
@classmethod
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) < 1:
raise ValueError("Must have at least 1 sub-question")
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v

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

@ -201,75 +201,6 @@ Scanned documents were excluded to avoid rotation and skewing issues.
"""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_with_qa_history_ranking(allow_model_requests, temp_db_path):
"""Test chat agent uses similarity ranking for qa_history.
This test verifies that when qa_history has more than 5 entries,
the ranking function is applied and the agent can still process requests.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a simple document
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
# Build session state with pre-populated qa_history (>5 items to trigger ranking)
session_state = ChatSessionState(
session_id="test-ranking",
qa_history=[
QAResponse(
question="What are the 11 class labels in DocLayNet?",
answer="The 11 class labels are: Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, and Title.",
),
QAResponse(
question="How was the annotation process organized?",
answer="The annotation was organized into 4 phases.",
),
QAResponse(
question="What data sources were used?",
answer="Sources include arXiv and government offices.",
),
QAResponse(
question="How were pages selected?",
answer="By selective subsampling.",
),
QAResponse(
question="What is the agreement metric?",
answer="The mAP metric was used.",
),
QAResponse(
question="What is machine learning?",
answer="A field of AI.",
),
],
)
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
)
# Ask a question - the key test is that ranking is applied without error
result = await agent.run(
"What class labels are defined in the dataset?",
deps=deps,
)
# Verify the agent produced a response (ranking didn't break anything)
assert result.output is not None
assert len(result.output) > 0
# Verify qa_history was updated (new Q&A was added)
assert len(session_state.qa_history) == 7 # 6 original + 1 new
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):

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

@ -41,14 +41,10 @@ class TestSessionContext:
assert ctx.render_markdown() == ""
def test_render_markdown_with_summary(self):
"""Test render_markdown returns formatted markdown."""
ctx = SessionContext(
summary="## Key Facts\n- Authentication uses JWT\n- Rate limit is 100/min"
)
result = ctx.render_markdown()
assert "## Prior Conversation Context" in result
assert "Authentication uses JWT" in result
assert "Rate limit is 100/min" in result
"""Test render_markdown returns the summary directly."""
summary = "## Key Facts\n- Authentication uses JWT\n- Rate limit is 100/min"
ctx = SessionContext(summary=summary)
assert ctx.render_markdown() == summary
def test_session_context_serialization_roundtrip(self):
"""Test SessionContext serializes and deserializes correctly."""

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

@ -1,224 +1,8 @@
from pathlib import Path
import pytest
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
CitationInfo,
QAResponse,
_embedding_cache,
_qa_cache_key,
build_document_filter,
format_conversation_context,
rank_qa_history_by_similarity,
)
from haiku.rag.client import HaikuRAG
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_state")
@pytest.mark.asyncio
async def test_rank_qa_history_empty():
"""Test empty history returns empty list."""
# Create a mock embedder - we won't actually call it
result = await rank_qa_history_by_similarity(
current_question="What is this?",
qa_history=[],
embedder=None, # type: ignore - won't be called for empty list
top_k=5,
)
assert result == []
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rank_qa_history_small_list(temp_db_path, allow_model_requests):
"""Test with history smaller than top_k returns all entries."""
async with HaikuRAG(temp_db_path, create=True) as client:
embedder = client.chunk_repository.embedder
# Create 3 Q&A pairs (less than top_k=5)
qa_history = [
QAResponse(question="What is Python?", answer="A programming language"),
QAResponse(question="What is Java?", answer="Another programming language"),
QAResponse(
question="What is Rust?", answer="A systems programming language"
),
]
result = await rank_qa_history_by_similarity(
current_question="Tell me about Python",
qa_history=qa_history,
embedder=embedder,
top_k=5,
)
# Should return all 3 entries since history < top_k
assert len(result) == 3
# All original entries should be present
assert set(qa.question for qa in result) == {
"What is Python?",
"What is Java?",
"What is Rust?",
}
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rank_qa_history_returns_top_k(temp_db_path, allow_model_requests):
"""Test ranking returns top-K most similar entries and populates cache."""
async with HaikuRAG(temp_db_path, create=True) as client:
embedder = client.chunk_repository.embedder
# Create 10 Q&A pairs on different topics
qa_history = [
QAResponse(
question="What are the 11 class labels in DocLayNet?",
answer="Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, and Title",
),
QAResponse(
question="How was the annotation process organized?",
answer="The process had 4 phases with 40 dedicated annotators",
),
QAResponse(
question="What data sources were used?",
answer="arXiv, government offices, company websites, financial reports and patents",
),
QAResponse(
question="How were pages selected?",
answer="By selective subsampling with bias towards pages with figures or tables",
),
QAResponse(
question="What is the inter-annotator agreement?",
answer="Computed as mAP@0.5-0.95 metric between pairwise annotations",
),
QAResponse(
question="What is machine learning?",
answer="A field of AI that enables systems to learn from data",
),
QAResponse(
question="How does neural network training work?",
answer="Through backpropagation and gradient descent",
),
QAResponse(
question="What is deep learning?",
answer="A subset of ML using neural networks with many layers",
),
]
# Clear cache to verify it gets populated
_embedding_cache.clear()
# Ask a question related to class labels (Q1)
result = await rank_qa_history_by_similarity(
current_question="Which class label has the highest count in DocLayNet?",
qa_history=qa_history,
embedder=embedder,
top_k=5,
)
# Should return exactly 5 entries
assert len(result) == 5
# The class labels Q&A should be in the top 5 (it's most semantically similar)
result_questions = [qa.question for qa in result]
assert "What are the 11 class labels in DocLayNet?" in result_questions
# Verify cache is populated for all Q/A pairs
for qa in qa_history:
cache_key = _qa_cache_key(qa.question, qa.answer)
assert cache_key in _embedding_cache
assert len(_embedding_cache[cache_key]) > 0
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rank_qa_history_preserves_order(temp_db_path, allow_model_requests):
"""Test that ranking preserves original order among selected items."""
async with HaikuRAG(temp_db_path, create=True) as client:
embedder = client.chunk_repository.embedder
# Create Q&A pairs where multiple are similar
qa_history = [
QAResponse(
question="What is Python?",
answer="A programming language",
citations=[
CitationInfo(
index=1,
document_id="doc1",
chunk_id="chunk1",
document_uri="python.md",
content="Python content",
)
],
),
QAResponse(
question="What is Java?",
answer="Another programming language",
),
QAResponse(
question="How to use Python for data science?",
answer="Use pandas, numpy, and scikit-learn",
),
]
result = await rank_qa_history_by_similarity(
current_question="Tell me about Python programming",
qa_history=qa_history,
embedder=embedder,
top_k=3,
)
# All should be returned
assert len(result) == 3
# The two Python-related questions should be in the results
result_questions = [qa.question for qa in result]
assert "What is Python?" in result_questions
assert "How to use Python for data science?" in result_questions
def test_format_conversation_context_empty():
"""Test format_conversation_context with empty history."""
result = format_conversation_context([])
assert result == ""
def test_format_conversation_context_with_history():
"""Test format_conversation_context formats qa_history as XML."""
citation = CitationInfo(
index=1,
document_id="doc-123",
chunk_id="chunk-456",
document_uri="test.md",
document_title="Test Document",
content="Test content",
)
qa_history = [
QAResponse(
question="What is Python?",
answer="A programming language",
citations=[citation],
),
QAResponse(
question="What is Java?",
answer="Another programming language",
),
]
result = format_conversation_context(qa_history)
assert "<conversation_context>" in result
assert "previous_qa" in result
assert "What is Python?" in result
assert "A programming language" in result
assert "What is Java?" in result
assert "Another programming language" in result
assert "Test Document" in result # source from first citation
def test_build_document_filter_simple():

18
tests/agents/research/test_research_graph.py
View file

@ -103,3 +103,21 @@ def test_format_conversational_context_for_prompt_excludes_background_when_none(
context = ResearchContext(original_question="What is X?")
result = format_conversational_context_for_prompt(context)
assert "<background>" not in result
def test_research_plan_allows_empty_sub_questions():
"""Test ResearchPlan accepts empty sub_questions when context is sufficient."""
from haiku.rag.agents.research.models import ResearchPlan
plan = ResearchPlan(sub_questions=[])
assert plan.sub_questions == []
def test_research_plan_rejects_too_many_sub_questions():
"""Test ResearchPlan rejects more than 12 sub_questions."""
from pydantic import ValidationError
from haiku.rag.agents.research.models import ResearchPlan
with pytest.raises(ValidationError, match="Cannot have more than 12"):
ResearchPlan(sub_questions=[f"q{i}" for i in range(13)])

1581
tests/cassettes/test_chat_agent/test_chat_agent_with_qa_history_ranking.yaml
View file

File diff suppressed because one or more lines are too long

126
tests/cassettes/test_chat_state/test_rank_qa_history_returns_top_k.yaml
View file

File diff suppressed because one or more lines are too long