Merge pull request #577 from ggozad/feat/from_spec

Support Pydantic AI agent specs via from_spec

CHANGELOG.md

@@ -1,8 +1,16 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- The four capabilities support Pydantic AI agent specs via `from_spec`, registered with
+  `Agent.from_spec(..., custom_capability_types=[...])`. `RAGCapability` and `AnalysisCapability` take
+  `db_path`, `config`, `defer_loading`, `request_limit` and `vision`.
+
 ### Fixed
 
+- `Store`, `HaikuRAG` and `create_capability` coerce a string `db_path` to `Path`, as the documented
+  `HaikuRAG("knowledge.lancedb")` and `rag(db_path="my.lancedb")` forms require.
 - A capability search that matches nothing returns `No results found.` instead of an empty string.
 - Repeating a search query within a question accumulates the results of both calls instead of replacing the earlier ones.
 - `file://` URIs resolve to a Windows path through `url2pathname`: `file:///C:/docs/a.pdf` was read as `\C:\docs\a.pdf`, so ingestion reported `File does not exist` for every discovered file. A URI authority is kept as a UNC server/share (`file://server/share/a.pdf`) except `localhost`, which is dropped.

docs/capabilities/index.md

@@ -77,6 +77,66 @@ same way with either one, and neither exposes tools or takes configuration.
     search too. If you need computation, register the analysis capability alone rather
     than adding it to the RAG one.
 
