Refactor to make agents a top-level module. Bring in the conversational agent from the app

app/backend/main.py

@@ -2,7 +2,6 @@ import logging
 import os
 from pathlib import Path
 
-from agent import ChatDeps, ChatSessionState, QAResponse, create_chat_agent
 from dotenv import load_dotenv
 from pydantic_ai.ui import SSE_CONTENT_TYPE
 from pydantic_ai.ui.ag_ui import AGUIAdapter
@@ -13,6 +12,12 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route
 
+from haiku.rag.agents.chat import (
+    ChatDeps,
+    ChatSessionState,
+    QAResponse,
+    create_chat_agent,
+)
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import load_yaml_config
 from haiku.rag.config.models import AppConfig

docs/agents.md

@@ -33,7 +33,7 @@ haiku-rag ask "What are the main features of haiku.rag?" --deep
 
 ```python
 from haiku.rag.client import HaikuRAG
-from haiku.rag.qa.agent import QuestionAnswerAgent
+from haiku.rag.agents.qa.agent import QuestionAnswerAgent
 
 async with HaikuRAG(path_to_db) as client:
     agent = QuestionAnswerAgent(
@@ -105,9 +105,9 @@ haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'"
 ```python
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.graph import build_research_graph
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_research_graph
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 
 async with HaikuRAG(path_to_db) as client:
     graph = build_research_graph(config=Config)
