From 0468d214625562fa1c069a902d5c307fb6672582 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Mon, 15 Jun 2026 10:55:59 +0200 Subject: [PATCH] Wired agent skills into request flow --- pyproject.toml | 3 +- src/configuration.py | 8 +++ src/utils/agents/query.py | 2 +- src/utils/pydantic_ai.py | 45 ++++++++++++++- tests/unit/conftest.py | 43 ++++++++++++++ tests/unit/utils/agents/test_query.py | 12 ++++ tests/unit/utils/test_pydantic_ai.py | 82 ++++++++++++++++++++++++++- uv.lock | 32 ++++++++--- 8 files changed, 212 insertions(+), 15 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index af2cbc46c..3faf37058 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,7 +80,8 @@ dependencies = [ # Used for token estimation before LLM calls (LCORE-1569 / conversation compaction) "tiktoken>=0.8.0", # Used for Pydantic AI - "pydantic-ai>=1.99.0" + "pydantic-ai>=1.99.0", + "pydantic-ai-skills>=0.11.0", ] diff --git a/src/configuration.py b/src/configuration.py index 0de369892..2ede87d27 100644 --- a/src/configuration.py +++ b/src/configuration.py @@ -32,6 +32,7 @@ RerankerConfiguration, RlsapiV1Configuration, ServiceConfiguration, + SkillsConfiguration, SplunkConfiguration, UserDataCollection, ) @@ -504,6 +505,13 @@ def reranker(self) -> "RerankerConfiguration": raise LogicError("logic error: configuration is not loaded") return self._configuration.reranker + @property + def skills(self) -> Optional[SkillsConfiguration]: + """Return agent skills configuration, or None if not provided.""" + if self._configuration is None: + raise LogicError("logic error: configuration is not loaded") + return self._configuration.skills + @property def rag_id_mapping(self) -> dict[str, str]: """Return mapping from vector_db_id to rag_id from BYOK and OKP RAG config. diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index 0322c43ad..cb84cf7ac 100644 --- a/src/utils/agents/query.py +++ b/src/utils/agents/query.py @@ -310,7 +310,7 @@ async def retrieve_agent_response( llm_response=moderation_result.message, ) try: - agent = build_agent(client, responses_params) + agent = build_agent(client, responses_params, configuration.skills) logger.debug("Starting agent non-streaming response processing") run_result = await agent.run(cast(str, responses_params.input)) except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as exc: diff --git a/src/utils/pydantic_ai.py b/src/utils/pydantic_ai.py index 5df570dc9..c655e67a5 100644 --- a/src/utils/pydantic_ai.py +++ b/src/utils/pydantic_ai.py @@ -2,14 +2,16 @@ from __future__ import annotations -from typing import Any, Final, cast +from typing import Any, Final, Optional, cast from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient from llama_stack_client import AsyncLlamaStackClient -from pydantic_ai import Agent +from pydantic_ai import Agent, AgentCapability from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIResponsesModelSettings +from pydantic_ai_skills import SkillsCapability from models.common.responses.responses_api_params import ResponsesApiParams +from models.config import SkillsConfiguration from pydantic_ai_lightspeed.llamastack import LlamaStackProvider _LLS_RESPONSES_EXTRA_FIELDS: Final[frozenset[str]] = frozenset( @@ -70,9 +72,46 @@ def _model_settings_from_responses_params( return cast(OpenAIResponsesModelSettings, settings_dict) +def _skills_capability( + skills_config: Optional[SkillsConfiguration], +) -> Optional[SkillsCapability]: + """Return a skills capability when skill paths are configured. + + Args: + skills_config: Agent skills configuration from LCS, or None when skills are disabled. + + Returns: + SkillsCapability when skill paths are configured, or None when skills are disabled. + """ + if skills_config is None or not skills_config.paths: + return None + return SkillsCapability( + directories=[str(path) for path in skills_config.paths], + validate=False, + ) + + +def _agent_capabilities( + skills: Optional[SkillsConfiguration], +) -> Optional[list[AgentCapability[None]]]: + """Assemble pydantic-ai capabilities for an LCS agent. + + Args: + skills: Agent skills configuration from LCS, or None when skills are disabled. + + Returns: + Configured capabilities, or None when no capabilities are enabled. + """ + capabilities: list[AgentCapability[None]] = [] + if skills_capability := _skills_capability(skills): + capabilities.append(skills_capability) + return capabilities or None + + def build_agent( client: AsyncLlamaStackClient | AsyncLlamaStackAsLibraryClient, responses_params: ResponsesApiParams, + skills: Optional[SkillsConfiguration], ) -> Agent[None, str]: """Build a Pydantic AI agent that mirrors ``responses_params`` on the Llama Stack backend. @@ -84,6 +123,7 @@ def build_agent( Parameters: client: Initialized Llama Stack client from ``AsyncLlamaStackClientHolder().get_client()``. responses_params: Parameters produced by ``prepare_responses_params`` for this turn. + skills: Agent skills configuration from LCS, or None when skills are disabled. Returns: ``Agent`` configured for ``await agent.run(...)`` (or streaming) against the same @@ -100,5 +140,6 @@ def build_agent( return Agent( model, instructions=responses_params.instructions, + capabilities=_agent_capabilities(skills), defer_model_check=True, ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index cf741b9e8..374db348d 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -3,11 +3,16 @@ from __future__ import annotations from collections.abc import Generator +from pathlib import Path +import httpx import pytest +from llama_stack_client import AsyncLlamaStackClient from pytest_mock import AsyncMockType, MockerFixture from configuration import AppConfig +from models.common.responses.responses_api_params import ResponsesApiParams +from models.config import SkillsConfiguration type AgentFixtures = Generator[ tuple[ @@ -72,3 +77,41 @@ def minimal_config_fixture() -> AppConfig: } ) return cfg + + +@pytest.fixture(name="mock_client") +def mock_client_fixture( # pylint: disable=protected-access + mocker: MockerFixture, +) -> AsyncLlamaStackClient: + """Remote Llama Stack client mock for build_agent tests.""" + client = mocker.Mock(spec=AsyncLlamaStackClient) + client.base_url = "http://localhost:8321" + client.api_key = "test-key" + client._client = mocker.Mock(spec=httpx.AsyncClient) + return client + + +@pytest.fixture(name="mock_params") +def mock_params_fixture() -> ResponsesApiParams: + """Minimal ResponsesApiParams for build_agent and similar utils tests.""" + return ResponsesApiParams( + model="provider/my-model", + input="test", + conversation="conv-test", + instructions="Be helpful.", + store=False, + stream=False, + ) + + +@pytest.fixture(name="mock_skills_configuration") +def mock_skills_configuration_fixture(tmp_path: Path) -> SkillsConfiguration: + """Filesystem-backed SkillsConfiguration with a single test skill.""" + skills_root = tmp_path / "skills" + skill_dir = skills_root / "test-skill" + skill_dir.mkdir(parents=True) + (skill_dir / "SKILL.md").write_text( + "---\nname: test-skill\ndescription: Test skill.\n---\n\nDo the thing.\n", + encoding="utf-8", + ) + return SkillsConfiguration(paths=[skills_root]) diff --git a/tests/unit/utils/agents/test_query.py b/tests/unit/utils/agents/test_query.py index 6baa1380f..37bcca73a 100644 --- a/tests/unit/utils/agents/test_query.py +++ b/tests/unit/utils/agents/test_query.py @@ -116,10 +116,20 @@ def blocked_moderation_fixture() -> ShieldModerationBlocked: ) +@pytest.fixture(name="patch_query_configuration") +def patch_query_configuration_fixture(mocker: MockerFixture) -> None: + """Patch query module configuration for isolated agent query tests.""" + mock_config = mocker.MagicMock() + mock_config.skills = None + mock_config.rag_id_mapping = {} + mocker.patch("utils.agents.query.configuration", mock_config) + + @pytest.fixture(name="patch_recording_metrics") def patch_recording_metrics_fixture(mocker: MockerFixture) -> None: """Patch LLM recording helpers so token usage tests stay isolated.""" mock_config = mocker.MagicMock() + mock_config.skills = None mock_config.rag_id_mapping = {} mocker.patch("utils.agents.query.configuration", mock_config) mocker.patch( @@ -422,6 +432,7 @@ async def test_success_returns_turn_summary( assert summary.id == "resp-success" @pytest.mark.asyncio + @pytest.mark.usefixtures("patch_query_configuration") async def test_agent_connection_error_raises_http_exception( self, mocker: MockerFixture, @@ -448,6 +459,7 @@ async def test_agent_connection_error_raises_http_exception( assert exc_info.value.status_code == 503 @pytest.mark.asyncio + @pytest.mark.usefixtures("patch_query_configuration") async def test_api_status_error_raises_http_exception( self, mocker: MockerFixture, diff --git a/tests/unit/utils/test_pydantic_ai.py b/tests/unit/utils/test_pydantic_ai.py index c7acc37e3..2a9a59d52 100644 --- a/tests/unit/utils/test_pydantic_ai.py +++ b/tests/unit/utils/test_pydantic_ai.py @@ -5,12 +5,18 @@ import httpx import pytest from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient +from llama_stack_client import AsyncLlamaStackClient +from pydantic_ai_skills import SkillsCapability from pytest_mock import MockerFixture +from models.common.responses.responses_api_params import ResponsesApiParams +from models.config import SkillsConfiguration from utils.pydantic_ai import ( _LLS_RESPONSES_EXTRA_FIELDS, + _agent_capabilities, _llama_stack_provider_from_client, _model_settings_from_responses_params, + _skills_capability, build_agent, ) @@ -196,6 +202,45 @@ def test_contains_expected_fields(self) -> None: assert expected == _LLS_RESPONSES_EXTRA_FIELDS +class TestSkillsCapability: + """Tests for _skills_capability.""" + + def test_returns_none_when_skills_not_configured(self) -> None: + """Test that missing skills configuration returns None.""" + assert _skills_capability(None) is None + + def test_returns_none_when_paths_empty(self) -> None: + """Test that an empty paths list returns None.""" + assert _skills_capability(SkillsConfiguration(paths=[])) is None + + def test_returns_capability_for_configured_paths( + self, mock_skills_configuration: SkillsConfiguration + ) -> None: + """Test that configured paths produce a SkillsCapability.""" + capability = _skills_capability(mock_skills_configuration) + + assert isinstance(capability, SkillsCapability) + assert list(capability.toolset.skills) == ["test-skill"] + + +class TestAgentCapabilities: + """Tests for _agent_capabilities.""" + + def test_returns_none_when_no_capabilities_configured(self) -> None: + """Test that missing configuration yields None for Agent construction.""" + assert _agent_capabilities(None) is None + assert _agent_capabilities(SkillsConfiguration(paths=[])) is None + + def test_returns_skills_capability_when_configured( + self, mock_skills_configuration: SkillsConfiguration + ) -> None: + """Test that configured skills are included in the capability list.""" + capabilities = _agent_capabilities(mock_skills_configuration) or [] + + assert len(capabilities) == 1 + assert isinstance(capabilities[0], SkillsCapability) + + class TestBuildAgent: """Tests for the build_agent factory function.""" @@ -220,7 +265,7 @@ def test_returns_agent_with_correct_model(self, mocker: MockerFixture) -> None: mock_params.store = False mock_params.previous_response_id = None - agent = build_agent(mock_client, mock_params) + agent = build_agent(mock_client, mock_params, None) assert agent is not None @@ -242,7 +287,7 @@ def test_agent_has_instructions(self, mocker: MockerFixture) -> None: mock_params.store = False mock_params.previous_response_id = None - agent = build_agent(mock_client, mock_params) + agent = build_agent(mock_client, mock_params, None) assert "You are a helpful assistant." in agent._instructions @@ -265,6 +310,37 @@ def test_agent_with_library_client(self, mocker: MockerFixture) -> None: mock_params.store = True mock_params.previous_response_id = None - agent = build_agent(mock_lib_client, mock_params) + agent = build_agent(mock_lib_client, mock_params, None) assert agent is not None + + def test_agent_includes_skills_capability_when_configured( + self, + mock_client: AsyncLlamaStackClient, + mock_params: ResponsesApiParams, + mock_skills_configuration: SkillsConfiguration, + ) -> None: + """Test that build_agent attaches SkillsCapability when skills are passed.""" + agent = build_agent( + mock_client, + mock_params, + mock_skills_configuration, + ) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert SkillsCapability in capability_types + + def test_agent_has_no_skills_capability_when_not_configured( + self, + mock_client: AsyncLlamaStackClient, + mock_params: ResponsesApiParams, + ) -> None: + """Test that build_agent omits SkillsCapability when skills are not passed.""" + agent = build_agent(mock_client, mock_params, None) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert SkillsCapability not in capability_types diff --git a/uv.lock b/uv.lock index 57ddf06dc..6b559ac8d 100644 --- a/uv.lock +++ b/uv.lock @@ -620,14 +620,14 @@ name = "cohere" version = "7.0.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "fastavro" }, - { name = "httpx" }, - { name = "pydantic" }, - { name = "pydantic-core" }, - { name = "requests" }, - { name = "tokenizers" }, - { name = "types-requests" }, - { name = "typing-extensions" }, + { name = "fastavro", marker = "sys_platform != 'emscripten'" }, + { name = "httpx", marker = "sys_platform != 'emscripten'" }, + { name = "pydantic", marker = "sys_platform != 'emscripten'" }, + { name = "pydantic-core", marker = "sys_platform != 'emscripten'" }, + { name = "requests", marker = "sys_platform != 'emscripten'" }, + { name = "tokenizers", marker = "sys_platform != 'emscripten'" }, + { name = "types-requests", marker = "sys_platform != 'emscripten'" }, + { name = "typing-extensions", marker = "sys_platform != 'emscripten'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cf/3c/670631ee223d7b64d157dc3f309bf93bde65efe0bb1a8341d9b575f407d3/cohere-7.0.4.tar.gz", hash = "sha256:35b6a397d35ae6eafa1a02921f42c2a98309a990874533e5238efaf3426b6a21", size = 208794, upload-time = "2026-06-11T15:17:52.994Z" } wheels = [ @@ -1996,6 +1996,7 @@ dependencies = [ { name = "psycopg2-binary" }, { name = "pyasn1" }, { name = "pydantic-ai" }, + { name = "pydantic-ai-skills" }, { name = "python-dotenv" }, { name = "pyyaml" }, { name = "requests" }, @@ -2103,6 +2104,7 @@ requires-dist = [ { name = "psycopg2-binary", specifier = ">=2.9.10" }, { name = "pyasn1", specifier = ">=0.6.3" }, { name = "pydantic-ai", specifier = ">=1.99.0" }, + { name = "pydantic-ai-skills", specifier = ">=0.11.0" }, { name = "python-dotenv", specifier = ">=1.2.2" }, { name = "pyyaml", specifier = ">=6.0.0" }, { name = "requests", specifier = ">=2.33.0" }, @@ -3478,6 +3480,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/38/db37ab59fced191a75babbbbe99470e512c91fe730b27d113e3280fdbe44/pydantic_ai-1.107.0-py3-none-any.whl", hash = "sha256:e031880b44ad7ce3836b2f6aa8ce2a0bd733cdb0b89a34adba647e96ddcba788", size = 7588, upload-time = "2026-06-10T14:53:00.57Z" }, ] +[[package]] +name = "pydantic-ai-skills" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "pydantic-ai-slim" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/94/d1/f8fbc1c792ba8d73fc424ddab143ab0e1d95f9e982be5a949aaa3c84d64e/pydantic_ai_skills-0.11.0.tar.gz", hash = "sha256:d4040f0b81da34e25b8f14dac5e1895e59d00390db8896448aac1ed1a4d0cf90", size = 9023711, upload-time = "2026-05-26T01:52:55.375Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/80/d9b4bf0f5a3e8d256d62aa30b1c4589dc492efdb2a2d0da0334b4e71f25c/pydantic_ai_skills-0.11.0-py3-none-any.whl", hash = "sha256:af8d78d451ce192dd2ef33abe86ad900bd51d9fd10c81a11abee82f62e8daf30", size = 61218, upload-time = "2026-05-26T01:52:53.773Z" }, +] + [[package]] name = "pydantic-ai-slim" version = "1.107.0"