diff --git a/CHANGELOG.md b/CHANGELOG.md index dbb43357..7fcd2e2c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/app/backend/main.py b/app/backend/main.py index 3ce50f43..ff297a30 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -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, ) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 941530e2..8e3b1c52 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/agents/chat/context.py b/haiku_rag_slim/haiku/rag/agents/chat/context.py index 5411a95a..27086a78 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/context.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/context.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 2c91c081..04c68f7c 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index eefb084f..bd5f08f0 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/haiku_rag_slim/haiku/rag/agents/research/graph.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 7dc40b79..f2cd3456 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -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, ) diff --git a/tests/agents/chat/test_context.py b/tests/agents/chat/test_context.py index 03d3e894..0d7ffb9b 100644 --- a/tests/agents/chat/test_context.py +++ b/tests/agents/chat/test_context.py @@ -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" diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index c62a943b..12d22d40 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -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"