@@ -126,9 +126,9 @@ async with HaikuRAG(path_to_db) as client:
 ```python
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig, ResearchConfig
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.graph import build_research_graph
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_research_graph
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 
 custom_config = AppConfig(
     research=ResearchConfig(

evaluations/evaluations/benchmark.py

@@ -20,7 +20,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
 from haiku.rag.config.models import ModelConfig
 from haiku.rag.logging import configure_cli_logging
-from haiku.rag.qa import get_qa_agent
+from haiku.rag.agents.qa import get_qa_agent
 from haiku.rag.utils import get_model
 
 load_dotenv()

haiku_rag_slim/haiku/rag/agents/__init__.py

@@ -0,0 +1,48 @@
+from haiku.rag.agents.chat import (
+    ChatDeps,
+    ChatSessionState,
+    CitationInfo,
+    QAResponse,
+    SearchAgent,
+    SearchDeps,
+    create_chat_agent,
+)
+from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent
+from haiku.rag.agents.research import (
+    EvaluationResult,
+    ResearchContext,
+    ResearchDependencies,
+    ResearchReport,
+    SearchAnswer,
+)
+from haiku.rag.agents.research.graph import (
+    build_conversational_graph,
+    build_research_graph,
+)
+from haiku.rag.agents.research.models import Citation
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
+
+__all__ = [
+    # QA
+    "get_qa_agent",
+    "QuestionAnswerAgent",
+    # Research
+    "build_research_graph",
+    "build_conversational_graph",
+    "ResearchContext",
+    "ResearchDependencies",
+    "ResearchDeps",
+    "ResearchState",
+    "ResearchReport",
+    "Citation",
+    "SearchAnswer",
+    "EvaluationResult",
+    # Chat
+    "create_chat_agent",
+    "SearchAgent",
+    "ChatDeps",
+    "ChatSessionState",
+    "CitationInfo",
+    "QAResponse",
+    "SearchDeps",
+]

haiku_rag_slim/haiku/rag/agents/chat/__init__.py

@@ -0,0 +1,23 @@
+from haiku.rag.agents.chat.agent import create_chat_agent
+from haiku.rag.agents.chat.search import SearchAgent
+from haiku.rag.agents.chat.state import (
+    ChatDeps,
+    ChatSessionState,
+    CitationInfo,
+    QAResponse,
+    SearchDeps,
+    build_document_filter,
+    format_conversation_context,
+)
+
+__all__ = [
+    "create_chat_agent",
+    "SearchAgent",
+    "ChatDeps",
+    "ChatSessionState",
+    "CitationInfo",
+    "QAResponse",
+    "SearchDeps",
+    "build_document_filter",
+    "format_conversation_context",
+]

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -1,120 +1,24 @@
-from dataclasses import dataclass
-
 from ag_ui.core import EventType, StateSnapshotEvent
-from pydantic import BaseModel
-from pydantic_ai import Agent, RunContext, ToolReturn, format_as_xml
+from pydantic_ai import Agent, RunContext, ToolReturn
 
-from haiku.rag.client import HaikuRAG
+from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
+from haiku.rag.agents.chat.search import SearchAgent
+from haiku.rag.agents.chat.state import (
+    ChatDeps,
+    ChatSessionState,
+    CitationInfo,
+    QAResponse,
+    build_document_filter,
+    format_conversation_context,
+)
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_conversational_graph
+from haiku.rag.agents.research.models import Citation, SearchAnswer
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.config.models import AppConfig
-from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import get_model
 
 
-class CitationInfo(BaseModel):
-    """Citation info for frontend display."""
-
-    index: int
-    document_id: str
-    chunk_id: str
-    document_uri: str
-    document_title: str | None = None
-    page_numbers: list[int] = []
-    headings: list[str] | None = None
-    content: str
-
-
-class QAResponse(BaseModel):
-    """A Q&A pair from conversation history with citations."""
-
-    question: str
-    answer: str
-    confidence: float = 0.9
-    citations: list[CitationInfo] = []
-
-    @property
-    def sources(self) -> list[str]:
-        """Source names for display."""
-        return list(
-            dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
-        )
-
-
-class ChatSessionState(BaseModel):
-    """State shared between frontend and agent via AG-UI."""
-
-    session_id: str = ""
-    citations: list[CitationInfo] = []
-    qa_history: list[QAResponse] = []
-
-
-def format_conversation_context(qa_history: list[QAResponse]) -> str:
-    """Format conversation history as XML for inclusion in prompts."""
-    if not qa_history:
-        return ""
-
-    context_data = {
-        "previous_qa": [
-            {
-                "question": qa.question,
-                "answer": qa.answer,
-                "sources": qa.sources,
-            }
-            for qa in qa_history
-        ],
-    }
-    return format_as_xml(context_data, root_tag="conversation_context")
-
-
-@dataclass
-class ChatDeps:
-    """Dependencies for chat agent."""
-
-    client: HaikuRAG
-    config: AppConfig
-    search_results: list[SearchResult] | None = None
-    session_state: ChatSessionState | None = None
-
-
-def build_document_filter(document_name: str) -> str:
-    """Build SQL filter for document name matching."""
-    escaped = document_name.replace("'", "''")
-    no_spaces = escaped.replace(" ", "")
-    return (
-        f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
-        f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
-    )
-
-
-CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
-
-You have access to a knowledge base of documents. Use your tools to search and answer questions.
-
-CRITICAL RULES:
-1. For greetings or casual chat: respond directly WITHOUT using any tools
-2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
-3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
-4. NEVER call the same tool multiple times for a single user message
-5. NEVER make up information - always use tools to get facts from the knowledge base
-
-How to decide which tool to use:
-- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
-- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
-- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
-
-IMPORTANT - When user mentions a document in search/ask:
-- If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
-  - Extract the TOPIC as `query`/`question`
-  - Extract the DOCUMENT NAME as `document_name`
-- Examples for search:
-  - "search for latrines in TB MED 593" → query="latrines", document_name="TB MED 593"
-  - "find waste disposal in the army manual" → query="waste disposal", document_name="army manual"
-- Examples for ask:
-  - "what does TB MED 593 say about latrines?" → question="what are the guidelines for latrines?", document_name="TB MED 593"
-  - "answer from the army manual about sanitation" → question="what are the sanitation guidelines?", document_name="army manual"
-
-Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user."""
-
-
 def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
     """Create the chat agent with search and ask tools."""
     model = get_model(config.qa.model, config)
@@ -141,8 +45,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             query: The search query (what to search for)
             document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
         """
-        from search_agent import SearchAgent
-
         # Build context from conversation history
         context = None
         if ctx.deps.session_state and ctx.deps.session_state.qa_history:
@@ -228,11 +130,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             question: The question to answer
             document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
         """
-        from haiku.rag.graph.research.dependencies import ResearchContext
-        from haiku.rag.graph.research.graph import build_conversational_graph
-        from haiku.rag.graph.research.models import Citation, SearchAnswer
-        from haiku.rag.graph.research.state import ResearchDeps, ResearchState
-
         # Build filter from document_name
         doc_filter = build_document_filter(document_name) if document_name else None
 

haiku_rag_slim/haiku/rag/agents/chat/prompts.py

@@ -0,0 +1,41 @@
+CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
+
+You have access to a knowledge base of documents. Use your tools to search and answer questions.
+
+CRITICAL RULES:
+1. For greetings or casual chat: respond directly WITHOUT using any tools
+2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
+3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
+4. NEVER call the same tool multiple times for a single user message
+5. NEVER make up information - always use tools to get facts from the knowledge base
+
+How to decide which tool to use:
+- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
+- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
+- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
+
+IMPORTANT - When user mentions a document in search/ask:
+- If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
+  - Extract the TOPIC as `query`/`question`
+  - Extract the DOCUMENT NAME as `document_name`
+- Examples for search:
+  - "search for latrines in TB MED 593" → query="latrines", document_name="TB MED 593"
+  - "find waste disposal in the army manual" → query="waste disposal", document_name="army manual"
+- Examples for ask:
+  - "what does TB MED 593 say about latrines?" → question="what are the guidelines for latrines?", document_name="TB MED 593"
+  - "answer from the army manual about sanitation" → question="what are the sanitation guidelines?", document_name="army manual"
+
+Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user."""
+
+SEARCH_SYSTEM_PROMPT = """You are a search query optimizer for a document knowledge base.
+
+Given a user's search request:
+1. ALWAYS run the original query first as-is
+2. Then generate 1-2 alternative queries using different keywords or phrasings
+3. Keep queries SHORT (2-5 words) - use keywords, not full sentences
+4. After all searches, respond with "Search complete"
+
+Example: User asks "latrines" → queries: "latrines", "latrine sanitation", "field toilet"
+Example: User asks "waste disposal" → queries: "waste disposal", "garbage management", "refuse handling"
+
+Do NOT generate long verbose queries like "environmental impact of waste disposal methods" - keep it simple."""

haiku_rag_slim/haiku/rag/agents/chat/search.py

@@ -1,37 +1,13 @@
-from dataclasses import dataclass, field
-
 from pydantic_ai import Agent, RunContext
 
+from haiku.rag.agents.chat.prompts import SEARCH_SYSTEM_PROMPT
+from haiku.rag.agents.chat.state import SearchDeps
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import get_model
 
 
-@dataclass
-class SearchDeps:
-    """Dependencies for search agent."""
-
-    client: HaikuRAG
-    config: AppConfig
-    filter: str | None = None
-    search_results: list[SearchResult] = field(default_factory=list)
-
-
-SEARCH_SYSTEM_PROMPT = """You are a search query optimizer for a document knowledge base.
-
-Given a user's search request:
-1. ALWAYS run the original query first as-is
-2. Then generate 1-2 alternative queries using different keywords or phrasings
-3. Keep queries SHORT (2-5 words) - use keywords, not full sentences
-4. After all searches, respond with "Search complete"
-
-Example: User asks "latrines" → queries: "latrines", "latrine sanitation", "field toilet"
-Example: User asks "waste disposal" → queries: "waste disposal", "garbage management", "refuse handling"
-
-Do NOT generate long verbose queries like "environmental impact of waste disposal methods" - keep it simple."""
-
-
 class SearchAgent:
     """Agent that generates multiple queries and consolidates results."""
 

haiku_rag_slim/haiku/rag/agents/chat/state.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel
+from pydantic_ai import format_as_xml
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config.models import AppConfig
+from haiku.rag.store.models import SearchResult
+
+
+class CitationInfo(BaseModel):
+    """Citation info for frontend display."""
+
+    index: int
+    document_id: str
+    chunk_id: str
+    document_uri: str
+    document_title: str | None = None
+    page_numbers: list[int] = []
+    headings: list[str] | None = None
+    content: str
+
+
+class QAResponse(BaseModel):
+    """A Q&A pair from conversation history with citations."""
+
+    question: str
+    answer: str
+    confidence: float = 0.9
+    citations: list[CitationInfo] = []
+
+    @property
+    def sources(self) -> list[str]:
+        """Source names for display."""
+        return list(
+            dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
+        )
+
+
+class ChatSessionState(BaseModel):
+    """State shared between frontend and agent via AG-UI."""
+
+    session_id: str = ""
+    citations: list[CitationInfo] = []
+    qa_history: list[QAResponse] = []
+
+
+def format_conversation_context(qa_history: list[QAResponse]) -> str:
+    """Format conversation history as XML for inclusion in prompts."""
+    if not qa_history:
+        return ""
+
+    context_data = {
+        "previous_qa": [
+            {
+                "question": qa.question,
+                "answer": qa.answer,
+                "sources": qa.sources,
+            }
+            for qa in qa_history
+        ],
+    }
+    return format_as_xml(context_data, root_tag="conversation_context")
+
+
+@dataclass
+class ChatDeps:
+    """Dependencies for chat agent."""
+
+    client: HaikuRAG
+    config: AppConfig
+    search_results: list[SearchResult] | None = None
+    session_state: ChatSessionState | None = None
+
+
+@dataclass
+class SearchDeps:
+    """Dependencies for search agent."""
+
+    client: HaikuRAG
+    config: AppConfig
+    filter: str | None = None
+    search_results: list[SearchResult] = field(default_factory=list)
+
+
+def build_document_filter(document_name: str) -> str:
+    """Build SQL filter for document name matching."""
+    escaped = document_name.replace("'", "''")
+    no_spaces = escaped.replace(" ", "")
+    return (
+        f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
+        f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
+    )

haiku_rag_slim/haiku/rag/agents/qa/__init__.py

@@ -1,7 +1,7 @@
+from haiku.rag.agents.qa.agent import QuestionAnswerAgent
+from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.qa.agent import QuestionAnswerAgent
-from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.utils import build_prompt
 
 

haiku_rag_slim/haiku/rag/agents/qa/agent.py

@@ -2,10 +2,14 @@ from pydantic import BaseModel
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.output import ToolOutput
 
+from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
+from haiku.rag.agents.research.models import (
+    Citation,
+    RawSearchAnswer,
+    resolve_citations,
+)
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig, ModelConfig
-from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
-from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import get_model
 

haiku_rag_slim/haiku/rag/agents/research/__init__.py

@@ -0,0 +1,6 @@
+from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
+from haiku.rag.agents.research.models import (
+    EvaluationResult,
+    ResearchReport,
+    SearchAnswer,
+)

haiku_rag_slim/haiku/rag/agents/research/dependencies.py

@@ -6,7 +6,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.store.models import SearchResult
 
 if TYPE_CHECKING:
-    from haiku.rag.graph.research.models import SearchAnswer
+    from haiku.rag.agents.research.models import SearchAnswer
 
 
 class ResearchContext(BaseModel):

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -5,10 +5,8 @@ from pydantic_ai.output import ToolOutput
 from pydantic_graph.beta import Graph, GraphBuilder, StepContext
 from pydantic_graph.beta.join import reduce_list_append
 
-from haiku.rag.config import Config
-from haiku.rag.config.models import AppConfig
-from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.graph.research.models import (
+from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
+from haiku.rag.agents.research.models import (
     Citation,
     ConversationalAnswer,
     EvaluationResult,
@@ -17,7 +15,7 @@ from haiku.rag.graph.research.models import (
     ResearchReport,
     SearchAnswer,
 )
-from haiku.rag.graph.research.prompts import (
+from haiku.rag.agents.research.prompts import (
     CONVERSATIONAL_SYNTHESIS_PROMPT,
     DECISION_PROMPT,
     PLAN_PROMPT,
@@ -25,7 +23,9 @@ from haiku.rag.graph.research.prompts import (
     SEARCH_PROMPT,
     SYNTHESIS_PROMPT,
 )
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
+from haiku.rag.config import Config
+from haiku.rag.config.models import AppConfig
 from haiku.rag.utils import build_prompt, get_model
 
 

haiku_rag_slim/haiku/rag/agents/research/state.py

@@ -4,9 +4,9 @@ from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, Field
 
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.models import EvaluationResult
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.models import EvaluationResult
 
 if TYPE_CHECKING:
     from haiku.rag.config.models import AppConfig

haiku_rag_slim/haiku/rag/app.py

@@ -18,11 +18,11 @@ from rich.progress import (
     TransferSpeedColumn,
 )
 
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_research_graph
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.client import HaikuRAG, RebuildMode
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.graph import build_research_graph
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 from haiku.rag.mcp import create_mcp_server
 from haiku.rag.monitor import FileWatcher
 from haiku.rag.store.models.document import Document

haiku_rag_slim/haiku/rag/client.py

@@ -28,7 +28,7 @@ from haiku.rag.store.repositories.settings import SettingsRepository
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
 
-    from haiku.rag.graph.research.models import Citation
+    from haiku.rag.agents.research.models import Citation
 
 logger = logging.getLogger(__name__)
 
@@ -1275,7 +1275,7 @@ class HaikuRAG:
         Returns:
             Tuple of (answer text, list of resolved citations).
         """
-        from haiku.rag.qa import get_qa_agent
+        from haiku.rag.agents.qa import get_qa_agent
 
         qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
         return await qa_agent.answer(question, filter=filter)

haiku_rag_slim/haiku/rag/graph/__init__.py

@@ -1,5 +0,0 @@
-from haiku.rag.graph.research.graph import build_research_graph
-
-__all__ = [
-    "build_research_graph",
-]

haiku_rag_slim/haiku/rag/graph/research/__init__.py

@@ -1,6 +0,0 @@
-from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.graph.research.models import (
-    EvaluationResult,
-    ResearchReport,
-    SearchAnswer,
-)

haiku_rag_slim/haiku/rag/mcp.py

@@ -4,9 +4,9 @@ from typing import Any
 from fastmcp import FastMCP
 from pydantic import BaseModel
 
+from haiku.rag.agents.research.models import ResearchReport
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.graph.research.models import ResearchReport
 from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import format_citations
 
@@ -186,9 +186,9 @@ def create_mcp_server(
         try:
             async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 if deep:
-                    from haiku.rag.graph.research.dependencies import ResearchContext
-                    from haiku.rag.graph.research.graph import build_research_graph
-                    from haiku.rag.graph.research.state import (
+                    from haiku.rag.agents.research.dependencies import ResearchContext
+                    from haiku.rag.agents.research.graph import build_research_graph
+                    from haiku.rag.agents.research.state import (
                         ResearchDeps,
                         ResearchState,
                     )
@@ -230,9 +230,9 @@ def create_mcp_server(
             A research report with findings, or None if an error occurred.
         """
         try:
-            from haiku.rag.graph.research.dependencies import ResearchContext
-            from haiku.rag.graph.research.graph import build_research_graph
-            from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+            from haiku.rag.agents.research.dependencies import ResearchContext
+            from haiku.rag.agents.research.graph import build_research_graph
+            from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 
             async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 graph = build_research_graph(config=config)

haiku_rag_slim/haiku/rag/utils.py

@@ -10,8 +10,8 @@ from packaging.version import Version, parse
 if TYPE_CHECKING:
     from rich.console import RenderableType
 
+    from haiku.rag.agents.research.models import Citation
     from haiku.rag.config.models import AppConfig, ModelConfig
-    from haiku.rag.graph.research.models import Citation
 
 
 def parse_datetime(s: str) -> datetime:

tests/agents/chat/test_chat_agent.py

@@ -0,0 +1,105 @@
+from haiku.rag.agents.chat import (
+    ChatDeps,
+    ChatSessionState,
+    CitationInfo,
+    QAResponse,
+    SearchAgent,
+    create_chat_agent,
+)
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+
+
+def test_create_chat_agent():
+    """Test that create_chat_agent returns a properly configured agent."""
+    agent = create_chat_agent(Config)
+    assert agent is not None
+    assert agent.name == "chat_agent" or agent.name is None
+
+
+def test_chat_deps_initialization(temp_db_path):
+    """Test ChatDeps can be initialized with required fields."""
+    client = HaikuRAG(temp_db_path, create=True)
+    deps = ChatDeps(client=client, config=Config)
+
+    assert deps.client is client
+    assert deps.config is Config
+    assert deps.search_results is None
+    assert deps.session_state is None
+
+    client.close()
+
+
+def test_chat_session_state():
+    """Test ChatSessionState model."""
+    state = ChatSessionState(session_id="test-session")
+    assert state.session_id == "test-session"
+    assert state.citations == []
+    assert state.qa_history == []
+
+
+def test_citation_info():
+    """Test CitationInfo model."""
+    citation = CitationInfo(
+        index=1,
+        document_id="doc-123",
+        chunk_id="chunk-456",
+        document_uri="test.md",
+        document_title="Test Document",
+        page_numbers=[1, 2],
+        headings=["Section 1"],
+        content="Test content",
+    )
+    assert citation.index == 1
+    assert citation.document_id == "doc-123"
+    assert citation.chunk_id == "chunk-456"
+    assert citation.content == "Test content"
+
+
+def test_qa_response():
+    """Test QAResponse model."""
+    citation = CitationInfo(
+        index=1,
+        document_id="doc-123",
+        chunk_id="chunk-456",
+        document_uri="test.md",
+        document_title="Test Document",
+        content="Test content",
+    )
+    qa = QAResponse(
+        question="What is this?",
+        answer="This is a test",
+        confidence=0.95,
+        citations=[citation],
+    )
+    assert qa.question == "What is this?"
+    assert qa.answer == "This is a test"
+    assert qa.confidence == 0.95
+    assert len(qa.citations) == 1
+    assert qa.sources == ["Test Document"]
+
+
+def test_qa_response_sources_with_uri_fallback():
+    """Test QAResponse.sources falls back to URI when title is None."""
+    citation = CitationInfo(
+        index=1,
+        document_id="doc-123",
+        chunk_id="chunk-456",
+        document_uri="test.md",
+        document_title=None,
+        content="Test content",
+    )
+    qa = QAResponse(
+        question="What is this?",
+        answer="This is a test",
+        citations=[citation],
+    )
+    assert qa.sources == ["test.md"]
+
+
+def test_search_agent_initialization(temp_db_path):
+    """Test SearchAgent can be initialized."""
+    client = HaikuRAG(temp_db_path, create=True)
+    search_agent = SearchAgent(client, Config)
+    assert search_agent is not None
+    client.close()

tests/agents/qa/test_qa.py

@@ -5,16 +5,16 @@ import pytest
 from datasets import Dataset
 from evaluations.evaluators import LLMJudge
 
+from haiku.rag.agents.qa.agent import QuestionAnswerAgent
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import ModelConfig
-from haiku.rag.qa.agent import QuestionAnswerAgent
 
 HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
 
 
 @pytest.fixture(scope="module")
 def vcr_cassette_dir():
-    return str(Path(__file__).parent / "cassettes" / "test_qa")
+    return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
 
 
 @pytest.mark.vcr()

tests/agents/research/test_research_graph.py

@@ -1,10 +1,10 @@
 import pytest
 
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_research_graph
+from haiku.rag.agents.research.models import ResearchReport
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.graph import build_research_graph
-from haiku.rag.graph.research.models import ResearchReport
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 
 
 @pytest.mark.vcr()

tests/agents/research/test_search_filter.py

@@ -1,9 +1,9 @@
 import pytest
 
+from haiku.rag.agents.research.dependencies import ResearchContext
+from haiku.rag.agents.research.graph import build_research_graph
+from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.research.dependencies import ResearchContext
-from haiku.rag.graph.research.graph import build_research_graph
-from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 
 
 @pytest.fixture

tests/graph/__init__.py

@@ -1 +0,0 @@
-"""Tests for haiku.rag.graph module."""

tests/test_app.py

@@ -310,7 +310,7 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
 @pytest.mark.asyncio
 async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
     """Test asking a question with citations."""
-    from haiku.rag.graph.research.models import Citation
+    from haiku.rag.agents.research.models import Citation
 
     mock_answer = "Test answer with citations"
     mock_citations = [
@@ -342,7 +342,7 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
 async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
     """Test asking a question with deep mode uses research graph."""
     import haiku.rag.app as app_module
-    from haiku.rag.graph.research.models import ResearchReport
+    from haiku.rag.agents.research.models import ResearchReport
 
     mock_output = ResearchReport(
         title="Test",
@@ -380,7 +380,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
 async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
     """Test asking a question with deep mode (cite is ignored for research graph)."""
     import haiku.rag.app as app_module
-    from haiku.rag.graph.research.models import ResearchReport
+    from haiku.rag.agents.research.models import ResearchReport
 
     mock_output = ResearchReport(
         title="Test",

tests/test_mcp.py

@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch
 
 import pytest
 
-from haiku.rag.graph.research.models import ResearchReport
+from haiku.rag.agents.research.models import ResearchReport
 from haiku.rag.mcp import create_mcp_server
 from haiku.rag.store.models.document import Document
 
@@ -251,7 +251,7 @@ async def test_mcp_ask_question_deep():
         with (
             patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
             patch(
-                "haiku.rag.graph.research.graph.build_research_graph"
+                "haiku.rag.agents.research.graph.build_research_graph"
             ) as mock_graph_builder,
         ):
             mock_rag = AsyncMock()
@@ -295,7 +295,7 @@ async def test_mcp_research_question():
         with (
             patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
             patch(
-                "haiku.rag.graph.research.graph.build_research_graph"
+                "haiku.rag.agents.research.graph.build_research_graph"
             ) as mock_graph_builder,
         ):
             mock_rag = AsyncMock()