LCORE-3315: Adapt to new OGX OpenApi client

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

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:

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",

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,

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

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"]

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. "

src/pydantic_ai_lightspeed/llamastack/_provider.py

@@ -6,14 +6,14 @@
 
 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
 from pydantic_ai.models import create_async_http_client
 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,

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",

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])

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()
     )
 

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

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()
     )
 

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()
     )
 

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()
     )