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

Clean up toolset exports, simplify update_session_context(), fix health check

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-11 11:51:58 +02:00
parent 87bcb7a5dc
commit c916356a42
No known key found for this signature in database
13 changed files with 361 additions and 259 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
app/backend/main.py
View file

@ -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),

93
docs/agents.md
View file

@ -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 `<background>` 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 `<prior_answers>` 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 `<background>` 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

33
docs/architecture.md
View file

@ -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:

23
docs/python.md
View file

@ -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.

268
docs/tools.md Normal file
View file

@ -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`.

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

@ -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

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

@ -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",
]

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

@ -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

1
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

@ -47,7 +47,6 @@ __all__ = [
"QASessionState",
"QAHistoryEntry",
"create_qa_toolset",
"run_qa_core",
"create_analysis_toolset",
"SESSION_NAMESPACE",
"SessionState",

1
mkdocs.yml
View file

@ -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

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

@ -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

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

@ -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"

0
tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_updates_state.yaml → tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_returns_context.yaml
View file