From dff489f75fc468caa930c74f457f166a3d2d6cb8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 16 Jan 2026 10:54:02 +0200 Subject: [PATCH] Add initial_context support to chat agent --- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 22 +++++++ haiku_rag_slim/haiku/rag/agents/chat/state.py | 1 + .../haiku/rag/agents/research/dependencies.py | 4 ++ .../haiku/rag/agents/research/graph.py | 7 ++- tests/agents/chat/test_chat_agent.py | 11 ++++ tests/agents/chat/test_state.py | 19 ++++++ tests/agents/research/test_research_graph.py | 59 +++++++++++++++++++ 7 files changed, 122 insertions(+), 1 deletion(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 4417c5b6..f676c691 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -32,6 +32,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: retries=3, ) + @agent.system_prompt + async def add_initial_context(ctx: RunContext[ChatDeps]) -> str: + """Add initial_context to system prompt when available.""" + if ctx.deps.session_state and ctx.deps.session_state.initial_context: + return f"\nBACKGROUND CONTEXT:\n{ctx.deps.session_state.initial_context}" + return "" + @agent.tool async def search( ctx: RunContext[ChatDeps], @@ -86,6 +93,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: qa_history=( ctx.deps.session_state.qa_history if ctx.deps.session_state else [] ), + initial_context=( + ctx.deps.session_state.initial_context + if ctx.deps.session_state + else None + ), ) # Return detailed results for the agent to present @@ -182,9 +194,14 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: # Build and run the conversational research graph graph = build_conversational_graph(config=ctx.deps.config) + initial_context = ( + ctx.deps.session_state.initial_context if ctx.deps.session_state else None + ) + context = ResearchContext( original_question=question, qa_responses=existing_qa, + initial_context=initial_context, ) state = ResearchState( context=context, @@ -238,6 +255,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: qa_history=( ctx.deps.session_state.qa_history if ctx.deps.session_state else [] ), + initial_context=( + ctx.deps.session_state.initial_context + if ctx.deps.session_state + else None + ), ) # Format answer with citation references and confidence diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 7708df89..f60169df 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -61,6 +61,7 @@ class ChatSessionState(BaseModel): session_id: str = "" citations: list[CitationInfo] = [] qa_history: list[QAResponse] = [] + initial_context: str | None = None def format_conversation_context(qa_history: list[QAResponse]) -> str: diff --git a/haiku_rag_slim/haiku/rag/agents/research/dependencies.py b/haiku_rag_slim/haiku/rag/agents/research/dependencies.py index 12f1f154..1f5559ae 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/agents/research/dependencies.py @@ -19,6 +19,10 @@ class ResearchContext(BaseModel): qa_responses: list[Any] = Field( default_factory=list, description="Structured QA pairs used during research" ) + initial_context: str | None = Field( + default=None, + description="Optional background context provided at session start", + ) def add_qa_response(self, qa: "SearchAnswer") -> None: """Add a structured QA response.""" diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index 6e57c582..905da165 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/haiku_rag_slim/haiku/rag/agents/research/graph.py @@ -31,7 +31,7 @@ from haiku.rag.utils import build_prompt, get_model def format_context_for_prompt(context: ResearchContext) -> str: """Format the research context as XML for inclusion in prompts.""" - context_data = { + context_data: dict[str, object] = { "original_question": context.original_question, "unanswered_questions": context.sub_questions, "qa_responses": [ @@ -52,6 +52,8 @@ def format_context_for_prompt(context: ResearchContext) -> str: for qa in context.qa_responses ], } + if context.initial_context: + context_data["initial_context"] = context.initial_context return format_as_xml(context_data, root_tag="research_context") @@ -61,6 +63,9 @@ def format_conversational_context_for_prompt(context: ResearchContext) -> str: "question": context.original_question, } + if context.initial_context: + context_data["initial_context"] = context.initial_context + # Only include conversation_history if there are qa_responses if context.qa_responses: context_data["conversation_history"] = [ diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 7bb2ac43..77e8d0a8 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -76,6 +76,17 @@ def test_chat_session_state(): assert state.qa_history == [] +def test_chat_agent_has_dynamic_system_prompt(): + """Test that chat agent registers a dynamic system prompt for initial_context.""" + agent = create_chat_agent(Config) + # The agent should have at least one system prompt function registered + # (the add_initial_context function) + assert len(agent._system_prompt_functions) >= 1 + # Verify it's the add_initial_context function + func_names = [r.function.__name__ for r in agent._system_prompt_functions] + assert "add_initial_context" in func_names + + def test_citation_info(): """Test CitationInfo model.""" citation = CitationInfo( diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index df37728d..60f59855 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -248,3 +248,22 @@ def test_build_document_filter_escapes_quotes(): def test_max_qa_history_constant(): """Test MAX_QA_HISTORY constant value.""" assert MAX_QA_HISTORY == 50 + + +def test_chat_session_state_initial_context(): + """Test ChatSessionState accepts initial_context.""" + from haiku.rag.agents.chat.state import ChatSessionState + + state = ChatSessionState( + session_id="test-session", + initial_context="This is background knowledge about the topic.", + ) + assert state.initial_context == "This is background knowledge about the topic." + + +def test_chat_session_state_initial_context_defaults_to_none(): + """Test ChatSessionState initial_context defaults to None.""" + from haiku.rag.agents.chat.state import ChatSessionState + + state = ChatSessionState(session_id="test-session") + assert state.initial_context is None diff --git a/tests/agents/research/test_research_graph.py b/tests/agents/research/test_research_graph.py index f2f292ef..c077f73c 100644 --- a/tests/agents/research/test_research_graph.py +++ b/tests/agents/research/test_research_graph.py @@ -44,3 +44,62 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus): assert result.executive_summary client.close() + + +def test_research_context_initial_context(): + """Test ResearchContext accepts initial_context.""" + context = ResearchContext( + original_question="What is X?", + initial_context="Background: X is a concept in domain Y.", + ) + assert context.initial_context == "Background: X is a concept in domain Y." + + +def test_research_context_initial_context_defaults_to_none(): + """Test ResearchContext initial_context defaults to None.""" + context = ResearchContext(original_question="What is X?") + assert context.initial_context is None + + +def test_format_context_for_prompt_includes_initial_context(): + """Test format_context_for_prompt includes initial_context in output.""" + from haiku.rag.agents.research.graph import format_context_for_prompt + + context = ResearchContext( + original_question="What is X?", + initial_context="Background: X is a concept in domain Y.", + ) + result = format_context_for_prompt(context) + assert "Background: X is a concept in domain Y." in result + assert "initial_context" in result + + +def test_format_context_for_prompt_excludes_initial_context_when_none(): + """Test format_context_for_prompt excludes initial_context when None.""" + from haiku.rag.agents.research.graph import format_context_for_prompt + + context = ResearchContext(original_question="What is X?") + result = format_context_for_prompt(context) + assert "initial_context" not in result + + +def test_format_conversational_context_for_prompt_includes_initial_context(): + """Test format_conversational_context_for_prompt includes initial_context.""" + from haiku.rag.agents.research.graph import format_conversational_context_for_prompt + + context = ResearchContext( + original_question="What is X?", + initial_context="Background: X is a concept in domain Y.", + ) + result = format_conversational_context_for_prompt(context) + assert "Background: X is a concept in domain Y." in result + assert "initial_context" in result + + +def test_format_conversational_context_for_prompt_excludes_initial_context_when_none(): + """Test format_conversational_context_for_prompt excludes initial_context when None.""" + from haiku.rag.agents.research.graph import format_conversational_context_for_prompt + + context = ResearchContext(original_question="What is X?") + result = format_conversational_context_for_prompt(context) + assert "initial_context" not in result