diff --git a/CHANGELOG.md b/CHANGELOG.md index 20827af2..51c1c6dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,12 @@ # 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 diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index f4b0e105..4988e84d 100644 --- a/docs/capabilities/index.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index e7cfeb3f..b0f6ba3f 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 05b57667..0528c1b5 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/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=()) diff --git a/haiku_rag_slim/haiku/rag/capabilities/policy.py b/haiku_rag_slim/haiku/rag/capabilities/policy.py index ce559c18..5dea32fb 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/policy.py +++ b/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], diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 6c7b2fca..46401483 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/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 diff --git a/tests/capabilities/test_agent_spec.py b/tests/capabilities/test_agent_spec.py new file mode 100644 index 00000000..c8c46423 --- /dev/null +++ b/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"]