Add initial_context support to chat agent

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

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:

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

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"] = [

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(

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

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