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

Add feature-based chat agent & prompt composition

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-10 13:39:12 +02:00
parent da47e3e345
commit ce3c3e7ed5
No known key found for this signature in database
5 changed files with 346 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
docs/agents.md
View file

@ -102,6 +102,40 @@ async with HaikuRAG(path_to_db) as client:
print(result.output)
```
### Feature Selection
By default, `create_chat_agent` enables search, documents, and QA toolsets. You can customize which capabilities the agent has via the `features` parameter:
```python
from haiku.rag.agents.chat import (
create_chat_agent,
FEATURE_SEARCH,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_ANALYSIS,
)
# Search-only agent
agent = create_chat_agent(config, client, context, features=[FEATURE_SEARCH])
# All features including code analysis
agent = create_chat_agent(
config, client, context,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
```
Available features:
| Feature | Constant | Tools added |
|---------|----------|-------------|
| Search | `FEATURE_SEARCH` | `search` |
| Documents | `FEATURE_DOCUMENTS` | `list_documents`, `get_document`, `summarize_document` |
| QA | `FEATURE_QA` | `ask` |
| Analysis | `FEATURE_ANALYSIS` | `analyze` |
The system prompt is automatically composed to match the selected features — only guidance for active tools is included. `SessionState` is always registered (shared by all features). `QASessionState` is only registered when the QA feature is active.
### Session State
The `ChatSessionState` maintains:

12
haiku_rag_slim/haiku/rag/agents/chat/__init__.py
View file

@ -1,4 +1,9 @@
from haiku.rag.agents.chat.agent import (
DEFAULT_FEATURES,
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
ChatDeps,
create_chat_agent,
run_chat_agent,
@ -8,6 +13,7 @@ from haiku.rag.agents.chat.context import (
summarize_session,
update_session_context,
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
@ -19,6 +25,12 @@ from haiku.rag.tools.qa import QAHistoryEntry
__all__ = [
"AGUI_STATE_KEY",
"DEFAULT_FEATURES",
"FEATURE_ANALYSIS",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_SEARCH",
"build_chat_prompt",
"create_chat_agent",
"run_chat_agent",
"trigger_background_summarization",

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

@ -1,6 +1,6 @@
import uuid
from dataclasses import dataclass
from typing import Any, cast
from typing import Any
from pydantic_ai import Agent
@ -10,7 +10,7 @@ from haiku.rag.agents.chat.context import (
from haiku.rag.agents.chat.context import (
trigger_background_summarization as _trigger_summarization,
)
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
@ -29,6 +29,13 @@ from haiku.rag.tools.search import create_search_toolset
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.utils import get_model
FEATURE_SEARCH = "search"
FEATURE_DOCUMENTS = "documents"
FEATURE_QA = "qa"
FEATURE_ANALYSIS = "analysis"
DEFAULT_FEATURES = [FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA]
@dataclass
class ChatDeps:
@ -162,6 +169,7 @@ def create_chat_agent(
config: AppConfig,
client: HaikuRAG,
context: ToolContext,
features: list[str] | None = None,
) -> Agent[ChatDeps, str]:
"""Create the chat agent with composed toolsets.
@ -169,8 +177,11 @@ def create_chat_agent(
config: Application configuration.
client: HaikuRAG client for database operations.
context: ToolContext for shared state across toolsets.
Should have SessionState and QASessionState registered
(will be auto-registered if not present).
SessionState is always registered. QASessionState is
registered only when the QA feature is active.
features: List of features to enable. Defaults to DEFAULT_FEATURES
(search, documents, qa). Available features: "search",
"documents", "qa", "analysis".
Returns:
The configured chat agent.
@ -182,37 +193,46 @@ def create_chat_agent(
deps = ChatDeps(config=config, tool_context=context)
result = await agent.run("Search for X", deps=deps)
"""
# Ensure session states are registered with proper AG-UI state key
if features is None:
features = DEFAULT_FEATURES
# SessionState is always registered (shared by all features)
existing = context.get(SESSION_NAMESPACE, SessionState)
if existing is None:
context.register(SESSION_NAMESPACE, SessionState(state_key=AGUI_STATE_KEY))
elif existing.state_key is None:
existing.state_key = AGUI_STATE_KEY
if context.get(QA_SESSION_NAMESPACE, QASessionState) is None:
context.register(QA_SESSION_NAMESPACE, QASessionState())
# Create toolsets - these capture client, config, and context in closures
search_toolset = create_search_toolset(client, config, context=context)
document_toolset = create_document_toolset(client, config, context=context)
qa_toolset = create_qa_toolset(client, config, context=context)
# QASessionState only when QA feature is active
if FEATURE_QA in features:
if context.get(QA_SESSION_NAMESPACE, QASessionState) is None:
context.register(QA_SESSION_NAMESPACE, QASessionState())
# Create toolsets conditionally based on features
toolsets = []
if FEATURE_SEARCH in features:
toolsets.append(create_search_toolset(client, config, context=context))
if FEATURE_DOCUMENTS in features:
toolsets.append(create_document_toolset(client, config, context=context))
if FEATURE_QA in features:
toolsets.append(create_qa_toolset(client, config, context=context))
if FEATURE_ANALYSIS in features:
from haiku.rag.tools.analysis import create_analysis_toolset
toolsets.append(create_analysis_toolset(client, config, context=context))
# Create the agent with composed toolsets
model = get_model(config.qa.model, config)
agent = cast(
Agent[ChatDeps, str],
Agent(
model,
deps_type=ChatDeps,
output_type=str,
instructions=CHAT_SYSTEM_PROMPT,
toolsets=[search_toolset, document_toolset, qa_toolset], # type: ignore[arg-type]
retries=3,
),
return Agent(
model,
deps_type=ChatDeps,
output_type=str,
instructions=build_chat_prompt(features),
toolsets=toolsets,
retries=3,
)
return agent
def trigger_background_summarization(deps: ChatDeps) -> None:
"""Trigger background session summarization if qa_history has entries.
@ -269,4 +289,9 @@ __all__ = [
"ChatSessionState",
"SessionContext",
"AGUI_STATE_KEY",
"FEATURE_SEARCH",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_ANALYSIS",
"DEFAULT_FEATURES",
]

105
haiku_rag_slim/haiku/rag/agents/chat/prompts.py
View file

@ -1,33 +1,112 @@
CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
_PROMPT_BASE = """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 automatically uses prior conversation context
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
2. NEVER call the same tool multiple times for a single user message
3. NEVER make up information - always use tools to get facts from the knowledge base"""
How to decide which tool to use:
_PROMPT_QA_RULES = """
4. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context"""
_PROMPT_SEARCH_RULES = """
5. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally"""
_PROMPT_TOOL_HEADER = """
How to decide which tool to use:"""
_PROMPT_DOCUMENTS = """
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return 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.
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document")."""
_PROMPT_QA = """
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations."""
_PROMPT_SEARCH = """
- "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."""
_PROMPT_ANALYSIS = """
- "analyze" - Use when the user asks for computation, data analysis, or quantitative tasks that require code execution (e.g., "calculate the average", "compare the numbers", "plot the data"). It runs Python code in a sandbox to produce results."""
_PROMPT_DOCUMENT_NAME_HEADER = """
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`
- Extract the DOCUMENT NAME as `document_name`"""
_PROMPT_SEARCH_EXAMPLES = """
- Examples for search:
- "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
- "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566"
- "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566" """
_PROMPT_QA_EXAMPLES = """
- Examples for ask:
- "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
- "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566"
- "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566" """
Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user."""
_PROMPT_CLOSING = """
Be friendly and conversational."""
_PROMPT_QA_CLOSING = (
""" When you use the "ask" tool, summarize the key findings for the user."""
)
def build_chat_prompt(features: list[str]) -> str:
"""Build a chat system prompt from the given feature list.
Each feature adds its relevant tool guidance to the prompt.
The base identity, critical rules, and closing are always included.
Args:
features: List of feature names (e.g., ["search", "documents", "qa"]).
Returns:
The composed system prompt string.
"""
parts = [_PROMPT_BASE]
# Add feature-specific critical rules
if "qa" in features:
parts.append(_PROMPT_QA_RULES)
if "search" in features:
parts.append(_PROMPT_SEARCH_RULES)
# Tool guidance header + per-feature sections
tool_sections = []
if "documents" in features:
tool_sections.append(_PROMPT_DOCUMENTS)
if "qa" in features:
tool_sections.append(_PROMPT_QA)
if "search" in features:
tool_sections.append(_PROMPT_SEARCH)
if "analysis" in features:
tool_sections.append(_PROMPT_ANALYSIS)
if tool_sections:
parts.append(_PROMPT_TOOL_HEADER)
parts.extend(tool_sections)
# Document name examples (relevant when search or qa is active)
if "search" in features or "qa" in features:
parts.append(_PROMPT_DOCUMENT_NAME_HEADER)
if "search" in features:
parts.append(_PROMPT_SEARCH_EXAMPLES)
if "qa" in features:
parts.append(_PROMPT_QA_EXAMPLES)
parts.append(_PROMPT_CLOSING)
if "qa" in features:
parts.append(_PROMPT_QA_CLOSING)
return "".join(parts)
CHAT_SYSTEM_PROMPT = build_chat_prompt(["search", "documents", "qa"])
SESSION_SUMMARY_PROMPT = """You are a session summarizer. Given a conversation history of Q&A pairs (and optionally existing context), produce a structured summary that captures key information for future context.

160
tests/agents/chat/test_features.py Normal file
View file

@ -0,0 +1,160 @@
from pydantic_ai import FunctionToolset
from haiku.rag.agents.chat.agent import (
DEFAULT_FEATURES,
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
ChatDeps,
create_chat_agent,
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
def _count_function_toolsets(agent) -> int:
"""Count FunctionToolset instances in an agent (excludes internal toolsets)."""
return sum(1 for t in agent.toolsets if type(t) is FunctionToolset)
# =============================================================================
# Feature Selection Tests
# =============================================================================
def test_default_features(temp_db_path):
"""Default features create search + document + qa toolsets and register both states."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
# Should have 3 toolsets (search, document, qa)
assert _count_function_toolsets(agent) == 3
# Both SessionState and QASessionState should be registered
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is not None
client.close()
def test_search_only(temp_db_path):
"""features=["search"] creates only search toolset, no QASessionState."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
agent = create_chat_agent(Config, client, context, features=[FEATURE_SEARCH])
assert _count_function_toolsets(agent) == 1
# SessionState always registered, but QASessionState should NOT be
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
client.close()
def test_search_and_documents(temp_db_path):
"""features=["search", "documents"] creates both toolsets, no QASessionState."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
agent = create_chat_agent(
Config, client, context, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS]
)
assert _count_function_toolsets(agent) == 2
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
client.close()
def test_all_features(temp_db_path):
"""All four features create four toolsets."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
agent = create_chat_agent(
Config,
client,
context,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
assert _count_function_toolsets(agent) == 4
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is not None
client.close()
def test_no_qa_skips_qa_session_state(temp_db_path):
"""Without QA feature, QASessionState is not registered."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
create_chat_agent(
Config, client, context, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS]
)
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
client.close()
def test_chat_deps_state_without_qa(temp_db_path):
"""ChatDeps.state getter omits qa_history/session_context when QASessionState absent."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
create_chat_agent(Config, client, context, features=[FEATURE_SEARCH])
deps = ChatDeps(config=Config, tool_context=context, session_id="test")
state = deps.state
assert "session_id" in state
# SessionState fields should be present
assert "document_filter" in state
assert "citation_registry" in state
assert "citations" in state
# QA fields should NOT be present
assert "qa_history" not in state
assert "session_context" not in state
client.close()
# =============================================================================
# Prompt Composition Tests
# =============================================================================
def test_build_chat_prompt_default():
"""Default features produce prompt mentioning all standard tools."""
prompt = build_chat_prompt(DEFAULT_FEATURES)
assert "list_documents" in prompt
assert "get_document" in prompt
assert "summarize_document" in prompt
assert "ask" in prompt
assert "search" in prompt
assert "analyze" not in prompt
def test_build_chat_prompt_search_only():
"""Search-only prompt doesn't mention ask or document tools."""
prompt = build_chat_prompt([FEATURE_SEARCH])
assert "search" in prompt
assert '"ask"' not in prompt
assert '"list_documents"' not in prompt
assert '"get_document"' not in prompt
assert '"summarize_document"' not in prompt
def test_build_chat_prompt_includes_analysis():
"""Analysis feature adds analyze guidance to prompt."""
prompt = build_chat_prompt(
[FEATURE_SEARCH, FEATURE_QA, FEATURE_DOCUMENTS, FEATURE_ANALYSIS]
)
assert "analyze" in prompt
assert "search" in prompt
assert "ask" in prompt