From a02221ea6f1ed817cd32e65b77006858c7a3df32 Mon Sep 17 00:00:00 2001 From: Andrei Tava Date: Wed, 9 Sep 2026 13:28:30 +0300 Subject: [PATCH 1/2] feat: surface the LLM Gateway per-request dollar cost instead of dropping it Callers can opt in with X-UiPath-LlmGateway-IncludeAssociatedDollarCost: true (Passthrough API). The gateway then adds associated_dollar_cost to non-streaming JSON bodies and appends a trailing frame to streams (SSE data event, AWS event-stream costMetadata message). The vendor SDKs re-parse the body or stop reading at their terminal event, so the value never reached callers. Core: UiPathHttpxClient/UiPathHttpxAsyncClient capture it into a ContextVar for opted-in requests. Non-streaming bodies are read directly; streaming responses are wrapped so the trailing frame is captured on exhaustion, or drained on close when the SDK stops at the vendor's terminal event ([DONE], Bedrock's metadata / invocation-metrics frame). LangChain: invoke()/ainvoke() set response_metadata["associated_dollar_cost"]; stream()/astream() yield one extra empty chunk carrying it after the vendor's last chunk. The Bedrock shim hides the gateway's costMetadata event from langchain-aws. Absence means "not priced", never $0. --- CHANGELOG.md | 5 + packages/uipath_langchain_client/CHANGELOG.md | 8 + .../uipath_langchain_client/pyproject.toml | 2 +- .../uipath_langchain_client/__version__.py | 2 +- .../uipath_langchain_client/base_client.py | 44 +- .../clients/bedrock/utils.py | 6 + src/uipath/llm_client/__version__.py | 2 +- src/uipath/llm_client/httpx_client.py | 26 ++ src/uipath/llm_client/utils/dollar_cost.py | 339 ++++++++++++++ tests/aws_event_stream.py | 37 ++ tests/conftest.py | 9 + tests/core/features/test_dollar_cost.py | 431 ++++++++++++++++++ tests/core/features/test_httpx_client.py | 182 ++++++++ .../clients/bedrock/test_dollar_cost.py | 217 +++++++++ .../bedrock/test_wrapped_boto_client.py | 17 + .../features/test_captured_headers.py | 179 +++++++- tests/lazy_stream.py | 27 ++ 17 files changed, 1519 insertions(+), 14 deletions(-) create mode 100644 src/uipath/llm_client/utils/dollar_cost.py create mode 100644 tests/aws_event_stream.py create mode 100644 tests/core/features/test_dollar_cost.py create mode 100644 tests/langchain/clients/bedrock/test_dollar_cost.py create mode 100644 tests/lazy_stream.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 61b90576..5eea4277 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 0bbe841a..0d5fba94 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 4dbdd9fd..c58d5eb9 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 c1498552..c1171734 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 7c2a4ba6..9cc0082d 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,18 @@ def _generate( logger=self.logger, ) set_captured_response_headers({}) + 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 +565,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 +602,7 @@ def _stream( logger=self.logger, ) set_captured_response_headers({}) + set_captured_dollar_cost(None) try: first = True with wrap_provider_errors(): @@ -602,8 +613,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 +644,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 +655,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 +683,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 67a93a45..3e280c7d 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 8dacbda1..73c73458 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 6df7a443..f67c135e 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 00000000..a5431e02 --- /dev/null +++ b/src/uipath/llm_client/utils/dollar_cost.py @@ -0,0 +1,339 @@ +"""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() + if data is None or data == _SSE_DONE_DATA: + return None + return _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 00000000..49d2b176 --- /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 ffcdcd7f..081a447d 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 00000000..0e9fb2ba --- /dev/null +++ b/tests/core/features/test_dollar_cost.py @@ -0,0 +1,431 @@ +"""Tests for the opt-in per-request dollar cost helpers.""" + +from httpx import AsyncByteStream, Request, Response, SyncByteStream + +from tests.aws_event_stream import bedrock_chunk, event_frame +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, + set_captured_dollar_cost, +) + + +def _json_response(content: bytes, content_type: str = "application/json") -> Response: + return Response(200, headers={"content-type": content_type}, content=content) + + +class TestRequestsDollarCost: + def test_true_header_opts_in(self): + request = Request( + "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"} + ) + assert requests_dollar_cost(request) is True + + def test_header_value_is_parsed_as_bool_like_the_gateway(self): + request = Request( + "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: " True "} + ) + assert requests_dollar_cost(request) is True + + def test_false_header_does_not_opt_in(self): + request = Request( + "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "false"} + ) + assert requests_dollar_cost(request) is False + + def test_missing_header_does_not_opt_in(self): + assert requests_dollar_cost(Request("POST", "https://example.com")) 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_non_numeric_field_returns_none(self): + response = _json_response(b'{"associated_dollar_cost": "n/a"}') + assert extract_associated_dollar_cost(response) 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_is_logged(self, caplog): + """A present but unusable value must not look like "not priced".""" + import logging + + response = _json_response(b'{"associated_dollar_cost": "n/a"}') + with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + assert extract_associated_dollar_cost(response) is None + assert "associated_dollar_cost='n/a'" in caplog.text + + +class TestDollarCostFromBody: + def test_reads_numeric_field(self): + assert dollar_cost_from_body({"associated_dollar_cost": 3}) == 3.0 + + def test_non_dict_returns_none(self): + assert dollar_cost_from_body(["associated_dollar_cost"]) is None + + +class TestCapturedDollarCostContextVar: + def test_default_is_none(self): + assert get_captured_dollar_cost() is None + + def test_set_and_get(self): + set_captured_dollar_cost(0.002145) + assert get_captured_dollar_cost() == 0.002145 + + +# ============================================================================ +# 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" + + +class _CountingSyncStream(SyncByteStream): + """Records how many chunks were pulled and whether close() was called.""" + + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.pulled = 0 + self.closed = False + + def __iter__(self): + for chunk in self._chunks: + self.pulled += 1 + yield chunk + + def close(self): + self.closed = True + + +class _CountingAsyncStream(AsyncByteStream): + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.pulled = 0 + self.closed = False + + async def __aiter__(self): + for chunk in self._chunks: + self.pulled += 1 + yield chunk + + async def aclose(self): + self.closed = True + + +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()) + + +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_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_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(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + assert frames.dollar_cost() == 0.002145 + assert frames.ends_with_done() is False + + def test_frame_split_across_chunks(self): + frame = event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) + frames = EventStreamCostFrames() + for i in range(0, len(frame), 7): + frames.feed(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", {}) + + event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) + ) + assert frames.dollar_cost() == 0.002145 + + def test_last_frame_not_cost_reads_none(self): + frames = EventStreamCostFrames() + frames.feed(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + 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(bedrock_chunk({"type": "message_stop", "amazon-bedrock-invocationMetrics": {}})) + 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_buffering_and_warns(self, caplog): + import logging + + frames = EventStreamCostFrames() + with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + frames.feed(b"\xff\xff\xff\xff" + b"\x00" * 12) + frames.feed(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + assert frames.dollar_cost() is None + assert frames._buffer == b"" + assert "Lost AWS event-stream framing" in caplog.text + + +class TestDollarCostSyncStream: + def test_full_read_captures_trailing_cost_frame(self): + inner = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + assert b"".join(_sse_stream(inner)) == b"".join([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + assert get_captured_dollar_cost() == 0.002145 + + def test_bytes_pass_through_unchanged(self): + chunks = [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] + assert list(_sse_stream(_CountingSyncStream(chunks))) == chunks + + 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 = _CountingSyncStream([*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 = _CountingSyncStream([*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_early_close_mid_stream_does_not_drain(self): + """Cancelling mid-stream must not keep downloading the model's output.""" + inner = _CountingSyncStream([*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_cost_frame_split_across_chunks(self): + inner = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME[:10], COST_FRAME[10:]]) + list(_sse_stream(inner)) + assert get_captured_dollar_cost() == 0.002145 + + def test_unpriced_stream_reads_none(self): + list(_sse_stream(_CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT]))) + assert get_captured_dollar_cost() is None + + def test_close_without_reading_only_closes_inner(self): + inner = _CountingSyncStream([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(_CountingSyncStream): + 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_full_read_captures_cost(self): + chunks = [ + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame("costMetadata", {"associated_dollar_cost": 0.002145}), + ] + inner = _CountingSyncStream(chunks) + assert list(DollarCostSyncStream(inner, EventStreamCostFrames())) == chunks + assert get_captured_dollar_cost() == 0.002145 + + def test_event_stream_close_after_terminal_frame_drains_the_cost_frame(self): + """langchain-aws's invoke adapter stops at the vendor's stop event and closes.""" + stop = bedrock_chunk({"type": "message_stop", "amazon-bedrock-invocationMetrics": {}}) + inner = _CountingSyncStream( + [stop, event_frame("costMetadata", {"associated_dollar_cost": 0.002145})] + ) + 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_trailing_cost_frame(self): + inner = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + chunks = [chunk async for chunk in DollarCostAsyncStream(inner, SseCostFrames())] + assert chunks == [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] + assert get_captured_dollar_cost() == 0.002145 + + async def test_close_after_done_drains_the_cost_frame(self): + inner = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = DollarCostAsyncStream(inner, SseCostFrames()) + 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_early_close_mid_stream_does_not_drain(self): + inner = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = DollarCostAsyncStream(inner, SseCostFrames()) + 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 = _CountingAsyncStream([COST_FRAME]) + await DollarCostAsyncStream(inner, SseCostFrames()).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(_CountingAsyncStream): + async def __aiter__(self): + yield DONE_EVENT + raise ConnectionError("boom") + + inner = _Failing([]) + stream = DollarCostAsyncStream(inner, SseCostFrames()) + async for chunk in stream: + if DONE_EVENT in chunk: + break + await stream.aclose() + assert inner.closed + assert get_captured_dollar_cost() is None + + async def test_event_stream_full_read_captures_cost(self): + chunks = [event_frame("costMetadata", {"associated_dollar_cost": 0.002145})] + inner = _CountingAsyncStream(chunks) + assert [c async for c in DollarCostAsyncStream(inner, EventStreamCostFrames())] == chunks + assert get_captured_dollar_cost() == 0.002145 + + +class TestAttachStreamingDollarCostCapture: + def test_wraps_sync_sse_stream(self): + response = Response(200, headers={"content-type": "text/event-stream; charset=utf-8"}) + response.stream = _CountingSyncStream([]) + attach_streaming_dollar_cost_capture(response) + assert isinstance(response.stream, DollarCostSyncStream) + + def test_wraps_async_sse_stream(self): + response = Response(200, headers={"content-type": "text/event-stream"}) + response.stream = _CountingAsyncStream([]) + attach_streaming_dollar_cost_capture(response) + assert isinstance(response.stream, DollarCostAsyncStream) + + def test_wraps_aws_event_stream(self): + response = Response(200, headers={"content-type": "application/vnd.amazon.eventstream"}) + response.stream = _CountingSyncStream([]) + attach_streaming_dollar_cost_capture(response) + assert isinstance(response.stream, DollarCostSyncStream) + + def test_leaves_other_content_types_alone(self): + inner = _CountingSyncStream([]) + 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(self, caplog): + """Compressed raw bytes cannot be parsed; warn instead of reading "not priced".""" + import logging + + inner = _CountingSyncStream([]) + response = Response( + 200, headers={"content-type": "text/event-stream", "content-encoding": "gzip"} + ) + response.stream = inner + with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + 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 91d00683..cec3c7ad 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 00000000..12b32f59 --- /dev/null +++ b/tests/langchain/clients/bedrock/test_dollar_cost.py @@ -0,0 +1,217 @@ +"""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. +""" + +import os +from unittest.mock import patch + +import httpx +import pytest +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 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}) + + +@pytest.fixture +def llmgw_settings(): + 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", + } + with patch.dict(os.environ, env, clear=True): + return LLMGatewaySettings() + + +def _make_converse_chat( + settings: LLMGatewaySettings, frames: list[bytes] +) -> UiPathChatBedrockConverse: + chat = UiPathChatBedrockConverse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + settings=settings, + model_details={}, + default_headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, + ) + # The wrapped boto client holds this same httpx client. + chat.uipath_sync_client._transport = httpx.MockTransport( # type: ignore[attr-defined] + lambda request: event_stream_response(*frames) + ) + chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] + return chat + + +def test_converse_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_converse_chat(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_converse_chat(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 + + +def test_converse_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None: + chat = _make_converse_chat(llmgw_settings, CONVERSE_STREAM_FRAMES) + chunks = list(chat.stream("Hello")) + assert all("associated_dollar_cost" not in c.response_metadata for c in chunks) + + +# 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 _make_anthropic_bedrock_chat( + settings: LLMGatewaySettings, frames: list[bytes] +) -> UiPathChatAnthropicBedrock: + chat = UiPathChatAnthropicBedrock( + 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)) + chat.uipath_sync_client._transport = transport # type: ignore[attr-defined] + chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] + chat.uipath_async_client._transport = transport # type: ignore[attr-defined] + chat.uipath_async_client._mounts = {} # type: ignore[attr-defined] + return chat + + +def test_anthropic_bedrock_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_anthropic_bedrock_chat(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_anthropic_bedrock_chat(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 _make_invoke_chat(settings: LLMGatewaySettings, frames: list[bytes]) -> UiPathChatBedrock: + chat = UiPathChatBedrock( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + settings=settings, + model_details={}, + default_headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, + ) + chat.uipath_sync_client._transport = httpx.MockTransport( # type: ignore[attr-defined] + lambda request: event_stream_response(*frames) + ) + chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] + return chat + + +def test_invoke_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_invoke_chat(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_invoke_chat_astream_surfaces_gateway_cost(llmgw_settings) -> None: + chat = _make_invoke_chat(llmgw_settings, [*ANTHROPIC_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 + + +def test_invoke_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None: + chat = _make_invoke_chat(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 47189587..df59e0c2 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/features/test_captured_headers.py b/tests/langchain/features/test_captured_headers.py index 04187792..45e9f226 100644 --- a/tests/langchain/features/test_captured_headers.py +++ b/tests/langchain/features/test_captured_headers.py @@ -1,23 +1,30 @@ -"""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 typing import Any from unittest.mock import patch 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.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, + set_captured_dollar_cost, +) from uipath.llm_client.utils.headers import ( extract_matching_headers, get_captured_response_headers, @@ -94,21 +101,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 +134,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) @@ -436,3 +448,150 @@ 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, + ) + chat.uipath_sync_client._transport = MockTransport( # type: ignore[attr-defined] + response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost + ) + chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] + chat.uipath_async_client._transport = MockAsyncTransport( # type: ignore[attr-defined] + response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost + ) + chat.uipath_async_client._mounts = {} # type: ignore[attr-defined] + 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 + + def test_inject_dollar_cost_populates_result(self, llmgw_settings): + chat = _make_normalized_chat(llmgw_settings) + set_captured_dollar_cost(0.002145) + + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, ChatResult + + result = ChatResult( + generations=[ChatGeneration(message=AIMessage(content="test", response_metadata={}))] + ) + chat._inject_dollar_cost(result.generations) + assert result.generations[0].message.response_metadata["associated_dollar_cost"] == 0.002145 + + def test_inject_dollar_cost_skipped_when_absent(self, llmgw_settings): + chat = _make_normalized_chat(llmgw_settings) + + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, ChatResult + + result = ChatResult( + generations=[ChatGeneration(message=AIMessage(content="test", response_metadata={}))] + ) + chat._inject_dollar_cost(result.generations) + assert "associated_dollar_cost" not in result.generations[0].message.response_metadata diff --git a/tests/lazy_stream.py b/tests/lazy_stream.py new file mode 100644 index 00000000..cfc1a7f8 --- /dev/null +++ b/tests/lazy_stream.py @@ -0,0 +1,27 @@ +"""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 + + def __iter__(self) -> Iterator[bytes]: + yield from self._chunks + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + def close(self) -> None: + pass + + async def aclose(self) -> None: + pass From 128681e22362d5ad1810273180b3cda1c85d24db Mon Sep 17 00:00:00 2001 From: Andrei Tava Date: Wed, 9 Sep 2026 14:36:03 +0300 Subject: [PATCH 2/2] test: compact dollar-cost tests and share their helpers Drop tests with no unique failure mode (verified by mutation), fold the counting byte streams into LazyByteStream, use one Bedrock model builder and one llmgw_settings fixture, and move the transport swap into a helper. Add the missing checks for the drain bound and for event-stream dispatch. Remove the dead [DONE] guard in the SSE cost parser. --- .../uipath_langchain_client/base_client.py | 2 + src/uipath/llm_client/utils/dollar_cost.py | 4 +- tests/core/features/test_dollar_cost.py | 252 +++++++----------- .../clients/bedrock/test_dollar_cost.py | 97 ++----- tests/langchain/conftest.py | 22 ++ .../features/test_captured_headers.py | 59 +--- tests/langchain/transports.py | 19 ++ tests/lazy_stream.py | 11 +- 8 files changed, 179 insertions(+), 287 deletions(-) create mode 100644 tests/langchain/transports.py 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 9cc0082d..d6eec3c8 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 @@ -528,6 +528,8 @@ 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(): diff --git a/src/uipath/llm_client/utils/dollar_cost.py b/src/uipath/llm_client/utils/dollar_cost.py index a5431e02..b3560f49 100644 --- a/src/uipath/llm_client/utils/dollar_cost.py +++ b/src/uipath/llm_client/utils/dollar_cost.py @@ -133,9 +133,7 @@ def ends_with_done(self) -> bool: def dollar_cost(self) -> float | None: data = self._last_complete_event_data() - if data is None or data == _SSE_DONE_DATA: - return None - return _dollar_cost_from_json_bytes(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), diff --git a/tests/core/features/test_dollar_cost.py b/tests/core/features/test_dollar_cost.py index 0e9fb2ba..57210d96 100644 --- a/tests/core/features/test_dollar_cost.py +++ b/tests/core/features/test_dollar_cost.py @@ -1,8 +1,11 @@ """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, @@ -14,35 +17,29 @@ extract_associated_dollar_cost, get_captured_dollar_cost, requests_dollar_cost, - set_captured_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) -class TestRequestsDollarCost: - def test_true_header_opts_in(self): - request = Request( - "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"} - ) - assert requests_dollar_cost(request) is True +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): - request = Request( - "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: " True "} - ) - assert requests_dollar_cost(request) is True + assert requests_dollar_cost(_request(" True ")) is True def test_false_header_does_not_opt_in(self): - request = Request( - "POST", "https://example.com", headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "false"} - ) - assert requests_dollar_cost(request) is False + assert requests_dollar_cost(_request("false")) is False def test_missing_header_does_not_opt_in(self): - assert requests_dollar_cost(Request("POST", "https://example.com")) is False + assert requests_dollar_cost(_request(None)) is False class TestExtractAssociatedDollarCost: @@ -66,40 +63,20 @@ def test_malformed_json_returns_none(self): def test_non_object_json_returns_none(self): assert extract_associated_dollar_cost(_json_response(b"[1, 2]")) is None - def test_non_numeric_field_returns_none(self): - response = _json_response(b'{"associated_dollar_cost": "n/a"}') - assert extract_associated_dollar_cost(response) 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_is_logged(self, caplog): + def test_non_numeric_field_returns_none_and_warns(self, caplog): """A present but unusable value must not look like "not priced".""" - import logging - response = _json_response(b'{"associated_dollar_cost": "n/a"}') - with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + 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 - -class TestDollarCostFromBody: - def test_reads_numeric_field(self): + def test_integer_cost_reads_as_float(self): assert dollar_cost_from_body({"associated_dollar_cost": 3}) == 3.0 - def test_non_dict_returns_none(self): - assert dollar_cost_from_body(["associated_dollar_cost"]) is None - - -class TestCapturedDollarCostContextVar: - def test_default_is_none(self): - assert get_captured_dollar_cost() is None - - def test_set_and_get(self): - set_captured_dollar_cost(0.002145) - assert get_captured_dollar_cost() == 0.002145 - # ============================================================================ # Streaming capture @@ -111,38 +88,10 @@ def test_set_and_get(self): b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "lo"}}]}\n\n', ] DONE_EVENT = b"data: [DONE]\n\n" - - -class _CountingSyncStream(SyncByteStream): - """Records how many chunks were pulled and whether close() was called.""" - - def __init__(self, chunks: list[bytes]): - self._chunks = chunks - self.pulled = 0 - self.closed = False - - def __iter__(self): - for chunk in self._chunks: - self.pulled += 1 - yield chunk - - def close(self): - self.closed = True - - -class _CountingAsyncStream(AsyncByteStream): - def __init__(self, chunks: list[bytes]): - self._chunks = chunks - self.pulled = 0 - self.closed = False - - async def __aiter__(self): - for chunk in self._chunks: - self.pulled += 1 - yield chunk - - async def aclose(self): - self.closed = True +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: @@ -156,12 +105,22 @@ 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') @@ -178,7 +137,7 @@ def test_event_and_comment_lines_are_skipped(self): frames.feed(b': keep-alive\nevent: cost\ndata: {"associated_dollar_cost": 1}\n\n') assert frames.dollar_cost() == 1.0 - def test_done_is_not_a_cost(self): + def test_done_is_terminal_and_not_a_cost(self): frames = SseCostFrames() frames.feed(DONE_EVENT) assert frames.ends_with_done() @@ -189,28 +148,24 @@ class TestEventStreamCostFrames: def test_cost_in_last_frame(self): frames = EventStreamCostFrames() frames.feed(event_frame("messageStop", {"stopReason": "end_turn"})) - frames.feed(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + 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): - frame = event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) frames = EventStreamCostFrames() - for i in range(0, len(frame), 7): - frames.feed(frame[i : i + 7]) + 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", {}) - + event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) - ) + 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(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + frames.feed(EVENTSTREAM_COST_FRAME) frames.feed(event_frame("messageStop", {"associated_dollar_cost": 0.002145})) assert frames.dollar_cost() is None @@ -221,7 +176,7 @@ def test_converse_metadata_event_is_terminal(self): def test_invoke_chunk_with_invocation_metrics_is_terminal(self): frames = EventStreamCostFrames() - frames.feed(bedrock_chunk({"type": "message_stop", "amazon-bedrock-invocationMetrics": {}})) + frames.feed(EVENTSTREAM_STOP_FRAME) assert frames.ends_with_done() is True def test_ordinary_frames_are_not_terminal(self): @@ -231,31 +186,24 @@ def test_ordinary_frames_are_not_terminal(self): frames.feed(event_frame("chunk", {"bytes": "not base64!"})) assert frames.ends_with_done() is False - def test_lost_framing_stops_buffering_and_warns(self, caplog): - import logging - + def test_lost_framing_stops_tracking_and_warns(self, caplog): frames = EventStreamCostFrames() - with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + with caplog.at_level(logging.WARNING, logger=LOGGER): frames.feed(b"\xff\xff\xff\xff" + b"\x00" * 12) - frames.feed(event_frame("costMetadata", {"associated_dollar_cost": 0.002145})) + frames.feed(EVENTSTREAM_COST_FRAME) assert frames.dollar_cost() is None - assert frames._buffer == b"" assert "Lost AWS event-stream framing" in caplog.text class TestDollarCostSyncStream: - def test_full_read_captures_trailing_cost_frame(self): - inner = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) - assert b"".join(_sse_stream(inner)) == b"".join([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) - assert get_captured_dollar_cost() == 0.002145 - - def test_bytes_pass_through_unchanged(self): + def test_full_read_captures_cost_and_passes_chunks_through(self): chunks = [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] - assert list(_sse_stream(_CountingSyncStream(chunks))) == chunks + 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 = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) stream = _sse_stream(inner) _read_until_done(stream) assert get_captured_dollar_cost() is None @@ -265,16 +213,26 @@ def test_close_after_done_drains_the_cost_frame(self): assert inner.closed def test_close_after_done_without_cost_frame_reads_none(self): - inner = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT]) + 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 = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) stream = _sse_stream(inner) next(iter(stream)) stream.close() @@ -282,24 +240,15 @@ def test_early_close_mid_stream_does_not_drain(self): assert get_captured_dollar_cost() is None assert inner.closed - def test_cost_frame_split_across_chunks(self): - inner = _CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME[:10], COST_FRAME[10:]]) - list(_sse_stream(inner)) - assert get_captured_dollar_cost() == 0.002145 - - def test_unpriced_stream_reads_none(self): - list(_sse_stream(_CountingSyncStream([*CONTENT_EVENTS, DONE_EVENT]))) - assert get_captured_dollar_cost() is None - def test_close_without_reading_only_closes_inner(self): - inner = _CountingSyncStream([COST_FRAME]) + 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(_CountingSyncStream): + class _Failing(LazyByteStream): def __iter__(self): yield DONE_EVENT raise ConnectionError("boom") @@ -311,21 +260,9 @@ def __iter__(self): assert inner.closed assert get_captured_dollar_cost() is None - def test_event_stream_full_read_captures_cost(self): - chunks = [ - event_frame("messageStop", {"stopReason": "end_turn"}), - event_frame("costMetadata", {"associated_dollar_cost": 0.002145}), - ] - inner = _CountingSyncStream(chunks) - assert list(DollarCostSyncStream(inner, EventStreamCostFrames())) == chunks - assert get_captured_dollar_cost() == 0.002145 - def test_event_stream_close_after_terminal_frame_drains_the_cost_frame(self): - """langchain-aws's invoke adapter stops at the vendor's stop event and closes.""" - stop = bedrock_chunk({"type": "message_stop", "amazon-bedrock-invocationMetrics": {}}) - inner = _CountingSyncStream( - [stop, event_frame("costMetadata", {"associated_dollar_cost": 0.002145})] - ) + """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() @@ -335,15 +272,14 @@ def test_event_stream_close_after_terminal_frame_drains_the_cost_frame(self): class TestDollarCostAsyncStream: - async def test_full_read_captures_trailing_cost_frame(self): - inner = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) - chunks = [chunk async for chunk in DollarCostAsyncStream(inner, SseCostFrames())] - assert chunks == [*CONTENT_EVENTS, DONE_EVENT, COST_FRAME] + 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 = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) - stream = DollarCostAsyncStream(inner, SseCostFrames()) + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_astream(inner) async for chunk in stream: if DONE_EVENT in chunk: break @@ -352,9 +288,18 @@ async def test_close_after_done_drains_the_cost_frame(self): 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 = _CountingAsyncStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) - stream = DollarCostAsyncStream(inner, SseCostFrames()) + inner = LazyByteStream([*CONTENT_EVENTS, DONE_EVENT, COST_FRAME]) + stream = _sse_astream(inner) async for _ in stream: break await stream.aclose() @@ -362,20 +307,20 @@ async def test_early_close_mid_stream_does_not_drain(self): assert get_captured_dollar_cost() is None async def test_close_without_reading_only_closes_inner(self): - inner = _CountingAsyncStream([COST_FRAME]) - await DollarCostAsyncStream(inner, SseCostFrames()).aclose() + 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(_CountingAsyncStream): + class _Failing(LazyByteStream): async def __aiter__(self): yield DONE_EVENT raise ConnectionError("boom") inner = _Failing([]) - stream = DollarCostAsyncStream(inner, SseCostFrames()) + stream = _sse_astream(inner) async for chunk in stream: if DONE_EVENT in chunk: break @@ -383,49 +328,48 @@ async def __aiter__(self): assert inner.closed assert get_captured_dollar_cost() is None - async def test_event_stream_full_read_captures_cost(self): - chunks = [event_frame("costMetadata", {"associated_dollar_cost": 0.002145})] - inner = _CountingAsyncStream(chunks) - assert [c async for c in DollarCostAsyncStream(inner, EventStreamCostFrames())] == chunks - assert get_captured_dollar_cost() == 0.002145 - class TestAttachStreamingDollarCostCapture: - def test_wraps_sync_sse_stream(self): + def test_sse_stream_captures_cost(self): response = Response(200, headers={"content-type": "text/event-stream; charset=utf-8"}) - response.stream = _CountingSyncStream([]) + 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"" - def test_wraps_async_sse_stream(self): response = Response(200, headers={"content-type": "text/event-stream"}) - response.stream = _CountingAsyncStream([]) + response.stream = _AsyncOnly() attach_streaming_dollar_cost_capture(response) assert isinstance(response.stream, DollarCostAsyncStream) - def test_wraps_aws_event_stream(self): + def test_aws_event_stream_captures_cost(self): response = Response(200, headers={"content-type": "application/vnd.amazon.eventstream"}) - response.stream = _CountingSyncStream([]) + response.stream = LazyByteStream([EVENTSTREAM_STOP_FRAME, EVENTSTREAM_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_leaves_other_content_types_alone(self): - inner = _CountingSyncStream([]) + 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(self, caplog): + def test_leaves_compressed_streams_alone_and_warns(self, caplog): """Compressed raw bytes cannot be parsed; warn instead of reading "not priced".""" - import logging - - inner = _CountingSyncStream([]) + inner = LazyByteStream([]) response = Response( 200, headers={"content-type": "text/event-stream", "content-encoding": "gzip"} ) response.stream = inner - with caplog.at_level(logging.WARNING, logger="uipath.llm_client.utils.dollar_cost"): + 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/langchain/clients/bedrock/test_dollar_cost.py b/tests/langchain/clients/bedrock/test_dollar_cost.py index 12b32f59..56115e0c 100644 --- a/tests/langchain/clients/bedrock/test_dollar_cost.py +++ b/tests/langchain/clients/bedrock/test_dollar_cost.py @@ -3,11 +3,9 @@ Frames arrive one per chunk, the shape in which the gateway's trailing frame shows up. """ -import os -from unittest.mock import patch +from typing import Any import httpx -import pytest from uipath_langchain_client.clients.bedrock.chat_models import ( UiPathChatAnthropicBedrock, UiPathChatBedrock, @@ -15,6 +13,7 @@ ) 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 @@ -34,39 +33,22 @@ COST_FRAME = event_frame("costMetadata", {"associated_dollar_cost": 0.002145}) -@pytest.fixture -def llmgw_settings(): - 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", - } - with patch.dict(os.environ, env, clear=True): - return LLMGatewaySettings() - - -def _make_converse_chat( - settings: LLMGatewaySettings, frames: list[bytes] -) -> UiPathChatBedrockConverse: - chat = UiPathChatBedrockConverse( +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"}, ) - # The wrapped boto client holds this same httpx client. - chat.uipath_sync_client._transport = httpx.MockTransport( # type: ignore[attr-defined] - lambda request: event_stream_response(*frames) - ) - chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] + 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_converse_chat(llmgw_settings, [*CONVERSE_STREAM_FRAMES, COST_FRAME]) + 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 @@ -76,19 +58,15 @@ def test_converse_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: 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_converse_chat(llmgw_settings, [*CONVERSE_STREAM_FRAMES, COST_FRAME]) + 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 -def test_converse_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None: - chat = _make_converse_chat(llmgw_settings, CONVERSE_STREAM_FRAMES) - chunks = list(chat.stream("Hello")) - assert all("associated_dollar_cost" not in c.response_metadata for c in chunks) - - # UiPathChatAnthropicBedrock: the anthropic SDK's own event-stream decoder drops unknown # events, so the cost must come from the raw bytes. @@ -141,25 +119,10 @@ def test_converse_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None ] -def _make_anthropic_bedrock_chat( - settings: LLMGatewaySettings, frames: list[bytes] -) -> UiPathChatAnthropicBedrock: - chat = UiPathChatAnthropicBedrock( - 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)) - chat.uipath_sync_client._transport = transport # type: ignore[attr-defined] - chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] - chat.uipath_async_client._transport = transport # type: ignore[attr-defined] - chat.uipath_async_client._mounts = {} # type: ignore[attr-defined] - return chat - - def test_anthropic_bedrock_stream_surfaces_gateway_cost(llmgw_settings) -> None: - chat = _make_anthropic_bedrock_chat(llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME]) + 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] @@ -168,7 +131,9 @@ def test_anthropic_bedrock_stream_surfaces_gateway_cost(llmgw_settings) -> None: async def test_anthropic_bedrock_astream_surfaces_gateway_cost(llmgw_settings) -> None: - chat = _make_anthropic_bedrock_chat(llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME]) + 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] @@ -180,22 +145,8 @@ async def test_anthropic_bedrock_astream_surfaces_gateway_cost(llmgw_settings) - # EOF, so the cost frame must be drained on close. -def _make_invoke_chat(settings: LLMGatewaySettings, frames: list[bytes]) -> UiPathChatBedrock: - chat = UiPathChatBedrock( - model="anthropic.claude-haiku-4-5-20251001-v1:0", - settings=settings, - model_details={}, - default_headers={INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER: "true"}, - ) - chat.uipath_sync_client._transport = httpx.MockTransport( # type: ignore[attr-defined] - lambda request: event_stream_response(*frames) - ) - chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] - return chat - - def test_invoke_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: - chat = _make_invoke_chat(llmgw_settings, [*ANTHROPIC_STREAM_FRAMES, COST_FRAME]) + 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] @@ -203,15 +154,7 @@ def test_invoke_chat_stream_surfaces_gateway_cost(llmgw_settings) -> None: assert priced[0].response_metadata["associated_dollar_cost"] == 0.002145 -async def test_invoke_chat_astream_surfaces_gateway_cost(llmgw_settings) -> None: - chat = _make_invoke_chat(llmgw_settings, [*ANTHROPIC_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 - - def test_invoke_chat_stream_omits_cost_when_not_priced(llmgw_settings) -> None: - chat = _make_invoke_chat(llmgw_settings, ANTHROPIC_STREAM_FRAMES) + 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/conftest.py b/tests/langchain/conftest.py index 7461ef59..f4478e18 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 45e9f226..d54ba4dd 100644 --- a/tests/langchain/features/test_captured_headers.py +++ b/tests/langchain/features/test_captured_headers.py @@ -5,15 +5,14 @@ """ import json -import os from typing import Any -from unittest.mock import patch 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, @@ -23,7 +22,6 @@ from uipath.llm_client.settings.utils import SingletonMeta from uipath.llm_client.utils.dollar_cost import ( INCLUDE_ASSOCIATED_DOLLAR_COST_HEADER, - set_captured_dollar_cost, ) from uipath.llm_client.utils.headers import ( extract_matching_headers, @@ -35,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", @@ -150,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, @@ -470,14 +453,15 @@ def _make_passthrough_chat( default_headers=default_headers, **model_kwargs, ) - chat.uipath_sync_client._transport = MockTransport( # type: ignore[attr-defined] - response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost - ) - chat.uipath_sync_client._mounts = {} # type: ignore[attr-defined] - chat.uipath_async_client._transport = MockAsyncTransport( # type: ignore[attr-defined] - response_json=response_json, stream_chunks=STREAM_CHUNKS, stream_cost=stream_cost + 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 + ), ) - chat.uipath_async_client._mounts = {} # type: ignore[attr-defined] return chat @@ -570,28 +554,3 @@ def test_invoke_routed_through_streaming_surfaces_cost(self, llmgw_settings): result = chat.invoke("Hello") assert result.content == "Hello!" assert result.response_metadata["associated_dollar_cost"] == 0.002145 - - def test_inject_dollar_cost_populates_result(self, llmgw_settings): - chat = _make_normalized_chat(llmgw_settings) - set_captured_dollar_cost(0.002145) - - from langchain_core.messages import AIMessage - from langchain_core.outputs import ChatGeneration, ChatResult - - result = ChatResult( - generations=[ChatGeneration(message=AIMessage(content="test", response_metadata={}))] - ) - chat._inject_dollar_cost(result.generations) - assert result.generations[0].message.response_metadata["associated_dollar_cost"] == 0.002145 - - def test_inject_dollar_cost_skipped_when_absent(self, llmgw_settings): - chat = _make_normalized_chat(llmgw_settings) - - from langchain_core.messages import AIMessage - from langchain_core.outputs import ChatGeneration, ChatResult - - result = ChatResult( - generations=[ChatGeneration(message=AIMessage(content="test", response_metadata={}))] - ) - chat._inject_dollar_cost(result.generations) - assert "associated_dollar_cost" not in result.generations[0].message.response_metadata diff --git a/tests/langchain/transports.py b/tests/langchain/transports.py new file mode 100644 index 00000000..1ca5ed39 --- /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 index cfc1a7f8..6490e3cf 100644 --- a/tests/lazy_stream.py +++ b/tests/lazy_stream.py @@ -12,16 +12,21 @@ class LazyByteStream(SyncByteStream, AsyncByteStream): def __init__(self, chunks: list[bytes]): self._chunks = chunks + self.pulled = 0 + self.closed = False def __iter__(self) -> Iterator[bytes]: - yield from self._chunks + 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: - pass + self.closed = True async def aclose(self) -> None: - pass + self.closed = True