From 696f8598eb5742f1790152178e1c92c6796affc9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 20 Feb 2026 11:37:55 +0200 Subject: [PATCH] Remove get_session_context, handle it automatically in ask() --- docs/apps.md | 2 +- docs/tools.md | 1 - haiku_rag_slim/haiku/rag/skills/rag.py | 112 ++++++------- haiku_rag_slim/haiku/rag/skills/rag/SKILL.md | 1 - tests/skills/test_rag.py | 160 +++++++++++-------- 5 files changed, 148 insertions(+), 128 deletions(-) diff --git a/docs/apps.md b/docs/apps.md index f3574ff9..2f0ac690 100644 --- a/docs/apps.md +++ b/docs/apps.md @@ -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 diff --git a/docs/tools.md b/docs/tools.md index 0f6fee83..eac09358 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py index 8eaa9186..296c0520 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag.py +++ b/haiku_rag_slim/haiku/rag/skills/rag.py @@ -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", diff --git a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md index 374c8a0b..f82ea967 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md +++ b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md @@ -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. diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index 1cd657f3..fb08576e 100644 --- a/tests/skills/test_rag.py +++ b/tests/skills/test_rag.py @@ -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