mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 10:55:35 +00:00
feat(logfire): add ManagedAgentSpec for whole-spec resolution
Build a `pydantic_ai.Agent` from a Logfire-managed agent-spec variable (`agentspec__<slug>`). The Logfire backend compiles `agentspec__*` server-side so the SDK consumes the already-assembled payload (instructions, model, model_settings, tools, mcp_servers, skills, capabilities). Callers pass their tool implementations and any harness capability classes they want the spec to be able to enable; everything else is sourced from Logfire. Unregistered capability `type` keys are silently dropped so the harness package stays optional. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
f460861435
commit
2d2f5a015d
@@ -5,13 +5,14 @@ from typing import TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
from .code_mode import CodeMode
|
||||
from .filesystem import FileSystem
|
||||
from .logfire import ManagedPrompt
|
||||
from .logfire import ManagedAgentSpec, ManagedPrompt
|
||||
from .shell import LLM_API_KEY_ENV_PATTERNS, Shell
|
||||
|
||||
__all__ = [
|
||||
'CodeMode',
|
||||
'FileSystem',
|
||||
'LLM_API_KEY_ENV_PATTERNS',
|
||||
'ManagedAgentSpec',
|
||||
'ManagedPrompt',
|
||||
'Shell',
|
||||
]
|
||||
@@ -26,6 +27,10 @@ def __getattr__(name: str) -> object:
|
||||
from .filesystem import FileSystem
|
||||
|
||||
return FileSystem
|
||||
if name == 'ManagedAgentSpec':
|
||||
from .logfire import ManagedAgentSpec
|
||||
|
||||
return ManagedAgentSpec
|
||||
if name == 'ManagedPrompt':
|
||||
from .logfire import ManagedPrompt
|
||||
|
||||
|
||||
@@ -9,6 +9,32 @@ Install the extra:
|
||||
pip install 'pydantic-ai-harness[logfire]'
|
||||
```
|
||||
|
||||
## `ManagedAgentSpec`
|
||||
|
||||
Build a whole `Agent` from a Logfire-managed spec — instructions, model, model_settings, tool/MCP/skill metadata, and harness-capability config all sourced from one variable. Pass your tool implementations and any harness capability classes you want the spec to be able to enable; the spec is the source of truth for everything else.
|
||||
|
||||
```python
|
||||
from pydantic_ai_harness import CodeMode
|
||||
from pydantic_ai_harness.logfire import ManagedAgentSpec
|
||||
|
||||
|
||||
def fetch_weather(city: str) -> str:
|
||||
return f'sunny and 72°F in {city}'
|
||||
|
||||
|
||||
spec = ManagedAgentSpec(
|
||||
'checkout_assistant',
|
||||
tools={'tool__weather': fetch_weather},
|
||||
capability_classes={'code_mode': CodeMode},
|
||||
)
|
||||
agent = await spec.build()
|
||||
result = await agent.run('What is the weather in Lisbon?')
|
||||
```
|
||||
|
||||
The Logfire backend compiles `agentspec__*` variables server-side, so one variable read returns the full assembled spec. Tool implementations stay in your code (function bodies aren't config) and capability classes stay in your code (the harness package is optional). Capability `type` keys in the spec that you haven't registered are silently dropped — a spec referencing `code_mode` doesn't break callers who haven't installed the extra.
|
||||
|
||||
See `_managed_agent_spec.py` for the full option surface (label / targeting_key / attributes / default_model / logfire_instance / `Variable[dict]` for typed specs).
|
||||
|
||||
## `ManagedPrompt`
|
||||
|
||||
Back an agent's instructions with a Logfire-managed
|
||||
@@ -17,8 +43,9 @@ Back an agent's instructions with a Logfire-managed
|
||||
> A broader, first-party `Managed` capability is in flight in
|
||||
> [pydantic-ai#5107](https://github.com/pydantic/pydantic-ai/pull/5107) and will eventually be
|
||||
> importable as `pydantic_ai.managed.logfire.Managed` -- covering instructions, model settings,
|
||||
> and whole-spec variables. Until then, `ManagedPrompt` is the supported path for backing
|
||||
> instructions with a Logfire-managed prompt.
|
||||
> and whole-spec variables. `ManagedAgentSpec` above is the harness path for
|
||||
> whole-spec resolution today; `ManagedPrompt` is the supported path for
|
||||
> backing instructions alone.
|
||||
|
||||
### The problem
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Logfire-backed capabilities: drive agent configuration from Logfire managed variables."""
|
||||
|
||||
from pydantic_ai_harness.logfire._managed_agent_spec import ManagedAgentSpec
|
||||
from pydantic_ai_harness.logfire._managed_prompt import ManagedPrompt
|
||||
|
||||
__all__ = ['ManagedPrompt']
|
||||
__all__ = ['ManagedAgentSpec', 'ManagedPrompt']
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
"""Back an agent's configuration with a Logfire-managed agent spec."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from dataclasses import KW_ONLY, dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import logfire
|
||||
from logfire.variables import Variable
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.exceptions import UserError
|
||||
from pydantic_ai.models import Model
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from logfire import Logfire
|
||||
from pydantic_ai.mcp import MCPToolset
|
||||
|
||||
|
||||
# Logfire stores agent specs as variables named `agentspec__<slug>`. The
|
||||
# OFREP-side compiler returns the assembled spec shape (instructions, model,
|
||||
# model_settings, tools, mcp_servers, skills, capabilities), so this module
|
||||
# never has to follow refs itself -- it just consumes the compiled payload.
|
||||
_AGENT_SPEC_VARIABLE_PREFIX = 'agentspec__'
|
||||
|
||||
# Empty default for the logfire Variable. When the variable hasn't been
|
||||
# published yet, the SDK returns this dict and the builder produces a
|
||||
# minimum-viable Agent (using the code-default model below).
|
||||
_EMPTY_SPEC: dict[str, Any] = {}
|
||||
|
||||
# Last-resort model when the spec doesn't pin one. Matches the snippet shown
|
||||
# in the Logfire UI's "Use this spec" card so the two stay in sync.
|
||||
_DEFAULT_MODEL = 'openai:gpt-5'
|
||||
|
||||
|
||||
def _empty_tools() -> dict[str, Callable[..., Any]]:
|
||||
return {}
|
||||
|
||||
|
||||
def _empty_capabilities() -> dict[str, type[AbstractCapability[Any]]]:
|
||||
return {}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManagedAgentSpec:
|
||||
"""Build a `pydantic_ai.Agent` from a Logfire-managed agent spec.
|
||||
|
||||
Pass the spec name (a `<slug>` from the Logfire UI's Agent Specs surface)
|
||||
plus your app's tool implementations and any harness capability classes
|
||||
you want the spec to be able to enable. Each `build()` call resolves the
|
||||
spec fresh and returns a ready-to-run Agent -- changes published from
|
||||
the Logfire UI take effect on the next call, no redeploy.
|
||||
|
||||
```python
|
||||
from pydantic_ai_harness import CodeMode
|
||||
from pydantic_ai_harness.logfire import ManagedAgentSpec
|
||||
|
||||
|
||||
def fetch_weather(city: str) -> str:
|
||||
return f'sunny and 72°F in {city}'
|
||||
|
||||
|
||||
spec = ManagedAgentSpec(
|
||||
'checkout_assistant',
|
||||
tools={'tool__weather': fetch_weather},
|
||||
capability_classes={'code_mode': CodeMode},
|
||||
)
|
||||
agent = await spec.build()
|
||||
result = await agent.run('What is the weather in Lisbon?')
|
||||
```
|
||||
|
||||
The Logfire UI owns instructions, model, model_settings, tool/MCP/skill
|
||||
metadata, and harness-capability config. Tool *implementations* and
|
||||
capability *classes* stay in your code -- tools because the function
|
||||
body is application logic, capabilities because the harness package is
|
||||
optional.
|
||||
|
||||
Pass an existing [`logfire.variables.Variable`][logfire.variables.Variable]
|
||||
as `name` when you want to use a variable you defined yourself (for
|
||||
example a typed `Variable[YourSpecModel]`).
|
||||
"""
|
||||
|
||||
name: str | Variable[dict[str, Any]]
|
||||
"""The agent-spec name (declared as the variable `agentspec__<name>`),
|
||||
or a pre-built `logfire.Variable[dict]`."""
|
||||
|
||||
_: KW_ONLY
|
||||
|
||||
tools: Mapping[str, Callable[..., Any]] = field(default_factory=_empty_tools)
|
||||
"""Map from a managed tool's variable name (e.g. `tool__weather`) to the
|
||||
Python callable that implements it. Tools the spec lists but you haven't
|
||||
mapped are silently skipped -- the model just won't see them."""
|
||||
|
||||
capability_classes: Mapping[str, type[AbstractCapability[Any]]] = field(default_factory=_empty_capabilities)
|
||||
"""Map from a capability `type` key in the spec (e.g. `code_mode`) to the
|
||||
class to instantiate. The class is called with `**config` from the spec.
|
||||
Unknown keys in the spec are silently dropped so the harness package can
|
||||
stay optional -- a spec referencing CodeMode doesn't break callers who
|
||||
haven't installed it."""
|
||||
|
||||
targeting_key: str | None = None
|
||||
"""`targetingKey` sent to Logfire's rollout/condition evaluation. When
|
||||
`None`, Logfire falls back to its own targeting context and then the
|
||||
active trace id."""
|
||||
|
||||
attributes: Mapping[str, Any] | None = None
|
||||
"""Attributes for condition-based targeting rules on the agent spec."""
|
||||
|
||||
label: str | None = None
|
||||
"""Explicit label on the managed spec to resolve (e.g. `'production'`).
|
||||
When `None`, the targeting rules pick the label."""
|
||||
|
||||
default_model: str | Model = _DEFAULT_MODEL
|
||||
"""Model to use when the spec doesn't pin one. Lets a Logfire-side
|
||||
rollback to an empty spec keep your app running. Accepts either a
|
||||
provider-prefixed model id or a pydantic-ai `Model` instance (e.g.
|
||||
`TestModel()` in tests)."""
|
||||
|
||||
logfire_instance: Logfire | None = None
|
||||
"""Logfire instance to resolve the variable on. When `None`, the global
|
||||
default instance is used. Ignored when `name` is a `Variable`."""
|
||||
|
||||
_variable: Variable[dict[str, Any]] = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.name, str):
|
||||
if self.logfire_instance is not None:
|
||||
warnings.warn(
|
||||
'`logfire_instance` is ignored when `name` is a `Variable`; '
|
||||
'the variable already carries its own Logfire instance.',
|
||||
stacklevel=2,
|
||||
)
|
||||
self._variable = self.name
|
||||
return
|
||||
|
||||
name = self.name
|
||||
if name.startswith(_AGENT_SPEC_VARIABLE_PREFIX):
|
||||
warnings.warn(
|
||||
f'The {_AGENT_SPEC_VARIABLE_PREFIX!r} prefix is added automatically; '
|
||||
f'pass the bare spec name rather than {name!r}.',
|
||||
stacklevel=2,
|
||||
)
|
||||
name = name[len(_AGENT_SPEC_VARIABLE_PREFIX) :]
|
||||
|
||||
variable_name = f'{_AGENT_SPEC_VARIABLE_PREFIX}{name.replace("-", "_")}'
|
||||
if not variable_name.isidentifier():
|
||||
raise ValueError(
|
||||
f'Agent-spec name {self.name!r} produces an invalid variable name {variable_name!r}; '
|
||||
'names may only contain letters, digits, hyphens, and underscores.'
|
||||
)
|
||||
|
||||
# Construct the Variable directly rather than via `logfire.var` so the
|
||||
# same spec can be built repeatedly without tripping the per-instance
|
||||
# duplicate-registration check.
|
||||
instance = self.logfire_instance if self.logfire_instance is not None else logfire.DEFAULT_LOGFIRE_INSTANCE
|
||||
self._variable = Variable(variable_name, type=dict, default=_EMPTY_SPEC, logfire_instance=instance)
|
||||
|
||||
async def build(self) -> Agent[None, str]:
|
||||
"""Resolve the spec and return a ready-to-run `Agent`.
|
||||
|
||||
Hits Logfire fresh on every call. Cache the returned `Agent` between
|
||||
requests if you want to amortise the resolve cost.
|
||||
"""
|
||||
resolved = self._variable.get(targeting_key=self.targeting_key, attributes=self.attributes, label=self.label)
|
||||
with resolved:
|
||||
# `resolved.value` is the dict the variable resolves to; cast it
|
||||
# to a typed `Mapping[str, Any]` so the wirers don't propagate
|
||||
# Unknowns through pyright.
|
||||
spec: Mapping[str, Any] = cast('Mapping[str, Any]', resolved.value)
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
'model': spec.get('model', self.default_model),
|
||||
'instructions': _instructions_with_skills(spec),
|
||||
'tools': _wire_tools(spec, self.tools),
|
||||
}
|
||||
model_settings = spec.get('model_settings')
|
||||
if isinstance(model_settings, dict) and model_settings:
|
||||
kwargs['model_settings'] = model_settings
|
||||
toolsets = _wire_mcp_servers(spec)
|
||||
if toolsets:
|
||||
kwargs['toolsets'] = toolsets
|
||||
capabilities = _wire_capabilities(spec, self.capability_classes)
|
||||
if capabilities:
|
||||
kwargs['capabilities'] = capabilities
|
||||
|
||||
return Agent(**kwargs)
|
||||
|
||||
|
||||
def _iter_dicts(raw: object) -> Iterable[dict[str, Any]]:
|
||||
"""Yield each `dict[str, Any]` entry from `raw`, skipping anything else.
|
||||
|
||||
`raw` comes from a resolved `Variable[dict]`, so it's `object`. The
|
||||
helper centralises the list-and-dict narrowing the four wirers all
|
||||
repeat, and keeps pyright happy with one cast instead of four.
|
||||
"""
|
||||
if not isinstance(raw, list):
|
||||
return
|
||||
for entry in cast('list[Any]', raw):
|
||||
if isinstance(entry, dict):
|
||||
yield cast('dict[str, Any]', entry)
|
||||
|
||||
|
||||
def _wire_tools(spec: Mapping[str, Any], impls: Mapping[str, Callable[..., Any]]) -> list[Callable[..., Any]]:
|
||||
"""Resolve each spec tool to its registered Python callable."""
|
||||
out: list[Callable[..., Any]] = []
|
||||
for entry in _iter_dicts(spec.get('tools')):
|
||||
name = entry.get('name')
|
||||
if not isinstance(name, str):
|
||||
continue
|
||||
impl = impls.get(name)
|
||||
if impl is not None:
|
||||
out.append(impl)
|
||||
return out
|
||||
|
||||
|
||||
def _wire_mcp_servers(spec: Mapping[str, Any]) -> list[MCPToolset]:
|
||||
"""Build `MCPToolset` for each compiled-spec MCP entry.
|
||||
|
||||
Lazy-imports `pydantic_ai.mcp` only when there's at least one valid
|
||||
entry to instantiate, so a spec without (usable) MCP servers doesn't
|
||||
require the `pydantic-ai-slim[mcp]` extra. Callers who include MCP
|
||||
servers in their spec must install the extra themselves; the import
|
||||
error fires at build time with a clear `pydantic-ai` suggestion.
|
||||
"""
|
||||
configs: list[tuple[str, dict[str, Any] | None]] = []
|
||||
for entry in _iter_dicts(spec.get('mcp_servers')):
|
||||
url = entry.get('url')
|
||||
if not isinstance(url, str) or not url:
|
||||
continue
|
||||
headers = entry.get('headers')
|
||||
configs.append((url, cast('dict[str, Any]', headers) if isinstance(headers, dict) else None))
|
||||
if not configs:
|
||||
return []
|
||||
|
||||
# Tested via a happy-path test that mocks `pydantic_ai.mcp.MCPToolset` --
|
||||
# the harness's default install doesn't include the `[mcp]` extra so we
|
||||
# can't actually instantiate one here.
|
||||
from pydantic_ai.mcp import MCPToolset
|
||||
|
||||
return [MCPToolset(url, headers=headers) for url, headers in configs]
|
||||
|
||||
|
||||
def _wire_capabilities(
|
||||
spec: Mapping[str, Any],
|
||||
classes: Mapping[str, type[AbstractCapability[Any]]],
|
||||
) -> list[AbstractCapability[Any]]:
|
||||
"""Instantiate each compiled-spec capability the caller has registered.
|
||||
|
||||
Capability entries the caller hasn't registered are silently dropped --
|
||||
`CodeMode`/`FileSystem`/`Shell` are optional installs, and a spec
|
||||
referencing one shouldn't fail to build just because the caller hasn't
|
||||
enabled it in their app.
|
||||
"""
|
||||
out: list[AbstractCapability[Any]] = []
|
||||
for entry in _iter_dicts(spec.get('capabilities')):
|
||||
type_key = entry.get('type')
|
||||
if not isinstance(type_key, str):
|
||||
continue
|
||||
cls = classes.get(type_key)
|
||||
if cls is None:
|
||||
continue
|
||||
config = entry.get('config')
|
||||
if config is None:
|
||||
out.append(cls())
|
||||
elif isinstance(config, dict):
|
||||
try:
|
||||
out.append(cls(**config))
|
||||
except TypeError as exc:
|
||||
raise UserError(f'ManagedAgentSpec: capability {type_key!r} rejected config {config!r}: {exc}') from exc
|
||||
return out
|
||||
|
||||
|
||||
def _instructions_with_skills(spec: Mapping[str, Any]) -> str | None:
|
||||
"""Append the skill catalog to the spec's instructions.
|
||||
|
||||
Skills are progressively-loaded: the model picks one by description, then
|
||||
the harness loads its full instructions on demand. The catalog has to
|
||||
land in the system prompt so the model can ASK for a skill in the first
|
||||
place. (Full-skill on-demand loading is future work -- today only the
|
||||
descriptions reach the model.)
|
||||
"""
|
||||
base = spec.get('instructions')
|
||||
base_text = base if isinstance(base, str) else None
|
||||
|
||||
blurbs: list[str] = []
|
||||
for entry in _iter_dicts(spec.get('skills')):
|
||||
name = entry.get('name')
|
||||
description = entry.get('description')
|
||||
if isinstance(name, str) and isinstance(description, str) and description:
|
||||
blurbs.append(f'- {name}: {description}')
|
||||
if not blurbs:
|
||||
return base_text
|
||||
|
||||
catalog = 'Available skills:\n' + '\n'.join(blurbs)
|
||||
return f'{base_text}\n\n{catalog}' if base_text else catalog
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Tests for the `ManagedAgentSpec` capability (source package `pydantic_ai_harness.logfire`).
|
||||
|
||||
Resolution runs against the code default (no Logfire provider) or against a
|
||||
locally-overridden variable value. The build is then exercised against
|
||||
TestModel so the assembled Agent actually runs end-to-end.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import logfire
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.exceptions import UserError
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness import ManagedAgentSpec
|
||||
from pydantic_ai_harness.logfire import ManagedAgentSpec as ManagedAgentSpecFromPackage
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope='module')
|
||||
def _configure_logfire() -> None:
|
||||
"""Configure Logfire once so variable resolution does not warn (warnings are errors)."""
|
||||
logfire.configure(send_to_logfire=False, console=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
async def test_build_falls_back_to_defaults_when_spec_unpublished() -> None:
|
||||
spec = ManagedAgentSpec('unpublished_slug', default_model=TestModel())
|
||||
|
||||
agent = await spec.build()
|
||||
|
||||
# The default-spec dict produces an Agent that runs with the default
|
||||
# model and no instructions / settings / tools / mcp servers.
|
||||
assert isinstance(agent, Agent)
|
||||
result = await agent.run('hello', model=TestModel())
|
||||
assert result.output == 'success (no tool calls)'
|
||||
|
||||
|
||||
async def test_re_exported_at_package_root() -> None:
|
||||
assert ManagedAgentSpec is ManagedAgentSpecFromPackage
|
||||
|
||||
|
||||
async def test_strips_prefix_with_warning() -> None:
|
||||
# Users sometimes copy the full `agentspec__` variable name out of the
|
||||
# UI; warn but accept rather than constructing the doubly-prefixed
|
||||
# variable `agentspec__agentspec__foo`.
|
||||
with pytest.warns(UserWarning, match='prefix is added automatically'):
|
||||
spec = ManagedAgentSpec('agentspec__leading_prefix')
|
||||
|
||||
assert spec._variable.name == 'agentspec__leading_prefix'
|
||||
|
||||
|
||||
async def test_invalid_name_rejected() -> None:
|
||||
with pytest.raises(ValueError, match='produces an invalid variable name'):
|
||||
ManagedAgentSpec('has spaces')
|
||||
|
||||
|
||||
async def test_accepts_prebuilt_variable() -> None:
|
||||
var = logfire.var(name='agentspec__prebuilt', type=dict, default={})
|
||||
|
||||
spec = ManagedAgentSpec(var, default_model=TestModel())
|
||||
agent = await spec.build()
|
||||
|
||||
assert isinstance(agent, Agent)
|
||||
|
||||
|
||||
async def test_override_value_flows_into_agent_kwargs() -> None:
|
||||
spec = ManagedAgentSpec(
|
||||
'override_slug',
|
||||
tools={'tool__weather': _fetch_weather},
|
||||
default_model=TestModel(call_tools=['_fetch_weather']),
|
||||
)
|
||||
|
||||
with spec._variable.override(
|
||||
{
|
||||
'instructions': 'Be brief.',
|
||||
'model_settings': {'temperature': 0.4},
|
||||
'tools': [{'name': 'tool__weather', 'description': 'Look up weather'}],
|
||||
}
|
||||
):
|
||||
agent = await spec.build()
|
||||
|
||||
assert agent.model_settings == {'temperature': 0.4}
|
||||
result = await agent.run('What is the weather in Paris?')
|
||||
assert '_fetch_weather' in result.output
|
||||
|
||||
|
||||
async def test_unmapped_tools_are_silently_skipped() -> None:
|
||||
# Spec lists two tools, caller registered one — only the registered
|
||||
# one reaches the Agent. Model just doesn't see the other.
|
||||
spec = ManagedAgentSpec(
|
||||
'skipped_tools_slug',
|
||||
tools={'tool__weather': _fetch_weather},
|
||||
default_model=TestModel(),
|
||||
)
|
||||
with spec._variable.override(
|
||||
{
|
||||
'tools': [
|
||||
{'name': 'tool__weather', 'description': 'weather'},
|
||||
{'name': 'tool__news', 'description': 'unmapped'},
|
||||
],
|
||||
}
|
||||
):
|
||||
agent = await spec.build()
|
||||
|
||||
# No assertion on agent.tools (private), but the agent runs without raising.
|
||||
result = await agent.run('hello', model=TestModel())
|
||||
assert isinstance(result.output, str)
|
||||
|
||||
|
||||
async def test_skills_fold_into_instructions() -> None:
|
||||
spec = ManagedAgentSpec('skills_slug', default_model=TestModel())
|
||||
|
||||
with spec._variable.override(
|
||||
{
|
||||
'instructions': 'You are a helpful assistant.',
|
||||
'skills': [
|
||||
{'name': 'skill__refund', 'description': 'Refund eligibility lookup.'},
|
||||
{'name': 'skill__shipping', 'description': 'Shipping status.'},
|
||||
],
|
||||
}
|
||||
):
|
||||
agent = await spec.build()
|
||||
|
||||
instructions = '\n'.join(i for i in (agent._instructions or []) if isinstance(i, str))
|
||||
assert 'You are a helpful assistant.' in instructions
|
||||
assert 'Available skills:' in instructions
|
||||
assert 'skill__refund: Refund eligibility lookup.' in instructions
|
||||
assert 'skill__shipping: Shipping status.' in instructions
|
||||
|
||||
|
||||
async def test_skills_without_base_instructions_emit_catalog_only() -> None:
|
||||
spec = ManagedAgentSpec('skills_only_slug', default_model=TestModel())
|
||||
|
||||
with spec._variable.override({'skills': [{'name': 'skill__refund', 'description': 'Refund lookup.'}]}):
|
||||
agent = await spec.build()
|
||||
|
||||
instructions = '\n'.join(i for i in (agent._instructions or []) if isinstance(i, str))
|
||||
assert instructions.startswith('Available skills:')
|
||||
|
||||
|
||||
async def test_capability_classes_instantiated_from_spec() -> None:
|
||||
instantiated: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
class StubCapability(AbstractCapability[Any]):
|
||||
def __init__(self, **config: Any) -> None:
|
||||
instantiated.append(('stub', config))
|
||||
|
||||
spec = ManagedAgentSpec('caps_slug', capability_classes={'stub_cap': StubCapability}, default_model=TestModel())
|
||||
|
||||
with spec._variable.override({'capabilities': [{'type': 'stub_cap', 'config': {'foo': 'bar'}}]}):
|
||||
await spec.build()
|
||||
|
||||
assert instantiated == [('stub', {'foo': 'bar'})]
|
||||
|
||||
|
||||
async def test_unregistered_capability_dropped_so_optional_deps_dont_break_callers() -> None:
|
||||
# Spec references `code_mode` but caller never installed the harness
|
||||
# extra, so they don't pass it in capability_classes. The build must
|
||||
# succeed (rather than raising) and just drop that capability.
|
||||
spec = ManagedAgentSpec('missing_cap_slug', default_model=TestModel())
|
||||
|
||||
with spec._variable.override({'capabilities': [{'type': 'code_mode'}]}):
|
||||
agent = await spec.build()
|
||||
|
||||
assert isinstance(agent, Agent)
|
||||
|
||||
|
||||
async def test_capability_with_bad_config_raises_user_error() -> None:
|
||||
class StrictCapability(AbstractCapability[Any]):
|
||||
def __init__(self, allowed_arg: int) -> None:
|
||||
self.allowed_arg = allowed_arg # pragma: no cover
|
||||
|
||||
spec = ManagedAgentSpec(
|
||||
'strict_cap_slug', capability_classes={'strict': StrictCapability}, default_model=TestModel()
|
||||
)
|
||||
|
||||
with spec._variable.override({'capabilities': [{'type': 'strict', 'config': {'unknown_arg': 1}}]}):
|
||||
with pytest.raises(UserError, match="capability 'strict' rejected config"):
|
||||
await spec.build()
|
||||
|
||||
|
||||
async def test_capability_with_non_dict_config_dropped() -> None:
|
||||
# `config` of e.g. a string is malformed; the entry is dropped rather than
|
||||
# crashing the build. Capability class is never instantiated.
|
||||
instantiated: list[str] = []
|
||||
|
||||
class NoArgCapability(AbstractCapability[Any]):
|
||||
def __init__(self) -> None:
|
||||
instantiated.append('called') # pragma: no cover
|
||||
|
||||
spec = ManagedAgentSpec(
|
||||
'bad_config_slug', capability_classes={'no_arg': NoArgCapability}, default_model=TestModel()
|
||||
)
|
||||
|
||||
with spec._variable.override({'capabilities': [{'type': 'no_arg', 'config': 'not-a-dict'}]}):
|
||||
await spec.build()
|
||||
|
||||
assert instantiated == []
|
||||
|
||||
|
||||
async def test_malformed_entries_silently_skipped() -> None:
|
||||
# Defending against a hand-edited spec value — non-dict entries in the
|
||||
# tools/mcp/skills/capabilities lists shouldn't crash the build.
|
||||
spec = ManagedAgentSpec('malformed_slug', default_model=TestModel())
|
||||
|
||||
with spec._variable.override(
|
||||
{
|
||||
'tools': ['not-a-dict', {'no_name': True}, {'name': 42}],
|
||||
'mcp_servers': [None, {'url': ''}, {'no_url': True}],
|
||||
'skills': [{'name': 'skill__x'}, 'string'], # missing description → dropped
|
||||
'capabilities': [{'type': 42}, None, {'no_type': True}],
|
||||
}
|
||||
):
|
||||
agent = await spec.build()
|
||||
|
||||
# No skills with a usable description → no catalog appended.
|
||||
assert not agent._instructions
|
||||
|
||||
|
||||
async def test_logfire_instance_ignored_with_prebuilt_variable_warns() -> None:
|
||||
# The Variable carries its own Logfire instance, so passing one explicitly
|
||||
# is ambiguous; warn rather than silently use the wrong instance.
|
||||
var = logfire.var(name='agentspec__instance_warn', type=dict, default={})
|
||||
|
||||
with pytest.warns(UserWarning, match='`logfire_instance` is ignored'):
|
||||
ManagedAgentSpec(var, logfire_instance=logfire.DEFAULT_LOGFIRE_INSTANCE)
|
||||
|
||||
|
||||
async def test_capability_without_config_instantiates_with_no_args() -> None:
|
||||
instantiated: list[str] = []
|
||||
|
||||
class NoArgCapability(AbstractCapability[Any]):
|
||||
def __init__(self) -> None:
|
||||
instantiated.append('no-arg')
|
||||
|
||||
spec = ManagedAgentSpec('noarg_cap_slug', capability_classes={'no_arg': NoArgCapability}, default_model=TestModel())
|
||||
|
||||
# Entry omits `config` entirely; we should fall through to the `cls()` branch.
|
||||
with spec._variable.override({'capabilities': [{'type': 'no_arg'}]}):
|
||||
await spec.build()
|
||||
|
||||
assert instantiated == ['no-arg']
|
||||
|
||||
|
||||
async def test_capability_with_valid_config_instantiates() -> None:
|
||||
# Counter-test to the bad-config one above: when the config matches the
|
||||
# capability's __init__ signature, the instance reaches the Agent.
|
||||
instantiated: list[int] = []
|
||||
|
||||
class StrictCapability(AbstractCapability[Any]):
|
||||
def __init__(self, allowed_arg: int) -> None:
|
||||
instantiated.append(allowed_arg)
|
||||
|
||||
spec = ManagedAgentSpec(
|
||||
'valid_cap_slug', capability_classes={'strict': StrictCapability}, default_model=TestModel()
|
||||
)
|
||||
|
||||
with spec._variable.override({'capabilities': [{'type': 'strict', 'config': {'allowed_arg': 42}}]}):
|
||||
await spec.build()
|
||||
|
||||
assert instantiated == [42]
|
||||
|
||||
|
||||
async def test_valid_mcp_server_reaches_toolset_kwarg(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# The harness's default install lacks the `[mcp]` extra, so `pydantic_ai.mcp`
|
||||
# would fail to import. Mock just the symbol we need so the wirer's
|
||||
# construction path runs end-to-end and the resulting toolset lands on the
|
||||
# Agent kwargs.
|
||||
fake_module = MagicMock()
|
||||
fake_toolset = MagicMock()
|
||||
fake_module.MCPToolset = fake_toolset
|
||||
monkeypatch.setitem(__import__('sys').modules, 'pydantic_ai.mcp', fake_module)
|
||||
|
||||
spec = ManagedAgentSpec('mcp_slug', default_model=TestModel())
|
||||
with spec._variable.override(
|
||||
{
|
||||
'mcp_servers': [
|
||||
{'url': 'https://example.com/mcp', 'headers': {'Authorization': 'Bearer x'}},
|
||||
{'url': 'https://other.com/mcp'}, # no headers → None
|
||||
],
|
||||
}
|
||||
):
|
||||
await spec.build()
|
||||
|
||||
# Both servers reached MCPToolset(); headers preserved per-entry.
|
||||
assert fake_toolset.call_args_list[0].args == ('https://example.com/mcp',)
|
||||
assert fake_toolset.call_args_list[0].kwargs == {'headers': {'Authorization': 'Bearer x'}}
|
||||
assert fake_toolset.call_args_list[1].args == ('https://other.com/mcp',)
|
||||
assert fake_toolset.call_args_list[1].kwargs == {'headers': None}
|
||||
|
||||
|
||||
def _fetch_weather(city: str) -> str:
|
||||
return f'sunny and 72°F in {city}'
|
||||
Reference in New Issue
Block a user