diff --git a/README.md b/README.md index f7a55a45..ce63dd72 100644 --- a/README.md +++ b/README.md @@ -11,10 +11,10 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p - **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion - **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query -- **Question answering** — RAG skill with citations (page numbers, section headings) +- **Question answering** — RAG capability with citations (page numbers, section headings) - **Vision QA** — Vision-capable models receive figure bytes alongside chunk text - **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM -- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis) +- **Analysis capability** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis) - **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory - **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion - **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, Cohere, LM Studio, vLLM (multimodal via `multimodal: true` on vLLM/VoyageAI/Cohere). QA: any model supported by Pydantic AI diff --git a/docs/chat.md b/docs/chat.md index ca53876c..a0a0a05e 100644 --- a/docs/chat.md +++ b/docs/chat.md @@ -17,7 +17,7 @@ haiku-rag chat --model openai:gpt-4o ## How it works -The chat is a Pydantic AI agent with the deferred [RAG capability](capabilities/rag.md) attached by default. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and native tool events directly from Pydantic AI. +The chat is a Pydantic AI agent with the [RAG capability](capabilities/rag.md) attached by default. A single capability loads eagerly; when both RAG and analysis are enabled, they remain deferred until the model chooses which one to load. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and native tool events directly from Pydantic AI. The session is in-memory for the lifetime of the TUI. Conversation history is kept across turns so follow-up questions reuse prior context. Citations are tracked per turn and inspectable via the command palette. Clearing the chat resets the session and the agent's memory. diff --git a/docs/cli.md b/docs/cli.md index dbdd4dc7..7c4a086a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -195,7 +195,7 @@ Launch an interactive chat session for multi-turn conversations: haiku-rag chat haiku-rag chat --db /path/to/database.lancedb -# Enable analysis skill (code execution) +# Enable the analysis capability (code execution) haiku-rag chat -s rag -s analysis ``` diff --git a/docs/configuration/index.md b/docs/configuration/index.md index ba1103ee..ad55be29 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -121,7 +121,7 @@ doctor: min_chunks: 3 # documents with fewer chunks are excluded prompts: - domain_preamble: "" # Prepended to skill instructions + domain_preamble: "" # Prepended to capability instructions processing: converter: docling-local # docling-local or docling-serve diff --git a/docs/configuration/prompts.md b/docs/configuration/prompts.md index 19b828bd..b495d445 100644 --- a/docs/configuration/prompts.md +++ b/docs/configuration/prompts.md @@ -6,7 +6,7 @@ Customize the prompts used by haiku.rag's capabilities to match your domain. ```yaml prompts: - # Domain context prepended to skill instructions + # Domain context prepended to capability instructions domain_preamble: | This knowledge base contains technical documentation for the Helios solar panel system, including installation manuals, maintenance procedures, and safety guidelines. @@ -18,7 +18,7 @@ prompts: ## Domain Preamble -The `domain_preamble` field provides **domain context** prepended to the rag and rag-analysis skill instructions. Use this to: +The `domain_preamble` field provides **domain context** prepended to the RAG and analysis capability instructions. Use this to: - Describe what the knowledge base contains - Clarify domain-specific terminology diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index 29b04485..601a0165 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -20,7 +20,7 @@ Context expansion is automatic and section-aware. For structured documents (with ## Question Answering Configuration -Configure the rag skill (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool): +Configure the RAG capability (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool): ```yaml qa: @@ -34,15 +34,15 @@ qa: ``` - **model**: LLM configuration (see [Providers](providers.md#model-settings)) -- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The skill's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix. -- **max_searches**: Maximum number of search tool calls the rag skill can make per question (default: 3) +- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The capability's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix. +- **max_searches**: Maximum number of search tool calls the RAG capability can make per question (default: 3) !!! note "Thinking on vLLM" `enable_thinking` only applies to models with a pydantic-ai reasoning profile (o-series, gpt-5, gpt-oss). For other vLLM-served models such as Qwen3 or the Gemma family, the field is a silent no-op — set the chat template switch via [`extra_body`](providers.md#raw-provider-pass-through) instead. ## Analysis Configuration -Configure the analysis skill: +Configure the analysis capability: ```yaml analysis: @@ -58,6 +58,6 @@ analysis: - **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`. - **code_timeout**: Maximum seconds for each code execution (default: 60) - **max_output_chars**: Truncate code output after this many characters (default: 50000) -- **max_executions**: Maximum `execute_code` calls per question before the skill is told to answer from what it has (default: 15) +- **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15) See [Analysis capability](../capabilities/analysis.md) for usage details. diff --git a/docs/overview.md b/docs/overview.md index f5fc1afb..83170fb9 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -27,7 +27,7 @@ The chat TUI is one way to interact with the database. `haiku-rag ask` and `haik **Search.** Hybrid retrieval (vector + full-text with reciprocal rank fusion), optional cross-encoder reranking, structure-aware context expansion. Image-as-query and cross-modal retrieval when configured with a multimodal embedder. -**Answer.** RAG skill with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis skill with a sandboxed Python interpreter for aggregation and computation across documents. +**Answer.** RAG capability with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis capability with a sandboxed Python interpreter for aggregation and computation across documents. **Integrate.** Use it from Python, the CLI, the [MCP server](mcp.md), or through composable native Pydantic AI [capabilities](capabilities/index.md). diff --git a/docs/python.md b/docs/python.md index 1b7c9461..2fb175a6 100644 --- a/docs/python.md +++ b/docs/python.md @@ -360,7 +360,7 @@ See [Analysis capability](capabilities/analysis.md) for details and configuratio `client.ask` and `client.analyze` are convenience wrappers. To build your own Pydantic AI agent, attach the native RAG and analysis capabilities directly. See [Capabilities](capabilities/index.md). -For the low-level toolset factories under `haiku.rag.tools` (one rung below the skill abstraction), see [Toolsets](tools.md). +For the low-level toolset factories under `haiku.rag.tools` (one rung below the capability abstraction), see [Toolsets](tools.md). ## Importing Pre-Processed Documents diff --git a/docs/tuning.md b/docs/tuning.md index 79ffc657..1843dfa9 100644 --- a/docs/tuning.md +++ b/docs/tuning.md @@ -34,7 +34,7 @@ Context expansion is automatic and section-aware. Search results are expanded to Model and temperature selection affect answer quality directly. See [Providers](configuration/providers.md#model-settings) for options. -`domain_preamble` prepends domain context to the rag and rag-analysis skill instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md). +`domain_preamble` prepends domain context to the RAG and analysis capability instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md). ## What Requires a Rebuild @@ -46,7 +46,7 @@ Model and temperature selection affect answer quality directly. See [Providers]( ## Inspector -The inspector shows what your model would actually receive for a given query. Run it against your database and step through the same hybrid search, context expansion, and chunk previews the rag skill uses at runtime. Press `c` on a chunk and you see the exact context the LLM would get back from a search hit. +The inspector shows what your model would actually receive for a given query. Run it against your database and step through the same hybrid search, context expansion, and chunk previews the RAG capability uses at runtime. Press `c` on a chunk and you see the exact context the LLM would get back from a search hit. ```bash haiku-rag inspect @@ -81,11 +81,11 @@ Mouse: click to select, scroll to view content. ### Search -Press `/` to open the search modal. Type a query and press `Enter`. The left panel lists results with relevance scores like `[0.95] content preview`. The right panel shows the full chunk and its metadata. `↑` / `↓` navigates results, `Enter` jumps to the document and chunk, `Esc` closes the modal. Search uses the same hybrid (vector + full-text) retrieval the rag skill uses. +Press `/` to open the search modal. Type a query and press `Enter`. The left panel lists results with relevance scores like `[0.95] content preview`. The right panel shows the full chunk and its metadata. `↑` / `↓` navigates results, `Enter` jumps to the document and chunk, `Esc` closes the modal. Search uses the same hybrid (vector + full-text) retrieval the RAG capability uses. ### Context expansion (`c`) -Press `c` on a chunk to see the expanded context that would be fed to the rag skill. This is where you find out whether your `chunk_size`, `chunker_type`, and `max_context_chars` settings actually deliver the surrounding content the model needs. The modal shows: +Press `c` on a chunk to see the expanded context that would be fed to the RAG capability. This is where you find out whether your `chunk_size`, `chunker_type`, and `max_context_chars` settings actually deliver the surrounding content the model needs. The modal shows: - The expanded text. Section-aware expansion stays within section boundaries on structured documents and fills `max_context_chars` outward on unstructured ones. - Source document, content type, and relevance score. diff --git a/evaluations/README.md b/evaluations/README.md index c8bc6c6f..4241ca52 100644 --- a/evaluations/README.md +++ b/evaluations/README.md @@ -59,7 +59,7 @@ alongside QA accuracy from the URIs the capability registered via the `cite` too ### Debugging runs in Logfire With `LOGFIRE_TOKEN` set, runs ship spans under `service_name = 'evals'`. The -`debug-evals` capability in `.claude/capabilities/` turns these into ready-made Logfire +`debug-evals` skill in `.claude/skills/` turns these into ready-made Logfire queries (recent runs, per-case pass rate and `cited_map`, failing and slowest cases) for use from Claude Code. diff --git a/examples/README.md b/examples/README.md index a9e3deb7..51a9f3ec 100644 --- a/examples/README.md +++ b/examples/README.md @@ -16,7 +16,7 @@ See `docker/README.md` for setup instructions. **Script:** `custom_agent.py` -Uses the deferred RAG capability to build a conversational agent. +Uses the eagerly loaded RAG capability to build a conversational agent. ```bash uv run python examples/custom_agent.py /path/to/db.lancedb diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index cbb1e593..e645bfca 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -119,6 +119,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): [analysis_search, analysis_execute_code, analysis_cite], id=_CAPABILITY_ID, max_retries=3, + sequential=True, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index f46b2f41..4ef80293 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -57,7 +57,12 @@ class RAGCapability(RAGCapabilityBase[RAGState]): """Register exact search-result chunk IDs as citations for the answer.""" return await self._with_state(self._cite(chunk_ids)) - return FunctionToolset([rag_search, rag_cite], id=_CAPABILITY_ID, max_retries=3) + return FunctionToolset( + [rag_search, rag_cite], + id=_CAPABILITY_ID, + max_retries=3, + sequential=True, + ) def create_capability( diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index dcce7f99..ca4e048a 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -36,19 +36,24 @@ def run_chat( enabled = capabilities or ["rag"] capability_list = [] + defer_loading = len(enabled) > 1 if "rag" in enabled: from haiku.rag.capabilities.rag import create_capability capability_list.append( - create_capability(db_path=db_path, config=config, defer_loading=False) + create_capability( + db_path=db_path, config=config, defer_loading=defer_loading + ) ) if "analysis" in enabled: from haiku.rag.capabilities.analysis import create_capability capability_list.append( - create_capability(db_path=db_path, config=config, defer_loading=False) + create_capability( + db_path=db_path, config=config, defer_loading=defer_loading + ) ) app = ChatApp( diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 6d5f0ea1..a7eaa888 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -17,6 +17,7 @@ from pydantic_ai.messages import ( TextPartDelta, ) from pydantic_ai.run import AgentRunResultEvent +from pydantic_ai.usage import UsageLimits from textual.app import App, SystemCommand from textual.binding import Binding from textual.widgets import Footer, Header, Input @@ -85,6 +86,14 @@ class ChatApp(App): super().__init__() self.db_path = db_path self._capabilities = capabilities + request_limits = [ + capability.default_request_limit + for capability in capabilities + if capability.default_request_limit is not None + ] + self._usage_limits = ( + UsageLimits(request_limit=min(request_limits)) if request_limits else None + ) self.read_only = read_only self._model = model self.client: HaikuRAG | None = None @@ -191,6 +200,7 @@ class ChatApp(App): message_history=self._messages, conversation_id=self._conversation_id, deps=deps, + usage_limits=self._usage_limits, ) as stream: async for event in stream: if isinstance(event, PartStartEvent) and isinstance( diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 6b7cb7e1..42d9c4a4 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -44,7 +44,9 @@ def test_rag_capability_api(temp_db_path): assert capability.id == "haiku-rag" assert capability.defer_loading is True assert set(capability.get_toolset().tools) == {"rag_search", "rag_cite"} - assert capability.get_toolset().max_retries == 3 + toolset = capability.get_toolset() + assert toolset.max_retries == 3 + assert toolset.sequential is True assert capability.state_type is RAGState assert capability.state_namespace == "rag" @@ -60,7 +62,9 @@ def test_analysis_capability_api(temp_db_path): "analysis_execute_code", "analysis_cite", } - assert capability.get_toolset().max_retries == 3 + toolset = capability.get_toolset() + assert toolset.max_retries == 3 + assert toolset.sequential is True assert capability.state_type is AnalysisState diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 72b7bb30..645d9ad5 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, patch import pytest from typer.testing import CliRunner +from haiku.rag.capabilities.analysis import create_capability as create_analysis from haiku.rag.capabilities.rag import RAGState, create_capability from haiku.rag.cli import _cli as cli @@ -22,12 +23,28 @@ def test_chat_command(): def test_run_chat_creates_app_and_runs(temp_db_path: Path): - """Test run_chat() creates a ChatApp and calls run().""" - with patch("haiku.rag.chat.app.ChatApp.run") as mock_run: + """Test run_chat() eagerly attaches one capability and runs the app.""" + with patch("haiku.rag.chat.app.ChatApp") as mock_app: from haiku.rag.chat import run_chat run_chat(db_path=temp_db_path) - mock_run.assert_called_once() + + mock_app.return_value.run.assert_called_once() + attached = mock_app.call_args.kwargs["capabilities"] + assert len(attached) == 1 + assert attached[0].defer_loading is False + + +def test_run_chat_defers_multiple_capabilities(temp_db_path: Path): + """Test chat only defers capabilities when routing between multiple choices.""" + with patch("haiku.rag.chat.app.ChatApp") as mock_app: + from haiku.rag.chat import run_chat + + run_chat(db_path=temp_db_path, capabilities=["rag", "analysis"]) + + attached = mock_app.call_args.kwargs["capabilities"] + assert len(attached) == 2 + assert all(capability.defer_loading for capability in attached) def _make_mock_client(): @@ -66,6 +83,23 @@ def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None): ), mock_client +def test_chat_uses_capability_request_limit(temp_db_path: Path): + """Test chat applies the strictest request guard from its capabilities.""" + from haiku.rag.chat.app import ChatApp + + app = ChatApp( + db_path=temp_db_path, + capabilities=[ + create_capability(db_path=temp_db_path), + create_analysis(db_path=temp_db_path), + ], + read_only=True, + ) + + assert app._usage_limits is not None + assert app._usage_limits.request_limit == 30 + + @pytest.mark.asyncio async def test_chat_app_has_required_widgets(temp_db_path: Path): """Test that ChatApp has the required widgets: ChatHistory, Input."""