diff --git a/pyproject.toml b/pyproject.toml index 5f324b18c..b9e4ad2ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,9 +28,9 @@ dependencies = [ # Used by authentication/k8s integration "kubernetes>=30.1.0", # Used to call Llama Stack APIs - "ogx==1.2.2", - "ogx-client==1.2.2", - "ogx-api==1.2.2", + "ogx==1.3.0", + "ogx-client==1.3.0", + "ogx-api==1.3.0", # Used by Logger "rich>=14.0.0", # Used by JWK token auth handler diff --git a/src/app/endpoints/models.py b/src/app/endpoints/models.py index 37491f700..6a9dd6262 100644 --- a/src/app/endpoints/models.py +++ b/src/app/endpoints/models.py @@ -93,7 +93,7 @@ async def models_endpoint_handler( # try to get Llama Stack client client = AsyncOgxClientHolder().get_client() # retrieve and normalize models across OpenAI/Anthropic/Google list shapes - parsed_models = parse_model_list_response(await client.models.list()) + parsed_models = parse_model_list_response(await client.openai.list()) # optional filtering by model type if model_type.model_type is not None: diff --git a/src/app/endpoints/rlsapi_v1.py b/src/app/endpoints/rlsapi_v1.py index 8e9a3d446..39bdfcfb4 100644 --- a/src/app/endpoints/rlsapi_v1.py +++ b/src/app/endpoints/rlsapi_v1.py @@ -188,7 +188,7 @@ async def _get_default_model_id() -> str: ) client = AsyncOgxClientHolder().get_client() try: - models = parse_model_list_response(await client.models.list()) + models = parse_model_list_response(await client.openai.list()) except APIConnectionError as e: error_response = ServiceUnavailableResponse( backend_name="OGX", diff --git a/src/app/endpoints/vector_stores.py b/src/app/endpoints/vector_stores.py index e68244c3e..b50a25ada 100644 --- a/src/app/endpoints/vector_stores.py +++ b/src/app/endpoints/vector_stores.py @@ -184,7 +184,7 @@ async def create_vector_store( return VectorStoreResponse( id=vector_store.id, - name=vector_store.name, + name=vector_store.name or "", created_at=vector_store.created_at, last_active_at=vector_store.last_active_at, expires_at=vector_store.expires_at, @@ -236,7 +236,7 @@ async def list_vector_stores( data = [ VectorStoreResponse( id=vs.id, - name=vs.name, + name=vs.name or "", created_at=vs.created_at, last_active_at=vs.last_active_at, expires_at=vs.expires_at or None, @@ -294,7 +294,7 @@ async def get_vector_store( return VectorStoreResponse( id=vector_store.id, - name=vector_store.name, + name=vector_store.name or "", created_at=vector_store.created_at, last_active_at=vector_store.last_active_at, expires_at=vector_store.expires_at, @@ -358,7 +358,7 @@ async def update_vector_store( return VectorStoreResponse( id=vector_store.id, - name=vector_store.name, + name=vector_store.name or "", created_at=vector_store.created_at, last_active_at=vector_store.last_active_at, expires_at=vector_store.expires_at, diff --git a/src/client.py b/src/client.py index 7fd4a1e5c..1e6be32bb 100644 --- a/src/client.py +++ b/src/client.py @@ -3,7 +3,7 @@ import json import os import tempfile -from typing import Optional, cast +from typing import Any, Optional import yaml from fastapi import HTTPException @@ -29,6 +29,31 @@ logger = get_logger(__name__) +def read_provider_data(client: AsyncOgxClient) -> dict[str, Any]: + """Read provider data from a library or service client. + + Library clients keep provider data on ``provider_data``. Service + clients store it as JSON in ``api_client.default_headers``. + + Parameters: + client: Initialized OGX client (library or service). + + Returns: + A mutable copy of the current provider data dict (empty if unset). + """ + if isinstance(client, AsyncOGXAsLibraryClient): + return dict(client.provider_data or {}) + + raw = client.api_client.default_headers.get("X-OGX-Provider-Data") + if not raw: + return {} + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, TypeError): + return {} + return parsed if isinstance(parsed, dict) else {} + + class AsyncOgxClientHolder(metaclass=Singleton): """Container for an initialised AsyncOgxClient.""" @@ -235,7 +260,7 @@ async def check_model_available(self, model_id: str) -> tuple[bool, str]: """ try: client = self.get_client() - models = parse_model_list_response(await client.models.list()) + models = parse_model_list_response(await client.openai.list()) except RuntimeError as e: logger.warning("Client not initialized, skipping model check: %s", e) return False, f"Client not initialized: {e!s}" @@ -257,7 +282,7 @@ async def check_model_available(self, model_id: str) -> tuple[bool, str]: try: await self.reload_library_client() client = self.get_client() - reloaded_models = parse_model_list_response(await client.models.list()) + reloaded_models = parse_model_list_response(await client.openai.list()) if any(m.identifier == model_id for m in reloaded_models): logger.info( "Model %s found after client reload", @@ -290,46 +315,32 @@ async def update_azure_token(self) -> AsyncOgxClient: if not updates: return self.get_client() + current_client = self.get_client() + provider_data = read_provider_data(current_client) + provider_data.update(updates) + if self.is_library_client: if not self._config_path: logger.warning("Cannot update Azure token: config path not set") - return self.get_client() + return current_client - current_provider_data = dict( - cast(AsyncOGXAsLibraryClient, self._lsc).provider_data or {} - ) - current_provider_data.update(updates) - client = AsyncOGXAsLibraryClient( - self._config_path, provider_data=current_provider_data + updated_client = AsyncOGXAsLibraryClient( + self._config_path, provider_data=provider_data ) - await client.initialize() - self._lsc = client + await updated_client.initialize() + self._lsc = updated_client # Re-apply logging configuration after ogx's setup_logging() is called. # This ensures the desired logging configuration is applied when # using AsyncOGXAsLibraryClient. setup_logging() - - return client - - # Service client mode - current_client = self.get_client() - current_headers = current_client.default_headers or {} - provider_data_json = current_headers.get("X-OGX-Provider-Data") - - try: - provider_data = json.loads(provider_data_json) if provider_data_json else {} - except (json.JSONDecodeError, TypeError): - provider_data = {} - - provider_data.update(updates) - - updated_headers = { - **current_headers, - "X-OGX-Provider-Data": json.dumps(provider_data), - } - - updated_client = current_client.copy( - set_default_headers=updated_headers # type: ignore[arg-type] + return updated_client + + # Service client: AsyncOgxClient has no .copy(); rebuild with provider_data. + updated_client = AsyncOgxClient( + base_url=str(current_client.base_url) if current_client.base_url else None, + api_key=current_client.api_key, + timeout=current_client.configuration.timeout, + provider_data=provider_data, ) self._lsc = updated_client return updated_client diff --git a/src/metrics/utils.py b/src/metrics/utils.py index 62b5fd97e..6e5954419 100644 --- a/src/metrics/utils.py +++ b/src/metrics/utils.py @@ -19,7 +19,7 @@ async def setup_model_metrics() -> None: logger.info("Setting up model metrics") check_configuration_loaded(configuration) model_list = parse_model_list_response( - await AsyncOgxClientHolder().get_client().models.list() + await AsyncOgxClientHolder().get_client().openai.list() ) models = [model for model in model_list if model.model_type == "llm"] diff --git a/src/models/api/responses/successful/vector_stores.py b/src/models/api/responses/successful/vector_stores.py index 08c21ae0d..74158c6af 100644 --- a/src/models/api/responses/successful/vector_stores.py +++ b/src/models/api/responses/successful/vector_stores.py @@ -1,5 +1,6 @@ """Successful responses for vector stores and vector store files.""" +from collections.abc import Mapping from typing import Any, ClassVar, Optional from pydantic import Field @@ -231,7 +232,7 @@ class VectorStoreFileResponse(AbstractSuccessfulResponse): id: str = Field(..., description="Vector store file ID") vector_store_id: str = Field(..., description="ID of the vector store") status: str = Field(..., description="File processing status") - attributes: Optional[dict[str, str | float | bool]] = Field( + attributes: Optional[Mapping[str, Any]] = Field( None, description=( "Set of up to 16 key-value pairs for storing additional information. " diff --git a/src/pydantic_ai_lightspeed/llamastack/_provider.py b/src/pydantic_ai_lightspeed/llamastack/_provider.py index 3ead58d6a..edaa73c96 100644 --- a/src/pydantic_ai_lightspeed/llamastack/_provider.py +++ b/src/pydantic_ai_lightspeed/llamastack/_provider.py @@ -6,7 +6,6 @@ import httpx from ogx.core.library_client import AsyncOGXAsLibraryClient -from ogx.core.request_headers import parse_request_provider_data from ogx_client import AsyncOgxClient from openai import AsyncOpenAI from pydantic_ai import ModelProfile @@ -14,6 +13,7 @@ from pydantic_ai.profiles.openai import openai_model_profile from pydantic_ai.providers import Provider +from client import read_provider_data from pydantic_ai_lightspeed.llamastack._transport import ( OgxLibraryTransport, wrap_http_client_with_provider_data, @@ -77,14 +77,8 @@ def from_ogx_client( api_key = client.api_key or "not-needed" base = str(client.base_url).rstrip("/") base_url = base if base.endswith("/v1") else f"{base}/v1" - raw_headers = client.default_headers - default_headers = { - str(key): str(value) - for key, value in raw_headers.items() - if isinstance(value, str) - } - provider_data = parse_request_provider_data(default_headers) - http_client = client._client # pylint: disable=protected-access + provider_data = read_provider_data(client) + http_client = client.api_client.async_client http_client = wrap_http_client_with_provider_data(http_client, provider_data) return OgxProvider( base_url=base_url, diff --git a/src/utils/responses.py b/src/utils/responses.py index 80b68c398..c6c95d88f 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -1329,7 +1329,7 @@ async def check_model_configured( HTTPException: If there's a connection error or other API error """ try: - models = parse_model_list_response(await client.models.list()) + models = parse_model_list_response(await client.openai.list()) for model in models: if model.identifier == model_id: return True @@ -1395,7 +1395,7 @@ async def select_model_for_responses( # 3. Fetch models list and select the first LLM model (model_type="llm") try: - models = parse_model_list_response(await client.models.list()) + models = parse_model_list_response(await client.openai.list()) except APIConnectionError as e: error_response = ServiceUnavailableResponse( backend_name="OGX", diff --git a/src/utils/types.py b/src/utils/types.py index 0af0296bc..97d3ac18e 100644 --- a/src/utils/types.py +++ b/src/utils/types.py @@ -1,14 +1,16 @@ """Common types for the project.""" from re import Pattern -from typing import Any +from typing import Any, TypeVar, cast from ogx_api import ImageContentItem, TextContentItem -type SingletonInstances = dict[type, Any] +type SingletonInstances = dict[type, object] CompiledPatterns = list[tuple[Pattern[str], str]] +T = TypeVar("T") + def content_to_str(content: Any) -> str: """Convert content (str, TextContentItem, ImageContentItem, or list) to string. @@ -43,13 +45,14 @@ class Singleton(type): _instances: SingletonInstances = {} - def __call__(cls, *args: Any, **kwargs: Any) -> Any: + def __call__(cls: type[T], *args: object, **kwargs: object) -> T: """ - Return the single cached instance of the class, creating and caching it on first call. + Return the cached singleton instance, creating it if necessary. Returns: - object: The singleton instance for this class. + The singleton instance for this class. """ - if cls not in cls._instances: - cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs) - return cls._instances[cls] + if cls not in Singleton._instances: + Singleton._instances[cls] = type.__call__(cls, *args, **kwargs) + + return cast(T, Singleton._instances[cls]) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index e683bced2..9739a6fdf 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -71,7 +71,7 @@ def make_openai_models_list_response( *models: OpenAIModel, ) -> ListModelsV1ModelsGet200Response: - """Build a ``client.models.list()`` response in the OpenAI OneOf shape. + """Build a ``client.openai.list()`` response in the OpenAI OneOf shape. Parameters: *models: OpenAI-style model entries for ``data``. @@ -90,7 +90,7 @@ def make_openai_model( provider_id: str = TEST_PROVIDER, model_type: str = "llm", ) -> OpenAIModel: - """Build an ``OpenAIModel`` for integration ``models.list`` mocks. + """Build an ``OpenAIModel`` for integration ``openai.list`` mocks. Parameters: model_id: Full model identifier (provider/name). @@ -815,8 +815,8 @@ def mock_ogx_client_fixture( mock_client.responses.create.return_value = mock_response - # Mock models.list - mock_client.models.list.return_value = make_openai_models_list_response( + # Mock openai.list + mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model() ) diff --git a/tests/integration/endpoints/test_model_list.py b/tests/integration/endpoints/test_model_list.py index 362100453..ef1f95a51 100644 --- a/tests/integration/endpoints/test_model_list.py +++ b/tests/integration/endpoints/test_model_list.py @@ -42,7 +42,7 @@ def mock_ogx_client_fixture( mock_client = mocker.AsyncMock() # Mock models list (required for model selection) - mock_client.models.list.return_value = make_openai_models_list_response( + mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model(model_id="test-provider/test-model-1"), make_openai_model( model_id="test-provider/test-model-2", model_type="embedding" @@ -78,7 +78,7 @@ def mock_ogx_client_failing_fixture( mock_client = mocker.AsyncMock() - mock_client.models.list.side_effect = APIConnectionError(request=mocker.Mock()) + mock_client.openai.list.side_effect = APIConnectionError(request=mocker.Mock()) # Create a mock holder instance mock_holder_instance = mock_holder_class.return_value diff --git a/tests/integration/endpoints/test_query_byok_integration.py b/tests/integration/endpoints/test_query_byok_integration.py index 79ad54945..033907c24 100644 --- a/tests/integration/endpoints/test_query_byok_integration.py +++ b/tests/integration/endpoints/test_query_byok_integration.py @@ -100,7 +100,7 @@ def _build_base_mock_client(mocker: MockerFixture) -> Any: mock_client = mocker.AsyncMock() # Model list - mock_client.models.list.return_value = make_openai_models_list_response( + mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model() ) diff --git a/tests/integration/endpoints/test_responses_integration.py b/tests/integration/endpoints/test_responses_integration.py index bfc58807a..652f8ed22 100644 --- a/tests/integration/endpoints/test_responses_integration.py +++ b/tests/integration/endpoints/test_responses_integration.py @@ -72,7 +72,7 @@ def _build_mock_client(mocker: MockerFixture) -> Any: """Build a mock Llama Stack client for responses integration tests. Returns a fully-configured AsyncMock client with sensible defaults for - responses.create, models.list, shields.list, vector_stores.list, and + responses.create, openai.list, shields.list, vector_stores.list, and conversations.create. """ mock_client = mocker.AsyncMock() @@ -93,7 +93,7 @@ def _build_mock_client(mocker: MockerFixture) -> Any: mock_response.model_dump.return_value = _RESPONSE_DUMP.copy() mock_client.responses.create = mocker.AsyncMock(return_value=mock_response) - mock_client.models.list.return_value = make_openai_models_list_response( + mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model() ) diff --git a/tests/integration/endpoints/test_streaming_query_integration.py b/tests/integration/endpoints/test_streaming_query_integration.py index ceeda3cc5..2ead8da39 100644 --- a/tests/integration/endpoints/test_streaming_query_integration.py +++ b/tests/integration/endpoints/test_streaming_query_integration.py @@ -37,7 +37,7 @@ def mock_llama_stack_streaming_fixture( ) mock_client = mocker.AsyncMock() - mock_client.models.list.return_value = make_openai_models_list_response( + mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model() ) diff --git a/tests/unit/app/endpoints/test_a2a.py b/tests/unit/app/endpoints/test_a2a.py index b226b64b8..b7bf30bb6 100644 --- a/tests/unit/app/endpoints/test_a2a.py +++ b/tests/unit/app/endpoints/test_a2a.py @@ -23,7 +23,6 @@ from a2a.utils import new_agent_text_message from fastapi import HTTPException, Request from ogx_client import APIConnectionError -from ogx_client.models.list_models_response import ListModelsResponse from pydantic_ai import AgentRunResultEvent from pydantic_ai.exceptions import AgentRunError from pydantic_ai.messages import ( @@ -38,6 +37,8 @@ from pydantic_ai.messages import TextPart as PydanticTextPart from pytest_mock import MockerFixture +from tests.unit.conftest import make_openai_model, make_openai_models_list_response + from app.endpoints.a2a import ( A2AAgentExecutor, TaskResultAggregator, @@ -686,7 +687,7 @@ async def test_process_task_streaming_handles_api_connection_error_on_models_lis mocker: MockerFixture, setup_configuration: AppConfig, # pylint: disable=unused-argument ) -> None: - """Test _process_task_streaming handles APIConnectionError from models.list().""" + """Test _process_task_streaming handles APIConnectionError from openai.list().""" executor = A2AAgentExecutor(auth_token="test-token") # Mock the context with valid input @@ -716,11 +717,11 @@ async def test_process_task_streaming_handles_api_connection_error_on_models_lis "app.endpoints.a2a._get_context_store", return_value=mock_context_store ) - # Mock the client to raise APIConnectionError on models.list() + # Mock the client to raise APIConnectionError on openai.list() mock_client = mocker.AsyncMock() # Create a mock httpx.Request for APIConnectionError mock_request = httpx.Request("GET", "http://test-llama-stack/models") - mock_client.models.list.side_effect = APIConnectionError( + mock_client.openai.list.side_effect = APIConnectionError( message="Connection refused: unable to reach Llama Stack", request=mock_request, ) @@ -776,9 +777,10 @@ async def test_process_task_streaming_handles_api_connection_error( # pylint: d # Mock the client mock_client = mocker.AsyncMock() - mock_models = [mocker.MagicMock()] - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct(data=mock_models) + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model(model_id="test-model") + ) ) mocker.patch( "app.endpoints.a2a.AsyncOgxClientHolder" @@ -863,8 +865,8 @@ async def test_process_task_streaming_handles_agent_run_error( # pylint: disabl ) mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response(mocker.MagicMock()) ) mocker.patch( "app.endpoints.a2a.AsyncOgxClientHolder" @@ -940,8 +942,8 @@ async def test_process_task_streaming_applies_compaction( # pylint: disable=too ) mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response(mocker.MagicMock()) ) mocker.patch( "app.endpoints.a2a.AsyncOgxClientHolder" diff --git a/tests/unit/app/endpoints/test_models.py b/tests/unit/app/endpoints/test_models.py index 4d1b92f44..231ee7fe0 100644 --- a/tests/unit/app/endpoints/test_models.py +++ b/tests/unit/app/endpoints/test_models.py @@ -5,8 +5,6 @@ import pytest from fastapi import HTTPException, Request, status from ogx_client import APIConnectionError -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pytest_mock import MockerFixture from pytest_subtests import SubTests @@ -14,23 +12,10 @@ from authentication.interface import AuthTuple from configuration import AppConfig from models.api.requests import ModelFilter +from tests.unit.conftest import make_openai_model, make_openai_models_list_response from tests.unit.utils.auth_helpers import mock_authorization_resolvers -def _make_model(model_id: str, provider_id: str, model_type: str) -> Model: - """Build an OGX Model for models-endpoint tests.""" - return Model.model_construct( - id=model_id, - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": model_type, - "provider_id": provider_id, - }, - ) - - @pytest.mark.asyncio async def test_models_endpoint_handler_configuration_not_loaded( mocker: MockerFixture, @@ -162,7 +147,7 @@ async def test_models_endpoint_handler_unable_to_retrieve_models_list( # Mock the LlamaStack client mock_client = mocker.AsyncMock() - mock_client.models.list.return_value = ListModelsResponse.model_construct(data=[]) + mock_client.openai.list.return_value = make_openai_models_list_response() mock_lsc = mocker.patch("app.endpoints.models.AsyncOgxClientHolder.get_client") mock_lsc.return_value = mock_client mock_config = mocker.Mock() @@ -219,7 +204,7 @@ async def test_models_endpoint_handler_model_type_query_parameter( # Mock the LlamaStack client mock_client = mocker.AsyncMock() - mock_client.models.list.return_value = ListModelsResponse.model_construct(data=[]) + mock_client.openai.list.return_value = make_openai_models_list_response() mock_lsc = mocker.patch("app.endpoints.models.AsyncOgxClientHolder.get_client") mock_lsc.return_value = mock_client mock_config = mocker.Mock() @@ -275,13 +260,15 @@ async def test_models_endpoint_handler_model_list_retrieved( # Mock the LlamaStack client mock_client = mocker.AsyncMock() - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[ - _make_model("model1", "provider1", "llm"), - _make_model("model2", "provider2", "embedding"), - _make_model("model3", "provider3", "llm"), - _make_model("model4", "provider4", "embedding"), - ] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id="model1", provider_id="provider1", model_type="llm"), + make_openai_model( + model_id="model2", provider_id="provider2", model_type="embedding" + ), + make_openai_model(model_id="model3", provider_id="provider3", model_type="llm"), + make_openai_model( + model_id="model4", provider_id="provider4", model_type="embedding" + ), ) mock_lsc = mocker.patch("app.endpoints.models.AsyncOgxClientHolder.get_client") mock_lsc.return_value = mock_client @@ -349,13 +336,15 @@ async def test_models_endpoint_handler_model_list_retrieved_with_query_parameter # Mock the LlamaStack client mock_client = mocker.AsyncMock() - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[ - _make_model("model1", "provider1", "llm"), - _make_model("model2", "provider2", "embedding"), - _make_model("model3", "provider3", "llm"), - _make_model("model4", "provider4", "embedding"), - ] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id="model1", provider_id="provider1", model_type="llm"), + make_openai_model( + model_id="model2", provider_id="provider2", model_type="embedding" + ), + make_openai_model(model_id="model3", provider_id="provider3", model_type="llm"), + make_openai_model( + model_id="model4", provider_id="provider4", model_type="embedding" + ), ) mock_lsc = mocker.patch("app.endpoints.models.AsyncOgxClientHolder.get_client") mock_lsc.return_value = mock_client @@ -441,9 +430,9 @@ async def test_models_endpoint_llama_stack_connection_error( } # mock AsyncOgxClientHolder to raise APIConnectionError - # when models.list() method is called + # when openai.list() method is called mock_client = mocker.AsyncMock() - mock_client.models.list.side_effect = APIConnectionError(request=None) # type: ignore + mock_client.openai.list.side_effect = APIConnectionError(request=None) # type: ignore mock_client_holder = mocker.patch("app.endpoints.models.AsyncOgxClientHolder") mock_client_holder.return_value.get_client.return_value = mock_client diff --git a/tests/unit/app/endpoints/test_rlsapi_v1.py b/tests/unit/app/endpoints/test_rlsapi_v1.py index 61adf6eae..68c955a71 100644 --- a/tests/unit/app/endpoints/test_rlsapi_v1.py +++ b/tests/unit/app/endpoints/test_rlsapi_v1.py @@ -14,11 +14,11 @@ import pytest from fastapi import HTTPException, status from ogx_client import APIConnectionError, APIStatusError -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pydantic import ValidationError from pytest_mock import MockerFixture +from tests.unit.conftest import make_openai_model, make_openai_models_list_response + import constants from app.endpoints.rlsapi_v1 import ( AUTH_DISABLED, @@ -362,23 +362,21 @@ async def test_get_default_model_id_errors( """Test _get_default_model_id fallback failures raise 503 responses.""" mocker.patch("app.endpoints.rlsapi_v1.configuration", minimal_config) - mock_embedding_model = Model.model_construct( - id="sentence-transformers/all-mpnet-base-v2", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "embedding"}, + mock_embedding_model = make_openai_model( + model_id="sentence-transformers/all-mpnet-base-v2", + provider_id="sentence-transformers", + model_type="embedding", ) mock_client = mocker.Mock() mock_client.models = mocker.Mock() if failure_mode == "no_llm_models": - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct(data=[mock_embedding_model]) + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response(mock_embedding_model) ) else: - mock_client.models.list = mocker.AsyncMock( + mock_client.openai.list = mocker.AsyncMock( side_effect=APIConnectionError(request=mocker.Mock()) ) @@ -410,18 +408,16 @@ async def test_config_error_503_matches_llm_error_503_shape( """ mocker.patch("app.endpoints.rlsapi_v1.configuration", minimal_config) - mock_embedding_model = Model.model_construct( - id="sentence-transformers/all-mpnet-base-v2", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "embedding"}, + mock_embedding_model = make_openai_model( + model_id="sentence-transformers/all-mpnet-base-v2", + provider_id="sentence-transformers", + model_type="embedding", ) mock_client = mocker.Mock() mock_client.models = mocker.Mock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct(data=[mock_embedding_model]) + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response(mock_embedding_model) ) mock_client_holder = mocker.Mock() @@ -454,27 +450,23 @@ async def test_get_default_model_id_auto_discovery_success( """Test _get_default_model_id returns first discovered LLM model ID.""" mocker.patch("app.endpoints.rlsapi_v1.configuration", minimal_config) - mock_llm_model = Model.model_construct( - id="openai/gpt-4o-mini", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "llm"}, + mock_llm_model = make_openai_model( + model_id="openai/gpt-4o-mini", + provider_id="openai", + model_type="llm", ) - mock_embedding_model = Model.model_construct( - id="sentence-transformers/all-mpnet-base-v2", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "embedding"}, + mock_embedding_model = make_openai_model( + model_id="sentence-transformers/all-mpnet-base-v2", + provider_id="sentence-transformers", + model_type="embedding", ) mock_client = mocker.Mock() mock_client.models = mocker.Mock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[mock_embedding_model, mock_llm_model] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + mock_embedding_model, mock_llm_model ) ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index b16ea951e..eac22cd9b 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -5,11 +5,16 @@ import logging from collections.abc import Callable, Generator from pathlib import Path -from typing import Optional +from typing import Any, Optional import httpx import pytest from ogx_client import AsyncOgxClient +from ogx_client.models.list_models_v1_models_get200_response import ( + ListModelsV1ModelsGet200Response, +) +from ogx_client.models.open_ai_list_models_response import OpenAIListModelsResponse +from ogx_client.models.open_ai_model import OpenAIModel from pytest_mock import AsyncMockType, MockerFixture from configuration import AppConfig @@ -27,6 +32,58 @@ ] +def make_openai_model( + *, + model_id: str = "provider/model", + provider_id: str = "provider", + model_type: str = "llm", + provider_resource_id: Optional[str] = None, + **extra_metadata: Any, +) -> OpenAIModel: + """Build an ``OpenAIModel`` for ``client.openai.list()`` mocks.""" + custom_metadata: dict[str, Any] = { + "provider_id": provider_id, + "model_type": model_type, + "provider_resource_id": provider_resource_id or model_id, + **extra_metadata, + } + return OpenAIModel.model_construct( + id=model_id, + created=0, + owned_by="test", + object="model", + custom_metadata=custom_metadata, + ) + + +def make_openai_models_list_response( + *models: OpenAIModel, +) -> ListModelsV1ModelsGet200Response: + """Build a ``client.openai.list()`` response in the OpenAI OneOf shape.""" + return ListModelsV1ModelsGet200Response( + OpenAIListModelsResponse.model_construct(data=list(models)) + ) + + +def attach_mock_api_client( + mocker: MockerFixture, + client: Any, + *, + default_headers: Optional[dict[str, str]] = None, + async_http_client: Optional[httpx.AsyncClient] = None, +) -> Any: + """Attach ``api_client`` with headers and async httpx client to a mock OGX client.""" + api_client = mocker.Mock() + api_client.default_headers = default_headers if default_headers is not None else {} + api_client.async_client = ( + async_http_client + if async_http_client is not None + else mocker.Mock(spec=httpx.AsyncClient) + ) + client.api_client = api_client + return api_client + + @pytest.fixture(autouse=True) def reset_logging_state() -> Generator[None, None, None]: """Reset logging state before and after each test. @@ -115,8 +172,7 @@ def mock_client_fixture( # pylint: disable=protected-access client = mocker.Mock(spec=AsyncOgxClient) client.base_url = "http://localhost:8321" client.api_key = "test-key" - client._client = mocker.Mock(spec=httpx.AsyncClient) - client.default_headers = {} + attach_mock_api_client(mocker, client) return client diff --git a/tests/unit/metrics/test_utis.py b/tests/unit/metrics/test_utis.py index 86e6945b8..1b341f4e8 100644 --- a/tests/unit/metrics/test_utis.py +++ b/tests/unit/metrics/test_utis.py @@ -1,30 +1,15 @@ """Unit tests for functions defined in metrics/utils.py""" import pytest -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pytest_mock import MockerFixture from metrics.utils import setup_model_metrics - - -def _make_model(model_id: str, provider_id: str, model_type: str) -> Model: - """Build an OGX Model for metrics tests.""" - return Model.model_construct( - id=model_id, - created=0, - owned_by="test", - object="model", - custom_metadata={"provider_id": provider_id, "model_type": model_type}, - ) +from tests.unit.conftest import make_openai_model, make_openai_models_list_response @pytest.mark.asyncio async def test_setup_model_metrics(mocker: MockerFixture) -> None: """Test the setup_model_metrics function.""" - # Mock the OGXAsLibraryClient - mock_client = mocker.patch("client.AsyncOgxClientHolder.get_client").return_value - # Make sure the client is an AsyncMock for async methods mock_client = mocker.AsyncMock() mocker.patch("client.AsyncOgxClientHolder.get_client", return_value=mock_client) mocker.patch( @@ -37,19 +22,24 @@ async def test_setup_model_metrics(mocker: MockerFixture) -> None: ) mock_metric = mocker.patch("metrics.provider_model_configuration") - model_default = _make_model("default_model", "default_provider", "llm") - model_0 = _make_model("test_model-0", "test_provider-0", "llm") - model_1 = _make_model("test_model-1", "test_provider-1", "llm") - not_llm_model = _make_model("not-llm-model", "not-llm-provider", "not-llm") + model_default = make_openai_model( + model_id="default_model", provider_id="default_provider", model_type="llm" + ) + model_0 = make_openai_model( + model_id="test_model-0", provider_id="test_provider-0", model_type="llm" + ) + model_1 = make_openai_model( + model_id="test_model-1", provider_id="test_provider-1", model_type="llm" + ) + not_llm_model = make_openai_model( + model_id="not-llm-model", provider_id="not-llm-provider", model_type="not-llm" + ) - # Mock the list of models returned by the client - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[ - model_0, - model_default, - not_llm_model, - model_1, - ] + mock_client.openai.list.return_value = make_openai_models_list_response( + model_0, + model_default, + not_llm_model, + model_1, ) await setup_model_metrics() diff --git a/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py b/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py index 67b385a95..84f0e1107 100644 --- a/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py +++ b/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py @@ -14,6 +14,7 @@ OgxProvider, ) from pydantic_ai_lightspeed.llamastack._transport import OgxServerTransport +from tests.unit.conftest import attach_mock_api_client class TestOgxProviderProperties: @@ -168,8 +169,7 @@ def test_server_client_extracts_base_url_with_v1( mock_client = mocker.Mock(spec=AsyncOgxClient) mock_client.base_url = "http://my-server:8321/v1" mock_client.api_key = "test-key" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) provider = OgxProvider.from_ogx_client(mock_client) @@ -181,8 +181,7 @@ def test_server_client_appends_v1_when_missing(self, mocker: MockerFixture) -> N mock_client = mocker.Mock(spec=AsyncOgxClient) mock_client.base_url = "http://my-server:8321" mock_client.api_key = "test-key" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) provider = OgxProvider.from_ogx_client(mock_client) @@ -195,8 +194,7 @@ def test_server_client_strips_trailing_slash_before_appending_v1( mock_client = mocker.Mock(spec=AsyncOgxClient) mock_client.base_url = "http://my-server:8321/" mock_client.api_key = "test-key" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) provider = OgxProvider.from_ogx_client(mock_client) @@ -208,8 +206,7 @@ def test_server_client_uses_provided_api_key(self, mocker: MockerFixture) -> Non mock_client = mocker.Mock(spec=AsyncOgxClient) mock_client.base_url = "http://my-server:8321/v1" mock_client.api_key = "my-secret" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) provider = OgxProvider.from_ogx_client(mock_client) @@ -222,8 +219,7 @@ def test_server_client_defaults_api_key_when_none( mock_client = mocker.Mock(spec=AsyncOgxClient) mock_client.base_url = "http://my-server:8321/v1" mock_client.api_key = None - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) provider = OgxProvider.from_ogx_client(mock_client) @@ -235,8 +231,7 @@ def test_server_client_passes_http_client(self, mocker: MockerFixture) -> None: mock_client.base_url = "http://my-server:8321/v1" mock_client.api_key = "test-key" inner_http = mocker.Mock(spec=httpx.AsyncClient) - mock_client._client = inner_http - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client, async_http_client=inner_http) provider = OgxProvider.from_ogx_client(mock_client) @@ -250,10 +245,12 @@ def test_server_client_wraps_transport_with_provider_data( mock_client.base_url = "http://my-server:8321/v1" mock_client.api_key = "test-key" inner_http = httpx.AsyncClient() - mock_client._client = inner_http - mock_client.default_headers = { - "X-OGX-Provider-Data": '{"azure_api_key": "token"}' - } + attach_mock_api_client( + mocker, + mock_client, + async_http_client=inner_http, + default_headers={"X-OGX-Provider-Data": '{"azure_api_key": "token"}'}, + ) provider = OgxProvider.from_ogx_client(mock_client) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 0345ceb41..73c47afb7 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -10,14 +10,13 @@ import pytest from fastapi import HTTPException from ogx_client import APIConnectionError, APIStatusError -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pydantic import AnyHttpUrl, SecretStr from pytest_mock import MockerFixture from authorization.azure_token_manager import AzureEntraIDManager from client import AsyncOgxClientHolder from configuration import AzureEntraIdConfiguration +from tests.unit.conftest import make_openai_model, make_openai_models_list_response from models.config import LlamaStackConfiguration from utils.types import Singleton @@ -141,7 +140,9 @@ async def test_update_azure_token_service_client() -> None: assert updated_client is not original_client assert holder.get_client() is updated_client - provider_data_json = updated_client.default_headers.get("X-OGX-Provider-Data") + provider_data_json = updated_client.api_client.default_headers.get( + "X-OGX-Provider-Data" + ) provider_data = json.loads(provider_data_json) assert provider_data["azure_api_key"] == "fresh-token" assert provider_data["azure_api_base"] == "https://api.example.com" @@ -173,11 +174,13 @@ async def test_load_service_client_defers_azure_provider_data() -> None: holder = AsyncOgxClientHolder() await holder.load(cfg) - default_headers = holder.get_client().default_headers or {} + default_headers = holder.get_client().api_client.default_headers or {} assert "X-OGX-Provider-Data" not in default_headers updated_client = await holder.update_azure_token() - provider_data_json = updated_client.default_headers.get("X-OGX-Provider-Data") + provider_data_json = updated_client.api_client.default_headers.get( + "X-OGX-Provider-Data" + ) assert provider_data_json is not None provider_data = json.loads(provider_data_json) assert provider_data["azure_api_key"] == "startup-token" @@ -224,17 +227,6 @@ def holder_with_mock_client( holder._lsc = mock_client return holder, mock_client - def _make_model(self, mocker: MockerFixture, model_id: str) -> Model: - """Create an OGX Model with the given ID.""" - _ = mocker - return Model.model_construct( - id=model_id, - created=0, - owned_by="test", - object="model", - custom_metadata={}, - ) - @pytest.mark.asyncio async def test_model_available( self, @@ -243,8 +235,8 @@ async def test_model_available( ) -> None: """Test returns True when the model is found in the registry.""" holder, mock_client = holder_with_mock_client - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[self._make_model(mocker, self.EXPECTED_MODEL_ID)] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id=self.EXPECTED_MODEL_ID) ) available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) @@ -260,8 +252,8 @@ async def test_model_not_found_service_client( ) -> None: """Test returns False and skips reload for non-library (service) clients.""" holder, mock_client = holder_with_mock_client - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[self._make_model(mocker, "other/model")] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id="other/model") ) available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) @@ -305,7 +297,7 @@ async def test_api_error( ) -> None: """Test returns False when model list fails with API errors.""" _, mock_client = holder_with_mock_client - mock_client.models.list.side_effect = exception_factory(mocker) + mock_client.openai.list.side_effect = exception_factory(mocker) holder = AsyncOgxClientHolder() available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) @@ -329,11 +321,11 @@ async def test_model_found_after_reload( ) holder.reload_library_client = mocker.AsyncMock() - wrong_model = self._make_model(mocker, "other/model") - correct_model = self._make_model(mocker, self.EXPECTED_MODEL_ID) - mock_client.models.list.side_effect = [ - ListModelsResponse.model_construct(data=[wrong_model]), - ListModelsResponse.model_construct(data=[correct_model]), + wrong_model = make_openai_model(model_id="other/model") + correct_model = make_openai_model(model_id=self.EXPECTED_MODEL_ID) + mock_client.openai.list.side_effect = [ + make_openai_models_list_response(wrong_model), + make_openai_models_list_response(correct_model), ] available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) @@ -359,8 +351,8 @@ async def test_reload_fails_returns_not_found( holder.reload_library_client = mocker.AsyncMock( side_effect=RuntimeError("Cannot reload: config path not set") ) - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[self._make_model(mocker, "other/model")] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id="other/model") ) available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) @@ -385,8 +377,8 @@ async def test_reload_http_exception_returns_not_found( holder.reload_library_client = mocker.AsyncMock( side_effect=HTTPException(status_code=503, detail="Llama Stack unavailable") ) - mock_client.models.list.return_value = ListModelsResponse.model_construct( - data=[self._make_model(mocker, "other/model")] + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model(model_id="other/model") ) available, reason = await holder.check_model_available(self.EXPECTED_MODEL_ID) diff --git a/tests/unit/utils/test_pydantic_ai.py b/tests/unit/utils/test_pydantic_ai.py index bc477ce8e..82068743a 100644 --- a/tests/unit/utils/test_pydantic_ai.py +++ b/tests/unit/utils/test_pydantic_ai.py @@ -4,7 +4,6 @@ from collections.abc import Callable -import httpx import pytest from fastapi import HTTPException from ogx.core.library_client import AsyncOGXAsLibraryClient @@ -30,6 +29,7 @@ build_agent, get_agent_capability_tools, ) +from tests.unit.conftest import attach_mock_api_client _QUESTION_VALIDITY_MODULE = ( "pydantic_ai_lightspeed.capabilities.question_validity._capability" @@ -178,8 +178,7 @@ def test_returns_agent_with_correct_model( mock_client = mocker.Mock() mock_client.base_url = "http://localhost:8321" mock_client.api_key = "test-key" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) mock_params = mocker.Mock() mock_params.model = "provider/my-model" @@ -208,8 +207,7 @@ def test_agent_has_instructions( mock_client = mocker.Mock() mock_client.base_url = "http://localhost:8321" mock_client.api_key = "test-key" - mock_client._client = mocker.Mock(spec=httpx.AsyncClient) - mock_client.default_headers = {} + attach_mock_api_client(mocker, mock_client) mock_params = mocker.Mock() mock_params.model = "provider/my-model" diff --git a/tests/unit/utils/test_query.py b/tests/unit/utils/test_query.py index 1495ab9b1..ae5984a23 100644 --- a/tests/unit/utils/test_query.py +++ b/tests/unit/utils/test_query.py @@ -9,8 +9,6 @@ import psycopg2 import pytest from fastapi import HTTPException -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pydantic_ai.messages import ImageUrl from pytest_mock import MockerFixture from sqlalchemy.exc import SQLAlchemyError @@ -53,29 +51,6 @@ def mock_config_fixture() -> AppConfig: return cfg -@pytest.fixture(name="mock_models") -def mock_models_fixture() -> ListModelsResponse: - """Create an OpenAI-style OGX models list response.""" - return ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "llm", "provider_id": "provider1"}, - ), - Model.model_construct( - id="provider2/model2", - created=0, - owned_by="test", - object="model", - custom_metadata={"model_type": "llm", "provider_id": "provider2"}, - ), - ] - ) - - class TestStoreConversationIntoCache: """Tests for store_conversation_into_cache function.""" diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index 44acc2cb7..4879b60ff 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -51,11 +51,11 @@ OpenAIResponseOutputMessageWebSearchToolCall as WebSearchCall, ) from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient -from ogx_client.models.list_models_response import ListModelsResponse -from ogx_client.models.model import Model from pydantic import AnyUrl, BaseModel from pytest_mock import MockerFixture +from tests.unit.conftest import make_openai_model, make_openai_models_list_response + import constants from models.api.requests import QueryRequest from models.common.query import Attachment @@ -1853,20 +1853,13 @@ async def test_prepare_responses_params_with_conversation_id( ) -> None: """Test prepare_responses_params with existing conversation ID.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) @@ -1897,20 +1890,13 @@ async def test_prepare_responses_params_create_conversation( ) -> None: """Test prepare_responses_params creates new conversation when ID not provided.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) @@ -1941,7 +1927,7 @@ async def test_prepare_responses_params_connection_error_on_models( ) -> None: """Test prepare_responses_params raises HTTPException on connection error when fetching models.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( + mock_client.openai.list = mocker.AsyncMock( side_effect=APIConnectionError( message="Connection failed", request=mocker.Mock() ) @@ -1962,20 +1948,13 @@ async def test_prepare_responses_params_connection_error_on_conversation( ) -> None: """Test prepare_responses_params raises HTTPException on connection error when creating conversation.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) mock_client.conversations.create = mocker.AsyncMock( @@ -2003,7 +1982,7 @@ async def test_prepare_responses_params_api_status_error_on_models( ) -> None: """Test prepare_responses_params raises HTTPException on API status error when fetching models.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( + mock_client.openai.list = mocker.AsyncMock( side_effect=APIStatusError( message="API error", response=mocker.Mock(request=None), body=None ) @@ -2024,20 +2003,13 @@ async def test_prepare_responses_params_includes_mcp_provider_data_headers( ) -> None: """Test that extra_headers with x-llamastack-provider-data is set when MCP tools have headers.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) @@ -2100,20 +2072,13 @@ async def test_prepare_responses_params_no_extra_headers_without_mcp_tools( ) -> None: """Test that extra_headers is None when no MCP tools have headers.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) @@ -2145,20 +2110,13 @@ async def test_prepare_responses_params_api_status_error_on_conversation( ) -> None: """Test prepare_responses_params raises HTTPException on API status error when creating conversation.""" mock_client = mocker.AsyncMock() - mock_client.models.list = mocker.AsyncMock( - return_value=ListModelsResponse.model_construct( - data=[ - Model.model_construct( - id="provider1/model1", - created=0, - owned_by="test", - object="model", - custom_metadata={ - "model_type": "llm", - "provider_id": "provider1", - }, - ) - ] + mock_client.openai.list = mocker.AsyncMock( + return_value=make_openai_models_list_response( + make_openai_model( + model_id="provider1/model1", + provider_id="provider1", + model_type="llm", + ) ) ) mock_client.conversations.create = mocker.AsyncMock( diff --git a/uv.lock b/uv.lock index 92b4f539f..3a5596da6 100644 --- a/uv.lock +++ b/uv.lock @@ -1836,9 +1836,9 @@ requires-dist = [ { name = "jsonpath-ng", specifier = ">=1.6.1" }, { name = "kubernetes", specifier = ">=30.1.0" }, { name = "litellm", specifier = ">=1.83.7" }, - { name = "ogx", specifier = "==1.2.2" }, - { name = "ogx-api", specifier = "==1.2.2" }, - { name = "ogx-client", specifier = "==1.2.2" }, + { name = "ogx", specifier = "==1.3.0" }, + { name = "ogx-api", specifier = "==1.3.0" }, + { name = "ogx-client", specifier = "==1.3.0" }, { name = "openai", specifier = ">=1.99.9" }, { name = "opentelemetry-distro", specifier = ">=0.49b0" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.34.1" }, @@ -2375,7 +2375,7 @@ wheels = [ [[package]] name = "ogx" -version = "1.2.2" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiosqlite" }, @@ -2406,14 +2406,14 @@ dependencies = [ { name = "websockets" }, { name = "zstandard" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0f/a2/645000263e2f1dc1598dfdded9dda4a82280fda2e4b0d2b431aa5d0e128d/ogx-1.2.2.tar.gz", hash = "sha256:4ac77235ea884da00afca5209df8727dd876e671466f1525d7168db4ff5e3860", size = 18175152, upload-time = "2026-07-27T14:55:42.696Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/bd/5f8ce11504990bb25336a402acfca48c042f89c101130f63ed07292b8f0e/ogx-1.3.0.tar.gz", hash = "sha256:cf20a116fc5c93cbd88eed3c6d0d25eb52ca338a53a54b8d3c09917a9391340d", size = 18530591, upload-time = "2026-08-07T13:36:17.545Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c5/1e/17e6c2b0d6e00f66f5e9ff0dfe7e67e3076fa15f59a6f1c54c6495232e21/ogx-1.2.2-py3-none-any.whl", hash = "sha256:d47da083cb86452491d0d6b714733e7ef94df20f38a73a3c363b463faf03d02d", size = 778976, upload-time = "2026-07-27T14:55:39.874Z" }, + { url = "https://files.pythonhosted.org/packages/07/ea/a9490b815f2fd6c91d343ee878891391a2cf2907debd2dac161463e1bb70/ogx-1.3.0-py3-none-any.whl", hash = "sha256:e3e1f50afad849ff385e6c1e033d9116346a2166333fd4088fc778fcc9692b0d", size = 814873, upload-time = "2026-08-07T13:36:15.91Z" }, ] [[package]] name = "ogx-api" -version = "1.2.2" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fastapi" }, @@ -2424,14 +2424,14 @@ dependencies = [ { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/40/7f/e9c890ea8aac593e16d13800854d7855dc61d34761c84ae826b37982e643/ogx_api-1.2.2.tar.gz", hash = "sha256:9ff962c0550c6bdd310f3b4cac7be1b82ed8009126036dcc680f198afadc236f", size = 166463, upload-time = "2026-07-27T14:55:01.339Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/0a/1e5eee4ef6ce1b54828533150b4de23dc231d3a41da920428ce9aa3dcca1/ogx_api-1.3.0.tar.gz", hash = "sha256:4e67727df23aaaa1b6773de47440f1eb68790212013138e048247cc3cb2cf312", size = 167727, upload-time = "2026-08-07T13:35:19.678Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/25/2f/4e1e2b44c67737d49d943144a2ceccac6a1c4f0bf6462ca404a941c290b4/ogx_api-1.2.2-py3-none-any.whl", hash = "sha256:cbebca0520334f9e4d01ca57fc3fa39d3f155ea3467612e8781641724ea61361", size = 166522, upload-time = "2026-07-27T14:54:59.925Z" }, + { url = "https://files.pythonhosted.org/packages/e8/30/0be9638b7c521a9a0c22e1c87af2822410ef0d2e43bb65780413254e9920/ogx_api-1.3.0-py3-none-any.whl", hash = "sha256:f29b0a98dfb0f9cc694ca7d643efd1aa212b23290d5b02a3b543cf365a8f5299", size = 168296, upload-time = "2026-08-07T13:35:18.037Z" }, ] [[package]] name = "ogx-client" -version = "1.2.2" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fire" }, @@ -2441,9 +2441,9 @@ dependencies = [ { name = "requests" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4d/c5/a99d4d23447ddd2205e034f58e445252074a049ae2000c25b15d1c1b5b08/ogx_client-1.2.2.tar.gz", hash = "sha256:64ab1ada323acfb746bc99631f56b0470006d4fc59e05c89b83ddf7330718d6a", size = 588153, upload-time = "2026-07-27T14:54:17.411Z" } +sdist = { url = "https://files.pythonhosted.org/packages/50/e6/d45678153a2624870ce8571a1968c828c408f5ce029483c34019757af4b0/ogx_client-1.3.0.tar.gz", hash = "sha256:ebe5d0004433a7e1728b1f45e18bd3dabfddc6ac3d9f7bba51fb64e777fb921d", size = 648519, upload-time = "2026-08-07T13:34:42.221Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/2a/cc814caacc68dd1fa68b40fab6c3f31fc50aaa0bda357c802a940f1c4bf7/ogx_client-1.2.2-py3-none-any.whl", hash = "sha256:423e06b79508d86a2d9f0c4702c8cf564582106f8fd0853576092e1e12929590", size = 1544657, upload-time = "2026-07-27T14:54:15.536Z" }, + { url = "https://files.pythonhosted.org/packages/38/2b/45c20180b9664a94e38706c7c0fcee32eb56195b3b36924455b1900cf3d4/ogx_client-1.3.0-py3-none-any.whl", hash = "sha256:50d18b4952e814cc6785b396220e75780ef8a28ea01bec008d1f716fb8f08203", size = 1593011, upload-time = "2026-08-07T13:34:40.321Z" }, ] [[package]]