diff --git a/src/pydantic_ai_lightspeed/llamastack/_provider.py b/src/pydantic_ai_lightspeed/llamastack/_provider.py index 0e66f6d3c..e710da8ec 100644 --- a/src/pydantic_ai_lightspeed/llamastack/_provider.py +++ b/src/pydantic_ai_lightspeed/llamastack/_provider.py @@ -6,6 +6,7 @@ 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 @@ -13,7 +14,10 @@ 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,19 @@ 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" + 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 + 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 +141,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 index e5401bd68..7cd64c245 100644 --- 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/conftest.py b/tests/unit/conftest.py index b9d1ecaed..b654cd1e7 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -115,6 +115,7 @@ def mock_client_fixture( # pylint: disable=protected-access client.base_url = "http://localhost:8321" client.api_key = "test-key" client._client = mocker.Mock(spec=httpx.AsyncClient) + client.default_headers = {} return client diff --git a/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py b/tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py index dd49739af..3e02ad3a6 100644 --- 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 index 9dc2d112e..9a01c4454 100644 --- a/tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py +++ b/tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py @@ -12,7 +12,11 @@ from pydantic_ai_lightspeed.llamastack._transport import ( LlamaStackLibraryTransport, + LlamaStackServerTransport, _AsyncByteStream, + decode_request_headers, + inject_provider_data_into_headers, + request_with_provider_data_headers, ) @@ -70,6 +74,112 @@ async def gen() -> AsyncGenerator[bytes, None]: assert chunks == [] +class TestInjectProviderDataIntoHeaders: + """Tests for shared provider data header helpers.""" + + def test_injects_when_absent(self) -> None: + """Test provider data is serialized into the canonical header.""" + headers = inject_provider_data_into_headers({}, {"api_key": "test-key"}) + + assert headers["X-LlamaStack-Provider-Data"] == '{"api_key": "test-key"}' + + def test_does_not_override_existing_header(self) -> None: + """Test an existing provider data header is left unchanged.""" + headers = inject_provider_data_into_headers( + {"X-LlamaStack-Provider-Data": '{"existing": true}'}, + {"api_key": "ignored"}, + ) + + assert headers["X-LlamaStack-Provider-Data"] == '{"existing": true}' + + def test_request_with_provider_data_headers_returns_copy(self) -> None: + """Test request helper returns a new request when headers are injected.""" + request = httpx.Request( + "POST", + "http://localhost/v1/responses", + content=b"{}", + ) + + updated = request_with_provider_data_headers( + request, + {"api_key": "test-key"}, + ) + + assert updated is not request + assert ( + updated.headers["X-LlamaStack-Provider-Data"] == '{"api_key": "test-key"}' + ) + + +class TestLlamaStackServerTransport: + """Tests for LlamaStackServerTransport.""" + + @pytest.mark.asyncio + async def test_injects_provider_data_before_delegating( + self, mocker: MockerFixture + ) -> None: + """Test provider data is added to outbound HTTP requests.""" + wrapped = mocker.AsyncMock() + wrapped.handle_async_request.return_value = httpx.Response(200) + transport = LlamaStackServerTransport( + wrapped, + provider_data={"api_key": "test-key"}, + ) + + request = httpx.Request( + "POST", + "http://localhost/v1/responses", + content=b"{}", + ) + await transport.handle_async_request(request) + + delegated_request = wrapped.handle_async_request.await_args.args[0] + assert ( + delegated_request.headers["X-LlamaStack-Provider-Data"] + == '{"api_key": "test-key"}' + ) + + @pytest.mark.asyncio + async def test_preserves_existing_provider_data_header( + self, mocker: MockerFixture + ) -> None: + """Test per-request provider data headers are not overwritten.""" + wrapped = mocker.AsyncMock() + wrapped.handle_async_request.return_value = httpx.Response(200) + transport = LlamaStackServerTransport( + wrapped, + provider_data={"api_key": "ignored"}, + ) + + request = httpx.Request( + "POST", + "http://localhost/v1/responses", + content=b"{}", + headers={"X-LlamaStack-Provider-Data": '{"existing": true}'}, + ) + await transport.handle_async_request(request) + + delegated_request = wrapped.handle_async_request.await_args.args[0] + assert ( + delegated_request.headers["X-LlamaStack-Provider-Data"] + == '{"existing": true}' + ) + + +class TestDecodeRequestHeaders: # pylint: disable=too-few-public-methods + """Tests for decode_request_headers.""" + + def test_decodes_raw_headers(self) -> None: + """Test raw httpx headers are decoded to strings.""" + request = httpx.Request( + "GET", + "http://localhost/v1/models", + headers={"X-Test": "value"}, + ) + + assert decode_request_headers(request)["X-Test"] == "value" + + class TestLlamaStackLibraryTransportInit: # pylint: disable=too-few-public-methods """Tests for LlamaStackLibraryTransport initialization.""" diff --git a/tests/unit/utils/test_pydantic_ai.py b/tests/unit/utils/test_pydantic_ai.py index cb4e528c1..05809386f 100644 --- a/tests/unit/utils/test_pydantic_ai.py +++ b/tests/unit/utils/test_pydantic_ai.py @@ -66,6 +66,7 @@ def test_returns_agent_with_correct_model(self, mocker: MockerFixture) -> None: 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 = {} mock_params = mocker.Mock() mock_params.model = "provider/my-model" @@ -91,6 +92,7 @@ def test_agent_has_instructions(self, mocker: MockerFixture) -> None: 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 = {} mock_params = mocker.Mock() mock_params.model = "provider/my-model"