Patched transport layer for server client

asimurka · asimurka · commit 5cd7fc4b3ee8 · 2026-07-13T09:40:24.000+02:00
diff --git a/src/pydantic_ai_lightspeed/llamastack/_provider.py b/src/pydantic_ai_lightspeed/llamastack/_provider.py
@@ -6,14 +6,18 @@
 
 import httpx
 from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient
+from llama_stack.core.request_headers import parse_request_provider_data
 from llama_stack_client import AsyncLlamaStackClient
 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 pydantic_ai_lightspeed.llamastack._transport import LlamaStackLibraryTransport
+from pydantic_ai_lightspeed.llamastack._transport import (
+    LlamaStackLibraryTransport,
+    wrap_http_client_with_provider_data,
+)
 
 if TYPE_CHECKING:
     from llama_stack.core.library_client import (  # pylint: disable=reimported
@@ -73,10 +77,13 @@ def from_llama_stack_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"
+        provider_data = parse_request_provider_data(client.default_headers)
+        http_client = client._client  # pylint: disable=protected-access
+        http_client = wrap_http_client_with_provider_data(http_client, provider_data)
         return LlamaStackProvider(
             base_url=base_url,
             api_key=api_key,
-            http_client=client._client,  # pylint: disable=protected-access
+            http_client=http_client,
         )
 
     def __init__(
@@ -128,15 +135,11 @@ def __init__(
             base_url = base_url or DEFAULT_BASE_URL
             api_key = api_key or "not-needed"
 
-            if http_client is not None:
-                self._client = AsyncOpenAI(
-                    base_url=base_url, api_key=api_key, http_client=http_client
-                )
-            else:
-                oai_http_client = create_async_http_client()
-                self._client = AsyncOpenAI(
-                    base_url=base_url, api_key=api_key, http_client=oai_http_client
-                )
+            self._client = AsyncOpenAI(
+                base_url=base_url,
+                api_key=api_key,
+                http_client=http_client or create_async_http_client(),
+            )
 
     def __repr__(self) -> str:
         """Return a string representation of the provider."""
diff --git a/src/pydantic_ai_lightspeed/llamastack/_transport.py b/src/pydantic_ai_lightspeed/llamastack/_transport.py
@@ -1,9 +1,9 @@
-"""httpx transport that routes OpenAI-compatible requests through a Llama Stack library client."""
+"""httpx transports for Llama Stack library and server modes."""
 
 from __future__ import annotations as _annotations
 
 import json
-from collections.abc import AsyncGenerator, AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator, Mapping
 from typing import Any
 
 import httpx
@@ -19,6 +19,141 @@
 from llama_stack.core.utils.context import preserve_contexts_async_generator
 from starlette.responses import StreamingResponse
 
+_PROVIDER_DATA_HEADER_KEYS = (
+    "X-LlamaStack-Provider-Data",
+    "x-llamastack-provider-data",
+)
+
+
+def decode_request_headers(request: httpx.Request) -> dict[str, str]:
+    """Decode raw httpx request headers into a string-to-string mapping.
+
+    Args:
+        request: The outgoing httpx request.
+
+    Returns:
+        Request headers with byte keys and values decoded to UTF-8 strings.
+    """
+    return {
+        k.decode("utf-8") if isinstance(k, bytes) else k: (
+            v.decode("utf-8") if isinstance(v, bytes) else v
+        )
+        for k, v in request.headers.raw
+    }
+
+
+def inject_provider_data_into_headers(
+    headers: Mapping[str, str],
+    provider_data: Mapping[str, Any] | None,
+) -> dict[str, str]:
+    """Add ``X-LlamaStack-Provider-Data`` when provider data is configured.
+
+    Args:
+        headers: Existing request headers.
+        provider_data: Provider credentials/metadata to forward to Llama Stack.
+
+    Returns:
+        Headers with provider data injected when absent from the request.
+    """
+    if not provider_data:
+        return dict(headers)
+    if any(key in headers for key in _PROVIDER_DATA_HEADER_KEYS):
+        return dict(headers)
+    result = dict(headers)
+    result["X-LlamaStack-Provider-Data"] = json.dumps(provider_data)
+    return result
+
+
+def request_with_provider_data_headers(
+    request: httpx.Request,
+    provider_data: Mapping[str, Any] | None,
+) -> httpx.Request:
+    """Return a request copy with provider data headers injected when needed.
+
+    Args:
+        request: The outgoing httpx request.
+        provider_data: Provider credentials/metadata to forward to Llama Stack.
+
+    Returns:
+        The original request, or a copy with provider data headers added.
+    """
+    headers = inject_provider_data_into_headers(
+        decode_request_headers(request),
+        provider_data,
+    )
+    if headers == decode_request_headers(request):
+        return request
+    return httpx.Request(
+        method=request.method,
+        url=request.url,
+        headers=headers,
+        content=request.content,
+        extensions=request.extensions,
+    )
+
+
+def wrap_http_client_with_provider_data(
+    http_client: httpx.AsyncClient,
+    provider_data: Mapping[str, Any] | None,
+) -> httpx.AsyncClient:
+    """Wrap an httpx client so outbound requests include provider data headers.
+
+    Args:
+        http_client: The client whose transport will be wrapped.
+        provider_data: Provider credentials/metadata to forward to Llama Stack.
+
+    Returns:
+        The original client when ``provider_data`` is empty, otherwise a new
+        client with a wrapping transport.
+    """
+    if not provider_data:
+        return http_client
+
+    transport = LlamaStackServerTransport(
+        http_client._transport,  # pylint: disable=protected-access
+        provider_data=provider_data,
+    )
+    return httpx.AsyncClient(
+        transport=transport,
+        timeout=http_client.timeout,
+        follow_redirects=http_client.follow_redirects,
+    )
+
+
+class LlamaStackServerTransport(httpx.AsyncBaseTransport):
+    """httpx transport that injects provider data headers before delegating over HTTP."""
+
+    def __init__(
+        self,
+        transport: httpx.AsyncBaseTransport,
+        *,
+        provider_data: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Initialize the wrapping transport.
+
+        Args:
+            transport: The underlying transport used for real HTTP requests.
+            provider_data: Provider credentials/metadata to forward to Llama Stack.
+        """
+        self._transport = transport
+        self._provider_data = provider_data
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        """Inject provider data headers and delegate to the wrapped transport.
+
+        Args:
+            request: The outgoing httpx request.
+
+        Returns:
+            The response from the wrapped transport.
+        """
+        request = request_with_provider_data_headers(request, self._provider_data)
+        return await self._transport.handle_async_request(request)
+
+    async def aclose(self) -> None:
+        """Close the wrapped transport."""
+        await self._transport.aclose()
+
 
 class _AsyncByteStream(httpx.AsyncByteStream):
     """Wraps an async byte generator as an httpx AsyncByteStream."""
@@ -80,19 +215,10 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
 
         body = json.loads(request.content) if request.content else {}
 
-        headers: dict[str, str] = {
-            k.decode("utf-8") if isinstance(k, bytes) else k: (
-                v.decode("utf-8") if isinstance(v, bytes) else v
-            )
-            for k, v in request.headers.raw
-        }
-
-        if self._client.provider_data:
-            keys = ["X-LlamaStack-Provider-Data", "x-llamastack-provider-data"]
-            if all(key not in headers for key in keys):
-                headers["X-LlamaStack-Provider-Data"] = json.dumps(
-                    self._client.provider_data
-                )
+        headers = inject_provider_data_into_headers(
+            decode_request_headers(request),
+            self._client.provider_data,
+        )
 
         with request_provider_data_context(headers):
             is_stream = body.get("stream", False)
diff --git a/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py b/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py
@@ -13,6 +13,7 @@
     DEFAULT_BASE_URL,
     LlamaStackProvider,
 )
+from pydantic_ai_lightspeed.llamastack._transport import LlamaStackServerTransport
 
 
 class TestLlamaStackProviderProperties:
@@ -168,6 +169,7 @@ def test_server_client_extracts_base_url_with_v1(
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
@@ -180,6 +182,7 @@ def test_server_client_appends_v1_when_missing(self, mocker: MockerFixture) -> N
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
@@ -193,6 +196,7 @@ def test_server_client_strips_trailing_slash_before_appending_v1(
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
@@ -205,6 +209,7 @@ def test_server_client_uses_provided_api_key(self, mocker: MockerFixture) -> Non
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
@@ -218,23 +223,45 @@ def test_server_client_defaults_api_key_when_none(
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
         assert provider.client.api_key == "not-needed"
 
     def test_server_client_passes_http_client(self, mocker: MockerFixture) -> None:
-        """Test that the server client's internal httpx client is reused."""
+        """Test that the server client's internal httpx client is reused when no provider data."""
         mock_client = mocker.Mock(spec=AsyncLlamaStackClient)
         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 = {}
 
         provider = LlamaStackProvider.from_llama_stack_client(mock_client)
 
         assert provider._client._client is inner_http
 
+    def test_server_client_wraps_transport_with_provider_data(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test provider data from default_headers is forwarded in server mode."""
+        mock_client = mocker.Mock(spec=AsyncLlamaStackClient)
+        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-LlamaStack-Provider-Data": '{"azure_api_key": "token"}'
+        }
+
+        provider = LlamaStackProvider.from_llama_stack_client(mock_client)
+
+        assert isinstance(
+            provider._client._client._transport,  # pylint: disable=protected-access
+            LlamaStackServerTransport,
+        )
+
 
 class TestSetHttpClient:  # pylint: disable=too-few-public-methods
     """Tests for LlamaStackProvider._set_http_client."""
diff --git a/tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py b/tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py
