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

Allow initial_context to be set for the chat agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-27 10:36:23 +02:00
parent 7c9d70fa62
commit ee3cb5bd87
No known key found for this signature in database
9 changed files with 176 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -1,6 +1,14 @@
# Changelog
## [Unreleased]
### Added
- **Initial Context for Chat Sessions**: New `initial_context` field on `ChatSessionState` allows external clients to seed sessions with background context
- Static context set once at session creation, used as fallback when no cached session context exists
- Incorporated into first summarization, after which evolved `session_context` takes precedence
- Eliminates need for clients to import and call internal cache functions (`cache_session_context`, `get_cached_session_context`)
- `session_id` now auto-generates a UUID if not provided (previously defaulted to empty string)
## [0.27.0] - 2026-01-26
### Added

8
app/backend/main.py
View file

@ -1,6 +1,5 @@
import logging
import os
import uuid
from pathlib import Path
from dotenv import find_dotenv, load_dotenv
@ -94,18 +93,13 @@ async def stream_chat(request: Request) -> Response:
session_id = chat_state.get("session_id")
document_filter = chat_state.get("document_filter", [])
# Determine session_id: prefer state, fall back to thread_id, generate UUID if neither
thread_id = getattr(run_input, "thread_id", None)
if not session_id:
session_id = thread_id or str(uuid.uuid4())
deps = ChatDeps(
client=get_client(db_path),
config=Config,
session_state=ChatSessionState(
session_id=session_id,
qa_history=initial_qa_history,
document_filter=document_filter,
**({"session_id": session_id} if session_id else {}),
),
state_key=AGUI_STATE_KEY,
)

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

@ -217,12 +217,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
graph = build_conversational_graph(config=ctx.deps.config)
session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
# Get session context from server cache for planning
# Get session context from server cache for planning, fallback to initial_context
cached_context = get_cached_session_context(session_id) if session_id else None
session_context = (
cached_context.render_markdown()
if cached_context and cached_context.summary
else None
else (
ctx.deps.session_state.initial_context
if ctx.deps.session_state
else None
)
)
# Find relevant prior answers from qa_history

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

@ -82,10 +82,12 @@ async def update_session_context(
config: AppConfig for model selection.
session_state: The session state to update.
"""
# Use existing session_context summary if available
# 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
summary = await summarize_session(
qa_history, config, current_context=current_context

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

@ -1,3 +1,4 @@
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
@ -55,7 +56,8 @@ class SessionContext(BaseModel):
class ChatSessionState(BaseModel):
"""State shared between frontend and agent via AG-UI."""
session_id: str = ""
session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
initial_context: str | None = None
citations: list[Citation] = []
qa_history: list[QAResponse] = []
session_context: SessionContext | None = None

10
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -87,7 +87,7 @@ async def _plan_step_logic(
else plan_prompt
)
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent(
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=effective_plan_prompt,
@ -154,7 +154,7 @@ async def _search_one_step_logic(
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=search_prompt,
@ -280,7 +280,7 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, EvaluationResult] = Agent(
agent: Agent[ResearchDependencies, EvaluationResult] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=decision_prompt,
@ -341,7 +341,7 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ResearchReport] = Agent(
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
@ -479,7 +479,7 @@ def build_conversational_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[assignment]
model=get_model(config.research.model, config),
output_type=ConversationalAnswer,
instructions=conversational_prompt,

2
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -153,7 +153,6 @@ class ChatApp(App):
# Create agent and session state
self.agent = create_chat_agent(self.config)
self.session_state = ChatSessionState(
session_id=str(uuid.uuid4()),
document_filter=self._document_filter,
)
@ -300,7 +299,6 @@ class ChatApp(App):
self._message_history.clear()
# Reset session state for fresh conversation (preserve document filter)
self.session_state = ChatSessionState(
session_id=str(uuid.uuid4()),
document_filter=self._document_filter,
)

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

@ -388,3 +388,103 @@ class TestSessionContextCache:
# Nothing should be cached (empty 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."""
from unittest.mock import patch
from haiku.rag.agents.chat.context import (
_session_context_cache,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState
_session_context_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.",
)
qa_history = [
QAResponse(
question="What is JWT?",
answer="JSON Web Token for authentication.",
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,
)
# 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_context_cache,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState, SessionContext
_session_context_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 = [
QAResponse(
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"

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

@ -1,5 +1,8 @@
import uuid
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
ChatSessionState,
QAResponse,
build_document_filter,
build_multi_document_filter,
@ -565,3 +568,50 @@ def test_chat_deps_state_getter_includes_document_filter():
assert state is not None
assert AGUI_STATE_KEY in state
assert state[AGUI_STATE_KEY]["document_filter"] == ["doc1.pdf", "doc2.pdf"]
def test_chat_session_state_auto_generates_session_id():
"""New ChatSessionState should have a valid UUID session_id."""
state = ChatSessionState()
assert state.session_id
assert len(state.session_id) == 36 # UUID format
# Verify it's a valid UUID
uuid.UUID(state.session_id)
def test_chat_session_state_preserves_explicit_session_id():
"""Explicit session_id should be preserved."""
state = ChatSessionState(session_id="my-custom-id")
assert state.session_id == "my-custom-id"
def test_chat_session_state_each_instance_gets_unique_id():
"""Each new instance should get a unique session_id."""
state1 = ChatSessionState()
state2 = ChatSessionState()
assert state1.session_id != state2.session_id
def test_chat_session_state_initial_context_default_none():
"""Initial context should default to None."""
state = ChatSessionState()
assert state.initial_context is None
def test_chat_session_state_initial_context_preserved():
"""Explicit initial_context should be preserved."""
state = ChatSessionState(initial_context="Background info about the project")
assert state.initial_context == "Background info about the project"
def test_chat_session_state_initial_context_serialization():
"""initial_context should serialize and deserialize correctly."""
state = ChatSessionState(
session_id="test-123",
initial_context="User is working on authentication",
)
state_dict = state.model_dump()
assert state_dict["initial_context"] == "User is working on authentication"
restored = ChatSessionState.model_validate(state_dict)
assert restored.initial_context == "User is working on authentication"