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

Remove get_session_context, handle it automatically in ask()

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-20 11:37:55 +02:00
parent bf724ad1ef
commit 696f8598eb
No known key found for this signature in database
5 changed files with 148 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/apps.md
View file

@ -46,7 +46,7 @@ Press `Ctrl+P` to open the command palette:
### Session Management
- Conversation history is maintained in memory for the session
- Previous Q/A pairs are used as context for follow-up questions via the `get_session_context` tool
- Previous Q/A pairs are automatically used as context for follow-up questions via the `ask` tool
- Citations are tracked per response and can be inspected
- Document filter restricts all searches to selected documents
- Clearing chat resets session state

1
docs/tools.md
View file

@ -44,7 +44,6 @@ Creates a RAG skill instance.
| `ask(question)` | Q&A with citations via the QA agent |
| `analyze(question, document?, filter?)` | Computational analysis via code execution (requires Docker) |
| `research(question)` | Deep multi-agent research producing comprehensive reports |
| `get_session_context(query)` | Retrieve relevant prior Q&A from the session |
### State

112
haiku_rag_slim/haiku/rag/skills/rag.py
View file

@ -56,6 +56,40 @@ def create_skill(
path = Path(__file__).parent / "rag"
metadata, instructions = parse_skill_md(path / "SKILL.md")
async def _find_relevant_prior_qa(
state: RAGState, query: str
) -> list[QAHistoryEntry]:
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.utils import cosine_similarity
if not state.qa_history:
return []
embedder = get_embedder(config)
query_embedding = await embedder.embed_query(query)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
state.qa_history[idx].question_embedding = new_embeddings[i]
matches = []
for qa in state.qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(query_embedding, qa.question_embedding)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matches.append(qa)
return matches
async def search(
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
) -> str:
@ -168,8 +202,31 @@ def create_skill(
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import format_citations
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
ask_question = question
if state:
matches = await _find_relevant_prior_qa(state, question)
if matches:
prior_parts = []
for qa in matches:
part = f"Q: {qa.question}\nA: {qa.answer}"
if qa.citations:
part += "\n" + format_citations(qa.citations)
prior_parts.append(part)
ask_question = (
"Context from prior questions in this session:\n\n"
+ "\n\n---\n\n".join(prior_parts)
+ "\n\n---\n\nCurrent question: "
+ question
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
answer, citations = await rag.ask(question)
answer, citations = await rag.ask(ask_question)
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
next_index = len(ctx.deps.state.citations) + 1
@ -218,58 +275,6 @@ def create_skill(
return output
async def get_session_context(ctx: RunContext[SkillRunDeps], query: str) -> str:
"""Retrieve relevant prior Q&A from the current session.
Call this before other tools when there may be prior questions
in the session that are relevant to the current query.
Args:
query: The current question or topic to find relevant context for.
"""
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.utils import cosine_similarity
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
if state is None or not state.qa_history:
return "No prior questions in this session."
embedder = get_embedder(config)
query_embedding = await embedder.embed_query(query)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
state.qa_history[idx].question_embedding = new_embeddings[i]
matches = []
for qa in state.qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(query_embedding, qa.question_embedding)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matches.append(qa)
if not matches:
return "No relevant prior questions found for this query."
parts = []
for qa in matches:
parts.append(f"Q: {qa.question}\nA: {qa.answer}")
return "Relevant prior Q&A:\n\n" + "\n\n---\n\n".join(parts)
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Conduct deep multi-agent research on a question.
@ -333,7 +338,6 @@ def create_skill(
ask,
analyze,
research,
get_session_context,
],
state_type=RAGState,
state_namespace="rag",

1
haiku_rag_slim/haiku/rag/skills/rag/SKILL.md
View file

@ -10,7 +10,6 @@ Use your tools to search and answer questions. Never make up information — alw
## How to decide which tool to use
- **get_session_context** — Call this first when there have been prior questions in the session. It finds relevant prior Q&A so you can avoid redundant searches and give more informed answers.
- **list_documents** — Use when the user wants to browse or see what documents are available (e.g., "what documents do you have?", "show me the documents", "list available docs").
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
- **search** — Use when the user wants to find relevant passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns matching chunks with metadata.

160
tests/skills/test_rag.py
View file

@ -31,7 +31,6 @@ class TestRAGSkillCreation:
"ask",
"analyze",
"research",
"get_session_context",
}
def test_create_skill_has_state(self, temp_db_path):
@ -253,6 +252,95 @@ class TestAskTool:
await ask(ctx, question="Second question")
assert state.citations[2].index == 3
async def test_ask_includes_prior_qa_context(self, rag_db, monkeypatch):
import random
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
captured_questions = []
async def mock_ask(self, question, **kwargs):
captured_questions.append(question)
return ("Answer about AI.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
# Pre-compute the embedding the fake embedder will produce for "Tell me about AI"
query_text = "Tell me about AI"
random.seed(hash(query_text) % (2**32))
query_embedding = [random.random() for _ in range(VECTOR_DIM)]
prior_citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
document_title="AI Overview",
content="AI content from source.",
)
]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is artificial intelligence?",
answer="AI is the simulation of human intelligence by machines.",
question_embedding=query_embedding,
citations=prior_citations,
),
]
)
ctx = _make_ctx(state)
await ask(ctx, question=query_text)
# rag.ask() should receive augmented question with prior context
assert len(captured_questions) == 1
augmented = captured_questions[0]
assert "Context from prior questions" in augmented
assert "What is artificial intelligence?" in augmented
assert "AI is the simulation" in augmented
assert "AI Overview" in augmented
assert query_text in augmented
# State should store the original question, not the augmented one
assert state.qa_history[-1].question == query_text
async def test_ask_no_prior_qa_context_when_irrelevant(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
captured_questions = []
async def mock_ask(self, question, **kwargs):
captured_questions.append(question)
return ("Answer.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
# Use orthogonal embedding — won't match the fake embedder's output
orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is the weather?",
answer="It is sunny today.",
question_embedding=orthogonal,
),
]
)
ctx = _make_ctx(state)
await ask(ctx, question="Explain quantum computing")
# rag.ask() should receive the original question unchanged
assert len(captured_questions) == 1
assert captured_questions[0] == "Explain quantum computing"
class TestAnalyzeTool:
async def test_analyze_returns_result(self, rag_db, monkeypatch):
@ -289,76 +377,6 @@ class TestAnalyzeTool:
assert state.qa_history[0].question == "How many documents?"
class TestGetSessionContextTool:
async def test_no_prior_questions(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
state = RAGState()
ctx = _make_ctx(state)
result = await get_ctx(ctx, query="What is AI?")
assert "no prior" in result.lower()
async def test_returns_relevant_prior_qa(self, rag_db):
import random
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
# Pre-compute the embedding that the fake embedder will produce
# for the query, so we can set it on the prior entry for high similarity
query_text = "Tell me about artificial intelligence"
random.seed(hash(query_text) % (2**32))
query_embedding = [random.random() for _ in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is artificial intelligence?",
answer="AI is the simulation of human intelligence by machines.",
question_embedding=query_embedding,
),
]
)
ctx = _make_ctx(state)
result = await get_ctx(ctx, query=query_text)
assert "artificial intelligence" in result.lower()
assert "simulation" in result.lower()
async def test_no_relevant_matches(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
# Use alternating ±1 embedding which is near-orthogonal to the
# all-positive vectors produced by the fake embedder
orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is the weather?",
answer="It is sunny today.",
question_embedding=orthogonal,
),
]
)
ctx = _make_ctx(state)
result = await get_ctx(ctx, query="Explain quantum computing")
assert "no relevant" in result.lower()
async def test_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
ctx = _make_ctx(state=None)
result = await get_ctx(ctx, query="What is AI?")
assert "no prior" in result.lower()
class TestResearchTool:
async def test_research_returns_report(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill