Support Pydantic AI agent specs via from_spec
`Agent.from_spec` raised `TypeError` on `RAGCapability` and `AnalysisCapability`, whose constructors take a state class, packaged instruction text and a tool-name set, and silently omitted both from the generated spec schema. The two zero-configuration capabilities constructed but with `id=None`, so pydantic-ai's duplicate-id rejection no longer held and a spec could register two citation policies, defeating the single-decision-maker invariant. Override `from_spec` on all four, delegating to `create_capability()` so ids and instructions come from one place. The spec surface is `db_path`, `config`, `defer_loading`, `request_limit` and `vision`; a live `HaikuRAG` client stays out of it, and a `config` mapping is validated through `AppConfig`.
This commit is contained in:
parent
971ec0a5b0
commit
61a756da7f
7 changed files with 274 additions and 0 deletions
|
|
@ -1,6 +1,12 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
### Fixed
|
||||||
|
|
||||||
- `Store`, `HaikuRAG` and `create_capability` coerce a string `db_path` to `Path`, as the documented
|
- `Store`, `HaikuRAG` and `create_capability` coerce a string `db_path` to `Path`, as the documented
|
||||||
|
|
|
||||||
|
|
@ -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
|
search too. If you need computation, register the analysis capability alone rather
|
||||||
than adding it to the RAG one.
|
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
|
## 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.
|
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.
|
||||||
|
|
|
||||||
|
|
@ -132,6 +132,29 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
|
||||||
)
|
)
|
||||||
return result.stdout or "No output."
|
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]:
|
def get_toolset(self) -> FunctionToolset[Any]:
|
||||||
async def analysis_search(
|
async def analysis_search(
|
||||||
ctx: RunContext[Any], query: str, limit: int | None = None
|
ctx: RunContext[Any], query: str, limit: int | None = None
|
||||||
|
|
|
||||||
|
|
@ -327,6 +327,12 @@ class EvidenceCompactionCapability(AbstractCapability[Any]):
|
||||||
capsule: Capsule = field(default_factory=Capsule, repr=False)
|
capsule: Capsule = field(default_factory=Capsule, repr=False)
|
||||||
images: tuple[str | BinaryContent, ...] = field(default=(), 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":
|
async def for_run(self, ctx: RunContext[Any]) -> "EvidenceCompactionCapability":
|
||||||
"""Give the run its own build cache, so concurrent runs cannot share one."""
|
"""Give the run its own build cache, so concurrent runs cannot share one."""
|
||||||
return replace(self, built_for=None, capsule=Capsule(), images=())
|
return replace(self, built_for=None, capsule=Capsule(), images=())
|
||||||
|
|
|
||||||
|
|
@ -72,6 +72,12 @@ class CitationPolicyCapability(AbstractCapability[Any]):
|
||||||
would share this capability's id.
|
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(
|
async def after_model_request(
|
||||||
self,
|
self,
|
||||||
ctx: RunContext[Any],
|
ctx: RunContext[Any],
|
||||||
|
|
|
||||||
|
|
@ -44,6 +44,29 @@ def instructions() -> str:
|
||||||
class RAGCapability(RAGCapabilityBase[RAGState]):
|
class RAGCapability(RAGCapabilityBase[RAGState]):
|
||||||
"""Deferred, native Pydantic AI capability for grounded RAG queries."""
|
"""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]:
|
def get_toolset(self) -> FunctionToolset[Any]:
|
||||||
async def rag_search(
|
async def rag_search(
|
||||||
ctx: RunContext[Any], query: str, limit: int | None = None
|
ctx: RunContext[Any], query: str, limit: int | None = None
|
||||||
|
|
|
||||||
150
tests/capabilities/test_agent_spec.py
Normal file
150
tests/capabilities/test_agent_spec.py
Normal file
|
|
@ -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"]
|
||||||
Loading…
Reference in a new issue