+## Agent specs
+
+The capabilities can be declared in a Pydantic AI [agent spec](https://ai.pydantic.dev/agent-spec/):
+
+```yaml title="agent.yaml"
+model: openai:gpt-5
+instructions: You are a research assistant with access to a document knowledge base.
+capabilities:
+  - RAGCapability:
+      db_path: /data/kb.lancedb
+      defer_loading: false
+  - EvidenceCompactionCapability
+  - CitationPolicyCapability
+```
+
+Pydantic AI does not discover third-party capabilities, so the caller names the classes:
+
+```python
+from pydantic_ai import Agent
+
+from haiku.rag.capabilities.compaction import EvidenceCompactionCapability
+from haiku.rag.capabilities.policy import CitationPolicyCapability
+from haiku.rag.capabilities.rag import RAGCapability
+
+agent = Agent.from_file(
+    "agent.yaml",
+    deps_type=Deps,
+    custom_capability_types=[
+        RAGCapability,
+        EvidenceCompactionCapability,
+        CitationPolicyCapability,
+    ],
+)
+```
+
+`deps_type` stays a Python argument, since the capabilities read and write their state
+through `deps.state` (see [State](#state)). `Agent.from_file` reads YAML, which needs
+`pydantic-ai-slim[spec]`; `Agent.from_spec` takes a dict and needs no YAML parser.
+
+Set `defer_loading: false` when the agent registers a single evidence capability, so its
+tools are visible immediately. Leave it at the default when the model should route among
+multiple capabilities.
+
+A `config:` block accepts a whole `AppConfig`, for agents in one process that need
+different databases or embedding models:
+
+```yaml
+capabilities:
+  - RAGCapability:
+      db_path: /data/kb.lancedb
+      config:
+        embeddings:
+          model: {provider: ollama, name: embeddinggemma, vector_dim: 2048}
+```
+
+The block is read like a `haiku.rag.yaml` file: keys it omits take `AppConfig` defaults
+rather than values from the configuration file on disk. The embedding model must match the
+database; a mismatch may prevent opening it or produce invalid retrieval. Write the block in
+full or omit it and let the [configuration file](../configuration/index.md) apply.
+
 ## State
 
 Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol.

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -70,9 +70,9 @@ def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str:
     return match[0] if match else chunk_id
 
 
-def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
+def resolve_db_path(db_path: Path | str | None, config: AppConfig) -> Path:
     if db_path is not None:
-        return db_path
+        return Path(db_path)
     if env_db := os.environ.get("HAIKU_RAG_DB"):
         return Path(env_db).expanduser()
     return config.storage.data_dir / "haiku.rag.lancedb"

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -132,6 +132,29 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
             )
         return result.stdout or "No output."
 
+    @classmethod
+    def from_spec(
+        cls,
+        db_path: Path | None = None,
+        config: AppConfig | None = None,
+        *,
+        defer_loading: bool = True,
+        request_limit: int | None = 30,
+        vision: bool | None = None,
+    ) -> "AnalysisCapability":
+        """Build from an agent spec, mirroring the factory's serializable arguments.
+
+        A live ``HaikuRAG`` client cannot be written in a spec, so ``rag`` is
+        absent here. ``config`` arrives as a mapping and is validated.
+        """
+        return create_capability(
+            db_path,
+            AppConfig.model_validate(config) if config is not None else None,
+            defer_loading=defer_loading,
+            request_limit=request_limit,
+            vision=vision,
+        )
+
     def get_toolset(self) -> FunctionToolset[Any]:
         async def analysis_search(
             ctx: RunContext[Any], query: str, limit: int | None = None
@@ -156,7 +179,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
 
 
 def create_capability(
-    db_path: Path | None = None,
+    db_path: Path | str | None = None,
     config: AppConfig | None = None,
     *,
     defer_loading: bool = True,

haiku_rag_slim/haiku/rag/capabilities/compaction.py

@@ -327,6 +327,12 @@ class EvidenceCompactionCapability(AbstractCapability[Any]):
     capsule: Capsule = field(default_factory=Capsule, repr=False)
     images: tuple[str | BinaryContent, ...] = field(default=(), repr=False)
 
+    @classmethod
+    def from_spec(cls) -> "EvidenceCompactionCapability":
+        """Build from an agent spec. The factory takes no configuration, so
+        neither does the spec surface."""
+        return create_capability()
+
     async def for_run(self, ctx: RunContext[Any]) -> "EvidenceCompactionCapability":
         """Give the run its own build cache, so concurrent runs cannot share one."""
         return replace(self, built_for=None, capsule=Capsule(), images=())

haiku_rag_slim/haiku/rag/capabilities/policy.py

@@ -72,6 +72,12 @@ class CitationPolicyCapability(AbstractCapability[Any]):
     would share this capability's id.
     """
 
+    @classmethod
+    def from_spec(cls) -> "CitationPolicyCapability":
+        """Build from an agent spec. The factory takes no configuration, so
+        neither does the spec surface."""
+        return create_capability()
+
     async def after_model_request(
         self,
         ctx: RunContext[Any],

haiku_rag_slim/haiku/rag/capabilities/rag.py

@@ -44,6 +44,29 @@ def instructions() -> str:
 class RAGCapability(RAGCapabilityBase[RAGState]):
     """Deferred, native Pydantic AI capability for grounded RAG queries."""
 
+    @classmethod
+    def from_spec(
+        cls,
+        db_path: Path | None = None,
+        config: AppConfig | None = None,
+        *,
+        defer_loading: bool = True,
+        request_limit: int | None = 20,
+        vision: bool | None = None,
+    ) -> "RAGCapability":
+        """Build from an agent spec, mirroring the factory's serializable arguments.
+
+        A live ``HaikuRAG`` client cannot be written in a spec, so ``rag`` is
+        absent here. ``config`` arrives as a mapping and is validated.
+        """
+        return create_capability(
+            db_path,
+            AppConfig.model_validate(config) if config is not None else None,
+            defer_loading=defer_loading,
+            request_limit=request_limit,
+            vision=vision,
+        )
+
     def get_toolset(self) -> FunctionToolset[Any]:
         async def rag_search(
             ctx: RunContext[Any], query: str, limit: int | None = None
@@ -64,7 +87,7 @@ class RAGCapability(RAGCapabilityBase[RAGState]):
 
 
 def create_capability(
-    db_path: Path | None = None,
+    db_path: Path | str | None = None,
     config: AppConfig | None = None,
     *,
     defer_loading: bool = True,

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -66,7 +66,7 @@ class HaikuRAG:
 
     def __init__(
         self,
-        db_path: Path | None = None,
+        db_path: Path | str | None = None,
         config: AppConfig | None = None,
         skip_validation: bool = False,
         create: bool = False,
@@ -75,7 +75,8 @@ class HaikuRAG:
         """Initialize the RAG client with a database path.
 
         Args:
-            db_path: Path to the database file. If None, uses config.storage.data_dir.
+            db_path: Path or string path to the database file. If None, uses
+                config.storage.data_dir.
             config: Configuration to use. Defaults to the current global config.
             skip_validation: Whether to skip configuration validation on database load.
             create: Whether to create the database if it doesn't exist.

haiku_rag_slim/haiku/rag/store/engine.py

@@ -170,14 +170,14 @@ class TagInfo:
 class Store:
     def __init__(
         self,
-        db_path: Path,
+        db_path: Path | str,
         config: AppConfig | None = None,
         skip_validation: bool = False,
         create: bool = False,
         read_only: bool = False,
         skip_migration_check: bool = False,
     ):
-        self.db_path: Path = db_path
+        self.db_path: Path = Path(db_path)
         self._config = config if config is not None else get_config()
         self._read_only = read_only
         self._create = create
@@ -191,7 +191,7 @@ class Store:
         self._is_new_db = False
 
         if self._connection_mode == ConnectionMode.LOCAL:
-            if not db_path.exists():
+            if not self.db_path.exists():
                 if not create:
                     raise FileNotFoundError(
                         f"Database does not exist at {self.db_path.absolute()}. "
@@ -199,8 +199,8 @@ class Store:
                     )
                 self._is_new_db = True
                 # Ensure parent directories exist for new databases
-                if not db_path.parent.exists():
-                    Path.mkdir(db_path.parent, parents=True)
+                if not self.db_path.parent.exists():
+                    Path.mkdir(self.db_path.parent, parents=True)
 
         # Create embedder (sync — no LanceDB needed)
         self.embedder = get_embedder(config=self._config)

tests/capabilities/test_agent_spec.py

@@ -0,0 +1,150 @@
+import json
+from collections.abc import Sequence
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent
+from pydantic_ai.agent import AgentSpec
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.exceptions import UserError
+
+from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
+from haiku.rag.capabilities.compaction import CAPABILITY_ID as COMPACTION_ID
+from haiku.rag.capabilities.compaction import EvidenceCompactionCapability
+from haiku.rag.capabilities.policy import CAPABILITY_ID as POLICY_ID
+from haiku.rag.capabilities.policy import CitationPolicyCapability
+from haiku.rag.capabilities.rag import RAGCapability, RAGState
+
+ALL_CAPABILITIES = [
+    RAGCapability,
+    AnalysisCapability,
+    EvidenceCompactionCapability,
+    CitationPolicyCapability,
+]
+
+
+def _from_spec(
+    spec: dict, types: Sequence[type[AbstractCapability[Any]]]
+) -> list[AbstractCapability[Any]]:
+    """Build an agent from a spec and return the capabilities it declared.
+
+    A spec needs a model, and pydantic-ai injects capabilities of its own
+    alongside ours.
+    """
+    agent = Agent.from_spec({"model": "test", **spec}, custom_capability_types=types)
+    return [
+        capability
+        for capability in agent.root_capability.capabilities
+        if type(capability) in types
+    ]
+
+
+def test_rag_capability_is_built_from_a_spec(temp_db_path):
+    (capability,) = _from_spec(
+        {"capabilities": [{"RAGCapability": {"db_path": str(temp_db_path)}}]},
+        [RAGCapability],
+    )
+
+    assert isinstance(capability, RAGCapability)
+    assert capability.db_path == temp_db_path
+    assert capability.id == "haiku-rag"
+    assert capability.state_type is RAGState
+    assert capability.tool_names == {"rag_search", "rag_cite"}
+    assert capability.request_limit == 20
+
+
+def test_analysis_capability_is_built_from_a_spec(temp_db_path):
+    (capability,) = _from_spec(
+        {"capabilities": [{"AnalysisCapability": {"db_path": str(temp_db_path)}}]},
+        [AnalysisCapability],
+    )
+
+    assert isinstance(capability, AnalysisCapability)
+    assert capability.db_path == temp_db_path
+    assert capability.id == "haiku-rag-analysis"
+    assert capability.state_type is AnalysisState
+    assert capability.request_limit == 30
+
+
+def test_a_config_mapping_in_a_spec_is_validated(temp_db_path, temp_yaml_config):
+    """A `config:` block is validated into AppConfig rather than reaching
+    get_config()."""
+    (capability,) = _from_spec(
+        {
+            "capabilities": [
+                {
+                    "RAGCapability": {
+                        "db_path": str(temp_db_path),
+                        "config": {"qa": {"max_searches": 9}},
+                    }
+                }
+            ]
+        },
+        [RAGCapability],
+    )
+
+    assert isinstance(capability, RAGCapability)
+    assert capability.config.qa.max_searches == 9
+
+
+@pytest.mark.parametrize("form", ["bare", "empty-mapping"])
+def test_the_zero_argument_capabilities_are_built_from_a_spec(form):
+    """Their ids must be stamped: pydantic-ai rejects a duplicate id, which is
+    what keeps a single decision-maker per run."""
+    names = ["EvidenceCompactionCapability", "CitationPolicyCapability"]
+    entries: list[Any] = (
+        list(names) if form == "bare" else [{name: {}} for name in names]
+    )
+
+    compaction, policy = _from_spec(
+        {"capabilities": entries},
+        [EvidenceCompactionCapability, CitationPolicyCapability],
+    )
+
+    assert compaction.id == COMPACTION_ID
+    assert policy.id == POLICY_ID
+
+
+def test_a_spec_cannot_register_two_citation_policies():
+    with pytest.raises(UserError, match=POLICY_ID):
+        _from_spec(
+            {"capabilities": ["CitationPolicyCapability"] * 2},
+            [CitationPolicyCapability],
+        )
+
+
+def test_the_generated_spec_schema_describes_every_capability():
+    schema = AgentSpec.model_json_schema_with_capabilities(ALL_CAPABILITIES)
+    serialized = json.dumps(schema)
+
+    for capability in ALL_CAPABILITIES:
+        assert capability.__name__ in serialized
+
+    params = schema["$defs"]["spec_params_RAGCapability"]["properties"]
+    assert set(params) == {
+        "config",
+        "db_path",
+        "defer_loading",
+        "request_limit",
+        "vision",
+    }
+    assert params["config"] == {
+        "anyOf": [{"$ref": "#/$defs/AppConfig"}, {"type": "null"}]
+    }
+    assert "AppConfig" in schema["$defs"]
+    assert {"format": "path", "type": "string"} in params["db_path"]["anyOf"]
+
+    # Internal constructor wiring must not become a spec surface.
+    for internal in (
+        "state_type",
+        "instruction_text",
+        "tool_names",
+        "state_namespace",
+        "borrowed_rag",
+        "rag_lock",
+    ):
+        assert internal not in serialized
+
+    # A zero-argument from_spec leaves no params object at all, so the per-run
+    # build caches cannot be set from a spec.
+    assert "spec_params_EvidenceCompactionCapability" not in schema["$defs"]

tests/capabilities/test_capabilities.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, cast
 from unittest.mock import AsyncMock, patch
@@ -102,11 +103,33 @@ def test_capability_factories_resolve_environment_and_defaults(
         config.storage.data_dir / "haiku.rag.lancedb"
     )
 
+    for factory in (create_rag, create_analysis):
+        db_path = factory(db_path=str(temp_db_path), config=config).db_path
+        assert db_path == temp_db_path
+        assert isinstance(db_path, Path)
+
     with patch("haiku.rag.config.get_config", return_value=config):
         assert create_rag().config is config
         assert create_analysis().config is config
 
 
+@pytest.mark.asyncio
+async def test_a_string_db_path_opens_a_store(temp_db_path):
+    """Store calls `absolute()` and `exists()` on db_path, which a str lacks."""
+    from haiku.rag.client import HaikuRAG
+
+    config = AppConfig()
+    async with HaikuRAG(temp_db_path, config, create=True):
+        pass
+
+    capability = create_rag(db_path=str(temp_db_path), config=config)
+    try:
+        rag = await capability._ensure_rag()
+        assert rag.store.db_path == temp_db_path
+    finally:
+        await capability._close()
+
+
 def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
     config = AppConfig(
         prompts=PromptsConfig(domain_preamble="The corpus contains solar manuals.")

tests/test_client.py

@@ -30,6 +30,15 @@ def vcr_cassette_dir():
     return str(Path(__file__).parent / "cassettes" / "test_client")
 
 
+@pytest.mark.asyncio
+async def test_a_string_db_path_is_accepted(temp_db_path):
+    """The documented `HaikuRAG("knowledge.lancedb")` form: Store calls
+    `exists()` and `absolute()` on db_path, which a str lacks."""
+    async with HaikuRAG(str(temp_db_path), create=True) as client:
+        assert client.store.db_path == temp_db_path
+        assert isinstance(client.store.db_path, Path)
+
+
 @pytest.mark.asyncio
 async def test_prepare_document_from_docling_runs_off_event_loop_thread(monkeypatch):
     import haiku.rag.client.documents as documents