diff --git a/CHANGELOG.md b/CHANGELOG.md index 61b9057..5eea427 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ All notable changes to `uipath_llm_client` (core package) will be documented in this file. +## [1.19.0] - 2026-09-09 + +### Added +- `UiPathHttpxClient` / `UiPathHttpxAsyncClient` now capture the per-request dollar cost the LLM Gateway reports for opted-in calls instead of dropping it. When a request carries `X-UiPath-LlmGateway-IncludeAssociatedDollarCost: true` (Passthrough API only), the value is exposed via `uipath.llm_client.utils.dollar_cost.get_captured_dollar_cost()` for the duration of the calling context. Non-streaming JSON responses are read from the top-level `associated_dollar_cost` field. Streaming responses are wrapped so the trailing frame the gateway appends after the vendor's terminal event — an SSE `data: {"associated_dollar_cost": ...}` event, or a `costMetadata` message on AWS event streams (Bedrock) — is picked up once the consumer finishes reading (including when the vendor SDK stops at `data: [DONE]` and closes the response without reading further). `None` means "not priced" and must never be read as $0; when a cost was requested but could not be read (compressed stream, lost event-stream framing, non-numeric value) a warning is logged so the gap is visible. Known residual: the gateway's SSE frame also reaches google-genai as a candidate-less chunk, so opted-in Gemini streams log langchain-google-genai's "Gemini produced an empty response" warning; the cost is still captured. + ## [1.18.5] - 2026-09-07 ### Fixed diff --git a/packages/uipath_langchain_client/CHANGELOG.md b/packages/uipath_langchain_client/CHANGELOG.md index 0bbe841..0d5fba9 100644 --- a/packages/uipath_langchain_client/CHANGELOG.md +++ b/packages/uipath_langchain_client/CHANGELOG.md @@ -2,6 +2,14 @@ All notable changes to `uipath_langchain_client` will be documented in this file. +## [1.19.0] - 2026-09-09 + +### Added +- `UiPathBaseChatModel` now surfaces the LLM Gateway's per-request dollar cost as `response_metadata["associated_dollar_cost"]` when the caller opted in by sending `X-UiPath-LlmGateway-IncludeAssociatedDollarCost: true` (e.g. via `default_headers`) on a passthrough model. `invoke()` / `ainvoke()` set it on the result message. `stream()` / `astream()` (and `invoke()` routed through streaming) yield one extra empty chunk carrying it after the vendor's last chunk, since the gateway only reports the cost after the terminal event; merging the chunks folds it into `response_metadata`. Covers all passthrough chat models, including the Bedrock ones: the gateway's trailing `costMetadata` event is decoded by the core client and hidden from langchain-aws. Opted-in Gemini streams additionally log langchain-google-genai's "Gemini produced an empty response" warning for the gateway's trailing frame; the cost is still surfaced. The key is absent when the gateway did not price the call — absence means "not priced", never $0. + +### Changed +- Requires `uipath-llm-client>=1.19.0`. + ## [1.18.5] - 2026-09-07 ### Changed diff --git a/packages/uipath_langchain_client/pyproject.toml b/packages/uipath_langchain_client/pyproject.toml index 4dbdd9f..c58d5eb 100644 --- a/packages/uipath_langchain_client/pyproject.toml +++ b/packages/uipath_langchain_client/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "langchain>=1.2.15,<2.0.0", - "uipath-llm-client>=1.18.5,<2.0.0", + "uipath-llm-client>=1.19.0,<2.0.0", ] [project.optional-dependencies] diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py index c149855..c117173 100644 --- a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py +++ b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py @@ -1,3 +1,3 @@ __title__ = "UiPath LangChain Client" __description__ = "A Python client for interacting with UiPath's LLM services via LangChain." -__version__ = "1.18.5" +__version__ = "1.19.0" diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py b/packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py index 7c2a4ba..d6eec3c 100644 --- a/packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py +++ b/packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py @@ -36,7 +36,7 @@ ) from langchain_core.embeddings import Embeddings from langchain_core.language_models.chat_models import BaseChatModel -from langchain_core.messages import BaseMessage +from langchain_core.messages import AIMessageChunk, BaseMessage from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult from pydantic import ( AliasChoices, @@ -52,6 +52,10 @@ UiPathHttpxAsyncClient, UiPathHttpxClient, ) +from uipath.llm_client.utils.dollar_cost import ( + get_captured_dollar_cost, + set_captured_dollar_cost, +) from uipath.llm_client.utils.exceptions import wrap_provider_errors from uipath.llm_client.utils.headers import ( UIPATH_DEFAULT_REQUEST_HEADERS, @@ -524,15 +528,20 @@ def _generate( logger=self.logger, ) set_captured_response_headers({}) + # Models that bypass the UiPath httpx client (litellm) would otherwise + # inherit the value a previous request left in this context. + set_captured_dollar_cost(None) try: with wrap_provider_errors(): result = self._uipath_generate( messages, stop=stop, run_manager=run_manager, **kwargs ) self._inject_gateway_headers(result.generations) + self._inject_dollar_cost(result.generations) return result finally: set_captured_response_headers({}) + set_captured_dollar_cost(None) def _uipath_generate( self, @@ -558,15 +567,18 @@ async def _agenerate( logger=self.logger, ) set_captured_response_headers({}) + set_captured_dollar_cost(None) try: with wrap_provider_errors(): result = await self._uipath_agenerate( messages, stop=stop, run_manager=run_manager, **kwargs ) self._inject_gateway_headers(result.generations) + self._inject_dollar_cost(result.generations) return result finally: set_captured_response_headers({}) + set_captured_dollar_cost(None) async def _uipath_agenerate( self, @@ -592,6 +604,7 @@ def _stream( logger=self.logger, ) set_captured_response_headers({}) + set_captured_dollar_cost(None) try: first = True with wrap_provider_errors(): @@ -602,8 +615,12 @@ def _stream( self._inject_gateway_headers([chunk]) first = False yield chunk + cost_chunk = self._dollar_cost_chunk() + if cost_chunk is not None: + yield cost_chunk finally: set_captured_response_headers({}) + set_captured_dollar_cost(None) def _uipath_stream( self, @@ -629,6 +646,7 @@ async def _astream( logger=self.logger, ) set_captured_response_headers({}) + set_captured_dollar_cost(None) try: first = True with wrap_provider_errors(): @@ -639,8 +657,12 @@ async def _astream( self._inject_gateway_headers([chunk]) first = False yield chunk + cost_chunk = self._dollar_cost_chunk() + if cost_chunk is not None: + yield cost_chunk finally: set_captured_response_headers({}) + set_captured_dollar_cost(None) async def _uipath_astream( self, @@ -663,6 +685,28 @@ def _inject_gateway_headers(self, generations: Sequence[ChatGeneration]) -> None for generation in generations: generation.message.response_metadata["headers"] = headers + def _inject_dollar_cost(self, generations: Sequence[ChatGeneration]) -> None: + """Absent (None) means "not priced" and is never injected as $0.""" + cost = get_captured_dollar_cost() + if cost is None: + return + for generation in generations: + generation.message.response_metadata["associated_dollar_cost"] = cost + + def _dollar_cost_chunk(self) -> ChatGenerationChunk | None: + """Trailing empty chunk carrying the cost, or None if not priced. + + The cost is only known after the last content chunk was yielded. Holding + chunks back would delay every token, so it rides on an extra empty chunk + (like langchain-openai's usage chunk); merging folds it into response_metadata. + """ + cost = get_captured_dollar_cost() + if cost is None: + return None + return ChatGenerationChunk( + message=AIMessageChunk(content="", response_metadata={"associated_dollar_cost": cost}) + ) + class UiPathBaseEmbeddings(UiPathBaseLLMClient, Embeddings): pass diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/utils.py b/packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/utils.py index 67a93a4..3e280c7 100644 --- a/packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/utils.py +++ b/packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/utils.py @@ -5,6 +5,8 @@ from httpx import Client +from uipath.llm_client.utils.dollar_cost import COST_METADATA_EVENT_TYPE + try: from botocore.eventstream import EventStreamBuffer except ImportError as e: @@ -70,6 +72,10 @@ def _stream_generator( event_as_dict = event.to_response_dict() dict_key = event_as_dict["headers"][":event-type"] dict_value = json.loads(event_as_dict["body"].decode("utf-8")) + if dict_key == COST_METADATA_EVENT_TYPE: + # Gateway cost frame, already captured from the raw bytes by the + # httpx client; langchain-aws raises on unknown events. + continue if "bytes" in dict_value: dict_value["bytes"] = base64.b64decode(dict_value["bytes"]) yield {dict_key: dict_value} diff --git a/src/uipath/llm_client/__version__.py b/src/uipath/llm_client/__version__.py index 8dacbda..73c7345 100644 --- a/src/uipath/llm_client/__version__.py +++ b/src/uipath/llm_client/__version__.py @@ -1,3 +1,3 @@ __title__ = "UiPath LLM Client" __description__ = "A Python client for interacting with UiPath's LLM services." -__version__ = "1.18.5" +__version__ = "1.19.0" diff --git a/src/uipath/llm_client/httpx_client.py b/src/uipath/llm_client/httpx_client.py index 6df7a44..f67c135 100644 --- a/src/uipath/llm_client/httpx_client.py +++ b/src/uipath/llm_client/httpx_client.py @@ -49,6 +49,12 @@ ) from uipath.llm_client.settings.base import UiPathAPIConfig, UiPathBaseSettings +from uipath.llm_client.utils.dollar_cost import ( + attach_streaming_dollar_cost_capture, + extract_associated_dollar_cost, + requests_dollar_cost, + set_captured_dollar_cost, +) from uipath.llm_client.utils.exceptions import patch_raise_for_status from uipath.llm_client.utils.headers import ( UIPATH_DEFAULT_REQUEST_HEADERS, @@ -74,6 +80,24 @@ _DEFAULT_MAX_RETRIES: typing.Final[int] = 5 +def _capture_dollar_cost(request: Request, response: Response, stream: bool) -> None: + """Record the gateway-reported dollar cost of this exchange in the current context. + + Gated on the opt-in header: the field can only exist when asked for, and + re-parsing large (embeddings) bodies on every call is not free. Streaming + bodies are not received yet and carry the cost as a trailing frame, so the + stream is wrapped instead. Reset first so an unpriced response reads as None, + not as the previous request's cost. + """ + set_captured_dollar_cost(None) + if not requests_dollar_cost(request): + return + if stream: + attach_streaming_dollar_cost_capture(response) + else: + set_captured_dollar_cost(extract_associated_dollar_cost(response)) + + class UiPathHttpxClient(Client): """Synchronous HTTP client configured for UiPath LLM services. @@ -271,6 +295,7 @@ def send(self, request: Request, *, stream: bool = False, **kwargs: Any) -> Resp captured = extract_matching_headers(response.headers, self._captured_headers) if captured: set_captured_response_headers(captured) + _capture_dollar_cost(request, response, stream) return patch_raise_for_status(response) @@ -433,4 +458,5 @@ async def send(self, request: Request, *, stream: bool = False, **kwargs: Any) - captured = extract_matching_headers(response.headers, self._captured_headers) if captured: set_captured_response_headers(captured) + _capture_dollar_cost(request, response, stream) return patch_raise_for_status(response) diff --git a/src/uipath/llm_client/utils/dollar_cost.py b/src/uipath/llm_client/utils/dollar_cost.py new file mode 100644 index 0000000..b3560f4 --- /dev/null +++ b/src/uipath/llm_client/utils/dollar_cost.py @@ -0,0 +1,337 @@ +"""Opt-in per-request dollar cost reported by the LLM Gateway. + +The gateway prices a call only on the Passthrough API and only when the request +carries ``X-UiPath-LlmGateway-IncludeAssociatedDollarCost: true``. Non-streaming +JSON responses get a top-level ``associated_dollar_cost`` field. Streaming +responses get a trailing frame after the vendor's terminal event: an SSE +``data: {"associated_dollar_cost": n}`` event, or an AWS event-stream message +with ``:event-type=costMetadata``. + +A missing value means "not priced", never $0. +""" + +import base64 +import binascii +import contextvars +import json +import logging +import re +from collections.abc import AsyncIterator, Iterator +from typing import Protocol + +from httpx import AsyncByteStream, Request, Response, SyncByteStream + +logger = logging.getLogger(__name__) + +INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER = "X-UiPath-LlmGateway-IncludeAssociatedDollarCost" +ASSOCIATED_DOLLAR_COST_FIELD = "associated_dollar_cost" +COST_METADATA_EVENT_TYPE = "costMetadata" + +_CAPTURED_DOLLAR_COST: contextvars.ContextVar[float | None] = contextvars.ContextVar( + "_captured_dollar_cost", default=None +) + + +def get_captured_dollar_cost() -> float | None: + """Dollar cost captured from the most recent response in this context, or None if it was not priced.""" + return _CAPTURED_DOLLAR_COST.get() + + +def set_captured_dollar_cost(cost: float | None) -> contextvars.Token[float | None]: + """Set the captured per-request dollar cost for the current context.""" + return _CAPTURED_DOLLAR_COST.set(cost) + + +def requests_dollar_cost(request: Request) -> bool: + """Whether the request opted in; same bool parse as the gateway, so we only parse bodies it could have priced.""" + return request.headers.get(INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER, "").strip().lower() == "true" + + +def dollar_cost_from_body(body: object) -> float | None: + """Read ``associated_dollar_cost`` from a decoded JSON payload, or None if absent.""" + if not isinstance(body, dict) or ASSOCIATED_DOLLAR_COST_FIELD not in body: + return None + cost = body[ASSOCIATED_DOLLAR_COST_FIELD] # pyright: ignore[reportUnknownVariableType] + # bool is excluded: True would otherwise read as $1. + if isinstance(cost, bool) or not isinstance(cost, int | float): + # Unlike absence ("not priced"), a present but unusable value is contract drift. + logger.warning( + "Ignoring non-numeric %s=%r in gateway response", ASSOCIATED_DOLLAR_COST_FIELD, cost + ) + return None + return float(cost) + + +def extract_associated_dollar_cost(response: Response) -> float | None: + """Read ``associated_dollar_cost`` from a buffered JSON response body, or None if absent.""" + if "application/json" not in response.headers.get("content-type", ""): + return None + try: + body = response.json() + except ValueError: + return None + return dollar_cost_from_body(body) + + +def _dollar_cost_from_json_bytes(payload: bytes | str) -> float | None: + try: + body = json.loads(payload) + except ValueError: + return None + return dollar_cost_from_body(body) + + +# --- Streaming --------------------------------------------------------------------- + + +class CostFrames(Protocol): + """Tracks the last complete frame of one wire format.""" + + def feed(self, chunk: bytes) -> None: ... + + def ends_with_done(self) -> bool: + """Whether the last frame is the vendor's terminal one, where some SDKs stop reading.""" + ... + + def dollar_cost(self) -> float | None: + """The cost if the last frame is the gateway's cost frame, else None.""" + ... + + +# The cost frame is ~50 bytes and always last, so a small tail is enough. +_SSE_TAIL_BYTES = 4096 +_SSE_EVENT_SEPARATOR = re.compile(rb"\r?\n\r?\n") +# The openai SDK stops reading here and closes, leaving the gateway's frame unread. +_SSE_DONE_DATA = "[DONE]" + + +def _sse_event_data(event: bytes) -> str | None: + lines = [ + line[5:].removeprefix(" ") + for line in event.decode("utf-8", errors="replace").splitlines() + if line.startswith("data:") + ] + return "\n".join(lines) if lines else None + + +class SseCostFrames: + """``text/event-stream``: the cost is a trailing ``data: {...}`` event.""" + + def __init__(self) -> None: + self._tail = b"" + + def feed(self, chunk: bytes) -> None: + self._tail = (self._tail + chunk)[-_SSE_TAIL_BYTES:] + + def _last_complete_event_data(self) -> str | None: + # Ignore an unterminated remainder (no closing blank line yet). + events = [e for e in _SSE_EVENT_SEPARATOR.split(self._tail)[:-1] if e.strip()] + return _sse_event_data(events[-1]) if events else None + + def ends_with_done(self) -> bool: + return self._last_complete_event_data() == _SSE_DONE_DATA + + def dollar_cost(self) -> float | None: + data = self._last_complete_event_data() + return None if data is None else _dollar_cost_from_json_bytes(data) + + +# AWS event-stream frame: 12-byte prelude (total length, headers length, CRC), +# headers, payload, 4-byte CRC. Headers are matched as their exact encoded bytes +# (name length, name, type 7 = string, value length, value) instead of being parsed. +_EVENTSTREAM_PRELUDE_BYTES = 12 +_EVENTSTREAM_MIN_FRAME_BYTES = _EVENTSTREAM_PRELUDE_BYTES + 4 +# Above the AWS maximum we have lost framing; stop buffering rather than grow. +_EVENTSTREAM_MAX_FRAME_BYTES = 16 * 1024 * 1024 + + +def _event_type_header(event_type: str) -> bytes: + return ( + bytes([len(b":event-type")]) + + b":event-type" + + b"\x07" + + len(event_type).to_bytes(2, "big") + + event_type.encode() + ) + + +_EVENTSTREAM_COST_HEADER = _event_type_header(COST_METADATA_EVENT_TYPE) +# Bedrock's terminal markers: converse ends with a `metadata` event; invoke tags the +# vendor's last chunk with invocation metrics. +_EVENTSTREAM_METADATA_HEADER = _event_type_header("metadata") +_BEDROCK_INVOCATION_METRICS = b"amazon-bedrock-invocationMetrics" + + +class EventStreamCostFrames: + """``application/vnd.amazon.eventstream``: the cost is a trailing ``costMetadata`` message.""" + + def __init__(self) -> None: + self._buffer = b"" + self._last_frame = b"" + self._broken = False + + def feed(self, chunk: bytes) -> None: + if self._broken: + return + self._buffer += chunk + while len(self._buffer) >= _EVENTSTREAM_PRELUDE_BYTES: + total = int.from_bytes(self._buffer[:4], "big") + if total < _EVENTSTREAM_MIN_FRAME_BYTES or total > _EVENTSTREAM_MAX_FRAME_BYTES: + logger.warning( + "Lost AWS event-stream framing (frame length %d); dollar cost not tracked", + total, + ) + self._broken = True + self._buffer = b"" + return + if len(self._buffer) < total: + return + self._last_frame, self._buffer = self._buffer[:total], self._buffer[total:] + + def _last_frame_parts(self) -> tuple[bytes, bytes] | None: + frame = self._last_frame + if len(frame) < _EVENTSTREAM_MIN_FRAME_BYTES: + return None + headers_end = _EVENTSTREAM_PRELUDE_BYTES + int.from_bytes(frame[4:8], "big") + return frame[_EVENTSTREAM_PRELUDE_BYTES:headers_end], frame[headers_end:-4] + + def ends_with_done(self) -> bool: + """langchain-aws's invoke adapter stops at the vendor's stop event and closes without reading to EOF.""" + parts = self._last_frame_parts() + if parts is None: + return False + headers, payload = parts + if _EVENTSTREAM_METADATA_HEADER in headers: + return True + # invoke wraps the vendor JSON as base64 under "bytes". + try: + body = json.loads(payload) + raw = body.get("bytes") if isinstance(body, dict) else None + decoded = base64.b64decode(raw) if isinstance(raw, str) else b"" + except (ValueError, binascii.Error): + return False + return _BEDROCK_INVOCATION_METRICS in decoded + + def dollar_cost(self) -> float | None: + parts = self._last_frame_parts() + if parts is None: + return None + headers, payload = parts + if _EVENTSTREAM_COST_HEADER not in headers: + return None + return _dollar_cost_from_json_bytes(payload) + + +# Safety net for the drain in close(): only the cost frame should follow the terminal one. +_DRAIN_LIMIT_BYTES = 64 * 1024 + + +class DollarCostSyncStream(SyncByteStream): + """Pass-through byte stream that captures the gateway's trailing cost frame. + + Recorded on exhaustion or on ``close()``. Some consumers (openai SDK at + ``[DONE]``, langchain-aws at the vendor's stop event) close without reading + further, so ``close()`` first drains what follows the terminal frame. A close + mid-stream drains nothing: the frame cannot have arrived yet. + + The drain is bounded in bytes only; httpcore fixes the read timeout before the + body loop, so it inherits the client's. The gateway ends the response right + after the frame, and the anthropic/google SDKs already read to EOF anyway. + """ + + def __init__(self, inner: SyncByteStream, frames: CostFrames) -> None: + self._inner = inner + self._frames = frames + self._iterator: Iterator[bytes] | None = None + self._done = False + + def __iter__(self) -> Iterator[bytes]: + self._iterator = iter(self._inner) + for chunk in self._iterator: + self._frames.feed(chunk) + yield chunk + self._finish() + + def _finish(self) -> None: + if not self._done: + self._done = True + set_captured_dollar_cost(self._frames.dollar_cost()) + + def close(self) -> None: + try: + if self._iterator is not None and not self._done and self._frames.ends_with_done(): + drained = 0 + for chunk in self._iterator: + self._frames.feed(chunk) + drained += len(chunk) + if drained > _DRAIN_LIMIT_BYTES: + break + except Exception: # noqa: BLE001 - cost capture must never break closing the response + logger.debug("Failed to drain the trailing cost frame", exc_info=True) + finally: + if self._iterator is not None: + self._finish() + self._inner.close() + + +class DollarCostAsyncStream(AsyncByteStream): + """Async counterpart of :class:`DollarCostSyncStream`.""" + + def __init__(self, inner: AsyncByteStream, frames: CostFrames) -> None: + self._inner = inner + self._frames = frames + self._iterator: AsyncIterator[bytes] | None = None + self._done = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + self._iterator = self._inner.__aiter__() + async for chunk in self._iterator: + self._frames.feed(chunk) + yield chunk + self._finish() + + def _finish(self) -> None: + if not self._done: + self._done = True + set_captured_dollar_cost(self._frames.dollar_cost()) + + async def aclose(self) -> None: + try: + if self._iterator is not None and not self._done and self._frames.ends_with_done(): + drained = 0 + async for chunk in self._iterator: + self._frames.feed(chunk) + drained += len(chunk) + if drained > _DRAIN_LIMIT_BYTES: + break + except Exception: # noqa: BLE001 - cost capture must never break closing the response + logger.debug("Failed to drain the trailing cost frame", exc_info=True) + finally: + if self._iterator is not None: + self._finish() + await self._inner.aclose() + + +def _cost_frames_for(response: Response) -> CostFrames | None: + content_type = response.headers.get("content-type", "") + if "text/event-stream" in content_type: + return SseCostFrames() + if "application/vnd.amazon.eventstream" in content_type: + return EventStreamCostFrames() + return None + + +def attach_streaming_dollar_cost_capture(response: Response) -> None: + """Wrap a streaming response so its trailing cost frame is captured as it is consumed.""" + frames = _cost_frames_for(response) + if frames is None: + return + # The wrapper sees raw bytes; a compressed stream would silently read as "not priced". + encoding = response.headers.get("content-encoding", "identity").lower() + if encoding not in ("", "identity"): + logger.warning("Not capturing dollar cost from a %s-encoded stream", encoding) + return + if isinstance(response.stream, SyncByteStream): + response.stream = DollarCostSyncStream(response.stream, frames) + elif isinstance(response.stream, AsyncByteStream): + response.stream = DollarCostAsyncStream(response.stream, frames) diff --git a/tests/aws_event_stream.py b/tests/aws_event_stream.py new file mode 100644 index 0000000..49d2b17 --- /dev/null +++ b/tests/aws_event_stream.py @@ -0,0 +1,37 @@ +"""Hand-encoded AWS event-stream frames for tests (botocore only ships a decoder).""" + +import base64 +import json +import struct +import zlib + +import httpx + +from tests.lazy_stream import LazyByteStream + + +def event_frame(event_type: str, payload: dict) -> bytes: + """Encode one AWS event-stream message (prelude + string header + payload + CRCs).""" + name = b":event-type" + value = event_type.encode() + headers = bytes([len(name)]) + name + bytes([7]) + struct.pack(">H", len(value)) + value + body = json.dumps(payload).encode() + total_length = 12 + len(headers) + len(body) + 4 + prelude = struct.pack(">II", total_length, len(headers)) + prelude += struct.pack(">I", zlib.crc32(prelude)) + message = prelude + headers + body + return message + struct.pack(">I", zlib.crc32(message)) + + +def bedrock_chunk(event: dict) -> bytes: + """An invoke-with-response-stream ``chunk`` frame carrying one base64-encoded vendor event.""" + return event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +def event_stream_response(*frames: bytes) -> httpx.Response: + """A lazily delivered event-stream response, one frame per chunk.""" + return httpx.Response( + 200, + stream=LazyByteStream(list(frames)), + headers={"content-type": "application/vnd.amazon.eventstream"}, + ) diff --git a/tests/conftest.py b/tests/conftest.py index ffcdcd7..081a447 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ from uipath.llm_client.settings import UiPathBaseSettings from uipath.llm_client.settings.llmgateway import LLMGatewaySettings +from uipath.llm_client.utils.dollar_cost import set_captured_dollar_cost @pytest.fixture @@ -27,6 +28,14 @@ def _activate(design_time_id: str, connection_id: str, folder_key: str = "test-f _resource_overwrites.reset(token) +@pytest.fixture(autouse=True) +def reset_captured_dollar_cost(): + """The cost ContextVar would otherwise leak between tests.""" + set_captured_dollar_cost(None) + yield + set_captured_dollar_cost(None) + + @pytest.fixture(autouse=True, scope="session") def setup_env(): from dotenv import find_dotenv, load_dotenv diff --git a/tests/core/features/test_dollar_cost.py b/tests/core/features/test_dollar_cost.py new file mode 100644 index 0000000..57210d9 --- /dev/null +++ b/tests/core/features/test_dollar_cost.py @@ -0,0 +1,375 @@ +"""Tests for the opt-in per-request dollar cost helpers.""" + +import logging + +from httpx import AsyncByteStream, Request, Response, SyncByteStream + +from tests.aws_event_stream import bedrock_chunk, event_frame +from tests.lazy_stream import LazyByteStream +from uipath.llm_client.utils.dollar_cost import ( + INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER, + DollarCostAsyncStream, + DollarCostSyncStream, + EventStreamCostFrames, + SseCostFrames, + attach_streaming_dollar_cost_capture, + dollar_cost_from_body, + extract_associated_dollar_cost, + get_captured_dollar_cost, + requests_dollar_cost, +) + +LOGGER = "uipath.llm_client.utils.dollar_cost" + + +def _json_response(content: bytes, content_type: str = "application/json") -> Response: + return Response(200, headers={"content-type": content_type}, content=content) + + +def _request(header_value: str | None) -> Request: + headers = {} if header_value is None else {INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: header_value} + return Request("POST", "https://example.com", headers=headers) + + +class TestRequestsDollarCost: + def test_header_value_is_parsed_as_bool_like_the_gateway(self): + assert requests_dollar_cost(_request(" True ")) is True + + def test_false_header_does_not_opt_in(self): + assert requests_dollar_cost(_request("false")) is False + + def test_missing_header_does_not_opt_in(self): + assert requests_dollar_cost(_request(None)) is False + + +class TestExtractAssociatedDollarCost: + def test_extracts_cost_from_json_body(self): + response = _json_response(b'{"choices": [], "associated_dollar_cost": 0.002145}') + assert extract_associated_dollar_cost(response) == 0.002145 + + def test_absent_field_returns_none(self): + """Absence must read as 'not priced', never $0.""" + assert extract_associated_dollar_cost(_json_response(b'{"choices": []}')) is None + + def test_non_json_content_type_returns_none(self): + response = _json_response( + b'data: {"associated_dollar_cost": 0.002145}\n', content_type="text/event-stream" + ) + assert extract_associated_dollar_cost(response) is None + + def test_malformed_json_returns_none(self): + assert extract_associated_dollar_cost(_json_response(b"not json")) is None + + def test_non_object_json_returns_none(self): + assert extract_associated_dollar_cost(_json_response(b"[1, 2]")) is None + + def test_boolean_field_returns_none(self): + response = _json_response(b'{"associated_dollar_cost": true}') + assert extract_associated_dollar_cost(response) is None + + def test_non_numeric_field_returns_none_and_warns(self, caplog): + """A present but unusable value must not look like "not priced".""" + response = _json_response(b'{"associated_dollar_cost": "n/a"}') + with caplog.at_level(logging.WARNING, logger=LOGGER): + assert extract_associated_dollar_cost(response) is None + assert "associated_dollar_cost='n/a'" in caplog.text + + def test_integer_cost_reads_as_float(self): + assert dollar_cost_from_body({"associated_dollar_cost": 3}) == 3.0 + + +# ============================================================================ +# Streaming capture +# ============================================================================ + +COST_FRAME = b'data: {"associated_dollar_cost": 0.002145}\n\n' +CONTENT_EVENTS = [ + b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "Hel"}}]}\n\n', + b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "lo"}}]}\n\n', +] +DONE_EVENT = b"data: [DONE]\n\n" +EVENTSTREAM_COST_FRAME = event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) +EVENTSTREAM_STOP_FRAME = bedrock_chunk( + {"type": "message_stop", "amazon-bedrock-invocationMetrics": {}} +) + + +def _read_until_done(stream: DollarCostSyncStream) -> None: + """Consume like the openai SDK: stop as soon as the [DONE] event is seen.""" + for chunk in stream: + if DONE_EVENT in chunk: + break + + +def _sse_stream(inner: SyncByteStream) -> DollarCostSyncStream: + return DollarCostSyncStream(inner, SseCostFrames()) + + +def _sse_astream(inner: AsyncByteStream) -> DollarCostAsyncStream: + return DollarCostAsyncStream(inner, SseCostFrames()) + + +class TestSseCostFrames: + def test_cost_in_last_event(self): + frames = SseCostFrames() + frames.feed(b"".join([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME])) + assert frames.dollar_cost() == 0.002145 + + def test_cost_frame_split_across_feeds(self): + frames = SseCostFrames() + for chunk in [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME[:10], COST_FRAME[10:]]: + frames.feed(chunk) + assert frames.dollar_cost() == 0.002145 + + def test_crlf_and_multi_line_data(self): + frames = SseCostFrames() + frames.feed(b'data: {"associated_dollar_cost":\r\ndata: 0.002145}\r\n\r\n') + assert frames.dollar_cost() == 0.002145 + + def test_unterminated_trailing_event_is_ignored(self): + frames = SseCostFrames() + frames.feed(DONE_EVENT + COST_FRAME.rstrip(b"\n")) + assert frames.ends_with_done() + assert frames.dollar_cost() is None + + def test_event_and_comment_lines_are_skipped(self): + frames = SseCostFrames() + frames.feed(b': keep-alive\nevent: cost\ndata: {"associated_dollar_cost": 1}\n\n') + assert frames.dollar_cost() == 1.0 + + def test_done_is_terminal_and_not_a_cost(self): + frames = SseCostFrames() + frames.feed(DONE_EVENT) + assert frames.ends_with_done() + assert frames.dollar_cost() is None + + +class TestEventStreamCostFrames: + def test_cost_in_last_frame(self): + frames = EventStreamCostFrames() + frames.feed(event_frame("messageStop", {"stopReason": "end_turn"})) + frames.feed(EVENTSTREAM_COST_FRAME) + assert frames.dollar_cost() == 0.002145 + assert frames.ends_with_done() is False + + def test_frame_split_across_chunks(self): + frames = EventStreamCostFrames() + for i in range(0, len(EVENTSTREAM_COST_FRAME), 7): + frames.feed(EVENTSTREAM_COST_FRAME[i : i + 7]) + assert frames.dollar_cost() == 0.002145 + + def test_two_frames_in_one_chunk(self): + frames = EventStreamCostFrames() + frames.feed(event_frame("messageStop", {}) + EVENTSTREAM_COST_FRAME) + assert frames.dollar_cost() == 0.002145 + + def test_last_frame_not_cost_reads_none(self): + frames = EventStreamCostFrames() + frames.feed(EVENTSTREAM_COST_FRAME) + frames.feed(event_frame("messageStop", {"associated_dollar_cost": 0.002145})) + assert frames.dollar_cost() is None + + def test_converse_metadata_event_is_terminal(self): + frames = EventStreamCostFrames() + frames.feed(event_frame("metadata", {"usage": {}})) + assert frames.ends_with_done() is True + + def test_invoke_chunk_with_invocation_metrics_is_terminal(self): + frames = EventStreamCostFrames() + frames.feed(EVENTSTREAM_STOP_FRAME) + assert frames.ends_with_done() is True + + def test_ordinary_frames_are_not_terminal(self): + frames = EventStreamCostFrames() + frames.feed(bedrock_chunk({"type": "content_block_delta"})) + assert frames.ends_with_done() is False + frames.feed(event_frame("chunk", {"bytes": "not base64!"})) + assert frames.ends_with_done() is False + + def test_lost_framing_stops_tracking_and_warns(self, caplog): + frames = EventStreamCostFrames() + with caplog.at_level(logging.WARNING, logger=LOGGER): + frames.feed(b"\xff\xff\xff\xff" + b"\x00" * 12) + frames.feed(EVENTSTREAM_COST_FRAME) + assert frames.dollar_cost() is None + assert "Lost AWS event-stream framing" in caplog.text + + +class TestDollarCostSyncStream: + def test_full_read_captures_cost_and_passes_chunks_through(self): + chunks = [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] + assert list(_sse_stream(LazyByteStream(chunks))) == chunks + assert get_captured_dollar_cost() == 0.002145 + + def test_close_after_done_drains_the_cost_frame(self): + """The openai SDK stops at [DONE] and closes without reading the frame after it.""" + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_stream(inner) + _read_until_done(stream) + assert get_captured_dollar_cost() is None + stream.close() + assert get_captured_dollar_cost() == 0.002145 + assert inner.pulled == 4 + assert inner.closed + + def test_close_after_done_without_cost_frame_reads_none(self): + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT]) + stream = _sse_stream(inner) + _read_until_done(stream) + stream.close() + assert get_captured_dollar_cost() is None + assert inner.closed + + def test_drain_is_bounded(self): + """Only the cost frame should follow the terminal one; anything else must not be downloaded whole.""" + inner = LazyByteStream([DONE_EVENT, *([b"x" * 1024] * 200)]) + stream = _sse_stream(inner) + _read_until_done(stream) + stream.close() + assert 64 <= inner.pulled - 1 <= 66 + assert get_captured_dollar_cost() is None + assert inner.closed + + def test_early_close_mid_stream_does_not_drain(self): + """Cancelling mid-stream must not keep downloading the model's output.""" + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_stream(inner) + next(iter(stream)) + stream.close() + assert inner.pulled == 1 + assert get_captured_dollar_cost() is None + assert inner.closed + + def test_close_without_reading_only_closes_inner(self): + inner = LazyByteStream([COST_FRAME]) + _sse_stream(inner).close() + assert inner.pulled == 0 + assert inner.closed + assert get_captured_dollar_cost() is None + + def test_drain_failure_still_closes_inner(self): + class _Failing(LazyByteStream): + def __iter__(self): + yield DONE_EVENT + raise ConnectionError("boom") + + inner = _Failing([]) + stream = _sse_stream(inner) + _read_until_done(stream) + stream.close() + assert inner.closed + assert get_captured_dollar_cost() is None + + def test_event_stream_close_after_terminal_frame_drains_the_cost_frame(self): + """langchain-aws stops at the vendor's stop event and closes.""" + inner = LazyByteStream([EVENTSTREAM_STOP_FRAME, EVENTSTREAM_COST_FRAME]) + stream = DollarCostSyncStream(inner, EventStreamCostFrames()) + next(iter(stream)) + stream.close() + assert get_captured_dollar_cost() == 0.002145 + assert inner.pulled == 2 + assert inner.closed + + +class TestDollarCostAsyncStream: + async def test_full_read_captures_cost_and_passes_chunks_through(self): + chunks = [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] + assert [c async for c in _sse_astream(LazyByteStream(chunks))] == chunks + assert get_captured_dollar_cost() == 0.002145 + + async def test_close_after_done_drains_the_cost_frame(self): + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_astream(inner) + async for chunk in stream: + if DONE_EVENT in chunk: + break + await stream.aclose() + assert get_captured_dollar_cost() == 0.002145 + assert inner.pulled == 4 + assert inner.closed + + async def test_drain_is_bounded(self): + inner = LazyByteStream([DONE_EVENT, *([b"x" * 1024] * 200)]) + stream = _sse_astream(inner) + async for _ in stream: + break + await stream.aclose() + assert 64 <= inner.pulled - 1 <= 66 + assert get_captured_dollar_cost() is None + + async def test_early_close_mid_stream_does_not_drain(self): + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_astream(inner) + async for _ in stream: + break + await stream.aclose() + assert inner.pulled == 1 + assert get_captured_dollar_cost() is None + + async def test_close_without_reading_only_closes_inner(self): + inner = LazyByteStream([COST_FRAME]) + await _sse_astream(inner).aclose() + assert inner.pulled == 0 + assert inner.closed + assert get_captured_dollar_cost() is None + + async def test_drain_failure_still_closes_inner(self): + class _Failing(LazyByteStream): + async def __aiter__(self): + yield DONE_EVENT + raise ConnectionError("boom") + + inner = _Failing([]) + stream = _sse_astream(inner) + async for chunk in stream: + if DONE_EVENT in chunk: + break + await stream.aclose() + assert inner.closed + assert get_captured_dollar_cost() is None + + +class TestAttachStreamingDollarCostCapture: + def test_sse_stream_captures_cost(self): + response = Response(200, headers={"content-type": "text/event-stream; charset=utf-8"}) + response.stream = LazyByteStream([DONE_EVENT, COST_FRAME]) + attach_streaming_dollar_cost_capture(response) + assert isinstance(response.stream, DollarCostSyncStream) + list(response.stream) + assert get_captured_dollar_cost() == 0.002145 + + def test_async_sse_stream_is_wrapped(self): + class _AsyncOnly(AsyncByteStream): + async def __aiter__(self): + yield b"" + + response = Response(200, headers={"content-type": "text/event-stream"}) + response.stream = _AsyncOnly() + attach_streaming_dollar_cost_capture(response) + assert isinstance(response.stream, DollarCostAsyncStream) + + def test_aws_event_stream_captures_cost(self): + response = Response(200, headers={"content-type": "application/vnd.amazon.eventstream"}) + response.stream = LazyByteStream([EVENTSTREAM_STOP_FRAME, EVENTSTREAM_COST_FRAME]) + attach_streaming_dollar_cost_capture(response) + list(response.stream) + assert get_captured_dollar_cost() == 0.002145 + + def test_leaves_other_content_types_alone(self): + inner = LazyByteStream([]) + response = Response(200, headers={"content-type": "application/json"}) + response.stream = inner + attach_streaming_dollar_cost_capture(response) + assert response.stream is inner + + def test_leaves_compressed_streams_alone_and_warns(self, caplog): + """Compressed raw bytes cannot be parsed; warn instead of reading "not priced".""" + inner = LazyByteStream([]) + response = Response( + 200, headers={"content-type": "text/event-stream", "content-encoding": "gzip"} + ) + response.stream = inner + with caplog.at_level(logging.WARNING, logger=LOGGER): + attach_streaming_dollar_cost_capture(response) + assert response.stream is inner + assert "gzip" in caplog.text diff --git a/tests/core/features/test_httpx_client.py b/tests/core/features/test_httpx_client.py index 91d0068..cec3c7a 100644 --- a/tests/core/features/test_httpx_client.py +++ b/tests/core/features/test_httpx_client.py @@ -1,12 +1,15 @@ """Tests for HTTPX client functionality.""" +from typing import Any from unittest.mock import MagicMock, patch import pytest from httpx import Auth, Client, Headers, MockTransport, Request, Response +from tests.lazy_stream import LazyByteStream from uipath.llm_client.settings import UiPathAPIConfig from uipath.llm_client.settings.constants import ApiType, RoutingMode +from uipath.llm_client.utils.dollar_cost import INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER from uipath.llm_client.utils.retry import ( RetryableAsyncHTTPTransport, RetryableHTTPTransport, @@ -387,6 +390,131 @@ def test_response_patched_with_raise_for_status(self): assert result.raise_for_status is not original_raise client.close() + def _opted_in_request(self) -> Request: + return Request( + "POST", + "https://example.com/test", + headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, + ) + + def _json_response(self, body: dict[str, Any]) -> MagicMock: + mock_response = MagicMock(spec=Response) + mock_response.headers = Headers({"content-type": "application/json"}) + mock_response.json.return_value = body + mock_response.is_error = False + mock_response.raise_for_status = MagicMock(return_value=mock_response) + return mock_response + + def test_dollar_cost_captured_when_opted_in(self): + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + client = UiPathHttpxClient(base_url="https://example.com") + mock_response = self._json_response({"associated_dollar_cost": 0.002145}) + + with patch.object(Client, "send", return_value=mock_response): + client.send(self._opted_in_request(), stream=False) + assert get_captured_dollar_cost() == 0.002145 + client.close() + + def test_dollar_cost_body_not_parsed_when_not_opted_in(self): + """Without the opt-in header the field cannot exist, so the body is not re-parsed.""" + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + client = UiPathHttpxClient(base_url="https://example.com") + mock_response = self._json_response({"associated_dollar_cost": 0.002145}) + mock_response.json.side_effect = AssertionError("body must not be parsed without opt-in") + + with patch.object(Client, "send", return_value=mock_response): + client.send(Request("POST", "https://example.com/test"), stream=False) + assert get_captured_dollar_cost() is None + client.close() + + def test_dollar_cost_json_body_not_read_for_streaming_response(self): + """A streamed body is not buffered yet; a non-SSE stream reads as None.""" + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import ( + get_captured_dollar_cost, + set_captured_dollar_cost, + ) + + client = UiPathHttpxClient(base_url="https://example.com") + mock_response = self._json_response({}) + mock_response.json.side_effect = AssertionError("must not be read on a streamed response") + + set_captured_dollar_cost(0.5) # stale value from an earlier request + with patch.object(Client, "send", return_value=mock_response): + client.send(self._opted_in_request(), stream=True) + assert get_captured_dollar_cost() is None + client.close() + + def test_dollar_cost_captured_from_trailing_sse_frame(self): + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + events = [ + b'data: {"id": "1"}\n\n', + b"data: [DONE]\n\n", + b'data: {"associated_dollar_cost": 0.002145}\n\n', + ] + + def handler(request: Request) -> Response: + return Response( + 200, + request=request, + headers={"content-type": "text/event-stream"}, + stream=LazyByteStream(events), + ) + + client = UiPathHttpxClient(base_url="https://example.com", transport=MockTransport(handler)) + with client.stream( + "POST", "/test", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"} + ) as response: + assert b"".join(response.iter_bytes()) == b"".join(events) + assert get_captured_dollar_cost() == 0.002145 + client.close() + + def test_streaming_body_untouched_when_not_opted_in(self): + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import ( + DollarCostSyncStream, + get_captured_dollar_cost, + ) + + events = [b"data: [DONE]\n\n", b'data: {"associated_dollar_cost": 0.002145}\n\n'] + + def handler(request: Request) -> Response: + return Response( + 200, + request=request, + headers={"content-type": "text/event-stream"}, + stream=LazyByteStream(events), + ) + + client = UiPathHttpxClient(base_url="https://example.com", transport=MockTransport(handler)) + with client.stream("POST", "/test") as response: + assert not isinstance(response.stream, DollarCostSyncStream) + response.read() + assert get_captured_dollar_cost() is None + client.close() + + def test_unpriced_response_resets_previous_dollar_cost(self): + """An unpriced response must not read as the previous request's cost.""" + from uipath.llm_client.httpx_client import UiPathHttpxClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + client = UiPathHttpxClient(base_url="https://example.com") + priced = self._json_response({"associated_dollar_cost": 0.002145}) + unpriced = self._json_response({"choices": []}) + + with patch.object(Client, "send", side_effect=[priced, unpriced]): + client.send(self._opted_in_request(), stream=False) + assert get_captured_dollar_cost() == 0.002145 + client.send(self._opted_in_request(), stream=False) + assert get_captured_dollar_cost() is None + client.close() + class TestUiPathHttpxAsyncClientSend: @pytest.mark.asyncio @@ -450,3 +578,57 @@ async def handler(request: Request) -> Response: assert auth.signed_user_agents == ["custom-httpx-client"] await client.aclose() + + @pytest.mark.asyncio + async def test_dollar_cost_captured_when_opted_in(self): + from uipath.llm_client.httpx_client import UiPathHttpxAsyncClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + client = UiPathHttpxAsyncClient(base_url="https://example.com") + request = Request( + "POST", + "https://example.com/test", + headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, + ) + + async def handler(request: Request) -> Response: + return Response( + 200, + request=request, + headers={"content-type": "application/json"}, + content=b'{"associated_dollar_cost": 0.002145}', + ) + + client._transport = MockTransport(handler) + await client.send(request, stream=False) + assert get_captured_dollar_cost() == 0.002145 + await client.aclose() + + @pytest.mark.asyncio + async def test_dollar_cost_captured_from_trailing_sse_frame(self): + from uipath.llm_client.httpx_client import UiPathHttpxAsyncClient + from uipath.llm_client.utils.dollar_cost import get_captured_dollar_cost + + events = [ + b'data: {"id": "1"}\n\n', + b"data: [DONE]\n\n", + b'data: {"associated_dollar_cost": 0.002145}\n\n', + ] + + async def handler(request: Request) -> Response: + return Response( + 200, + request=request, + headers={"content-type": "text/event-stream"}, + stream=LazyByteStream(events), + ) + + client = UiPathHttpxAsyncClient( + base_url="https://example.com", transport=MockTransport(handler) + ) + async with client.stream( + "POST", "/test", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"} + ) as response: + assert b"".join([chunk async for chunk in response.aiter_bytes()]) == b"".join(events) + assert get_captured_dollar_cost() == 0.002145 + await client.aclose() diff --git a/tests/langchain/clients/bedrock/test_dollar_cost.py b/tests/langchain/clients/bedrock/test_dollar_cost.py new file mode 100644 index 0000000..56115e0 --- /dev/null +++ b/tests/langchain/clients/bedrock/test_dollar_cost.py @@ -0,0 +1,160 @@ +"""Gateway dollar-cost capture on the Bedrock chat models, through their real SDK decoders. + +Frames arrive one per chunk, the shape in which the gateway's trailing frame shows up. +""" + +from typing import Any + +import httpx +from uipath_langchain_client.clients.bedrock.chat_models import ( + UiPathChatAnthropicBedrock, + UiPathChatBedrock, + UiPathChatBedrockConverse, +) + +from tests.aws_event_stream import bedrock_chunk, event_frame, event_stream_response +from tests.langchain.transports import route_transport +from uipath.llm_client.settings import LLMGatewaySettings +from uipath.llm_client.utils.dollar_cost import INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER + +CONVERSE_STREAM_FRAMES = [ + event_frame("messageStart", {"role": "assistant"}), + event_frame("contentBlockDelta", {"delta": {"text": "Hello"}, "contentBlockIndex": 0}), + event_frame("contentBlockStop", {"contentBlockIndex": 0}), + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame( + "metadata", + { + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + }, + ), +] +COST_FRAME = event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) + + +def _make_chat(cls: type, settings: LLMGatewaySettings, frames: list[bytes]) -> Any: + chat = cls( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + settings=settings, + model_details={}, + default_headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, + ) + transport = httpx.MockTransport(lambda request: event_stream_response(*frames)) + route_transport(chat, transport, transport) + return chat + + +def test_converse_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_chat( + UiPathChatBedrockConverse, llmgw_settings, [*CONVERSE_STREAM_FRAMES, COST_FRAME] + ) + chunks = list(chat.stream("Hello")) + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + assert "".join(c.text for c in chunks) == "Hello" + + +async def test_converse_chat_astream_surfaces_gateway_cost(llmgw_settings) -> None: + """langchain-core bridges the sync stream per-next() under copied contexts; the cost must survive.""" + chat = _make_chat( + UiPathChatBedrockConverse, llmgw_settings, [*CONVERSE_STREAM_FRAMES, COST_FRAME] + ) + chunks = [c async for c in chat.astream("Hello")] + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + + +# UiPathChatAnthropicBedrock: the anthropic SDK's own event-stream decoder drops unknown +# events, so the cost must come from the raw bytes. + + +ANTHROPIC_STREAM_FRAMES = [ + bedrock_chunk( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude", + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + } + ), + bedrock_chunk( + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + ), + bedrock_chunk( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + } + ), + bedrock_chunk({"type": "content_block_stop", "index": 0}), + bedrock_chunk( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + } + ), + bedrock_chunk( + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 1, + "outputTokenCount": 1, + "invocationLatency": 1, + "firstByteLatency": 1, + }, + } + ), +] + + +def test_anthropic_bedrock_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_chat( + UiPathChatAnthropicBedrock, llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME] + ) + chunks = list(chat.stream("Hello")) + assert "".join(c.text for c in chunks) == "Hello" + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + + +async def test_anthropic_bedrock_astream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_chat( + UiPathChatAnthropicBedrock, llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME] + ) + chunks = [c async for c in chat.astream("Hello")] + assert "".join(c.text for c in chunks) == "Hello" + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + + +# UiPathChatBedrock: langchain-aws stops at the vendor's stop event without reading to +# EOF, so the cost frame must be drained on close. + + +def test_invoke_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_chat(UiPathChatBedrock, llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME]) + chunks = list(chat.stream("Hello")) + assert "".join(c.text for c in chunks) == "Hello" + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + + +def test_invoke_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None: + chat = _make_chat(UiPathChatBedrock, llmgw_settings, ANTHROPIC_STREAM_FRAMES) + chunks = list(chat.stream("Hello")) + assert all("associated_dollar_cost" not in c.response_metadata for c in chunks) diff --git a/tests/langchain/clients/bedrock/test_wrapped_boto_client.py b/tests/langchain/clients/bedrock/test_wrapped_boto_client.py index 4718958..df59e0c 100644 --- a/tests/langchain/clients/bedrock/test_wrapped_boto_client.py +++ b/tests/langchain/clients/bedrock/test_wrapped_boto_client.py @@ -13,6 +13,8 @@ import pytest from uipath_langchain_client.clients.bedrock.utils import WrappedBotoClient +from tests.aws_event_stream import event_frame, event_stream_response + _ERROR_BODY = { "title": "License not available", "status": 403, @@ -58,3 +60,18 @@ def test_invoke_model_with_response_stream_raises_on_http_error() -> None: stream = client.invoke_model_with_response_stream(body=json.dumps({"prompt": "hi"}))["body"] with pytest.raises(httpx.HTTPStatusError): list(stream) + + +def test_converse_stream_hides_gateway_cost_event_from_langchain_aws() -> None: + # langchain-aws raises on unknown events; the httpx client captures the cost itself. + body = event_stream_response( + event_frame("messageStart", {"role": "assistant"}), + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame("costMetadata", {"associated_dollar_cost": 0.002145}), + ) + client = _wrapped(lambda request: body) + events = list(client.converse_stream(messages=[])["stream"]) + assert events == [ + {"messageStart": {"role": "assistant"}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] diff --git a/tests/langchain/conftest.py b/tests/langchain/conftest.py index 7461ef5..f4478e1 100644 --- a/tests/langchain/conftest.py +++ b/tests/langchain/conftest.py @@ -4,6 +4,10 @@ distributed to per-provider conftest files under tests/langchain/clients/. """ +import os +from unittest.mock import patch + +import pytest from uipath_langchain_client.clients.anthropic.chat_models import UiPathChatAnthropic from uipath_langchain_client.clients.bedrock.chat_models import ( UiPathChatAnthropicBedrock, @@ -27,6 +31,24 @@ ) from uipath_langchain_client.clients.vertexai.chat_models import UiPathChatAnthropicVertex +from uipath.llm_client.settings import LLMGatewaySettings + +LLMGW_ENV = { + "LLMGW_URL": "https://cloud.uipath.com", + "LLMGW_SEMANTIC_ORG_ID": "test-org-id", + "LLMGW_SEMANTIC_TENANT_ID": "test-tenant-id", + "LLMGW_REQUESTING_PRODUCT": "test-product", + "LLMGW_REQUESTING_FEATURE": "test-feature", + "LLMGW_ACCESS_TOKEN": "test-access-token", +} + + +@pytest.fixture +def llmgw_settings(): + with patch.dict(os.environ, LLMGW_ENV, clear=True): + return LLMGatewaySettings() + + COMPLETION_CLIENTS_CLASSES = [ UiPathChat, UiPathChatOpenAI, diff --git a/tests/langchain/features/test_captured_headers.py b/tests/langchain/features/test_captured_headers.py index 0418779..d54ba4d 100644 --- a/tests/langchain/features/test_captured_headers.py +++ b/tests/langchain/features/test_captured_headers.py @@ -1,23 +1,28 @@ -"""Tests for captured response headers functionality. +"""Tests for LLM Gateway response capture into LangChain's response_metadata. -Tests that LLM Gateway response headers are captured and surfaced -in LangChain's response_metadata on AIMessage objects. +Covers captured response headers and the opt-in per-request dollar cost, for +both the non-streaming and the streaming paths. """ import json -import os -from unittest.mock import patch +from typing import Any import httpx import pytest from uipath_langchain_client.clients.normalized.chat_models import UiPathChat +from uipath_langchain_client.clients.openai.chat_models import UiPathChatOpenAI +from tests.langchain.transports import route_transport +from tests.lazy_stream import LazyByteStream from uipath.llm_client.httpx_client import ( UiPathHttpxAsyncClient, UiPathHttpxClient, ) from uipath.llm_client.settings import LLMGatewaySettings from uipath.llm_client.settings.utils import SingletonMeta +from uipath.llm_client.utils.dollar_cost import ( + INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER, +) from uipath.llm_client.utils.headers import ( extract_matching_headers, get_captured_response_headers, @@ -28,15 +33,6 @@ # Fixtures # ============================================================================ -LLMGW_ENV = { - "LLMGW_URL": "https://cloud.uipath.com", - "LLMGW_SEMANTIC_ORG_ID": "test-org-id", - "LLMGW_SEMANTIC_TENANT_ID": "test-tenant-id", - "LLMGW_REQUESTING_PRODUCT": "test-product", - "LLMGW_REQUESTING_FEATURE": "test-feature", - "LLMGW_ACCESS_TOKEN": "test-access-token", -} - SAMPLE_GATEWAY_HEADERS = { "X-UiPath-RequestId": "req-123", "X-UiPath-TraceId": "trace-456", @@ -94,21 +90,25 @@ def __init__( response_json: dict | None = None, response_headers: dict[str, str] | None = None, stream_chunks: list[dict] | None = None, + stream_cost: float | None = None, ): self._response_json = response_json or CHAT_RESPONSE_JSON self._response_headers = response_headers or SAMPLE_GATEWAY_HEADERS self._stream_chunks = stream_chunks + self._stream_cost = stream_cost def handle_request(self, request: httpx.Request) -> httpx.Response: headers = dict(self._response_headers) if self._stream_chunks and request.headers.get("X-UiPath-Streaming-Enabled") == "true": - lines = [] - for chunk in self._stream_chunks: - lines.append(f"data: {json.dumps(chunk)}\n") - lines.append("data: [DONE]\n") - content = "\n".join(lines).encode() + events = [f"data: {json.dumps(chunk)}\n\n".encode() for chunk in self._stream_chunks] + events.append(b"data: [DONE]\n\n") + if self._stream_cost is not None: + # As the gateway does: the cost frame trails the terminal event. + cost_json = json.dumps({"associated_dollar_cost": self._stream_cost}) + events.append(f"data: {cost_json}\n\n".encode()) headers["content-type"] = "text/event-stream" - return httpx.Response(200, content=content, headers=headers) + # One lazy chunk per event, so the SDK can stop at [DONE] before the cost frame. + return httpx.Response(200, stream=LazyByteStream(events), headers=headers) content = json.dumps(self._response_json).encode() headers["content-type"] = "application/json" @@ -123,8 +123,9 @@ def __init__( response_json: dict | None = None, response_headers: dict[str, str] | None = None, stream_chunks: list[dict] | None = None, + stream_cost: float | None = None, ): - self._sync = MockTransport(response_json, response_headers, stream_chunks) + self._sync = MockTransport(response_json, response_headers, stream_chunks, stream_cost) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: return self._sync.handle_request(request) @@ -138,12 +139,6 @@ def clear_singletons(): SingletonMeta._instances.clear() -@pytest.fixture -def llmgw_settings(): - with patch.dict(os.environ, LLMGW_ENV, clear=True): - return LLMGatewaySettings() - - def _make_normalized_chat( settings: LLMGatewaySettings, response_headers: dict[str, str] | None = None, @@ -436,3 +431,126 @@ def test_inject_gateway_headers_skipped_when_disabled(self, llmgw_settings): chat._inject_gateway_headers(result.generations) assert "headers" not in result.generations[0].message.response_metadata set_captured_response_headers({}) + + +# ============================================================================ +# Test opt-in per-request dollar cost capture +# ============================================================================ + + +def _make_passthrough_chat( + settings: LLMGatewaySettings, + response_json: dict, + default_headers: dict[str, str] | None = None, + stream_cost: float | None = None, + **model_kwargs: Any, +) -> UiPathChatOpenAI: + """Passthrough model with the real openai SDK in the loop, since that is what re-parses the body.""" + chat = UiPathChatOpenAI( + model="gpt-4o-2024-11-20", + settings=settings, + model_details={}, + default_headers=default_headers, + **model_kwargs, + ) + route_transport( + chat, + MockTransport( + response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost + ), + MockAsyncTransport( + response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost + ), + ) + return chat + + +OPT_IN_HEADERS = {INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"} +PRICED_RESPONSE_JSON = {**CHAT_RESPONSE_JSON, "associated_dollar_cost": 0.002145} + + +class TestDollarCostCapture: + """The gateway's opt-in per-request dollar cost ends up in response_metadata.""" + + def test_invoke_surfaces_cost_when_opted_in(self, llmgw_settings): + chat = _make_passthrough_chat( + llmgw_settings, response_json=PRICED_RESPONSE_JSON, default_headers=OPT_IN_HEADERS + ) + result = chat.invoke("Hello") + assert result.response_metadata["associated_dollar_cost"] == 0.002145 + + @pytest.mark.asyncio + async def test_ainvoke_surfaces_cost_when_opted_in(self, llmgw_settings): + chat = _make_passthrough_chat( + llmgw_settings, response_json=PRICED_RESPONSE_JSON, default_headers=OPT_IN_HEADERS + ) + result = await chat.ainvoke("Hello") + assert result.response_metadata["associated_dollar_cost"] == 0.002145 + + def test_invoke_omits_cost_when_gateway_did_not_price_the_call(self, llmgw_settings): + """Not priced must never surface as $0: the key must be absent.""" + chat = _make_passthrough_chat( + llmgw_settings, response_json=CHAT_RESPONSE_JSON, default_headers=OPT_IN_HEADERS + ) + result = chat.invoke("Hello") + assert "associated_dollar_cost" not in result.response_metadata + + def test_invoke_omits_cost_when_not_opted_in(self, llmgw_settings): + """Without the opt-in header the field cannot be gateway-authored.""" + chat = _make_passthrough_chat(llmgw_settings, response_json=PRICED_RESPONSE_JSON) + result = chat.invoke("Hello") + assert "associated_dollar_cost" not in result.response_metadata + + def test_stream_surfaces_cost_on_trailing_empty_chunk(self, llmgw_settings): + """The cost is only known after the last vendor chunk, so it rides on one extra chunk.""" + chat = _make_passthrough_chat( + llmgw_settings, + response_json=CHAT_RESPONSE_JSON, + default_headers=OPT_IN_HEADERS, + stream_cost=0.002145, + ) + chunks = list(chat.stream("Hello")) + assert "".join(str(c.content) for c in chunks) == "Hello!" + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].content == "" + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + merged = chunks[0] + for chunk in chunks[1:]: + merged = merged + chunk + assert merged.response_metadata["associated_dollar_cost"] == 0.002145 + + @pytest.mark.asyncio + async def test_astream_surfaces_cost_on_trailing_empty_chunk(self, llmgw_settings): + chat = _make_passthrough_chat( + llmgw_settings, + response_json=CHAT_RESPONSE_JSON, + default_headers=OPT_IN_HEADERS, + stream_cost=0.002145, + ) + chunks = [chunk async for chunk in chat.astream("Hello")] + assert "".join(str(c.content) for c in chunks) == "Hello!" + priced = [c for c in chunks if "associated_dollar_cost" in c.response_metadata] + assert len(priced) == 1 + assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 + + def test_stream_omits_cost_when_gateway_did_not_price_the_call(self, llmgw_settings): + chat = _make_passthrough_chat( + llmgw_settings, response_json=CHAT_RESPONSE_JSON, default_headers=OPT_IN_HEADERS + ) + chunks = list(chat.stream("Hello")) + assert "".join(str(c.content) for c in chunks) == "Hello!" + assert all("associated_dollar_cost" not in c.response_metadata for c in chunks) + + def test_invoke_routed_through_streaming_surfaces_cost(self, llmgw_settings): + """streaming=True routes invoke() through _stream.""" + chat = _make_passthrough_chat( + llmgw_settings, + response_json=CHAT_RESPONSE_JSON, + default_headers=OPT_IN_HEADERS, + stream_cost=0.002145, + streaming=True, + ) + result = chat.invoke("Hello") + assert result.content == "Hello!" + assert result.response_metadata["associated_dollar_cost"] == 0.002145 diff --git a/tests/langchain/transports.py b/tests/langchain/transports.py new file mode 100644 index 0000000..1ca5ed3 --- /dev/null +++ b/tests/langchain/transports.py @@ -0,0 +1,19 @@ +"""Route a chat model's UiPath httpx clients to a test transport.""" + +from typing import Any + +import httpx + + +def route_transport( + chat: Any, + transport: httpx.BaseTransport, + async_transport: httpx.AsyncBaseTransport | None = None, +) -> None: + # Vendor SDKs hold these very client objects, so swapping the transport is + # enough. Mounts are cleared because they would bypass the swapped transport. + chat.uipath_sync_client._transport = transport + chat.uipath_sync_client._mounts = {} + if async_transport is not None: + chat.uipath_async_client._transport = async_transport + chat.uipath_async_client._mounts = {} diff --git a/tests/lazy_stream.py b/tests/lazy_stream.py new file mode 100644 index 0000000..6490e3c --- /dev/null +++ b/tests/lazy_stream.py @@ -0,0 +1,32 @@ +"""A byte stream that delivers chunks lazily, like a network transport. + +``httpx.Response(content=...)`` pre-buffers the body, so ``iter_bytes()`` never +touches ``response.stream``; tests of stream wrappers need the lazy shape. +""" + +from collections.abc import AsyncIterator, Iterator + +from httpx import AsyncByteStream, SyncByteStream + + +class LazyByteStream(SyncByteStream, AsyncByteStream): + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.pulled = 0 + self.closed = False + + def __iter__(self) -> Iterator[bytes]: + for chunk in self._chunks: + self.pulled += 1 + yield chunk + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + self.pulled += 1 + yield chunk + + def close(self) -> None: + self.closed = True + + async def aclose(self) -> None: + self.closed = True