diff --git a/app/backend/main.py b/app/backend/main.py index e91968b0..0c8ebb80 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -15,12 +15,12 @@ from starlette.routing import Route from haiku.rag.agents.chat import ( AGUI_STATE_KEY, ChatDeps, - ToolContext, create_chat_agent, ) from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig +from haiku.rag.tools import ToolContext load_dotenv(find_dotenv(usecwd=True)) @@ -109,13 +109,9 @@ async def stream_chat(request: Request) -> Response: async def health_check(_: Request) -> JSONResponse: """Health check endpoint.""" - # Create a temporary agent just for health check - context = ToolContext() - agent = create_chat_agent(Config, get_client(), context) return JSONResponse( { "status": "healthy", - "agent_model": str(agent.model), "qa_provider": Config.qa.model.provider, "qa_model": Config.qa.model.name, "db_path": str(db_path), diff --git a/docs/agents.md b/docs/agents.md index 68a6f14b..5b3f47ac 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -51,25 +51,14 @@ async with HaikuRAG(path_to_db) as client: ## Chat Agent -The chat agent enables multi-turn conversational RAG. It maintains session state including Q/A history and uses that context to improve follow-up answers. +The chat agent enables multi-turn conversational RAG. It is built from composable [toolsets](tools.md) and maintains session state to improve follow-up answers. Key features: -- **Composable toolsets**: Built from reusable `FunctionToolset` factories in `haiku.rag.tools` -- **Semantic prior answer recall**: The `ask` tool embeds each question and matches it against conversation history — relevant prior answers are passed to the research planner, which can skip searching when they suffice -- **Background summarization**: After each `ask` call, an async background task summarizes the QA history into a compact session context, cached server-side for the next request -- **Session context injection**: The session context summary flows into the research planner as `` XML, letting it resolve ambiguous references ("How does *it* handle X?") -- **Document filtering**: Natural language document filtering ("search in document X about...") - -### Tools - -The chat agent composes five tools from `haiku.rag.tools`: - -- `list_documents` — Browse available documents in the knowledge base -- `summarize_document` — Generate a summary of a specific document -- `get_document` — Retrieve a specific document by title or URI -- `search` — Hybrid search with query expansion and optional document filter -- `ask` — Answer questions using the conversational research graph +- **Composable toolsets**: Built from reusable `FunctionToolset` factories — see [Toolsets](tools.md) +- **Semantic prior answer recall**: Similar prior Q/A pairs are retrieved and passed to the research planner, which can skip searching when they suffice +- **Background summarization**: After each `ask` call, the QA history is summarized into a compact session context for the next request +- **Document filtering**: Session-level or per-query document filtering ### CLI Usage @@ -134,93 +123,31 @@ Available features: | QA | `FEATURE_QA` | `ask` | | Analysis | `FEATURE_ANALYSIS` | `analyze` | -The system prompt is automatically composed to match the selected features — only guidance for active tools is included. `SessionState` is always registered (shared by all features). `QASessionState` is only registered when the QA feature is active. +The system prompt is automatically composed to match the selected features. See [Toolsets](tools.md) for details on each toolset's parameters and behavior. ### Session State The `ChatSessionState` maintains: - `session_id` — Unique identifier for the session -- `qa_history` — List of previous Q/A pairs (FIFO, max 50) +- `qa_history` — List of previous Q/A pairs - `session_context` — Automatically maintained session context summary - `document_filter` — List of document titles/URIs to restrict searches - `citation_registry` — Stable mapping of chunk IDs to citation indices **Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source. -```python -# Example: citation indices are stable across calls -state = ChatSessionState() - -# First call returns citations [1], [2], [3] -# Second call reuses [1] if same chunk, assigns [4], [5] for new chunks -# User can reference [1] in follow-up and it still refers to original source -``` - ### Conversational Memory -The chat agent maintains two layers of conversational memory that work together: +The chat agent maintains two layers of conversational memory: **1. Semantic prior answer recall** -When the `ask` tool receives a question, it: - -1. Embeds the new question -2. Computes cosine similarity against all cached `qa_history` embeddings -3. Selects prior answers above a 0.7 similarity threshold -4. Passes them as `prior_answers` to the research graph's `ResearchContext` - -The research planner sees these as `` in its prompt. If the prior answers already cover the question, the planner marks research as complete and skips directly to synthesis — no new searches needed. - -Question embeddings are cached per-session to avoid re-embedding on every turn. Uncached embeddings are batch-computed. +When the `ask` tool receives a question, it embeds the question and compares it against prior Q/A embeddings. Sufficiently similar prior answers are passed to the research planner, which can skip searching entirely if they already cover the question. **2. Background session summarization** -After each `ask` call completes, `trigger_background_summarization()` spawns an async task that: - -1. Formats the full `qa_history` (questions, answers, confidence, sources) as markdown -2. Sends it to an LLM with the current session context (if any) as input -3. Produces a compact summary capturing key facts, entities, and document references -4. Caches the result server-side under the `session_id` - -If a new `ask` fires before the previous summarization finishes, the old task is cancelled. On the next request, the cached summary is picked up and injected as `session_context` into the research planner's `` XML. - -This means follow-up questions like "Tell me more about the authentication part" resolve correctly even though the planner never saw the original conversation — it has the summary. - -### AG-UI Integration - -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 -from haiku.rag.tools import ToolContext - -# AGUI_STATE_KEY = "haiku.rag.chat" - -context = ToolContext() -agent = create_chat_agent(config, client, context) -deps = ChatDeps( - config=config, - tool_context=context, - state_key=AGUI_STATE_KEY, # Enables namespaced state emission -) -``` - -The emitted state structure: - -```json -{ - "haiku.rag.chat": { - "session_id": "", - "citations": [...], - "qa_history": [...], - "document_filter": [...], - "citation_registry": {"chunk-id-1": 1, "chunk-id-2": 2} - } -} -``` - -Frontend clients should extract state from under this key. See the [Web Application](apps.md#web-application) for a complete implementation example. +After each `ask` call, a background task summarizes the full QA history into a compact session context. This summary is injected into the research planner on the next request, allowing it to resolve ambiguous references ("Tell me more about the authentication part") without having seen the full conversation. ## Research Graph diff --git a/docs/architecture.md b/docs/architecture.md index 3de8ba9f..00417156 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -108,20 +108,19 @@ flowchart TB end subgraph Chat["Chat Agent"] - Q2[Question] --> Expand[Query Expansion] - Expand --> S2[Search/Ask] + Q2[Question] --> Tools[Tool Selection] + Tools --> S2[Search / Ask / Get] S2 --> A2[Answer] A2 --> History[Session History] History -.-> Q2 end subgraph Research["Research Graph"] - Q3[Question] --> Plan[Plan] - Plan --> Batch[Get Batch] - Batch --> SearchN[Search × N] - SearchN --> Evaluate[Evaluate] - Evaluate -->|Continue| Batch - Evaluate -->|Done| Synthesize[Synthesize] + Q3[Question] --> Plan[Plan Next] + Plan --> SearchOne[Search One] + SearchOne --> Eval[Evaluate] + Eval -->|Continue| Plan + Eval -->|Done| Synthesize[Synthesize] end subgraph RLM["RLM Agent"] @@ -141,17 +140,17 @@ flowchart TB **Chat Agent** - Multi-turn conversational RAG: -- Maintains session history -- Uses previous Q/A pairs as context -- Query expansion for better recall -- Natural language document filtering +- Composed from reusable [toolsets](tools.md) (search, documents, QA, analysis) +- Maintains session history with prior answer recall +- Background summarization for context continuity +- Session-level document filtering -**Research Graph** - Multi-step research workflow: +**Research Graph** - Iterative research workflow: -- Decomposes questions into sub-questions -- Parallel search execution -- Iterative refinement based on confidence -- Synthesizes structured research report +- Proposes one question at a time, evaluates the answer, then decides whether to continue +- Session context resolves ambiguous references +- Prior answers let the planner skip redundant searches +- Synthesizes structured report or conversational answer **RLM Agent** - Complex analytical tasks via code execution: diff --git a/docs/python.md b/docs/python.md index c1a148b4..a047b8e1 100644 --- a/docs/python.md +++ b/docs/python.md @@ -426,3 +426,26 @@ result = await client.rlm( The RLM agent writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis. See [RLM Agent](rlm.md) for details on capabilities and configuration. + +## Building Custom Agents + +haiku.rag provides composable toolset factories that can be mixed into any pydantic-ai agent. This lets you build custom agents with exactly the capabilities you need — search, document management, Q&A, or code analysis — sharing state across tools via `ToolContext`. + +```python +from pydantic_ai import Agent +from haiku.rag.tools import ToolContext, create_search_toolset, create_qa_toolset + +async with HaikuRAG("path/to/db.lancedb") as client: + context = ToolContext() + agent = Agent( + "openai:gpt-4o", + instructions="You are a helpful assistant.", + toolsets=[ + create_search_toolset(client, config, context=context), + create_qa_toolset(client, config, context=context), + ], + ) + result = await agent.run("What are the main findings?") +``` + +See [Toolsets](tools.md) for the full API reference and composition guide. diff --git a/docs/tools.md b/docs/tools.md new file mode 100644 index 00000000..dfd5b18d --- /dev/null +++ b/docs/tools.md @@ -0,0 +1,268 @@ +# Toolsets + +haiku.rag provides composable `FunctionToolset` factories in `haiku.rag.tools`. Each factory creates a pydantic-ai `FunctionToolset` that can be mixed into any agent. A shared `ToolContext` lets toolsets accumulate state (search results, citations, QA history) across invocations. + +## ToolContext + +`ToolContext` is a namespace-based state container. Toolsets register Pydantic models under string namespaces, and any toolset sharing the same context can read or write the same state. + +```python +from haiku.rag.tools import ToolContext + +context = ToolContext() +``` + +### Registering and accessing state + +```python +from pydantic import BaseModel + +class MyState(BaseModel): + count: int = 0 + +context.register("my_namespace", MyState()) + +# Get state (returns None if not registered) +state = context.get("my_namespace") + +# Get with type checking (returns None if wrong type) +state = context.get("my_namespace", MyState) + +# Get or create (creates default if not registered) +state = context.get_or_create("my_namespace", MyState) +``` + +### Serialization + +The entire context can be serialized and restored: + +```python +# Serialize all namespaces +data = context.dump_namespaces() +# {"my_namespace": {"count": 0}} + +# Restore a namespace from serialized data +context.load_namespace("my_namespace", MyState, data["my_namespace"]) +``` + +## Search Toolset + +`create_search_toolset()` provides hybrid search (vector + full-text) with context expansion and citation tracking. + +```python +from haiku.rag.tools import ToolContext, create_search_toolset + +context = ToolContext() +search = create_search_toolset(client, config, context=context) +``` + +**Parameters:** + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `client` | required | HaikuRAG client | +| `config` | required | AppConfig | +| `context` | `None` | ToolContext for state accumulation | +| `expand_context` | `True` | Expand results with surrounding chunks | +| `base_filter` | `None` | SQL WHERE clause applied to all searches | +| `tool_name` | `"search"` | Name of the tool exposed to the agent | + +**Tool: `search(query, limit?, filter?)`** + +Searches the knowledge base and returns formatted results. When a `ToolContext` with `SessionState` is registered, citations get stable indices via `citation_registry`. + +**State:** Search results accumulate in `SearchState.results` under the `haiku.rag.search` namespace. + +## Document Toolset + +`create_document_toolset()` provides document browsing, retrieval, and summarization. + +```python +from haiku.rag.tools import ToolContext, create_document_toolset + +context = ToolContext() +docs = create_document_toolset(client, config, context=context) +``` + +**Parameters:** + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `client` | required | HaikuRAG client | +| `config` | required | AppConfig (used for summarization LLM) | +| `context` | `None` | ToolContext for session filtering | +| `base_filter` | `None` | SQL WHERE clause for list operations | + +**Tools:** + +- `list_documents(page?)` — Paginated document listing (50 per page). Returns `DocumentListResponse` with document titles, URIs, and pagination info. +- `get_document(query)` — Retrieve a document by title or URI. Uses `find_document()` which tries exact URI match, then partial URI match, then partial title match. +- `summarize_document(query)` — Generate an LLM summary of a document's content. + +## QA Toolset + +`create_qa_toolset()` provides question answering via the research graph, with prior answer recall and background summarization. + +```python +from haiku.rag.tools import ToolContext, create_qa_toolset + +context = ToolContext() +qa = create_qa_toolset(client, config, context=context) +``` + +**Parameters:** + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `client` | required | HaikuRAG client | +| `config` | required | AppConfig | +| `context` | `None` | ToolContext for state accumulation | +| `base_filter` | `None` | SQL WHERE clause applied to searches | +| `tool_name` | `"ask"` | Name of the tool exposed to the agent | +| `session_context` | `None` | Session context for the research graph | +| `prior_answers` | `None` | Prior answers for context | + +**Tool: `ask(question, document_name?)`** + +Runs the research graph in conversational mode and returns a `QAResult`. When a `ToolContext` is provided: + +- Prior answers from `QASessionState.qa_history` are matched via embedding similarity +- The answer is appended to `qa_history` +- Background summarization is triggered +- Citations get stable indices via `SessionState.citation_registry` + +**State:** QA history accumulates in `QASessionState` under the `haiku.rag.qa_session` namespace. + +### Using `run_qa_core()` directly + +For programmatic use without an agent, `run_qa_core()` provides the same QA flow: + +```python +from haiku.rag.tools.qa import run_qa_core + +result = await run_qa_core( + client=client, + config=config, + question="What are the main features?", + document_name="User Guide", # optional document filter + context=context, # optional ToolContext + session_context="User is building a web app", # optional +) + +print(result.answer) +print(result.confidence) +for citation in result.citations: + print(f" [{citation.index}] {citation.document_title}") +``` + +## Analysis Toolset + +`create_analysis_toolset()` provides computational analysis via the RLM agent, which writes and executes Python code in a Docker sandbox. + +```python +from haiku.rag.tools import create_analysis_toolset + +analysis = create_analysis_toolset(client, config, context=context) +``` + +**Parameters:** + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `client` | required | HaikuRAG client | +| `config` | required | AppConfig | +| `context` | `None` | ToolContext for session filtering | +| `base_filter` | `None` | SQL WHERE clause applied to searches | +| `tool_name` | `"analyze"` | Name of the tool exposed to the agent | + +**Tool: `analyze(task, document_name?)`** + +Executes a computational task via code execution and returns an `AnalysisResult`. Requires Docker — see [RLM Agent](rlm.md) for setup. + +## Composing Custom Agents + +Toolsets are designed to be composed into custom pydantic-ai agents: + +```python +from pydantic_ai import Agent +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.tools import ( + ToolContext, + create_search_toolset, + create_qa_toolset, + create_document_toolset, +) + +async with HaikuRAG("path/to/db.lancedb") as client: + # Shared context across all toolsets + context = ToolContext() + + # Pick the toolsets you need + search = create_search_toolset(client, Config, context=context) + qa = create_qa_toolset(client, Config, context=context) + docs = create_document_toolset(client, Config, context=context) + + agent = Agent( + "openai:gpt-4o", + instructions="You are a helpful research assistant.", + toolsets=[search, qa, docs], + ) + + result = await agent.run("What documents do we have about climate?") + print(result.output) + + # Access accumulated state + from haiku.rag.tools import SearchState, SEARCH_NAMESPACE + search_state = context.get(SEARCH_NAMESPACE, SearchState) + if search_state: + print(f"Total search results: {len(search_state.results)}") +``` + +All toolsets respect session-level document filters when a `SessionState` is registered in the context. This means setting `SessionState.document_filter` restricts all tools simultaneously. + +## AG-UI State Management + +When using the chat agent with [AG-UI](https://docs.ag-ui.com) streaming, `ChatDeps` implements the `StateHandler` protocol. State is emitted under a namespaced key via `state_key`: + +```python +from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, create_chat_agent +from haiku.rag.tools import ToolContext + +context = ToolContext() +agent = create_chat_agent(config, client, context) +deps = ChatDeps( + config=config, + tool_context=context, + state_key=AGUI_STATE_KEY, # "haiku.rag.chat" +) +``` + +The emitted state structure: + +```json +{ + "haiku.rag.chat": { + "session_id": "uuid", + "citations": [], + "qa_history": [], + "session_context": null, + "document_filter": [], + "citation_registry": {} + } +} +``` + +State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). See the [Web Application](apps.md#web-application) for a complete implementation. + +## Filter Helpers + +`haiku.rag.tools.filters` provides utilities for building SQL filters: + +**`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593"). + +**`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic. + +**`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`. + +**`get_session_filter(context, base_filter?)`** — Extracts `document_filter` from `SessionState` in the `ToolContext`, builds a SQL filter from it, and combines with an optional `base_filter`. diff --git a/haiku_rag_slim/haiku/rag/agents/__init__.py b/haiku_rag_slim/haiku/rag/agents/__init__.py index 2703d977..bb866996 100644 --- a/haiku_rag_slim/haiku/rag/agents/__init__.py +++ b/haiku_rag_slim/haiku/rag/agents/__init__.py @@ -1,7 +1,6 @@ from haiku.rag.agents.chat import ( ChatDeps, ChatSessionState, - QAHistoryEntry, create_chat_agent, ) from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent @@ -15,6 +14,7 @@ from haiku.rag.agents.research import ( ) from haiku.rag.agents.research.graph import build_research_graph from haiku.rag.agents.research.state import ResearchDeps, ResearchState +from haiku.rag.tools.qa import QAHistoryEntry __all__ = [ # QA diff --git a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py index 1f9d9d47..fb1892eb 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py @@ -9,19 +9,12 @@ from haiku.rag.agents.chat.agent import ( run_chat_agent, trigger_background_summarization, ) -from haiku.rag.agents.chat.context import ( - summarize_session, - update_session_context, -) from haiku.rag.agents.chat.prompts import build_chat_prompt from haiku.rag.agents.chat.state import ( AGUI_STATE_KEY, ChatSessionState, 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", @@ -36,11 +29,5 @@ __all__ = [ "trigger_background_summarization", "ChatDeps", "ChatSessionState", - "DocumentInfo", - "DocumentListResponse", - "QAHistoryEntry", "SessionContext", - "ToolContext", - "summarize_session", - "update_session_context", ] diff --git a/haiku_rag_slim/haiku/rag/agents/chat/context.py b/haiku_rag_slim/haiku/rag/agents/chat/context.py index e3195094..9fbaf9d4 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/context.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/context.py @@ -6,7 +6,7 @@ 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, SessionContext +from haiku.rag.agents.chat.state import SessionContext from haiku.rag.config.models import AppConfig from haiku.rag.utils import get_model @@ -115,32 +115,30 @@ async def summarize_session( async def update_session_context( qa_history: list["QAHistoryEntry"], config: AppConfig, - session_state: ChatSessionState, -) -> None: - """Update session context in the session state. + session_id: str = "", + current_context: str | None = None, +) -> SessionContext: + """Summarize qa_history and cache the resulting session context. Args: qa_history: List of Q&A pairs from the conversation. config: AppConfig for model selection. - session_state: The session state to update. - """ - # Use existing session_context summary if available, else initial_context - current_context: str | None = None - if session_state.session_context and session_state.session_context.summary: - current_context = session_state.session_context.summary - elif session_state.initial_context: - current_context = session_state.initial_context + session_id: Session ID for caching. If empty, result is not cached. + current_context: Previous summary to incorporate. + Returns: + The new SessionContext with summary and timestamp. + """ summary = await summarize_session( qa_history, config, current_context=current_context ) - session_state.session_context = SessionContext( + context = SessionContext( summary=summary, last_updated=datetime.now(), ) - # Also cache for next-run delivery in stateless contexts - if session_state.session_id: - cache_session_context(session_state.session_id, session_state.session_context) + if session_id: + cache_session_context(session_id, context) + return context def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str: @@ -165,23 +163,15 @@ async def _update_context_background( ) -> None: """Background task to update session context after an ask.""" try: - qa_history = list(qa_session_state.qa_history) - - session_state = ChatSessionState( - session_id=session_id, - qa_history=qa_history, - ) - - await update_session_context( - qa_history=qa_history, + result = await update_session_context( + qa_history=list(qa_session_state.qa_history), config=config, - session_state=session_state, + session_id=session_id, + current_context=qa_session_state.session_context, ) - # Update the QASessionState with the new context - cached = get_cached_session_context(session_id) - if cached and cached.summary: - qa_session_state.session_context = cached.summary + if result.summary: + qa_session_state.session_context = result.summary except asyncio.CancelledError: pass diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py index d417c318..633cd7b1 100644 --- a/haiku_rag_slim/haiku/rag/tools/__init__.py +++ b/haiku_rag_slim/haiku/rag/tools/__init__.py @@ -47,7 +47,6 @@ __all__ = [ "QASessionState", "QAHistoryEntry", "create_qa_toolset", - "run_qa_core", "create_analysis_toolset", "SESSION_NAMESPACE", "SessionState", diff --git a/mkdocs.yml b/mkdocs.yml index 0efc7e8a..b5497574 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -72,6 +72,7 @@ nav: - Custom Pipelines: custom-pipelines.md - Tuning: tuning.md - Agents: agents.md + - Toolsets: tools.md - RLM Agent: rlm.md - Applications: apps.md - Server: server.md diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index e9662b9c..25b06a1c 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -7,8 +7,6 @@ from haiku.rag.agents.chat import ( AGUI_STATE_KEY, ChatDeps, ChatSessionState, - QAHistoryEntry, - ToolContext, create_chat_agent, ) from haiku.rag.agents.chat.context import ( @@ -18,7 +16,8 @@ from haiku.rag.agents.chat.context import ( 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 import ToolContext +from haiku.rag.tools.qa import MAX_QA_HISTORY, QAHistoryEntry from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState diff --git a/tests/agents/chat/test_context.py b/tests/agents/chat/test_context.py index e4501ff6..02f9b5bf 100644 --- a/tests/agents/chat/test_context.py +++ b/tests/agents/chat/test_context.py @@ -195,14 +195,11 @@ class TestUpdateSessionContext: @pytest.mark.asyncio @pytest.mark.vcr() - async def test_update_session_context_updates_state( + async def test_update_session_context_returns_context( self, allow_model_requests, temp_db_path ): - """Test update_session_context updates the session_state.""" + """Test update_session_context returns a populated SessionContext.""" from haiku.rag.agents.chat.context import update_session_context - from haiku.rag.agents.chat.state import ChatSessionState - - session_state = ChatSessionState(session_id="test-session") qa_history = [ QAHistoryEntry( @@ -212,34 +209,26 @@ class TestUpdateSessionContext: ) ] - await update_session_context( + result = await update_session_context( qa_history=qa_history, config=Config, - session_state=session_state, + session_id="test-session", ) - # session_context should now be populated - assert session_state.session_context is not None - assert session_state.session_context.summary != "" - assert session_state.session_context.last_updated is not None + assert result.summary != "" + assert result.last_updated is not None @pytest.mark.asyncio async def test_update_session_context_with_empty_history(self): - """Test update_session_context with empty history sets empty context.""" + """Test update_session_context with empty history returns empty summary.""" from haiku.rag.agents.chat.context import update_session_context - from haiku.rag.agents.chat.state import ChatSessionState - session_state = ChatSessionState(session_id="test-session") - - await update_session_context( + result = await update_session_context( qa_history=[], config=Config, - session_state=session_state, ) - # session_context should exist but have empty summary - assert session_state.session_context is not None - assert session_state.session_context.summary == "" + assert result.summary == "" class TestSessionContextCache: @@ -319,12 +308,9 @@ class TestSessionContextCache: get_cached_session_context, update_session_context, ) - from haiku.rag.agents.chat.state import ChatSessionState _session_cache.clear() - session_state = ChatSessionState(session_id="cache-test-session") - qa_history = [ QAHistoryEntry( question="What is Python?", @@ -333,23 +319,21 @@ class TestSessionContextCache: ) ] - # Mock summarize_session to avoid LLM call (we're testing caching, not summarization) with patch( "haiku.rag.agents.chat.context.summarize_session", new=AsyncMock(return_value="Mocked summary"), ): - await update_session_context( + result = await update_session_context( qa_history=qa_history, config=Config, - session_state=session_state, + session_id="cache-test-session", ) - # Should be cached + assert result.summary == "Mocked summary" + cached = get_cached_session_context("cache-test-session") assert cached is not None assert cached.summary == "Mocked summary" - assert session_state.session_context is not None - assert cached.summary == session_state.session_context.summary @pytest.mark.asyncio async def test_update_session_context_no_cache_without_session_id(self): @@ -359,41 +343,24 @@ class TestSessionContextCache: get_cached_session_context, update_session_context, ) - from haiku.rag.agents.chat.state import ChatSessionState _session_cache.clear() - # No session_id - session_state = ChatSessionState() - await update_session_context( qa_history=[], config=Config, - session_state=session_state, ) - # Nothing should be cached (empty session_id) + # Nothing should be cached (no session_id) cached = get_cached_session_context("") assert cached is None @pytest.mark.asyncio - async def test_update_session_context_uses_initial_context_as_fallback(self): - """Test update_session_context uses initial_context when no session_context exists.""" + async def test_update_session_context_passes_current_context(self): + """Test update_session_context passes current_context to summarizer.""" from unittest.mock import patch - from haiku.rag.agents.chat.context import ( - _session_cache, - update_session_context, - ) - from haiku.rag.agents.chat.state import ChatSessionState - - _session_cache.clear() - - # Create session_state with initial_context but no session_context - session_state = ChatSessionState( - session_id="initial-context-test", - initial_context="User is working on a Python web application with FastAPI.", - ) + from haiku.rag.agents.chat.context import update_session_context qa_history = [ QAHistoryEntry( @@ -403,7 +370,6 @@ class TestSessionContextCache: ) ] - # Mock summarize_session to capture what gets passed as current_context captured_current_context = [] async def mock_summarize(qa_history, config, current_context=None): @@ -417,61 +383,8 @@ class TestSessionContextCache: await update_session_context( qa_history=qa_history, config=Config, - session_state=session_state, + current_context="Previous session summary", ) - # initial_context should have been passed as current_context assert len(captured_current_context) == 1 - assert ( - captured_current_context[0] - == "User is working on a Python web application with FastAPI." - ) - - @pytest.mark.asyncio - async def test_update_session_context_session_context_takes_precedence(self): - """Test session_context.summary takes precedence over initial_context.""" - from unittest.mock import patch - - from haiku.rag.agents.chat.context import ( - _session_cache, - update_session_context, - ) - from haiku.rag.agents.chat.state import ChatSessionState, SessionContext - - _session_cache.clear() - - # Create session_state with BOTH initial_context and session_context - session_state = ChatSessionState( - session_id="precedence-test", - initial_context="Initial background info", - session_context=SessionContext(summary="Evolved session summary"), - ) - - qa_history = [ - QAHistoryEntry( - question="What is JWT?", - answer="JSON Web Token.", - confidence=0.95, - ) - ] - - # Mock summarize_session to capture what gets passed as current_context - captured_current_context = [] - - async def mock_summarize(qa_history, config, current_context=None): - captured_current_context.append(current_context) - return "Mocked summary" - - with patch( - "haiku.rag.agents.chat.context.summarize_session", - new=mock_summarize, - ): - await update_session_context( - qa_history=qa_history, - config=Config, - session_state=session_state, - ) - - # session_context.summary should take precedence over initial_context - assert len(captured_current_context) == 1 - assert captured_current_context[0] == "Evolved session summary" + assert captured_current_context[0] == "Previous session summary" diff --git a/tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_updates_state.yaml b/tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_returns_context.yaml similarity index 100% rename from tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_updates_state.yaml rename to tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_returns_context.yaml