Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]


Expand Down
8 changes: 8 additions & 0 deletions src/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
RerankerConfiguration,
RlsapiV1Configuration,
ServiceConfiguration,
SkillsConfiguration,
SplunkConfiguration,
UserDataCollection,
)
Expand Down Expand Up @@ -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
Comment thread
asimurka marked this conversation as resolved.

@property
def rag_id_mapping(self) -> dict[str, str]:
"""Return mapping from vector_db_id to rag_id from BYOK and OKP RAG config.
Expand Down
2 changes: 1 addition & 1 deletion src/utils/agents/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
45 changes: 43 additions & 2 deletions src/utils/pydantic_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -70,9 +72,46 @@ def _model_settings_from_responses_params(
return cast(OpenAIResponsesModelSettings, settings_dict)


def _skills_capability(
skills_config: Optional[SkillsConfiguration],
) -> Optional[SkillsCapability]:
Comment thread
asimurka marked this conversation as resolved.
"""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.

Expand All @@ -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
Expand All @@ -100,5 +140,6 @@ def build_agent(
return Agent(
model,
instructions=responses_params.instructions,
capabilities=_agent_capabilities(skills),
defer_model_check=True,
)
43 changes: 43 additions & 0 deletions tests/unit/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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[
Expand Down Expand Up @@ -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])
12 changes: 12 additions & 0 deletions tests/unit/utils/agents/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
82 changes: 79 additions & 3 deletions tests/unit/utils/test_pydantic_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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."""

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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
32 changes: 24 additions & 8 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading