diff --git a/README.md b/README.md index f672d87..b3d034c 100644 --- a/README.md +++ b/README.md @@ -187,6 +187,26 @@ message = client.message( print(message.content) ``` +### Response metadata + +Use `with_response_metadata` to access the gateway's `X-Otari-Request-ID`: + +```python +result = client.with_response_metadata.message( + model="anthropic:claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello!"}], + max_tokens=256, +) +print(result.request_id, result.data.content) +``` + +For sync and async streams, `stream.request_id` is populated when iteration +starts, before the first event is yielded. + +Only a gateway running in platform mode sends this header, so `request_id` is +`None` when you call a standalone self-hosted gateway. On a stream, `None` also +means iteration has not opened the response yet. + ### Embeddings ```python diff --git a/src/otari/__init__.py b/src/otari/__init__.py index b42d8eb..df7523f 100644 --- a/src/otari/__init__.py +++ b/src/otari/__init__.py @@ -36,6 +36,7 @@ UnsupportedCapabilityError, UpstreamProviderError, ) +from otari.response_metadata import AsyncOtariStream, OtariResponse, OtariStream from otari.types import ( BatchRequestItem, BatchResult, @@ -62,6 +63,7 @@ __all__ = [ "AsyncOtariClient", + "AsyncOtariStream", "AuthenticationError", "BatchNotCompleteError", "BatchRequestItem", @@ -83,6 +85,8 @@ "OtariClient", "OtariClientOptions", "OtariError", + "OtariResponse", + "OtariStream", "RateLimitError", "RerankResponse", "TranscriptionResult", diff --git a/src/otari/async_client.py b/src/otari/async_client.py index c37f101..b594936 100644 --- a/src/otari/async_client.py +++ b/src/otari/async_client.py @@ -27,11 +27,11 @@ import asyncio from functools import cached_property -from typing import TYPE_CHECKING, Any, cast, overload +from typing import TYPE_CHECKING, Any, Literal, cast, overload import httpx -from otari._base import _BaseOtariClient, build_request +from otari._base import _BaseOtariClient, _header_get, build_request from otari._client import ApiClient, Configuration from otari._client.api.batches_api import BatchesApi from otari._client.api.chat_api import ChatApi @@ -54,6 +54,7 @@ from otari._streaming import aiter_sse from otari.control_plane import ControlPlane from otari.errors import OtariError +from otari.response_metadata import AsyncOtariStream, OtariResponse if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable @@ -63,6 +64,7 @@ from otari._client.models.count_tokens_response import CountTokensResponse from otari._client.models.create_embedding_response import CreateEmbeddingResponse from otari._client.models.images_response import ImagesResponse + from otari._client.models.message_response import MessageResponse from otari._client.models.model_object import ModelObject from otari._client.models.moderation_response import ModerationResponse from otari._client.models.rerank_response import RerankResponse @@ -148,6 +150,11 @@ def control_plane(self) -> ControlPlane: raise OtariError(msg) return ControlPlane(self._gateway_root_url, self._admin_token) + @cached_property + def with_response_metadata(self) -> AsyncOtariClientWithResponseMetadata: + """Inference methods that return per-request Otari response metadata.""" + return AsyncOtariClientWithResponseMetadata(self) + # -- Chat completions --------------------------------------------------- @overload @@ -470,6 +477,17 @@ async def _call(self, fn: Callable[[], Any]) -> Any: except ApiException as exc: raise self._map_api_exception(exc) from exc + async def _call_with_response_metadata( + self, + fn: Callable[[], Any], + ) -> OtariResponse[Any]: + """Run a generated HTTP-info call and preserve its Otari request ID.""" + response = await self._call(fn) + return OtariResponse( + data=response.data, + request_id=_header_get(response.headers, "X-Otari-Request-ID"), + ) + async def _post( self, path: str, @@ -494,6 +512,26 @@ async def _post( async def _stream(self, path: str, body: dict[str, Any], kind: Any) -> AsyncIterator[Any]: """Open a raw async streaming POST and yield parsed SSE chunks.""" + async for chunk in self._iter_stream(path, body, kind): + yield chunk + + def _stream_with_response_metadata( + self, + path: str, + body: dict[str, Any], + kind: Any, + ) -> AsyncOtariStream[Any]: + """Open a stream that exposes metadata for its individual request.""" + return AsyncOtariStream(lambda stream: self._iter_stream(path, body, kind, stream)) + + async def _iter_stream( + self, + path: str, + body: dict[str, Any], + kind: Any, + stream: AsyncOtariStream[Any] | None = None, + ) -> AsyncIterator[Any]: + """Issue and parse the raw HTTP streaming request.""" url = f"{self._base_url}{path}" headers = { "Content-Type": "application/json", @@ -504,6 +542,8 @@ async def _stream(self, path: str, body: dict[str, Any], kind: Any) -> AsyncIter if response.status_code >= 400: raw = await response.aread() raise self._map_streaming_response(response, raw) + if stream is not None: + stream._set_request_id(_header_get(response.headers, "X-Otari-Request-ID")) async for chunk in aiter_sse(response, kind): yield chunk @@ -519,3 +559,163 @@ async def __aenter__(self) -> AsyncOtariClient: async def __aexit__(self, *args: Any) -> None: await self.close() + + +class AsyncOtariClientWithResponseMetadata: + """Opt-in async inference API that retains metadata for each HTTP response.""" + + def __init__(self, client: AsyncOtariClient) -> None: + self._client = client + + @overload + async def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion]: ... + + @overload + async def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: Literal[True], + **kwargs: Any, + ) -> AsyncOtariStream[ChatCompletionChunk]: ... + + @overload + async def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion] | AsyncOtariStream[ChatCompletionChunk]: ... + + async def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion] | AsyncOtariStream[ChatCompletionChunk]: + """Create a chat completion and retain its Otari request ID.""" + body = {"model": model, "messages": messages, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/chat/completions", body, "chat") + request = build_request(ChatCompletionRequest, body) + return await self._client._call_with_response_metadata( + lambda: self._client._chat.chat_completions_v1_chat_completions_post_with_http_info( + request + ) + ) + + @overload + async def response( + self, + *, + model: str, + input: Any, + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[Any]: ... + + @overload + async def response( + self, + *, + model: str, + input: Any, + stream: Literal[True], + **kwargs: Any, + ) -> AsyncOtariStream[dict[str, Any]]: ... + + @overload + async def response( + self, + *, + model: str, + input: Any, + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[Any] | AsyncOtariStream[dict[str, Any]]: ... + + async def response( + self, + *, + model: str, + input: Any, # noqa: A002 + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[Any] | AsyncOtariStream[dict[str, Any]]: + """Create an OpenAI-style response and retain its Otari request ID.""" + body = {"model": model, "input": input, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/responses", body, "responses") + return await self._client._call_with_response_metadata( + lambda: self._client._responses.create_response_v1_responses_post_with_http_info( + body # type: ignore[arg-type] + ) + ) + + @overload + async def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse]: ... + + @overload + async def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: Literal[True], + **kwargs: Any, + ) -> AsyncOtariStream[dict[str, Any]]: ... + + @overload + async def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse] | AsyncOtariStream[dict[str, Any]]: ... + + async def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse] | AsyncOtariStream[dict[str, Any]]: + """Create an Anthropic-style message and retain its Otari request ID.""" + body = {"model": model, "messages": messages, "max_tokens": max_tokens, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/messages", body, "messages") + request = build_request(MessagesRequest, body) + return await self._client._call_with_response_metadata( + lambda: self._client._messages.create_message_v1_messages_post_with_http_info( + request + ) + ) diff --git a/src/otari/client.py b/src/otari/client.py index 7d73ab1..5a10349 100644 --- a/src/otari/client.py +++ b/src/otari/client.py @@ -27,11 +27,11 @@ from __future__ import annotations from functools import cached_property -from typing import TYPE_CHECKING, Any, cast, overload +from typing import TYPE_CHECKING, Any, Literal, cast, overload import httpx -from otari._base import _BaseOtariClient, build_request +from otari._base import _BaseOtariClient, _header_get, build_request from otari._client import ApiClient, Configuration from otari._client.api.batches_api import BatchesApi from otari._client.api.chat_api import ChatApi @@ -54,6 +54,7 @@ from otari._streaming import iter_sse from otari.control_plane import ControlPlane from otari.errors import OtariError +from otari.response_metadata import OtariResponse, OtariStream if TYPE_CHECKING: from collections.abc import Callable, Iterator @@ -63,6 +64,7 @@ from otari._client.models.count_tokens_response import CountTokensResponse from otari._client.models.create_embedding_response import CreateEmbeddingResponse from otari._client.models.images_response import ImagesResponse + from otari._client.models.message_response import MessageResponse from otari._client.models.model_object import ModelObject from otari._client.models.moderation_response import ModerationResponse from otari._client.models.rerank_response import RerankResponse @@ -155,6 +157,11 @@ def control_plane(self) -> ControlPlane: raise OtariError(msg) return ControlPlane(self._gateway_root_url, self._admin_token) + @cached_property + def with_response_metadata(self) -> OtariClientWithResponseMetadata: + """Inference methods that return per-request Otari response metadata.""" + return OtariClientWithResponseMetadata(self) + # -- Chat completions --------------------------------------------------- @overload @@ -494,6 +501,14 @@ def _call(self, fn: Callable[[], Any]) -> Any: except ApiException as exc: raise self._map_api_exception(exc) from exc + def _call_with_response_metadata(self, fn: Callable[[], Any]) -> OtariResponse[Any]: + """Run a generated HTTP-info call and preserve its Otari request ID.""" + response = self._call(fn) + return OtariResponse( + data=response.data, + request_id=_header_get(response.headers, "X-Otari-Request-ID"), + ) + def _post( self, path: str, @@ -517,11 +532,26 @@ def _post( return response def _stream(self, path: str, body: dict[str, Any], kind: Any) -> Iterator[Any]: - """Open a raw streaming POST and yield parsed SSE chunks. + """Open a raw streaming POST and yield parsed SSE chunks.""" + yield from self._iter_stream(path, body, kind) - The generated core buffers responses, so streaming is hand-written here: - a raw httpx streaming request parsed by :mod:`otari._streaming`. - """ + def _stream_with_response_metadata( + self, + path: str, + body: dict[str, Any], + kind: Any, + ) -> OtariStream[Any]: + """Open a stream that exposes metadata for its individual request.""" + return OtariStream(lambda stream: self._iter_stream(path, body, kind, stream)) + + def _iter_stream( + self, + path: str, + body: dict[str, Any], + kind: Any, + stream: OtariStream[Any] | None = None, + ) -> Iterator[Any]: + """Issue and parse the raw HTTP streaming request.""" url = f"{self._base_url}{path}" headers = { "Content-Type": "application/json", @@ -532,6 +562,8 @@ def _stream(self, path: str, body: dict[str, Any], kind: Any) -> Iterator[Any]: if response.status_code >= 400: raw = response.read() raise self._map_streaming_response(response, raw) + if stream is not None: + stream._set_request_id(_header_get(response.headers, "X-Otari-Request-ID")) yield from iter_sse(response, kind) # -- Cleanup ------------------------------------------------------------ @@ -546,3 +578,163 @@ def __enter__(self) -> OtariClient: def __exit__(self, *args: Any) -> None: self.close() + + +class OtariClientWithResponseMetadata: + """Opt-in inference API that retains metadata for each HTTP response.""" + + def __init__(self, client: OtariClient) -> None: + self._client = client + + @overload + def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion]: ... + + @overload + def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: Literal[True], + **kwargs: Any, + ) -> OtariStream[ChatCompletionChunk]: ... + + @overload + def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion] | OtariStream[ChatCompletionChunk]: ... + + def completion( + self, + *, + model: str, + messages: list[dict[str, Any]], + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[ChatCompletion] | OtariStream[ChatCompletionChunk]: + """Create a chat completion and retain its Otari request ID.""" + body = {"model": model, "messages": messages, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/chat/completions", body, "chat") + request = build_request(ChatCompletionRequest, body) + return self._client._call_with_response_metadata( + lambda: self._client._chat.chat_completions_v1_chat_completions_post_with_http_info( + request + ) + ) + + @overload + def response( + self, + *, + model: str, + input: Any, + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[Any]: ... + + @overload + def response( + self, + *, + model: str, + input: Any, + stream: Literal[True], + **kwargs: Any, + ) -> OtariStream[dict[str, Any]]: ... + + @overload + def response( + self, + *, + model: str, + input: Any, + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[Any] | OtariStream[dict[str, Any]]: ... + + def response( + self, + *, + model: str, + input: Any, # noqa: A002 + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[Any] | OtariStream[dict[str, Any]]: + """Create an OpenAI-style response and retain its Otari request ID.""" + body = {"model": model, "input": input, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/responses", body, "responses") + return self._client._call_with_response_metadata( + lambda: self._client._responses.create_response_v1_responses_post_with_http_info( + body # type: ignore[arg-type] + ) + ) + + @overload + def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: Literal[False] | None = None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse]: ... + + @overload + def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: Literal[True], + **kwargs: Any, + ) -> OtariStream[dict[str, Any]]: ... + + @overload + def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: bool | None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse] | OtariStream[dict[str, Any]]: ... + + def message( + self, + *, + model: str, + messages: list[dict[str, Any]], + max_tokens: int, + stream: bool | None = None, + **kwargs: Any, + ) -> OtariResponse[MessageResponse] | OtariStream[dict[str, Any]]: + """Create an Anthropic-style message and retain its Otari request ID.""" + body = {"model": model, "messages": messages, "max_tokens": max_tokens, **kwargs} + if stream: + body["stream"] = True + return self._client._stream_with_response_metadata("/messages", body, "messages") + request = build_request(MessagesRequest, body) + return self._client._call_with_response_metadata( + lambda: self._client._messages.create_message_v1_messages_post_with_http_info( + request + ) + ) diff --git a/src/otari/response_metadata.py b/src/otari/response_metadata.py new file mode 100644 index 0000000..7caae6e --- /dev/null +++ b/src/otari/response_metadata.py @@ -0,0 +1,76 @@ +"""Per-request response metadata wrappers for the Otari SDK.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Generic, TypeVar + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable, Iterator + +_T = TypeVar("_T") + + +@dataclass(frozen=True, slots=True) +class OtariResponse(Generic[_T]): + """A non-streaming response paired with its Otari request identifier.""" + + data: _T + request_id: str | None + + +class OtariStream(Generic[_T]): + """A synchronous response stream with metadata for its individual request. + + ``request_id`` is populated when iteration opens the HTTP response, before + the first event is yielded. It is therefore ``None`` before iteration starts. + """ + + def __init__(self, iterator_factory: Callable[[OtariStream[_T]], Iterator[_T]]) -> None: + self.request_id: str | None = None + self._iterator = iterator_factory(self) + + def __iter__(self) -> OtariStream[_T]: + return self + + def __next__(self) -> _T: + return next(self._iterator) + + def close(self) -> None: + """Close the underlying response stream.""" + close = getattr(self._iterator, "close", None) + if close is not None: + close() + + def _set_request_id(self, request_id: str | None) -> None: + self.request_id = request_id + + +class AsyncOtariStream(Generic[_T]): + """An asynchronous response stream with metadata for its individual request. + + ``request_id`` is populated when iteration opens the HTTP response, before + the first event is yielded. It is therefore ``None`` before iteration starts. + """ + + def __init__( + self, + iterator_factory: Callable[[AsyncOtariStream[_T]], AsyncIterator[_T]], + ) -> None: + self.request_id: str | None = None + self._iterator = iterator_factory(self) + + def __aiter__(self) -> AsyncOtariStream[_T]: + return self + + async def __anext__(self) -> _T: + return await self._iterator.__anext__() + + async def aclose(self) -> None: + """Close the underlying response stream.""" + close = getattr(self._iterator, "aclose", None) + if close is not None: + await close() + + def _set_request_id(self, request_id: str | None) -> None: + self.request_id = request_id diff --git a/tests/unit/test_async_client.py b/tests/unit/test_async_client.py index cdfb9fb..116ba2e 100644 --- a/tests/unit/test_async_client.py +++ b/tests/unit/test_async_client.py @@ -14,6 +14,7 @@ import pytest import respx +from otari import AsyncOtariStream, OtariResponse from otari._client.models.chat_completion import ChatCompletion from otari._client.models.chat_completion_chunk import ChatCompletionChunk from otari.async_client import AsyncOtariClient @@ -34,6 +35,7 @@ MESSAGE_RESPONSE, MODELS_RESPONSE, RERANK_RESPONSE, + RESPONSES_RESPONSE, TRANSCRIPTION_RESPONSE, _sse, ) @@ -69,6 +71,25 @@ async def test_completion_returns_typed(self, mock_rest: Any) -> None: assert mock.last.url.endswith("/v1/chat/completions") assert mock.last.headers.get("Otari-Key") == "Bearer vk" + async def test_response_with_metadata_exposes_request_id_without_changing_response( + self, + mock_rest: Any, + ) -> None: + mock_rest( + status=200, + body=RESPONSES_RESPONSE, + headers={"X-Otari-Request-ID": "req-async-response-123"}, + ) + client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") + + result = await client.with_response_metadata.response( + model="openai:gpt-4o-mini", + input="Hi", + ) + + assert result.data == RESPONSES_RESPONSE + assert result.request_id == "req-async-response-123" + async def test_embedding_returns_typed(self, mock_rest: Any) -> None: mock_rest(status=200, body=EMBEDDING_RESPONSE) client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") @@ -90,6 +111,24 @@ async def test_message_returns_typed(self, mock_rest: Any) -> None: assert result.id == "msg-1" assert mock.last.url.endswith("/v1/messages") + async def test_message_with_response_metadata_exposes_request_id(self, mock_rest: Any) -> None: + mock_rest( + status=200, + body=MESSAGE_RESPONSE, + headers={"x-otari-request-id": "req-async-message-123"}, + ) + client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") + + result = await client.with_response_metadata.message( + model="anthropic:claude", + messages=[{"role": "user", "content": "Hi"}], + max_tokens=8, + ) + + assert isinstance(result, OtariResponse) + assert result.data.id == "msg-1" + assert result.request_id == "req-async-message-123" + async def test_count_tokens_returns_typed(self, mock_rest: Any) -> None: mock = mock_rest(status=200, body=COUNT_TOKENS_RESPONSE) client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") @@ -157,6 +196,59 @@ async def test_yields_typed_chunks_and_stops_on_done(self) -> None: assert route.calls.last.request.headers["accept"] == "text/event-stream" assert route.calls.last.request.headers["otari-key"] == "Bearer vk" + @respx.mock + async def test_response_stream_metadata_exposes_request_id_without_changing_events(self) -> None: + event = '{"type":"response.completed","response":{"id":"resp-1"}}' + respx.post("http://localhost:8000/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "X-Otari-Request-ID": "req-async-response-stream-123", + }, + content=_sse(event), + ) + ) + client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") + + stream = await client.with_response_metadata.response( + model="openai:gpt-4o-mini", + input="Hi", + stream=True, + ) + + assert [event async for event in stream] == [ + {"type": "response.completed", "response": {"id": "resp-1"}} + ] + assert stream.request_id == "req-async-response-stream-123" + + @respx.mock + async def test_message_stream_metadata_exposes_request_id_without_mutating_events(self) -> None: + event = '{"type":"message_stop"}' + respx.post("http://localhost:8000/v1/messages").mock( + return_value=httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "X-Otari-Request-ID": "req-async-stream-123", + }, + content=_sse(event), + ) + ) + client = AsyncOtariClient(api_base="http://localhost:8000", api_key="vk") + + stream = await client.with_response_metadata.message( + model="anthropic:claude", + messages=[{"role": "user", "content": "Hi"}], + max_tokens=8, + stream=True, + ) + + assert isinstance(stream, AsyncOtariStream) + assert stream.request_id is None + assert [event async for event in stream] == [{"type": "message_stop"}] + assert stream.request_id == "req-async-stream-123" + @respx.mock async def test_streaming_error_maps(self) -> None: respx.post("http://localhost:8000/v1/chat/completions").mock( diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 9b70716..f5640cb 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -15,6 +15,7 @@ import pytest import respx +from otari import OtariResponse, OtariStream from otari._client.models.chat_completion import ChatCompletion from otari._client.models.chat_completion_chunk import ChatCompletionChunk from otari.client import OtariClient @@ -64,6 +65,13 @@ "usage": {"input_tokens": 1, "output_tokens": 1}, } +RESPONSES_RESPONSE: dict[str, Any] = { + "id": "resp-1", + "object": "response", + "status": "completed", + "output": [], +} + COUNT_TOKENS_RESPONSE: dict[str, Any] = {"input_tokens": 42} MODERATION_RESPONSE: dict[str, Any] = { @@ -187,6 +195,45 @@ def test_platform_mode_sends_bearer(self, mock_rest: Any) -> None: assert mock.last.headers.get("Authorization") == "Bearer tk" +class TestResponseMetadata: + def test_non_streaming_exposes_request_id_without_changing_response(self, mock_rest: Any) -> None: + mock_rest( + status=200, + body=RESPONSES_RESPONSE, + headers={"X-Otari-Request-ID": "req-response-123"}, + ) + client = OtariClient(api_base="http://localhost:8000", api_key="vk") + + result = client.with_response_metadata.response(model="openai:gpt-4o-mini", input="Hi") + + assert result.data == RESPONSES_RESPONSE + assert result.request_id == "req-response-123" + + @respx.mock + def test_streaming_exposes_request_id_without_changing_events(self) -> None: + event = '{"type":"response.completed","response":{"id":"resp-1"}}' + respx.post("http://localhost:8000/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "X-Otari-Request-ID": "req-response-stream-123", + }, + content=_sse(event), + ) + ) + client = OtariClient(api_base="http://localhost:8000", api_key="vk") + + stream = client.with_response_metadata.response( + model="openai:gpt-4o-mini", + input="Hi", + stream=True, + ) + + assert list(stream) == [{"type": "response.completed", "response": {"id": "resp-1"}}] + assert stream.request_id == "req-response-stream-123" + + class TestEmbedding: def test_returns_typed_embedding(self, mock_rest: Any) -> None: mock = mock_rest(status=200, body=EMBEDDING_RESPONSE) @@ -222,6 +269,41 @@ def test_returns_typed_message_response(self, mock_rest: Any) -> None: assert body["max_tokens"] == 64 assert body["model"] == "anthropic:claude-3-5-sonnet" + def test_with_response_metadata_exposes_request_id(self, mock_rest: Any) -> None: + mock_rest( + status=200, + body=MESSAGE_RESPONSE, + headers={"x-otari-request-id": "req-message-123"}, + ) + client = OtariClient(api_base="http://localhost:8000", api_key="vk") + + result = client.with_response_metadata.message( + model="anthropic:claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hi"}], + max_tokens=64, + ) + + assert isinstance(result, OtariResponse) + assert result.data.id == "msg-1" + assert result.request_id == "req-message-123" + + def test_with_response_metadata_is_available_for_chat(self, mock_rest: Any) -> None: + mock_rest( + status=200, + body=CHAT_RESPONSE, + headers={"X-Otari-Request-ID": "req-chat-123"}, + ) + client = OtariClient(api_base="http://localhost:8000", api_key="vk") + + result = client.with_response_metadata.completion( + model="openai:gpt-4o-mini", + messages=[{"role": "user", "content": "Hi"}], + ) + + assert isinstance(result, OtariResponse) + assert isinstance(result.data, ChatCompletion) + assert result.request_id == "req-chat-123" + def test_count_tokens_returns_typed_response(self, mock_rest: Any) -> None: mock = mock_rest(status=200, body=COUNT_TOKENS_RESPONSE) client = OtariClient(api_base="http://localhost:8000", api_key="vk") @@ -378,6 +460,33 @@ def test_platform_mode_streaming_sends_bearer(self) -> None: ) assert route.calls.last.request.headers["authorization"] == "Bearer tk" + @respx.mock + def test_message_stream_metadata_exposes_request_id_without_mutating_events(self) -> None: + event = '{"type":"message_stop"}' + respx.post("http://localhost:8000/v1/messages").mock( + return_value=httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "X-Otari-Request-ID": "req-stream-123", + }, + content=_sse(event), + ) + ) + client = OtariClient(api_base="http://localhost:8000", api_key="vk") + + stream = client.with_response_metadata.message( + model="anthropic:claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hi"}], + max_tokens=64, + stream=True, + ) + + assert isinstance(stream, OtariStream) + assert stream.request_id is None + assert list(stream) == [{"type": "message_stop"}] + assert stream.request_id == "req-stream-123" + # --------------------------------------------------------------------------- # Images (generated core) + audio (raw httpx)