From ff22f3de068d51fb370e68e5f2a65d36907ef619 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:51:53 -0700 Subject: [PATCH 01/58] build(deps): bump langgraph and tornado and add dated scan ignores to clear osv-scan and image-scan (#48302) * build(deps): suppress filelock Windows advisories and drop stale certifi entry to clear osv-scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): suppress dev-only msal-node advisory in vscode-extension until azure identity allows 7.1 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): bump ddtrace and msgpack to clear osv-scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): bump pylint to 4.0.9 to clear osv-scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): lock dnspython 2.9.0 on Python 3.11+ and defer 3.10 to clear osv-scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): bump qs in vscode-extension to clear osv-scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): extend the pillow 12.3.0 override for Python 3.10 to 2026-11-27 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): extend the diskcache ignore, drop the withdrawn braces entry, and name the filelock cap in its reasons * build(deps): bump langgraph and add dated ignores for the 2026-10-28 filelock advisories * build(deps): move tornado to 6.6 on Python 3.11+, name huggingface_hub's filelock use, and scope a grype rule to the filelock High * build(deps): say which packages import tornado in its ignore reasons --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .grype.yaml | 9 +++++ osv-scanner.toml | 103 ++++++++++++++++++++++++++++++++++++++++------- pyproject.toml | 2 + uv.lock | 84 +++++++++++++++++++++++++++++++++----- 4 files changed, 173 insertions(+), 25 deletions(-) diff --git a/.grype.yaml b/.grype.yaml index c5e49851dc9..dedc8fd528e 100644 --- a/.grype.yaml +++ b/.grype.yaml @@ -11,3 +11,12 @@ ignore: package: name: zlib type: apk + - vulnerability: GHSA-v4x8-3wr5-3wq8 + reason: >- + filelock ReadWriteLock advisory, and no package in the image takes ReadWriteLock; + the fix (4.2.1) is blocked by litellm's own filelock<4.0 cap until #48635 or #48802 lifts it, + and the image then moves off 3.32.7, which retires this version-scoped rule + package: + name: filelock + version: 3.32.7 + type: python diff --git a/osv-scanner.toml b/osv-scanner.toml index 33e2f614edc..32a4f65fdf9 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -1,27 +1,102 @@ -[[IgnoredVulns]] -id = "GHSA-w8v5-vhqr-4h9v" -ignoreUntil = 2026-11-01 -reason = "diskcache has no fixed release published; remove this entry once one exists" - [[IgnoredVulns]] id = "GHSA-v4x8-3wr5-3wq8" ignoreUntil = 2026-11-20 -reason = "the only fixed filelock is 4.2.1, a new major that requires Python >= 3.11 while litellm still supports 3.10, and the advisory is in ReadWriteLock, which litellm never uses (only FileLock); bump filelock and remove this entry once 3.10 is dropped or a 3.x fix ships" +reason = "filelock ReadWriteLock advisory, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" [[IgnoredVulns]] -id = "GHSA-vfj7-8cjw-p6xm" -ignoreUntil = 2026-11-03 -reason = "braces has no fixed release published (3.0.3, 2024-05-21, is latest and last_affected); dev-only via knip > fast-glob > micromatch; remove this entry once one exists" +id = "GHSA-567q-5qvw-pf92" +ignoreUntil = 2026-11-20 +reason = "Windows-only local DoS (CVSS 2.5) in filelock soft locks; litellm never takes SoftFileLock, but huggingface_hub falls back to it in the Hugging Face cache when the filesystem has no flock, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting, and the attack needs another local user who can write that cache; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" [[IgnoredVulns]] -id = "GHSA-6q48-6c7w-9gw2" -ignoreUntil = 2026-10-23 -reason = "certifi 2026.10.19 (published 2026-10-20) is inside the 3-day uv exclude-newer window until 2026-10-23 02:24 UTC; bump certifi and remove this entry once it clears" +id = "GHSA-9vhv-jxww-xfm5" +ignoreUntil = 2026-11-20 +reason = "Windows-only local DoS (CVSS 3.3) in filelock ReadWriteLock, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-6627-m3xw-93mp" +ignoreUntil = 2026-11-20 +reason = "filelock StrictSoftFileLock claim directories ignore the lock mode, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-69hj-f68j-2879" +ignoreUntil = 2026-11-20 +reason = "filelock UnixFileLock truncates and chmods a file another user renames or hard-links in at the lock path; litellm's CLI takes FileLock only inside ~/.litellm, which it creates owner-only, and huggingface_hub takes it in the Hugging Face cache, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting; the attack needs another local user who can rename or link files in that cache; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-8mx5-7m5m-r77g" +ignoreUntil = 2026-11-20 +reason = "filelock WindowsFileLock ignores mode; the proxy image is Linux, so only Windows installs run it; there litellm's CLI takes FileLock only inside the owner-only ~/.litellm, whose lock file holds no data, and huggingface_hub takes it in the Hugging Face cache, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-9c4x-gm8v-wgc4" +ignoreUntil = 2026-11-20 +reason = "filelock StrictSoftFileLock and SoftReadWriteLock adopt another user's protocol directory, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-9j6q-wvj5-jgjg" +ignoreUntil = 2026-11-20 +reason = "filelock SoftFileLock symlink swap leaks whether a path exists; litellm never takes SoftFileLock, but huggingface_hub falls back to it in the Hugging Face cache when the filesystem has no flock, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting, and the attack needs another local user who can write that cache; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-cwqr-hh33-gh44" +ignoreUntil = 2026-11-20 +reason = "filelock ReadWriteLock stalls on a crafted lock database, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-fq56-3cx3-fjxp" +ignoreUntil = 2026-11-20 +reason = "filelock SoftFileLock stale-lock break removes a live lock on a reused inode; litellm never takes SoftFileLock, but huggingface_hub falls back to it in the Hugging Face cache when the filesystem has no flock, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting, and the attack needs another local user who can write that cache; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-j8f7-rjxc-mr56" +ignoreUntil = 2026-11-20 +reason = "filelock ReadWriteLock follows a symlink into SQLite, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed in 4.0.9, which litellm's own filelock<4.0 cap blocks on every Python until it is lifted (#48635, #48802); remove once that lands" + +[[IgnoredVulns]] +id = "GHSA-m462-7rxw-wvpg" +ignoreUntil = 2026-11-20 +reason = "filelock StrictSoftFileLock follows a symlink swapped in at .filelock, which no package in the lock takes (litellm and huggingface_hub, filelock's only dependents, use just FileLock and SoftFileLock); fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-rfpc-qcw5-3pc4" +ignoreUntil = 2026-11-20 +reason = "filelock SoftFileLock removes a live holder's marker when another user forges its PID; litellm never takes SoftFileLock, but huggingface_hub falls back to it in the Hugging Face cache when the filesystem has no flock, which litellm reaches only to download the Kimi K3 tokenizer for Bedrock token counting, and the attack needs another local user who can write that cache; fixed only in 4.2.1, and filelock 4.1.1+ needs Python >= 3.11; litellm's own filelock<4.0 cap keeps every Python on 3.32.7 until it is lifted for 3.11+ (#48635, #48802); remove once that lands and 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-pcg2-g6cj-q92c" +ignoreUntil = 2026-11-27 +reason = "tornado 6.6 (published 2026-10-21) is the only fix and requires Python >= 3.11, so the Python 3.10 resolution stays on 6.5.10 while 3.11+ locks 6.6; litellm starts no tornado server or client, since semantic-router and aurelio-sdk only declare it and tenacity and sentry-sdk import it only to wrap tornado code when it is present, so tornado's HTTP header parsing is never reached; remove this entry once 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-pw9p-v7r6-gx4p" +ignoreUntil = 2026-11-27 +reason = "tornado 6.6 (published 2026-10-21) is the only fix and requires Python >= 3.11, so the Python 3.10 resolution stays on 6.5.10 while 3.11+ locks 6.6; litellm starts no tornado server or client, since semantic-router and aurelio-sdk only declare it and tenacity and sentry-sdk import it only to wrap tornado code when it is present, so tornado's RedirectHandler is never reached; remove this entry once 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-r4j5-9xpp-3985" +ignoreUntil = 2026-11-27 +reason = "tornado 6.6 (published 2026-10-21) is the only fix and requires Python >= 3.11, so the Python 3.10 resolution stays on 6.5.10 while 3.11+ locks 6.6; litellm starts no tornado server or client, since semantic-router and aurelio-sdk only declare it and tenacity and sentry-sdk import it only to wrap tornado code when it is present, so tornado's chunked body reading is never reached; remove this entry once 3.10 is dropped" + +[[IgnoredVulns]] +id = "GHSA-w8v5-vhqr-4h9v" +ignoreUntil = 2026-11-27 +reason = "diskcache still has no fixed release published (checked 2026-10-28), and it is reached only through the optional caching extra; remove this entry once one exists" + +[[IgnoredVulns]] +id = "GHSA-45c3-73f7-pv4m" +ignoreUntil = 2026-11-27 +reason = "dnspython 2.9.0 (published 2026-10-09) is the only fix and requires Python >= 3.11, so the Python 3.10 resolution stays on 2.8.0 while 3.11+ locks 2.9.0; the advisory is CPU use in dns.name.from_wire_parser and litellm only reaches dnspython through pydantic EmailStr, which never does DNS lookups; remove this entry once 3.10 is dropped or a 3.10-compatible fix ships" + +[[IgnoredVulns]] +id = "GHSA-xcx6-jx8p-x3h5" +ignoreUntil = 2026-11-24 +reason = "@azure/msal-node is fixed only in 7.1.0 (published 2026-10-13), but it is dev-only in vscode-extension via @vscode/vsce > @azure/identity, and the latest @azure/identity (4.14.0) still requires msal-node ^6.0.0, so the lock cannot move without forcing a major override; the advisory is in the auth_code flow, which extension packaging does not run; bump and remove this entry once @azure/identity accepts msal-node >= 7.1.0" [[PackageOverrides]] name = "pillow" version = "12.3.0" ecosystem = "PyPI" vulnerability.ignore = true -effectiveUntil = 2026-10-27 -reason = "pillow 13.0.0 requires Python >= 3.11, so the Python 3.10 resolution of the ci and dev groups stays on 12.3.0; 3.11+ is on 13.0.0" +effectiveUntil = 2026-11-27 +reason = "pillow 13.0.0 requires Python >= 3.11, so the Python 3.10 resolution of the ci and dev groups and of the mlflow extra stays on 12.3.0; 3.11+ is on 13.0.0; pillow has never shipped an older-major release after the next major (no 10.x after 11.0.0, no 11.x after 12.0.0), so 13.x is the only fix line; extended to 2026-11-27 with maintainer approval, same date as the dnspython 3.10 entry; remove once 3.10 is dropped" diff --git a/pyproject.toml b/pyproject.toml index 00cdb1587aa..200bb463313 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -360,11 +360,13 @@ cache-keys = [ ] constraint-dependencies = [ "tornado>=6.5.8", + "tornado>=6.6; python_version >= '3.11'", "aiohttp>=3.14.2,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", "httplib2>=0.32.0", "setuptools>=84.1.0", + "dnspython>=2.9.0; python_version >= '3.11'", ] override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. diff --git a/uv.lock b/uv.lock index b40490112cf..5be97af9ed2 100644 --- a/uv.lock +++ b/uv.lock @@ -25,11 +25,13 @@ members = [ ] constraints = [ { name = "aiohttp", specifier = ">=3.14.2,<4.0" }, + { name = "dnspython", marker = "python_full_version >= '3.11'", specifier = ">=2.9.0" }, { name = "httplib2", specifier = ">=0.32.0" }, { name = "packaging", specifier = ">=24.0" }, { name = "setuptools", specifier = ">=84.1.0" }, { name = "soupsieve", specifier = ">=2.8.4" }, { name = "tornado", specifier = ">=6.5.8" }, + { name = "tornado", marker = "python_full_version >= '3.11'", specifier = ">=6.6" }, ] overrides = [ { name = "cryptography", specifier = ">=50.0.0,<51.0" }, @@ -579,7 +581,8 @@ dependencies = [ { name = "python-dotenv" }, { name = "requests" }, { name = "requests-toolbelt" }, - { name = "tornado" }, + { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.14'" }, + { name = "tornado", version = "6.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/27/0e/c2e369ad173fb3d76448e46d10beb3dcc53388318933ddf8169a3f21a810/aurelio_sdk-0.0.19.tar.gz", hash = "sha256:14107e7440ff2efd0b4a08c52fb595e7680bd4bc973a0ddfb3b64157c6666b91", size = 15258, upload-time = "2025-03-24T14:37:32.203Z" } wheels = [ @@ -4702,7 +4705,7 @@ wheels = [ [[package]] name = "langgraph" -version = "1.2.6" +version = "1.2.17" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, @@ -4713,22 +4716,22 @@ dependencies = [ { name = "pydantic", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/7a/ea09b05bb0cbddfa43bd34fc581357e87fc3f21a751cc0d419688c3106da/langgraph-1.2.6.tar.gz", hash = "sha256:f9b45a34f13930c94d96cdb76277447ad2cc70ec2d18cd2764d7fdadb36cdc1b", size = 714400, upload-time = "2026-06-18T20:58:21.514Z" } +sdist = { url = "https://files.pythonhosted.org/packages/75/39/38b60bfc6fd064bc172693e19b3a6dccdaf9141bc9b984a9f83cdb06d014/langgraph-1.2.17.tar.gz", hash = "sha256:311516e313abd5424596d6c5250507c89be2e1973cbe813061696728803b0e10", size = 774002, upload-time = "2026-10-24T01:07:50.567Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/89/32/772db1b00a9fe42f50320d1aa20caefb76e621eff1f7218b9918093d631d/langgraph-1.2.6-py3-none-any.whl", hash = "sha256:1cf94d3ca124f84f77ce408fa1b06c3dee680a8aafffe364a8fd5d7d03eb8695", size = 246132, upload-time = "2026-06-18T20:58:20.335Z" }, + { url = "https://files.pythonhosted.org/packages/00/d9/2e5b99319e2e8d39bc44d84836497c2e8f616554bb47843c8c819c8fed28/langgraph-1.2.17-py3-none-any.whl", hash = "sha256:e022a9fea64631fed946557a23f3fa5ac2ff38c72dcb999f426fc77d55ccde60", size = 260489, upload-time = "2026-10-24T01:07:49.155Z" }, ] [[package]] name = "langgraph-checkpoint" -version = "4.1.1" +version = "4.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/69/31fdbdc65a85bbd6178afa193c772bb926620f47b4869638bc2bc80afaaa/langgraph_checkpoint-4.3.0.tar.gz", hash = "sha256:c75965d84cc2c1d549163e910a15bcb577758001b141619d05297c463280b018", size = 182652, upload-time = "2026-10-12T22:26:31.478Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/1f/0c/84747e340bf4f29291c84cdd5733fc8d0a822f3d33bb24e664a18afa4a7c/langgraph_checkpoint-4.3.0-py3-none-any.whl", hash = "sha256:bedfafe2f997ded60e4fa593e79f56f436a6e45586392dc382aa810d0c751c64", size = 58063, upload-time = "2026-10-12T22:26:30.429Z" }, ] [[package]] @@ -4746,7 +4749,7 @@ wheels = [ [[package]] name = "langgraph-sdk" -version = "0.4.4" +version = "0.4.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, @@ -4755,9 +4758,9 @@ dependencies = [ { name = "orjson" }, { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3f/ae/91446c1fffa04a2dc1f81afbfb5bfff3590452891a72bd9098a656ab2657/langgraph_sdk-0.4.4.tar.gz", hash = "sha256:4e651ffa09de695681579396375377bdde23bedbc8e35b070e615cbd5af7da8b", size = 345789, upload-time = "2026-08-27T21:25:22.593Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/ef/e8f63b62dea6ecb75bef9d336f62e3e8f853b17a9943e4727c1bbeb09b11/langgraph_sdk-0.4.9.tar.gz", hash = "sha256:24d775ac1382b94c2595b8eb4eb4e0aa8dde95c905f18e017dc6574e1d79070d", size = 358868, upload-time = "2026-10-23T14:19:05.963Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/43/d5/2ad7f3a835c6fcebf58b6669552536ecd52d2dfe4cf49c6c36648fd83572/langgraph_sdk-0.4.4-py3-none-any.whl", hash = "sha256:39afe416c91742925e6f8a93715f566d499b36e1b636b804a4ffe3190e4f4e64", size = 162270, upload-time = "2026-08-27T21:25:21.072Z" }, + { url = "https://files.pythonhosted.org/packages/57/7f/3102130eab099da9653ebe4a49123087a5889f4070eeee99b402255b1786/langgraph_sdk-0.4.9-py3-none-any.whl", hash = "sha256:c009243f4784758a07287fb2fd21fe47381c3d85f30479df2ae7f345c8703bb1", size = 167613, upload-time = "2026-10-23T14:19:04.82Z" }, ] [[package]] @@ -10228,7 +10231,8 @@ dependencies = [ { name = "pyyaml" }, { name = "regex" }, { name = "tiktoken" }, - { name = "tornado" }, + { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.14'" }, + { name = "tornado", version = "6.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, { name = "urllib3" }, ] sdist = { url = "https://files.pythonhosted.org/packages/dc/a9/1a689e916e8b280f1fd8fb335cc059be626a22fe4533baa045d32fcd6de5/semantic_router-0.1.15.tar.gz", hash = "sha256:328256ddc3c2b713101ec69561d6585aecbf1198ea3461e1486289d8c3a35288", size = 95605, upload-time = "2026-05-23T12:58:15.444Z" } @@ -10920,6 +10924,9 @@ wheels = [ name = "tornado" version = "6.5.10" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/06/61/53d562a57b28c08eda40b258c0f975e360541943ad7c7bef897a40caafda/tornado-6.5.10.tar.gz", hash = "sha256:a6b1ccd08c04b4a06fb5aeb381be99de5ad1e5375c1785e31d78c880feb57687", size = 537910, upload-time = "2026-09-15T13:47:48.73Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/cd/5b/ff5fc58fa2427c30dea74c90053f4fc5eda1e7f3833ed3ecc7147fe2b311/tornado-6.5.10-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9261783640e23258694a9ff0795df430a5a7b0a651d3dd53dd0969ad6be16da7", size = 465883, upload-time = "2026-09-15T13:47:35.463Z" }, @@ -10933,6 +10940,61 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f3/1a/fd497f3a7f7b74bb04f4b94536b5c9f80742b5d50501fd27977652ddec16/tornado-6.5.10-cp39-abi3-win_arm64.whl", hash = "sha256:ce045d3c298fddd30e89a2777f97039d1b641eb9518ac7b26a4721903539c694", size = 467847, upload-time = "2026-09-15T13:47:47.283Z" }, ] +[[package]] +name = "tornado" +version = "6.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/02/5f549b1ad5da6267ed1e1fac2eef69257b1979cc812497002499b7895ff9/tornado-6.6.tar.gz", hash = "sha256:34c4aa916d7f2b2d5cadee2c6850f5a4a7f3d688cbf58399ea9bd3057af4abd5", size = 617998, upload-time = "2026-10-21T15:31:04.733Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/f5/d3e6f76b945f9812842db742d0708a3b6287e850e81be3ca10e290c103c6/tornado-6.6-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:01abf15429d64f04ccfafc238b191cc2c177108b7a331f4461707ad77198c8c2", size = 528326, upload-time = "2026-10-21T15:30:05.528Z" }, + { url = "https://files.pythonhosted.org/packages/30/69/b20cf3a8bf641ec792a5a6050aa153e8f120604f9c23471142de535113e5/tornado-6.6-cp311-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:11bb97c00981eb8254e55bd6a5ad960c983ae343fef8ae3ae5584cd1fb70236a", size = 529544, upload-time = "2026-10-21T15:30:07.126Z" }, + { url = "https://files.pythonhosted.org/packages/0f/f4/d9d8d3ac345aa9df3a132ad03bd443160c2ab24ad4003b4e8d6d5af06371/tornado-6.6-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e152a3b2d3eae113b836d979823f7bb2320f7e334a84449729964921f374a4c", size = 530523, upload-time = "2026-10-21T15:30:08.478Z" }, + { url = "https://files.pythonhosted.org/packages/ad/82/1296ade1c14dee5a4906e5ca092997622ece7c4dc768c68ed9caaf895d1b/tornado-6.6-cp311-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e912d4c4babd04db6f34b25b748174a4a176debe86c8610beedf1815f9af7fdd", size = 530739, upload-time = "2026-10-21T15:30:09.98Z" }, + { url = "https://files.pythonhosted.org/packages/66/30/6ac1eab66c81af46dafd93c9b24bd311f82e199d9d7b8ce2a64a91b6a200/tornado-6.6-cp311-abi3-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:33aabe8a03d426cef4e0fe5669f0643cadd983b54957f1a7ac2789ec326e9644", size = 529738, upload-time = "2026-10-21T15:30:11.498Z" }, + { url = "https://files.pythonhosted.org/packages/24/b4/6c1f6aa1fd148e4d0f500a10cd0c9949e7bf385e161ff497c81de5b3ae4f/tornado-6.6-cp311-abi3-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9aaf0e1c2fb9f046f30e2036b8d2060666399e8464d0a72769e43c754089567b", size = 527908, upload-time = "2026-10-21T15:30:13.298Z" }, + { url = "https://files.pythonhosted.org/packages/54/5b/67d3d452b1179bc9e265cebd381a56ab4cf2d7b2a54607aedd8d6e1ebc2e/tornado-6.6-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:61fc5513cea97720ab9bf54b04827ac0f4500c63f0b072b963d6145c798b4661", size = 530367, upload-time = "2026-10-21T15:30:15.183Z" }, + { url = "https://files.pythonhosted.org/packages/74/14/4e8e9a14f77bcafb2408634b2d4c8c676ad7a43bf8dac4ee2d48370698f9/tornado-6.6-cp311-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:31e1a0d5cd704d2fe8a008c1cc27a0e2b67269ca263630a27985dddc160b5537", size = 530456, upload-time = "2026-10-21T15:30:16.741Z" }, + { url = "https://files.pythonhosted.org/packages/56/12/e3ddef99a2089ccfc86f32bd3f39caaad87bee4868be1bacbf213a0b07a2/tornado-6.6-cp311-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:74416e41e9bdf70fbc2fb242f6fbbc7a58ec11e59d39f15fc5a421ecb6845118", size = 529610, upload-time = "2026-10-21T15:30:18.031Z" }, + { url = "https://files.pythonhosted.org/packages/35/3d/62a692343d46b37e852e5078e0f6513534d925f70a04580072e3c127b129/tornado-6.6-cp311-abi3-musllinux_1_2_s390x.whl", hash = "sha256:8cdd69104ae9d705b3f04d2563aa96e357043e267a455802daae9a1e633a37c8", size = 529803, upload-time = "2026-10-21T15:30:19.326Z" }, + { url = "https://files.pythonhosted.org/packages/49/04/6c59a9413adfe4810f648a4dd9fdf476f13b0495f9cc25ef856b71f9ce1e/tornado-6.6-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:fc7c8a3263501fbe97518cd699101bc3c5f499a7dd4f75fc6cbb3bd9ee1db1f1", size = 529751, upload-time = "2026-10-21T15:30:20.826Z" }, + { url = "https://files.pythonhosted.org/packages/31/36/a1b3b84cd347ae43f982d3bf95628de6c6365668dba2543af47c6a1c888e/tornado-6.6-cp311-abi3-win32.whl", hash = "sha256:b54ca69ff93d0374efb083a12a1e9ec2dd2036e8496d197e0e8ef77b8d9135c3", size = 530643, upload-time = "2026-10-21T15:30:22.428Z" }, + { url = "https://files.pythonhosted.org/packages/b7/50/dc813fe35a672e14ee8bed8b5c634eeb343a05c98eb455318a169de24567/tornado-6.6-cp311-abi3-win_amd64.whl", hash = "sha256:48f26ed9d7beb64e5f06adc864524e581ab6d62ba8ff46eaaa1ac2d0e6d753fd", size = 530886, upload-time = "2026-10-21T15:30:23.975Z" }, + { url = "https://files.pythonhosted.org/packages/0d/71/62a9b9e980026d38874740d0a7c162f53eb0563f6a441e3aff84c48db7e8/tornado-6.6-cp311-abi3-win_arm64.whl", hash = "sha256:29114ffd4e4de869d390393df9e9a0f191065aae1a96f38a58d35e464c04a867", size = 530899, upload-time = "2026-10-21T15:30:25.264Z" }, + { url = "https://files.pythonhosted.org/packages/ac/4c/633e3568eb574df555d6ddd2dac44b7793a5c26cd63b444e653fcfb1d6ec/tornado-6.6-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:63fc8d243b52fd8dd9ed4ad02d0df08f2024f9aea632a433a8a6c91075076d44", size = 528432, upload-time = "2026-10-21T15:30:26.578Z" }, + { url = "https://files.pythonhosted.org/packages/37/eb/b525d1870ad28a56f47b664dd8fd3fab4f7ddd237c87c14e97248bf4762f/tornado-6.6-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d4b85cac795bfc6a7f20ba81871abb8c2ae8fba019061b8a9ae1c2f97e29637", size = 532586, upload-time = "2026-10-21T15:30:28.027Z" }, + { url = "https://files.pythonhosted.org/packages/b7/35/909a33c253352346ca081a27ca66b45c1398468adf784b767597ca773c28/tornado-6.6-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e035c065edadfa696529d05eebcbeb25154c129b5aee6dbfb697a6abf5d4f697", size = 533641, upload-time = "2026-10-21T15:30:29.318Z" }, + { url = "https://files.pythonhosted.org/packages/6f/64/100202666d6e3d96288ab184827762415456c139798a15c962547f9fabf9/tornado-6.6-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5b94e589ab855090c448b5c0937b3a75f1c1a604b7b36207e8236d209f720b9c", size = 533869, upload-time = "2026-10-21T15:30:30.882Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e0/af3ed62bc15830d60328d5b4a1a9076b35ea0ba142cb17b1c1772cb15c59/tornado-6.6-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1a394c8c9db9137f2a2b42869d9f7b91398fcf19da120d0dde4c4546172710f7", size = 532767, upload-time = "2026-10-21T15:30:32.273Z" }, + { url = "https://files.pythonhosted.org/packages/24/d7/1e2e1a6932eb62d63324b64dba2f62e6eaf79ad883dd209352c297b7848c/tornado-6.6-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0215d1df6e4b5e736ddad5adf0af4138cf7437fdd39c77c5e69768e37850a632", size = 530956, upload-time = "2026-10-21T15:30:33.746Z" }, + { url = "https://files.pythonhosted.org/packages/70/65/3b9a0e53d7a4c39f2ebe95486540fa8eeb67bfb6e8c437d603c00d35b128/tornado-6.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:4d711790e0a9576ad924d942c4607e4e401d697588ea6f30f6662824008a7add", size = 533372, upload-time = "2026-10-21T15:30:35.085Z" }, + { url = "https://files.pythonhosted.org/packages/44/48/90979cbfaa3ff941324a2605eb794af0f087c670cd99f18a69ab40498073/tornado-6.6-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:034a66bae0f55421d007803e2385d2a59707d8829b2f026373d42f08256e9922", size = 533425, upload-time = "2026-10-21T15:30:36.425Z" }, + { url = "https://files.pythonhosted.org/packages/bb/c4/ce743a86d91f85c255b77a4282eb5b2716b61e74430eb9d1df8578a539d2/tornado-6.6-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5cb38483b077129056a94dd415a71d29ab4e42b46da894fee8b4657003f10b21", size = 532545, upload-time = "2026-10-21T15:30:37.789Z" }, + { url = "https://files.pythonhosted.org/packages/02/5e/a82131e35778540fc19139229d02aa8cf2453a8022e202be50c88f0f8ac3/tornado-6.6-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:46bbbd107efa37777ac4b83601866d19b4dbd5b770a7a9123582633d95959ad3", size = 532700, upload-time = "2026-10-21T15:30:39.336Z" }, + { url = "https://files.pythonhosted.org/packages/3a/09/6dac8c2c677c65d60b3d7dae94176dbfe392f314e7c8e2d709623edd6a26/tornado-6.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b85e9d0263cdbc027b95947c8b873a7bc7fd6a08d16bfb4f5047350b0808e464", size = 532712, upload-time = "2026-10-21T15:30:40.698Z" }, + { url = "https://files.pythonhosted.org/packages/b2/ac/71e5776cbcdc3264e81aa6db28c69717ba50607a9eb8d5b766f11a4b020e/tornado-6.6-cp314-cp314t-win32.whl", hash = "sha256:4276eef0880e5fe4488e6799c2842f59af3b06637e9313ab4fb506f9c9c2d6cb", size = 530491, upload-time = "2026-10-21T15:30:42.115Z" }, + { url = "https://files.pythonhosted.org/packages/d8/0f/0b6a6908a5b5aa8320009dfdb84434b1cf160b577a00456e59ab9e6ea598/tornado-6.6-cp314-cp314t-win_amd64.whl", hash = "sha256:10d4af98c6df09f8429b1340ccaea3a72887b86f695dfb8056bda517c8115a00", size = 530780, upload-time = "2026-10-21T15:30:43.414Z" }, + { url = "https://files.pythonhosted.org/packages/3c/39/c152e1cc5371f75f5f9f0a4f44bf2dc9396be9407d384de7428446fb4651/tornado-6.6-cp314-cp314t-win_arm64.whl", hash = "sha256:b1f0ef20a5278854f5fdac31637bafadb055405df28b7fc945e54add12b172e3", size = 530781, upload-time = "2026-10-21T15:30:44.741Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c5/f96c7d7a8895745bb37dfb6e6107e1ee40ad02a4b162d7d6c143644455ab/tornado-6.6-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:8f8e382da30fe042f27f27a871e21aefc2613870e5a41648de18446c72d6ea59", size = 528446, upload-time = "2026-10-21T15:30:46.112Z" }, + { url = "https://files.pythonhosted.org/packages/c0/dd/ab8f6500d274b9c55b025798eb54e92591bcff27515799bec443bd54e35a/tornado-6.6-cp315-cp315t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8c7a8a328396c2195454a2ecd9499512441bbce99505544a717e6a87d16ff6f9", size = 532774, upload-time = "2026-10-21T15:30:47.393Z" }, + { url = "https://files.pythonhosted.org/packages/87/23/b78e7d1bf329160b8157284005395421ea941864c16c7f04c3e48c98bcf9/tornado-6.6-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ca025bb0da81c294aef26f27427316df1525220820315f4bf32dc1e84a4d2232", size = 533872, upload-time = "2026-10-21T15:30:48.884Z" }, + { url = "https://files.pythonhosted.org/packages/e0/4f/ab66b614ad0bc9a219484be6bc2c002c8de9916a7b6c7b27fe6273838e5c/tornado-6.6-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:59c22ebc905ac2a267f8d9d8a878f86afbcbbde329d75c485f1422fc45a81b83", size = 534027, upload-time = "2026-10-21T15:30:50.157Z" }, + { url = "https://files.pythonhosted.org/packages/da/63/501fa15a36bace2f4b9158228cf922d606403261bf71c960c433ec3ca0a4/tornado-6.6-cp315-cp315t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f89802bc0bcd15f988d9325b4e6d2171bb04ea68d188732556335e10d3fd00bd", size = 532936, upload-time = "2026-10-21T15:30:51.45Z" }, + { url = "https://files.pythonhosted.org/packages/7f/b7/82d643f635f0300bda98613ca2660b4f3134b6fb594c6fddfa4de706806c/tornado-6.6-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:70043256d5958f9c655066153d23d60ea3e29bf9854f7fb75850e223c477758a", size = 531254, upload-time = "2026-10-21T15:30:52.865Z" }, + { url = "https://files.pythonhosted.org/packages/d4/53/92c40891a94750bc09257310af964bf25c59e3531d6199736c318baa864f/tornado-6.6-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:82b49495e3849d644b841ef33b35491160055997deaa93720ff647ef019421db", size = 533566, upload-time = "2026-10-21T15:30:54.15Z" }, + { url = "https://files.pythonhosted.org/packages/7e/88/b82f5c5efde731b88d839ad99cb6ace2a521f03bb6d7346d2de089c76da8/tornado-6.6-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:07b2bcaa4bf447bdbfb664e4c28344226540db3f0ebccd74886528338a240765", size = 533567, upload-time = "2026-10-21T15:30:55.451Z" }, + { url = "https://files.pythonhosted.org/packages/02/2c/e9a7f5d0fca0bdbcd48b2d9c7884dd0b32669c9945d98bc6d303a80ecfea/tornado-6.6-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:c36b921374c302e54909bae17d9f354061d80e7ae7d8c94f73abbcc217866460", size = 532849, upload-time = "2026-10-21T15:30:56.729Z" }, + { url = "https://files.pythonhosted.org/packages/d0/f5/4ade737340ca6ce9c677ea0c2843591c8ad065d07f0913c28d922c7e078e/tornado-6.6-cp315-cp315t-musllinux_1_2_s390x.whl", hash = "sha256:7412d9e37aa9f7048e4b591f520a95f7a097278ffab15f1ddf3f5c2320cf6648", size = 532873, upload-time = "2026-10-21T15:30:58.067Z" }, + { url = "https://files.pythonhosted.org/packages/8e/04/6acdc937d0bb6189ddb34ceded517d7bd1c7b3bded5d4e1b3442160e6b6f/tornado-6.6-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:15bc95112b2d24b0bfaee9a004e3f974da929e33b4244a7f695a22f45b08a5e5", size = 532883, upload-time = "2026-10-21T15:30:59.395Z" }, + { url = "https://files.pythonhosted.org/packages/ab/27/0a766fac6b1cc629d96346329a0485287fd2857c51a887a3a154f2a69938/tornado-6.6-cp315-cp315t-win32.whl", hash = "sha256:f661c21117b875d612371f3eb243d0c69f56d9bdd6aee38b7c24453419dfe554", size = 530488, upload-time = "2026-10-21T15:31:00.77Z" }, + { url = "https://files.pythonhosted.org/packages/5c/80/3e4142570883da27b3f60ae82c638ac96e60997eac684740613557519c96/tornado-6.6-cp315-cp315t-win_amd64.whl", hash = "sha256:ee8853ab35da510dc952879ebd9995c82bc10ffb7925fd87930dc0d6916f4b39", size = 530777, upload-time = "2026-10-21T15:31:02.087Z" }, + { url = "https://files.pythonhosted.org/packages/17/57/4a85c617b6349962320b4afbe1386d5b05fd57218ee8ad6e82e0110ec9ee/tornado-6.6-cp315-cp315t-win_arm64.whl", hash = "sha256:7de62e2dfdf559fd7154d2b9aa58ff3d9cff8304ae0519341feb082618bc2a33", size = 530780, upload-time = "2026-10-21T15:31:03.425Z" }, +] + [[package]] name = "tqdm" version = "4.67.3" From c068d1350b6092554262c5a70beefc9add1a10fd Mon Sep 17 00:00:00 2001 From: Zabir Raihan Date: Thu, 29 Oct 2026 15:03:00 +0800 Subject: [PATCH 02/58] fix(responses): stream every citation batch before completion (#48747) * fix(responses): stream every citation batch before completion * test(responses): cover citations arriving before answer text * fix(responses): close reasoning before a citation opens the message item * test(responses): build citation batches without rebinding locals * fix(responses): hold citations until reasoning closes Citations that arrive while reasoning streams now wait and stream once the message item opens, so a citation no longer ends reasoning early. Annotation events also stop sending sequence_number, matching the text deltas. * fix(responses): keep held citations when reasoning ends without answer text A citation that arrives before the answer message opens now waits for the message, then streams once text arrives or at stream end when the final response has a message. A disconnect keeps the held citation in the billed response, and a tool call with no answer text no longer streams a message item that response.completed lacks. * fix(responses): keep a held citation's message ahead of the next tool call When the answer already has content, a citation held before the message opens now opens it before the next tool call takes its output index, so the stream and response.completed keep the message ahead of the function call like the non-streaming response. * test(responses): simplify citation comprehensions * test(responses): set the new citation state on the hand-built streaming iterator * fix(responses): keep held citations ahead of tool calls and in the final message A citation that arrived before a tool call with no answer text was dropped from response.completed and the non-streaming response, because the final transform built a message only when content was set. Keep an empty-text message for those citations, and flush a held citation before the next tool call opens, so the streamed message stays ahead of the function call. * test(responses): annotate the cited tool call test locals as Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../streaming_iterator.py | 128 +++++---- .../transformation.py | 6 +- .../test_litellm_completion_responses.py | 41 +++ .../test_streaming_iterator_transformation.py | 270 ++++++++++++++++++ 4 files changed, 391 insertions(+), 54 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 982158040c4..1e8ca85fdbb 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -30,6 +30,7 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, BaseLiteLLMOpenAIResponseObject, + ChatCompletionAnnotation, ContentPartAddedEvent, ContentPartDoneEvent, ContentPartDonePartOutputText, @@ -56,6 +57,7 @@ from litellm.types.llms.openai import ( WebSearchCallInProgressEvent, WebSearchCallSearchingEvent, ) +from litellm.types.responses.main import GenericResponseOutputItemContentAnnotation from litellm.types.utils import Delta as ChatCompletionDelta from litellm.types.utils import ( ModelResponse, @@ -105,7 +107,6 @@ def _delta_carries_output(delta: ChatCompletionDelta) -> bool: getattr(delta, "reasoning_content", None), delta.tool_calls, delta.function_call, - getattr(delta, "annotations", None), getattr(delta, "images", None), getattr(delta, "audio", None), ) @@ -150,7 +151,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.sent_output_text_done_event: bool = False self.sent_output_content_part_done_event: bool = False self.sent_output_item_done_event: bool = False - self.sent_annotation_events: bool = False + self._annotation_index: int = 0 + self._annotations_awaiting_message: tuple[ChatCompletionAnnotation, ...] = () self.litellm_model_response: ModelResponse | TextCompletionResponse | None = None self.completed_response = None self.final_text: str = "" @@ -339,7 +341,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): namespace_kwargs: Final = {"namespace": tool_namespace} if tool_namespace else {} return {**item_kwargs, "name": tool_name, **namespace_kwargs} - def _is_reasoning_end(self, chunk): + def _is_reasoning_end(self, chunk: ModelResponseStream) -> bool: if not chunk.choices: return False delta: Final = chunk.choices[0].delta @@ -350,7 +352,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if hasattr(delta, "thinking_blocks") and delta.thinking_blocks: return False - return delta.content or delta.function_call or delta.tool_calls or chunk.choices[0].finish_reason is not None + return bool( + delta.content or delta.function_call or delta.tool_calls or chunk.choices[0].finish_reason is not None + ) def _reserve_web_search_indexes(self, provider_fields: object) -> None: if not isinstance(provider_fields, dict): @@ -428,6 +432,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: fn_name = str(getattr(fn, "name", "") or "") fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) + if call_id not in self._tool_output_index_by_call_id: + self._queue_held_annotations() output_index = self._get_or_assign_tool_output_index(call_id) if self._is_tool_search_call(call_id, fn_name): @@ -782,8 +788,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return event def _queue_message_item_added_events(self) -> None: - if self._cached_item_id is None: - self._cached_item_id = f"msg_{uuid.uuid4()}" + item_id: Final = self._message_item_id() self.sent_message_item_added_event = True self.sent_content_part_added_event = True self._sequence_number += 1 @@ -792,7 +797,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): output_index=self._message_index(), item=BaseLiteLLMOpenAIResponseObject( **{ - "id": self._cached_item_id, + "id": item_id, "type": "message", "role": "assistant", "status": "in_progress", @@ -1081,6 +1086,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self.litellm_model_response: self._queue_compaction_item_added() self._queue_compaction_item_done() + self._queue_held_annotations() if self._pending_response_events: return self._pending_response_events.pop(0) # If tool calls exist, emit tool events before finishing/response.completed. @@ -1163,7 +1169,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): delta: Final = chunk.choices[0].delta if chunk.choices else None if delta and getattr(delta, "reasoning_content", None): self._accumulated_reasoning_content_parts.append(delta.reasoning_content) - if not self._is_reasoning_end(chunk): + if self._is_reasoning_end(chunk): + self._queue_reasoning_done_events() + + def _queue_reasoning_done_events(self) -> None: + if not self._reasoning_active or self._reasoning_done_emitted: return reasoning_content: Final = "".join(self._accumulated_reasoning_content_parts) self._cached_reasoning_item_id = ( @@ -1237,9 +1247,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._snapshot_chunk_for_stream_chunk_builder(chunk) ) self._queue_reasoning_done_events_on_reasoning_end(chunk) - response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) - if response_api_chunk: - self._pending_response_events.append(response_api_chunk) + self._queue_chunk_events(chunk) if self._pending_response_events: return self._pending_response_events.pop(0) @@ -1300,9 +1308,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) self._queue_reasoning_done_events_on_reasoning_end(chunk) - response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) - if response_api_chunk: - self._pending_response_events.append(response_api_chunk) + self._queue_chunk_events(chunk) if self._pending_response_events: return self._pending_response_events.pop(0) # Otherwise, loop to next chunk @@ -1313,6 +1319,58 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.finished = True raise e + def _queue_chunk_events(self, chunk: ModelResponseStream) -> None: + self._hold_chunk_annotations(chunk) + response_api_chunk: Final = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) + if response_api_chunk: + self._pending_response_events.append(response_api_chunk) + if self.sent_message_item_added_event: + self._queue_held_annotations() + + def _hold_chunk_annotations(self, chunk: ModelResponseStream) -> None: + chunk_annotations: Final = (getattr(chunk.choices[0].delta, "annotations", None) or ()) if chunk.choices else () + self._annotations_awaiting_message = (*self._annotations_awaiting_message, *chunk_annotations) + + def _queue_held_annotations(self) -> None: + annotations: Final = self._annotations_awaiting_message + if not annotations: + return + self._annotations_awaiting_message = () + self._queue_reasoning_done_events() + self._queue_annotation_added_events(annotations) + + def _queue_annotation_added_events(self, annotations: Sequence[ChatCompletionAnnotation]) -> None: + if not self.sent_message_item_added_event: + self._queue_message_item_added_events() + item_id: Final = self._message_item_id() + response_annotations: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations + ) + ) + self._pending_response_events.extend( + self._annotation_added_event(item_id, annotation) for annotation in response_annotations + ) + + def _message_item_id(self) -> str: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{uuid.uuid4()}" + return self._cached_item_id + + def _annotation_added_event( + self, item_id: str, annotation: GenericResponseOutputItemContentAnnotation + ) -> OutputTextAnnotationAddedEvent: + event: Final = OutputTextAnnotationAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + item_id=item_id, + output_index=self._message_index(), + content_index=0, + annotation_index=self._annotation_index, + annotation=annotation.model_dump(), + ) + self._annotation_index += 1 + return event + def _transform_chat_completion_chunk_to_response_api_chunk( self, chunk: ModelResponseStream ) -> ResponsesAPIStreamingResponse | None: @@ -1321,37 +1379,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): This currently handles emitting the OutputTextDeltaEvent, which is used by other tools using the responses API and the ReasoningSummaryTextDeltaEvent, which is used by the responses API to emit reasoning content. - It also handles emitting annotation.added events when annotations are detected in the chunk. """ - if self._cached_item_id is None: - self._cached_item_id = f"msg_{uuid.uuid4()}" - item_id: Final = self._cached_item_id + item_id: Final = self._message_item_id() - # Check if this chunk has annotations first (before processing text/reasoning) - # This ensures we detect and queue annotation events from the annotation chunk - if chunk.choices and hasattr(chunk.choices[0].delta, "annotations"): - annotations: Final = chunk.choices[0].delta.annotations - if annotations and self.sent_annotation_events is False: - self.sent_annotation_events = True - # Store annotation events to emit them one by one - if not hasattr(self, "_pending_annotation_events"): - response_annotations = LiteLLMCompletionResponsesConfig.transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations - ) - self._pending_annotation_events = [] - for idx, annotation in enumerate(response_annotations): - annotation_dict = ( - annotation.model_dump() if hasattr(annotation, "model_dump") else dict(annotation) - ) - event = OutputTextAnnotationAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, - item_id=item_id, - output_index=self._message_index(), - content_index=0, - annotation_index=idx, - annotation=annotation_dict, - ) - self._pending_annotation_events.append(event) # Priority 1: Handle reasoning content (highest priority) if ( chunk.choices @@ -1395,13 +1425,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._pending_tool_events: return self._pending_tool_events.pop(0) - # Priority 4: If we have pending annotation events, emit the next one - # This happens when the current chunk has no text/reasoning content - if hasattr(self, "_pending_annotation_events") and self._pending_annotation_events: - event = self._pending_annotation_events.pop(0) - return event - - # Priority 5: If we have pending tool events (from earlier chunk), emit the next one if self._pending_tool_events: return self._pending_tool_events.pop(0) @@ -1444,7 +1467,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _is_unstreamed_empty_message(self, item: object) -> bool: if getattr(item, "type", None) != "message" or self._message_output_index is not None: return False - return not any(getattr(part, "text", None) for part in attribute_items(item, "content")) + return not any( + getattr(part, "text", None) or getattr(part, "annotations", None) + for part in attribute_items(item, "content") + ) def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: if not isinstance(self.litellm_model_response, ModelResponse): diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index a342d704a9d..8b076ae922e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -3555,7 +3555,7 @@ class LiteLLMCompletionResponsesConfig: choice=choice, ) message_output_items.extend(image_generation_items) - elif choice.message.content is not None: + elif choice.message.content is not None or getattr(choice.message, "annotations", None): message_output_items.append( GenericResponseOutputItem( type="message", @@ -3629,13 +3629,13 @@ class LiteLLMCompletionResponsesConfig: return OutputText( type="output_text", - text=message.content, + text=message.content or "", annotations=transformed_annotations, ) @staticmethod def transform_chat_completion_annotations_to_response_output_annotations( - annotations: list[ChatCompletionAnnotation] | None, + annotations: Sequence[ChatCompletionAnnotation] | None, ) -> list[GenericResponseOutputItemContentAnnotation]: response_output_annotations: Final[list[GenericResponseOutputItemContentAnnotation]] = [] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 3fd4be6f722..475628569a0 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -663,6 +663,46 @@ class TestLiteLLMCompletionResponsesConfig: assert "message" not in output_types assert "function_call" in output_types + def test_cited_tool_call_only_response_keeps_its_citations_in_a_message_item(self) -> None: + citation: Final = {"start_index": 0, "end_index": 0, "url": "https://example.com", "title": "Example"} + chat_completion_response: Final = ModelResponse( + id="test-response-id", + created=1234567890, + model="gemini-3.8-flash", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + annotations=[{"type": "url_citation", "url_citation": citation}], + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_cited", + type="function", + function=Function(name="save_museum", arguments='{"name": "Louvre"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="which museum?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) + ) + + output: Final = responses_api_response.output + assert [item.type for item in output] == ["message", "function_call"] + assert output[0].content[0].text == "" + assert [annotation.url for annotation in output[0].content[0].annotations] == [citation["url"]] + def test_content_bearing_response_still_emits_message_item(self): """Turns with real text content must keep their message output item.""" chat_completion_response = ModelResponse( @@ -4306,6 +4346,7 @@ class TestEnsureOutputItemContentPartAdded: iterator._reasoning_active = False iterator._pending_response_events = [] iterator._pending_tool_events = [] + iterator._annotations_awaiting_message = () iterator._tool_output_index_by_call_id = {} iterator._tool_args_by_call_id = {} iterator._tool_item_id_by_call_id = {} diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index d9a0033b179..5fd0512fbab 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -13,6 +13,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve import itertools import json from collections.abc import Callable +from itertools import accumulate, chain, islice, product from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -989,6 +990,15 @@ def _with_reasoning(chunk: ModelResponseStream, reasoning: str) -> ModelResponse ) +def _with_citation(chunk: ModelResponseStream) -> ModelResponseStream: + choice: Final = chunk.choices[0] + annotations: Final = _annotation_only_chunk().choices[0].delta.annotations + delta: Final = Delta(**{**choice.delta.model_dump(), "annotations": annotations}) + return chunk.model_copy( + update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason=choice.finish_reason)]} + ) + + def _reasoning_then_text_chunks() -> tuple[ModelResponseStream, ...]: return (_reasoning_chunk("let me think"), _chunk("Hello"), _chunk("!", finish_reason="stop")) @@ -1662,6 +1672,266 @@ async def test_a_compaction_block_without_its_summary_streams_no_compaction_item assert [item["type"] for item in events[-1]["response"]["output"]] == ["message"] +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("batches", [(5,), (2, 3)]) +@pytest.mark.parametrize("with_text", [True, False]) +@pytest.mark.parametrize("reasoning_first", [True, False]) +@pytest.mark.asyncio +async def test_all_citation_batches_stream_before_message_done( + sync_mode: bool, batches: tuple[int, ...], with_text: bool, reasoning_first: bool +) -> None: + annotations: Final = [ + { + "type": "url_citation", + "url_citation": { + "start_index": 0, + "end_index": 5, + "url": f"https://example.com/{index}", + "title": f"Source {index}", + }, + } + for index in range(sum(batches)) + ] + citation_chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content="!" if with_text else "", + annotations=annotations[end - size : end], + ), + finish_reason="stop" if end == len(annotations) else None, + ) + ], + ) + for size, end in zip(batches, accumulate(batches), strict=True) + ) + reasoning_chunks: Final = _reasoning_then_text_chunks()[:1] if reasoning_first else () + chunks: Final = (*reasoning_chunks, _chunk("Claim"), *citation_chunks) + + events: Final = await _collect_events(_build_iterator(chunks), sync_mode) + + added: Final = [event for event in events if event.type == "response.output_text.annotation.added"] + terminal: Final = events[-1] + message: Final = next(item for item in terminal.response.output if item.type == "message") + final_annotations: Final = [annotation.model_dump() for annotation in message.content[0].annotations] + assert [event.annotation for event in added] == final_annotations + assert len(added) == len(annotations) + assert [event.annotation_index for event in added] == list(range(len(annotations))) + message_index: Final = next( + event.output_index + for event in events + if event.type == "response.output_item.added" and event.item.type == "message" + ) + assert all(event.item_id == message.id and event.output_index == message_index for event in added) + done_index: Final = next(index for index, event in enumerate(events) if event.type == "response.output_text.done") + assert all(events.index(event) < done_index for event in added) + assert message.content[0].text == "Claim" + ("!" * len(batches) if with_text else "") + + +def _citation_after_reasoning_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me think"), _annotation_only_chunk()) + + +def _citation_with_reasoning_chunks() -> tuple[ModelResponseStream, ...]: + return (_with_reasoning(_annotation_only_chunk(), "let me think"),) + + +def _citation_between_reasoning_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me "), _annotation_only_chunk(), _reasoning_chunk("think")) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "stream_chunks", + [_citation_after_reasoning_chunks, _citation_with_reasoning_chunks, _citation_between_reasoning_chunks], + ids=["citation_after_reasoning", "citation_with_reasoning", "citation_between_reasoning"], +) +@pytest.mark.asyncio +async def test_citation_during_reasoning_streams_once_reasoning_closes( + stream_chunks: StreamChunks, sync_mode: bool +) -> None: + chunks: Final = (*stream_chunks(), _chunk("Claim", finish_reason="stop")) + events: Final = await _collect_events(_build_iterator(chunks), sync_mode) + + output: Final = events[-1].response.output + message_index: Final = next(index for index, item in enumerate(output) if item.type == "message") + message: Final = output[message_index] + reasoning: Final = next(item for item in output if item.type == "reasoning") + reasoning_done: Final = next( + index + for index, event in enumerate(events) + if event.type == "response.output_item.done" and event.item.type == "reasoning" + ) + message_added: Final = next(index for index, event in enumerate(events) if _is_message_item(event)) + part_added: Final = next(index for index, event in enumerate(events) if event.type == "response.content_part.added") + text_done: Final = next(index for index, event in enumerate(events) if event.type == "response.output_text.done") + citations: Final = [event for event in events if event.type == "response.output_text.annotation.added"] + reasoning_deltas: Final = [ + index for index, event in enumerate(events) if event.type == "response.reasoning_summary_text.delta" + ] + summary_done: Final = next(event for event in events if event.type == "response.reasoning_summary_text.done") + text_delta: Final = next(event for event in events if event.type == "response.output_text.delta") + + assert len(citations) == 1 + assert reasoning_done < message_added < part_added < events.index(citations[0]) < text_done + assert all(index < reasoning_done for index in reasoning_deltas) + assert summary_done.text == reasoning.content[0].text == "let me think" + assert citations[0].item_id == message.id + assert citations[0].output_index == message_index + assert citations[0].annotation_index == 0 + assert [event.annotation for event in citations] == [entry.model_dump() for entry in message.content[0].annotations] + assert ("sequence_number" in citations[0].model_dump()) == ("sequence_number" in text_delta.model_dump()) + assert message.content[0].text == "Claim" + + +def _citation_then_tool_call_finishing_chunks() -> tuple[ModelResponseStream, ...]: + return (_annotation_only_chunk(), _tool_call_chunk(finish_reason="tool_calls")) + + +def _reasoning_citation_then_tool_call_finishing_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me think"), *_citation_then_tool_call_finishing_chunks()) + + +def _reasoning_citation_tool_call_then_empty_finish_chunks() -> tuple[ModelResponseStream, ...]: + return ( + _reasoning_chunk("let me think"), + _annotation_only_chunk(), + _tool_call_chunk(), + _chunk("", finish_reason="tool_calls"), + ) + + +def _citation_between_reasoning_then_empty_finish_chunks() -> tuple[ModelResponseStream, ...]: + return (*_citation_between_reasoning_chunks(), _chunk("", finish_reason="stop")) + + +def _tool_call_then_citation_finishing_chunks() -> tuple[ModelResponseStream, ...]: + return (_tool_call_chunk(), _annotation_only_chunk()) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "stream_chunks", + [ + _citation_then_tool_call_finishing_chunks, + _reasoning_citation_then_tool_call_finishing_chunks, + _reasoning_citation_tool_call_then_empty_finish_chunks, + _citation_between_reasoning_then_empty_finish_chunks, + _citation_between_reasoning_chunks, + _tool_call_then_citation_finishing_chunks, + ], + ids=[ + "citation_then_tool_call", + "reasoning_citation_then_tool_call", + "reasoning_citation_tool_call_then_empty_finish", + "citation_between_reasoning_then_empty_finish", + "citation_between_reasoning_without_finish", + "tool_call_then_citation_without_text", + ], +) +@pytest.mark.asyncio +async def test_streamed_citations_match_the_completed_message(stream_chunks: StreamChunks, sync_mode: bool) -> None: + events: Final = await _collect_events(_build_iterator(stream_chunks()), sync_mode) + + output: Final = events[-1].response.output + completed_message_indexes: Final = [index for index, item in enumerate(output) if item.type == "message"] + completed_citations: Final = chain.from_iterable( + output[index].content[0].annotations for index in completed_message_indexes + ) + completed_annotations: Final = [annotation.model_dump() for annotation in completed_citations] + streamed_message_indexes: Final = [event.output_index for event in events if _is_message_item(event)] + citations: Final = [event for event in events if event.type == "response.output_text.annotation.added"] + reasoning_done: Final = [ + index + for index, event in enumerate(events) + if event.type == "response.output_item.done" and event.item.type == "reasoning" + ] + message_added: Final = [index for index, event in enumerate(events) if _is_message_item(event)] + + assert len(citations) == 1 + assert [event.annotation for event in citations] == completed_annotations + assert sorted(set(streamed_message_indexes)) == completed_message_indexes + assert all(event.output_index in completed_message_indexes for event in citations) + assert len(reasoning_done) == sum(item.type == "reasoning" for item in output) + assert all(index < added for index, added in product(reasoning_done, message_added)) + + +def _empty_text_citation_then_tool_call_chunks() -> tuple[ModelResponseStream, ...]: + return (_chunk(""), *_citation_then_tool_call_finishing_chunks()) + + +def _reasoning_empty_text_citation_then_tool_call_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me think"), *_empty_text_citation_then_tool_call_chunks()) + + +def _empty_text_then_cited_tool_call_chunks() -> tuple[ModelResponseStream, ...]: + return (_chunk(""), _with_citation(_tool_call_chunk(finish_reason="tool_calls"))) + + +def _citation_tool_call_then_text_chunks() -> tuple[ModelResponseStream, ...]: + return (_annotation_only_chunk(), _tool_call_chunk(), _chunk("Claim", finish_reason="tool_calls")) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + ("stream_chunks", "expected_order"), + [ + (_empty_text_citation_then_tool_call_chunks, ("message", "function_call")), + (_reasoning_empty_text_citation_then_tool_call_chunks, ("reasoning", "message", "function_call")), + (_empty_text_then_cited_tool_call_chunks, ("message", "function_call")), + (_citation_tool_call_then_text_chunks, ("message", "function_call")), + (_citation_then_tool_call_finishing_chunks, ("message", "function_call")), + (_reasoning_citation_tool_call_then_empty_finish_chunks, ("reasoning", "message", "function_call")), + ], + ids=[ + "empty_text_citation_then_tool_call", + "reasoning_empty_text_citation_then_tool_call", + "empty_text_then_cited_tool_call", + "citation_tool_call_then_text", + "citation_then_tool_call_without_text", + "reasoning_citation_tool_call_then_empty_finish", + ], +) +@pytest.mark.asyncio +async def test_citation_before_a_tool_call_keeps_the_message_ahead_of_it( + stream_chunks: StreamChunks, expected_order: tuple[str, ...], sync_mode: bool +) -> None: + events: Final = await _collect_events(_build_iterator(stream_chunks()), sync_mode) + + added: Final = [event for event in events if event.type == "response.output_item.added"] + output: Final = events[-1].response.output + message_index: Final = expected_order.index("message") + citations: Final = [event for event in events if event.type == "response.output_text.annotation.added"] + + assert tuple(event.item.type for event in added) == expected_order + assert [event.output_index for event in added] == list(range(len(expected_order))) + assert tuple(item.type for item in output) == expected_order + assert [event.output_index for event in citations] == [message_index] + assert [event.annotation for event in citations] == [ + annotation.model_dump() for annotation in output[message_index].content[0].annotations + ] + + +def test_disconnect_after_a_held_citation_still_bills_it() -> None: + iterator: Final = _build_iterator((*_citation_between_reasoning_chunks(), _chunk("Claim", finish_reason="stop"))) + reasoning_deltas: Final = (event for event in iterator if event.type == "response.reasoning_summary_text.delta") + assert len(list(islice(reasoning_deltas, 2))) == 2 + + incomplete: Final = iterator.partial_response_on_disconnect() + + assert incomplete is not None + messages: Final = [item for item in incomplete.response.output if item.type == "message"] + annotations: Final = chain.from_iterable(message.content[0].annotations for message in messages) + assert [annotation.url for annotation in annotations] == ["https://example.com"] + + def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_bill_as_hidden_response_cost() -> None: from litellm.litellm_core_utils.hidden_params import get_hidden_params From 5eb1042bbf57f19d0d698b5c7b67b133000fbf5f Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:08:12 -0700 Subject: [PATCH 03/58] fix(token_counter): count OpenAI models newer than tiktoken's table with o200k_base (#48728) * fix(token_counter): count OpenAI models newer than tiktoken's table with o200k_base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): keep tiktoken's encoding for suffixed variants of models it knows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(token_counter): move the o200k provenance note into the assert message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): strip openai/ and azure/ prefixes before picking the encoding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust_bridge): azure/gpt-4o now resolves to o200k_base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): check the provider prefix through _fix_model_name to keep the type gate flat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(token_counter): check o200k on the gpt-6 models tiktoken 0.14 still does not know * fix(token_counter): count Azure deployment names and OpenAI alias groups with their OpenAI tokenizer * fix(token_counter): count mixed, alias, and base-model deployments with every served model's tokenizer A group that serves OpenAI and non-OpenAI models now reserves at least each served model's count, model_group_alias names count as their target group, custom Azure deployment names count with their base model, and azure_ai/, chatgpt/ and similar prefixes resolve to the listed OpenAI model. * fix(token_counter): keep mixed groups on OpenAI tokenizers, cache served models, and count context checks, failed rows and stream fallback with the served model * fix(token-counter): keep mixed groups at their old reservation and count new OpenAI release names with o200k A group that serves an OpenAI model next to another vendor now counts its own name the way tiktoken reads it as written, plus its OpenAI models, so it reserves what main did instead of the lower o200k count base_model picks the tokenizer for estimated usage only when it names an OpenAI model, now also in the Cursor stream counter and the output-limit response Unlisted gpt-5+ and o-series names count with o200k_base, and wildcard routes run the context-window check with the matched OpenAI model * fix(token-counter): count an alias to a group no OpenAI model serves by the alias as written Keeps the count main gives such an alias, so its budget reservation and rate limit match main and no other vendor's tokenizer loads for it. Also sorts the import the main merge left unsorted. * test: type the new token-counting tests' parameters * fix(token-counter): count a release-named group no OpenAI model serves as main does, and snapshot the OpenAI model set before iterating it * fix(token-counter): read served models through the router's allowlisted deployments_named, and coordinate the registration race test * perf(token-counter): classify each distinct served model once when walking a model group Budgeted requests walk every deployment in their model group, since the router's cached copy is outside the router surface the proxy may use. Classifying each distinct served model once instead of once per deployment cuts a 200-deployment group from 2.3-4.5 ms to about 1.4 ms with the same result. --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- AGENTS.md | 1 + litellm/__init__.py | 2 +- .../streaming_chunk_builder_utils.py | 10 +- litellm/litellm_core_utils/token_counter.py | 139 +++++++++++++++++- litellm/llms/azure/azure.py | 20 ++- litellm/llms/openai/common_utils.py | 10 +- litellm/llms/openai/openai.py | 15 +- .../hooks/parallel_request_limiter_v3.py | 28 +++- litellm/proxy/proxy_server.py | 8 +- .../proxy/response_api_endpoints/endpoints.py | 7 +- .../spend_tracking/budget_reservation.py | 17 ++- litellm/proxy/spend_tracking/input_tokens.py | 26 +++- litellm/proxy/utils.py | 6 +- litellm/router.py | 27 +++- .../io_token_rate_limit_check.py | 12 +- .../router_code_coverage.py | 1 + .../test_streaming_chunk_builder_utils.py | 36 +++++ .../litellm_core_utils/test_token_counter.py | 93 ++++++++++++ .../llms/openai/test_openai_common_utils.py | 23 +++ .../hooks/test_parallel_request_limiter_v3.py | 41 ++++++ .../response_api_endpoints/test_endpoints.py | 28 +++- .../spend_tracking/test_budget_reservation.py | 116 +++++++++++++++ tests/unit/proxy/test_proxy_token_counter.py | 26 ++++ ..._utils_model_creation_and_error_logging.py | 24 +++ tests/unit/rust_bridge/test_token_counter.py | 2 +- .../test_router/test_io_token_rate_limits.py | 32 ++++ tests/unit/test_router/test_router.py | 80 ++++++++++ 27 files changed, 788 insertions(+), 42 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 355e300f129..ce4b63173d2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -109,6 +109,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega - No file sprawl: deliberate file and folder structure - Standard over hand-rolled: use the official SDK or a library where one exists. Where none does, follow industry standards instead of inventing local conventions - API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module. If none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration +- Current default for new model names: logic that picks behavior from a model-name table (tokenizer, message framing, supported params) gives a provider's model the table does not list yet that provider's current behavior, never its oldest, since providers ship names faster than any table is updated Follow conventional commits for commit names and PR titles diff --git a/litellm/__init__.py b/litellm/__init__.py index fe11884a484..b7d46c197f2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -591,7 +591,7 @@ vertex_ai_safety_settings: Optional[dict] = None ####### COMPLETION MODELS ################### from typing import Set -open_ai_chat_completion_models: Set = set() +open_ai_chat_completion_models: Set[str] = set() open_ai_text_completion_models: Set = set() cohere_models: Set = set() cohere_chat_models: Set = set() diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 8d243e572fa..719f75546b9 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRAC from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) +from litellm.litellm_core_utils.token_counter import usage_tokenizer_model from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantContentValue, @@ -51,7 +52,7 @@ from litellm.types.utils import ( UnreportedUsage, Usage, ) -from litellm.utils import print_verbose, token_counter +from litellm.utils import get_base_model_from_metadata, print_verbose, token_counter if TYPE_CHECKING: from openai.types.completion_usage import CompletionUsage @@ -1552,6 +1553,9 @@ def stream_chunk_builder( return stream_chunk_builder_text_completion(chunks=chunks, messages=messages) model: Final = chunks[0]["model"] + usage_model: Final = usage_tokenizer_model( + model, get_base_model_from_metadata(logging_obj.model_call_details) if logging_obj is not None else None + ) # Initialize the response dictionary response: Final = processor.build_base_response(chunks) @@ -1575,7 +1579,7 @@ def stream_chunk_builder( completion_output = get_content_from_model_response(response) usage = processor.calculate_usage( chunks=chunks, - model=model, + model=usage_model, completion_output=completion_output, messages=messages, reasoning_tokens=0, @@ -1766,7 +1770,7 @@ def stream_chunk_builder( usage = processor.calculate_usage( chunks=chunks, - model=model, + model=usage_model, completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index f42e4eae08f..2fe840dc863 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -15,6 +15,7 @@ import anyio import anyio.lowlevel import httpx import tiktoken +from tiktoken.model import MODEL_TO_ENCODING from typing_extensions import ParamSpec, TypeVar import litellm @@ -49,6 +50,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value +from litellm.router_utils.deployment_pricing import deployment_pricing_models from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.llms.anthropic import ( AnthropicContentParamSource, @@ -1005,17 +1007,144 @@ def openai_tokenizer_encoding(model: str) -> "Encoding": def openai_tokenizer_encoding_name(model: str) -> str: - """The tiktoken encoding name for `model`, without loading the encoding.""" - from litellm.utils import print_verbose + """The tiktoken encoding name for `model`, without loading the encoding. - model_to_use: Final = _fix_model_name(model) + A listed OpenAI or Azure model that tiktoken does not know is either a suffixed variant of a model it does + know (`text-embedding-ada-002-v2`) or a newer release than its table, which gets o200k_base, the encoding every + OpenAI model since gpt-4o uses. Provider and region prefixes (`azure/eu/`, `azure_ai/`) are dropped, a name + listed only under Azure (`gpt-41`, `openai/o1-mini`) counts as listed through its `azure/` cost-map key, and an + unlisted name tiktoken counts with o200k_base (`openai/gpt-5.9`) or named like a gpt-5 or o-series release + (`openai/gpt-6.2`) gets it.""" + model_to_use: Final = _openai_model_name(model) or _fix_model_name(model) if "gpt-4o" in model_to_use: return "o200k_base" try: return tiktoken.encoding_name_for_model(model_to_use) except KeyError: - print_verbose("Warning: model not found. Using cl100k_base encoding.") - return "cl100k_base" + return _suffixed_variant_encoding_name(model_to_use) or "o200k_base" + + +def unresolved_tokenizer_model(model: str) -> str: + """`model`, or `gpt-3.5-turbo` when resolving the OpenAI model `model` names moves it off cl100k_base, the encoding + tiktoken gives the name as written, so a count with the result never drops below that one.""" + as_written: Final = _fix_model_name(model) + as_written_encoding: Final = ( + "o200k_base" if "gpt-4o" in as_written else (_tiktoken_encoding_name(as_written) or "cl100k_base") + ) + moved_off_cl100k: Final = ( + as_written_encoding == "cl100k_base" and openai_tokenizer_encoding_name(model) != "cl100k_base" + ) + return "gpt-3.5-turbo" if moved_off_cl100k else model + + +def usage_tokenizer_model(model: str, base_model: str | None) -> str: + """The model to estimate missing usage with: a deployment's `base_model` when it names an OpenAI model, since a + custom deployment name says nothing about its tokenizer, else `model`.""" + return base_model if base_model is not None and is_openai_model(base_model) else model + + +def is_openai_model(model: str) -> bool: + """Whether `model` names an OpenAI or Azure OpenAI model, listed in the cost map or counted with o200k_base by + tiktoken, with or without its prefixes.""" + return _openai_model_name(model) is not None + + +def deployment_openai_model(deployment: object) -> str | None: + """The OpenAI model a router deployment serves, read from its base model first, since a custom deployment name + (`azure/my-gpt-5`) says nothing about its tokenizer, or None when it serves no OpenAI model.""" + return _as_openai_model(_served_model(deployment)) + + +def served_openai_models(deployments: Iterable[object]) -> tuple[str | None, ...]: + """The distinct OpenAI models `deployments` serve, with None standing for any deployment serving another model.""" + served: Final = dict.fromkeys(_served_model(deployment) for deployment in deployments) + return tuple(dict.fromkeys(_as_openai_model(model) for model in served)) + + +def _served_model(deployment: object) -> str | None: + return next(iter(deployment_pricing_models(deployment)), None) + + +def _as_openai_model(model: str | None) -> str | None: + return model if model is not None and is_openai_model(model) else None + + +def single_encoding_openai_model(served: Sequence[str | None]) -> str | None: + """An OpenAI model to count a request with, when every `served_openai_models` entry is an OpenAI model on the same + encoding, else None.""" + if not served or None in served: + return None + encodings: Final = frozenset(openai_tokenizer_encoding_name(model) for model in served if model is not None) + return served[0] if len(encodings) == 1 else None + + +def _openai_model_name(model: str) -> str | None: + return _listed_openai_model_name(model) or _tiktoken_o200k_model_name(model) or _versioned_openai_model_name(model) + + +_VERSIONED_OPENAI_MODEL: Final = re.compile(r"(?:gpt-(?:[5-9]|[1-9]\d+)|o[1-9]\d*)(?:[.-]|$)") + + +def _versioned_openai_model_name(model: str) -> str | None: + name: Final = model.rsplit("/", 1)[-1] + return name if _VERSIONED_OPENAI_MODEL.match(name) else None + + +def _tiktoken_o200k_model_name(model: str) -> str | None: + name: Final = model.rsplit("/", 1)[-1] + return name if _tiktoken_encoding_name(name) == "o200k_base" else None + + +def _listed_openai_model_name(model: str) -> str | None: + segments: Final = model.split("/") + suffixes: Final = ("/".join(segments[start:]) for start in range(len(segments))) + listed: Final = any( + _is_openai_cost_map_model(suffix) or _is_openai_cost_map_model(f"azure/{suffix}") for suffix in suffixes + ) + return segments[-1] if listed else None + + +def _is_openai_cost_map_model(model: str) -> bool: + return ( + model in litellm.azure_llms + or model in litellm.open_ai_chat_completion_models + or _is_azure_openai_text_model(model) + ) + + +_TEXT_MODEL_MODES: Final = frozenset({"chat", "completion", "responses", "realtime"}) + + +def _is_azure_openai_text_model(model: str) -> bool: + cost_info: Final[Mapping[str, object]] = litellm.model_cost.get(model) or {} + return ( + cost_info.get("litellm_provider") == "azure" + and cost_info.get("mode") in _TEXT_MODEL_MODES + and _is_openai_model_name(model.rsplit("/", 1)[-1]) + ) + + +def _is_openai_model_name(name: str) -> bool: + return ( + name in litellm.open_ai_chat_completion_models + or _tiktoken_encoding_name(name) is not None + or any( + name.startswith(f"{listed}-") or listed.startswith(f"{name}-") + for listed in tuple(litellm.open_ai_chat_completion_models) + ) + ) + + +def _tiktoken_encoding_name(model: str) -> str | None: + try: + return tiktoken.encoding_name_for_model(model) + except KeyError: + return None + + +def _suffixed_variant_encoding_name(model: str) -> str | None: + known: Final = tuple(name for name in MODEL_TO_ENCODING if model.startswith(f"{name}-")) + return MODEL_TO_ENCODING[max(known, key=len)] if known else None def uses_legacy_message_accounting(model: str) -> bool: diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 5a181faad11..a1ba8de6d7f 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -33,6 +33,7 @@ from litellm.types.utils import ( from litellm.utils import ( CustomStreamWrapper, convert_to_model_response_object, + get_base_model_from_metadata, ) from ...types.llms.openai import HttpxBinaryResponseContent @@ -145,6 +146,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): azure_client: AzureOpenAI | OpenAI, data: dict, timeout: float | httpx.Timeout, + logging_obj: LiteLLMLoggingObj, ): """ Helper to: @@ -160,7 +162,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except BadRequestError as e: if not is_output_token_limit_error(e): raise - return build_output_token_limit_response(e=e, data=data, is_async=False) + return build_output_token_limit_response( + e=e, + data=data, + is_async=False, + base_model=get_base_model_from_metadata(logging_obj.model_call_details), + ) except Exception as e: raise e @@ -192,7 +199,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except BadRequestError as e: if not is_output_token_limit_error(e): raise - return build_output_token_limit_response(e=e, data=data, is_async=True) + return build_output_token_limit_response( + e=e, + data=data, + is_async=True, + base_model=get_base_model_from_metadata(logging_obj.model_call_details), + ) except Exception as e: raise e @@ -350,7 +362,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) headers, response = self.make_sync_azure_openai_chat_completion_request( - azure_client=azure_client, data=data, timeout=timeout + azure_client=azure_client, data=data, timeout=timeout, logging_obj=logging_obj ) if not isinstance(response, BaseModel): raise AzureOpenAIError( @@ -564,7 +576,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) headers, response = self.make_sync_azure_openai_chat_completion_request( - azure_client=azure_client, data=data, timeout=timeout + azure_client=azure_client, data=data, timeout=timeout, logging_obj=logging_obj ) logging_obj.model_call_details["response_headers"] = headers streamwrapper: Final = CustomStreamWrapper( diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 12ead582975..216725b9449 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: import litellm from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object -from litellm.litellm_core_utils.token_counter import token_counter +from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( _DEFAULT_TTL_FOR_HTTPX_CLIENTS, @@ -242,7 +242,7 @@ async def _aiter_once(chunk: ChatCompletionChunk) -> AsyncIterator[ChatCompletio def build_output_token_limit_response( - e: openai.BadRequestError, data: Mapping[str, object], is_async: bool + e: openai.BadRequestError, data: Mapping[str, object], is_async: bool, base_model: str | None ) -> tuple[httpx.Headers, ChatCompletion | Iterator[ChatCompletionChunk] | AsyncIterator[ChatCompletionChunk]]: """Synthesize the length-truncated response the provider itself returns for slightly larger output budgets. @@ -251,7 +251,11 @@ def build_output_token_limit_response( """ model: Final[str] = str(data.get("model", "")) messages: Final = data.get("messages") - prompt_tokens: Final = token_counter(model=model, messages=messages) if isinstance(messages, list) else 0 + prompt_tokens: Final = ( + token_counter(model=usage_tokenizer_model(model, base_model), messages=messages) + if isinstance(messages, list) + else 0 + ) if not data.get("stream"): return e.response.headers, _output_token_limit_completion(model, prompt_tokens) chunk: Final = _output_token_limit_chunk(model) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 9d4df354a30..5a3ab8aebea 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -44,6 +44,7 @@ from litellm.utils import ( CustomStreamWrapper, ProviderConfigManager, convert_to_model_response_object, + get_base_model_from_metadata, ) from ...types.llms.openai import * @@ -557,7 +558,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.BadRequestError as e: if not is_output_token_limit_error(e): raise - return build_output_token_limit_response(e=e, data=data, is_async=True) + return build_output_token_limit_response( + e=e, + data=data, + is_async=True, + base_model=get_base_model_from_metadata(logging_obj.model_call_details), + ) except Exception as e: raise e @@ -604,7 +610,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.BadRequestError as e: if not is_output_token_limit_error(e): raise - return build_output_token_limit_response(e=e, data=data, is_async=False) + return build_output_token_limit_response( + e=e, + data=data, + is_async=False, + base_model=get_base_model_from_metadata(logging_obj.model_call_details), + ) except Exception as e: if raw_response is not None: raise Exception( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index e44d9304ac1..a3a6f582dce 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -76,6 +76,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.proxy.spend_tracking.input_tokens import tokenizer_models from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, @@ -159,6 +160,12 @@ def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None: return resolve_model_group_alias(llm_router.model_group_alias, model) +def _tokenizer_models_via_proxy_router(model: str) -> tuple[str, ...]: + from litellm.proxy.proxy_server import llm_router + + return tokenizer_models(model, llm_router) + + def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) -> PTUTeamCeiling | None: from litellm.proxy.proxy_server import llm_router @@ -884,6 +891,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] = _resolve_ptu_team_ceiling_via_proxy_router, fail_closed_resolver: Callable[[], bool] = _fail_closed_rate_limit_enforcement_from_general_settings, force_hash_tag_grouping_resolver: Callable[[], bool] = _force_redis_hash_tag_grouping_from_litellm_settings, + tokenizer_models_resolver: Callable[[str], tuple[str, ...]] = _tokenizer_models_via_proxy_router, ): self.internal_usage_cache = internal_usage_cache self._time_provider = time_provider or datetime.now @@ -893,6 +901,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver self._fail_closed_resolver = fail_closed_resolver self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver + self._tokenizer_models_resolver = tokenizer_models_resolver if self.internal_usage_cache.dual_cache.redis_cache is not None: self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( BATCH_RATE_LIMITER_SCRIPT @@ -4009,15 +4018,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: estimate: Final = max( 0, - int( - token_counter( - model=model or "", - messages=countable_messages, - text=selected_text, - tools=countable_tools, - tool_choice=countable_tool_choice, - use_default_image_token_count=True, + *( + int( + token_counter( + model=tokenizer_model, + messages=countable_messages, + text=selected_text, + tools=countable_tools, + tool_choice=countable_tool_choice, + use_default_image_token_count=True, + ) ) + for tokenizer_model in self._tokenizer_models_resolver(model or "") ), ) return estimate + audio_token_estimate diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bb8fffd85bc..5e37ff1ba65 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -100,7 +100,11 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count +from litellm.litellm_core_utils.token_counter import ( + UnsupportedContentTypeError, + deployment_openai_model, + offload_token_count, +) from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, KimiK3Encodings, @@ -15542,7 +15546,7 @@ async def _count_request_tokens( count_source="local_estimate", ) estimate: Final = await _local_token_estimate( - model=model_to_use, + model=deployment_openai_model(deployment_view) or model_to_use, prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 4f8b1697f7e..5737421d4d4 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.token_counter import usage_tokenizer_model from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -96,6 +97,7 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import DeleteResponseResult from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse +from litellm.utils import get_base_model_from_metadata if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging @@ -537,7 +539,10 @@ async def cursor_model_list( def _dispatched_prompt_token_counter(logging_obj: LiteLLMLoggingObj) -> Callable[[], int]: return lambda: count_request_input_tokens( - model=str(logging_obj.model_call_details.get("model") or ""), + model=usage_tokenizer_model( + str(logging_obj.model_call_details.get("model") or ""), + get_base_model_from_metadata(logging_obj.model_call_details), + ), request_input=messages_since_newest_compaction(logging_obj.model_call_details.get("messages")), system_input=None, ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 0b91e57f01e..08840b44532 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,6 +7,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import chain from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast @@ -39,7 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) -from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model +from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model, tokenizer_models from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -1599,7 +1600,19 @@ async def count_request_input_tokens( models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router) if not models: return MappingProxyType({}) - return await count_input_tokens(request_body=request_body, raw_body=raw_body, models=models) + counted_as: Final = MappingProxyType({model: tokenizer_models(model, llm_router) for model in models}) + counts: Final = await count_input_tokens( + request_body=request_body, + raw_body=raw_body, + models=tuple(dict.fromkeys(chain.from_iterable(counted_as.values()))), + ) + return MappingProxyType( + { + model: count + for model, names in counted_as.items() + if (count := max((counts[name] for name in names if name in counts), default=None)) is not None + } + ) def _estimate_input_tokens( diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 927152296e9..0139869c6c6 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -15,14 +15,16 @@ import asyncio import json from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.token_counter import served_openai_models, unresolved_tokenizer_model from litellm.responses.litellm_completion_transformation.compaction_items import ( billable_input_since_newest_compaction, holds_compaction_summary, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -33,6 +35,9 @@ from litellm.rust_bridge.token_counter import ( rust_tokenizer, ) +if TYPE_CHECKING: + from litellm.router import Router + TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 _INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") @@ -47,6 +52,25 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: return sum(len(str(request_body.get(field, ""))) for field in _INPUT_SIZE_FIELDS) +def tokenizer_models(model: str, llm_router: Router | None) -> tuple[str, ...]: + """The names to count a request for `model` with, keeping the largest count. + + An alias resolves to its target group. A group served only by OpenAI models counts as those models, so it gets + their tokenizer. Any other group counts as its OpenAI models, if any, and as the requested name with the encoding + tiktoken gives that name as written, so its other deployments keep that count without loading their vendor's + tokenizer. A name the router does not index counts as written.""" + if llm_router is None: + return (model,) + group: Final = resolve_model_group_alias(llm_router.model_group_alias, model) or model + served: Final = served_openai_models(llm_router.deployments_named(group) or ()) + if not served: + return (model,) + openai_served: Final = tuple(served_model for served_model in served if served_model is not None) + if len(openai_served) == len(served): + return openai_served + return (unresolved_tokenizer_model(model), *openai_served) + + async def count_input_tokens( request_body: dict, raw_body: bytes | None, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..9f58c267231 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -168,7 +168,7 @@ from litellm.litellm_core_utils.served_output_texts import ( record_served_output_texts, served_stream_output_texts, ) -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -1139,7 +1139,9 @@ def _estimated_prompt_usage( None, ) return _estimate_dispatched_failure_usage( - model=str(model_call_details.get("model") or ""), + model=usage_tokenizer_model( + str(model_call_details.get("model") or ""), get_base_model_from_metadata(model_call_details) + ), request_input=( messages_since_newest_compaction(dispatched_input) if dispatched_call_type in _RESPONSES_CALL_TYPES diff --git a/litellm/router.py b/litellm/router.py index b63130c3313..e01f6474ded 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -162,7 +162,11 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, mask_sensitive_structure, ) -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import ( + offload_token_count, + served_openai_models, + single_encoding_openai_model, +) from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception from litellm.llms.base_llm.passthrough.transformation import replace_path_segment @@ -1533,6 +1537,7 @@ class Router: self._cached_deployment_tracks_usage = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( self._deployment_tracks_usage ) + self.served_openai_models = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._served_openai_models) self._discovered_model_info_cache: InMemoryCache = InMemoryCache( max_size_in_memory=max(len(model_list or ()), 1), default_ttl=2 * MODEL_INFO_REFRESH_SECONDS, @@ -11890,6 +11895,9 @@ class Router: indices: Final = self.model_name_to_deployment_indices.get(model_name) return None if indices is None else tuple(self.model_list[idx] for idx in indices) + def _served_openai_models(self, model_name: str) -> tuple[str | None, ...]: + return served_openai_models(self.deployments_named(model_name) or ()) + def team_public_model_names_of(self, team_id: str) -> frozenset[str]: """The public names `team_id` serves its own deployments under.""" with self._deployment_index_lock: @@ -14044,6 +14052,7 @@ class Router: self._ptu_share_teams.cache_clear() self._cached_team_ptu_ceiling.cache_clear() self._cached_deployment_tracks_usage.cache_clear() + self.served_openai_models.cache_clear() self._zero_cost_cache.clear() self._routing_group_rows = None @@ -14303,6 +14312,7 @@ class Router: messages: Sequence[Mapping[str, object]] | None, input: str | list[object] | None, request_kwargs: Mapping[str, object] | None = None, + tokenizer_model: str = "", ) -> int: """ Count input tokens for context-window pre-call checks. @@ -14333,7 +14343,7 @@ class Router: system_message: Final = anthropic_system_to_openai_message(extras.get("system")) if messages is not None: counted_messages: Final = (system_message, *messages) if system_message is not None else messages - return litellm.token_counter(messages=counted_messages, tools=tools) + return litellm.token_counter(model=tokenizer_model, messages=counted_messages, tools=tools) if input is not None: from openai.types.responses.response_create_params import ResponseInputParam @@ -14348,6 +14358,7 @@ class Router: only_since_compaction=True, ) return litellm.token_counter( + model=tokenizer_model, messages=cast(list, input_messages), # cast-ok: transformed chat messages tools=tools, ) @@ -14404,6 +14415,10 @@ class Router: messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, + tokenizer_model=single_encoding_openai_model( + self.served_openai_models(model) or served_openai_models(healthy_deployments) + ) + or "", ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( @@ -14418,6 +14433,8 @@ class Router: request_kwargs: Mapping[str, object] | None, input_token_count: int | None, count_inline: bool, + model: str, + deployments: Sequence[Mapping[str, object]], ) -> int | None: if input_token_count is not None: return input_token_count @@ -14428,6 +14445,10 @@ class Router: messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, + tokenizer_model=single_encoding_openai_model( + self.served_openai_models(model) or served_openai_models(deployments) + ) + or "", ) except Exception as e: verbose_router_logger.error( @@ -14622,6 +14643,8 @@ class Router: count_inline=has_countable_input and not skip_inline_token_count and any(limit is not None for limit in deployment_max_input_tokens), + model=model, + deployments=_returned_deployments, ) request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} requested_params: Final[frozenset[str]] = ( diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 2e350a2f87b..0b3d48c9e1d 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -22,7 +22,7 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET from litellm.types.router import RouterCacheEnum, RouterErrors @@ -127,6 +127,12 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | return itpm_key, otpm_key +def _tokenizer_model(deployment: object) -> str: + params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None + served: Final = params.get("model") if is_str_keyed_mapping(params) else None + return deployment_openai_model(deployment) or (served if isinstance(served, str) else "") + + def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int: if not request_kwargs: return 0 @@ -427,7 +433,7 @@ def io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" + _model: Final = _tokenizer_model(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) @@ -491,7 +497,7 @@ async def async_io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" + _model: Final = _tokenizer_model(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 9fb3633d53c..f936d23188c 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -175,6 +175,7 @@ ignored_function_names = [ "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py + "_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py "_drops_params", # Read by split_request_priority and priority_header_update, tested through both "_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 59946d54dc7..dd1f495b0d0 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2519,3 +2519,39 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] assert entry["guardrail_usage"] == {"topicPolicyUnits": 2} assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) + + +@pytest.mark.parametrize( + ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] +) +def test_stream_chunk_builder_counts_a_custom_named_deployments_missing_usage_with_its_base_model_only_when_openai( + base_model: str, counted_as: str +) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + chunks: Final = [ + ModelResponseStream(id="s", model="my-gpt-6", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]), + ModelResponseStream( + id="s", model="my-gpt-6", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")] + ), + ] + logging_obj: Final = Logging( + model="my-gpt-6", + messages=messages, + stream=True, + call_type="acompletion", + start_time=None, + litellm_call_id="call", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {"model_info": {"base_model": base_model}}} + + built: Final = stream_chunk_builder(chunks=chunks, messages=messages, logging_obj=logging_obj) + + assert built is not None + assert litellm.token_counter(model=base_model, messages=messages) != litellm.token_counter( + model="my-gpt-6", messages=messages + ) + assert built.usage.prompt_tokens == litellm.token_counter(model=counted_as, messages=messages) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index b02e0732817..9b0f6607708 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.token_counter import ( calculate_img_tokens, estimate_image_tokens, high_detail_image_token_upper_bound, + is_openai_model, offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new @@ -2600,3 +2601,95 @@ def test_get_modified_max_tokens_leaves_a_gemini_4_long_decode_budget_alone(monk ) == 200_000 ) + + +_MULTILINGUAL_SAMPLE: Final = "日本語のテキストとहिन्दी पाठ और русский текст. " * 20 + + +@pytest.mark.parametrize( + "model", + [ + "gpt-6-sol", + "gpt-6.1-sol", + "gpt-6.1-terra", + "azure/gpt-41", + "azure/gpt-5.1", + "openai/gpt-5.1", + "openai/gpt-4o", + "gpt-41", + "eu/gpt-6.1-sol", + "azure/eu/gpt-6.1-sol", + "azure/gpt-5.1-chat", + "azure/gpt-6.1-sol-2026-09-29", + "openai/gpt-6.1-sol-2026-09-29", + "azure_ai/gpt-6.1-sol", + "chatgpt/gpt-6.1-sol", + "perplexity/openai/gpt-5", + ], +) +def test_token_counter_counts_openai_models_newer_than_tiktokens_table_with_o200k(model: str): + expected: Final = len(tiktoken.get_encoding("o200k_base").encode(_MULTILINGUAL_SAMPLE)) + + assert expected != len(tiktoken.get_encoding("cl100k_base").encode(_MULTILINGUAL_SAMPLE)) + assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == expected, ( + f"{model} should count with o200k_base. OpenAI's /v1/responses/input_tokens gave gpt-5, gpt-5.1, gpt-6-sol, " + "gpt-6-luna, gpt-6-astra, gpt-6.1-sol and gpt-6.1-terra the same count on this sample, checked 2026-10-28" + ) + + +@pytest.mark.parametrize("model", ["openai/gpt-5.9", "chatgpt/gpt-5.3-instant", "azure/gpt-5.9-mini"]) +def test_token_counter_counts_unlisted_openai_names_with_the_encoding_tiktoken_gives_them(model: str): + assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len( + tiktoken.encoding_for_model(model.rsplit("/", 1)[-1]).encode(_MULTILINGUAL_SAMPLE) + ) + + +@pytest.mark.parametrize("model", ["gpt-6.2", "openai/gpt-6.2", "azure/gpt-6.2-mini", "azure/eu/o9-mini", "gpt-10"]) +def test_token_counter_counts_unlisted_names_of_newer_openai_releases_with_o200k(model: str): + assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len( + tiktoken.get_encoding("o200k_base").encode(_MULTILINGUAL_SAMPLE) + ) + + +@pytest.mark.parametrize( + "model", + [ + "groq/qwen3-32b", + "openai/my-llama", + "text-embedding-ada-002-v2", + "azure/gpt-4", + "azure/my-deployment", + "azure/ada", + "azure/mistral-large-latest", + "azure/command-r-plus", + "azure_ai/mistral-large-latest", + ], +) +def test_token_counter_keeps_cl100k_for_non_openai_models_and_variants_of_cl100k_models(model: str): + assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len( + tiktoken.get_encoding("cl100k_base").encode(_MULTILINGUAL_SAMPLE) + ) + + +def test_openai_model_check_holds_while_openai_models_are_registered_concurrently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(litellm.model_cost, "azure/azure-only-chat-model", {"litellm_provider": "azure", "mode": "chat"}) + registered: Final = tuple(f"concurrently-registered-model-{index}" for index in range(200)) + checked: Final = threading.Semaphore(0) + + def register() -> None: + for name in registered: + if not checked.acquire(timeout=10): + return + litellm.open_ai_chat_completion_models.add(name) + + try: + with ThreadPoolExecutor(max_workers=1) as pool: + registering: Final = pool.submit(register) + while not registering.done(): + assert not is_openai_model("azure-only-chat-model") + checked.release() + registering.result() + finally: + litellm.open_ai_chat_completion_models.difference_update(registered) diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 8dc92e23fdf..206fd13260c 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -322,6 +322,29 @@ def test_mapped_response_still_bills_the_prompt_the_provider_processed(provider) assert litellm.completion_cost(completion_response=response) > 0 +@pytest.mark.parametrize("provider", ["openai", "azure"]) +@pytest.mark.parametrize( + ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] +) +@pytest.mark.asyncio +async def test_mapped_response_bills_a_custom_named_deployment_with_its_base_model_only_when_it_is_openai( + provider: str, base_model: str, counted_as: str +) -> None: + messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + overrides: Final = {"model": f"{provider}/my-gpt-6", "messages": messages, "base_model": base_model} + + sync_response: Final = litellm.completion( + **_completion_kwargs(provider, _sync_client_raising(provider, OUTPUT_LIMIT_400_MESSAGE), **overrides) + ) + async_response: Final = await litellm.acompletion( + **_completion_kwargs(provider, _async_client_raising(provider, OUTPUT_LIMIT_400_MESSAGE), **overrides) + ) + + expected: Final = token_counter(model=counted_as, messages=messages) + assert token_counter(model=base_model, messages=messages) != token_counter(model="my-gpt-6", messages=messages) + assert (sync_response.usage.prompt_tokens, async_response.usage.prompt_tokens) == (expected, expected) + + @pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.asyncio async def test_async_output_limit_400_maps_to_length_truncated_response(provider): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index e056e723a04..a38f1a54944 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -10,6 +10,7 @@ import time from collections.abc import AsyncGenerator, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone +from functools import partial from typing import Any, Dict, Final, List, Optional from unittest.mock import patch @@ -45,6 +46,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) +from litellm.proxy.spend_tracking.input_tokens import tokenizer_models from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.router_utils.ptu_shares import PTUTeamCeiling @@ -9172,3 +9174,42 @@ async def test_a_successful_extraction_counts_its_input_tokens_toward_the_key_tp call_type="aextraction", ) assert refused.value.status_code == 429 + + +_MULTILINGUAL_REQUEST: Final[dict[str, object]] = { + "messages": [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] +} + + +def _project_itpm_estimate(router: Router, model: str) -> int: + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), + tokenizer_models_resolver=partial(tokenizer_models, llm_router=router), + ) + return handler._estimate_precise_input_tokens(_MULTILINGUAL_REQUEST, model, "acompletion") + + +def test_project_itpm_counts_a_group_mixing_in_a_non_openai_model_by_its_name_as_written_and_its_openai_models(): + router: Final = Router( + model_list=[ + {"model_name": "gpt-6.1-sol", "litellm_params": {"model": model, "api_key": "sk-test"}} + for model in ("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5") + ] + ) + messages: Final = _MULTILINGUAL_REQUEST["messages"] + + assert _project_itpm_estimate(router, "gpt-6.1-sol") == max( + litellm.token_counter(model="my-model", messages=messages), + litellm.token_counter(model="openai/gpt-6.1-sol", messages=messages), + ) + + +def test_project_itpm_counts_a_model_group_alias_as_its_target_group(): + router: Final = Router( + model_list=[ + {"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-test"}} + ], + model_group_alias={"prod": "gpt-6.1-sol"}, + ) + + assert _project_itpm_estimate(router, "prod") == _project_itpm_estimate(router, "gpt-6.1-sol") diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 7760b6b7fcc..f1aac24dd60 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -189,7 +189,11 @@ _TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's n async def _cursor_stream_failure_prompt_tokens( - monkeypatch: pytest.MonkeyPatch, request_input: list[Mapping[str, str]], fails_after_output: bool + monkeypatch: pytest.MonkeyPatch, + request_input: list[Mapping[str, str]], + fails_after_output: bool, + model: str = "gpt-6-astra", + model_info: Mapping[str, str] | None = None, ) -> int: import litellm.proxy.proxy_server as ps from litellm.integrations.custom_logger import CustomLogger @@ -208,7 +212,6 @@ async def _cursor_stream_failure_prompt_tokens( usage: Final = request_data.get("combined_usage_object") self.prompt_tokens = usage.prompt_tokens if isinstance(usage, litellm.Usage) else None - model: Final = "gpt-6-astra" upstream_url: Final = "https://streaming.example/v1" created: Final = { "type": "response.created", @@ -239,6 +242,7 @@ async def _cursor_stream_failure_prompt_tokens( { "model_name": model, "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + "model_info": dict(model_info or {}), } ], num_retries=0, @@ -276,6 +280,26 @@ async def test_a_failed_cursor_stream_counts_its_input_from_the_newest_compactio assert resent == dispatched_after_trim +@pytest.mark.asyncio +@pytest.mark.parametrize("fails_after_output", [False, True]) +@pytest.mark.parametrize( + ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] +) +async def test_a_failed_cursor_stream_counts_a_custom_named_deployment_with_its_base_model_only_when_it_is_openai( + monkeypatch: pytest.MonkeyPatch, base_model: str, counted_as: str, fails_after_output: bool +) -> None: + from litellm.proxy.utils import count_request_input_tokens + + turn: Final = {"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20} + + counted: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [turn], fails_after_output, model="my-gpt-6", model_info={"base_model": base_model} + ) + + assert count_request_input_tokens(base_model, [turn], None) != count_request_input_tokens("my-gpt-6", [turn], None) + assert counted == count_request_input_tokens(counted_as, [turn], None) + + @pytest.mark.asyncio async def test_responses_api_background_polling_rejects_missing_input(): from fastapi import Response as FastAPIResponse diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 024fe5229be..035b7d9dcd3 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -28,6 +28,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( estimate_request_max_cost, get_budget_window_start, release_unbound_budget_reservation, + count_request_input_tokens, reserve_budget_for_request, ) from litellm.proxy.utils import ProxyLogging @@ -426,3 +427,118 @@ async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spe ) def test_budget_window_start_reads_json_encoded_windows(window: str, expected: datetime | None) -> None: assert get_budget_window_start(window) == expected + + +_MULTILINGUAL_MESSAGES: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + + +def _router_with_groups(*groups: tuple[str, str]) -> Router: + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": model, "api_key": "sk-test"}} for name, model in groups + ] + ) + + +async def _input_tokens(router: Router, model: str) -> dict[str, int]: + return dict( + await count_request_input_tokens( + request_body={"model": model, "messages": _MULTILINGUAL_MESSAGES}, + route="/v1/chat/completions", + llm_router=router, + ) + ) + + +@pytest.mark.asyncio +async def test_alias_group_of_an_openai_model_reserves_on_that_models_token_count(): + router: Final = _router_with_groups(("gpt-6.1-sol", "openai/gpt-6.1-sol"), ("prod-gpt-6", "openai/gpt-6.1-sol")) + + by_name: Final = await _input_tokens(router, "gpt-6.1-sol") + + assert await _input_tokens(router, "prod-gpt-6") == {"prod-gpt-6": by_name["gpt-6.1-sol"]} + + +@pytest.mark.asyncio +async def test_model_group_alias_of_an_openai_group_reserves_on_that_groups_token_count(): + router: Final = Router( + model_list=[ + {"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-test"}} + ], + model_group_alias={"prod": "gpt-6.1-sol"}, + ) + + by_name: Final = await _input_tokens(router, "gpt-6.1-sol") + + assert await _input_tokens(router, "prod") == {"prod": by_name["gpt-6.1-sol"]} + + +@pytest.mark.asyncio +async def test_model_group_alias_of_a_group_no_openai_model_serves_reserves_on_the_alias_as_written(): + router: Final = Router( + model_list=[ + { + "model_name": "claude-sonnet-5-5", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"}, + } + ], + model_group_alias={"claude-alias": "claude-sonnet-5-5"}, + ) + + assert await _input_tokens(router, "claude-alias") == { + "claude-alias": litellm.token_counter(model="claude-alias", messages=_MULTILINGUAL_MESSAGES) + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", ["gpt-6-claude", "o4-alias"]) +async def test_release_named_group_or_alias_no_openai_model_serves_reserves_on_an_unknown_names_count( + name: str, +) -> None: + router: Final = Router( + model_list=[ + {"model_name": group, "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"}} + for group in ("claude-sonnet-5-5", "gpt-6-claude") + ], + model_group_alias={"o4-alias": "claude-sonnet-5-5"}, + ) + + assert await _input_tokens(router, name) == { + name: litellm.token_counter(model="my-model", messages=_MULTILINGUAL_MESSAGES) + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("group", "counted_as"), + [("gpt-6.1-sol", "my-model"), ("openai/gpt-6.1-sol", "my-model"), ("mixed", "my-model"), ("gpt-4o", "gpt-4o")], +) +async def test_group_mixing_in_a_non_openai_model_reserves_on_its_name_as_written_and_its_openai_models( + group: str, counted_as: str +): + served: Final = ("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5") + router: Final = _router_with_groups(*((group, model) for model in served)) + + reserved: Final = (await _input_tokens(router, group))[group] + + assert reserved == max( + litellm.token_counter(model=counted_as, messages=_MULTILINGUAL_MESSAGES), + litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_MESSAGES), + ) + + +@pytest.mark.asyncio +async def test_group_reserves_on_its_current_deployments_after_another_vendors_deployment_joins(): + router: Final = _router_with_groups(("prod", "openai/gpt-6.1-sol")) + openai_only: Final = (await _input_tokens(router, "prod"))["prod"] + + router.add_deployment( + Deployment( + model_name="prod", + litellm_params=LiteLLM_Params(model="anthropic/claude-sonnet-5-5", api_key="sk-test"), + ) + ) + + by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES) + assert openai_only != by_own_name + assert (await _input_tokens(router, "prod"))["prod"] == by_own_name diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 0e8fc6823bd..b21730c64f6 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3989,6 +3989,32 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form")) +@pytest.mark.asyncio +async def test_token_counter_counts_a_custom_azure_deployment_name_with_its_base_models_tokenizer( + monkeypatch: pytest.MonkeyPatch, +): + azure_params: Final = {"api_key": "sk-test", "api_base": "https://example.openai.azure.com"} + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "named", + "litellm_params": {"model": "azure/my-gpt-6", "base_model": "azure/gpt-6.1-sol", **azure_params}, + }, + {"model_name": "served", "litellm_params": {"model": "azure/gpt-6.1-sol", **azure_params}}, + ] + ), + ) + messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + + named: Final = await token_counter(request=TokenCountRequest(model="named", messages=messages)) + served: Final = await token_counter(request=TokenCountRequest(model="served", messages=messages)) + + assert named.total_tokens == served.total_tokens + + _COUNTED_HELLO: Final[Mapping[str, JsonValue]] = MappingProxyType( {"model": "count-log-gpt", "messages": [{"role": "user", "content": "Hello, how many tokens is this?"}]} ) diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index bc13595df6c..f2151572f49 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -2308,6 +2308,30 @@ def test_a_failed_dispatch_is_estimated_as_input_only(): assert usage.total_tokens == usage.prompt_tokens +@pytest.mark.parametrize( + ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] +) +def test_a_failed_dispatch_to_a_custom_named_deployment_is_counted_with_its_base_model_only_when_it_is_openai( + base_model: str, counted_as: str +): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _estimated_prompt_usage, count_request_input_tokens + + messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + logging_obj = _dispatched_failure_logging_obj( + "my-gpt-6", get_litellm_params(metadata={"model_info": {"base_model": base_model}}) + ) + logging_obj.model_call_details["messages"] = messages + + usage: Final = _estimated_prompt_usage(logging_obj.model_call_details, {}, "acompletion") + + assert usage is not None + assert count_request_input_tokens(base_model, messages, None) != count_request_input_tokens( + "my-gpt-6", messages, None + ) + assert usage.prompt_tokens == count_request_input_tokens(counted_as, messages, None) + + @pytest.mark.parametrize("request_input", [[], object()]) def test_nothing_is_estimated_when_there_is_nothing_to_count(request_input): from litellm.proxy.utils import _estimate_dispatched_failure_usage diff --git a/tests/unit/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py index 3513a0ef44c..068d389ebbd 100644 --- a/tests/unit/rust_bridge/test_token_counter.py +++ b/tests/unit/rust_bridge/test_token_counter.py @@ -144,7 +144,7 @@ async def test_each_tokenizer_gets_its_own_cached_counter(fake_tokenizers: None) ("gemini/gemini-2.5-pro", "cl100k_base"), ("mistral/mistral-large-latest", "cl100k_base"), ("my-router-alias", "cl100k_base"), - ("azure/gpt-4o", "cl100k_base"), + ("azure/gpt-4o", "o200k_base"), ("command-r-plus", "cl100k_base"), ("gpt-4o", "o200k_base"), ("gpt-4o-mini", "o200k_base"), diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 0ae22534278..05921aa70b3 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -3,6 +3,7 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). """ import asyncio +from typing import Final import pytest @@ -1039,3 +1040,34 @@ class TestContextSlotRetention: assert deployment is not None router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) assert get_io_token_rate_limit_request_kwargs() is kwargs + + +@pytest.mark.asyncio +async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_token_count() -> None: + async def reserved_input(deployment: dict[str, object]) -> object: + metadata: Final[dict[str, object]] = {} + request_kwargs: Final[dict[str, object]] = { + "messages": [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}], + "max_tokens": 5, + "metadata": metadata, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + await async_io_token_pre_call_check(DualCache(), deployment) + return metadata[ITPM_RESERVED_KEY] + + named = await reserved_input( + { + "litellm_params": {"model": "azure/my-gpt-6", "base_model": "azure/gpt-6.1-sol", "itpm": 100_000}, + "model_info": {"id": "named-deployment"}, + "model_name": "named", + } + ) + served = await reserved_input( + { + "litellm_params": {"model": "azure/gpt-6.1-sol", "itpm": 100_000}, + "model_info": {"id": "served-deployment"}, + "model_name": "served", + } + ) + + assert named == served diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 300c9cff6ff..d59fb69dcaf 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7938,6 +7938,86 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch assert calls == [1] +_MULTILINGUAL_PROMPT: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] + + +def _context_limited_router(*served: str) -> litellm.Router: + limit: Final = ( + litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) + + litellm.token_counter(model="prod", messages=_MULTILINGUAL_PROMPT) + ) // 2 + return litellm.Router( + model_list=[ + { + "model_name": "prod", + "litellm_params": {"model": model, "api_key": "sk-test"}, + "model_info": {"max_input_tokens": limit}, + } + for model in served + ], + enable_pre_call_checks=True, + ) + + +def test_pre_call_checks_count_a_group_of_one_openai_model_with_that_models_tokenizer(): + router: Final = _context_limited_router("openai/gpt-6.1-sol") + + kept: Final = router._pre_call_checks( + model="prod", healthy_deployments=router.get_model_list(model_name="prod"), messages=_MULTILINGUAL_PROMPT + ) + + assert len(kept) == 1 + + +def test_pre_call_checks_count_a_group_mixing_in_another_vendor_by_its_own_name(): + router: Final = _context_limited_router("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5") + + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="prod", healthy_deployments=router.get_model_list(model_name="prod"), messages=_MULTILINGUAL_PROMPT + ) + + +async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer(): + router: Final = _context_limited_router("openai/gpt-6.1-sol") + + count: Final = await router._acount_pre_call_check_tokens( + model="prod", + healthy_deployments=router.get_model_list(model_name="prod"), + messages=_MULTILINGUAL_PROMPT, + input=None, + request_kwargs=None, + ) + + assert count == litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) + + +async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer(): + limit: Final = ( + litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) + + litellm.token_counter(model="my-model", messages=_MULTILINGUAL_PROMPT) + ) // 2 + router: Final = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-test"}, + "model_info": {"max_input_tokens": limit}, + } + ], + enable_pre_call_checks=True, + ) + + sync_response: Final = router.completion( + model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT, mock_response="ok" + ) + async_response: Final = await router.acompletion( + model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT, mock_response="ok" + ) + + assert (sync_response.choices[0].message.content, async_response.choices[0].message.content) == ("ok", "ok") + + def test_pre_call_checks_uses_precounted_tokens(monkeypatch): """ An async caller counts off the event loop and passes the result in. _pre_call_checks From ec7162aea2079c4616d9eda5f0738956d59a3102 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:23:11 -0700 Subject: [PATCH 04/58] fix(cost_map): mark fireworks_ai/minimax-m3 as vision and pin capabilities verified against live calls (#43390) * fix(cost_map): mark fireworks_ai/minimax-m3 as vision and pin capabilities verified against live calls Fireworks' serverless minimax-m3 answers image inputs, but its model listings say supports_image_input false, so the last two registry syncs flipped supports_vision back to false and the Fireworks pre-flight turned every image request into a 400 without ever calling the provider. Flip both minimax-m3 rows back to supports_vision true and record the live-call evidence in ci_cd/cost_map_pins.json. The cost map guard now fails any PR that lowers a pinned capability, quoting the evidence, and the sync bot may not touch the pins file. * fix(cost_map): fail the guard when the pins file is missing * fix(cost_map_guard): skip the pin checks on a branch cut before the pins file existed * fix(ci): check cost map pins against the merge result and keep price sync PRs off the pins file * fix(ci): run the cost map guard against the live base branch tip and re-run it when main's pins change * fix(ci): pass the live base tip to the cost map guard through env and re-run every guard run on a pin change main now runs the guard on pull_request with the base branch's copy of the script, so the workflow passes the live tip in COST_MAP_GUARD_TARGET, which the older script ignores. The recheck reads past GraphQL's 100-file cap, re-runs every guard run for a shared head, and fails when a re-run cannot start. The guard tells a merge conflict apart from a failed merge-tree call. * fix(ci): run the cost map pins recheck on litellm_ base branches too * fix(ci): bar bot-opened PRs from the pins, recheck pin PRs on cost map changes, and skip the merge for untouched PRs * fix(ci): read the pushed files from the compare API in the cost map pins recheck The push event GitHub hands to Actions leaves each commit's added, modified and removed lists null, so the recheck never saw a pins change and re-ran only the PRs that touch the pins. Compare the push's before and after commits instead, and re-check every cost map PR when the branch is new, the push was forced, the compare fails, or it lists 300 files, the most it returns. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/CODEOWNERS | 1 + .github/workflows/cost-map-guard.yml | 12 +- .github/workflows/cost-map-pins-recheck.yml | 91 ++++++ ci_cd/cost_map_guard.py | 147 +++++++++- ci_cd/cost_map_pins.json | 16 ++ ...odel_prices_and_context_window_backup.json | 4 +- model_prices_and_context_window.json | 4 +- tests/unit/test_cost_map_guard.py | 270 +++++++++++++++++- 8 files changed, 516 insertions(+), 29 deletions(-) create mode 100644 .github/workflows/cost-map-pins-recheck.yml create mode 100644 ci_cd/cost_map_pins.json diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 158110cbbab..2826098fd1f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,5 +1,6 @@ /model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri +/ci_cd/cost_map_pins.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /helm/ @yassin-berriai /terraform/ @yassin-berriai /.github/workflows/zizmor.yml @yuneng-berri diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml index 05b22017cbb..6519b65e71b 100644 --- a/.github/workflows/cost-map-guard.yml +++ b/.github/workflows/cost-map-guard.yml @@ -21,16 +21,18 @@ jobs: with: ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - - name: Fetch the pull request head and its merge base + - name: Fetch the pull request head, its merge base, and the base branch tip id: revisions env: GH_TOKEN: ${{ github.token }} - BASE_SHA: ${{ github.event.pull_request.base.sha }} + BASE_REF: ${{ github.event.pull_request.base.ref }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} run: | - merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')" - git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" + target="$(gh api "repos/${GITHUB_REPOSITORY}/branches/${BASE_REF}" --jq '.commit.sha')" + merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${target}...${HEAD_SHA}" --jq '.merge_base_commit.sha')" + git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" "$target" echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT" + echo "target=$target" >> "$GITHUB_OUTPUT" - name: Set up uv uses: ./.github/actions/setup-uv-with-retries with: @@ -38,6 +40,8 @@ jobs: - name: Run the guard env: MERGE_BASE: ${{ steps.revisions.outputs.merge_base }} + COST_MAP_GUARD_TARGET: ${{ steps.revisions.outputs.target }} + COST_MAP_GUARD_AUTHOR_TYPE: ${{ github.event.pull_request.user.type }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} HEAD_REF: ${{ github.event.pull_request.head.ref }} run: | diff --git a/.github/workflows/cost-map-pins-recheck.yml b/.github/workflows/cost-map-pins-recheck.yml new file mode 100644 index 00000000000..f2572efbb31 --- /dev/null +++ b/.github/workflows/cost-map-pins-recheck.yml @@ -0,0 +1,91 @@ +name: Cost map pins recheck + +on: + push: + branches: + - main + - "litellm_**" + paths: + - model_prices_and_context_window.json + - litellm/model_prices_and_context_window_backup.json + - ci_cd/cost_map_pins.json + +permissions: {} + +jobs: + rerun-cost-map-guard: + runs-on: ubuntu-latest + timeout-minutes: 30 + permissions: + actions: write + contents: read + pull-requests: read + steps: + - name: Re-run the cost map guard on open pull requests the push could break + env: + GH_TOKEN: ${{ github.token }} + BASE_REF: ${{ github.ref_name }} + BEFORE_SHA: ${{ github.event.before }} + AFTER_SHA: ${{ github.event.after }} + FORCED: ${{ github.event.forced }} + COST_MAP_PATHS: '["model_prices_and_context_window.json", "litellm/model_prices_and_context_window_backup.json", "ci_cd/cost_map_pins.json"]' + PINS_PATHS: '["ci_cd/cost_map_pins.json"]' + run: | + if [[ "$BEFORE_SHA" =~ ^0+$ || "$FORCED" == true ]] \ + || ! gh api "repos/${GITHUB_REPOSITORY}/compare/${BEFORE_SHA}...${AFTER_SHA}?per_page=1" \ + --jq '(env.PINS_PATHS | fromjson) as $pins | (.files | length) < 300 and all(.files[]; [.filename, .previous_filename] | any(IN($pins[])) | not)' \ + | grep -qx true; then + export RECHECK_PATHS="$COST_MAP_PATHS" + else + export RECHECK_PATHS="$PINS_PATHS" + fi + rerun_cutoff="$(date -u -d '30 days ago' +%Y-%m-%dT%H:%M:%SZ)" + candidates="$(gh api graphql --paginate -F owner="${GITHUB_REPOSITORY_OWNER}" -F name="${GITHUB_REPOSITORY#*/}" -F base="$BASE_REF" \ + -f query='query($owner: String!, $name: String!, $base: String!, $endCursor: String) { + repository(owner: $owner, name: $name) { + pullRequests(states: OPEN, baseRefName: $base, first: 25, after: $endCursor) { + nodes { number headRefOid changedFiles files(first: 100) { nodes { path } } } + pageInfo { hasNextPage endCursor } + } + } + }' \ + --jq '(env.RECHECK_PATHS | fromjson) as $paths | .data.repository.pullRequests.nodes[] + | {number, headRefOid, listed: any(.files.nodes[].path; IN($paths[])), unlisted: (.changedFiles > 100)} + | select(.listed or .unlisted) | "\(.number) \(.headRefOid) \(.listed)"')" + failed=0 + declare -A rerun_heads + while read -r number head_sha listed; do + [ -n "$number" ] || continue + if [ "$listed" = false ]; then + if ! changed="$(gh api --paginate "repos/${GITHUB_REPOSITORY}/pulls/${number}/files?per_page=100" \ + --jq '(env.RECHECK_PATHS | fromjson) as $paths | .[].filename | select(IN($paths[]))')"; then + echo "::error::could not list the files of #$number" + failed=1 + continue + fi + [ -n "$changed" ] || continue + fi + [ -z "${rerun_heads[$head_sha]:-}" ] || continue + rerun_heads[$head_sha]=1 + if ! runs="$(gh run list --repo "$GITHUB_REPOSITORY" --workflow cost-map-guard.yml --commit "$head_sha" --limit 100 --json databaseId,createdAt --jq '.[] | "\(.databaseId) \(.createdAt)"')"; then + echo "::error::could not list the cost map guard runs of #$number" + failed=1 + continue + fi + if [ -z "$runs" ]; then + echo "::warning::no cost map guard run found for #$number" + continue + fi + while read -r run_id created_at; do + if [[ "$created_at" < "$rerun_cutoff" ]]; then + echo "::warning::cost map guard run $run_id for #$number is over 30 days old, so GitHub cannot re-run it; #$number gets the pin checks on its next push" + continue + fi + gh run watch "$run_id" --repo "$GITHUB_REPOSITORY" > /dev/null || true + if ! gh run rerun "$run_id" --repo "$GITHUB_REPOSITORY"; then + echo "::error::could not re-run cost map guard run $run_id for #$number" + failed=1 + fi + done <<< "$runs" + done <<< "$candidates" + exit "$failed" diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index c1aeaabc5d2..d5fd755ba18 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -1,37 +1,67 @@ """Guard the cost map on pull requests. -Every pull request whose diff against its merge base touches one of the two cost map files gets the file -checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched -skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the -branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always -get the file checks and additionally may only touch those two files and may only add or update models. +Every pull request whose diff against its merge base touches one of the two cost map files or the pins file gets the +file checks: the files parse, the backup copy matches the root file, and every capability pinned in +ci_cd/cost_map_pins.json (a value verified against a live provider call where the provider's model listing says +otherwise) still holds in the tree that merging the pull request into its base branch produces, so a branch cut before +a pin was added cannot lower it, and a pull request that deletes the pins file fails too. A pull request that leaves +all three untouched skips them, since merging it keeps the base branch's copies and its head tree only carries +whatever state the branch was cut from. Pull requests from the cost map sync bot (branches named +litellm_cost_map_sync_*) always get the file checks and additionally may only touch the two cost map files and may +only add or update models, so a sync can never lower a pinned capability or edit the pins. Price sync pull requests +(branches named litellm-providers/price-sync*) and pull requests a bot account opened (COST_MAP_GUARD_AUTHOR_TYPE is +Bot) get the human checks and may not touch the pins file either. """ from __future__ import annotations import argparse import json +import os import subprocess import sys -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" +PINS_PATH: Final = "ci_cd/cost_map_pins.json" GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH) +CHECKED_PATHS: Final = (*GUARDED_PATHS, PINS_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" +PRICE_SYNC_BRANCH_PREFIX: Final = "litellm-providers/price-sync" +TARGET_ENV: Final = "COST_MAP_GUARD_TARGET" +AUTHOR_TYPE_ENV: Final = "COST_MAP_GUARD_AUTHOR_TYPE" +BOT_AUTHOR_TYPE: Final = "Bot" +MERGE_CONFLICT_EXIT_CODE: Final = 1 SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({}) CostMap = dict[str, object] +_MISSING: Final = object() @dataclass(frozen=True, slots=True) class Snapshot: cost_map: str backup: str + pins: str + + +@dataclass(frozen=True, slots=True) +class MergeFailed: + stderr: str + + +@dataclass(frozen=True, slots=True) +class Pin: + key: str + field: str + value: object + checked: str + source: str def _parse_object(text: str, path: str) -> CostMap | str: @@ -91,6 +121,48 @@ def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]: return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)} +def _pin(key: str, field: str, spec: object) -> Pin | str: + malformed: Final = f"{PINS_PATH}: {key}.{field} needs value, checked, and source" + if not isinstance(spec, dict) or "value" not in spec: + return malformed + checked: Final = spec.get("checked") + source: Final = spec.get("source") + if not isinstance(checked, str) or not isinstance(source, str) or not checked or not source: + return malformed + return Pin(key, field, spec["value"], checked, source) + + +def _pins(parsed: CostMap) -> Iterator[Pin | str]: + for key, fields in parsed.items(): + if not isinstance(fields, dict): + yield f"{PINS_PATH}: {key} must map field names to pins" + continue + for field, spec in fields.items(): + yield _pin(key, field, spec) + + +def _pin_violations(parsed: CostMap, entries: Mapping[str, Mapping[str, object]]) -> Iterator[str]: + for pin in _pins(parsed): + if isinstance(pin, str): + yield pin + elif pin.key not in entries: + yield f"{PINS_PATH} pins {pin.key}.{pin.field} but {COST_MAP_PATH} has no {pin.key} entry" + elif entries[pin.key].get(pin.field, _MISSING) != pin.value: + yield ( + f"{COST_MAP_PATH}: {pin.key}.{pin.field} must stay {json.dumps(pin.value)} " + f"(verified {pin.checked}: {pin.source}); change {PINS_PATH} with new evidence first" + ) + + +def pin_failures(pins_text: str, head_map: CostMap) -> tuple[str, ...]: + if not pins_text: + return (f"{PINS_PATH} is missing; restore it, its pins were verified against live provider calls",) + parsed: Final = _parse_object(pins_text, PINS_PATH) + if isinstance(parsed, str): + return (parsed,) + return tuple(_pin_violations(parsed, _entries(head_map))) + + def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]: base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH) if isinstance(base_map, str): @@ -120,8 +192,18 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str ) +def pin_author_failures(head_ref: str, author_type: str, changed_files: Sequence[str]) -> tuple[str, ...]: + if PINS_PATH not in changed_files: + return () + if head_ref.startswith(PRICE_SYNC_BRANCH_PREFIX): + return (f"price sync PRs may not change {PINS_PATH}; a pin moves only in a human PR with new evidence",) + if author_type == BOT_AUTHOR_TYPE: + return (f"PRs opened by a bot may not change {PINS_PATH}; a pin moves only in a human PR with new evidence",) + return () + + def touches_cost_map(changed_files: Sequence[str]) -> bool: - return any(path in GUARDED_PATHS for path in changed_files) + return any(path in CHECKED_PATHS for path in changed_files) def contract_for(bot: bool, changed_files: Sequence[str]) -> str: @@ -130,9 +212,21 @@ def contract_for(bot: bool, changed_files: Sequence[str]) -> str: return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched" +def _merged_pin_failures(base: Snapshot, merged: Snapshot | None) -> tuple[str, ...]: + if merged is None: + return ("the pull request does not merge cleanly into its base branch; merge the base branch in first",) + if not merged.pins and not base.pins: + return () + merged_map: Final = _parse_object(merged.cost_map, COST_MAP_PATH) + if isinstance(merged_map, str): + return (f"after merging into the base branch, {merged_map}",) + return pin_failures(merged.pins, merged_map) + + def guard_failures( base: Snapshot, head: Snapshot, + merged: Snapshot | None, changed_files: Sequence[str], bot: bool, *, @@ -147,6 +241,7 @@ def guard_failures( *_file_failures( head, head_map, evaluation_output_price_sources=evaluation_output_price_sources ), + *_merged_pin_failures(base, merged), *(_bot_failures(base, head_map, changed_files) if bot else ()), ) @@ -157,7 +252,22 @@ def _git(*args: str) -> str | None: def snapshot(revision: str) -> Snapshot: - return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in GUARDED_PATHS)) + return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in CHECKED_PATHS)) + + +def merged_tree(merge_base: str, target: str, head: str) -> str | MergeFailed | None: + result: Final = subprocess.run( + ("git", "merge-tree", "--write-tree", f"--merge-base={merge_base}", target, head), + check=False, + capture_output=True, + text=True, + ) + tree: Final = result.stdout.split("\n", 1)[0] + if result.returncode == 0: + return tree + if result.returncode == MERGE_CONFLICT_EXIT_CODE and tree: + return None + return MergeFailed(result.stderr.strip()) def changed_files(base: str, head: str) -> tuple[str, ...] | None: @@ -169,6 +279,13 @@ def main(argv: Sequence[str]) -> int: parser: Final = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base", required=True, help="merge base of the pull request") parser.add_argument("--head", required=True, help="head commit of the pull request") + env_target: Final = os.environ.get(TARGET_ENV) + parser.add_argument( + "--target", + default=env_target, + required=env_target is None, + help=f"tip of the base branch the pull request merges into, defaulting to ${TARGET_ENV}", + ) parser.add_argument("--head-ref", required=True, help="head branch name of the pull request") args: Final = parser.parse_args(argv) bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX) @@ -176,7 +293,19 @@ def main(argv: Sequence[str]) -> int: if changed is None: print(f"cost map guard failed: git diff {args.base} {args.head} failed, so the changed files are unknown") return 1 - failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed, bot) + if not bot and not touches_cost_map(changed): + print(f"cost map guard passed ({contract_for(bot, changed)})") + return 0 + merged: Final = merged_tree(args.base, args.target, args.head) + if isinstance(merged, MergeFailed): + print(f"cost map guard failed: git merge-tree into {args.target} failed: {merged.stderr}") + return 1 + failures: Final = ( + *guard_failures( + snapshot(args.base), snapshot(args.head), None if merged is None else snapshot(merged), changed, bot + ), + *pin_author_failures(args.head_ref, os.environ.get(AUTHOR_TYPE_ENV, ""), changed), + ) contract: Final = contract_for(bot, changed) if failures: print(f"cost map guard failed ({contract}):") diff --git a/ci_cd/cost_map_pins.json b/ci_cd/cost_map_pins.json new file mode 100644 index 00000000000..780cd5f1276 --- /dev/null +++ b/ci_cd/cost_map_pins.json @@ -0,0 +1,16 @@ +{ + "fireworks_ai/accounts/fireworks/models/minimax-m3": { + "supports_vision": { + "value": true, + "checked": "2026-09-26", + "source": "POST https://api.fireworks.ai/inference/v1/chat/completions read a 512x512 PNG back as digit=7, shape=circle, color=red while GET https://api.fireworks.ai/inference/v1/models still said supports_image_input false and GET https://api.fireworks.ai/v1/serverless/models?format=nested carried no supportsImageInput on the row" + } + }, + "fireworks_ai/minimax-m3": { + "supports_vision": { + "value": true, + "checked": "2026-09-26", + "source": "POST https://api.fireworks.ai/inference/v1/chat/completions read a 512x512 PNG back as digit=7, shape=circle, color=red while GET https://api.fireworks.ai/inference/v1/models still said supports_image_input false and GET https://api.fireworks.ai/v1/serverless/models?format=nested carried no supportsImageInput on the row" + } + } +} diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 23c6ccf9c0a..3838f7510f0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28071,7 +28071,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -28345,7 +28345,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 23c6ccf9c0a..3838f7510f0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28071,7 +28071,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -28345,7 +28345,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index 578d5482822..2602149b623 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -1,5 +1,6 @@ import importlib.util import json +import os import subprocess import sys from collections.abc import Mapping @@ -51,12 +52,16 @@ def _serialize(cost_map: dict[str, object]) -> str: return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" -def _snapshot(cost_map: dict[str, object], backup: str | None = None) -> object: +NO_PINS: Final = "{}\n" + + +def _snapshot(cost_map: dict[str, object], backup: str | None = None, pins: str = NO_PINS) -> object: text = _serialize(cost_map) - return guard.Snapshot(cost_map=text, backup=text if backup is None else backup) + return guard.Snapshot(cost_map=text, backup=text if backup is None else backup, pins=pins) BASE: Final = _snapshot(BASE_MAP) +UP_TO_DATE: Final = object() def _failures( @@ -64,10 +69,13 @@ def _failures( changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True, evaluation_output_price_sources: Mapping[str, str] = guard.EVALUATION_OUTPUT_PRICE_SOURCES, + base: object = BASE, + merged: object = UP_TO_DATE, ) -> tuple[str, ...]: return guard.guard_failures( - BASE, + base, head, + head if merged is UP_TO_DATE else merged, changed_files, bot, evaluation_output_price_sources=evaluation_output_price_sources, @@ -129,14 +137,14 @@ def test_bot_may_add_and_reprice_models() -> None: def test_broken_json_is_reported() -> None: - head = guard.Snapshot(cost_map="{not json", backup="{not json") + head = guard.Snapshot(cost_map="{not json", backup="{not json", pins=NO_PINS) assert _failures(head, bot=False) == ( f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", ) def test_non_object_root_is_reported() -> None: - head = guard.Snapshot(cost_map="[]", backup="[]") + head = guard.Snapshot(cost_map="[]", backup="[]", pins=NO_PINS) assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) @@ -154,15 +162,15 @@ CODE_ONLY: Final = ( def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None: - unparseable: Final = guard.Snapshot(cost_map="{not json", backup="") + unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", pins=NO_PINS) assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == () assert _failures(STALE_HEAD, changed_files=(), bot=False) == () assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == () -@pytest.mark.parametrize("guarded_path", guard.GUARDED_PATHS) -def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None: - failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False) +@pytest.mark.parametrize("checked_path", guard.CHECKED_PATHS) +def test_touching_any_cost_map_file_keeps_the_file_checks(checked_path: str) -> None: + failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, checked_path), bot=False) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] @@ -204,8 +212,155 @@ def test_bot_may_not_change_special_root_keys() -> None: assert _failures(head) == ("bot PRs may not change fallback_generalizations",) -def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: +PINNED_SOURCE: Final = "a 512x512 probe image was read back over the provider's chat endpoint" +PINS: Final = _serialize( + {"openrouter/a": {"supports_vision": {"value": True, "checked": "2026-09-26", "source": PINNED_SOURCE}}} +) + + +def test_a_pinned_capability_that_holds_passes_for_humans_and_bots() -> None: + head = _snapshot(BASE_MAP, pins=PINS) + assert _failures(head, bot=False) == () + assert _failures(head, bot=True) == () + + +LOWERED: Final = ( + f"{guard.COST_MAP_PATH}: openrouter/a.supports_vision must stay true (verified 2026-09-26: {PINNED_SOURCE}); " + f"change {guard.PINS_PATH} with new evidence first" +) +NO_VISION_MAP: Final = {**BASE_MAP, "openrouter/a": _entry(supports_vision=False)} + + +@pytest.mark.parametrize("bot", [False, True]) +def test_lowering_a_pinned_capability_fails_with_its_evidence(bot: bool) -> None: + assert _failures(_snapshot(NO_VISION_MAP, pins=PINS), bot=bot) == (LOWERED,) + + +def test_dropping_a_pinned_field_or_model_is_reported() -> None: + no_field = _snapshot({**BASE_MAP, "openrouter/a": _entry()}, pins=PINS) + assert [failure for failure in _failures(no_field, bot=False) if "must stay true" in failure] + no_model = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/a"}, pins=PINS) + assert _failures(no_model, bot=False) == ( + f"{guard.PINS_PATH} pins openrouter/a.supports_vision but {guard.COST_MAP_PATH} has no openrouter/a entry", + ) + + +@pytest.mark.parametrize( + ("pins", "expected"), + [ + ( + "{not json", + f"{guard.PINS_PATH} is not valid JSON: Expecting property name enclosed in double quotes: " + "line 1 column 2 (char 1)", + ), + ( + _serialize({"openrouter/a": {"supports_vision": {"value": True}}}), + f"{guard.PINS_PATH}: openrouter/a.supports_vision needs value, checked, and source", + ), + (_serialize({"openrouter/a": True}), f"{guard.PINS_PATH}: openrouter/a must map field names to pins"), + ], +) +def test_malformed_pins_are_reported(pins: str, expected: str) -> None: + assert _failures(_snapshot(BASE_MAP, pins=pins), bot=False) == (expected,) + + +MISSING_PINS: Final = f"{guard.PINS_PATH} is missing; restore it, its pins were verified against live provider calls" + + +@pytest.mark.parametrize("bot", [False, True]) +def test_a_pr_that_deletes_the_pins_file_fails(bot: bool) -> None: + head: Final = _snapshot(BASE_MAP, pins="") + assert _failures(head, bot=bot, base=_snapshot(BASE_MAP, pins=PINS)) == (MISSING_PINS,) + + +@pytest.mark.parametrize("bot", [False, True]) +def test_pin_checks_are_skipped_when_no_side_of_the_merge_has_the_pins_file(bot: bool) -> None: + assert _failures(_snapshot(NO_VISION_MAP, pins=""), bot=bot, base=_snapshot(NO_VISION_MAP, pins="")) == () + + +@pytest.mark.parametrize("bot", [False, True]) +def test_a_branch_cut_before_a_pin_was_added_cannot_lower_it(bot: bool) -> None: + head: Final = _snapshot(NO_VISION_MAP, pins="") + merged: Final = _snapshot(NO_VISION_MAP, pins=PINS) + assert _failures(head, bot=bot, base=_snapshot(BASE_MAP, pins=""), merged=merged) == (LOWERED,) + + +@pytest.mark.parametrize("bot", [False, True]) +def test_a_branch_cut_before_a_pin_was_added_passes_when_merging_keeps_the_pinned_value(bot: bool) -> None: + repriced: Final = {"openrouter/b": _entry(3e-06)} + head: Final = _snapshot({**NO_VISION_MAP, **repriced}, pins="") + merged: Final = _snapshot({**BASE_MAP, **repriced}, pins=PINS) + assert _failures(head, bot=bot, base=_snapshot(NO_VISION_MAP, pins=""), merged=merged) == () + + +def test_a_pr_that_does_not_merge_cleanly_is_told_to_merge_the_base_branch_in() -> None: + assert _failures(BASE, bot=False, merged=None) == ( + "the pull request does not merge cleanly into its base branch; merge the base branch in first", + ) + + +def test_a_merge_result_that_is_not_valid_json_is_reported() -> None: + merged: Final = guard.Snapshot(cost_map="{broken", backup="{broken", pins=PINS) + assert _failures(BASE, bot=False, merged=merged) == ( + f"after merging into the base branch, {guard.COST_MAP_PATH} is not valid JSON: " + "Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", + ) + + +def test_a_pr_that_adds_the_pins_file_gets_its_pins_checked() -> None: + head: Final = _snapshot({**BASE_MAP, "openrouter/a": _entry(supports_vision=False)}, pins=PINS) + changed: Final = (*MAP_FILES, guard.PINS_PATH) + failures: Final = _failures(head, changed_files=changed, bot=False, base=_snapshot(BASE_MAP, pins="")) + assert [failure for failure in failures if "openrouter/a.supports_vision must stay true" in failure] + + +def test_bot_may_not_touch_the_pins_file() -> None: + head = _snapshot(BASE_MAP, pins=PINS) + changed = (*guard.GUARDED_PATHS, guard.PINS_PATH) + assert _failures(head, changed_files=changed, bot=False) == () + assert _failures(head, changed_files=changed) == ( + f"bot PRs may only change the cost map files, not {guard.PINS_PATH}", + ) + + +PRICE_SYNC_PIN_FAILURE: Final = ( + f"price sync PRs may not change {guard.PINS_PATH}; a pin moves only in a human PR with new evidence" +) +BOT_AUTHOR_PIN_FAILURE: Final = ( + f"PRs opened by a bot may not change {guard.PINS_PATH}; a pin moves only in a human PR with new evidence" +) + + +@pytest.mark.parametrize( + ("head_ref", "author_type", "changed", "expected"), + [ + ("litellm-providers/price-sync-fireworks-ai", "User", (*MAP_FILES, guard.PINS_PATH), (PRICE_SYNC_PIN_FAILURE,)), + ("litellm-providers/price-sync-fireworks-ai", "Bot", MAP_FILES, ()), + ("litellm_azure_ai_mai_image_2_5_pro", "Bot", (*MAP_FILES, guard.PINS_PATH), (BOT_AUTHOR_PIN_FAILURE,)), + ("litellm_azure_ai_mai_image_2_5_pro", "Bot", MAP_FILES, ()), + ("fix-openrouter-price-sync", "User", (*MAP_FILES, guard.PINS_PATH), ()), + ("fix-openrouter-price-sync", "", (*MAP_FILES, guard.PINS_PATH), ()), + ], +) +def test_only_human_opened_non_price_sync_prs_may_touch_the_pins_file( + head_ref: str, author_type: str, changed: tuple[str, ...], expected: tuple[str, ...] +) -> None: + assert guard.pin_author_failures(head_ref, author_type, changed) == expected + + +def test_checked_in_pins_hold_in_the_checked_in_cost_map() -> None: + cost_map = json.loads((ROOT / guard.COST_MAP_PATH).read_text()) + pins = (ROOT / guard.PINS_PATH).read_text() + assert json.loads(pins), "the pins file must pin at least one capability" + assert guard.pin_failures(pins, cost_map) == () + + +def _commit(repo: Path, cost_map: dict[str, object], message: str, pins: str | None = None) -> str: text = _serialize(cost_map) + pins_file = repo / guard.PINS_PATH + if pins is not None or not pins_file.exists(): + pins_file.parent.mkdir(exist_ok=True) + pins_file.write_text(NO_PINS if pins is None else pins) (repo / guard.COST_MAP_PATH).write_text(text) (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) (repo / guard.BACKUP_PATH).write_text(text) @@ -230,13 +385,38 @@ def _git_commit(repo: Path, message: str) -> str: ).stdout.strip() -def _run_guard(repo: Path, base: str, head: str, head_ref: str) -> subprocess.CompletedProcess[str]: +def _run_guard( + repo: Path, + base: str, + head: str, + head_ref: str, + *, + target: str | None = None, + target_env: str | None = None, + author_type: str | None = None, +) -> subprocess.CompletedProcess[str]: + target_args: Final = () if target_env is not None else ("--target", target or base) + env_overrides: Final = { + **({} if target_env is None else {guard.TARGET_ENV: target_env}), + **({} if author_type is None else {guard.AUTHOR_TYPE_ENV: author_type}), + } return subprocess.run( - (sys.executable, str(CI_CD / "cost_map_guard.py"), "--base", base, "--head", head, "--head-ref", head_ref), + ( + sys.executable, + str(CI_CD / "cost_map_guard.py"), + "--base", + base, + "--head", + head, + *target_args, + "--head-ref", + head_ref, + ), cwd=repo, capture_output=True, text=True, check=False, + env={**os.environ, **env_overrides}, ) @@ -298,3 +478,69 @@ def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None: head = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head") assert _run_guard(tmp_path, base, head, BOT_REF).returncode == 1 assert _run_guard(tmp_path, base, head, "litellm_fix_pricing").returncode == 0 + + +@pytest.mark.parametrize("head_ref", [BOT_REF, "litellm_fix_pricing"]) +def test_main_reads_the_pins_from_the_head_revision(tmp_path: Path, head_ref: str) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base: Final = _commit(tmp_path, BASE_MAP, "base", pins=PINS) + head: Final = _commit(tmp_path, {**BASE_MAP, "openrouter/a": _entry(supports_vision=False)}, "sync lowers it") + result: Final = _run_guard(tmp_path, base, head, head_ref) + assert result.returncode == 1, result.stdout + result.stderr + pin_line: Final = f"- {guard.COST_MAP_PATH}: openrouter/a.supports_vision must stay true (verified 2026-09-26: " + assert [line for line in result.stdout.splitlines() if line.startswith(pin_line)] + + +@pytest.mark.parametrize("head_ref", [BOT_REF, "litellm_fix_pricing"]) +def test_main_holds_a_pin_the_base_branch_added_after_the_pr_branched(tmp_path: Path, head_ref: str) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + merge_base: Final = _commit(tmp_path, BASE_MAP, "merge base") + target: Final = _commit(tmp_path, BASE_MAP, "pin on the base branch", pins=PINS) + subprocess.run(("git", "checkout", "-q", merge_base), cwd=tmp_path, check=True) + head: Final = _commit(tmp_path, NO_VISION_MAP, "branch cut before the pin lowers it", pins=NO_PINS) + result: Final = _run_guard(tmp_path, merge_base, head, head_ref, target_env=target) + assert result.returncode == 1, result.stdout + result.stderr + pin_line: Final = f"- {guard.COST_MAP_PATH}: openrouter/a.supports_vision must stay true (verified 2026-09-26: " + assert [line for line in result.stdout.splitlines() if line.startswith(pin_line)] + + +def test_main_reports_a_failed_merge_apart_from_a_conflict(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base: Final = _commit(tmp_path, BASE_MAP, "base") + head: Final = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head") + result: Final = _run_guard(tmp_path, base, head, "litellm_fix_pricing", target="0" * 40) + assert result.returncode == 1, result.stdout + result.stderr + assert result.stdout.startswith(f"cost map guard failed: git merge-tree into {'0' * 40} failed: ") + + +def test_main_tells_a_conflicting_pr_to_merge_the_base_branch_in(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + merge_base: Final = _commit(tmp_path, BASE_MAP, "merge base") + target: Final = _commit(tmp_path, {**BASE_MAP, "openrouter/b": _entry(3e-06)}, "base branch reprices b") + subprocess.run(("git", "checkout", "-q", merge_base), cwd=tmp_path, check=True) + head: Final = _commit(tmp_path, {**BASE_MAP, "openrouter/b": _entry(4e-06)}, "pr reprices b differently") + result: Final = _run_guard(tmp_path, merge_base, head, "litellm_fix_pricing", target=target) + assert result.returncode == 1, result.stdout + result.stderr + conflict_line: Final = "- the pull request does not merge cleanly into its base branch; merge the base branch in first" + assert conflict_line in result.stdout.splitlines() + + +def test_main_skips_the_merge_for_a_pr_that_leaves_the_cost_map_alone(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base: Final = _commit(tmp_path, BASE_MAP, "base") + head: Final = _commit_code_only(tmp_path, "code change") + result: Final = _run_guard(tmp_path, base, head, "litellm_fix_pricing", target="0" * 40) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.splitlines() == ["cost map guard passed (human PR, cost map untouched)"] + + +@pytest.mark.parametrize(("author_type", "expected_code"), [("Bot", 1), ("User", 0)]) +def test_main_lets_only_a_human_opened_pr_change_the_pins(tmp_path: Path, author_type: str, expected_code: int) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base: Final = _commit(tmp_path, BASE_MAP, "base", pins=PINS) + head: Final = _commit(tmp_path, NO_VISION_MAP, "lower the flag and drop its pin", pins=NO_PINS) + result: Final = _run_guard( + tmp_path, base, head, "litellm_azure_ai_mai_image_2_5_pro", author_type=author_type + ) + assert result.returncode == expected_code, result.stdout + result.stderr + assert (f"- {BOT_AUTHOR_PIN_FAILURE}" in result.stdout.splitlines()) == (author_type == "Bot") From 65d912fc2e1e0ca81b228836859c70bbf0d05081 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:29:40 -0700 Subject: [PATCH 05/58] feat(responses): support Codex client tool_search on Anthropic and other bridged providers (#48972) * feat(responses): support Codex client tool_search on Anthropic and other bridged providers * test(router): cover get_search_tool_support provider defaults and model_info overrides * fix(responses): keep tools loaded by earlier Codex tool searches across previous_response_id * perf(responses): read stored turns once on client tool search continuations * fix(router): avoid OAuth refreshes when listing search capabilities * test(cli): expect disabled search on fallback Codex models --- .../proxy/common_utils/codex_model_catalog.py | 13 + .../client_tool_search.py | 268 +++++++++++++ .../handler.py | 109 ++++- .../session_handler.py | 70 +++- .../streaming_iterator.py | 49 ++- .../tool_search.py | 47 ++- .../transformation.py | 26 +- litellm/router.py | 32 ++ .../router_code_coverage.py | 1 + .../router_surface_allowlist.json | 1 + .../coverage_registry/llm_conversational.yaml | 1 + .../test_responses_codex_tool_search_e2e.py | 205 ++++++++++ tests/unit/proxy/client/cli/test_agents.py | 4 +- .../common_utils/test_codex_model_catalog.py | 60 ++- .../proxy/proxy_server/test_routes_models.py | 53 +-- .../test_anthropic_responses_bridge.py | 377 +++++++++++++++++- .../test_client_tool_search.py | 68 ++++ .../test_litellm_completion_responses.py | 50 ++- tests/unit/test_router/test_router.py | 75 ++++ 19 files changed, 1385 insertions(+), 124 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/client_tool_search.py create mode 100644 tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py create mode 100644 tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 397523288cf..47cd1db9f39 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -102,6 +102,7 @@ class CodexFallbackModel(LiteLLMBaseModel): experimental_supported_tools: tuple[()] = () context_window: int | None base_instructions: str + supports_search_tool: bool = False class CodexStockUpgrade(LiteLLMBaseModel): @@ -134,6 +135,7 @@ class CodexStockModel(LiteLLMBaseModel): upgrade: CodexStockUpgrade | None = None service_tiers: tuple[CodexServiceTier, ...] = () default_service_tier: str | None = None + supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None @@ -160,6 +162,7 @@ class CodexCatalogRow: reasoning_effort: str | None = None display_name: str | None = None service_tiers: tuple[object, ...] = () + search_tool_support: tuple[bool | None, ...] = () @cache @@ -252,6 +255,12 @@ def configured_service_tiers( return offered +def offered_search_tool(per_deployment: Sequence[bool | None]) -> bool | None: + if False in per_deployment: + return False + return True if per_deployment and None not in per_deployment else None + + def _codex_can_drive(row: CodexCatalogRow) -> bool: return "*" not in row.id and (row.mode is None or row.mode in CODEX_CHAT_MODES) @@ -306,6 +315,7 @@ def _entry( {tier.id: tier for tier in (stock_model.service_tiers if stock_model is not None else ())} ) configured: Final = configured_service_tiers(row.service_tiers, row.id, stock_tiers) + search_tool: Final = offered_search_tool(row.search_tool_support) if stock_model is None: return CodexFallbackModel( slug=row.id, @@ -314,6 +324,7 @@ def _entry( service_tiers=configured if configured is not None else (), context_window=row.max_input_tokens, base_instructions=instructions, + supports_search_tool=search_tool is True, ) upgrade: Final = ( stock_model.upgrade if stock_model.upgrade is not None and stock_model.upgrade.model in served else None @@ -327,6 +338,7 @@ def _entry( "visibility": "list", "supported_in_api": True, "upgrade": upgrade, + **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), } if configured is None: @@ -433,6 +445,7 @@ def _catalog_row( upstream_model=llm_router.get_routable_upstream_model(lookup_id, team_id) if llm_router is not None else None, display_name=llm_router.get_configured_display_name(label_id) if llm_router is not None else None, service_tiers=llm_router.get_configured_service_tiers(lookup_id, team_id) if llm_router is not None else (), + search_tool_support=llm_router.get_search_tool_support(lookup_id, team_id) if llm_router is not None else (), ) diff --git a/litellm/responses/litellm_completion_transformation/client_tool_search.py b/litellm/responses/litellm_completion_transformation/client_tool_search.py new file mode 100644 index 00000000000..030bb014b7a --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/client_tool_search.py @@ -0,0 +1,268 @@ +import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, cast # noqa: TID251 # loaded tool specs are forwarded as the client sent them + +from openai.types.responses.response_input_param import ( + FunctionCallOutput, + ResponseFunctionToolCallParam, + ResponseInputItemParam, +) +from pydantic import ConfigDict, TypeAdapter, ValidationError + +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.openai import ( + ALL_RESPONSES_API_TOOL_PARAMS, + ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, + ResponseInputParam, +) + +from .custom_tools import function_chat_tool +from .tool_search import CLIENT_TOOL_SEARCH_FUNCTION_NAME, has_client_tool_search + +__all__ = [ + "ClientToolSearchBridge", + "bridge_client_tool_search", + "client_tool_search_chat_tool", + "client_tool_search_loaded_tools", +] + +_DEFAULT_SEARCH_DESCRIPTION: Final = ( + "Search the deferred tools by keyword. The matching tools are loaded and become callable on the next turn." +) +_DEFAULT_SEARCH_PARAMETERS: Final = MappingProxyType( + { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query for deferred tools."}, + "limit": {"type": "integer", "description": "Maximum number of tools to return. Defaults to 8."}, + }, + "required": ["query"], + "additionalProperties": False, + } +) + + +class _ClientToolSearchItem(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + execution: str | None = None + call_id: str | None = None + arguments: object = None + tools: tuple[Mapping[str, object], ...] = () + + +class _ToolSpec(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + name: str = "" + description: str = "" + defer_loading: object = None + tools: tuple[Mapping[str, object], ...] = () + + +class _ClientToolSearchTool(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + description: str = "" + parameters: Mapping[str, object] | None = None + + +@dataclass(frozen=True, slots=True) +class _LoadedTool: + namespace: str | None + namespace_description: str + spec: Mapping[str, object] + + @property + def callable_name(self) -> str: + name: Final = str(self.spec.get("name") or "") + return name if self.namespace is None else f"{self.namespace}__{name}" + + +@dataclass(frozen=True, slots=True) +class ClientToolSearchBridge: + input: str | ResponseInputParam + tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + bridged: bool + + +def client_tool_search_chat_tool(tool: Mapping[str, object]) -> ChatCompletionToolParam: + spec: Final = _ClientToolSearchTool.model_validate(tool) + function: Final = ChatCompletionToolParamFunctionChunk( + name=CLIENT_TOOL_SEARCH_FUNCTION_NAME, + description=spec.description or _DEFAULT_SEARCH_DESCRIPTION, + parameters=dict(spec.parameters or _DEFAULT_SEARCH_PARAMETERS), + ) + return function_chat_tool(function, None, False) + + +def _client_tool_search_item(item: object) -> _ClientToolSearchItem | None: + try: + parsed: Final = _ClientToolSearchItem.model_validate(item) + except ValidationError: + return None + if parsed.type not in ("tool_search_call", "tool_search_output") or parsed.execution != "client": + return None + return parsed + + +_INPUT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def _searched_tools(item: object) -> tuple[Mapping[str, object], ...]: + parsed: Final = _client_tool_search_item(item) + if parsed is None or parsed.type != "tool_search_output" or not parsed.call_id: + return () + return parsed.tools + + +def client_tool_search_loaded_tools(input: object) -> tuple[Mapping[str, object], ...]: + if not isinstance(input, (list, tuple)): + return () + return tuple(chain.from_iterable(map(_searched_tools, _INPUT_ITEMS.validate_python(input)))) + + +def _tool_spec(tool: object) -> _ToolSpec | None: + try: + return _ToolSpec.model_validate(tool) + except ValidationError: + return None + + +def _as_tool_param(spec: Mapping[str, object]) -> ALL_RESPONSES_API_TOOL_PARAMS: + return cast("ALL_RESPONSES_API_TOOL_PARAMS", spec) # cast-ok: client tool specs are forwarded as sent + + +def _without_deferral(spec: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in spec.items() if key != "defer_loading"} + + +def _loaded_tools(tool: Mapping[str, object]) -> tuple[_LoadedTool, ...]: + spec: Final = _tool_spec(tool) + if spec is None: + return () + if spec.type != "namespace": + return (_LoadedTool(namespace=None, namespace_description="", spec=_without_deferral(tool)),) + return tuple( + _LoadedTool(namespace=spec.name, namespace_description=spec.description, spec=_without_deferral(member)) + for member in spec.tools + ) + + +def _all_loaded_tools(tools: Iterable[Mapping[str, object]]) -> tuple[_LoadedTool, ...]: + return tuple(chain.from_iterable(map(_loaded_tools, tools))) + + +def _callable_names(tool: object) -> tuple[str, ...]: + spec: Final = _tool_spec(tool) + if spec is None: + return () + if spec.type != "namespace": + return (spec.name,) + return tuple(f"{spec.name}__{member.get('name') or ''}" for member in spec.tools) + + +def _first_occurrences(loaded: Sequence[_LoadedTool], taken: frozenset[str]) -> tuple[_LoadedTool, ...]: + first_index: Final = {tool.callable_name: index for index, tool in reversed(tuple(enumerate(loaded)))} + return tuple( + tool + for index, tool in enumerate(loaded) + if first_index[tool.callable_name] == index and tool.callable_name not in taken + ) + + +def _regrouped(loaded: Sequence[_LoadedTool]) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + namespaces: Final = {tool.namespace: tool.namespace_description for tool in reversed(loaded) if tool.namespace} + ordered_namespaces: Final = tuple(dict.fromkeys(tool.namespace for tool in loaded if tool.namespace)) + top_level: Final = tuple(tool.spec for tool in loaded if tool.namespace is None) + grouped: Final = tuple( + { + "type": "namespace", + "name": namespace, + "description": namespaces[namespace], + "tools": tuple(tool.spec for tool in loaded if tool.namespace == namespace), + } + for namespace in ordered_namespaces + ) + return tuple(_as_tool_param(tool) for tool in (*top_level, *grouped)) + + +def _is_deferred(tool: object) -> bool: + spec: Final = _tool_spec(tool) + return spec is not None and spec.defer_loading is True + + +def _without_deferred(tool: ALL_RESPONSES_API_TOOL_PARAMS) -> ALL_RESPONSES_API_TOOL_PARAMS | None: + spec: Final = _tool_spec(tool) + if spec is None: + return tool + if spec.type != "namespace": + return None if spec.defer_loading is True else tool + eager_members: Final = tuple(member for member in spec.tools if not _is_deferred(member)) + if len(eager_members) == len(spec.tools): + return tool + if not eager_members: + return None + return _as_tool_param({**tool, "tools": eager_members}) + + +def _search_output_text(loaded: Sequence[_LoadedTool]) -> str: + if not loaded: + return "No tools matched the search." + return "Loaded these tools, which are now callable: " + ", ".join(tool.callable_name for tool in loaded) + + +def _search_arguments_json(arguments: object) -> str: + return arguments if isinstance(arguments, str) else json.dumps(arguments or {}) + + +def _bridged_item(item: ResponseInputItemParam) -> ResponseInputItemParam | None: + parsed: Final = _client_tool_search_item(item) + if parsed is None: + return item + if not parsed.call_id: + return None + if parsed.type == "tool_search_call": + return ResponseFunctionToolCallParam( + type="function_call", + call_id=parsed.call_id, + name=CLIENT_TOOL_SEARCH_FUNCTION_NAME, + arguments=_search_arguments_json(parsed.arguments), + ) + return FunctionCallOutput( + type="function_call_output", call_id=parsed.call_id, output=_search_output_text(_all_loaded_tools(parsed.tools)) + ) + + +def bridge_client_tool_search( + input: str | ResponseInputParam, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, + loaded_earlier: Sequence[Mapping[str, object]] = (), +) -> ClientToolSearchBridge: + existing: Final = tuple(tools or ()) + searches_on_client: Final = has_client_tool_search(existing) + items: Final = () if isinstance(input, str) else tuple(input) + search_items: Final = tuple(parsed for item in items if (parsed := _client_tool_search_item(item)) is not None) + if not searches_on_client and not search_items: + return ClientToolSearchBridge(input=input, tools=existing, bridged=False) + eager: Final = ( + tuple(kept for tool in existing if (kept := _without_deferred(tool)) is not None) + if searches_on_client + else existing + ) + loaded: Final = _all_loaded_tools((*loaded_earlier, *client_tool_search_loaded_tools(input))) + taken: Final = frozenset(chain.from_iterable(map(_callable_names, eager))) + bridged_input: Final = ( + input if isinstance(input, str) else [bridged for item in input if (bridged := _bridged_item(item)) is not None] + ) + return ClientToolSearchBridge( + input=bridged_input, + tools=(*eager, *_regrouped(_first_occurrences(loaded, taken))), + bridged=True, + ) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index ea3bda1b8b5..141e97d902d 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -2,17 +2,21 @@ Handler for transforming responses api requests to litellm.completion requests """ -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass from typing import Final from pydantic import TypeAdapter import litellm -from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools +from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, StoredTurns from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.litellm_completion_transformation.tool_search import has_client_tool_search from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, with_compaction_items_as_params, @@ -28,6 +32,28 @@ from litellm.types.utils import ModelResponse _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +@dataclass(frozen=True, slots=True) +class _BridgedRequest: + input: str | ResponseInputParam + responses_api_request: ResponsesAPIOptionalRequestParams + + +def _bridged_request( + hoisted: HoistedAdditionalTools, + responses_api_request: ResponsesAPIOptionalRequestParams, + loaded_earlier: Sequence[Mapping[str, object]], +) -> _BridgedRequest: + searched: Final = bridge_client_tool_search(hoisted.input, hoisted.tools, loaded_earlier) + return _BridgedRequest( + input=searched.input, + responses_api_request=( + {**responses_api_request, "tools": list(searched.tools)} + if hoisted.hoisted or searched.bridged + else responses_api_request + ), + ) + + class LiteLLMCompletionTransformationHandler: def response_api_handler( self, @@ -48,10 +74,22 @@ class LiteLLMCompletionTransformationHandler: | Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] ): hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools")) - bridged_input: Final = hoisted.input - bridged_request: Final[ResponsesAPIOptionalRequestParams] = ( - {**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request - ) + previous_response_id: Final = responses_api_request.get("previous_response_id") + if _is_async and previous_response_id and has_client_tool_search(hoisted.tools): + return self._async_response_api_handler_after_earlier_searches( + previous_response_client_ids, + previous_response_id, + hoisted, + model=model, + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ) + bridged: Final = _bridged_request(hoisted, responses_api_request, ()) + bridged_input: Final = bridged.input + bridged_request: Final = bridged.responses_api_request litellm_completion_request: Final[dict] = ( LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, @@ -73,13 +111,8 @@ class LiteLLMCompletionTransformationHandler: **kwargs, ) - completion_args: Final = {} - completion_args.update(kwargs) - completion_args.update(litellm_completion_request) - completion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = litellm.completion( - **completion_args, + **{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}, ) if isinstance(litellm_completion_response, ModelResponse): @@ -104,9 +137,53 @@ class LiteLLMCompletionTransformationHandler: ) raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") + async def _async_response_api_handler_after_earlier_searches( + self, + previous_response_client_ids: tuple[str, ...], + previous_response_id: str, + hoisted: HoistedAdditionalTools, + /, + *, + model: str, + responses_api_request: ResponsesAPIOptionalRequestParams, + custom_llm_provider: str | None, + stream: bool | None, + extra_headers: Mapping[str, object] | None, + request_kwargs: Mapping[str, object], + ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: + stored_turns: Final = await ResponsesSessionHandler.get_stored_turns_for_previous_response_id( + previous_response_id, previous_response_client_ids + ) + bridged: Final = _bridged_request( + hoisted, + responses_api_request, + ResponsesSessionHandler.client_tool_search_tools_loaded_before( + stored_turns, previous_response_id, previous_response_client_ids + ), + ) + return await self.async_response_api_handler( + previous_response_client_ids, + stored_turns, + litellm_completion_request=( + LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=bridged.input, + responses_api_request=bridged.responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + **request_kwargs, + ) + ), + request_input=bridged.input, + responses_api_request=bridged.responses_api_request, + **request_kwargs, + ) + async def async_response_api_handler( self, previous_response_client_ids: tuple[str, ...] = (), + stored_turns: StoredTurns | None = None, /, *, litellm_completion_request: dict, @@ -122,16 +199,12 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, + stored_turns=stored_turns, input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) - acompletion_args: Final = {} - acompletion_args.update(kwargs) - acompletion_args.update(litellm_completion_request) - acompletion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **acompletion_args, + **{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}, ) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 1ab6293b432..ce2396355a2 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -9,6 +9,10 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.responses.litellm_completion_transformation.client_tool_search import ( + bridge_client_tool_search, + client_tool_search_loaded_tools, +) from litellm.responses.litellm_completion_transformation.compaction_items import ( holds_compaction_summary, input_items, @@ -88,6 +92,12 @@ class _ReplayStart: at_output: bool +@dataclass(frozen=True, slots=True) +class StoredTurns: + spend_logs: tuple["SpendLogsPayload", ...] + proxy_server_requests: tuple[Mapping[str, object] | None, ...] + + def _compaction_point( index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None ) -> _ReplayStart | None: @@ -190,6 +200,7 @@ class ResponsesSessionHandler: previous_response_id: str, item_cache_control: bool = False, client_response_ids: tuple[str, ...] = (), + stored_turns: StoredTurns | None = None, ) -> ChatCompletionSession: """ Return the chat completion message history for a previous response id @@ -200,23 +211,21 @@ class ResponsesSessionHandler: ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") - all_spend_logs: Final = [ - spend_log - for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + turns: Final = ( + stored_turns + if stored_turns is not None + else await ResponsesSessionHandler.get_stored_turns_for_previous_response_id( previous_response_id, client_response_ids ) - if _is_a_conversation_turn(spend_log) - ] + ) + all_spend_logs: Final = turns.spend_logs verbose_proxy_logger.debug("found %s spend logs for this response id", len(all_spend_logs)) litellm_session_id: str | None = None if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") - proxy_server_requests: Final = [ - await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) - for spend_log in all_spend_logs - ] + proxy_server_requests: Final = turns.proxy_server_requests replayable_turns: Final = _turns_through_referenced_response( all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) ) @@ -252,6 +261,45 @@ class ResponsesSessionHandler: instructions=_latest_replayable_instructions(proxy_server_requests), ) + @staticmethod + async def get_stored_turns_for_previous_response_id( + previous_response_id: str, + client_response_ids: tuple[str, ...] = (), + ) -> StoredTurns: + spend_logs: Final = tuple( + spend_log + for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id, client_response_ids + ) + if _is_a_conversation_turn(spend_log) + ) + return StoredTurns( + spend_logs=spend_logs, + proxy_server_requests=tuple( + [ + await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) + for spend_log in spend_logs + ] + ), + ) + + @staticmethod + def client_tool_search_tools_loaded_before( + stored_turns: StoredTurns, + previous_response_id: str, + client_response_ids: tuple[str, ...] = (), + ) -> tuple[Mapping[str, object], ...]: + turns: Final = _turns_through_referenced_response( + stored_turns.spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) + ) + return tuple( + chain.from_iterable( + client_tool_search_loaded_tools(request.get("input")) + for request in stored_turns.proxy_server_requests[:turns] + if request is not None + ) + ) + @staticmethod def _history_of_turns( spend_logs: Sequence["SpendLogsPayload"], @@ -310,7 +358,7 @@ class ResponsesSessionHandler: if response_input_param: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=response_input_param, + input=bridge_client_tool_search(response_input_param, None).input, responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, @@ -325,7 +373,7 @@ class ResponsesSessionHandler: if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=responses_api_output, + input=bridge_client_tool_search(responses_api_output, None).input, responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 1e8ca85fdbb..9f0058a5958 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,7 +1,8 @@ import time import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from itertools import filterfalse +from types import MappingProxyType from typing import Any, Final, cast import litellm @@ -18,9 +19,9 @@ from litellm.responses.litellm_completion_transformation.custom_tools import ( ) from litellm.responses.litellm_completion_transformation.reasoning_items import mint_reasoning_item_id from litellm.responses.litellm_completion_transformation.tool_search import ( + ToolSearchExecution, build_tool_search_call, - has_server_tool_search, - is_anthropic_tool_search_use, + tool_search_routing, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -192,16 +193,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) self._web_search_calls: dict[str, object] = {} # mutable-ok: latest call by provider id self._queued_web_search_call_ids: set[str] = set() # mutable-ok: emitted call ids - self._tool_search_requested: bool = has_server_tool_search(request_tools) - self._tool_search_call_ids: frozenset[str] = frozenset() + self._tool_search = tool_search_routing(request_tools) + self._tool_search_executions: Mapping[str, ToolSearchExecution] = MappingProxyType({}) - def _is_tool_search_call(self, call_id: str, fn_name: str) -> bool: - if call_id in self._tool_search_call_ids: - return True - if not self._tool_search_requested or not is_anthropic_tool_search_use(call_id, fn_name): - return False - self._tool_search_call_ids = self._tool_search_call_ids | {call_id} - return True + def _tool_search_execution(self, call_id: str, fn_name: str) -> ToolSearchExecution | None: + known: Final = self._tool_search_executions.get(call_id) + if known is not None: + return known + execution: Final = self._tool_search.execution_of(call_id, fn_name) + if execution is None: + return None + self._tool_search_executions = MappingProxyType({**self._tool_search_executions, call_id: execution}) + return execution def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: self._sequence_number += 1 @@ -216,18 +219,22 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: self._pending_tool_events.append(self._next_output_item_added(output_index, item)) - def _queue_tool_search_call_added(self, call_id: str, output_index: int, args_delta: str) -> None: + def _queue_tool_search_call_added( + self, call_id: str, output_index: int, args_delta: str, execution: ToolSearchExecution + ) -> None: if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" - item: Final = build_tool_search_call(call_id, "", "in_progress") + item: Final = build_tool_search_call(call_id, "", "in_progress", execution) self._tool_item_id_by_call_id[call_id] = item.id self._queue_output_item_added(output_index, BaseLiteLLMOpenAIResponseObject(**item.model_dump())) self._tool_args_by_call_id[call_id] += args_delta - def _queue_tool_search_call_done(self, call_id: str, output_index: int, final_args: str) -> None: - self._queue_tool_search_call_added(call_id, output_index, "") + def _queue_tool_search_call_done( + self, call_id: str, output_index: int, final_args: str, execution: ToolSearchExecution + ) -> None: + self._queue_tool_search_call_added(call_id, output_index, "", execution) item: Final = build_tool_search_call( - call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed" + call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution ) self._sequence_number += 1 self._pending_tool_events.append( @@ -436,8 +443,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_held_annotations() output_index = self._get_or_assign_tool_output_index(call_id) - if self._is_tool_search_call(call_id, fn_name): - self._queue_tool_search_call_added(call_id, output_index, fn_args_delta) + if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None: + self._queue_tool_search_call_added(call_id, output_index, fn_args_delta, search_execution) continue self._started_tool_call = True @@ -505,8 +512,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_web_search_events(call_id, web_search_call) self._queued_web_search_call_ids.add(call_id) continue - if self._is_tool_search_call(call_id, fn_name): - self._queue_tool_search_call_done(call_id, output_index, fn_args) + if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None: + self._queue_tool_search_call_done(call_id, output_index, fn_args, search_execution) continue # Track if this is a new tool call that wasn't streamed diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 780a23e02b2..5a1b27f9092 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -1,4 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence, Set +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -18,25 +19,31 @@ if TYPE_CHECKING: __all__ = [ "ANTHROPIC_TOOL_SEARCH_TOOL_NAME", "ANTHROPIC_TOOL_SEARCH_TOOL_TYPE", + "CLIENT_TOOL_SEARCH_FUNCTION_NAME", "CompletionRoute", "ToolSearchCallStatus", + "ToolSearchExecution", + "ToolSearchRouting", "answers_tool_search", "anthropic_tool_search_chat_tool", "build_tool_search_call", "chat_config_builds_anthropic_request", "chat_messages_without_tool_search_calls", "chat_route_rejects_unanswered_tool_calls", + "has_client_tool_search", "has_server_tool_search", "holds_tool_search_use", "input_without_unanswered_anthropic_tool_search_calls", "is_anthropic_tool_search_chat_tool", "is_anthropic_tool_search_use", + "is_client_tool_search_tool", "is_openai_tool_search_tool", "is_server_tool_search_tool", "provider_builds_anthropic_request", "provider_chat_config", "runs_tool_search_on_the_server", "tool_search_call_ids", + "tool_search_routing", "tools_with_anthropic_tool_search_mapped", "tools_without_anthropic_tool_search", "unanswered_tool_search_call_ids", @@ -46,7 +53,9 @@ __all__ = [ ANTHROPIC_TOOL_SEARCH_TOOL_TYPE: Final = "tool_search_tool_regex_20251119" ANTHROPIC_TOOL_SEARCH_TOOL_NAME: Final = "tool_search_tool_regex" +CLIENT_TOOL_SEARCH_FUNCTION_NAME: Final = "tool_search" ToolSearchCallStatus: TypeAlias = Literal["in_progress", "completed", "incomplete"] +ToolSearchExecution: TypeAlias = Literal["server", "client"] _ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) _MESSAGE_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -79,6 +88,35 @@ def has_server_tool_search(tools: Sequence[object] | None) -> bool: return any(is_server_tool_search_tool(tool) for tool in tools or ()) +def is_client_tool_search_tool(tool: object) -> bool: + try: + parsed: Final = _ToolSearchFields.model_validate(tool) + except ValidationError: + return False + return parsed.type == "tool_search" and parsed.execution == "client" + + +def has_client_tool_search(tools: Sequence[object] | None) -> bool: + return any(is_client_tool_search_tool(tool) for tool in tools or ()) + + +@dataclass(frozen=True, slots=True) +class ToolSearchRouting: + server: bool + client: bool + + def execution_of(self, tool_id: str, tool_name: str) -> ToolSearchExecution | None: + if self.server and is_anthropic_tool_search_use(tool_id, tool_name): + return "server" + if self.client and tool_name == CLIENT_TOOL_SEARCH_FUNCTION_NAME: + return "client" + return None + + +def tool_search_routing(tools: Sequence[object] | None) -> ToolSearchRouting: + return ToolSearchRouting(server=has_server_tool_search(tools), client=has_client_tool_search(tools)) + + def anthropic_tool_search_chat_tool() -> AnthropicToolSearchToolRegex: return AnthropicToolSearchToolRegex(type=ANTHROPIC_TOOL_SEARCH_TOOL_TYPE, name=ANTHROPIC_TOOL_SEARCH_TOOL_NAME) @@ -346,13 +384,18 @@ def _parsed_arguments(arguments_json: str) -> Mapping[str, object] | None: return None -def build_tool_search_call(tool_id: str, arguments_json: str, status: ToolSearchCallStatus) -> OutputToolSearchCall: +def build_tool_search_call( + tool_id: str, + arguments_json: str, + status: ToolSearchCallStatus, + execution: ToolSearchExecution = "server", +) -> OutputToolSearchCall: arguments: Final = _parsed_arguments(arguments_json) return OutputToolSearchCall( id=openai_shaped_tool_call_item_id("tool_search_call", tool_id), call_id=tool_id, arguments={} if arguments is None else arguments, - execution="server", + execution=execution, status="incomplete" if arguments is None else status, type="tool_search_call", ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 8b076ae922e..ce8841c9190 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -64,6 +64,7 @@ from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, + StoredTurns, ) from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex from litellm.types.llms.base import CachedTokensDetails @@ -124,6 +125,7 @@ from litellm.types.utils import ( Usage, ) +from .client_tool_search import client_tool_search_chat_tool from .compaction_items import ( carries_only_compaction, compaction_output_item, @@ -151,16 +153,17 @@ from .custom_tools import ( from .reasoning_items import decode_thinking_blocks, encode_thinking_blocks, mint_reasoning_item_id from .tool_search import ( CompletionRoute, + ToolSearchRouting, answers_tool_search, anthropic_tool_search_chat_tool, build_tool_search_call, has_server_tool_search, holds_tool_search_use, is_anthropic_tool_search_chat_tool, - is_anthropic_tool_search_use, is_server_tool_search_tool, provider_builds_anthropic_request, tool_search_call_ids, + tool_search_routing, without_tool_search_calls, without_tool_search_uses, ) @@ -1342,6 +1345,7 @@ class LiteLLMCompletionResponsesConfig: request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), input_replaces_history: bool = False, + stored_turns: StoredTurns | None = None, ) -> dict[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1357,6 +1361,7 @@ class LiteLLMCompletionResponsesConfig: ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, item_cache_control=keeps_item_marks, client_response_ids=previous_response_client_ids, + stored_turns=stored_turns, ) chat_completion_session: Final = ( await lookup(previous_response_id) @@ -2818,9 +2823,11 @@ class LiteLLMCompletionResponsesConfig: if tool_type == "custom": converted: Final = convert_custom_tool_to_function_tool(tool) return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) - if tool_type == "tool_search" and validated_tool_search_execution(tool.get("execution")) != "client": + if tool_type == "tool_search" and validated_tool_search_execution(tool.get("execution")) == "client": + return ResponsesToolChatForm(chat_tools=(client_tool_search_chat_tool(tool),), web_search_options=None) + if tool_type == "tool_search": return ResponsesToolChatForm(chat_tools=(anthropic_tool_search_chat_tool(),), web_search_options=None) - if tool_type in ("computer_use", "image_generation", "local_shell", "shell", "tool_search"): + if tool_type in ("computer_use", "image_generation", "local_shell", "shell"): verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", @@ -2960,7 +2967,7 @@ class LiteLLMCompletionResponsesConfig: request_tools: Final = responses_api_request.get("tools") if responses_api_request is not None else None custom_tool_names: Final = extract_custom_tool_names(request_tools) namespace_tool_names: Final = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) - tool_search_requested: Final = has_server_tool_search(request_tools) + tool_search: Final = tool_search_routing(request_tools) web_search_calls: Final = LiteLLMCompletionResponsesConfig._web_search_calls_by_call_id( chat_completion_response @@ -2977,7 +2984,7 @@ class LiteLLMCompletionResponsesConfig: if ( server_tool_call := LiteLLMCompletionResponsesConfig._server_tool_call_item( - tool_id, tool_name, tool_arguments, web_search_calls, tool_search_requested + tool_id, tool_name, tool_arguments, web_search_calls, tool_search ) ) is not None: responses_tools.append(server_tool_call) @@ -3042,14 +3049,15 @@ class LiteLLMCompletionResponsesConfig: tool_name: str, tool_arguments: str, web_search_calls: Mapping[str, ResponseFunctionWebSearch], - tool_search_requested: bool, + tool_search: ToolSearchRouting, ) -> ResponseFunctionWebSearch | OutputToolSearchCall | None: web_search_call: Final = web_search_calls.get(tool_id) if web_search_call is not None: return web_search_call - if tool_search_requested and is_anthropic_tool_search_use(tool_id, tool_name): - return build_tool_search_call(tool_id, tool_arguments, "completed") - return None + execution: Final = tool_search.execution_of(tool_id, tool_name) + if execution is None: + return None + return build_tool_search_call(tool_id, tool_arguments, "completed", execution) @staticmethod def _web_search_calls_by_call_id( diff --git a/litellm/router.py b/litellm/router.py index e01f6474ded..02ad6275b82 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -723,6 +723,29 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: + from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request + + configured: Final = _configured_model_info(deployment).get("supports_search_tool") + if isinstance(configured, bool): + return configured + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + model: Final = params.get("model") + provider: Final = params.get("custom_llm_provider") + if not isinstance(model, str): + return None + try: + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=provider if isinstance(provider, str) else None + ) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=resolved_model, provider=LlmProviders(resolved_provider) + ) + except (litellm.BadRequestError, ValueError): + return None + return True if chat_config_builds_anthropic_request(chat_config) else None + + def _every_deployment_refused( deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] ) -> bool: @@ -12234,6 +12257,15 @@ class Router: for deployment in self._routable_deployments(model_name, team_id) ) + def get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: + """ + Whether each deployment a request from team_id to model_name can route + to (see _routable_deployments) answers Codex's client `tool_search`: + its model_info.supports_search_tool when that is a bool, else True when + its route builds an Anthropic Messages request, else None (unknown). + """ + return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) + def get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: """ Return the `litellm_params.model` of the first deployment a request diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index f936d23188c..c7781819ae4 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -107,6 +107,7 @@ ignored_function_names = [ "_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py "_request_fallback_list", # Tested through every mid-stream retry test in test_router.py "_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call + "_search_tool_support", # Tested through the test_get_search_tool_support_* tests "_mid_stream_retry_trigger", # Tested through the retry policy mid-stream budget test in test_router.py "_anthropic_messages_group_retry_policy", # Tested through the retry budget precedence test in test_router.py "_anthropic_messages_resolved_retry_policy", # Tested through the malformed retry policy tests in test_router.py diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json index 428fd9b4011..52e5bb43625 100644 --- a/tests/code_coverage_tests/router_surface_allowlist.json +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -102,6 +102,7 @@ "get_routable_upstream_model", "get_router_deployment_budget_limiter", "get_routing_group", + "get_search_tool_support", "get_settings", "get_valid_args", "guardrail_list", diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index c3c5330c440..0e450b08ebe 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -155,3 +155,4 @@ - {id: llm.responses.anthropic.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses streams Responses events"} - {id: llm.responses.anthropic.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses: cost header and spend row agree"} - {id: llm.responses.anthropic.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic function_call_output round trip over /v1/responses"} +- {id: llm.responses.anthropic.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: tool_search, streaming: nonstream, assertions: [works], source: "llm_translation/test_responses_codex_tool_search_e2e.py", rationale: "Codex client tool_search over /v1/responses finds and loads a deferred tool on Claude"} diff --git a/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py b/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py new file mode 100644 index 00000000000..79fd61e247a --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py @@ -0,0 +1,205 @@ +"""Live e2e: Codex's client tool_search over /v1/responses on Claude. + +Codex holds MCP tools back behind its client-executed `tool_search` tool when the model catalog says +`supports_search_tool`. Over the completion bridge the proxy offers that search to Claude as a plain function, +hands Claude's call back as a `tool_search_call` Codex runs itself, and loads the tools Codex returns in +`tool_search_output` so Claude can call them on the next turn +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from openai.types.responses import ( + FunctionToolParam, + Response, + ResponseFunctionToolCall, + ResponseInputParam, + ResponseToolSearchCall, + ToolParam, + ToolSearchToolParam, +) +from proxy_client import ProxyClient +from pydantic import BaseModel, ConfigDict +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +ANTHROPIC_BACKEND: Final = "anthropic/claude-haiku-5-5" +DEFERRED_TOOL_COUNT: Final = 50 +CALENDAR_TOOL_NAME: Final = "mcp__codex_apps__calendar_create_event" + +CODEX_TOOL_SEARCH: Final[ToolSearchToolParam] = { + "type": "tool_search", + "execution": "client", + "description": ( + "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next " + "model call. Some of the tools may not have been provided to you upfront, and you should use this tool " + "(`tool_search`) to search for the required tools." + ), + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query for deferred tools."}, + "limit": {"type": "number", "description": "Maximum number of tools to return. Defaults to 8."}, + }, + "required": ["query"], + "additionalProperties": False, + }, +} + +SHELL_TOOL: Final[FunctionToolParam] = { + "type": "function", + "name": "shell", + "description": "Runs a shell command.", + "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"]}, + "strict": False, +} + +CALENDAR_TOOL: Final[FunctionToolParam] = { + "type": "function", + "name": CALENDAR_TOOL_NAME, + "description": "Create a calendar event.", + "parameters": { + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, + }, + "strict": False, + "defer_loading": True, +} + +PROMPT: Final = "Create a calendar event titled Standup. Find the calendar tool with tool_search first, then call it." + + +class _CalendarArguments(BaseModel): + model_config = ConfigDict(frozen=True) + + title: str + + +def _filler_tool(index: int) -> FunctionToolParam: + return { + "type": "function", + "name": f"mcp__codex_apps__records_{index:02d}", + "description": f"List, inspect, update and archive records of family {index} in the connected workspace.", + "parameters": { + "type": "object", + "properties": { + "title": {"type": "string", "description": "Title of the record to act on"}, + "notes": {"type": "string", "description": "Free-form notes and context for the operation"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "priority": {"type": "integer", "minimum": 0, "maximum": 5}, + }, + "required": ["title"], + "additionalProperties": False, + }, + "strict": False, + "defer_loading": True, + } + + +DEFERRED_TOOLS: Final[tuple[FunctionToolParam, ...]] = ( + *(_filler_tool(index) for index in range(DEFERRED_TOOL_COUNT - 1)), + CALENDAR_TOOL, +) + + +def _eager(tool: FunctionToolParam) -> ToolParam: + return FunctionToolParam( + type="function", + name=tool["name"], + description=tool.get("description"), + parameters=tool.get("parameters"), + strict=False, + ) + + +def _input_tokens(response: Response) -> int: + assert response.usage is not None, f"response carried no usage: {response.id}" + return response.usage.input_tokens + + +def _search_call(response: Response) -> ResponseToolSearchCall: + calls: Final = tuple(item for item in response.output if isinstance(item, ResponseToolSearchCall)) + assert len(calls) == 1, f"expected one tool_search_call, got {response.output!r}" + return calls[0] + + +class TestResponsesCodexToolSearchOverAnthropic: + @pytest.mark.covers("llm.responses.anthropic.tool_search.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.ANTHROPIC,), + models=(ANTHROPIC_BACKEND,), + capabilities=(Capability.TOOL_SEARCH, Capability.FUNCTION_CALLING), + mode=Mode.NONSTREAM, + ) + ) + def test_claude_finds_a_deferred_tool_through_codex_tool_search_and_calls_it( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = f"e2e-codex-tool-search-{unique_marker()}" + model_id: Final = proxy.create_model( + model, LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY") + ) + resources.defer(lambda: proxy.delete_model(model_id)) + client: Final = sdk.openai(resources.key()) + searchable_tools: Final[list[ToolParam]] = [SHELL_TOOL, CODEX_TOOL_SEARCH] + + eager: Final = client.responses.create( + model=model, + input=PROMPT, + tools=[SHELL_TOOL, *(_eager(tool) for tool in DEFERRED_TOOLS)], + extra_body=NO_PROXY_CACHE, + ) + first: Final = client.responses.create( + model=model, input=PROMPT, tools=searchable_tools, extra_body=NO_PROXY_CACHE + ) + + search: Final = _search_call(first) + assert search.execution == "client", f"tool_search ran on the wrong side: {search!r}" + assert search.call_id, f"tool_search_call had no call_id: {search!r}" + assert not any(isinstance(item, ResponseFunctionToolCall) for item in first.output), first.output + assert _input_tokens(first) * 4 < _input_tokens(eager), ( + f"deferred first turn used {_input_tokens(first)} input tokens against {_input_tokens(eager)} eager" + ) + + follow_up: Final[ResponseInputParam] = [ + {"role": "user", "content": PROMPT}, + { + "type": "tool_search_call", + "call_id": search.call_id, + "execution": "client", + "arguments": search.arguments, + }, + { + "type": "tool_search_output", + "call_id": search.call_id, + "status": "completed", + "execution": "client", + "tools": [CALENDAR_TOOL], + }, + ] + second: Final = client.responses.create( + model=model, input=follow_up, tools=searchable_tools, extra_body=NO_PROXY_CACHE + ) + + calendar_call: Final = next( + ( + item + for item in second.output + if isinstance(item, ResponseFunctionToolCall) and item.name == CALENDAR_TOOL_NAME + ), + None, + ) + assert calendar_call is not None, f"Claude did not call the loaded tool: {second.output!r}" + assert "standup" in _CalendarArguments.model_validate_json(calendar_call.arguments).title.lower() diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3716162692d..6334fc1482e 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -780,7 +780,7 @@ class TestCodexModelSync: "availability_nux": None, "default_reasoning_level": None, } - assert "supports_search_tool" not in entries["m"] + assert entries["m"]["supports_search_tool"] is False def test_alias_gets_no_upgrade_nudge_away_from_its_own_name(self, tmp_path): listing = self._listing(self._row("gpt-5.6-terra"), self._row("api-54")) @@ -939,7 +939,7 @@ class TestCodexModelSync: assert result == ModelSyncArgs(("-c", f"model_catalog_json={json.dumps(str(catalog))}")) entry = self._entries(tmp_path)["api-terra"] assert entry["slug"] == "api-terra" - assert "supports_search_tool" not in entry + assert entry["supports_search_tool"] is False def test_stock_catalog_is_decoded_as_utf8_regardless_of_locale(self, tmp_path): description = "Modelo equilibrado para el trabajo diario, con acentos y ñ." diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index b43e43df7c9..75b3ee70aab 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -507,8 +507,11 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): upstream_model="openai/gpt-5.5", display_name="Team 5.5", service_tiers=(["ultrafast"], None), + search_tool_support=(None, None), + ), + CodexCatalogRow( + id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) ), - CodexCatalogRow(id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,)), ) @@ -530,7 +533,12 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): (row,) = codex_catalog_rows(listing, (("gpt-6", "gpt-6"),), router) assert row == CodexCatalogRow( - id="gpt-6", mode="chat", max_input_tokens=9, upstream_model="openai/gpt-6-astra", service_tiers=(["ultrafast"],) + id="gpt-6", + mode="chat", + max_input_tokens=9, + upstream_model="openai/gpt-6-astra", + service_tiers=(["ultrafast"],), + search_tool_support=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == ( @@ -611,3 +619,51 @@ def test_catalog_rows_without_a_router_carry_only_the_listing(): listing = ({"id": "plain", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) assert codex_catalog_rows(listing, (("plain", "plain"),), None) == (CodexCatalogRow(id="plain", mode="chat"),) + + +def test_codex_tool_search_is_offered_off_the_route_and_model_info_and_only_when_every_deployment_answers_it(): + claude = "anthropic/claude-sonnet-5-5" + converse = "bedrock/us.anthropic.claude-sonnet-5-5" + router = litellm.Router( + model_list=[ + {"model_name": "claude", "litellm_params": {"model": claude}}, + { + "model_name": "claude-off", + "litellm_params": {"model": claude}, + "model_info": {"supports_search_tool": False}, + }, + {"model_name": "converse", "litellm_params": {"model": converse}}, + { + "model_name": "converse-on", + "litellm_params": {"model": converse}, + "model_info": {"supports_search_tool": True}, + }, + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5"}}, + { + "model_name": "gpt-5.5-off", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"supports_search_tool": False}, + }, + {"model_name": "mixed", "litellm_params": {"model": claude}}, + {"model_name": "mixed", "litellm_params": {"model": converse}}, + ] + ) + names = ("claude", "claude-off", "converse", "converse-on", "gpt-5.5", "gpt-5.5-off", "mixed") + listing = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "x", "mode": "chat", "max_input_tokens": 1000000} + for name in names + ) + + models = json.loads(codex_model_list_body(listing, tuple((name, name) for name in names), router))["models"] + + stock_support = bundled_codex_models()["gpt-5.5"].model_dump()["supports_search_tool"] + assert {entry["slug"]: entry["supports_search_tool"] for entry in models} == { + "claude": True, + "claude-off": False, + "converse": False, + "converse-on": True, + "gpt-5.5": stock_support, + "gpt-5.5-off": False, + "mixed": False, + } + assert models[0]["context_window"] == 1000000 diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 36f2a8d5acd..5c6ee0ffe5d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -27,9 +27,7 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from .conftest import normalize # type: ignore[import-not-found] -def _stub_model_info_response( - model_id: str = "gpt-4", provider: str = "openai" -) -> dict: +def _stub_model_info_response(model_id: str = "gpt-4", provider: str = "openai") -> dict: return { "id": model_id, "object": "model", @@ -54,6 +52,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_search_tool_support = MagicMock(return_value=()) monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -70,9 +69,7 @@ def patched_models(monkeypatch): def _fake_create_model_info_response(model_id, provider="openai", **kwargs): return _stub_model_info_response(model_id=model_id, provider=provider) - monkeypatch.setattr( - proxy_utils, "create_model_info_response", _fake_create_model_info_response - ) + monkeypatch.setattr(proxy_utils, "create_model_info_response", _fake_create_model_info_response) monkeypatch.setattr(proxy_utils, "validate_model_access", lambda **kwargs: None) @@ -111,9 +108,7 @@ def test_get_models_happy_path(client, auth_as, patched_models, path): @pytest.mark.parametrize("path", ["/v1/models", "/models"]) -def test_get_models_anthropic_format_when_header_present( - client, auth_as, patched_models, path -): +def test_get_models_anthropic_format_when_header_present(client, auth_as, patched_models, path): """Pins: ``GET /v1/models`` returns the Anthropic-native models shape when the caller sends an ``anthropic-version`` header (Claude Code gateway discovery), while the default OpenAI shape is unchanged without it.""" @@ -133,9 +128,7 @@ def test_get_models_anthropic_format_when_header_present( @pytest.mark.parametrize("path", ["/v1/models", "/models"]) -def test_anthropic_format_exposes_token_limits( - client, auth_as, patched_models, monkeypatch, path -): +def test_anthropic_format_exposes_token_limits(client, auth_as, patched_models, monkeypatch, path): """Claude Code sizes requests off the listing, so the Anthropic-native entries carry the same token limits the OpenAI listing resolves, with the output budget named max_tokens as the Messages API names it.""" @@ -150,9 +143,7 @@ def test_anthropic_format_exposes_token_limits( "max_output_tokens": 64000, } - monkeypatch.setattr( - proxy_utils, "create_model_info_response", _create_model_info_response - ) + monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) with auth_as(): response = client.get(path, headers={"anthropic-version": "2023-06-01"}) @@ -227,9 +218,7 @@ def test_anthropic_format_uses_configured_display_name(client, auth_as, patched_ @pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) -def test_anthropic_display_name_resolved_via_internal_team_key( - client, auth_as, patched_models, monkeypatch, params -): +def test_anthropic_display_name_resolved_via_internal_team_key(client, auth_as, patched_models, monkeypatch, params): """For a team-scoped row the configured display name must be looked up by the internal routing key while the entry itself is keyed by the public name, so the clean name lands on the id the client actually sees.""" @@ -263,14 +252,10 @@ def test_anthropic_display_name_resolved_via_internal_team_key( "get_available_models_for_user", _fake_get_available_models_for_user, ) - monkeypatch.setattr( - model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] - ) + monkeypatch.setattr(model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]) with auth_as(): - response = client.get( - "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} - ) + response = client.get("/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}) assert response.status_code == 200 (entry,) = response.json()["data"] @@ -311,9 +296,7 @@ def test_get_model_by_id_not_found(client, auth_as, patched_models, path): @pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) -def test_anthropic_format_returns_public_team_model_name( - client, auth_as, patched_models, monkeypatch, params -): +def test_anthropic_format_returns_public_team_model_name(client, auth_as, patched_models, monkeypatch, params): """Regression: the Anthropic-native listing must go through the same team name translation as the OpenAI listing, so a caller never sees the internal ``model_name_{team_id}_{uuid}`` routing key.""" @@ -344,14 +327,10 @@ def test_anthropic_format_returns_public_team_model_name( "get_available_models_for_user", _fake_get_available_models_for_user, ) - monkeypatch.setattr( - model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] - ) + monkeypatch.setattr(model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]) with auth_as(): - response = client.get( - "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} - ) + response = client.get("/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}) assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] @@ -382,7 +361,9 @@ def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( patched_models.model_group_alias = {} patched_models.has_model_id.return_value = False - patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: ( + frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + ) monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) with auth_as(): @@ -402,7 +383,9 @@ def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( def test_anthropic_format_keeps_served_ids_for_other_anthropic_clients(client, auth_as, patched_models, path): """An Anthropic SDK asking for the vendor shape gets the served ids: the view is Claude Code's alone.""" with auth_as(): - response = client.get(path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"}) + response = client.get( + path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"} + ) assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4", "claude-sonnet"] diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 36df738c0bb..ef198b487b9 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -17,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.litellm_completion_transformation.session_handler import StoredTurns from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -55,6 +56,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): request_kwargs: Mapping[str, object] | None = None, previous_response_client_ids: tuple[str, ...] = (), input_replaces_history: bool = False, + stored_turns: StoredTurns | None = None, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -118,8 +120,10 @@ async def test_aresponses_forwards_timeout_to_acompletion(): class _FakeSpendLogsDB: def __init__(self, spend_logs: Sequence[Mapping[str, object]]) -> None: self._spend_logs = spend_logs + self.queries = 0 async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + self.queries += 1 return self._spend_logs @@ -1010,9 +1014,7 @@ class _RecordingCustomLLM(CustomLLM): def completion(self, *args: object, **kwargs: object) -> ModelResponse: self.sent_messages.extend(cast(list[object], kwargs["messages"])) - return litellm.completion( - model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog." - ) + return litellm.completion(model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog.") def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text( @@ -1198,3 +1200,372 @@ def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_th ], } ] + + +_CODEX_TOOL_SEARCH: Final = { + "type": "tool_search", + "execution": "client", + "description": "Search the deferred app tools.", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}, "limit": {"type": "integer"}}, + "required": ["query"], + "additionalProperties": False, + }, +} +_SHELL_TOOL: Final = { + "type": "function", + "name": "shell", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}, +} +_CALENDAR_PARAMETERS: Final = { + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, +} +_CODEX_SEARCH_CALL: Final = { + "type": "tool_search_call", + "call_id": "search-1", + "execution": "client", + "arguments": {"query": "calendar create", "limit": 1}, +} +_CODEX_SEARCH_OUTPUT: Final = { + "type": "tool_search_output", + "call_id": "search-1", + "status": "completed", + "execution": "client", + "tools": [ + { + "type": "function", + "name": "mcp__codex_apps__calendar_create_event", + "description": "Create a calendar event.", + "defer_loading": True, + "parameters": _CALENDAR_PARAMETERS, + } + ], +} + + +def _deferred_namespace(name: str, *members: str) -> Mapping[str, object]: + return { + "type": "namespace", + "name": name, + "description": f"{name} tools", + "tools": [ + {"type": "function", "name": member, "defer_loading": True, "parameters": _CALENDAR_PARAMETERS} + for member in members + ], + } + + +def _claude_tool_use(tool_id: str, name: str, tool_input: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_tool_use", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "tool_use", "id": tool_id, "name": name, "input": tool_input}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +def _sent(route: respx.Route) -> Mapping[str, object]: + return json.loads(route.calls.last.request.content) + + +def test_responses_over_anthropic_offer_codex_tool_search_as_a_function_and_hold_back_deferred_tools( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + + litellm.responses( + model=_MODEL, + input="Create a calendar event", + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL, _deferred_namespace("mcp__vercel__", "list_projects")], + api_key="mock_api_key", + ) + + sent_tools: Final = _sent(route)["tools"] + assert [tool["name"] for tool in sent_tools] == ["tool_search", "shell"] + assert sent_tools[0]["description"] == _CODEX_TOOL_SEARCH["description"] + assert sent_tools[0]["input_schema"] == _CODEX_TOOL_SEARCH["parameters"] + assert "defer_loading" not in json.dumps(sent_tools) + + +@pytest.mark.parametrize("tool_input", [{"query": "calendar create", "limit": 1}, {"query": "calendar create"}]) +def test_claude_calling_tool_search_comes_back_as_a_client_tool_search_call( + respx_mock: respx.MockRouter, tool_input: Mapping[str, object] +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_search", "tool_search", tool_input) + ) + + response: Final = litellm.responses( + model=_MODEL, input="Create a calendar event", tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], api_key="mock_api_key" + ) + + (item,) = response.output + assert (item.type, item.execution, item.call_id, item.arguments) == ( + "tool_search_call", + "client", + "toolu_search", + tool_input, + ) + + +def test_claude_calling_a_function_named_tool_search_stays_a_function_call_without_codex_tool_search( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_search", "tool_search", {"query": "calendar"}) + ) + own_tool_search: Final = {**_SHELL_TOOL, "name": "tool_search"} + + response: Final = litellm.responses(model=_MODEL, input="Search", tools=[own_tool_search], api_key="mock_api_key") + + (item,) = response.output + assert (item.type, item.name, json.loads(item.arguments)) == ("function_call", "tool_search", {"query": "calendar"}) + + +def test_streamed_claude_tool_search_is_one_client_tool_search_call_item(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_search_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_search", "name": "tool_search", "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "calendar create",'}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": ' "limit": 1}'}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input="Create a calendar event", + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + stream=True, + api_key="mock_api_key", + ) + ) + + done: Final = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] + assert [(item.type, item.execution, item.call_id, item.arguments) for item in done] == [ + ("tool_search_call", "client", "toolu_search", {"query": "calendar create", "limit": 1}) + ] + assert not any(event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA for event in events) + assert [item.type for item in events[-1].response.output] == ["tool_search_call"] + + +def test_codex_tool_search_output_loads_the_found_tool_and_answers_the_search_on_anthropic( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_event", "mcp__codex_apps__calendar_create_event", {"title": "Standup"}) + ) + + response: Final = litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "Create a calendar event called Standup"}, + _CODEX_SEARCH_CALL, + _CODEX_SEARCH_OUTPUT, + ], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + api_key="mock_api_key", + ) + + sent: Final = _sent(route) + assert [(tool["name"], tool["input_schema"]) for tool in sent["tools"][1:]] == [ + ("shell", _SHELL_TOOL["parameters"]), + ("mcp__codex_apps__calendar_create_event", _CALENDAR_PARAMETERS), + ] + assert "defer_loading" not in json.dumps(sent["tools"]) + (assistant, answer) = sent["messages"][1:] + assert [block for block in assistant["content"] if block["type"] == "tool_use"] == [ + {"type": "tool_use", "id": "search-1", "name": "tool_search", "input": {"query": "calendar create", "limit": 1}} + ] + (result,) = answer["content"] + assert result["type"] == "tool_result" and result["tool_use_id"] == "search-1" + assert "mcp__codex_apps__calendar_create_event" in json.dumps(result["content"]) + (item,) = response.output + assert (item.type, item.name, json.loads(item.arguments)) == ( + "function_call", + "mcp__codex_apps__calendar_create_event", + {"title": "Standup"}, + ) + + +def test_a_namespace_codex_loads_through_tool_search_is_callable_and_its_calls_keep_the_namespace( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_projects", "mcp__vercel____list_projects", {"title": "x"}) + ) + loaded: Final = _deferred_namespace("mcp__vercel__", "list_projects", "get_project") + + response: Final = litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "List my Vercel projects"}, + {**_CODEX_SEARCH_CALL, "arguments": {"query": "vercel projects"}}, + {**_CODEX_SEARCH_OUTPUT, "tools": [loaded]}, + ], + tools=[_CODEX_TOOL_SEARCH, loaded], + api_key="mock_api_key", + ) + + assert [tool["name"] for tool in _sent(route)["tools"]] == [ + "tool_search", + "mcp__vercel____list_projects", + "mcp__vercel____get_project", + ] + (item,) = response.output + assert (item.type, item.namespace, item.name) == ("function_call", "mcp__vercel__", "list_projects") + + +def test_codex_tool_search_output_loading_the_same_tool_twice_sends_it_once(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + + litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "Create two events"}, + _CODEX_SEARCH_CALL, + _CODEX_SEARCH_OUTPUT, + {**_CODEX_SEARCH_CALL, "call_id": "search-2"}, + {**_CODEX_SEARCH_OUTPUT, "call_id": "search-2"}, + ], + tools=[_CODEX_TOOL_SEARCH], + api_key="mock_api_key", + ) + + assert [tool["name"] for tool in _sent(route)["tools"]] == ["tool_search", "mcp__codex_apps__calendar_create_event"] + + +def _stored_tool_call_turn( + request_id: str, request_input: object, tool_call_id: str, name: str, arguments: str +) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "codex-session", + "proxy_server_request": {"model": _MODEL, "input": request_input, "tools": [_CODEX_TOOL_SEARCH, _SHELL_TOOL]}, + "response": { + "id": request_id, + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": tool_call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + ], + }, + } + ], + }, + } + + +_SEARCH_TURN: Final = _stored_tool_call_turn( + "chatcmpl-search-turn", + "Create a calendar event titled Standup.", + "search-1", + "tool_search", + '{"query": "calendar"}', +) +_SEARCH_TURN_STORED_AS_RESPONSE: Final = { + **_SEARCH_TURN, + "response": {"id": "chatcmpl-search-turn", "object": "response", "output": [_CODEX_SEARCH_CALL]}, +} +_CALENDAR_TURN: Final = _stored_tool_call_turn( + "chatcmpl-calendar-turn", + [_CODEX_SEARCH_OUTPUT], + "toolu_standup", + "mcp__codex_apps__calendar_create_event", + '{"title": "Standup"}', +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search_turn", [_SEARCH_TURN, _SEARCH_TURN_STORED_AS_RESPONSE], ids=["chat_completion", "responses_output"] +) +async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search_loaded( + search_turn: Mapping[str, object], +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + prisma_client: Final = _FakePrismaClient([search_turn, _CALENDAR_TURN]) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma_client): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-calendar-turn", + input=[ + {"type": "function_call_output", "call_id": "toolu_standup", "output": "Created Standup."}, + {"role": "user", "content": "Now create another event titled Retro."}, + ], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + sent: Final = json.loads(anthropic.body) + assert [tool["name"] for tool in sent["tools"]] == [ + "tool_search", + "shell", + "mcp__codex_apps__calendar_create_event", + ] + assert [ + [(block["type"], block.get("id") or block.get("tool_use_id")) for block in message["content"]] + for message in sent["messages"][1:] + ] == [ + [("tool_use", "search-1")], + [("tool_result", "search-1")], + [("tool_use", "toolu_standup")], + [("tool_result", "toolu_standup"), ("text", None)], + ] + assert prisma_client.db.queries == 1 diff --git a/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py b/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py new file mode 100644 index 00000000000..dc975b3d112 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py @@ -0,0 +1,68 @@ +from typing import Final + +from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam + +_CODEX_TOOL_SEARCH: Final = {"type": "tool_search", "execution": "client"} +_SHELL: Final = {"type": "function", "name": "shell", "parameters": {"type": "object", "properties": {}}} +_SEARCH_CALL: Final = { + "type": "tool_search_call", + "call_id": "search-1", + "execution": "client", + "arguments": {"query": "shell"}, +} + + +def _search_output(*tools: dict[str, object]) -> dict[str, object]: + return { + "type": "tool_search_output", + "call_id": "search-1", + "status": "completed", + "execution": "client", + "tools": list(tools), + } + + +def test_requests_without_client_tool_search_pass_through_untouched() -> None: + request_input: Final[ResponseInputParam] = [ + {"type": "message", "role": "user", "content": "hi"}, + {"type": "tool_search_call", "call_id": "srv-1", "execution": "server", "arguments": {"query": "x"}}, + ] + tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = [_SHELL, {"type": "tool_search", "execution": "server"}] + + bridged: Final = bridge_client_tool_search(request_input, tools) + + assert bridged.bridged is False + assert bridged.input is request_input + assert bridged.tools == tuple(tools) + + +def test_a_search_that_found_nothing_answers_the_call_and_loads_no_tools() -> None: + bridged: Final = bridge_client_tool_search([_SEARCH_CALL, _search_output()], [_CODEX_TOOL_SEARCH, _SHELL]) + + assert bridged.tools == (_CODEX_TOOL_SEARCH, _SHELL) + assert bridged.input[1] == { + "type": "function_call_output", + "call_id": "search-1", + "output": "No tools matched the search.", + } + + +def test_a_loaded_tool_already_offered_eagerly_is_not_sent_twice() -> None: + loaded_shell: Final = {**_SHELL, "description": "loaded copy", "defer_loading": True} + + bridged: Final = bridge_client_tool_search( + [_SEARCH_CALL, _search_output(loaded_shell)], [_CODEX_TOOL_SEARCH, _SHELL] + ) + + assert bridged.tools == (_CODEX_TOOL_SEARCH, _SHELL) + + +def test_deferred_request_tools_are_held_back_only_when_the_client_searches() -> None: + deferred: Final = {**_SHELL, "name": "deploy", "defer_loading": True} + + with_search: Final = bridge_client_tool_search("hi", [_CODEX_TOOL_SEARCH, _SHELL, deferred]) + without_search: Final = bridge_client_tool_search("hi", [_SHELL, deferred]) + + assert with_search.tools == (_CODEX_TOOL_SEARCH, _SHELL) + assert without_search.tools == (_SHELL, deferred) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 475628569a0..4250c78c669 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2,6 +2,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from copy import deepcopy +from types import MappingProxyType from typing import Final, Literal, cast import pytest @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, ) from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks +from litellm.responses.litellm_completion_transformation.tool_search import tool_search_routing from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, ChatCompletionSession, @@ -1288,12 +1290,7 @@ class TestFunctionCallTransformation: def test_parallel_tool_calls_dropped_when_no_chat_tools_remain(self) -> None: transform: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request - codex_tool_search: Final = { - "type": "tool_search", - "execution": "client", - "description": "Searches over deferred tool metadata with BM25.", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, - } + hosted_tool_search: Final = {"type": "tool_search", "execution": "server"} function_tool: Final = { "type": "function", "name": "get_goal", @@ -1311,7 +1308,7 @@ class TestFunctionCallTransformation: hosted_only_result: Final = transform( model="azure/gpt-5.4-mini", input="Reply with just the word pong.", - responses_api_request={"tools": [codex_tool_search], "parallel_tool_calls": True}, + responses_api_request={"tools": [hosted_tool_search], "parallel_tool_calls": True}, custom_llm_provider="azure", ) function_tools_result: Final = transform( @@ -1770,7 +1767,7 @@ class TestToolTransformation: assert len(result_tools) == 0 assert web_search_options is None - def test_transform_codex_tools_drops_hosted_tool_search(self) -> None: + def test_transform_codex_tools_offers_client_tool_search_as_a_function(self) -> None: codex_tools: Final = [ { "type": "function", @@ -1823,16 +1820,22 @@ class TestToolTransformation: }, {"type": "web_search", "external_web_access": False, "search_content_types": ["text", "image"]}, ] - function_and_custom_count: Final = sum(1 for tool in codex_tools if tool["type"] in ("function", "custom")) + codex_search: Final = codex_tools[3] ( result_tools, web_search_options, ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=codex_tools) - assert not any(tool.get("type") == "tool_search" for tool in result_tools) assert all(tool.get("type") == "function" for tool in result_tools) - assert len(result_tools) == function_and_custom_count + assert [tool["function"]["name"] for tool in result_tools] == [ + "exec_command", + "write_stdin", + "apply_patch", + "tool_search", + ] + assert result_tools[3]["function"]["description"] == codex_search["description"] + assert result_tools[3]["function"]["parameters"] == codex_search["parameters"] assert web_search_options is not None @pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) @@ -1877,9 +1880,6 @@ class TestToolTransformation: assert [tool["type"] for tool in result["tools"]] == ["function", "function"] assert all("defer_loading" not in tool for tool in result["tools"]) - @pytest.mark.parametrize( - "search_tools", [[], [{"type": "tool_search", "execution": "client"}]], ids=["none", "client"] - ) @pytest.mark.parametrize( ("model", "provider"), [ @@ -1890,7 +1890,7 @@ class TestToolTransformation: ], ) def test_claude_routes_keep_the_callers_deferral_when_no_server_search_tool_is_sent( - self, model: str, provider: str, search_tools: list[dict[str, object]] + self, model: str, provider: str ) -> None: function_tools: Final = [ tool for tool in _deferred_tools_with_server_tool_search() if tool["type"] != "tool_search" @@ -1899,7 +1899,7 @@ class TestToolTransformation: result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, input="What is the weather in Paris right now?", - responses_api_request={"tools": [*search_tools, *function_tools]}, + responses_api_request={"tools": function_tools}, custom_llm_provider=provider, ) @@ -4361,8 +4361,8 @@ class TestEnsureOutputItemContentPartAdded: iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(None) iterator._web_search_calls = {} iterator._queued_web_search_call_ids = set() - iterator._tool_search_requested = False - iterator._tool_search_call_ids = set() + iterator._tool_search = tool_search_routing(None) + iterator._tool_search_executions = MappingProxyType({}) return iterator def _make_text_chunk(self): @@ -6754,7 +6754,6 @@ def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None ] - @pytest.mark.parametrize( ("call", "output"), [ @@ -6905,7 +6904,14 @@ def test_compaction_replay_invents_no_call_for_an_output_whose_call_was_never_st [ pytest.param( _COMPACTION_REQUESTED, - [("user", ()), ("assistant", ("call_0",)), ("tool", "Sunny."), ("assistant", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + [ + ("user", ()), + ("assistant", ("call_0",)), + ("tool", "Sunny."), + ("assistant", ()), + ("assistant", ("call_0",)), + ("tool", "Rain."), + ], id="native_compaction_block", ), pytest.param( @@ -7047,7 +7053,9 @@ async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chun assert incomplete_event.type == "response.incomplete" assert incomplete_event.response.status == "incomplete" assert incomplete_event.response.usage is not None - assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) + assert incomplete_event.response.usage.output_tokens == litellm.token_counter( + model="gpt-4o-mini", text=relayed_text + ) assert incomplete_event.response.usage.input_tokens > 0 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d59fb69dcaf..b55b3c5efc2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -59,6 +59,7 @@ from litellm.router import ( _is_retriable_anthropic_status, _responses_stream_holds_event, _without_line_breaks, + _search_tool_support, Span, ) from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request @@ -13517,6 +13518,80 @@ def test_get_configured_service_tiers_reads_only_the_deployments_a_request_can_r assert router.get_configured_service_tiers("gpt-6-astra") == (["ultrafast"],) +@pytest.mark.parametrize( + "litellm_params, model_info, expected", + [ + ({"model": "anthropic/claude-fable-5-1"}, None, True), + ({"model": "bedrock/invoke/us.anthropic.claude-fable-5-1"}, None, True), + ({"model": "openai/gpt-6-astra"}, None, None), + ({"model": "gemini/gemini-3.8-flash"}, None, None), + ({"model": "sonar-pro", "api_base": "https://api.perplexity.ai"}, None, None), + ({"model": "gemini/gemini-3.8-flash"}, {"supports_search_tool": True}, True), + ({"model": "anthropic/claude-fable-5-1"}, {"supports_search_tool": False}, False), + ], + ids=[ + "anthropic-defaults-on", + "bedrock-invoke-anthropic-defaults-on", + "openai-passthrough-unset", + "other-bridged-provider-unset", + "provider-inferred-from-api-base-unset", + "model-info-opt-in-wins", + "model-info-opt-out-wins", + ], +) +def test_get_search_tool_support_defaults_from_the_provider_and_lets_model_info_override( + litellm_params: dict[str, str], model_info: dict[str, bool] | None, expected: bool | None +): + deployment: Final = {"model_name": "codex-model", "litellm_params": litellm_params} + router: Final = litellm.Router(model_list=[{**deployment, "model_info": model_info} if model_info else deployment]) + + assert router.get_search_tool_support("codex-model") == (expected,) + + +def test_get_search_tool_support_reports_each_routable_deployment_of_the_group(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "mixed", "litellm_params": {"model": "anthropic/claude-fable-5-1"}}, + {"model_name": "mixed", "litellm_params": {"model": "openai/gpt-6-astra"}}, + { + "model_name": "mixed", + "litellm_params": {"model": "anthropic/claude-fable-5-1", "api_base": "https://paused.example"}, + "model_info": {"blocked": True, "supports_search_tool": False}, + }, + ] + ) + + assert router.get_search_tool_support("mixed") == (True, None) + assert router.get_search_tool_support("not-a-real-model") == () + + +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("provider", ["chatgpt", "github_copilot"]) +@pytest.mark.parametrize("explicit_provider", [False, True]) +async def test_search_tool_catalog_lookup_needs_no_credentials( + provider: str, explicit_provider: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "chatgpt")) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) + (tmp_path / "chatgpt").mkdir() + (tmp_path / "chatgpt" / "auth.json").write_text( + json.dumps({"access_token": "expired-access", "refresh_token": "test-refresh", "expires_at": 1}) + ) + (tmp_path / "copilot").mkdir() + (tmp_path / "copilot" / "api-key.json").write_text(json.dumps({"token": "expired-key", "expires_at": 1})) + (tmp_path / "copilot" / "access-token").write_text("test-access") + respx_mock.route().respond(401, json={"error": "invalid_grant"}) + params: Final = ( + {"model": "catalog-model", "custom_llm_provider": provider} + if explicit_provider + else {"model": f"{provider}/catalog-model"} + ) + + assert _search_tool_support({"model_name": "catalog-model", "litellm_params": params}) is None + assert len(respx_mock.calls) == 0 + + @pytest.mark.parametrize( "alias_value, expected_group", [ From 3c1eb73a0623d4191acace5d69f0acb21d579cc1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:31:04 -0700 Subject: [PATCH 06/58] fix(proxy)!: list every priced Vertex family under a vertex_ai/* grant (#41959) * fix(proxy): list every priced Vertex family under a vertex_ai/* grant A vertex_ai/* key listed none of the 68 cost map rows priced under the Mistral, Llama, Qwen, AI21, OpenAI, image, video and embedding labels because their sets were never part of the vertex_ai union, and Qwen rows had no ingestion branch at all. Wildcard expansion also dropped a meta/ or openai/ publisher segment as if it were a provider prefix, and the wildcard health check would have picked $0-per-token rows first once the union grew, so it now sorts paid rows ahead of unpriced ones. The VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST allowlist is gone so the cost map guard enforces every label from now on. * fix(health): route every wildcard health-check pick through the listed provider * fix(health-check): keep per-character-billed rows ahead of rows priced $0 everywhere in the wildcard pick * fix(health-check): break $0-per-token ties toward per-character-billed rows in the wildcard pick * fix(health): only health-check ids the router sends to the provider The wildcard health check picker now skips ids that get_llm_provider routes to a different provider or to none, so the longer vertex_ai union cannot pull partner rows into the check. Drops the forced custom_llm_provider and the per-character tie-break from earlier revisions, and pins the replicate and wandb wildcard listing to each row's full owner id. * fix(health): match wildcard health picks against the provider's own route The previous check kept a bare id only when get_llm_provider routed it to the exact provider name, which dropped every ai21 row because bare ai21 ids route to the ai21_chat alias, so ai21/* health checks sent no request. Compare the bare id's route with the route of the provider-prefixed id instead, so aliases match. Resolve routes with inferred_provider, which returns the declared provider for github_copilot and chatgpt without starting their device-code login. * fix(health): keep wildcard health picks that only route through a saved custom_llm_provider Rank picks routed to the checked provider first and ids routed to no provider second, each by cost, and skip only ids routed to a different provider, so a wildcard saved with custom_llm_provider still sends its health request * refactor(vertex_ai): mark the new Vertex registry sets Final * chore: re-run CI against current main --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 17 ++- .../litellm_core_utils/llm_request_utils.py | 30 +++-- litellm/proxy/auth/model_checks.py | 34 +++--- .../test_health_check_helpers.py | 101 ++++++++++++++++ .../test_llm_request_utils.py | 32 +++++ tests/unit/proxy/auth/test_model_checks.py | 109 ++++++++++++++++++ .../test_model_prices_and_context_window.py | 13 --- 7 files changed, 295 insertions(+), 41 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index b7d46c197f2..de8371e7a2e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -619,6 +619,8 @@ vertex_openai_models: Set = set() vertex_minimax_models: Set = set() vertex_moonshot_models: Set = set() vertex_zai_models: Set = set() +vertex_qwen_models: Final[Set[str]] = set() +vertex_ai_models: Final[Set[str]] = set() ai21_models: Set = set() ai21_chat_models: Set = set() nlp_cloud_models: Set = set() @@ -833,6 +835,10 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje elif value.get("litellm_provider") == "vertex_ai-zai_models": key = key.replace("vertex_ai/", "") vertex_zai_models.add(key) + elif value.get("litellm_provider") == "vertex_ai-qwen_models": + vertex_qwen_models.add(key.removeprefix("vertex_ai/")) + elif value.get("litellm_provider") == "vertex_ai": + vertex_ai_models.add(key.removeprefix("vertex_ai/")) elif value.get("litellm_provider") == "ai21": if value.get("mode") == "chat": ai21_chat_models.add(key) @@ -1204,7 +1210,16 @@ def _build_models_by_provider() -> dict: | vertex_deepseek_models | vertex_minimax_models | vertex_moonshot_models - | vertex_zai_models, + | vertex_zai_models + | vertex_qwen_models + | vertex_mistral_models + | vertex_llama3_models + | vertex_ai_ai21_models + | vertex_openai_models + | vertex_ai_image_models + | vertex_ai_video_models + | vertex_embedding_models + | vertex_ai_models, "ai21": ai21_models | ai21_chat_models, "bedrock": bedrock_models | bedrock_converse_models, "petals": petals_models, diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 0bf59a6d1f2..db881aa16c6 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,8 +1,9 @@ -from collections.abc import Mapping +from collections.abc import Iterable, Iterator, Mapping from typing import Final, cast import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider def _form_field_value(value: object) -> str: @@ -153,24 +154,33 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): if custom_llm_provider not in litellm.models_by_provider: return [] - known_models: Final = litellm.models_by_provider.get(custom_llm_provider, []) - model_costs: Final = [] + known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ()) + ranked_rows: Final = ( + (route_rank, cost, model) + for model, cost in _chat_row_costs(known_models, custom_llm_provider) + if (route_rank := _route_rank(model, custom_llm_provider)) is not None + ) + return [model for _, _, model in sorted(ranked_rows, key=lambda row: row[:2])[:n]] - for model in known_models: + +def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float]]: + for model in models: try: model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: continue if model_info.get("mode") != "chat": continue - _cost = (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0) - model_costs.append((model, _cost)) + yield model, (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0) - # Sort by cost (ascending) - model_costs.sort(key=lambda x: x[1]) - # Return the top n cheapest models - return [model for model, _ in model_costs[:n]] +def _route_rank(model: str, custom_llm_provider: str) -> int | None: + routed_provider: Final = inferred_provider(model) + if routed_provider is None: + return 1 + if routed_provider == inferred_provider(f"{custom_llm_provider}/{model}"): + return 0 + return None def get_proxy_server_request_headers(litellm_params: dict | None) -> dict: diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index c4c538083ee..48884d0fedc 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -20,6 +20,8 @@ from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager, get_valid_models +LLM_PROVIDER_PREFIXES: Final = frozenset(llm_provider.value for llm_provider in LlmProviders) + _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) @@ -458,23 +460,21 @@ def get_known_models_from_wildcard( # add model prefix to wildcard models wildcard_models = [f"{model_prefix}{model.removeprefix(provider_id_prefix)}" for model in wildcard_models] - known_providers: Final = {provider.value for provider in LlmProviders} - suffix_appended_wildcard_models: Final = [] - for model in wildcard_models: - if not model.startswith(f"{wildcard_provider_prefix}/"): - # `get_provider_models` returns provider-prefixed ids (e.g. "ollama/gemma3:1b"). - # When the wildcard uses a custom prefix (e.g. "ollama_server1/*" to distinguish - # multiple instances), replace that existing provider prefix instead of stacking - # both, which would otherwise yield an uncallable "ollama_server1/ollama/gemma3:1b". - # Only strip the leading segment when it is a known provider, so ids whose first - # segment is an org rather than a provider (e.g. "meta-llama/Llama-3-8B") keep it. - leading, sep, model_suffix = model.partition("/") - if sep and leading in known_providers: - model = f"{wildcard_provider_prefix}/{model_suffix}" - else: - model = f"{wildcard_provider_prefix}/{model}" - suffix_appended_wildcard_models.append(model) - return list(dict.fromkeys(suffix_appended_wildcard_models)) + return list( + dict.fromkeys(_under_wildcard_prefix(model, wildcard_provider_prefix, provider) for model in wildcard_models) + ) + + +def _under_wildcard_prefix(model: str, wildcard_provider_prefix: str, listed_provider: str) -> str: + if model.startswith(f"{wildcard_provider_prefix}/"): + return model + leading, sep, model_suffix = model.partition("/") + leading_is_listed_providers_own_prefix: Final = ( + sep != "" + and leading in LLM_PROVIDER_PREFIXES + and (leading == listed_provider or listed_provider.startswith(f"{leading}_")) + ) + return f"{wildcard_provider_prefix}/{model_suffix if leading_is_listed_providers_own_prefix else model}" def expand_wildcard_deployments_for_model_info( diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 88e2bc4b1fd..44b36a355c6 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -9,6 +9,7 @@ from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest +import responses import respx import litellm @@ -1544,3 +1545,103 @@ async def test_ahealth_check_probes_extraction_models_through_the_extraction_api sent: Final = json.loads(upstream.calls[0].request.content) assert sent["state"] == "health check" assert sent["fields"]["reachable"]["type"] == "text" + + +VERTEX_AUTHORIZED_USER: Final = json.dumps( + { + "type": "authorized_user", + "client_id": "health-client", + "client_secret": "health-secret", + "refresh_token": "health-refresh", + } +) +CHAT_COMPLETION_BODY: Final = MappingProxyType( + { + "id": "chatcmpl-health", + "object": "chat.completion", + "created": 0, + "model": "health", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } +) +GEMINI_GENERATE_CONTENT_BODY: Final = MappingProxyType( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } +) + + +async def test_vertex_wildcard_health_check_only_calls_vertex_hosts( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-not-vertex") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + vertex_route: Final = respx_mock.route(host__regex=r".*aiplatform\.googleapis\.com$").respond( + json=dict(GEMINI_GENERATE_CONTENT_BODY) + ) + other_hosts: Final = respx_mock.route().respond(json=dict(CHAT_COMPLETION_BODY)) + + with responses.RequestsMock() as google_token_endpoint: + google_token_endpoint.post( + "https://oauth2.googleapis.com/token", + json={"access_token": "health-access-token", "expires_in": 3600, "token_type": "Bearer"}, + ) + result: Final = await ahealth_check( + model_params={ + "model": "vertex_ai/*", + "vertex_project": "health-project", + "vertex_location": "us-central1", + "vertex_credentials": VERTEX_AUTHORIZED_USER, + }, + mode="chat", + ) + + assert "error" not in result, result + assert vertex_route.called + assert not other_hosts.called + + +async def test_ai21_wildcard_health_check_calls_ai21( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + ai21_route: Final = respx_mock.route(host="api.ai21.com").respond(json=dict(CHAT_COMPLETION_BODY)) + other_hosts: Final = respx_mock.route().respond(json=dict(CHAT_COMPLETION_BODY)) + + result: Final = await ahealth_check(model_params={"model": "ai21/*", "api_key": "health-ai21-key"}, mode="chat") + + assert "error" not in result, result + assert ai21_route.called + assert not other_hosts.called + + +async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_a_routing_prefix( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem( + litellm.model_cost, + "fake-together-chat", + {"input_cost_per_token": 1e-9, "output_cost_per_token": 1e-9, "mode": "chat", "litellm_provider": "together_ai"}, + ) + monkeypatch.setitem(litellm.models_by_provider, "together_ai", ("fake-together-chat",)) + litellm.in_memory_llm_clients_cache.flush_cache() + together_route: Final = respx_mock.route(host="api.together.ai").respond(json=dict(CHAT_COMPLETION_BODY)) + other_hosts: Final = respx_mock.route().respond(json=dict(CHAT_COMPLETION_BODY)) + + result: Final = await ahealth_check( + model_params={"model": "together_ai/*", "custom_llm_provider": "together_ai", "api_key": "health-together-key"}, + mode="chat", + ) + + assert "error" not in result, result + assert json.loads(together_route.calls.last.request.content)["model"] == "fake-together-chat" + assert not other_hosts.called + diff --git a/tests/unit/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py index 350419b49a5..0a2bd5cec1b 100644 --- a/tests/unit/litellm_core_utils/test_llm_request_utils.py +++ b/tests/unit/litellm_core_utils/test_llm_request_utils.py @@ -1,9 +1,13 @@ +from typing import Final + import httpx import pytest +import litellm from litellm.litellm_core_utils.llm_request_utils import ( ensure_extra_body_is_safe, flatten_form_field_values, + pick_cheapest_chat_models_from_llm_provider, serialize_multipart_form_fields, ) @@ -133,3 +137,31 @@ def test_ensure_extra_body_is_safe_returns_non_dict_unchanged(): result = ensure_extra_body_is_safe(cast(dict[str, object], extra_body)) assert result is extra_body + + +PICKER_PROVIDER: Final = "deepseek" + + +@pytest.fixture +def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyPatch) -> None: + rows: Final = { + "deepseek/fake-pricey-chat": {"input_cost_per_token": 3e-12, "output_cost_per_token": 3e-12, "mode": "chat"}, + "deepseek/fake-cheapest-embedding": {"input_cost_per_token": 1e-13, "mode": "embedding"}, + "openai/fake-free-chat": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, + "deepseek/fake-cheap-chat": {"input_cost_per_token": 1e-12, "output_cost_per_token": 1e-12, "mode": "chat"}, + "fake-free-chat-without-provider": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, + } + for name, row in rows.items(): + monkeypatch.setitem(litellm.model_cost, name, {**row, "litellm_provider": PICKER_PROVIDER}) + monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(rows)) + + +def test_pick_cheapest_chat_models_prefers_provider_routed_ids_and_skips_ids_routed_elsewhere( + deepseek_rows_whose_cheapest_ids_route_elsewhere: None, +) -> None: + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=10) == [ + "deepseek/fake-cheap-chat", + "deepseek/fake-pricey-chat", + "fake-free-chat-without-provider", + ] + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-cheap-chat"] diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index 83761620cc9..3f0f97e63b0 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -954,6 +954,115 @@ def test_wildcard_lists_a_model_once_when_registered_both_bare_and_provider_pref litellm.add_known_models(model_cost_map={}) +def _is_vertex_ai_label(label: object) -> bool: + return label == "vertex_ai" or (isinstance(label, str) and label.startswith("vertex_ai-")) + + +def test_vertex_wildcard_lists_every_vertex_cost_map_row_under_its_callable_id(): + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + vertex_rows: Final = { + f"vertex_ai/{name.removeprefix('vertex_ai/')}" + for name, info in litellm.model_cost.items() + if _is_vertex_ai_label(info.get("litellm_provider")) + } + assert vertex_rows + + listed: Final = set(get_known_models_from_wildcard("vertex_ai/*")) + + assert sorted(vertex_rows - listed) == [] + + +VERTEX_COST_MAP_LABELS: Final = ( + "vertex_ai", + "vertex_ai-language-models", + "vertex_ai-vision-models", + "vertex_ai-chat-models", + "vertex_ai-code-chat-models", + "vertex_ai-text-models", + "vertex_ai-code-text-models", + "vertex_ai-embedding-models", + "vertex_ai-anthropic_models", + "vertex_ai-llama_models", + "vertex_ai-deepseek_models", + "vertex_ai-mistral_models", + "vertex_ai-ai21_models", + "vertex_ai-image-models", + "vertex_ai-video-models", + "vertex_ai-openai_models", + "vertex_ai-minimax_models", + "vertex_ai-moonshot_models", + "vertex_ai-zai_models", + "vertex_ai-qwen_models", +) + + +@contextmanager +def _registered_rows(label_per_key: Mapping[str, str]) -> Generator[None, None, None]: + import litellm + + try: + litellm.register_model( + {key: {"litellm_provider": label, "mode": "chat"} for key, label in label_per_key.items()}, + persist_across_reloads=False, + ) + yield + finally: + provider_model_sets: Final = tuple(value for value in vars(litellm).values() if isinstance(value, set)) + for key in label_per_key: + for provider_models in provider_model_sets: + provider_models.discard(key) + provider_models.discard(key.removeprefix("vertex_ai/")) + litellm.add_known_models(model_cost_map={}) + + +@pytest.mark.parametrize("label", VERTEX_COST_MAP_LABELS) +def test_vertex_wildcard_lists_a_row_registered_under_any_vertex_label(label: str): + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + key: Final = f"vertex_ai/fake-{label}-row" + with _registered_rows({key: label}): + listed: Final = get_known_models_from_wildcard("vertex_ai/*") + + assert key in listed + + +@pytest.mark.parametrize( + ("wildcard_model", "listed_prefix"), + [("vertex_ai/*", "vertex_ai"), ("my_vertex/*", "my_vertex")], +) +def test_wildcard_keeps_a_publisher_segment_that_names_another_provider(wildcard_model: str, listed_prefix: str): + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + with _registered_rows( + { + "vertex_ai/openai/fake-publisher-maas": "vertex_ai-openai_models", + "vertex_ai/meta/fake-publisher-maas": "vertex_ai-llama_models", + } + ): + listed: Final = get_known_models_from_wildcard( + wildcard_model=wildcard_model, + litellm_params=LiteLLM_Params(model="vertex_ai/*"), + ) + + assert f"{listed_prefix}/openai/fake-publisher-maas" in listed + assert f"{listed_prefix}/meta/fake-publisher-maas" in listed + assert f"{listed_prefix}/fake-publisher-maas" not in listed + + +@pytest.mark.parametrize("provider", ["replicate", "wandb"]) +def test_wildcard_lists_each_model_once_under_its_full_owner_id(provider: str): + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + full_ids: Final = {f"{provider}/{model.removeprefix(f'{provider}/')}" for model in litellm.models_by_provider[provider]} + assert any(model.count("/") > 1 for model in full_ids) + + assert sorted(get_known_models_from_wildcard(f"{provider}/*")) == sorted(full_ids) + + def test_get_complete_model_list_drops_no_default_models_sentinel(): from litellm.proxy.auth.model_checks import get_complete_model_list diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 2aceed493a7..8ca678a94ad 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -781,18 +781,6 @@ class TestBedrockClaude35SonnetContextWindow: PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset({"sagemaker", "bedrock_converse"}) MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY: Final = frozenset({"search", "decisions", "evaluation"}) -VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST: Final = frozenset( - { - "vertex_ai-ai21_models", - "vertex_ai-embedding-models", - "vertex_ai-image-models", - "vertex_ai-llama_models", - "vertex_ai-mistral_models", - "vertex_ai-openai_models", - "vertex_ai-qwen_models", - "vertex_ai-video-models", - } -) def is_registered_provider(label: str, model_names: tuple[str, ...]) -> bool: @@ -814,7 +802,6 @@ def unregistered_providers(rows: Mapping[str, object]) -> list[str]: and "litellm_provider" in entry and entry.get("mode") not in MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY and entry["litellm_provider"] not in PROVIDER_LABELS_WITHOUT_A_MODEL_SET - and entry["litellm_provider"] not in VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST ) return sorted( label From 9a68573e9b6b134732e33c53eb40d68f1711aa72 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 29 Oct 2026 00:33:40 -0700 Subject: [PATCH 07/58] fix(responses): keep each Codex tool once when its tool list repeats across turns (#49006) * test(responses): cover Codex tool lists repeated across turns on the Bedrock chat bridge and Mantle * fix(responses): keep each Codex tool once when its tool list repeats across turns * test(responses): type the new tool list dedupe unit tests --- litellm/responses/additional_tools.py | 63 +++++- .../count_request.py | 13 +- ...es_codex_repeated_additional_tools_wire.py | 204 ++++++++++++++++++ .../test_count_request.py | 12 ++ tests/unit/responses/test_additional_tools.py | 76 ++++++- 5 files changed, 353 insertions(+), 15 deletions(-) create mode 100644 tests/integration/providers/test_responses_codex_repeated_additional_tools_wire.py diff --git a/litellm/responses/additional_tools.py b/litellm/responses/additional_tools.py index 4b897ff58c0..ce530447c0d 100644 --- a/litellm/responses/additional_tools.py +++ b/litellm/responses/additional_tools.py @@ -1,5 +1,8 @@ -from collections.abc import Sequence +import json +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass +from itertools import groupby +from types import MappingProxyType from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools from pydantic import ValidationError @@ -15,8 +18,8 @@ class _InputItemType(LiteLLMBaseModel): type: str = "" -class _AdditionalToolsItem(LiteLLMBaseModel): - tools: tuple[dict[str, object], ...] = () +class _ToolList(LiteLLMBaseModel): + tools: tuple[Mapping[str, object], ...] = () @dataclass(frozen=True, slots=True) @@ -33,12 +36,54 @@ def _is_additional_tools_item(item: object) -> bool: return False -def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: +def _tool_list_of(value: object) -> tuple[Mapping[str, object], ...]: try: - parsed: Final = _AdditionalToolsItem.model_validate(item) + return _ToolList.model_validate(value).tools except ValidationError: return () - return tuple(cast("ALL_RESPONSES_API_TOOL_PARAMS", tool) for tool in parsed.tools) + + +def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + return tuple(cast("ALL_RESPONSES_API_TOOL_PARAMS", tool) for tool in _tool_list_of(item)) + + +def _tool_identity(tool: Mapping[str, object]) -> str: + name: Final = tool.get("name") + if isinstance(name, str) and name: + return f"{tool.get('type')}:{name}" + return json.dumps(tool, sort_keys=True, default=str) + + +def _members_of_copies(copies: Sequence[Mapping[str, object]]) -> Iterator[Mapping[str, object]]: + for namespace in copies: + yield from _tool_list_of(namespace) + + +def _merged(copies: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + newest: Final = copies[-1] + if newest.get("type") != "namespace" or len(copies) == 1: + return newest + members: Final = tuple(_members_of_copies(copies)) + return {**newest, "tools": list(_unique_tools(members))} if members else newest + + +def _identity_of(keyed_tool: tuple[str, Mapping[str, object]]) -> str: + return keyed_tool[0] + + +def _unique_tools(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + """Codex re-sends its tool list in the input on every turn, so the same tool arrives many times. + Keep each tool once, at its first position, with the newest copy's definition, and merge the + members of a namespace sent more than once.""" + keyed: Final = tuple((_tool_identity(tool), tool) for tool in tools) + copies_by_identity: Final = MappingProxyType( + { + identity: tuple(tool for _, tool in group) + for identity, group in groupby(sorted(keyed, key=_identity_of), key=_identity_of) + } + ) + first_seen_order: Final = tuple(dict.fromkeys(identity for identity, _ in keyed)) + return tuple(_merged(copies_by_identity[identity]) for identity in first_seen_order) def hoist_additional_tools( @@ -58,4 +103,8 @@ def hoist_additional_tools( len(items), ) remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)] - return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted) + unique: Final = tuple( + cast("ALL_RESPONSES_API_TOOL_PARAMS", tool) # cast-ok: a merged namespace is rebuilt from the same tool dicts + for tool in _unique_tools((*existing, *hoisted)) + ) + return HoistedAdditionalTools(input=remaining_input, tools=unique, hoisted=hoisted) diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 7679b4ede71..8e8e4066bb8 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -20,7 +20,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.responses.main import previous_response_client_ids -from litellm.types.llms.openai import ResponseInputParam +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam from litellm.utils import ProviderConfigManager _COUNTED_CHAT_PARAMS: Final = ( @@ -119,14 +119,13 @@ def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> M if not isinstance(responses_input, list): return responses_request input_items: Final = cast(ResponseInputParam, responses_input) # cast-ok: forwarded as sent, as inference does - hoisted: Final = hoist_additional_tools(input_items, None) + top_level_tools: Final = cast( # cast-ok: tools are forwarded as sent, as inference does + "tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]", _ITEMS.validate_python(responses_request.get("tools") or ()) + ) + hoisted: Final = hoist_additional_tools(input_items, top_level_tools) if not hoisted.hoisted: return responses_request - return { - **responses_request, - "input": hoisted.input, - "tools": [*_ITEMS.validate_python(responses_request.get("tools") or ()), *hoisted.hoisted], - } + return {**responses_request, "input": hoisted.input, "tools": list(hoisted.tools)} def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mapping[str, object]: diff --git a/tests/integration/providers/test_responses_codex_repeated_additional_tools_wire.py b/tests/integration/providers/test_responses_codex_repeated_additional_tools_wire.py new file mode 100644 index 00000000000..e03b58188f8 --- /dev/null +++ b/tests/integration/providers/test_responses_codex_repeated_additional_tools_wire.py @@ -0,0 +1,204 @@ +import json +import uuid +from collections.abc import Mapping +from typing import Final + +import httpx +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +BEDROCK_GPT: Final = "us.openai.gpt-5.6-sol" +MANTLE_MODEL: Final = "bedrock_mantle/openai.gpt-5.6-sol" +TOKEN: Final = "synthetic-bedrock-bearer" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +TOOL_SPECS: Final = TypeAdapter(tuple[dict[str, dict[str, JsonValue]], ...]) + + +def function_tool(name: str, description: str = "") -> dict[str, JsonValue]: + return { + "type": "function", + "name": name, + "description": description or f"run {name}", + "parameters": {"type": "object", "properties": {"arg": {"type": "string"}}}, + } + + +def namespace_tool(name: str, *members: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "namespace", "name": name, "description": "", "tools": list(members)} + + +def tool_list_item(*tools: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "additional_tools", "role": "developer", "tools": list(tools)} + + +def user_turn(text: str) -> dict[str, JsonValue]: + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]} + + +def assistant_turn(text: str) -> dict[str, JsonValue]: + return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} + + +def converse_tool_specs(request: Request) -> tuple[dict[str, JsonValue], ...]: + tool_config: Final = JSON_OBJECT.validate_json(request.body).get("toolConfig") + assert isinstance(tool_config, dict), request.body + return tuple(entry["toolSpec"] for entry in TOOL_SPECS.validate_python(tool_config["tools"])) + + +def bedrock_converse_peer(request: Request) -> Reply: + names: Final = tuple(str(spec["name"]) for spec in converse_tool_specs(request)) + repeated: Final = tuple(name for index, name in enumerate(names) if name in names[:index]) + if repeated: + return Reply( + status=400, + body=json.dumps( + { + "message": f"The tool {repeated[0]} is already defined at toolConfig.tools.{names.index(repeated[0])}." + } + ).encode(), + ) + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "converse wire control"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 30, "outputTokens": 3, "totalTokens": 33}, + "metrics": {"latencyMs": 10}, + } + ).encode() + ) + + +def send_to_bedrock_converse(gateway: Gateway, history: list[JsonValue]) -> tuple[httpx.Response, tuple[Request, ...]]: + with wire_server(bedrock_converse_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/converse/{BEDROCK_GPT}", + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + response: Final = gateway.request( + "POST", "/v1/responses", {"model": model, "input": history, "store": False, "cache": {"no-cache": True}} + ) + return response, wire.drain() + + +def forwarded_specs_by_name(forwarded: tuple[Request, ...]) -> Mapping[str, dict[str, JsonValue]]: + assert len(forwarded) == 1, forwarded + return {str(spec["name"]): spec for spec in converse_tool_specs(forwarded[0])} + + +def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_converse_once(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + tools: Final = tool_list_item( + namespace_tool("collaboration", function_tool("followup_task")), function_tool("exec_command") + ) + history: Final[list[JsonValue]] = [ + tools, + user_turn(f"first {marker}"), + assistant_turn("ok"), + tools, + user_turn(f"second {marker}"), + assistant_turn("ok"), + tools, + user_turn(f"third {marker}"), + ] + + response, forwarded = send_to_bedrock_converse(gateway, history) + + assert response.status_code == 200, response.text + assert response.json()["output"][0]["content"][0]["text"] == "converse wire control", response.text + names: Final = [str(spec["name"]) for spec in converse_tool_specs(forwarded[0])] + assert sorted(names) == ["collaboration__followup_task", "exec_command"], names + + +def test_newest_copy_of_a_repeated_codex_tool_is_what_bedrock_converse_receives(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + tool_list_item(function_tool("exec_command", "old description")), + user_turn(f"first {marker}"), + assistant_turn("ok"), + tool_list_item(function_tool("exec_command", "new description")), + user_turn(f"second {marker}"), + ] + + response, forwarded = send_to_bedrock_converse(gateway, history) + + assert response.status_code == 200, response.text + specs: Final = forwarded_specs_by_name(forwarded) + assert list(specs) == ["exec_command"], specs + assert specs["exec_command"]["description"] == "new description", specs + + +def test_distinct_codex_tools_spread_across_turns_all_reach_bedrock_converse(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + tool_list_item(function_tool("exec_command"), namespace_tool("mcp__docs__", function_tool("search"))), + user_turn(f"first {marker}"), + assistant_turn("ok"), + tool_list_item(function_tool("apply_patch"), namespace_tool("mcp__docs__", function_tool("fetch"))), + user_turn(f"second {marker}"), + ] + + response, forwarded = send_to_bedrock_converse(gateway, history) + + assert response.status_code == 200, response.text + assert sorted(forwarded_specs_by_name(forwarded)) == [ + "apply_patch", + "exec_command", + "mcp__docs____fetch", + "mcp__docs____search", + ] + + +def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_mantle_once(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = function_tool("shell") + collaboration: Final = namespace_tool("collaboration", function_tool("followup_task")) + + def mantle_peer(request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "id": "resp_synthetic_mantle", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "openai.gpt-5.6-sol", + "output": [ + { + "type": "message", + "id": "msg_synthetic_mantle", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "mantle wire control", "annotations": []}], + } + ], + "usage": {"input_tokens": 21, "output_tokens": 4, "total_tokens": 25}, + } + ).encode() + ) + + with wire_server(mantle_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=MANTLE_MODEL, api_key=TOKEN, api_base=wire.url, aws_region_name="us-east-2") + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": [ + tool_list_item(shell, collaboration), + user_turn(f"first {marker}"), + assistant_turn("ok"), + tool_list_item(shell, collaboration), + user_turn(f"second {marker}"), + ], + "store": False, + }, + ) + forwarded: Final = wire.drain() + + assert response.status_code == 200, response.text + assert len(forwarded) == 1, forwarded + assert JSON_OBJECT.validate_json(forwarded[0].body)["tools"] == [shell, collaboration], forwarded[0].body diff --git a/tests/unit/responses/litellm_completion_transformation/test_count_request.py b/tests/unit/responses/litellm_completion_transformation/test_count_request.py index 0477ae15b8e..a07076fb64f 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_count_request.py +++ b/tests/unit/responses/litellm_completion_transformation/test_count_request.py @@ -94,6 +94,18 @@ async def test_claude_count_request_counts_additional_tools_input_items_as_tools assert hoisted == await _claude_count({"input": [_QUESTION], "tools": [_WEATHER_TOOL]}, policy) +@pytest.mark.asyncio +async def test_claude_count_request_counts_a_tool_list_repeated_on_every_turn_once() -> None: + policy: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset()) + tool_list: Final[dict[str, JsonValue]] = {"type": "additional_tools", "tools": [_WEATHER_TOOL]} + + repeated: Final = await _claude_count( + {"input": [tool_list, _QUESTION, tool_list, _QUESTION], "tools": [_WEATHER_TOOL]}, policy + ) + + assert repeated == await _claude_count({"input": [_QUESTION, _QUESTION], "tools": [_WEATHER_TOOL]}, policy) + + @pytest.mark.asyncio async def test_claude_count_request_drops_a_malformed_tool_flag_under_drop_params_like_inference() -> None: policy: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset()) diff --git a/tests/unit/responses/test_additional_tools.py b/tests/unit/responses/test_additional_tools.py index bef3b27eacd..8844d07b105 100644 --- a/tests/unit/responses/test_additional_tools.py +++ b/tests/unit/responses/test_additional_tools.py @@ -1,6 +1,14 @@ +from typing import Final + +from pydantic import JsonValue + from litellm.responses.additional_tools import hoist_additional_tools -_EXEC_TOOL = {"type": "custom", "name": "exec", "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}} +_EXEC_TOOL = { + "type": "custom", + "name": "exec", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, +} _WAIT_TOOL = {"type": "function", "name": "wait", "parameters": {"type": "object", "properties": {}}} _TOP_LEVEL_TOOL = {"type": "function", "name": "top_level", "parameters": {"type": "object", "properties": {}}} _USER_MESSAGE = {"type": "message", "role": "user", "content": "Run ls"} @@ -46,3 +54,69 @@ def test_additional_tools_item_without_a_tools_list_is_stripped_and_contributes_ assert hoisted.input == [_USER_MESSAGE] assert hoisted.tools == () assert hoisted.hoisted == () + + +def _tool_list(*tools: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "additional_tools", "role": "developer", "tools": list(tools)} + + +def _namespace(name: str, *members: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "namespace", "name": name, "description": "", "tools": list(members)} + + +def _function(name: str, description: str = "run it") -> dict[str, JsonValue]: + return {"type": "function", "name": name, "description": description, "parameters": {"type": "object"}} + + +def test_tool_list_repeated_on_every_turn_keeps_each_tool_once_with_the_newest_definition() -> None: + collaboration: Final = _namespace("collaboration", _function("followup_task")) + request_input: Final = [ + _tool_list(collaboration, _function("exec", "old")), + _USER_MESSAGE, + _tool_list(collaboration, _function("exec", "new")), + _USER_MESSAGE, + ] + + hoisted: Final = hoist_additional_tools(request_input, None) + + assert hoisted.tools == (collaboration, _function("exec", "new")) + + +def test_namespace_sent_in_several_turns_keeps_every_distinct_member_once() -> None: + request_input: Final = [ + _tool_list(_namespace("mcp__docs__", _function("search"), _function("fetch", "old"))), + _USER_MESSAGE, + _tool_list(_namespace("mcp__docs__", _function("fetch", "new"), _function("list"))), + ] + + hoisted: Final = hoist_additional_tools(request_input, None) + + assert hoisted.tools == ( + _namespace("mcp__docs__", _function("search"), _function("fetch", "new"), _function("list")), + ) + + +def test_unnamed_tools_collapse_only_when_identical() -> None: + docs_server: Final[dict[str, JsonValue]] = { + "type": "mcp", + "server_label": "docs", + "server_url": "https://docs.example", + } + wiki_server: Final[dict[str, JsonValue]] = { + "type": "mcp", + "server_label": "wiki", + "server_url": "https://wiki.example", + } + request_input: Final = [_tool_list(docs_server, wiki_server), _USER_MESSAGE, _tool_list(docs_server)] + + hoisted: Final = hoist_additional_tools(request_input, None) + + assert hoisted.tools == (docs_server, wiki_server) + + +def test_top_level_tool_repeated_in_a_tool_list_item_is_kept_once() -> None: + request_input: Final = [_tool_list(_TOP_LEVEL_TOOL, _WAIT_TOOL), _USER_MESSAGE] + + hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL]) + + assert hoisted.tools == (_TOP_LEVEL_TOOL, _WAIT_TOOL) From 904c092ffe9c6e39b7354f77648d7a18373084ad Mon Sep 17 00:00:00 2001 From: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:39:11 +0800 Subject: [PATCH 08/58] feat(providers): add BigModel support and model pricing (#48563) * feat: add BigModel provider with native Chat, Responses, and Messages APIs Co-authored-by: Bytechoreographer * test(providers): make BigModel unit tests a package Co-authored-by: Bytechoreographer * docs(providers): clarify BigModel parameter interface Co-authored-by: Bytechoreographer * fix(bigmodel): keep the Coding Plan base, surface Responses auth errors, and scope routes Chat Completions keeps a GLM Coding Plan api_base. Responses raises the status from BigModel's success=false body instead of returning an empty 200, and a JSON body answering a stream request is replayed as a one-shot stream. Messages streams keep ping and [DONE] lines. bigmodel leaves openai_compatible_providers so speech, image, and transcription routes send it no request * fix(responses): replay a JSON body on a stream request only when it has a Content-Length A gateway that labels a chunked SSE stream as application/json kept streaming before the replay, so only a sized body is read in full * fix(bigmodel): scope the JSON stream replay to BigModel, map its business error codes, and echo requested tools * fix(bigmodel): forward Codex's include and parallel_tool_calls, keep the OpenAI org header off BigModel, and pass unchanged Messages stream events through * fix(bigmodel): map disabled thinking to low effort, keep Messages temperature, raise context window and content policy errors, and default parallel_tool_calls to true * fix(bigmodel): keep the caller's reasoning_effort when Messages rewrites disabled thinking * fix(bigmodel): keep context-window and content-policy errors on streamed Responses replays --------- Co-authored-by: Bytechoreographer Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + examples/bigmodel.md | 63 + litellm/__init__.py | 5 + .../exception_mapping_utils.py | 3 + .../get_llm_provider_logic.py | 7 + .../llms/base_llm/responses/transformation.py | 3 + litellm/llms/bigmodel/__init__.py | 0 litellm/llms/bigmodel/chat/__init__.py | 0 litellm/llms/bigmodel/chat/transformation.py | 59 + litellm/llms/bigmodel/common_utils.py | 92 ++ litellm/llms/bigmodel/messages/__init__.py | 0 .../llms/bigmodel/messages/transformation.py | 206 ++++ litellm/llms/bigmodel/responses/__init__.py | 0 .../llms/bigmodel/responses/transformation.py | 263 ++++ litellm/llms/custom_httpx/llm_http_handler.py | 79 +- litellm/main.py | 8 +- ...odel_prices_and_context_window_backup.json | 76 ++ .../provider_endpoints_support_backup.json | 26 + .../provider_create_fields.json | 28 + litellm/types/utils.py | 1 + litellm/utils.py | 16 + model_prices_and_context_window.json | 76 ++ provider_endpoints_support.json | 26 + tests/unit/llms/bigmodel/__init__.py | 0 .../llms/bigmodel/test_bigmodel_provider.py | 1066 +++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 13 + .../public_endpoints/test_public_endpoints.py | 20 + .../public/assets/logos/bigmodel.png | Bin 0 -> 11339 bytes .../CredentialModal.integration.test.tsx | 42 + .../components/provider_info_helpers.test.tsx | 12 + .../src/components/provider_info_helpers.tsx | 5 + 31 files changed, 2170 insertions(+), 26 deletions(-) create mode 100644 examples/bigmodel.md create mode 100644 litellm/llms/bigmodel/__init__.py create mode 100644 litellm/llms/bigmodel/chat/__init__.py create mode 100644 litellm/llms/bigmodel/chat/transformation.py create mode 100644 litellm/llms/bigmodel/common_utils.py create mode 100644 litellm/llms/bigmodel/messages/__init__.py create mode 100644 litellm/llms/bigmodel/messages/transformation.py create mode 100644 litellm/llms/bigmodel/responses/__init__.py create mode 100644 litellm/llms/bigmodel/responses/transformation.py create mode 100644 tests/unit/llms/bigmodel/__init__.py create mode 100644 tests/unit/llms/bigmodel/test_bigmodel_provider.py create mode 100644 ui/litellm-dashboard/public/assets/logos/bigmodel.png diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4768eeb633e..e7865817048 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -236,6 +236,7 @@ jobs: tests/unit/llms/azure_ai tests/unit/llms/base_llm tests/unit/llms/baseten + tests/unit/llms/bigmodel tests/unit/llms/black_forest_labs tests/unit/llms/brave tests/unit/llms/bytez diff --git a/examples/bigmodel.md b/examples/bigmodel.md new file mode 100644 index 00000000000..88a69796ebe --- /dev/null +++ b/examples/bigmodel.md @@ -0,0 +1,63 @@ +# BigModel (Zhipu AI) + +Use `bigmodel/` for the domestic BigModel platform and `zai/` for international Z.AI. Each platform uses its own credentials and default URL + +```yaml +model_list: + - model_name: glm + litellm_params: + model: bigmodel/glm-5.3 + api_key: os.environ/BIGMODEL_API_KEY + api_base: https://open.bigmodel.cn +``` + +One model configuration supports all three native APIs. The provider selects the upstream endpoint from the requested API + +| Proxy endpoint | BigModel endpoint | +| --- | --- | +| `/v1/chat/completions` | `/api/paas/v4/chat/completions` | +| `/v1/responses` | `/api/v1/responses` | +| `/v1/messages` | `/api/anthropic/v1/messages` | + +`api_base` is optional and defaults to `https://open.bigmodel.cn`. Set `BIGMODEL_API_BASE` to override the default, or pass an explicit `api_base` to override the environment value. Existing API-specific base URLs and complete endpoint URLs are normalized, preserving a custom proxy's path prefix + +A GLM Coding Plan key uses `https://open.bigmodel.cn/api/coding/paas/v4` as its `api_base`, since Coding Plan and standard keys do not work on each other's endpoints. Chat Completions then calls the Coding Plan endpoint, and Responses and Messages keep the endpoints in the table above + +The Admin UI lists BigModel (Zhipu AI), displays its official logo, and provides a required password API Key field and an optional API Base field with the default URL. Authentication uses an explicit `api_key` or `BIGMODEL_API_KEY` + +```python +import litellm + +chat = litellm.completion( + model="bigmodel/glm-5.3", + messages=[{"role": "user", "content": "Hello"}], + thinking={"type": "enabled"}, + reasoning_effort="low", +) + +response = litellm.responses( + model="bigmodel/glm-5.3", + input="Hello", + reasoning={"effort": "low"}, + store=True, +) + +async def messages(): + return await litellm.anthropic.messages.acreate( + model="bigmodel/glm-5.3", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + thinking={"type": "adaptive"}, + output_config={"effort": "low"}, + ) +``` + +Each API retains its native response format and streaming events. The native Messages SDK entry point is asynchronous, as is the proxy's `/v1/messages` route + +GLM-5.3, GLM-5.3-Flash, and GLM-5.3-FlashX use always-on thinking. GLM-5.3 supports `low`, `high`, and `max` reasoning effort. Messages preserves native `thinking`, `output_config`, and `reasoning_effort` parameters + +Set `store=True` when creating a Responses request if you need continuation, retrieval, or deletion. Use the response ID returned by LiteLLM. BigModel documents a seven-day retention period and does not support cancel, compact, native WebSocket, or file_search + +The model catalog and packaged backup include GLM-5.3, GLM-5.3-Flash, and GLM-5.3-FlashX. Prices are expressed in USD/token using a static 7.0 CNY/USD conversion, with the source and verification date in each entry's metadata. Use existing `litellm_params` cost fields for contract-specific deployment prices + +API references: [Chat Completions](https://docs.bigmodel.cn/cn/guide/develop/openai/introduction), [Responses](https://docs.bigmodel.cn/cn/guide/develop/responses/introduction), [Messages](https://docs.bigmodel.cn/cn/guide/develop/claude/introduction), and [pricing](https://docs.bigmodel.cn/cn/guide/start/pricing), checked on 2026-10-27 diff --git a/litellm/__init__.py b/litellm/__init__.py index de8371e7a2e..651997321a3 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -637,6 +637,7 @@ watsonx_models: Set = set() gemini_models: Set = set() xai_models: Set = set() zai_models: Set = set() +bigmodel_models: Final[set[str]] = set() # mutable-ok: provider registry is populated and extended at runtime deepseek_models: Set = set() tencent_models: Set = set() runwayml_models: Set = set() @@ -883,6 +884,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje xai_models.add(key) elif value.get("litellm_provider") == "zai": zai_models.add(key) + elif value.get("litellm_provider") == "bigmodel": + bigmodel_models.add(key) elif value.get("litellm_provider") == "fal_ai": fal_ai_models.add(key) elif value.get("litellm_provider") == "deepseek": @@ -1113,6 +1116,7 @@ model_list = list( | text_completion_inception_models | xai_models | zai_models + | bigmodel_models | fal_ai_models | deepseek_models | modelscope_models @@ -1236,6 +1240,7 @@ def _build_models_by_provider() -> dict: "text-completion-inception": text_completion_inception_models, "xai": xai_models, "zai": zai_models, + "bigmodel": bigmodel_models, "fal_ai": fal_ai_models, "deepseek": deepseek_models, "tencent": tencent_models, diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 47ec8331c83..703550dd2b5 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string +from litellm.llms.bigmodel.common_utils import raise_bigmodel_business_error from litellm.types.utils import LlmProviders from ..exceptions import ( @@ -2762,6 +2763,8 @@ def exception_type( exception_provider=exception_provider, extra_information=extra_information, ) + if custom_llm_provider == "bigmodel": + raise_bigmodel_business_error(model, f"{exception_provider} - {error_str}", original_exception) # pyright: ignore[reportUnknownArgumentType] # exception_type takes untyped parameters if custom_llm_provider == "openrouter": _map_openrouter_exception( model=model, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index d9933087334..f8c8ff71f8e 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -794,6 +794,13 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, ) = litellm.XAIChatConfig().get_openai_compatible_provider_info(api_base, api_key) + elif custom_llm_provider == "bigmodel": + from litellm.llms.bigmodel.common_utils import get_bigmodel_api_key, get_bigmodel_url + + ( + api_base, # rebind-ok: existing provider-resolution pipeline resolves the API base in place + dynamic_api_key, # rebind-ok: existing provider-resolution pipeline resolves the API key in place + ) = get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key) elif custom_llm_provider == "zai": ( api_base, diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 01baeed9bef..7f45c1eba28 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -268,6 +268,9 @@ class BaseResponsesAPIConfig(ABC): """Returns True if litellm should fake a stream for the given model and stream value""" return False + def may_answer_stream_with_json(self) -> bool: + return False + def supports_native_websocket(self) -> bool: """ Returns True if the provider has a native WebSocket endpoint for Responses API. diff --git a/litellm/llms/bigmodel/__init__.py b/litellm/llms/bigmodel/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bigmodel/chat/__init__.py b/litellm/llms/bigmodel/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bigmodel/chat/transformation.py b/litellm/llms/bigmodel/chat/transformation.py new file mode 100644 index 00000000000..12d5e912c15 --- /dev/null +++ b/litellm/llms/bigmodel/chat/transformation.py @@ -0,0 +1,59 @@ +from collections.abc import Mapping +from typing import Final + +from litellm.llms.zai.chat.transformation import ZAIChatConfig + +from ..common_utils import get_bigmodel_api_key, get_bigmodel_url, is_unsupported_disabled_thinking + +_SUPPORTED_PARAMS: Final = ( + "max_tokens", + "max_completion_tokens", + "stream", + "stream_options", + "temperature", + "top_p", + "stop", + "tools", + "tool_choice", + "response_format", + "thinking", + "reasoning_effort", +) + + +class BigModelChatConfig(ZAIChatConfig): + @property + def custom_llm_provider(self) -> str: + return "bigmodel" + + def get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: + return get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key) + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited interface requires a list + return list(_SUPPORTED_PARAMS) + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body + mapped: Final = { + **optional_params, + **{ + ("max_tokens" if key == "max_completion_tokens" else key): value + for key, value in non_default_params.items() + if key in _SUPPORTED_PARAMS and key != "thinking" + }, + } + if "thinking" not in non_default_params: + return mapped + thinking: Final = non_default_params["thinking"] + if is_unsupported_disabled_thinking(model, thinking): + return { + **mapped, + "reasoning_effort": mapped.get("reasoning_effort", "low"), + "extra_body": {"thinking": {"type": "enabled"}}, + } + return {**mapped, "extra_body": {"thinking": thinking}} diff --git a/litellm/llms/bigmodel/common_utils.py b/litellm/llms/bigmodel/common_utils.py new file mode 100644 index 00000000000..0c9b6da5c8f --- /dev/null +++ b/litellm/llms/bigmodel/common_utils.py @@ -0,0 +1,92 @@ +import re +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit, urlunsplit + +import httpx +from pydantic import JsonValue, TypeAdapter, ValidationError + +import litellm +from litellm.exceptions import ContentPolicyViolationError, ContextWindowExceededError +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str + +BIGMODEL_API_BASE: Final = "https://open.bigmodel.cn" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_BUSINESS_CODE_IN_ERROR: Final = re.compile(r"""['"]code['"]\s*:\s*['"]?(\d{4})\b""") +_CODING_CHAT_PATH: Final = "/api/coding/paas/v4" +_API_PATHS: Final = ( + "/api/coding/paas/v4/chat/completions", + _CODING_CHAT_PATH, + "/api/paas/v4/chat/completions", + "/api/anthropic/v1/messages", + "/api/anthropic/v1", + "/api/v1/responses", + "/api/paas/v4", + "/api/anthropic", + "/api/v1", + "/api", +) +_ENDPOINT_PATHS: Final = MappingProxyType( + { + "chat": "/api/paas/v4", + "responses": "/api/v1/responses", + "messages": "/api/anthropic/v1/messages", + } +) + + +def get_bigmodel_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("BIGMODEL_API_KEY") + + +def get_bigmodel_url(api_base: str | None, endpoint: Literal["chat", "responses", "messages"]) -> str: + resolved_base: Final = api_base or get_secret_str("BIGMODEL_API_BASE") or BIGMODEL_API_BASE + parsed: Final = urlsplit(resolved_base) + path: Final = parsed.path.rstrip("/") + suffix: Final = next((candidate for candidate in _API_PATHS if path.endswith(candidate)), "") + endpoint_path: Final = ( + _CODING_CHAT_PATH if endpoint == "chat" and suffix.startswith(_CODING_CHAT_PATH) else _ENDPOINT_PATHS[endpoint] + ) + return urlunsplit( + (parsed.scheme, parsed.netloc, f"{path.removesuffix(suffix)}{endpoint_path}", parsed.query, parsed.fragment) + ) + + +def bigmodel_json_body(response: httpx.Response) -> Mapping[str, JsonValue]: + try: + return _JSON_OBJECT.validate_json(response.content) + except ValidationError: + raise BaseLLMException(message=response.text, status_code=response.status_code, headers=response.headers) + + +def is_unsupported_disabled_thinking(model: str, thinking: object) -> bool: + return thinking == {"type": "disabled"} and litellm.utils.supports_factory( + model=model, custom_llm_provider="bigmodel", key="thinking_always_on" + ) + + +def bigmodel_business_error( + code: JsonValue, message: str, model: str, response: httpx.Response | None +) -> ContextWindowExceededError | ContentPolicyViolationError | None: + match code: + case 1261 | "1261": + return ContextWindowExceededError(message=message, model=model, llm_provider="bigmodel", response=response) + case 1301 | "1301": + return ContentPolicyViolationError(message=message, model=model, llm_provider="bigmodel", response=response) + case _: + return None + + +def raise_bigmodel_business_error(model: object, error_str: str, original_exception: object) -> None: + found: Final = _BUSINESS_CODE_IN_ERROR.search(error_str) + response: Final = getattr(original_exception, "response", None) + error: Final = bigmodel_business_error( + found.group(1) if found else None, + error_str, + str(model), + response if isinstance(response, httpx.Response) else None, + ) + if error is not None: + raise error diff --git a/litellm/llms/bigmodel/messages/__init__.py b/litellm/llms/bigmodel/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bigmodel/messages/transformation.py b/litellm/llms/bigmodel/messages/transformation.py new file mode 100644 index 00000000000..82ce59631f1 --- /dev/null +++ b/litellm/llms/bigmodel/messages/transformation.py @@ -0,0 +1,206 @@ +import json +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Final + +import httpx +from openai._streaming import ServerSentEvent, SSEDecoder +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.openai_like.messages.transformation import ( + DEFAULT_ANTHROPIC_API_VERSION, + OpenAILikeAnthropicMessagesConfig, +) +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.router import GenericLiteLLMParams + +from ..common_utils import ( + bigmodel_json_body, + get_bigmodel_api_key, + get_bigmodel_url, + is_unsupported_disabled_thinking, +) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_OUTPUT_CONFIG: Final = TypeAdapter(dict[str, object]) + + +def _with_lowest_effort_thinking( + params: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: the parent transform takes a request dict + requested_output_config: Final = params.get("output_config") + output_config: Final = { + **({} if "reasoning_effort" in params else {"effort": "low"}), + **( + _OUTPUT_CONFIG.validate_python(requested_output_config) if isinstance(requested_output_config, dict) else {} + ), + } + return {**params, "thinking": {"type": "adaptive"}, **({"output_config": output_config} if output_config else {})} + + +def _normalize_message(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: # mutable-ok: native JSON response + usage: Final = body.get("usage") + if not isinstance(usage, dict): + return dict(body) + server_tools: Final = usage.get("server_tool_use") + if not isinstance(server_tools, dict) or "web_fetch_requests" in server_tools: + return dict(body) + normalized: Final[JsonValue] = { + **body, + "usage": {**usage, "server_tool_use": {"web_fetch_requests": 0, **server_tools}}, + } + return normalized + + +class _BigModelMessagesStream(httpx.AsyncByteStream): + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for event in SSEDecoder().aiter_bytes(self._response.aiter_bytes()): + yield _normalize_event(event) + + async def aclose(self) -> None: + await self._response.aclose() + + +def _json_object(data: str) -> Mapping[str, JsonValue] | None: + try: + return _JSON_OBJECT.validate_json(data) + except ValidationError: + return None + + +def _normalized_data(data: str) -> str: + body: Final = _json_object(data) + if body is None: + return data + message: Final = body.get("message") + normalized: Final[JsonValue] = ( + {**body, "message": _normalize_message(message)} + if body.get("type") == "message_start" and isinstance(message, dict) + else _normalize_message(body) + ) + return data if normalized == body else json.dumps(normalized, ensure_ascii=False) + + +def _normalize_event(event: ServerSentEvent) -> bytes: + metadata: Final = "".join( + f"{name}: {value}\n" + for name, value in (("event", event.event), ("id", event.id), ("retry", event.retry)) + if value is not None + ) + data_lines: Final = "".join(f"data: {line}\n" for line in _normalized_data(event.data).split("\n")) + return f"{metadata}{data_lines}\n".encode() + + +class BigModelMessagesConfig(OpenAILikeAnthropicMessagesConfig): + @property + def custom_llm_provider(self) -> str: + return "bigmodel" + + def is_native_messages_route(self) -> bool: + return True + + def should_strip_billing_metadata(self) -> bool: + return True + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict[str, object]], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict[str, object], # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: matches dict-typed base signature + ) -> dict[str, object]: # mutable-ok: matches dict-typed base signature + request_params: Final = ( + _with_lowest_effort_thinking(anthropic_messages_optional_request_params) + if is_unsupported_disabled_thinking(model, anthropic_messages_optional_request_params.get("thinking")) + else anthropic_messages_optional_request_params + ) + return super().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the parent transform returns a bare dict + model, messages, request_params, litellm_params, headers + ) + + def transform_anthropic_messages_response( + self, model: str, raw_response: httpx.Response, logging_obj: Logging + ) -> AnthropicMessagesResponse: + normalized: Final = httpx.Response( + raw_response.status_code, + json=_normalize_message(bigmodel_json_body(raw_response)), + request=raw_response.request, + ) + return super().transform_anthropic_messages_response(model, normalized, logging_obj) + + def get_async_streaming_response_iterator( + self, + model: str, + httpx_response: httpx.Response, + request_body: dict[str, object], # mutable-ok: inherited native stream API takes a request dict + litellm_logging_obj: Logging, + ) -> AsyncIterator[bytes]: + normalized: Final = httpx.Response( + httpx_response.status_code, + headers={ + key: value + for key, value in httpx_response.headers.items() + if key.lower() not in ("content-encoding", "content-length") + }, + stream=_BigModelMessagesStream(httpx_response), + request=httpx_response.request, + ) + return super().get_async_streaming_response_iterator( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # native parent yields bytes but its iterator annotation omits the element type + model, normalized, request_body, litellm_logging_obj + ) + + @staticmethod + def _translate_reasoning_effort_to_anthropic( + model: str, optional_params: Mapping[str, object], max_tokens: int | None, custom_llm_provider: str + ) -> None: + return None + + @staticmethod + def _translate_adaptive_effort_for_non_adaptive_model( + model: str, optional_params: Mapping[str, object], max_tokens: int | None, custom_llm_provider: str + ) -> None: + return None + + @staticmethod + def _drop_incompatible_temperature_for_thinking( + model: str, optional_params: Mapping[str, object], custom_llm_provider: str + ) -> None: + return None + + def validate_anthropic_messages_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: HTTP handler merges forwarded headers + resolved_key: Final = get_bigmodel_api_key(api_key) + if resolved_key is None: + raise AuthenticationError( + message="Set BIGMODEL_API_KEY or pass api_key", llm_provider="bigmodel", model=model + ) + return { + **{key: value for key, value in headers.items() if key.lower() not in ("x-api-key", "authorization")}, + "x-api-key": resolved_key, + "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION, + "content-type": "application/json", + }, api_base + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return get_bigmodel_url(api_base, "messages") diff --git a/litellm/llms/bigmodel/responses/__init__.py b/litellm/llms/bigmodel/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bigmodel/responses/transformation.py b/litellm/llms/bigmodel/responses/transformation.py new file mode 100644 index 00000000000..80c53682d16 --- /dev/null +++ b/litellm/llms/bigmodel/responses/transformation.py @@ -0,0 +1,263 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError +from typing_extensions import Never + +from litellm.exceptions import AuthenticationError, UnsupportedParamsError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.types.llms.openai import ( + ResponseInputParam, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +from ..common_utils import bigmodel_business_error, bigmodel_json_body, get_bigmodel_api_key, get_bigmodel_url + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SUPPORTED_PARAMS: Final = ( + "input", + "model", + "instructions", + "stream", + "temperature", + "top_p", + "max_output_tokens", + "tools", + "tool_choice", + "parallel_tool_calls", + "include", + "reasoning", + "text", + "prompt_cache_key", + "previous_response_id", + "store", + "extra_headers", + "extra_query", + "extra_body", + "timeout", +) + + +_STATUS_BY_BUSINESS_CODE: Final = MappingProxyType( + { + **dict.fromkeys((1000, 1001, 1003, 1005), 401), + **dict.fromkeys((1210, 1211, 1212, 1213, 1214, 1215, 1221, 1222, 1261, 1301), 400), + 1220: 403, + **dict.fromkeys( + (1113, 1302, 1305, 1308, 1309, 1310, 1311, 1313, 1314, 1315, 1316, 1317, 1318, 1319, 1320, 1321), + 429, + ), + } +) + + +class _RequestedTools(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + tools: JsonValue = None + tool_choice: JsonValue = None + parallel_tool_calls: JsonValue = None + + +def _requested_tools(logging_obj: Logging) -> _RequestedTools: + try: + return _RequestedTools.model_validate(logging_obj.model_call_details.get("optional_params") or {}) + except ValidationError: + return _RequestedTools() + + +def _envelope_status_code(code: JsonValue) -> int: + if not isinstance(code, int) or isinstance(code, bool): + return 500 + if 400 <= code <= 599: + return code + return _STATUS_BY_BUSINESS_CODE.get(code, 500) + + +def _normalize_output_item(item: JsonValue) -> JsonValue: + if not isinstance(item, dict) or item.get("type") not in ("message", "reasoning"): + return item + content: Final = item.get("content") + normalized: Final[JsonValue] = {"summary": [], **item} if item["type"] == "reasoning" else item + if not isinstance(content, dict): + return normalized + return {**normalized, "content": [content]} + + +def _normalize_response( + body: Mapping[str, JsonValue], + requested: _RequestedTools, +) -> dict[str, JsonValue]: # mutable-ok: JSON serialization boundary + usage: Final = body.get("usage") + output: Final = body.get("output") + tools: Final = body.get("tools") + tool_choice: Final = body.get("tool_choice") + parallel_tool_calls: Final = body.get("parallel_tool_calls") + input_tokens: Final = usage.get("input_tokens") if isinstance(usage, dict) else None + output_tokens: Final = usage.get("output_tokens") if isinstance(usage, dict) else None + normalized_usage: Final = ( + { + "total_tokens": input_tokens + output_tokens, + **usage, + } + if isinstance(usage, dict) and isinstance(input_tokens, int) and isinstance(output_tokens, int) + else usage + ) + normalized: Final[JsonValue] = { + **body, + "parallel_tool_calls": ( + requested.parallel_tool_calls is not False if parallel_tool_calls is None else parallel_tool_calls + ), + "tools": (requested.tools or []) if tools is None else tools, + "tool_choice": (requested.tool_choice or "auto") if tool_choice is None else tool_choice, + **({"usage": normalized_usage} if "usage" in body else {}), + **({"output": [_normalize_output_item(item) for item in output]} if isinstance(output, list) else {}), + } + return normalized + + +def _normalize_stream_part(part: JsonValue) -> JsonValue: + if not isinstance(part, dict): + return part + if part.get("type") == "output_text": + return { + "annotations": [], + "logprobs": None, + **part, + } + if part.get("type") == "reasoning_text": + return {"reasoning": part.get("text", ""), **part} + return part + + +class BigModelResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.BIGMODEL + + def supports_native_file_search(self) -> bool: + return False + + def supports_native_websocket(self) -> bool: + return False + + def should_fake_stream( + self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None + ) -> bool: + return False + + def may_answer_stream_with_json(self) -> bool: + return True + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited API requires a list + return list(_SUPPORTED_PARAMS) + + def map_openai_params( + self, + response_api_optional_params: ResponsesAPIOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body + return {key: value for key, value in response_api_optional_params.items() if key in _SUPPORTED_PARAMS} + + def validate_environment( + self, headers: Mapping[str, str], model: str, litellm_params: GenericLiteLLMParams | None + ) -> dict[str, str]: # mutable-ok: HTTP handler merges forwarded headers + resolved_key: Final = get_bigmodel_api_key(litellm_params.api_key if litellm_params else None) + if resolved_key is None: + raise AuthenticationError( + message="Set BIGMODEL_API_KEY or pass api_key", llm_provider="bigmodel", model=model + ) + return { + **headers, + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + } + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + return get_bigmodel_url(api_base, "responses") + + def transform_response_api_response( + self, model: str, raw_response: httpx.Response, logging_obj: Logging + ) -> ResponsesAPIResponse: + body: Final = bigmodel_json_body(raw_response) + if body.get("success") is False: + message: Final = str(body.get("msg") or raw_response.text) + raise bigmodel_business_error(body.get("code"), message, model, None) or BaseLLMException( + message=message, + status_code=_envelope_status_code(body.get("code")), + headers=raw_response.headers, + ) + normalized: Final = _normalize_response(body, _requested_tools(logging_obj)) + normalized_response: Final = httpx.Response( + status_code=raw_response.status_code, + headers={ + key: value + for key, value in raw_response.headers.items() + if key.lower() not in ("content-encoding", "content-length") + }, + json=normalized, + request=raw_response.request, + ) + return super().transform_response_api_response(model, normalized_response, logging_obj) + + def transform_get_response_api_response( + self, raw_response: httpx.Response, logging_obj: Logging + ) -> ResponsesAPIResponse: + return self.transform_response_api_response("", raw_response, logging_obj) + + def transform_streaming_response( + self, model: str, parsed_chunk: Mapping[str, object], logging_obj: Logging + ) -> ResponsesAPIStreamingResponse: + chunk: Final = _JSON_OBJECT.validate_python(parsed_chunk) + response: Final = chunk.get("response") + item: Final = chunk.get("item") + normalized: Final[JsonValue] = { + **chunk, + **( + {"response": _normalize_response(response, _requested_tools(logging_obj))} + if isinstance(response, dict) + else {} + ), + **({"item": _normalize_output_item(item)} if "item" in chunk else {}), + } + event_model: Final = self.get_event_model_class(str(normalized.get("type"))) + match chunk.get("type"): + case "response.output_text.done": + return event_model.model_validate({"logprobs": [], **normalized}) + case "response.content_part.done": + return event_model.model_validate({**normalized, "part": _normalize_stream_part(chunk.get("part"))}) + case "error": + error: Final = chunk.get("error") + error_fields: Final = error if isinstance(error, dict) else chunk + normalized_error: Final[JsonValue] = { + **normalized, + "error": {"type": "api_error", **error_fields, "code": error_fields.get("code") or "unknown_error"}, + } + return event_model.model_validate(normalized_error) + case _: + return event_model.model_validate(normalized) + + def transform_cancel_response_api_request( + self, response_id: str, api_base: str, litellm_params: GenericLiteLLMParams, headers: Mapping[str, str] + ) -> Never: + raise UnsupportedParamsError(message="BigModel does not support cancelling responses", status_code=400) + + def transform_compact_response_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: Mapping[str, object], + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> Never: + raise UnsupportedParamsError(message="BigModel does not support compacting responses", status_code=400) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e7da6586b97..dc10eff65b8 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -367,6 +367,19 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] +def _is_sized_json_body(response: httpx.Response) -> bool: + headers: Final = httpx.Headers(response.headers) + media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type")) + return "content-length" in headers and "application/json" in media_types + + +def _is_json_document(content: bytes) -> bool: + try: + return isinstance(json.loads(content), dict) + except ValueError: + return False + + def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -> _SignedBodyKwargs: if signed_body is not None: return {"data": signed_body} @@ -3025,6 +3038,17 @@ class BaseLLMHTTPHandler: }, ) + replay_as_stream: Final = partial( + MockResponsesAPIStreamingIterator, + model=model, + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, + ) + try: if is_stream_request: response = sync_httpx_client.post( @@ -3036,7 +3060,13 @@ class BaseLLMHTTPHandler: ) if caller_requested_stream: if fake_stream is True: - return MockResponsesAPIStreamingIterator( + return replay_as_stream(response=response) + if not ( + responses_api_provider_config.may_answer_stream_with_json() + and _is_sized_json_body(response) + and _is_json_document(response.read()) + ): + return SyncResponsesAPIStreamingIterator( response=response, model=model, logging_obj=logging_obj, @@ -3046,17 +3076,6 @@ class BaseLLMHTTPHandler: request_data=request_context, call_type=CallTypes.responses.value, ) - - return SyncResponsesAPIStreamingIterator( - response=response, - model=model, - logging_obj=logging_obj, - responses_api_provider_config=responses_api_provider_config, - litellm_metadata=litellm_metadata, - custom_llm_provider=custom_llm_provider, - request_data=request_context, - call_type=CallTypes.responses.value, - ) response.read() else: response = sync_httpx_client.post( @@ -3071,6 +3090,9 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if is_stream_request and caller_requested_stream: + return replay_as_stream(response=response) + initial_response: Final = responses_api_provider_config.transform_response_api_response( model=model, raw_response=response, @@ -3226,6 +3248,17 @@ class BaseLLMHTTPHandler: }, ) + replay_as_stream: Final = partial( + MockResponsesAPIStreamingIterator, + model=model, + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, + ) + try: if is_stream_request: response = await async_httpx_client.post( @@ -3239,7 +3272,13 @@ class BaseLLMHTTPHandler: if caller_requested_stream: if fake_stream is True: - return MockResponsesAPIStreamingIterator( + return replay_as_stream(response=response) + if not ( + responses_api_provider_config.may_answer_stream_with_json() + and _is_sized_json_body(response) + and _is_json_document(await response.aread()) + ): + return ResponsesAPIStreamingIterator( response=response, model=model, logging_obj=logging_obj, @@ -3249,17 +3288,6 @@ class BaseLLMHTTPHandler: request_data=request_context, call_type=CallTypes.responses.value, ) - - return ResponsesAPIStreamingIterator( - response=response, - model=model, - logging_obj=logging_obj, - responses_api_provider_config=responses_api_provider_config, - litellm_metadata=litellm_metadata, - custom_llm_provider=custom_llm_provider, - request_data=request_context, - call_type=CallTypes.responses.value, - ) await response.aread() else: response = await async_httpx_client.post( @@ -3276,6 +3304,9 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if is_stream_request and caller_requested_stream: + return replay_as_stream(response=response) + initial_response: Final = responses_api_provider_config.transform_response_api_response( model=model, raw_response=response, diff --git a/litellm/main.py b/litellm/main.py index 41c4302499e..b8ad28c26fe 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2672,10 +2672,13 @@ def _complete_custom_openai( # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there api_base = _resolve_openai_api_base(api_base) + if custom_llm_provider == "bigmodel" and not api_key: + raise litellm.AuthenticationError( + message="Set BIGMODEL_API_KEY or pass api_key", llm_provider="bigmodel", model=model + ) organization = ( organization - or litellm.organization - or get_secret("OPENAI_ORGANIZATION") + or (None if custom_llm_provider == "bigmodel" else litellm.organization or get_secret("OPENAI_ORGANIZATION")) or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 ) openai.organization = organization @@ -6299,6 +6302,7 @@ def completion( or custom_llm_provider == "nebius" or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" + or custom_llm_provider == "bigmodel" or custom_llm_provider in litellm.openai_compatible_providers or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3838f7510f0..50b3560f05b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -78995,5 +78995,81 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bigmodel/glm-5.3": { + "input_cost_per_token": 1.1428571428571428e-06, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 2.857142857142857e-07, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 8, output 28, cached input 2 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true + }, + "bigmodel/glm-5.3-flash": { + "input_cost_per_token": 1.142857142857143e-07, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 3.2857142857142855e-08, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "supports_vision": true, + "supports_video_input": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 0.8, output 2.8, cached input 0.23 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true + }, + "bigmodel/glm-5.3-flashx": { + "input_cost_per_token": 2.857142857142857e-07, + "output_cost_per_token": 1e-06, + "cache_read_input_token_cost": 8.142857142857142e-08, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "supports_vision": true, + "supports_video_input": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 2, output 7, cached input 0.57 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index e40cd134759..7b0234e28d7 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -5326,6 +5326,32 @@ "a2a": false, "generateContent": false } + }, + "bigmodel": { + "display_name": "BigModel (Zhipu AI) (`bigmodel`)", + "url": "https://docs.bigmodel.cn/cn/guide/develop/responses/introduction", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } } }, "endpoints": { diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 2e72cf3d984..fd17079b15f 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4228,5 +4228,33 @@ } ], "default_model_placeholder": "cursor/claude-4-sonnet" + }, + { + "provider": "BIGMODEL", + "provider_display_name": "BigModel (Zhipu AI)", + "litellm_provider": "bigmodel", + "credential_fields": [ + { + "key": "api_key", + "label": "BigModel (Zhipu AI) API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://open.bigmodel.cn", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": "https://open.bigmodel.cn" + } + ], + "default_model_placeholder": "bigmodel/glm-5.3" } ] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 70ca05904ef..266a6a20131 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4424,6 +4424,7 @@ class LlmProviders(str, Enum): JINA_AI = "jina_ai" XAI = "xai" ZAI = "zai" + BIGMODEL = "bigmodel" CUSTOM_OPENAI = "custom_openai" TEXT_COMPLETION_OPENAI = "text-completion-openai" COHERE = "cohere" diff --git a/litellm/utils.py b/litellm/utils.py index 72c37c3760e..6a9beeec916 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5674,6 +5674,7 @@ def add_provider_specific_params_to_optional_params( "openai", "azure", "text-completion-openai", + "bigmodel", *litellm.openai_compatible_providers, ] or JSONProviderRegistry.exists(custom_llm_provider) if calls_through_openai_sdk: @@ -9183,6 +9184,7 @@ class ProviderConfigManager: LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False), LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False), LlmProviders.ZAI: (lambda: litellm.ZAIChatConfig(), False), + LlmProviders.BIGMODEL: (ProviderConfigManager._get_bigmodel_chat_config, False), LlmProviders.OPENCODE_GO: (ProviderConfigManager.get_opencode_go_config, True), LlmProviders.LAMBDA_AI: (lambda: litellm.LambdaAIChatConfig(), False), LlmProviders.INCEPTION: (lambda: litellm.InceptionChatConfig(), False), @@ -9375,6 +9377,12 @@ class ProviderConfigManager: return get_bedrock_chat_config(model=model) + @staticmethod + def _get_bigmodel_chat_config() -> BaseConfig: + from litellm.llms.bigmodel.chat.transformation import BigModelChatConfig + + return BigModelChatConfig() + @staticmethod def get_opencode_go_config(model: str) -> BaseConfig: from litellm.llms.opencode.chat.transformation import go_chat_config @@ -9696,6 +9704,10 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> BaseAnthropicMessagesConfig | None: + if provider == LlmProviders.BIGMODEL: + from litellm.llms.bigmodel.messages.transformation import BigModelMessagesConfig + + return BigModelMessagesConfig() if provider == LlmProviders.OPENCODE_GO: from litellm.llms.opencode.common_utils import go_endpoint from litellm.llms.opencode.messages.transformation import OpenCodeGoMessagesConfig @@ -9927,6 +9939,10 @@ class ProviderConfigManager: if provider is None: return None + if provider == LlmProviders.BIGMODEL: + from litellm.llms.bigmodel.responses.transformation import BigModelResponsesAPIConfig + + return BigModelResponsesAPIConfig() if provider == LlmProviders.OPENCODE_GO: from litellm.llms.opencode.common_utils import go_endpoint from litellm.llms.opencode.responses.transformation import OpenCodeGoResponsesConfig diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3838f7510f0..50b3560f05b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -78995,5 +78995,81 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bigmodel/glm-5.3": { + "input_cost_per_token": 1.1428571428571428e-06, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 2.857142857142857e-07, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 8, output 28, cached input 2 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true + }, + "bigmodel/glm-5.3-flash": { + "input_cost_per_token": 1.142857142857143e-07, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 3.2857142857142855e-08, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "supports_vision": true, + "supports_video_input": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 0.8, output 2.8, cached input 0.23 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true + }, + "bigmodel/glm-5.3-flashx": { + "input_cost_per_token": 2.857142857142857e-07, + "output_cost_per_token": 1e-06, + "cache_read_input_token_cost": 8.142857142857142e-08, + "litellm_provider": "bigmodel", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "source": "https://docs.bigmodel.cn/cn/guide/start/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "thinking_always_on": true, + "supports_vision": true, + "supports_video_input": true, + "metadata": { + "comment": "Official API list price checked 2026-10-27: input 2, output 7, cached input 0.57 CNY per million tokens; converted at 7.0 CNY/USD. Cache storage is temporarily free and is not encoded." + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_native_streaming": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index e40cd134759..7b0234e28d7 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -5326,6 +5326,32 @@ "a2a": false, "generateContent": false } + }, + "bigmodel": { + "display_name": "BigModel (Zhipu AI) (`bigmodel`)", + "url": "https://docs.bigmodel.cn/cn/guide/develop/responses/introduction", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } } }, "endpoints": { diff --git a/tests/unit/llms/bigmodel/__init__.py b/tests/unit/llms/bigmodel/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bigmodel/test_bigmodel_provider.py b/tests/unit/llms/bigmodel/test_bigmodel_provider.py new file mode 100644 index 00000000000..cdf98c8f9a4 --- /dev/null +++ b/tests/unit/llms/bigmodel/test_bigmodel_provider.py @@ -0,0 +1,1066 @@ +import asyncio +import contextlib +import gzip +import io +import json +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from functools import partial +from itertools import chain +from typing import Final, Literal + +import httpx +import pytest +import respx +from openai.types.responses import Response as OpenAIResponse +from pydantic import JsonValue + +import litellm +from litellm.llms.bigmodel.chat.transformation import BigModelChatConfig +from litellm.llms.bigmodel.common_utils import get_bigmodel_url +from litellm.llms.bigmodel.messages.transformation import BigModelMessagesConfig +from litellm.llms.bigmodel.responses.transformation import BigModelResponsesAPIConfig +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +MODEL: Final = "glm-5.3" +BASE: Final = "https://open.bigmodel.cn" +KEY: Final = "test-bigmodel-key" +BASE_PATHS: Final = ( + "", + "/", + "/api", + "/api/paas/v4", + "/api/paas/v4/", + "/api/paas/v4/chat/completions", + "/api/v1", + "/api/v1/responses", + "/api/anthropic", + "/api/anthropic/v1", + "/api/anthropic/v1/messages", +) + + +class _ByteChunks(httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks = chunks + self.is_closed = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + async def aclose(self) -> None: + self.is_closed = True + + +@pytest.fixture(autouse=True) +def bigmodel_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("BIGMODEL_API_KEY", KEY) + monkeypatch.delenv("BIGMODEL_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + yield + + +def response_payload() -> Mapping[str, JsonValue]: + return { + "id": "resp_bigmodel_test", + "created_at": 1, + "object": "response", + "model": MODEL, + "status": "completed", + "parallel_tool_calls": False, + "tools": None, + "tool_choice": None, + "output": [ + {"id": "rs_test", "type": "reasoning", "content": {"type": "reasoning_text", "text": "Thinking"}}, + { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": {"type": "output_text", "text": "Hello", "annotations": []}, + }, + ], + "usage": { + "input_tokens": 20, + "output_tokens": 10, + "input_tokens_details": {"cached_tokens": 5}, + "output_tokens_details": {"reasoning_tokens": 3}, + }, + } + + +def test_provider_resolution_and_all_native_endpoints() -> None: + model, provider, key, base = litellm.get_llm_provider(f"bigmodel/{MODEL}") + assert (model, provider, key, base) == (MODEL, "bigmodel", KEY, f"{BASE}/api/paas/v4") + assert isinstance(ProviderConfigManager.get_provider_chat_config(MODEL, LlmProviders.BIGMODEL), BigModelChatConfig) + assert isinstance( + ProviderConfigManager.get_provider_responses_api_config("bigmodel", MODEL), BigModelResponsesAPIConfig + ) + assert isinstance( + ProviderConfigManager.get_provider_anthropic_messages_config(MODEL, LlmProviders.BIGMODEL), + BigModelMessagesConfig, + ) + assert "bigmodel" in litellm.provider_list + assert litellm.get_llm_provider(f"zai/{MODEL}")[3] == "https://api.z.ai/api/paas/v4" + + +@pytest.mark.parametrize("path", BASE_PATHS) +def test_one_base_selects_each_api_without_duplicate_segments(path: str) -> None: + assert get_bigmodel_url(f"{BASE}{path}", "chat") == f"{BASE}/api/paas/v4" + assert get_bigmodel_url(f"{BASE}{path}", "responses") == f"{BASE}/api/v1/responses" + assert get_bigmodel_url(f"{BASE}{path}", "messages") == f"{BASE}/api/anthropic/v1/messages" + + +@pytest.mark.parametrize( + "path", ("/api/coding/paas/v4", "/api/coding/paas/v4/", "/api/coding/paas/v4/chat/completions") +) +def test_coding_plan_base_keeps_its_chat_path_and_shares_the_other_apis(path: str) -> None: + assert get_bigmodel_url(f"{BASE}{path}", "chat") == f"{BASE}/api/coding/paas/v4" + assert get_bigmodel_url(f"{BASE}{path}", "responses") == f"{BASE}/api/v1/responses" + assert get_bigmodel_url(f"{BASE}{path}", "messages") == f"{BASE}/api/anthropic/v1/messages" + + +def test_custom_base_and_environment_precedence(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("BIGMODEL_API_BASE", "https://env.example/prefix/api/v1") + assert get_bigmodel_url(None, "messages") == "https://env.example/prefix/api/anthropic/v1/messages" + assert ( + get_bigmodel_url("https://override.example/prefix/api/anthropic/", "chat") + == "https://override.example/prefix/api/paas/v4" + ) + assert ( + BigModelResponsesAPIConfig().validate_environment({}, MODEL, GenericLiteLLMParams(api_key="explicit"))[ + "Authorization" + ] + == "Bearer explicit" + ) + assert BigModelResponsesAPIConfig().validate_environment({}, MODEL, None)["Authorization"] == f"Bearer {KEY}" + + +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.asyncio +@respx.mock +async def test_chat_uses_native_endpoint_and_maps_token_limit(is_async: bool) -> None: + route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "object": "chat.completion", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + } + ) + params: Final = { + "model": f"bigmodel/{MODEL}", + "messages": [{"role": "user", "content": "Hello"}], + "api_base": f"{BASE}/api/anthropic", + "max_completion_tokens": 10, + "thinking": {"type": "enabled"}, + "reasoning_effort": "low", + "extra_body": {"clear_thinking": False}, + } + response: Final = await litellm.acompletion(**params) if is_async else litellm.completion(**params) + assert response.choices[0].message.content == "Hello" + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert body["model"] == MODEL + assert body["max_tokens"] == 10 + assert "max_completion_tokens" not in body + assert body["thinking"] == {"type": "enabled"} + assert body["reasoning_effort"] == "low" + assert body["clear_thinking"] is False + assert request.headers["authorization"] == f"Bearer {KEY}" + assert route.call_count == 1 + if is_async: + health: Final = await litellm.ahealth_check(params, mode="chat") + assert "error" not in health + assert json.loads(route.calls.last.request.content)["thinking"] == {"type": "enabled"} + + +@respx.mock +def test_chat_does_not_send_the_openai_organization_to_bigmodel(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_ORGANIZATION", "org-openai-only") + route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "object": "chat.completion", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + _ = litellm.completion(model=f"bigmodel/{MODEL}", messages=[{"role": "user", "content": "Hello"}]) + assert "openai-organization" not in route.calls.last.request.headers + + +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.asyncio +@respx.mock +async def test_responses_native_body_text_usage_and_continuation(is_async: bool) -> None: + route: Final = respx.post(f"{BASE}/api/v1/responses").respond(json=response_payload()) + params: Final = { + "model": f"bigmodel/{MODEL}", + "input": "Hello", + "api_base": BASE, + "max_output_tokens": 1, + "reasoning": {"effort": "max"}, + } + response: Final = await litellm.aresponses(**params) if is_async else litellm.responses(**params) + assert isinstance(response, ResponsesAPIResponse) + assert response.output_text == "Hello" + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.total_tokens == 30 + assert response.usage.input_tokens_details.cached_tokens == 5 + assert response.usage.output_tokens_details.reasoning_tokens == 3 + assert response.tools == [] + assert response.tool_choice == "auto" + OpenAIResponse.model_validate(response.model_dump()) + first_body: Final = json.loads(route.calls.last.request.content) + assert first_body["model"] == MODEL + assert first_body["max_output_tokens"] == 1 + assert first_body["reasoning"] == {"effort": "max"} + assert first_body.get("store") is not True + assert "messages" not in first_body + assert route.calls.last.request.headers["authorization"] == f"Bearer {KEY}" + followup: Final = {**params, "previous_response_id": response.id, "store": True} + if is_async: + await litellm.aresponses(**followup) + else: + litellm.responses(**followup) + assert json.loads(route.calls.last.request.content)["previous_response_id"] == "resp_bigmodel_test" + assert json.loads(route.calls.last.request.content)["store"] is True + + +@pytest.mark.asyncio +@respx.mock +async def test_responses_decodes_compressed_response_once() -> None: + respx.post(f"{BASE}/api/v1/responses").respond( + content=gzip.compress(json.dumps(response_payload()).encode()), + headers={"content-encoding": "gzip", "content-type": "application/json"}, + ) + response: Final = await litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello") + assert response.output_text == "Hello" + assert response.usage.total_tokens == 30 + + +@pytest.mark.parametrize( + "reasoning_params", + ( + {"thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}, + {"thinking": {"type": "enabled", "budget_tokens": 1024}}, + {"reasoning_effort": "low"}, + ), +) +def test_messages_preserves_native_reasoning_params(reasoning_params: Mapping[str, JsonValue]) -> None: + payload: Final = BigModelMessagesConfig().transform_anthropic_messages_request( + model=MODEL, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 2048, **reasoning_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert {key: payload.get(key) for key in reasoning_params} == reasoning_params + + +@pytest.mark.parametrize( + ("model", "extra", "expected"), + ( + (MODEL, {}, {"thinking": {"type": "enabled"}, "reasoning_effort": "low"}), + (MODEL, {"reasoning_effort": "high"}, {"thinking": {"type": "enabled"}, "reasoning_effort": "high"}), + ("glm-4.7", {}, {"thinking": {"type": "disabled"}, "reasoning_effort": None}), + ), +) +@respx.mock +def test_chat_asks_an_always_thinking_model_for_its_lowest_effort_when_thinking_is_disabled( + model: str, extra: Mapping[str, JsonValue], expected: Mapping[str, JsonValue] +) -> None: + route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + } + ) + litellm.completion( + model=f"bigmodel/{model}", + messages=[{"role": "user", "content": "Hello"}], + thinking={"type": "disabled"}, + **extra, + ) + body: Final = json.loads(route.calls.last.request.content) + assert {key: body.get(key) for key in expected} == expected + + +@pytest.mark.parametrize( + ("model", "extra", "expected"), + ( + (MODEL, {}, {"thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}), + ( + MODEL, + {"output_config": {"effort": "high"}}, + {"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}, + ), + ( + MODEL, + {"reasoning_effort": "high"}, + {"thinking": {"type": "adaptive"}, "reasoning_effort": "high", "output_config": None}, + ), + ("glm-4.7", {}, {"thinking": {"type": "disabled"}, "output_config": None}), + ), +) +def test_messages_asks_an_always_thinking_model_for_its_lowest_effort_when_thinking_is_disabled( + model: str, extra: Mapping[str, JsonValue], expected: Mapping[str, JsonValue] +) -> None: + payload: Final = BigModelMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 64, "thinking": {"type": "disabled"}, **extra}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert {key: payload.get(key) for key in expected} == expected + + +@pytest.mark.parametrize( + "reasoning_params", + ({"thinking": {"type": "enabled", "budget_tokens": 1024}}, {"output_config": {"effort": "low"}}), +) +@pytest.mark.parametrize("temperature", (0, 0.2)) +def test_messages_keeps_the_callers_temperature_next_to_thinking( + reasoning_params: Mapping[str, JsonValue], temperature: float +) -> None: + payload: Final = BigModelMessagesConfig().transform_anthropic_messages_request( + model=MODEL, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 2048, "temperature": temperature, **reasoning_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload.get("temperature") == temperature + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize("web_fetch_requests", (None, 2)) +async def test_messages_forward_native_tools_thinking_and_usage(web_fetch_requests: int | None) -> None: + route: Final = respx.post(f"{BASE}/api/anthropic/v1/messages").respond( + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 20, + "output_tokens": 10, + "server_tool_use": { + "web_search_requests": 1, + **({"web_fetch_requests": web_fetch_requests} if web_fetch_requests is not None else {}), + }, + }, + } + ) + tools: Final = [{"name": "weather", "input_schema": {"type": "object", "properties": {}}}] + params: Final = { + "model": f"bigmodel/{MODEL}", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Hello"}], + "tools": tools, + "thinking": {"type": "enabled", "budget_tokens": 50}, + "api_base": f"{BASE}/api/paas/v4", + "provider_specific_header": { + "custom_llm_provider": "bigmodel", + "native_anthropic_messages_only": True, + "extra_headers": {"x-native-message": "preserved"}, + }, + } + response: Final = await litellm.anthropic.messages.acreate(**params) + assert response["content"][0]["text"] == "Hello" + assert response["usage"]["server_tool_use"] == { + "web_search_requests": 1, + "web_fetch_requests": web_fetch_requests or 0, + } + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == MODEL + assert body["tools"] == tools + assert body["thinking"]["type"] == "enabled" + assert "input" not in body + assert route.calls.last.request.headers["x-api-key"] == KEY + assert route.calls.last.request.headers["x-native-message"] == "preserved" + + +@pytest.mark.parametrize( + "terminal_type,status", + (("response.completed", "completed"), ("response.incomplete", "incomplete")), +) +@pytest.mark.parametrize("include_logprobs", (False, True)) +@pytest.mark.parametrize("content_type", ("text/event-stream", "application/json")) +@pytest.mark.asyncio +@respx.mock +async def test_responses_sse_without_done_preserves_terminal_usage( + terminal_type: str, status: str, include_logprobs: bool, content_type: str +) -> None: + logprobs: Final = [{"token": "Hello", "bytes": [72, 101, 108, 108, 111], "logprob": -0.1, "top_logprobs": []}] + events: Final = ( + { + "type": "response.created", + "sequence_number": 0, + "response": { + **{key: value for key, value in response_payload().items() if key != "parallel_tool_calls"}, + "status": "in_progress", + "output": [], + "usage": None, + }, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + }, + { + "type": "response.content_part.done", + "sequence_number": 2, + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "Hello"}, + }, + { + "type": "response.content_part.done", + "sequence_number": 3, + "item_id": "rs_test", + "output_index": 1, + "content_index": 0, + "part": {"type": "reasoning_text", "text": "Thinking"}, + }, + { + "type": "response.output_text.done", + "sequence_number": 4, + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "text": "Hello", + **({"logprobs": logprobs} if include_logprobs else {}), + }, + {"type": terminal_type, "sequence_number": 5, "response": {**response_payload(), "status": status}}, + ) + sse: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + respx.post(f"{BASE}/api/v1/responses").respond(content=sse, headers={"content-type": content_type}) + stream: Final = await litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello", stream=True) + chunks: Final = tuple([chunk async for chunk in stream]) + assert chunks[-1].type == terminal_type + assert chunks[-1].response.status == status + assert chunks[-1].response.usage.total_tokens == 30 + assert chunks[-1].response.output_text == "Hello" + OpenAIResponse.model_validate(chunks[-1].response.model_dump()) + assert chunks[0].response.parallel_tool_calls is True + assert chunks[-2].logprobs == (logprobs if include_logprobs else []) + assert tuple(chunk.delta for chunk in chunks if chunk.type == "response.output_text.delta") == ("Hello",) + assert tuple(chunk.part.text for chunk in chunks if chunk.type == "response.content_part.done") == ( + "Hello", + "Thinking", + ) + + +@respx.mock +def test_get_response_uses_bigmodel_url_and_normalizes_usage() -> None: + route: Final = respx.post(f"{BASE}/api/v1/responses").respond(json=response_payload()) + response: Final = litellm.responses(model=f"bigmodel/{MODEL}", input="Hello", store=True) + respx.get(f"{BASE}/api/v1/responses/resp_bigmodel_test").respond(json=response_payload()) + retrieved: Final = litellm.get_responses(response_id=response.id) + assert retrieved.output_text == "Hello" + assert retrieved.usage.total_tokens == 30 + assert route.call_count == 1 + + +@pytest.mark.parametrize("provider_api", ("chat", "responses", "messages")) +def test_missing_bigmodel_key_does_not_use_other_provider_credentials( + provider_api: Literal["chat", "responses", "messages"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("BIGMODEL_API_KEY") + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "unrelated-anthropic-key") + call: Final = { + "chat": partial(litellm.completion, model=f"bigmodel/{MODEL}", messages=[{"role": "user", "content": "Hello"}]), + "responses": partial(BigModelResponsesAPIConfig().validate_environment, {}, MODEL, None), + "messages": partial(BigModelMessagesConfig().validate_anthropic_messages_environment, {}, MODEL, (), {}, {}), + }[provider_api] + with pytest.raises(litellm.AuthenticationError, match="BIGMODEL_API_KEY"): + call() + + +def _audio_file() -> io.BytesIO: + audio: Final = io.BytesIO(b"RIFF0000WAVEfmt ") + audio.name = "audio.wav" + return audio + + +@pytest.mark.parametrize( + "call", + ( + partial(litellm.speech, model=f"bigmodel/{MODEL}", input="Hello", voice="tongtong"), + partial(litellm.image_generation, model=f"bigmodel/{MODEL}", prompt="Hello"), + lambda: litellm.transcription(model=f"bigmodel/{MODEL}", file=_audio_file()), + ), + ids=("speech", "image_generation", "transcription"), +) +@respx.mock +def test_openai_routes_bigmodel_does_not_serve_send_no_request( + call: Callable[[], object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + upstream: Final = respx.route().mock(return_value=httpx.Response(200, json={})) + with contextlib.suppress(Exception): + call() + assert not upstream.called + + +@pytest.mark.parametrize( + ("code", "error_class", "status_code"), + ( + (401, litellm.AuthenticationError, 401), + (1113, litellm.RateLimitError, 429), + (1301, litellm.BadRequestError, 400), + (1200, litellm.InternalServerError, 500), + ), +) +@respx.mock +def test_responses_failure_envelope_raises_its_status_instead_of_a_parse_error( + code: int, error_class: type[Exception], status_code: int +) -> None: + respx.post(f"{BASE}/api/v1/responses").respond(json={"code": code, "msg": "Token expired", "success": False}) + with pytest.raises(error_class, match="Token expired") as raised: + litellm.responses(model=f"bigmodel/{MODEL}", input="Hello") + assert getattr(raised.value, "status_code", None) == status_code + + +@pytest.mark.parametrize( + ("code", "error_class"), + ((1261, litellm.ContextWindowExceededError), (1301, litellm.ContentPolicyViolationError)), +) +@pytest.mark.parametrize("route", ("chat", "responses")) +@respx.mock +def test_prompt_too_long_and_unsafe_content_raise_the_errors_fallbacks_catch( + code: int, error_class: type[Exception], route: Literal["chat", "responses"] +) -> None: + if route == "chat": + respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + status_code=400, json={"error": {"code": str(code), "message": "rejected by BigModel"}} + ) + else: + respx.post(f"{BASE}/api/v1/responses").respond( + json={"code": code, "msg": "rejected by BigModel", "success": False} + ) + call: Final = ( + (lambda: litellm.completion(model=f"bigmodel/{MODEL}", messages=[{"role": "user", "content": "Hello"}])) + if route == "chat" + else (lambda: litellm.responses(model=f"bigmodel/{MODEL}", input="Hello")) + ) + with pytest.raises(error_class, match="rejected by BigModel") as raised: + call() + assert raised.value.status_code == 400 + + +@pytest.mark.parametrize( + ("code", "error_class"), + ((1261, litellm.ContextWindowExceededError), (1301, litellm.ContentPolicyViolationError)), +) +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.asyncio +@respx.mock +async def test_streamed_prompt_too_long_and_unsafe_content_raise_the_errors_fallbacks_catch( + code: int, error_class: type[Exception], is_async: bool +) -> None: + respx.post(f"{BASE}/api/v1/responses").respond(json={"code": code, "msg": "rejected by BigModel", "success": False}) + params: Final = {"model": f"bigmodel/{MODEL}", "input": "Hello", "stream": True} + + async def call() -> None: + if is_async: + await litellm.aresponses(**params) + return + litellm.responses(**params) + + with pytest.raises(error_class, match="rejected by BigModel") as raised: + await call() + assert raised.value.status_code == 400 + + +@pytest.mark.parametrize("route", ("responses", "messages")) +@pytest.mark.asyncio +@respx.mock +async def test_a_non_json_answer_surfaces_its_body_text(route: Literal["responses", "messages"]) -> None: + path: Final = "/api/v1/responses" if route == "responses" else "/api/anthropic/v1/messages" + respx.post(f"{BASE}{path}").respond( + content=b"gateway maintenance", headers={"content-type": "text/html"} + ) + call: Final = ( + (lambda: litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello")) + if route == "responses" + else ( + lambda: litellm.anthropic.messages.acreate( + model=f"bigmodel/{MODEL}", max_tokens=16, messages=[{"role": "user", "content": "Hello"}] + ) + ) + ) + with pytest.raises(Exception, match="gateway maintenance") as raised: + await call() + assert "validation error" not in str(raised.value) + + +@pytest.mark.parametrize("stream", (False, True)) +@respx.mock +def test_responses_forward_codex_tool_params_and_echo_them_when_the_answer_omits_them(stream: bool) -> None: + tools: Final = [{"type": "function", "name": "weather", "parameters": {"type": "object", "properties": {}}}] + completed: Final = { + "type": "response.completed", + "sequence_number": 0, + "response": {**response_payload(), "parallel_tool_calls": None}, + } + route: Final = respx.post(f"{BASE}/api/v1/responses").respond( + **( + {"content": f"event: response.completed\ndata: {json.dumps(completed)}\n\n"} + if stream + else {"json": {**response_payload(), "parallel_tool_calls": None}} + ) + ) + answer: Final = litellm.responses( + model=f"bigmodel/{MODEL}", + input="Hello", + tools=tools, + tool_choice="required", + parallel_tool_calls=True, + include=["reasoning.encrypted_content"], + stream=stream, + ) + response: Final = tuple(answer)[-1].response if stream else answer + validated: Final = OpenAIResponse.model_validate(response.model_dump()) + sent: Final = json.loads(route.calls.last.request.content) + assert sent["parallel_tool_calls"] is True + assert sent["include"] == ["reasoning.encrypted_content"] + assert validated.tools[0].name == "weather" + assert validated.tool_choice == "required" + assert validated.parallel_tool_calls is True + + +@pytest.mark.asyncio +@respx.mock +async def test_responses_stream_answered_with_a_failure_envelope_raises_its_status() -> None: + respx.post(f"{BASE}/api/v1/responses").respond(json={"code": 401, "msg": "Token expired", "success": False}) + with pytest.raises(litellm.AuthenticationError, match="Token expired"): + await litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello", stream=True) + with pytest.raises(litellm.AuthenticationError, match="Token expired"): + litellm.responses(model=f"bigmodel/{MODEL}", input="Hello", stream=True) + + +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.asyncio +@respx.mock +async def test_responses_stream_answered_with_one_json_response_replays_it_as_events(is_async: bool) -> None: + respx.post(f"{BASE}/api/v1/responses").respond(json=response_payload()) + request: Final = {"model": f"bigmodel/{MODEL}", "input": "Hello", "stream": True} + chunks: Final = ( + tuple([chunk async for chunk in await litellm.aresponses(**request)]) + if is_async + else tuple(litellm.responses(**request)) + ) + assert chunks[-1].type == "response.completed" + assert chunks[-1].response.output_text == "Hello" + assert chunks[-1].response.usage.total_tokens == 30 + + +class _GatedStream(httpx.AsyncByteStream): + def __init__(self, first: bytes, rest: bytes, gate: asyncio.Event) -> None: + self._first = first + self._rest = rest + self._gate = gate + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._first + await self._gate.wait() + yield self._rest + + +@pytest.mark.asyncio +@respx.mock +async def test_responses_stream_labeled_json_without_a_length_yields_events_before_it_ends() -> None: + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": {**response_payload(), "status": "in_progress", "output": [], "usage": None}, + } + completed: Final = {"type": "response.completed", "sequence_number": 1, "response": response_payload()} + gate: Final = asyncio.Event() + respx.post(f"{BASE}/api/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=_GatedStream( + f"event: response.created\ndata: {json.dumps(created)}\n\n".encode(), + f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(), + gate, + ), + ) + ) + + stream: Final = await asyncio.wait_for( + litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello", stream=True), timeout=10 + ) + first: Final = await asyncio.wait_for(anext(stream), timeout=10) + gate.set() + rest: Final = tuple([chunk async for chunk in stream]) + assert first.type == "response.created" + assert rest[-1].type == "response.completed" + assert rest[-1].response.output_text == "Hello" + + +@pytest.mark.parametrize("event_type", ("response.failed", "error")) +@pytest.mark.asyncio +@respx.mock +async def test_responses_failure_events_raise_instead_of_completing(event_type: str) -> None: + event: Final = ( + {"type": event_type, "sequence_number": 1, "code": "server_error", "message": "Provider failed"} + if event_type == "error" + else { + "type": event_type, + "sequence_number": 1, + "response": { + **response_payload(), + "status": "failed", + "error": {"code": "server_error", "message": "Provider failed"}, + }, + } + ) + respx.post(f"{BASE}/api/v1/responses").respond( + content=f"event: {event_type}\ndata: {json.dumps(event)}\n\n", + headers={"content-type": "text/event-stream"}, + ) + stream: Final = await litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello", stream=True) + with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="Provider failed"): + await anext(stream) + + +@respx.mock +def test_responses_preserve_tools_multimodal_input_and_existing_usage_total() -> None: + payload: Final = response_payload() + route: Final = respx.post(f"{BASE}/api/v1/responses").respond( + json={**payload, "usage": {**payload["usage"], "total_tokens": 31}} + ) + input_items: Final = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this"}, + {"type": "input_image", "image_url": "https://example.com/image.png"}, + ], + } + ] + tools: Final = [{"type": "function", "name": "weather", "parameters": {"type": "object", "properties": {}}}] + response: Final = litellm.responses( + model=f"bigmodel/{MODEL}", + input=input_items, + tools=tools, + tool_choice="auto", + text={"format": {"type": "json_object"}}, + ) + assert response.usage.total_tokens == 31 + assert response.output[0].content[0].text == "Thinking" + body: Final = json.loads(route.calls.last.request.content) + assert body["input"] == input_items + assert body["tools"] == tools + assert body["text"] == {"format": {"type": "json_object"}} + + +@respx.mock +def test_responses_preserve_explicit_tool_configuration() -> None: + tools: Final = [{"type": "function", "name": "weather", "parameters": {"type": "object"}, "strict": False}] + tool_choice: Final = {"type": "function", "name": "weather"} + reasoning: Final = { + "id": "rs_test", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Planning"}], + "content": [{"type": "reasoning_text", "text": "Thinking"}], + } + tool_call: Final = { + "type": "function_call", + "id": "fc_test", + "call_id": "call_test", + "name": "weather", + "arguments": "{}", + "status": "completed", + } + respx.post(f"{BASE}/api/v1/responses").respond( + json={ + **response_payload(), + "output": [reasoning, tool_call], + "tools": tools, + "tool_choice": tool_choice, + "parallel_tool_calls": True, + } + ) + response: Final = litellm.responses(model=f"bigmodel/{MODEL}", input="Hello", tools=tools, tool_choice=tool_choice) + validated: Final = OpenAIResponse.model_validate(response.model_dump()) + assert validated.tools[0].name == "weather" + assert validated.tool_choice.name == "weather" + assert validated.parallel_tool_calls is True + assert validated.output[0].summary[0].text == reasoning["summary"][0]["text"] + assert validated.output[0].content[0].text == reasoning["content"][0]["text"] + assert validated.output[1].model_dump(exclude_none=True) == tool_call + + +@pytest.mark.asyncio +@respx.mock +async def test_router_routes_one_model_configuration_to_three_native_endpoints() -> None: + chat: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "created": 1, + "model": MODEL, + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + } + ) + responses: Final = respx.post(f"{BASE}/api/v1/responses").respond(json=response_payload()) + messages: Final = respx.post(f"{BASE}/api/anthropic/v1/messages").respond( + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 10}, + } + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "glm-native", + "litellm_params": {"model": f"bigmodel/{MODEL}", "api_key": KEY, "api_base": BASE}, + } + ], + num_retries=0, + ) + chat_result: Final = await router.acompletion(model="glm-native", messages=[{"role": "user", "content": "Hello"}]) + responses_result: Final = await router.aresponses(model="glm-native", input="Hello") + messages_result: Final = await router.aanthropic_messages( + model="glm-native", max_tokens=10, messages=[{"role": "user", "content": "Hello"}] + ) + assert ( + chat_result.choices[0].message.content + == responses_result.output_text + == messages_result["content"][0]["text"] + == "Hello" + ) + assert tuple(route.call_count for route in (chat, responses, messages)) == (1, 1, 1) + assert all(json.loads(route.calls.last.request.content)["model"] == MODEL for route in (chat, responses, messages)) + assert ( + chat.calls.last.request.headers["authorization"] + == responses.calls.last.request.headers["authorization"] + == f"Bearer {KEY}" + ) + assert messages.calls.last.request.headers["x-api-key"] == KEY + + +@pytest.mark.asyncio +@respx.mock +async def test_chat_stream_preserves_deltas_and_final_usage() -> None: + chunks: Final = ( + { + "id": "chat_test", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "delta": {"reasoning_content": "Thinking"}, "finish_reason": None}], + }, + { + "id": "chat_test", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello"}, "finish_reason": None}], + }, + { + "id": "chat_test", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + }, + ) + respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + content="".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + stream: Final = await litellm.acompletion( + model=f"bigmodel/{MODEL}", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + stream_options={"include_usage": True}, + ) + output: Final = tuple([chunk async for chunk in stream]) + assert ( + "".join(choice.delta.content or "" for choice in chain.from_iterable(chunk.choices for chunk in output)) + == "Hello" + ) + assert ( + "".join( + getattr(choice.delta, "reasoning_content", "") or "" + for choice in chain.from_iterable(chunk.choices for chunk in output) + ) + == "Thinking" + ) + assert tuple(chunk.usage.total_tokens for chunk in output if hasattr(chunk, "usage") and chunk.usage) == (30,) + + +@pytest.mark.asyncio +@respx.mock +async def test_messages_stream_preserves_native_events() -> None: + events: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 20, + "output_tokens": 0, + "server_tool_use": {"web_search_requests": 0}, + }, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "先思考"}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "test-signature"}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "你好 🌏"}}, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 10}, + }, + {"type": "message_stop"}, + ) + native_events: Final = tuple( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n" for event in events + ) + content: Final = ( + "id: bigmodel-event\nretry: 1000\n" + + native_events[0] + + "event: ping\n\n" + + "".join(native_events[1:]) + + "data: [DONE]\n\n" + ).encode() + upstream: Final = _ByteChunks(tuple(content[index : index + 7] for index in range(0, len(content), 7))) + respx.post(f"{BASE}/api/anthropic/v1/messages").mock( + return_value=httpx.Response( + 200, + stream=upstream, + headers={"content-type": "text/event-stream"}, + ) + ) + stream: Final = await litellm.anthropic.messages.acreate( + model=f"bigmodel/{MODEL}", + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + output: Final = tuple([chunk async for chunk in stream]) + lines: Final = tuple(chain.from_iterable(chunk.decode().splitlines() for chunk in output)) + parsed: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")) + assert tuple(chunk["type"] for chunk in parsed) == tuple(event["type"] for event in events) + assert tuple(event["delta"] for event in parsed if event["type"] == "content_block_delta") == ( + {"type": "thinking_delta", "thinking": "先思考"}, + {"type": "signature_delta", "signature": "test-signature"}, + {"type": "text_delta", "text": "你好 🌏"}, + ) + assert parsed[-2]["usage"]["output_tokens"] == 10 + assert parsed[0]["message"]["usage"]["server_tool_use"] == { + "web_search_requests": 0, + "web_fetch_requests": 0, + } + assert tuple(line for line in lines if line.startswith("data: {"))[1:] == tuple( + f"data: {json.dumps(event, ensure_ascii=False)}" for event in events[1:] + ) + assert "event: ping" in lines + assert "data: [DONE]" in lines + assert b"id: bigmodel-event\n" in output[0] + assert b"retry: 1000\n" in output[0] + assert upstream.is_closed + + +@respx.mock +def test_stored_response_lifecycle_reuses_base_and_unwraps_id() -> None: + respx.post(f"{BASE}/api/v1/responses").respond(json=response_payload()) + response: Final = litellm.responses(model=f"bigmodel/{MODEL}", input="Hello", store=True) + listing: Final = respx.get(f"{BASE}/api/v1/responses/resp_bigmodel_test/input_items").respond( + json={"object": "list", "data": [], "has_more": False} + ) + deletion: Final = respx.delete(f"{BASE}/api/v1/responses/resp_bigmodel_test").respond( + json={"id": "resp_bigmodel_test", "object": "response.deleted", "deleted": True} + ) + assert litellm.list_input_items(response_id=response.id)["data"] == [] + assert litellm.delete_responses(response_id=response.id).deleted is True + assert listing.call_count == deletion.call_count == 1 + assert deletion.calls.last.request.headers["authorization"] == f"Bearer {KEY}" + + +@pytest.mark.parametrize("operation", ("cancel", "compact")) +@respx.mock +def test_unsupported_responses_operations_fail_without_an_upstream_request(operation: str) -> None: + call: Final = ( + partial(litellm.cancel_responses, response_id="resp_test", custom_llm_provider="bigmodel") + if operation == "cancel" + else partial(litellm.compact_responses, model=f"bigmodel/{MODEL}", input="Hello") + ) + with pytest.raises(litellm.UnsupportedParamsError, match="BigModel"): + call() + assert respx.calls.call_count == 0 + + +@pytest.mark.parametrize("model", ("glm-5.3", "glm-5.3-flash", "glm-5.3-flashx")) +def test_bigmodel_model_metadata_and_cached_token_costs(model: str) -> None: + entry: Final = litellm.model_cost[f"bigmodel/{model}"] + info: Final = litellm.get_model_info(model=model, custom_llm_provider="bigmodel") + assert info["litellm_provider"] == "bigmodel" + assert info["max_input_tokens"] == entry["max_input_tokens"] + assert info["max_output_tokens"] == entry["max_output_tokens"] + assert info["supports_reasoning"] + assert info["supports_function_calling"] + assert bool(info.get("supports_vision")) == bool(entry.get("supports_vision")) + prompt_cost, output_cost = litellm.cost_per_token( + model=model, + custom_llm_provider="bigmodel", + prompt_tokens=100, + completion_tokens=20, + cache_read_input_tokens=40, + ) + assert prompt_cost == pytest.approx(60 * entry["input_cost_per_token"] + 40 * entry["cache_read_input_token_cost"]) + assert output_cost == pytest.approx(20 * entry["output_cost_per_token"]) + assert prompt_cost > 0 + assert output_cost > 0 diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index f10f5a67b2f..90d1f6772a0 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5784,3 +5784,16 @@ async def test_non_chatgpt_responses_preserves_caller_headers_after_validation( await client.close() else: client.client.close() + + +@pytest.mark.asyncio +async def test_a_stream_answered_with_one_json_body_is_not_replayed_for_a_provider_that_does_not_opt_in( + httpx_transport: None, +) -> None: + with respx.mock(assert_all_called=True) as router: + router.post("https://api.openai.com/v1/responses").respond(json={"error": {"message": "quota exceeded"}}) + stream: Final = await litellm.aresponses( + model="openai/gpt-5.5", input="Hello", api_key="fake-key", stream=True + ) + with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="before any terminal event"): + _ = tuple([event async for event in stream]) diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 013d3e59125..072ef838e78 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -25,6 +25,26 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreate from litellm.types.utils import LlmProviders +def test_bigmodel_fields_require_private_key_and_share_default_base_with_adapter() -> None: + from litellm.llms.bigmodel.common_utils import BIGMODEL_API_BASE + + app: Final = FastAPI() + app.include_router(router) + response: Final = TestClient(app).get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = TypeAdapter(list[ProviderCreateInfo]).validate_python(response.json()) + provider: Final = next(item for item in providers if item.litellm_provider == "bigmodel") + fields: Final = {field.key: field for field in provider.credential_fields} + assert provider.provider_display_name == "BigModel (Zhipu AI)" + assert fields["api_key"].required is True + assert fields["api_key"].field_type == "password" + assert fields["api_base"].required is False + assert fields["api_base"].default_value == BIGMODEL_API_BASE + assert fields["api_base"].placeholder == BIGMODEL_API_BASE + assert provider.default_model_placeholder is not None + assert provider.default_model_placeholder.startswith("bigmodel/") + + def test_fuse_presets_route_serves_the_shared_catalog_without_authentication() -> None: app: Final = FastAPI() app.include_router(router) diff --git a/ui/litellm-dashboard/public/assets/logos/bigmodel.png b/ui/litellm-dashboard/public/assets/logos/bigmodel.png new file mode 100644 index 0000000000000000000000000000000000000000..f22030f85b8fe85456e8abd027952c44edf16227 GIT binary patch literal 11339 zcmWk!2|QHa7k{&07|S60zEjyLgc)nL5VA#xA}Pwg#MqaFC_7`xmO_by7?dAT_9cn2 zZ`rpoX8F(m``)|n^Ld~1?mhS1dzO3dO|~?@&cwjW0001!v5~$N^^N`i1A|b@d*DJ< z0Du6NW;O=b{mir+hPZ(o3WahfiK1lt+vCON)jP{d*24+!Ns5ktNMwrlt3B<2H5I4X z(=#%aj7XyRX6-57{=2?SbbaxUa86OUJ%4mc^2%AcQ?$FiPZG14yYhG$w?^DQCOL;~ z%Q?>J-dl7C+;=Fjnd*jyQt0U z)iZ5R;?`W(AYS+0ica>+s5eJm!TZSjD*>^TjsX^gdNYwGOsvRqkzLZ$QxKF}*(le#r`=k2z6 z zT^BzO-~+Nx65`Li(s!xf-(ydwapyO}HWS{SIVSGt6tC9(I@2jyjjAGeXB^&7J8+2K z`Z9iYr(& z?^hn!#cn_Ov~&B-Zd}}1Ml;c?@-V$;KR4%WWa&Jt;H+W%lsds|%6F_D{<~53&*{xx zNaU$c$l;UlV~@&%>-RSvRvp_t+BWmryzy+yF=5xBbVJvDP0MZ7CT3f~aq+hAPGmIR zFm%)9$*#Nqo@dZOKn*^kmf(`SFJp&`iz7V9IDV9M5|oQ~e}15eS-qcj2pP%fqN>E7 z04sBA00KG(!{B5ZG9w;FX5A;9u<>nhABykNm~E2CEB}OKIdw^*%+i|1N5_it#I;qj zmfarD)%}yP^~0rq?$R2^@5T?7kHAZKe_U$y(e3r|mH4@?-JQjwp@=!$IeyV%b+Myu ze`$HpduRJ-b$ev-6P*8jy^7M+H((vnhe9u0@H8BO zKLr`H@faXz@)#9ptyLffp%0R{SFH_Abs>sE%oEniO;5WA4^(sAF4F&cEup@0{;ug% z+Y@t*dl4aPHSf|sA558a3|Cv5`&-)M&ho`cZtCU6z>!0rhmyb-r#YRu1q|x}^V2Ws zX=ajH(E-rC{g)K$@nB)*62G~as1(Rk$va%#SS&zRR2f$to60!@^ut+%zA%2{1WEquhEFDNX;NWZ=cK;WG}Zf} zv8hQx@iE3h{J#fqM=2>vJlCCE)rZ}2^aDBdPG@uB%6Iq@VZCVAO5aMcv6{z70pmKo z+fp7ZSuC4Xv^wuXgU=35=OFC-Sg1RnmX@BLa-RA52VG`&mdzR>hj7&9_|Jl?zRV07 z&G&!itL}E3*P&570Y4)&K4qauHe@ez^2`_iwO_|hAD6M26yh3^KSp*CekH7R`<~U_ z%Z<#*o&lG(pzYaYb}cGC3#s|xYU>dTTEA#a`LswuZqZJa@5(_^k;(q)6?`+0HmaB)S93Kmt1eNv~4x}uP*`<1a>XSrX8yg!ZviQT4 z-}S1ys{i5c*0}Xw6=BzYxhwm7Du+A|nWK=H_e9zAS8CF882uAtQ60o&dt1bciGD-8 ziqi1J#MP_~BGZ7)m-K1%w!?!pt9@lL1zjg6Cr&8|C@n2*8IvejwF_l8mAW%k+!6V4 zUHttaO}c{(c`#4@B?`=y^ZhxyFnYmnWuZRX66zhJ*jgeV-U@%DE`{MEf62hh8?>tk zELLpjBSA1Y9Fp+nO%GkiBFu+X$HLA|sM*KDxZSC)Tr>*DknX_lieI{Kbu3;6L+R=s z2*W&pIEjaomtD*+bxoGFRlEfGFUTnL#gBNhbN#@yhRTiYE-$OhZQKJ;bS9jMtzVdU zRn|F$#l^*wpwKQ5YidIw#!#^T>=%8%-m*JEU*C+5v{WQK9Q{I0PC|mdgtCT->WGI% z8yml-Q-fYg0iYkHkWDX-N)PKMupEMd?2^z&#n-0Bl+!kXH>X)at7hpbPeF75&NE2| zPRFFna*-bF)}RYD730%#R{m{0JT%FqK%o8?b0wD8qUKdFRy8C-s24%mT*yY_0Ab;- zyD3SGJ-xCg{Xk4qw5)7(@X1pv4d;zMm=~=Mmjr-7 zipKCIHqe^cGf3usGy#(SGZU%L?#~T_dlhAn)QBEOIvrLro4f08E-)Dx-3O3l{8IWF zfR(*~;DmIuq4ofFLw^#Cxl8a8utoeUTXLGQ$iz@aZLa8*`HxE~8bLO9X}A5K7+<6T z0{|3I3*ITj6x{ZbMqHMJ>I~$BUzUJX{*6yKQlnq?==XK=-cp!h$i)0OO@-@_Oe(BG z&ki`GX#EV5nOj;~#Je;Zfx#H60(~huJoG<=Gs&{vfI03gdiv?s4;Fn%S|KOmP}U$^2l+-`d#VVe4KlClBz)JpEO8>*l!=rvwOf1FV4p zEIa2NkV(Uq!hZBMvmGA&A4*E-*{SA+@l)6=beN#{jNcG`kb^jryGezKXT`@4ANHS$ z=?D{E|ERuoe)3vU*2M8bZU!QO(}-U!{NyuVBm>gGxc_bp&Y0EqG6&yi`TMlKw1HGJ zCr-|0W?NgfGdbb;-wqp$2YfWp;V6fifICfZ-vsQJiPxj*^Le58n|?U{b>#UU->Z+M zFEF6|p3u|R9ErnX=-*Xs4!r$y;s<38AMDGfSx#@x|1jR;9%?!KK|;U7)W#;_)#&dp zHOXKV20>w96VrJw()EU}J@KaynToj7(g74~MB4Mk3iX#P&$A=^+a3QEte+NLbjUR~!6LdNteJ)=jLWg;; z5T)mT|2|6;%FllT6GeGDf=66be~-M>5FTLxI_L1mQ(IH=(uGS)OXVhIPENXux%7gU zXqC8TP{N7s_aUaHI%E?K;_D2egfxYYO&=*HOyDb8TGAUE^bL_U(GXhzdDQdMm$|y1~H`#KTiK+&q=0+`N-VPo^~fIudnxCPGx z79Xh|7d9b_W$QS}WctDA;oS9X@;ahqg}q4B_)Z$3S_`1HnUd> z7WM~X3$b+n?k9xlu5qo{#Sp(sNJ#9bz`5#NI{lpgLZVAX99Z+6mW22n?CtDYsFD^w zm;lM)wbL#fuTA}i;ETO!ugF(EYDsm;3XtdLgIC! z3HuZ-`g|LJw*VomcIByqv$ONx&A$FD0TA|v3hrtoQQq>Bn{9wCSP!azh;xu{N8P|( zg*zhN2726Y4ZC6}sdD}VUlYMiWLuv}fqQ8LqJBdBuq+vxAPA3 z1OL1;k;vd_Z?~XpivrkbnT5;>3mqrV_py_pCLEn0iZx%ecS6fkC&a+kAl@~)pg`z{9R#m~ zr+Hy=gh1fz8DW!;i!~lj!Xb1@?5>UF9H=R5E;G`uz1pR8ulz-URg$0)tK8L#rO+R} zgW6V@+b;u;ybSP4R{sIG|F)`*!=p|0CR5z>Ky>cPOOob=ZQ1z%yRtbw;i>E&Rb@TC z;q-&4;eW=C?Nr{pq&KTX!b&6o4z>R-&00Bm3e~Wjn$I?GSG$7L(9nw3%69X9{qG;YP(ADhvmS&^ifMbKCqmE> zVGOeIY=CDTu!1y#;`FqIL-G&^^a!c#Ht&OAZ0>!;I^Egh$B!c)`+!i0$FM72KSbBy zzZeRk)gUdx5x2&!7w4p+4EP42;2V^!xdJ`>VX5iiRIe#Lp=kl~B$2A$-535zgFi9v5HM$7>9N$3a`N!`T zyFI7z*g|(8jq`P7NLb*;ME#KEyQPWiSg;J*gzlCsAdZh^vxFO-UX@Ue4nhZ0{Dozr zi-D|)JqJVp{FZ9 zYC}K${DK;}KP~G6>ujsv+a!&=M>vgD(E7p>=3RtspFDW6IqsH5d=hO?#NBSIVd& zAnklp7L6Q^z7DurtO~v|1WM426_2B_77MEhgm_ixNV?8th4lXJd#X1? zwMl%ce5hVZ`r6DgD)%C~k{)DX6xp_4fax9@Vmk0tGIMK5Sf8Ld4&I(=+EO0U9DDC6 zU-%O(%niG+&$Z{#cJm&+xN6BZ5Y0yS$&*zBeJ-Fu_F(duD|g8$Rv?xL!8A%~0)ma{ z3x7+J5AOXW+|*k4g`;$*-{1bs4Vq6=F?-M+nU8uFcKnZav-7|8nI?97Z1a^5U1|MB zD4`uETTGW1y)t@jW)Uk}jMBSW21J$!9UxOTPt!{2#n&n1^E$8ShPpJSUL_#L@b!)5 z*Z-}cgf*~wUcW>;8zc8-wbadq)7x_8ynrzZ>C_A4F#hr<51TM{W#DbNLJheYjG%jS z#Y-1^#G`*CuSfa>(dyZtgG=j_+4q%tb=|w!N#Gw;s=mR@c$3Dm78)|MzDBzwzl_WS zmBfUS=7aBjwGrNvil2?+OgiYlk~=j&$;+*@%&CWWxCZvQws22dhoUKslY2|60RGtU zi@}A{7`7OV3S^h>7kV{+g)Sk9Q-M9w#<)CwjZ@A^OvSJ`~~Ip*vlhMiw)$GBYBUl!P1cFS=URV4i|WL|1rw9?w$2a$SamSVs3B+K8sWN zcf-?=T4>NvD;n4LNj^g#_XY)EnO)Edw)0NM6YH}ft)sMmey%tpNf&XP7C>SWIF3|6RrEVv#i9j2)Tv9HFJ-1|&0E>CB?uDG;-~$XH(Ua!`hWhc7xXgTovy z@&XYMu(;AlK3iutf_>u-c+Rn_bQ3u#?;PU|ASxkN{g$v?!&kTzm1MyxNM;&?`78sU z($hBzM2-w?y$XUVWY0lGK0(arMY`QZV~OWmYID)oK?wKd3=ucIZd|>{iCE{E1~QW<#hfbO$n3~mJ6@ochWP%dcLzO)X)eFXuU$733<7&aP`=QhBoa2%rK2# zDHBS*wp<`GkK$K<8rO$xWvXw>%&@3#> z&1H~swclqCF@S>9TH1?q)&+`hyF_`M_)ICfdpB2BRDd8lxq7T9im9@c`RptM8pMba zFWGb#5f!r&a`?F{CE1H2G)zEjeIEiz@HdDs)KyW~PNTCkiS%%k3KJ{4WSIF_zrS|K zmg{6nB5kf2kVNR4sQs`DKX`E5O-q(CB(^W*3m%l`5F<87G^9`k&ztvb4~LqQc6ncT8no_ z&QzHJR3t@VFnOAr^*YFq(gcBG6^BQjyeVJuu-~!y-rL*VW)~P3*wn-^t!{i!LcFpr z6*znYr9j7b@gnCq2Pfx!D%Y+(Z>j5^7x%sb^PRi6=+v=FYEz`Pw8ZzfB6vfb)GO)&=%Z5 zT1c$V;65}_@2zB)GB3namx`T|&~?zy((B;$e%qiJ(97Q+I=brW>J-p2RNTA$c`_t) zdzSRQj%n+9`<2IgE+FF8&!|oc`HvaWFf{b6^??0V$;o(4EL|tb)?TMkc|Fl?i1|6; z31N%DU-a)Y&$iaq)|qr;;|~@e%#DqC8YiO2*|Ke{*t%PQ7gi<7uJY@FN$({do|Ua1 z19^0)1}6GGK{mm=pWHrmrM!I{=>sO7&&}CM<-N05@uMu%*V5 zfxh*$if%`|X%KWV3#&7X$;w*Vkz4=V@i6LCj>uv+g*&_55iSVac_fQfG4?>gpe`7V zz|Ap>&LWin=^`dS`Qab*(b@;-6PX{#u!gecX1OQ_^EcPHUDKh#+lfqV*Lyh=|Sc>q5^o9*kP{iP`^ZDn0w!7IY|6H{}q0a$kh1xGv)?}e68#%c={73y7$wAF}YHDl^5pi`^Q^CQACia*`^NS z6Nc`=5228t%a>!?-9)5bM@Q`1o19>|2w;U$A0ftJv)0tdGmdj!Hr$TS34dD$q_fTyguKjK_6@j zvEtZ>C=VxZ1Y+_lUQ{@O*hL2hg*Kti=TeJ5z?U$&9K{9jgC~16Vc}bClas-Uf*42~ z@U&W_TtrMmLnG{OD>@K_MoUKg7^0uKftZovH^Ry>CLj}c&()qd+@62UHW>Lyw6Z*E z!>UIGUn^2xDw40H!0xC6R0ReQ`>$lK59yrM32T4Xi0|V9#mRY4V1RIusEUe;2iAp& zsmel37lk&$s*2sXtuh+rM|^PZMsRi;aeq|wbnT?*8kHj`j1^T#H;L3lmyky*Xc|O&=WfZY>ZS|dwj=H**<@TualkGoP;DQ%euO-~h&$Xuql0@LE4l&7y zhK_)EVJh10cS(t#DEh;gr5)x29SZ6ElpaJc_afo%_laRKV2{Aly^%0h#iukJi$=3= z2BHTTR7uM>75OL;cPLs6Sk02;Z;(M>c50X{N>M?m6lxbeNto<<{?_3Q2ykibl7eKh zA-j(}53G0ovV*E2!7U~o4}U_YwHN^@p~#l~M6aX-t2jfg`e(-fpieu_9F95bO5Cp@S{m3$nPJBWC-V>nsNvAmBTv$Ml#WMyT`Xij|(A9y_bkH8GJ zXGc4;?4m!BV83HWYd{y!QYD;zV~W~C8SDQobcXCN8ob4T5#UU{;2kt1!|3D@*wPnM z2C~sgJ{$q9eG$)BNe*|9VH*x!RU^sVqoLuVCRWg*Mw@6=F%{wOk%MA@-&qeIb{wM~ zMb>OYv+YS4zM*^hnX~i8AVU>N=;$+={9N`eR5PYYHIN#SXPyjg<|yNcd50rw&GmzF{eTZj>A%i;$Et3na~(+8cd)1P!^eQg9z5RiXTFWQsS{EHH5g1+upa*c{Rr8B+sqlw{eb6DdnXoxh*b@Fma_@>Ob49iK z@Y`|&CfdT%0lFG8rZ3$P-Wn+8%TWT#MO9xSrZufD&ja%FGl6K9NI;9?kckn0|~= zvsjb=fi9P~-|OY8UjqZr;$LZ_^Q5>fEiGqfBRC?BcwC;L>5cr(nKf_*Aji`5mtX$uOb5pAsBSmcz3|xIKzq?y2OSKkuLWVC?oQ z#oJ$DC|_v9i1I5f1a9D~YxF0|bLJ$H^?2G6Wo*|~N77z+T=7h0Yb3>S1l@k0O;^O} z!l+T83HCQG(&WeF`#X2qf_EJ$cYNEPxXY$KIvP}gaEoI1uXc)ameBDkCGBF^NGgeQ zeM1$uiIF5DumBViQ3JYoAYuc8vOK27Ih2lE9pA6_uTD>bD?zeAFNF3~Fk1cSBynZW zHTB-h0Mg~jdVSI^4i#3}a53IcA2}k+P3_2ld}k8ur?g8gy*0EB*p&c2`|s9+!R zi{PeSsE?uACk?$njFK1REfiM1n>8zP&y=%7R+3#G2`B=f653oQ6@8MdC~!l(Oi0PA zC#+A{v4pw522w@*dybRuFT~hGV_R) z8;x5w?usehYku}P(!9BrGN~1w!ixHYu6IAym3}C5DR%9s+9j8!Po?s#V2$yS;R^}% zqutZ_`SX)M1vr7&)YSD^<2Wl^Pr9Qc2Hw(!?-7?a5C=P>>*`?@~a zVA#dhw)bdvhGI({nGh!0$Xl-^0cE}{&|r-G{QLmo$9`VI4_?Rjks^<~F$Eh##hS)K zKz>1#sa@$ewC4EG5Bmx(XAaGlX)C`E=st_XgPCN4NNg`t^Y;qrh`;TIfR;P_6AoWM zA)n^#`ol)BSYHrc9SxfOhB=xC0kTQSdC+>JwKY7g>G5;u&+fOno4TNN3sb$bn8@`~ z`))@NWxIeC|>EP$*zJ^}~b!!-jO{gMd z05_uVDDq7T62Ct7?OVAQdMo#RHB!F55r>t}zuaOZ%)oH+uPY}`|7<#%HW#Y|22>K$ zcn~N~5LijL`%e-Y>djzA@za{iQT$6Y%ytRIFx-UxxI11jGJ>x{G%X*YGKqBiSAFLqP8is8+&S?IHrz2Nh}oTPc&l*G!_l`Py8eOkIE0?nV;8$R90E zp*dp|K6S^ZDckE0LpM{SIT7q{-g!@W)=|KQK?7ld+ zs_W%f;K}a8cWez8z+fCA={_qgvD0e@5$EY^9)0oG&!1Yi;JX=LtRdBrBKfx%N=v?! zw3KjvtCEdbm?wYd^1%1sf~yhzZpT?Pom`8<-S++&zyp?vU-H$Kzo3kI2ZzC68NXbv zI)D@)B`GdLKWRvT@g3xL5P#eoE>*6pTvQw$1Pjb`Qr>?KV<&|~W$*|@oyB%u9!*^y zk@lu`&{vy9P9fo`g3f4UNg%`K!rJdFsV$eGp?V1EcNcjaUDn~mKraac5f+$CEm_Jf z5~31{VtW9%X-gW@@)o{NC?t{n&M(~#h0z3MofJz!fG@=ObPbNy$iQGQ>GC&el_8c$ zwtz{LrIS;k$x0OH!0pI^fI9rIO{D3kD_h*AH}%9sUg>ap16x{W58}$BgTMNjM%B*A zgn}Val-5M27@rc1Jqc{&<7GSIVshOjBb~5t&sV@^k{};B(WMCnCs3VN+h^}FNns-v zM3&KLG&WjBYQAr1C~@j_V=ZwED5>_J??^=tmAMQjXb!tLz7?HMv#EbkTo-w}Lx z8v4Da=3)=B$E1Dvd&3LDLjENqF}|gn>)o|qucxZX8T2v2j4Y}MZtKZm(CVSV{{COT z?gOC)2G2s_8nI!yIfn-aR>?#JcchN?vi8?w(*g2l4iI}|o=7%W535Ny7@$0B4yMxq zSVlM2vyMucJyu(kNe^mDFlyL9QG5I4Qu64tb6wqLU&4`LR&O zoKv0aF^(mY?P9_nV&HvrljSNv-F|T!wX%Mm3wd_!dEoFD>H4GyZOF>lLgMYp%ea3( z_CCygPJnOUgwV1?9#61-i(5ld<~=dn)juks)y1{ZW?+S=jww+=;D|k@ti7ykeEb&M znP9GVGpD3q>41BObs}4;0P*4lWzh(-dVa||B_)Z7 zFxUS2ve#&6@5vq4Xpx+}h4x+os{#IBqrkN!dh1HF1$rfAfS4v4Qp{`C>=ZG4Pby7UYO@*`}p^8 z`hzguk1ax@pHsdLy;@%XZ0Hy!l5>vFh`;hao>h!|{_9#QTGWyS8~wYZiwsyY02r8z zc%LWXG&}^iUeJx>ThCuG7!@eC!DHW{zoRDx3$O_DQARL8RN(_#k}wf52Mgt)BE@P% z_VvB7*-%mUA-ZW4y}}DDR1jpA8}mqIZ0a>+{oHsy?HMu}MT#NP>mwOEi|cl*zvQt5 zlkb{Nq$)`GpZGxcE^-JUR{k?y`Lmd#$fvyn&;1m~4>K~24}EbYR?5i@dV}zHy)`0r zd37-QpBwcM$m!Miw!Xfcw@4b2!1Ba|63^vyLt{j7)E)7@Z;Yae3$=l4f;FbLE}la7o^DW{>aFG z#yWwvc;dBT7NA0e3N!85L0vhSm_Pi&XMxMwQ9XY<*qDRR-wt4`#S~)OZpR<0$1G%f zJTFN1SwEWSuqa@vM%-R53!cgF5)=U|2!qy7k2HJdpZDZjw2jPzm@}qFzSoL&*EzR! zpw)u@O)CaTF-s!9wtx7$zKe!BA$8-H_^$X&Q@7Jf^%BzrVs5eg;ClQmdiBu5SMgaj zRymX-j+Nrv8QDW-N@Nz8LArs(k*_VrkxlUymmrRBnFrx_RZj&K?~(||=OUWZB(F+* zlk~I(rEczbYa={m_*|CF^!u*| { const actual = await vi.importActual("../networking"); + const { readFileSync } = await import("node:fs"); + const { resolve } = await import("node:path"); + const metadata = JSON.parse( + readFileSync( + resolve(__dirname, "../../../../../litellm/proxy/public_endpoints/provider_create_fields.json"), + "utf8", + ), + ) as { provider: string }[]; return { ...actual, getProviderCreateMetadata: vi.fn().mockResolvedValue([ + ...metadata.filter((provider) => provider.provider === "BIGMODEL"), { provider: "OpenAI", provider_display_name: Providers.OpenAI, @@ -157,6 +166,39 @@ const chooseProvider = async (user: ReturnType, provider await user.click(options.find((option) => option.textContent === provider) ?? options[0]); }; +describe("CredentialModal with BigModel metadata", () => { + it("requires a masked API key and shows the default base", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal(); + await chooseProvider(user, Providers.BIGMODEL); + expect(await screen.findByLabelText("API Base")).toHaveValue("https://open.bigmodel.cn"); + expect(screen.getByLabelText("BigModel (Zhipu AI) API Key")).toHaveAttribute("type", "password"); + fill("Credential Name:", "bigmodel-key"); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + expect(await screen.findByText("Required")).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); + }); + + it("saves the key and an explicit base for BigModel", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal(); + await chooseProvider(user, Providers.BIGMODEL); + await screen.findByLabelText("BigModel (Zhipu AI) API Key"); + fill("Credential Name:", "bigmodel-key"); + fill("BigModel (Zhipu AI) API Key", "synthetic-bigmodel-key"); + fill("API Base", "https://bigmodel-proxy.example"); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + expect(onSubmit).toHaveBeenCalledWith( + { + credential_name: "bigmodel-key", + custom_llm_provider: "BIGMODEL", + credential_values: { api_key: "synthetic-bigmodel-key", api_base: "https://bigmodel-proxy.example" }, + }, + [], + ); + }); +}); + describe("CredentialModal with Anthropic workload identity federation", () => { it("deletes stored exchange fields when switching to a static token", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 53e54e2b032..810b010f25e 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -20,6 +20,18 @@ const BUNDLED_MODEL_MAP: Record = JSON.pa ); describe("provider_info_helpers", () => { + it("resolves BigModel metadata, name, logo and model placeholder", () => { + const metadata = PROVIDER_CREATE_FIELDS.find((provider) => provider.provider === "BIGMODEL"); + expect(provider_map.BIGMODEL).toBe(metadata?.litellm_provider); + expect(getProviderLogoAndName("bigmodel")).toEqual({ + displayName: "BigModel (Zhipu AI)", + logo: providerLogoMap[Providers.BIGMODEL], + }); + expect(getProviderLogoAndName("BIGMODEL")).toEqual(getProviderLogoAndName("bigmodel")); + expect(providerLogoMap[Providers.BIGMODEL]).toContain("bigmodel"); + expect(getPlaceholder(Providers.BIGMODEL)).toMatch(/^bigmodel\//); + }); + it("maps Microsoft 365 Copilot to its chat model provider and placeholder", () => { expect(provider_map.MICROSOFT_365_COPILOT).toBe("microsoft_365_copilot"); expect(getPlaceholder(Providers.MICROSOFT_365_COPILOT)).toBe("microsoft_365_copilot/chat"); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 1e1aafde1b5..1ed32c160e6 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -6,6 +6,7 @@ import anthropicLogo from "../../public/assets/logos/anthropic.svg"; import assemblyaiSmallLogo from "../../public/assets/logos/assemblyai_small.png"; import basetenLogo from "../../public/assets/logos/baseten.svg"; import bedrockLogo from "../../public/assets/logos/bedrock.svg"; +import bigmodelLogo from "../../public/assets/logos/bigmodel.png"; import cerebrasLogo from "../../public/assets/logos/cerebras.svg"; import cloudflareLogo from "../../public/assets/logos/cloudflare.svg"; import cohereLogo from "../../public/assets/logos/cohere.svg"; @@ -198,6 +199,7 @@ export enum Providers { xAI = "xAI", XINFERENCE = "Xinference", ZAI = "Z.AI (Zhipu AI)", + BIGMODEL = "BigModel (Zhipu AI)", ZeroGPU = "ZeroGPU", } @@ -328,6 +330,7 @@ export const provider_map: Record = { xAI: "xai", XINFERENCE: "xinference", ZAI: "zai", + BIGMODEL: "bigmodel", ZeroGPU: "zerogpu", }; @@ -337,6 +340,7 @@ export const resolveLitellmProviderSlug = (provider: string): string => provider_map[provider] ?? provider.toLowerCase(); export const providerLogoMap: Partial> = { + [Providers.BIGMODEL]: bigmodelLogo.src, [Providers.A2A_Agent]: a2aAgentLogo.src, [Providers.AI21]: ai21Logo.src, [Providers.AI21_CHAT]: ai21Logo.src, @@ -502,6 +506,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Voyage]: "voyage/", [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", [Providers.ZAI]: "zai/glm-4.5", + [Providers.BIGMODEL]: "bigmodel/glm-5.3", [Providers.ZeroGPU]: "zerogpu/gpt-oss-120b", }; From 3f01341f0a32c74115774cc732cde8b3f70b70c8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:40:02 -0700 Subject: [PATCH 09/58] fix(health)!: openai realtime health check reports healthy with a missing or invalid api key (#41233) * fix(health): report openai realtime deployments with a missing or invalid api key as unhealthy * fix(health): time out a silent openai realtime handshake and keep the first-event parsing in the openai handler * fix(health): treat invalid_model as healthy, read the first event only on OpenAI hosts, map rate limits to 429 * fix(health): respect the global api_base, report a missing key as 401, keep invalid_model healthy only for known models Drop the separate first-event timeout so model_info.health_check_timeout governs the wait, and keep the shared stream error table unchanged by raising AuthenticationError for invalid_api_key in the realtime check * fix(health): keep invalid_model healthy unless the cost map lists the model outside realtime, and load realtime types lazily * fix(health): read the cost map's listed endpoints, not its mode, since registering a realtime deployment overwrites the mode * fix(health): read a realtime gateway's first event and keep Live deployments listing /v1/live healthy Behind a custom api_base the OpenAI realtime health check now waits up to 3 seconds for the first event, so a gateway that passes on OpenAI's invalid_api_key error reports unhealthy while a silent gateway stays healthy. /v1/live counts as a realtime endpoint for the invalid_model rule. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/openai/realtime/handler.py | 56 +++++- litellm/realtime_api/main.py | 65 +++++- tests/unit/realtime_api/test_main.py | 256 +++++++++++++++++++++++- tests/unit/test_lazy_imports.py | 2 +- 4 files changed, 362 insertions(+), 17 deletions(-) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index b7e4cc6868f..5be488f0f6e 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -4,11 +4,17 @@ This file contains the calling OpenAI's `/v1/realtime` endpoint. This requires websockets, and is currently only supported on LiteLLM Proxy. """ +import asyncio import ssl -from typing import TYPE_CHECKING, Final, cast +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Literal, cast + +from pydantic import TypeAdapter from litellm._logging import redact_string, verbose_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.exceptions import AuthenticationError +from litellm.responses.streaming_iterator import map_stream_error_to_exception from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -25,6 +31,54 @@ from ..openai import OpenAIChatCompletion if TYPE_CHECKING: from fastapi import WebSocket + from openai.types.realtime import RealtimeError + from websockets.asyncio.client import ClientConnection + +_SERVER_EVENT_FIELDS: Final = TypeAdapter(Mapping[str, object]) +_CREDENTIALS_ACCEPTED_FOR_A_NON_REALTIME_MODEL: Final = "invalid_model" +_CREDENTIALS_REJECTED: Final = "invalid_api_key" + + +def first_event_error(first_event: str | bytes) -> "RealtimeError | None": + event: Final = _SERVER_EVENT_FIELDS.validate_json(first_event) + if event.get("type") != "error": + return None + from openai.types.realtime import RealtimeErrorEvent + + return RealtimeErrorEvent.model_validate(event).error + + +def _confirm_first_event( + first_event: str | bytes, model: str, sent_credentials: bool, model_listed_outside_realtime: bool +) -> Literal[True]: + error: Final = first_event_error(first_event) + if error is None or ( + error.code == _CREDENTIALS_ACCEPTED_FOR_A_NON_REALTIME_MODEL and not model_listed_outside_realtime + ): + return True + if error.code == _CREDENTIALS_REJECTED or not sent_credentials: + raise AuthenticationError(message=error.message, llm_provider="openai", model=model) + raise map_stream_error_to_exception(error, model=model, custom_llm_provider="openai") + + +async def confirm_session_started( + connection: "ClientConnection", model: str, sent_credentials: bool, model_listed_outside_realtime: bool +) -> Literal[True]: + return _confirm_first_event(await connection.recv(), model, sent_credentials, model_listed_outside_realtime) + + +async def confirm_gateway_session_started( + connection: "ClientConnection", + model: str, + sent_credentials: bool, + model_listed_outside_realtime: bool, + first_event_wait_seconds: float, +) -> Literal[True]: + try: + first_event: Final = await asyncio.wait_for(connection.recv(), first_event_wait_seconds) + except asyncio.TimeoutError: + return True + return _confirm_first_event(first_event, model, sent_credentials, model_listed_outside_realtime) class OpenAIRealtime(OpenAIChatCompletion): diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 7b9068ff95e..e9b534bdd89 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -31,7 +31,12 @@ from litellm.types.realtime import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders -from litellm.utils import ProviderConfigManager, resolve_api_key_secret_reference +from litellm.utils import ( + ModelInfoLookupMiss, + ProviderConfigManager, + lookup_model_info_for_registration, + resolve_api_key_secret_reference, +) from ..litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor @@ -42,8 +47,13 @@ from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed from ..llms.custom_httpx.http_handler import realtime_ssl_for_url +from ..llms.openai.common_utils import is_openai_backed_api_base from ..llms.openai.live.handler import OpenAILiveSessions -from ..llms.openai.realtime.handler import OpenAIRealtime +from ..llms.openai.realtime.handler import ( + OpenAIRealtime, + confirm_gateway_session_started, + confirm_session_started, +) from ..llms.openai.workload_identity import openai_static_api_key, resolve_openai_bearer_token_async from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config @@ -740,6 +750,18 @@ def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) return "/v1/realtime/transcription_sessions" in (model_info.get("supported_endpoints") or ()) +_REALTIME_SERVED_ENDPOINTS: Final = ("/v1/realtime", "/v1/audio/transcriptions", "/v1/live") +_GATEWAY_FIRST_EVENT_WAIT_SECONDS: Final = 3.0 + + +def _cost_map_lists_outside_realtime(model: str) -> bool: + model_info: Final = lookup_model_info_for_registration(model, "openai") + endpoints: Final = ( + () if isinstance(model_info, ModelInfoLookupMiss) else model_info.get("supported_endpoints") or () + ) + return bool(endpoints) and not any(endpoint.startswith(_REALTIME_SERVED_ENDPOINTS) for endpoint in endpoints) + + _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"} @@ -785,6 +807,7 @@ async def realtime_health_check( api_version: str | None = None, realtime_protocol: str | None = None, model_params: Mapping[str, object] | None = None, + gateway_first_event_wait_seconds: float = _GATEWAY_FIRST_EVENT_WAIT_SECONDS, ) -> bool: """ Health check for realtime API - tries connection to the realtime API websocket @@ -800,9 +823,16 @@ async def realtime_health_check( without the OpenAI-Beta header is bridged to, with transcription-only models adding intent=transcription Returns: - bool - True if connection is successful, False otherwise + bool - True once the connection is open, and on OpenAI's own hosts once the first server event is not an + error, since OpenAI accepts the websocket handshake with missing or invalid credentials and only reports + the failure in its first server event. Behind any other OpenAI api_base the first event is read for + `gateway_first_event_wait_seconds` only, and a gateway that stays silent that long counts as healthy. + An `invalid_model` error also counts as healthy unless litellm's + cost map lists only endpoints realtime never serves for the model: OpenAI checks the key first, and Live + and transcription models are served through other endpoints. The listed `mode` is not read, since a + deployment's own `mode: realtime` replaces it when the router registers the deployment Raises: - Exception - if the connection is not successful + Exception - if the connection is not successful or if OpenAI's first server event is any other error """ import websockets @@ -822,10 +852,11 @@ async def realtime_health_check( str | None, api_version or resolved_params.get("api_version") ) url: str | None = None + openai_api_base: Final = resolved_api_base or litellm.api_base or "https://api.openai.com/" health_check_api_key: Final = ( await resolve_openai_bearer_token_async( static_api_key=openai_static_api_key(resolved_api_key), - api_base=resolved_api_base, + api_base=openai_api_base, litellm_params=resolved_params, ) if custom_llm_provider == LlmProviders.OPENAI.value @@ -850,10 +881,7 @@ async def realtime_health_check( query_params=azure_query_params, ) elif custom_llm_provider == "openai": - url = openai_realtime.construct_url( - api_base=resolved_api_base or "https://api.openai.com/", - query_params={"model": model}, - ) + url = openai_realtime.construct_url(api_base=openai_api_base, query_params={"model": model}) elif custom_llm_provider == "xai": url = xai_realtime.construct_url( api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model} @@ -899,8 +927,23 @@ async def realtime_health_check( additional_headers=auth_headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, ssl=realtime_ssl_for_url(url), - ): - return True + ) as connection: + if custom_llm_provider != "openai": + return True + if not is_openai_backed_api_base(url): + return await confirm_gateway_session_started( + connection, + model, + sent_credentials=bool(auth_headers), + model_listed_outside_realtime=_cost_map_lists_outside_realtime(model), + first_event_wait_seconds=gateway_first_event_wait_seconds, + ) + return await confirm_session_started( + connection, + model, + sent_credentials=bool(auth_headers), + model_listed_outside_realtime=_cost_map_lists_outside_realtime(model), + ) _realtime_health_check = realtime_health_check diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 4d89ee497a5..5478b807691 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -4,12 +4,15 @@ import time from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from types import TracebackType -from typing import Final +from typing import Final, Protocol from unittest.mock import AsyncMock, MagicMock, patch import httpx +import openai import pytest import respx +from websockets.exceptions import ConnectionClosedError +from websockets.frames import Close import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -439,18 +442,96 @@ async def test_arealtime_without_live_start_preserves_openai_realtime_dispatch( assert captured.values["api_base"] == "https://realtime-gateway.example/v1" +class _ServerEvents(Protocol): + async def recv(self) -> str: ... + + +class _ScriptedConnection: + def __init__(self, *frames: str) -> None: + self._frames: Final = iter(frames) + + async def recv(self) -> str: + return next(self._frames) + + +class _SilentConnection: + async def recv(self) -> str: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +class _ConnectionClosedBeforeAnyEvent: + async def recv(self) -> str: + raise ConnectionClosedError(Close(3000, "invalid_api_key"), None) + + +_OPENAI_INVALID_API_KEY_EVENT: Final = ( + '{"type": "error", "event_id": "event_1", "error": {"type": "invalid_request_error", "code": "invalid_api_key", ' + '"message": "Incorrect API key provided: sk-proj-****0000. You can find your API key at ' + 'https://platform.openai.com/account/api-keys.", "param": null, "event_id": null}}' +) +_OPENAI_MISSING_AUTH_EVENT: Final = ( + '{"type": "error", "event_id": "event_2", "error": {"type": "invalid_request_error", "code": null, ' + '"message": "Missing bearer or basic authentication in header", "param": null, "event_id": null}}' +) +_OPENAI_SERVER_ERROR_EVENT: Final = ( + '{"type": "error", "event_id": "event_3", ' + '"error": {"type": "server_error", "code": null, "message": "The server had an error", "param": null}}' +) +_OPENAI_RATE_LIMIT_EVENT: Final = ( + '{"type": "error", "event_id": "event_5", "error": {"type": "tokens", "code": "rate_limit_exceeded", ' + '"message": "Rate limit reached for gpt-realtime", "param": null, "event_id": null}}' +) +_OPENAI_INSUFFICIENT_QUOTA_EVENT: Final = ( + '{"type": "error", "event_id": "event_6", "error": {"type": "insufficient_quota", "code": "insufficient_quota", ' + '"message": "You exceeded your current quota", "param": null, "event_id": null}}' +) +_OPENAI_LIVE_MODEL_ON_REALTIME_EVENT: Final = ( + '{"type": "error", "event_id": "event_7", "error": {"type": "invalid_request_error", "code": "invalid_model", ' + '"message": "Model \\"gpt-live-1\\" is not supported in realtime mode.", "param": null, "event_id": null}}' +) +_OPENAI_CHAT_MODEL_ON_REALTIME_EVENT: Final = ( + '{"type": "error", "event_id": "event_8", "error": {"type": "invalid_request_error", "code": "invalid_model", ' + '"message": "Model \\"gpt-5.5\\" is not supported in realtime mode.", "param": null, "event_id": null}}' +) +_OPENAI_MODEL_NOT_FOUND_EVENT: Final = ( + '{"type": "error", "event_id": "event_9", "error": {"type": "invalid_request_error", "code": "model_not_found", ' + '"message": "The model gpt-realtime-2.1x does not exist or you do not have access to it.", "param": null, ' + '"event_id": null}}' +) +_OPENAI_SESSION_CREATED_EVENT: Final = ( + '{"type": "session.created", "event_id": "event_4", "session": {"type": "realtime", "model": "gpt-realtime"}}' +) + + +def _register_as_realtime_deployment(model: str, *supported_endpoints: str) -> None: + litellm.Router( + model_list=[ + { + "model_name": "realtime-deployment", + "litellm_params": {"model": f"openai/{model}", "api_key": "sk-real"}, + "model_info": { + "mode": "realtime", + **({"supported_endpoints": list(supported_endpoints)} if supported_endpoints else {}), + }, + } + ] + ) + + class _CapturingConnect: - def __init__(self) -> None: + def __init__(self, connection: _ServerEvents | None = None) -> None: self.url: str | None = None self.kwargs: dict[str, object] = {} + self._connection: Final = connection or _ScriptedConnection(_OPENAI_SESSION_CREATED_EVENT) def __call__(self, url: str, **kwargs: object) -> "_CapturingConnect": self.url = url self.kwargs = kwargs return self - async def __aenter__(self) -> MagicMock: - return MagicMock() + async def __aenter__(self) -> _ServerEvents: + return self._connection async def __aexit__( self, @@ -652,6 +733,173 @@ async def test_azure_health_check_honors_env_realtime_protocol(monkeypatch): assert connect.url == _AZURE_BETA_HEALTH_URL +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_event", "expected_exception", "expected_status", "expected_message"), + [ + (_OPENAI_INVALID_API_KEY_EVENT, litellm.AuthenticationError, 401, "Incorrect API key provided"), + (_OPENAI_SERVER_ERROR_EVENT, litellm.InternalServerError, 500, "The server had an error"), + (_OPENAI_RATE_LIMIT_EVENT, litellm.RateLimitError, 429, "Rate limit reached"), + (_OPENAI_INSUFFICIENT_QUOTA_EVENT, litellm.RateLimitError, 429, "You exceeded your current quota"), + ], +) +async def test_openai_health_check_reports_the_first_error_event_as_unhealthy( + first_event: str, expected_exception: type[openai.APIStatusError], expected_status: int, expected_message: str +): + connect: Final = _CapturingConnect(_ScriptedConnection(first_event)) + with patch("websockets.connect", connect), pytest.raises(expected_exception) as raised: + await realtime_main.realtime_health_check( + model="gpt-realtime", custom_llm_provider="openai", api_key="sk-wrong" + ) + assert raised.value.status_code == expected_status + assert expected_message in str(raised.value) + assert connect.url == "wss://api.openai.com/v1/realtime?model=gpt-realtime" + + +@pytest.mark.asyncio +async def test_openai_health_check_without_an_api_key_sends_no_auth_header_and_fails_authentication(): + connect: Final = _CapturingConnect(_ScriptedConnection(_OPENAI_MISSING_AUTH_EVENT)) + with patch("websockets.connect", connect), pytest.raises(litellm.AuthenticationError) as raised: + await realtime_main.realtime_health_check(model="gpt-realtime", custom_llm_provider="openai", api_key=None) + assert connect.kwargs["additional_headers"] == {} + assert raised.value.status_code == 401 + assert "Missing bearer or basic authentication" in str(raised.value) + + +@pytest.mark.asyncio +async def test_openai_health_check_is_healthy_once_session_created_arrives(): + connect: Final = _CapturingConnect(_ScriptedConnection(_OPENAI_SESSION_CREATED_EVENT)) + with patch("websockets.connect", connect): + assert await realtime_main.realtime_health_check( + model="gpt-realtime", custom_llm_provider="openai", api_key="sk-real" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-live-1", "gpt-4o-transcribe", "gpt-audio", "gpt-live-unreleased"]) +async def test_openai_health_check_is_healthy_for_a_model_served_outside_realtime_once_the_key_is_accepted( + local_model_cost_map, model: str +): + connect: Final = _CapturingConnect(_ScriptedConnection(_OPENAI_LIVE_MODEL_ON_REALTIME_EVENT)) + with patch("websockets.connect", connect): + assert await realtime_main.realtime_health_check(model=model, custom_llm_provider="openai", api_key="sk-real") + + +@pytest.mark.asyncio +async def test_openai_health_check_is_healthy_for_a_live_deployment_that_lists_the_live_sessions_endpoint( + local_model_cost_map, +): + _register_as_realtime_deployment("gpt-live-1", "/v1/live/sessions") + connect: Final = _CapturingConnect(_ScriptedConnection(_OPENAI_LIVE_MODEL_ON_REALTIME_EVENT)) + with patch("websockets.connect", connect): + assert await realtime_main.realtime_health_check( + model="gpt-live-1", custom_llm_provider="openai", api_key="sk-real" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "first_event", "expected_exception", "expected_message"), + [ + ( + "gpt-5.5", + _OPENAI_CHAT_MODEL_ON_REALTIME_EVENT, + litellm.BadRequestError, + 'Model "gpt-5.5" is not supported in realtime mode', + ), + ( + "gpt-realtime-2.1x", + _OPENAI_MODEL_NOT_FOUND_EVENT, + litellm.NotFoundError, + "The model gpt-realtime-2.1x does not exist", + ), + ], +) +async def test_openai_health_check_is_unhealthy_for_a_model_name_every_session_would_fail_on( + local_model_cost_map, model: str, first_event: str, expected_exception: type[Exception], expected_message: str +): + _register_as_realtime_deployment(model) + connect: Final = _CapturingConnect(_ScriptedConnection(first_event)) + with patch("websockets.connect", connect), pytest.raises(expected_exception) as raised: + await realtime_main.realtime_health_check(model=model, custom_llm_provider="openai", api_key="sk-real") + assert expected_message in str(raised.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_api_base", "global_api_base"), + [("https://realtime-gateway.example/v1", None), (None, "https://realtime-gateway.example/v1")], +) +async def test_openai_health_check_through_a_compatible_gateway_is_healthy_when_it_sends_no_first_event( + monkeypatch: pytest.MonkeyPatch, deployment_api_base: str | None, global_api_base: str | None +): + monkeypatch.setattr(litellm, "api_base", global_api_base) + connect: Final = _CapturingConnect(_SilentConnection()) + with patch("websockets.connect", connect): + assert await asyncio.wait_for( + realtime_main.realtime_health_check( + model="gpt-realtime", + custom_llm_provider="openai", + api_key="sk-real", + api_base=deployment_api_base, + gateway_first_event_wait_seconds=0.01, + ), + timeout=5, + ) + assert connect.url == "wss://realtime-gateway.example/v1/realtime?model=gpt-realtime" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("connection", "expected_exception"), + [ + (_ScriptedConnection(_OPENAI_INVALID_API_KEY_EVENT), litellm.AuthenticationError), + (_ConnectionClosedBeforeAnyEvent(), ConnectionClosedError), + ], +) +async def test_openai_health_check_through_a_compatible_gateway_is_unhealthy_when_it_rejects_the_session( + connection: _ServerEvents, expected_exception: type[Exception] +): + connect: Final = _CapturingConnect(connection) + with patch("websockets.connect", connect), pytest.raises(expected_exception): + await realtime_main.realtime_health_check( + model="gpt-realtime", + custom_llm_provider="openai", + api_key="sk-wrong", + api_base="https://realtime-gateway.example/v1", + ) + assert connect.url == "wss://realtime-gateway.example/v1/realtime?model=gpt-realtime" + + +@pytest.mark.asyncio +async def test_openai_health_check_never_reports_healthy_before_the_first_event_arrives(): + connect: Final = _CapturingConnect(_SilentConnection()) + with patch("websockets.connect", connect), pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + realtime_main.realtime_health_check(model="gpt-realtime", custom_llm_provider="openai", api_key="sk-real"), + timeout=0.05, + ) + + +@pytest.mark.asyncio +async def test_openai_health_check_is_unhealthy_when_the_socket_closes_before_any_event(): + connect: Final = _CapturingConnect(_ConnectionClosedBeforeAnyEvent()) + with patch("websockets.connect", connect), pytest.raises(ConnectionClosedError): + await realtime_main.realtime_health_check( + model="gpt-realtime", custom_llm_provider="openai", api_key="sk-wrong" + ) + + +@pytest.mark.asyncio +async def test_xai_health_check_trusts_the_handshake_without_reading_a_first_event(): + connect: Final = _CapturingConnect(_ScriptedConnection()) + with patch("websockets.connect", connect): + assert await realtime_main.realtime_health_check( + model="grok-voice-latest", custom_llm_provider="xai", api_key="sk-wrong" + ) + assert connect.url == "wss://api.x.ai/v1/realtime?model=grok-voice-latest" + + class _ConnectThatStopsAfterCapturingTheUrl: url: str | None = None diff --git a/tests/unit/test_lazy_imports.py b/tests/unit/test_lazy_imports.py index 0416020a9b5..43aee4616d2 100644 --- a/tests/unit/test_lazy_imports.py +++ b/tests/unit/test_lazy_imports.py @@ -47,7 +47,7 @@ def test_import_litellm_does_not_load_fastapi_or_bpe_table(): [ sys.executable, "-c", - "import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding') if m in sys.modules))", + "import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding','openai.types.realtime') if m in sys.modules))", ], check=True, capture_output=True, From dff474e450b9f8cf94d3fad508ab6c789528a63e Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:56:02 -0700 Subject: [PATCH 10/58] fix(responses): replay the original Responses item order, phase, and boundaries on the next chat turn (#48912) * fix(responses): replay the original Responses item order, phase, and boundaries on the next chat turn * fix(responses): keep tool-call-only stream content empty and drop unreachable replay guards * fix(responses): reject an assistant phase OpenAI does not define unless drop_params is set * fix(responses): keep other providers' reasoning and stream output unchanged and fall back on records that do not fit * fix(responses): label a structured output turn with its returned message's phase and keep lone surrogates from failing the checksum * fix(responses): leave a turn without a phase when its content message cannot be read * test(responses): add black box tests for chat turn history replay * test(responses): expect the replay record on bridged chat replies * test(responses): expect the replay record in bridged chat translation cases * fix(streaming): keep the served service tier on the finish chunk the stream adds * test(responses): read the expected chat message without JSON validation so the replay matcher stays --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../handler.py | 2 + .../history.py | 432 ++++++++++++ .../transformation.py | 185 +++++- .../litellm_core_utils/streaming_handler.py | 10 + litellm/tool_loop.py | 4 +- .../integration/_support/responses_replay.py | 214 ++++++ ...st_responses_bridge_history_replay_wire.py | 233 +++++++ .../test_responses_bridge_incomplete.py | 15 +- ..._tool_loop_responses_history_replay_sdk.py | 37 ++ .../bases/openai_responses.py | 13 +- .../translation/response_scenarios.py | 12 +- ...itellm_responses_transformation_history.py | 627 ++++++++++++++++++ ...responses_transformation_transformation.py | 45 +- .../test_streaming_handler.py | 32 + tests/unit/test_tool_loop.py | 66 ++ 15 files changed, 1882 insertions(+), 45 deletions(-) create mode 100644 litellm/completion_extras/litellm_responses_transformation/history.py create mode 100644 tests/integration/_support/responses_replay.py create mode 100644 tests/integration/providers/test_responses_bridge_history_replay_wire.py create mode 100644 tests/integration/sdk/test_tool_loop_responses_history_replay_sdk.py create mode 100644 tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_history.py diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 3f7017eccc2..76d17ea9693 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -260,6 +260,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + replays_openai_history=custom_llm_provider == "openai", ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -363,6 +364,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + replays_openai_history=custom_llm_provider == "openai", ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/completion_extras/litellm_responses_transformation/history.py b/litellm/completion_extras/litellm_responses_transformation/history.py new file mode 100644 index 00000000000..4435de64cde --- /dev/null +++ b/litellm/completion_extras/litellm_responses_transformation/history.py @@ -0,0 +1,432 @@ +import hashlib +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate, chain, pairwise +from types import MappingProxyType +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +REPLAY_FIELD: Final = "openai_responses_replay" +PHASE_FIELD: Final = "phase" + +Phase: TypeAlias = Literal["commentary", "final_answer"] + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _OutputText(_Frozen): + type: Literal["output_text"] + text: str + + +class _OutputMessage(_Frozen): + type: Literal["message"] + role: Literal["assistant"] + content: tuple[_OutputText, ...] + phase: Phase | None = None + + +class _OutputReasoning(_Frozen): + type: Literal["reasoning"] + id: str + + +class _OutputToolCall(_Frozen): + type: Literal["function_call", "custom_tool_call"] + call_id: str + id: str | None = None + + +class _OutputToolSearch(_Frozen): + type: Literal["tool_search_call", "tool_search_output"] + + +_OutputItem: TypeAlias = Annotated[ + _OutputMessage | _OutputReasoning | _OutputToolCall | _OutputToolSearch, Field(discriminator="type") +] +_OUTPUT_ITEM: Final[TypeAdapter[_OutputItem]] = TypeAdapter(_OutputItem) + + +class _AnyMessage(_Frozen): + type: Literal["message"] + + +class _MessageEntry(_Frozen): + type: Literal["message"] + parts: tuple[Annotated[int, Field(ge=0)], ...] + phase: Phase | None = None + + +class _IndexedEntry(_Frozen): + type: Literal["reasoning", "function_call", "custom_tool_call", "tool_search_call", "tool_search_output"] + index: Annotated[int, Field(ge=0)] + + +_Entry: TypeAlias = Annotated[_MessageEntry | _IndexedEntry, Field(discriminator="type")] + + +class _Replay(_Frozen): + version: Literal[1] + items: tuple[_Entry, ...] + checksum: str + + +class _Namespaced(_Frozen): + namespace: str | None = None + + +class _ToolFunction(_Frozen): + name: str = "" + arguments: str = "" + provider_specific_fields: _Namespaced | None = None + + +class _ToolCall(_Frozen): + id: str + function: _ToolFunction + provider_specific_fields: _Namespaced | None = None + + +class _SummaryText(_Frozen): + type: str = "summary_text" + text: str = "" + + +class _Reasoning(_Frozen): + id: str | None = None + encrypted_content: str | None = None + summary: tuple[_SummaryText, ...] | None = None + + +class _ProviderFields(_Frozen): + phase: Phase | None = None + tool_search_items: tuple[Mapping[str, object], ...] | None = None + openai_responses_replay: object = None + + +class _ChatMessage(_Frozen): + role: str | None = None + content: str | None = None + tool_calls: tuple[_ToolCall, ...] | None = None + reasoning_items: tuple[_Reasoning, ...] | None = None + phase: Phase | None = None + provider_specific_fields: _ProviderFields | None = None + + +class _PhaseFields(_Frozen): + phase: Phase | None = None + + +class _RequestedPhaseFields(_Frozen): + phase: object = None + + +class _RequestedPhase(_Frozen): + phase: object = None + provider_specific_fields: _RequestedPhaseFields | None = None + + +@dataclass(frozen=True, slots=True) +class UnknownPhase: + value: object + + +_PHASE: Final[TypeAdapter[Phase]] = TypeAdapter(Phase) + + +class _RetainedFields(_PhaseFields): + openai_responses_replay: object = None + + +@dataclass(frozen=True, slots=True) +class ReplayedTurn: + items: tuple[Mapping[str, object], ...] + custom_call_ids: frozenset[str] + + +def _as_json(value: object) -> object: + return value.model_dump() if isinstance(value, BaseModel) else value + + +def _parsed_output_item(item: object) -> _OutputItem | None: + try: + return _OUTPUT_ITEM.validate_python(_as_json(item)) + except ValidationError: + return None + + +def _is_message(item: object) -> bool: + try: + _AnyMessage.model_validate(_as_json(item)) + except ValidationError: + return False + return True + + +def _chat_message(message: object) -> _ChatMessage | None: + try: + return _ChatMessage.model_validate(_as_json(message)) + except ValidationError: + return None + + +def requested_phase(message: object) -> Phase | UnknownPhase | None: + try: + requested: Final = _RequestedPhase.model_validate(_as_json(message)) + except ValidationError: + return None + fields: Final = requested.provider_specific_fields + value: Final = requested.phase or (fields.phase if fields else None) + if value is None: + return None + try: + return _PHASE.validate_python(value) + except ValidationError: + return UnknownPhase(value) + + +def reasoning_input_item( + item_id: str | None, summary: Sequence[object] | None, encrypted_content: str | None +) -> Mapping[str, object]: + return { + "type": "reasoning", + **({"id": item_id} if item_id else {}), + "summary": list(summary or ()), + **({"encrypted_content": encrypted_content} if encrypted_content else {}), + } + + +def retained_replay_fields(provider_specific_fields: object) -> Mapping[str, object] | None: + try: + retained: Final = _RetainedFields.model_validate(provider_specific_fields) + except ValidationError: + return None + return retained.model_dump(exclude_none=True) or None + + +def output_message_text(output_items: Sequence[object], *, first_message_only: bool) -> str: + parsed: Final = map(_parsed_output_item, output_items) + messages: Final = tuple(item for item in parsed if isinstance(item, _OutputMessage)) + return "".join(_texts(messages[:1] if first_message_only else messages)) + + +def _tool_search_entries(message: _ChatMessage) -> tuple[Mapping[str, object], ...]: + fields: Final = message.provider_specific_fields + items: Final = (fields.tool_search_items if fields else None) or () + return tuple(item for item in items if item.get("type") in ("tool_search_call", "tool_search_output")) + + +def _checksum(message: _ChatMessage, entries: Sequence[_MessageEntry | _IndexedEntry]) -> str: + fields: Final = message.provider_specific_fields + bound: Final = { + "content": message.content or "", + "tool_calls": [[call.id, call.function.name, call.function.arguments] for call in message.tool_calls or ()], + "reasoning": [ + [item.id, item.encrypted_content, [part.text for part in item.summary or ()]] + for item in message.reasoning_items or () + ], + "tool_search_items": list((fields.tool_search_items if fields else None) or ()), + "phase": message.phase or (fields.phase if fields else None), + "items": [entry.model_dump(mode="json", exclude_none=True) for entry in entries], + } + encoded: Final = json.dumps(bound, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + return hashlib.sha256(encoded.encode("utf-8", "surrogatepass")).hexdigest() + + +def _kind(item: _OutputItem) -> str: + match item: + case _OutputToolCall(): + return "tool_call" + case _OutputToolSearch(): + return "tool_search" + case _OutputMessage() | _OutputReasoning(): + return item.type + + +def _entry(items: Sequence[_OutputItem], position: int) -> _MessageEntry | _IndexedEntry: + item: Final = items[position] + if isinstance(item, _OutputMessage): + return _MessageEntry(type="message", parts=tuple(len(part.text) for part in item.content), phase=item.phase) + ordinal: Final = sum(1 for earlier in items[:position] if _kind(earlier) == _kind(item)) + return _IndexedEntry(type=item.type, index=ordinal) + + +def _texts(items: Sequence[_OutputItem]) -> Iterator[str]: + for item in items: + if isinstance(item, _OutputMessage): + yield from (part.text for part in item.content) + + +def _matches(items: Sequence[_OutputItem], message: _ChatMessage) -> bool: + reasoning_ids: Final = tuple(item.id for item in items if isinstance(item, _OutputReasoning)) + calls: Final = tuple(item for item in items if isinstance(item, _OutputToolCall)) + tool_calls: Final = message.tool_calls or () + searches: Final = sum(1 for item in items if isinstance(item, _OutputToolSearch)) + return ( + "".join(_texts(items)) == (message.content or "") + and reasoning_ids == tuple(item.id for item in message.reasoning_items or ()) + and len(calls) == len(tool_calls) + and all(chat.id in (call.call_id, call.id) for call, chat in zip(calls, tool_calls)) + and searches == len(_tool_search_entries(message)) + ) + + +def _content_phase( + output_items: Sequence[object], parsed: Sequence[_OutputItem | None], *, first_message_only: bool +) -> Phase | None: + messages: Final = tuple(item for raw, item in zip(output_items, parsed) if _is_message(raw)) + if not messages: + return None + content_message: Final = messages[0] if first_message_only else messages[-1] + return content_message.phase if isinstance(content_message, _OutputMessage) else None + + +def _with_phase(message: _ChatMessage, phase: Phase | None) -> _ChatMessage: + if phase is None: + return message + fields: Final = (message.provider_specific_fields or _ProviderFields()).model_copy(update={PHASE_FIELD: phase}) + return message.model_copy(update={"provider_specific_fields": fields}) + + +def replay_fields( + output_items: Sequence[object], message: object, *, completed: bool, first_message_only: bool +) -> Mapping[str, object]: + parsed: Final = tuple(map(_parsed_output_item, output_items)) + phase: Final = _content_phase(output_items, parsed, first_message_only=first_message_only) + phase_fields: Final[Mapping[str, object]] = MappingProxyType({PHASE_FIELD: phase} if phase else {}) + items: Final = tuple(item for item in parsed if item is not None) + parsed_message: Final = _chat_message(message) + chat: Final = _with_phase(parsed_message, phase) if parsed_message is not None else None + if not completed or not items or len(items) != len(parsed) or chat is None or not _matches(items, chat): + return phase_fields + entries: Final = tuple(_entry(items, position) for position in range(len(items))) + replay: Final = _Replay(version=1, items=entries, checksum=_checksum(chat, entries)) + return MappingProxyType({**phase_fields, REPLAY_FIELD: replay.model_dump(mode="json", exclude_none=True)}) + + +def _stored_replay(message: _ChatMessage) -> _Replay | None: + fields: Final = message.provider_specific_fields + if fields is None or fields.openai_responses_replay is None: + return None + try: + return _Replay.model_validate(fields.openai_responses_replay) + except ValidationError: + return None + + +def _message_texts(entries: Sequence[_MessageEntry], content: str) -> tuple[tuple[str, ...], ...]: + lengths: Final = tuple(chain.from_iterable(entry.parts for entry in entries)) + bounds: Final = tuple(accumulate(lengths, initial=0)) + parts: Final = tuple(content[start:end] for start, end in pairwise(bounds)) + groups: Final = tuple(accumulate((len(entry.parts) for entry in entries), initial=0)) + return tuple(parts[start:end] for start, end in pairwise(groups)) + + +def _message_item(texts: Sequence[str], phase: Phase | None) -> Mapping[str, object]: + return { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text} for text in texts], + **({PHASE_FIELD: phase} if phase else {}), + } + + +def _function_call_item(call: _ToolCall) -> Mapping[str, object]: + namespaced: Final = call.provider_specific_fields or call.function.provider_specific_fields + namespace: Final = namespaced.namespace if namespaced else None + return { + "type": "function_call", + "call_id": call.id, + "name": call.function.name, + "arguments": call.function.arguments, + **({"namespace": namespace} if namespace else {}), + } + + +def _custom_tool_call_item(call: _ToolCall) -> Mapping[str, object]: + return { + "type": "custom_tool_call", + "call_id": call.id, + "name": call.function.name, + "input": call.function.arguments, + } + + +def _reasoning_item(item: _Reasoning) -> Mapping[str, object]: + summary: Final = tuple(part.model_dump() for part in item.summary or ()) + return reasoning_input_item(item.id, summary, item.encrypted_content) + + +@dataclass(frozen=True, slots=True) +class _Resolved: + message: _ChatMessage + texts: Mapping[int, tuple[str, ...]] + tool_searches: tuple[Mapping[str, object], ...] + + +def _input_item(resolved: _Resolved, position: int, entry: _MessageEntry | _IndexedEntry) -> Mapping[str, object]: + if isinstance(entry, _MessageEntry): + return _message_item(resolved.texts[position], entry.phase) + tool_calls: Final = resolved.message.tool_calls or () + match entry.type: + case "reasoning": + return _reasoning_item((resolved.message.reasoning_items or ())[entry.index]) + case "function_call": + return _function_call_item(tool_calls[entry.index]) + case "custom_tool_call": + return _custom_tool_call_item(tool_calls[entry.index]) + case "tool_search_call" | "tool_search_output": + return dict(resolved.tool_searches[entry.index]) + + +def _fits(message: _ChatMessage, entries: Sequence[_MessageEntry | _IndexedEntry]) -> bool: + tool_calls: Final = len(message.tool_calls or ()) + tool_searches: Final = len(_tool_search_entries(message)) + sizes: Final = MappingProxyType( + { + "reasoning": len(message.reasoning_items or ()), + "function_call": tool_calls, + "custom_tool_call": tool_calls, + "tool_search_call": tool_searches, + "tool_search_output": tool_searches, + } + ) + parts: Final = chain.from_iterable(entry.parts for entry in entries if isinstance(entry, _MessageEntry)) + indexed: Final = (entry for entry in entries if isinstance(entry, _IndexedEntry)) + return sum(parts) == len(message.content or "") and all(entry.index < sizes[entry.type] for entry in indexed) + + +def _resolve( + message: _ChatMessage, entries: Sequence[_MessageEntry | _IndexedEntry], replay_tool_search: bool +) -> ReplayedTurn | None: + indexed: Final = tuple(entry for entry in entries if isinstance(entry, _IndexedEntry)) + if not replay_tool_search and any(entry.type in ("tool_search_call", "tool_search_output") for entry in indexed): + return None + messages: Final = tuple(entry for entry in entries if isinstance(entry, _MessageEntry)) + calls: Final = tuple(entry for entry in indexed if entry.type in ("function_call", "custom_tool_call")) + texts: Final = _message_texts(messages, message.content or "") + positions: Final = tuple(position for position, entry in enumerate(entries) if isinstance(entry, _MessageEntry)) + resolved: Final = _Resolved( + message=message, + texts=MappingProxyType(dict(zip(positions, texts, strict=True))), + tool_searches=_tool_search_entries(message), + ) + tool_calls: Final = message.tool_calls or () + return ReplayedTurn( + items=tuple(_input_item(resolved, position, entry) for position, entry in enumerate(entries)), + custom_call_ids=frozenset(tool_calls[entry.index].id for entry in calls if entry.type == "custom_tool_call"), + ) + + +def replay_items(message: object, *, replay_tool_search: bool) -> ReplayedTurn | None: + chat: Final = _chat_message(message) + stored: Final = _stored_replay(chat) if chat is not None and chat.role == "assistant" else None + if chat is None or stored is None or stored.checksum != _checksum(chat, stored.items): + return None + if not _fits(chat, stored.items): + return None + return _resolve(chat, stored.items, replay_tool_search) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index ee3639ac387..6173794a210 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -24,6 +24,15 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm import ModelResponse from litellm._logging import verbose_logger +from litellm.completion_extras.litellm_responses_transformation.history import ( + PHASE_FIELD, + UnknownPhase, + output_message_text, + reasoning_input_item, + replay_fields, + replay_items, + requested_phase, +) from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -490,7 +499,7 @@ def _fold_choices_into_one( choices: Sequence["Choices"], tool_calls: Sequence[Mapping[str, object]], trailing_reasoning_content: str | None, - trailing_reasoning_item: _BuiltReasoningItem | None, + trailing_reasoning_items: Sequence[_BuiltReasoningItem], tool_search_items: Sequence[Mapping[str, object]], structured_output: bool = False, ) -> "Choices": @@ -509,7 +518,7 @@ def _fold_choices_into_one( reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) reasoning_items: Final = ( *chain.from_iterable(_message_reasoning_items(m) for m in messages), - *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + *trailing_reasoning_items, ) return Choices( message=Message( @@ -526,6 +535,61 @@ def _fold_choices_into_one( ) +def _with_replay_fields( + choice: "Choices", output_items: Sequence[object], *, completed: bool, structured_output: bool +) -> "Choices": + fields: Final = { + **(choice.message.provider_specific_fields or {}), + **replay_fields(output_items, choice.message, completed=completed, first_message_only=structured_output), + } + message: Final = choice.message.model_copy(update={"provider_specific_fields": fields or None}) + return choice.model_copy(update={"message": message}) + + +_STREAMED_OUTPUT_ITEMS: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) + + +def completed_stream_history( + output_payload: object, + tool_search_fields: Mapping[str, object] | None, + *, + completed: bool, + structured_output: bool, +) -> dict[str, object] | None: # mutable-ok: Delta.provider_specific_fields is typed as a dict + items: Final = _STREAMED_OUTPUT_ITEMS.validate_python(output_payload) + calls: Final = tuple(item for item in items if item.get("type") in ("function_call", "custom_tool_call")) + streamed_message: Final = { + "role": "assistant", + "content": output_message_text(items, first_message_only=structured_output), + "tool_calls": [tool_call_dict_from_output_item(call, index) for index, call in enumerate(calls)], + "reasoning_items": list(_reasoning_items_from_output_items(items)), + "provider_specific_fields": tool_search_fields, + } + fields: Final = { + **(tool_search_fields or {}), + **replay_fields(items, streamed_message, completed=completed, first_message_only=structured_output), + } + return fields or None + + +def _assistant_phase_field(msg: object, replay_openai_history: bool, drop_params: bool) -> Mapping[str, object]: + match requested_phase(msg) if replay_openai_history else None: + case None: + return {} + case UnknownPhase(value=value) if drop_params: + verbose_logger.debug("Chat provider: dropping assistant phase=%r under drop_params", value) + return {} + case UnknownPhase(value=value): + raise litellm.UnsupportedParamsError( + message=f"phase must be 'commentary' or 'final_answer', got {value!r}. " + "To drop it, set `litellm.drop_params=True` or pass `drop_params=True`." + ) + case phase: + return {PHASE_FIELD: phase} + + def _responses_event_type(chunk: Mapping[str, object]) -> object: event_type: Final = chunk.get("type") return event_type.value if isinstance(event_type, ResponsesAPIStreamEvents) else event_type @@ -667,15 +731,7 @@ def _reasoning_item_to_response_input( An item without an id is sent without one: the Responses API accepts that and verifies the encrypted content on its own, while it rejects any id it did not mint. """ - item_id: Final = r_item.get("id") - encrypted_content: Final = r_item.get("encrypted_content") - return { - "type": "reasoning", - **({"id": item_id} if item_id else {}), - # summary is always required by the Responses API, even when empty - "summary": r_item.get("summary") or [], - **({"encrypted_content": encrypted_content} if encrypted_content else {}), - } + return dict(reasoning_input_item(r_item.get("id"), r_item.get("summary"), r_item.get("encrypted_content"))) def _requested_defer_loading(tool: Mapping[str, object], function_tool: Mapping[str, object]) -> object: @@ -804,12 +860,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): keep_prompt_cache_breakpoints: bool = False, replay_tool_search: bool = False, include_namespace: bool = False, + replay_openai_history: bool = False, ) -> tuple[list[object], str | None]: converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input( messages, drop_params=drop_params, replay_tool_search=replay_tool_search, include_namespace=include_namespace, + replay_openai_history=replay_openai_history, ) return ( converted_input_items @@ -825,10 +883,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): drop_params: bool = False, replay_tool_search: bool = False, include_namespace: bool = False, + replay_openai_history: bool = False, ) -> tuple[list[object], str | None]: input_items: Final[list[object]] = [] instructions: str | None = None - custom_tool_call_ids: Final = frozenset( + replayed_turns: Final = tuple( + replay_items(msg, replay_tool_search=replay_tool_search) if replay_openai_history else None + for msg in messages + ) + replayed_custom_call_ids: Final = frozenset( + chain.from_iterable(turn.custom_call_ids for turn in replayed_turns if turn is not None) + ) + custom_tool_call_ids: Final = replayed_custom_call_ids | frozenset( tool_call["id"] for msg in messages if msg.get("role") == "assistant" and isinstance(msg.get("tool_calls"), list) @@ -849,7 +915,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): tool_calls = msg.get("tool_calls") tool_call_id = msg.get("tool_call_id") - if role == "system": + if (replayed_turn := replayed_turns[index]) is not None: + input_items.extend(replayed_turn.items) + elif role == "system": # Extract system message as instructions if isinstance(content, str) and index < leading_system_count: if instructions: @@ -915,6 +983,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "content": self._convert_content_to_responses_format( content, "assistant", drop_params=drop_params ), + **_assistant_phase_field(msg, replay_openai_history, drop_params), } ) for tool_call in tool_calls: @@ -955,6 +1024,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "content": self._convert_content_to_responses_format( content, cast(str, role), drop_params=drop_params ), + **_assistant_phase_field(msg, replay_openai_history and role == "assistant", drop_params), } ) elif role == "assistant": @@ -1081,6 +1151,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint, replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast include_namespace=True, + replay_openai_history=provider == "openai", ) # OpenAI's Responses API rejects an empty input. For a system-only # request, carry the system message as a system-role input item instead @@ -1165,6 +1236,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items: Sequence[object], handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None, structured_output: bool = False, + replay_openai_history: bool = False, + completed: bool = False, ) -> list["Choices"]: """ Convert Responses API output items to chat completion choices. @@ -1174,6 +1247,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): handle_raw_dict_callback: Optional callback for handling raw dict items structured_output: The chat request carried a json response_format, so only the first message item becomes content + replay_openai_history: Record the output's item sequence for the next turn + completed: The response finished, so its item sequence is whole Returns: List of Choices objects @@ -1196,7 +1271,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices: Final[list[Choices]] = [] index = 0 reasoning_content: str | None = None - pending_reasoning_item: _BuiltReasoningItem | None = None + pending_reasoning_items: tuple[_BuiltReasoningItem, ...] = () # Collect all tool calls to put them in a single choice # (Chat Completions API expects all tool calls in one message) @@ -1205,12 +1280,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): for item in output_items: if isinstance(item, ResponseReasoningItem): - pending_reasoning_item = _build_reasoning_item( - item_id=item.id, - encrypted_content=getattr(item, "encrypted_content", None), - summary_raw=item.summary, + pending_reasoning_items = ( + *(pending_reasoning_items if replay_openai_history else ()), + _build_reasoning_item( + item_id=item.id, + encrypted_content=getattr(item, "encrypted_content", None), + summary_raw=item.summary, + ), + ) + reasoning_content = " ".join( + s["text"] + for s in chain.from_iterable(built["summary"] for built in pending_reasoning_items) + if s.get("text") ) - reasoning_content = " ".join(s["text"] for s in pending_reasoning_item["summary"] if s.get("text")) elif isinstance(item, ResponseOutputMessage): for content in item.content: @@ -1227,7 +1309,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotations=annotations, reasoning_items=cast( list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), + (list(pending_reasoning_items) or None), ), ) @@ -1240,7 +1322,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) reasoning_content = None # flush - pending_reasoning_item = None # flush + pending_reasoning_items = () # flush index += 1 elif isinstance(item, ResponseFunctionToolCall): @@ -1300,16 +1382,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items, "stop", tool_search_items=tool_search_items ) ] - return [ - _fold_choices_into_one( - choices=choices, - tool_calls=accumulated_tool_calls, - trailing_reasoning_content=reasoning_content, - trailing_reasoning_item=pending_reasoning_item, - tool_search_items=tool_search_items, - structured_output=structured_output, - ) - ] + folded: Final = _fold_choices_into_one( + choices=choices, + tool_calls=accumulated_tool_calls, + trailing_reasoning_content=reasoning_content, + trailing_reasoning_items=pending_reasoning_items, + tool_search_items=tool_search_items, + structured_output=structured_output, + ) + if not replay_openai_history: + return [folded] + return [_with_replay_fields(folded, output_items, completed=completed, structured_output=structured_output)] @staticmethod def _build_empty_choice( @@ -1443,6 +1526,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + replay_openai_history=_served_provider( + cast("Mapping[str, object]", litellm_params) # cast-ok: litellm_params is an untyped dict + ) + == "openai", + completed=raw_response.status == "completed", ) response_is_incomplete: Final = raw_response.status == "incomplete" or ( @@ -1506,9 +1594,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + replays_openai_history: bool = False, ) -> BaseModelResponseIterator: return OpenAiResponsesToChatCompletionStreamIterator( - streaming_response, sync_stream, json_mode, structured_output=structured_output + streaming_response, + sync_stream, + json_mode, + structured_output=structured_output, + replays_openai_history=replays_openai_history, ) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: @@ -1890,8 +1983,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + replays_openai_history: bool = False, ): super().__init__(streaming_response, sync_stream, json_mode) + self._replays_openai_history: Final = replays_openai_history self._chat_completion_id: str | None = None self._served_service_tier: str | None = None self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state @@ -1937,6 +2032,8 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): def translate_responses_chunk_to_openai_stream( parsed_chunk: dict | BaseModel, tool_call_index_map: dict[int, int] | None = None, # mutable-ok: per-stream state, remapped in place + replays_openai_history: bool = False, + structured_output: bool = False, ) -> "ModelResponseStream": """ Translate a Responses API streaming chunk to OpenAI chat completion streaming format. @@ -1944,6 +2041,8 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): Args: parsed_chunk: Dict containing the Responses API event chunk tool_call_index_map: Per-stream output_index -> sequential tool_call index map + replays_openai_history: Record the completed output's item sequence for the next turn + structured_output: Only the first message item became content Returns: ModelResponseStream: OpenAI-formatted streaming chunk @@ -2154,6 +2253,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): terminal_reasoning_items_typed: Final = _as_chat_reasoning_items( _reasoning_items_from_output_items(output_items) ) + output_payload: Final = cast(object, output_items) # cast-ok: untyped + tool_search_fields: Final = _tool_search_provider_fields( + _tool_search_items_from_output_items(output_payload) + ) usage = None if response_data.get("usage"): @@ -2167,10 +2270,17 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): StreamingChoices( index=0, delta=Delta( - content="", + content=None if replays_openai_history else "", reasoning_items=terminal_reasoning_items_typed, - provider_specific_fields=_tool_search_provider_fields( - _tool_search_items_from_output_items(cast(object, output_items)) # cast-ok: untyped + provider_specific_fields=( + completed_stream_history( + output_payload, + tool_search_fields, + completed=event_type == "response.completed", + structured_output=structured_output, + ) + if replays_openai_history + else tool_search_fields ), ), finish_reason=finish_reason, @@ -2216,7 +2326,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( - chunk, tool_call_index_map=self._tool_call_index_map + chunk, + tool_call_index_map=self._tool_call_index_map, + replays_openai_history=self._replays_openai_history, + structured_output=self._structured_output, ) ) if self._structured_output and _is_message_item_done(chunk): diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f2f06d801a8..5542249d2c7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1938,6 +1938,16 @@ class CustomStreamWrapper: model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} ) + served_tier: Final = next( + ( + tier + for chunk in reversed(self.chunks) + if isinstance(tier := getattr(chunk, "service_tier", None), str) and tier + ), + None, + ) + if served_tier is not None: + setattr(model_response, "service_tier", served_tier) # noqa: B010 # pydantic extra, not a declared field _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason if _finish_reason is not None: model_response.choices[0].finish_reason = _finish_reason diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index aaa0d62eea7..d0b2850bb5b 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -63,6 +63,8 @@ def _assistant_tool_call(tool_call: ChatCompletionMessageToolCall) -> ChatComple def _assistant_message( message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_provider_specific_fields: bool ) -> ChatCompletionAssistantMessage: + from litellm.completion_extras.litellm_responses_transformation.history import retained_replay_fields + return cast( # cast-ok: dict literal with the optional message fields spread in only when set "ChatCompletionAssistantMessage", { @@ -78,7 +80,7 @@ def _assistant_message( "provider_specific_fields", getattr(message, "provider_specific_fields", None) if replays_provider_specific_fields - else None, + else retained_replay_fields(getattr(message, "provider_specific_fields", None)), ), ) if value is not None diff --git a/tests/integration/_support/responses_replay.py b/tests/integration/_support/responses_replay.py new file mode 100644 index 00000000000..c4e73ffbf2b --- /dev/null +++ b/tests/integration/_support/responses_replay.py @@ -0,0 +1,214 @@ +import json +import uuid +from collections.abc import Callable, Sequence +from itertools import chain +from typing import Final + +from integration._support.client import list_value, object_value +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue + +MODEL: Final = "gpt-6.1-sol" +PLAN: Final = "I am starting by choosing which city to check." +PICK: Final = "My pick is Tokyo, so I will check its weather." +PATCH_NOTE: Final = "I will create hello.txt containing hi." +ANSWER: Final = "It is raining in Tokyo." + +WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +PATCH_TOOL: Final[dict[str, JsonValue]] = { + "type": "custom", + "custom": {"name": "apply_patch", "description": "Apply a V4A patch to the workspace"}, +} + + +def _message(item_id: str, text: str, phase: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "phase": phase, + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +def _reasoning(item_id: str, summary: str) -> dict[str, JsonValue]: + return { + "type": "reasoning", + "id": item_id, + "summary": [{"type": "summary_text", "text": summary}], + "encrypted_content": f"enc_{item_id}", + } + + +_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_weather", + "call_id": "call_weather", + "name": "get_weather", + "arguments": '{"city":"Tokyo"}', + "status": "completed", +} +WEATHER_TURN: Final[tuple[dict[str, JsonValue], ...]] = ( + _message("msg_plan", PLAN, "commentary"), + _reasoning("rs_pick", "Tokyo is often rainy this time of year."), + _message("msg_pick", PICK, "commentary"), + _WEATHER_CALL, +) +TWO_THOUGHTS_TURN: Final[tuple[dict[str, JsonValue], ...]] = ( + _reasoning("rs_scan", "Paris and Tokyo are the two candidates."), + _reasoning("rs_pick", "Tokyo is often rainy this time of year."), + _message("msg_pick", PICK, "commentary"), + _WEATHER_CALL, +) +CALL_ONLY_TURN: Final[tuple[dict[str, JsonValue], ...]] = (_WEATHER_CALL,) +PATCH_TURN: Final[tuple[dict[str, JsonValue], ...]] = ( + _message("msg_patch", PATCH_NOTE, "commentary"), + { + "type": "custom_tool_call", + "id": "ctc_patch", + "call_id": "call_patch", + "name": "apply_patch", + "input": "*** Begin Patch\n*** Add File: hello.txt\n+hi\n*** End Patch", + "status": "completed", + }, +) +WEATHER_REPLAY: Final = ( + ("message", "assistant", "commentary", PLAN), + ("reasoning", "rs_pick"), + ("message", "assistant", "commentary", PICK), + ("function_call", "call_weather", "get_weather"), + ("function_call_output", "call_weather"), +) +TWO_THOUGHTS_REPLAY: Final = ( + ("reasoning", "rs_scan"), + ("reasoning", "rs_pick"), + ("message", "assistant", "commentary", PICK), + ("function_call", "call_weather", "get_weather"), + ("function_call_output", "call_weather"), +) + + +def _response(identity: str, output: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + return { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": MODEL, + "output": list(output), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +def _event(payload: dict[str, JsonValue]) -> bytes: + return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n".encode() + + +def _opened(item: dict[str, JsonValue]) -> dict[str, JsonValue]: + match item["type"]: + case "message": + return {**item, "status": "in_progress", "content": []} + case "function_call": + return {**item, "status": "in_progress", "arguments": ""} + case _: + return item + + +def _deltas(index: int, item: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + match item["type"]: + case "message": + text: Final = object_value(list_value(item["content"])[0])["text"] + return ( + { + "type": "response.output_text.delta", + "item_id": item["id"], + "output_index": index, + "content_index": 0, + "delta": text, + }, + ) + case "function_call": + return ( + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": index, + "delta": item["arguments"], + }, + ) + case _: + return () + + +def _item_events(index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + return ( + _event({"type": "response.output_item.added", "output_index": index, "item": _opened(item)}), + *(_event(delta) for delta in _deltas(index, item)), + _event({"type": "response.output_item.done", "output_index": index, "item": item}), + ) + + +def _stream(identity: str, output: Sequence[dict[str, JsonValue]]) -> tuple[bytes, ...]: + response: Final = _response(identity, output) + return ( + _event({"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}), + *chain.from_iterable(_item_events(index, item) for index, item in enumerate(output)), + _event({"type": "response.completed", "response": response}), + ) + + +def _answered(body: dict[str, JsonValue]) -> bool: + return any( + object_value(item).get("type") in {"function_call_output", "custom_tool_call_output"} + for item in list_value(body["input"]) + ) + + +def responses_upstream(turn_one: Sequence[dict[str, JsonValue]]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + assert request.target == "/responses", request.target + body: Final = object_value(json.loads(request.body)) + identity: Final = uuid.uuid4().hex + output: Final = (_message("msg_answer", ANSWER, "final_answer"),) if _answered(body) else turn_one + if body.get("stream"): + return Reply(content_type="text/event-stream", chunks=_stream(identity, output)) + return Reply(body=json.dumps(_response(identity, output)).encode()) + + return respond + + +def _text(content: JsonValue) -> JsonValue: + if isinstance(content, str): + return content + return "".join(str(object_value(part).get("text")) for part in list_value(content)) + + +def _shape(item: JsonValue) -> tuple[JsonValue, ...]: + fields: Final = object_value(item) + match fields.get("type"): + case "message": + return ("message", fields.get("role"), fields.get("phase"), _text(fields.get("content"))) + case "reasoning": + return ("reasoning", fields.get("id")) + case "function_call" | "custom_tool_call" as kind: + return (kind, fields.get("call_id"), fields.get("name")) + case kind: + return (kind, fields.get("call_id")) + + +def posted_inputs(wire: Wire) -> tuple[tuple[tuple[JsonValue, ...], ...], ...]: + return tuple( + tuple(_shape(item) for item in list_value(object_value(json.loads(request.body))["input"])) + for request in wire.drain() + if request.method == "POST" + ) diff --git a/tests/integration/providers/test_responses_bridge_history_replay_wire.py b/tests/integration/providers/test_responses_bridge_history_replay_wire.py new file mode 100644 index 00000000000..f85dca67be4 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_history_replay_wire.py @@ -0,0 +1,233 @@ +import json +import uuid +from collections.abc import Sequence +from functools import reduce +from itertools import chain +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.responses_replay import ( + CALL_ONLY_TURN, + MODEL, + PATCH_NOTE, + PATCH_TOOL, + PATCH_TURN, + PICK, + PLAN, + TWO_THOUGHTS_REPLAY, + TWO_THOUGHTS_TURN, + WEATHER_REPLAY, + WEATHER_TOOL, + WEATHER_TURN, + posted_inputs, + responses_upstream, +) +from integration._support.wire import wire_server +from pydantic import JsonValue + +_API_KEY: Final = "synthetic-openai-key" + + +def _chunks(response: httpx.Response) -> list[dict[str, JsonValue]]: + return [ + object_value(json.loads(line[6:])) + for line in response.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ] + + +def _tool_call(index: JsonValue, calls: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + parts: Final = [call for call in calls if call["index"] == index] + functions: Final = [object_value(part["function"]) for part in parts] + return { + "id": next(part["id"] for part in parts if part.get("id")), + "type": "function", + "function": { + "name": next(function["name"] for function in functions if function.get("name")), + "arguments": "".join(str(function.get("arguments") or "") for function in functions), + }, + } + + +def _deltas(chunks: Sequence[dict[str, JsonValue]]) -> list[dict[str, JsonValue]]: + return [ + object_value(object_value(choice)["delta"]) + for choice in chain.from_iterable(list_value(chunk["choices"]) for chunk in chunks) + ] + + +def _assembled(chunks: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + deltas: Final = _deltas(chunks) + calls: Final = [object_value(call) for call in chain.from_iterable(d.get("tool_calls") or [] for d in deltas)] + fields: Final = reduce( + lambda merged, delta: {**merged, **object_value(delta.get("provider_specific_fields") or {})}, deltas, {} + ) + return { + "role": "assistant", + "content": "".join(str(delta["content"]) for delta in deltas if delta.get("content")) or None, + "tool_calls": [_tool_call(index, calls) for index in sorted({int(str(call["index"])) for call in calls})], + "reasoning_items": list(chain.from_iterable(list_value(d.get("reasoning_items") or []) for d in deltas)), + "provider_specific_fields": fields, + } + + +def _first_turn( + gateway: Gateway, model: str, prompt: str, tools: Sequence[dict[str, JsonValue]], *, stream: bool +) -> dict[str, JsonValue]: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "tools": list(tools), + "stream": stream, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + if stream: + return _assembled(_chunks(response)) + return object_value(object_value(list_value(response.json()["choices"])[0])["message"]) + + +def _next_turn( + gateway: Gateway, + model: str, + prompt: str, + tools: Sequence[dict[str, JsonValue]], + message: dict[str, JsonValue], + **extra: JsonValue, +) -> httpx.Response: + call_id: Final = object_value(list_value(message["tool_calls"])[0])["id"] + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [ + {"role": "user", "content": prompt}, + message, + {"role": "tool", "tool_call_id": call_id, "content": "Rain, 18C"}, + ], + "tools": list(tools), + "cache": {"no-cache": True}, + **extra, + }, + ) + + +def _with_phase(message: dict[str, JsonValue], phase: str) -> dict[str, JsonValue]: + fields: Final = object_value(message.get("provider_specific_fields") or {}) + return {**message, "provider_specific_fields": {**fields, "phase": phase}} + + +def _assistant_messages(items: Sequence[tuple[JsonValue, ...]]) -> list[tuple[JsonValue, ...]]: + return [item for item in items if item[:2] == ("message", "assistant")] + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_chat_next_turn_replays_responses_items_in_the_order_openai_produced(gateway: Gateway, stream: bool) -> None: + prompt: Final = f"Pick a city and check whether it is raining there {uuid.uuid4().hex}" + with wire_server(responses_upstream(WEATHER_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=stream) + response: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], message) + inputs: Final = posted_inputs(wire) + assert response.status_code == 200, response.text + assert message["content"] == PLAN + PICK, message + assert object_value(message["provider_specific_fields"])["phase"] == "commentary", message + assert inputs[1] == (("message", "user", None, prompt), *WEATHER_REPLAY), inputs + + +def test_chat_next_turn_sends_the_phase_the_app_set_on_the_assistant_message(gateway: Gateway) -> None: + prompt: Final = f"Pick a city and check whether it is raining there {uuid.uuid4().hex}" + with wire_server(responses_upstream(WEATHER_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=False) + response: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], _with_phase(message, "final_answer")) + inputs: Final = posted_inputs(wire) + assert response.status_code == 200, response.text + assert [item[2] for item in _assistant_messages(inputs[1])] == ["final_answer"], inputs + + +def test_chat_next_turn_rejects_an_undefined_phase_unless_drop_params_is_set(gateway: Gateway) -> None: + prompt: Final = f"Pick a city and check whether it is raining there {uuid.uuid4().hex}" + with wire_server(responses_upstream(WEATHER_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _with_phase(_first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=False), "analysis") + rejected: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], message) + dropped: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], message, drop_params=True) + inputs: Final = posted_inputs(wire) + assert rejected.status_code == 400, rejected.text + assert "'analysis'" in rejected.text, rejected.text + assert dropped.status_code == 200, dropped.text + assert len(inputs) == 2, inputs + assert [item[2] for item in _assistant_messages(inputs[1])] == [None], inputs + + +def test_chat_next_turn_sends_the_text_the_app_edited_instead_of_the_recorded_items(gateway: Gateway) -> None: + prompt: Final = f"Pick a city and check whether it is raining there {uuid.uuid4().hex}" + edited_text: Final = "I picked Tokyo." + with wire_server(responses_upstream(WEATHER_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=False) + response: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], {**message, "content": edited_text}) + inputs: Final = posted_inputs(wire) + assert response.status_code == 200, response.text + assert _assistant_messages(inputs[1]) == [("message", "assistant", "commentary", edited_text)], inputs + + +def test_chat_next_turn_replays_a_custom_tool_call_as_a_custom_tool_call(gateway: Gateway) -> None: + prompt: Final = f"Create hello.txt containing hi {uuid.uuid4().hex}" + with wire_server(responses_upstream(PATCH_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [PATCH_TOOL], stream=False) + response: Final = _next_turn(gateway, model, prompt, [PATCH_TOOL], message) + inputs: Final = posted_inputs(wire) + assert response.status_code == 200, response.text + assert inputs[1] == ( + ("message", "user", None, prompt), + ("message", "assistant", "commentary", PATCH_NOTE), + ("custom_tool_call", "call_patch", "apply_patch"), + ("custom_tool_call_output", "call_patch"), + ), inputs + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_chat_reply_keeps_every_reasoning_item_and_the_next_turn_replays_each_one( + gateway: Gateway, stream: bool +) -> None: + prompt: Final = f"Weigh two cities, then check whether it is raining in your pick {uuid.uuid4().hex}" + with wire_server(responses_upstream(TWO_THOUGHTS_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=stream) + response: Final = _next_turn(gateway, model, prompt, [WEATHER_TOOL], message) + inputs: Final = posted_inputs(wire) + assert response.status_code == 200, response.text + assert [object_value(item)["id"] for item in list_value(message["reasoning_items"])] == ["rs_scan", "rs_pick"] + assert inputs[1] == (("message", "user", None, prompt), *TWO_THOUGHTS_REPLAY), inputs + + +def test_chat_stream_of_a_tool_call_only_turn_carries_null_content_like_the_json_reply(gateway: Gateway) -> None: + prompt: Final = f"Check whether it is raining in Tokyo {uuid.uuid4().hex}" + with wire_server(responses_upstream(CALL_ONLY_TURN)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/responses/{MODEL}", api_base=wire.url, api_key=_API_KEY) + message: Final = _first_turn(gateway, model, prompt, [WEATHER_TOOL], stream=False) + streamed: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "tools": [WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert streamed.status_code == 200, streamed.text + deltas: Final = _deltas(_chunks(streamed)) + assert message["content"] is None, message + assert [delta["content"] for delta in deltas if delta.get("content") is not None] == [], deltas diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index e3978b4debd..396930e9971 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -1,6 +1,7 @@ import json import uuid from typing import Final +from unittest.mock import ANY import httpx import pytest @@ -288,6 +289,7 @@ def test_chat_over_responses_deployment_merges_message_and_function_call(gateway "index": 0, } ], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -397,6 +399,7 @@ def test_chat_over_responses_deployment_keeps_reasoning_with_merged_tool_call(ga "index": 0, } ], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -480,6 +483,7 @@ def test_chat_over_responses_deployment_returns_tool_call_only_reply_as_one_choi "index": 0, } ], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -570,6 +574,7 @@ def test_chat_over_responses_deployment_merges_function_call_followed_by_message "index": 0, } ], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -671,6 +676,7 @@ def test_chat_over_responses_deployment_folds_two_messages_and_function_call_int "index": 0, } ], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -729,6 +735,7 @@ def test_chat_over_responses_deployment_folds_commentary_and_final_answer_into_o "message": { "role": "assistant", "content": 'Let me answer that.{"city": "Paris", "summary": "Sunny, 21C"}', + "provider_specific_fields": {"phase": "final_answer", "openai_responses_replay": ANY}, }, } ], response.text @@ -843,6 +850,7 @@ def test_chat_over_responses_deployment_folds_reasoning_two_messages_and_functio "reasoning_content": "Checking the forecast.", "reasoning_items": [_REASONING_ITEM], "tool_calls": _TWO_CALLS_CHAT_SHAPE[:1], + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -884,6 +892,7 @@ def test_chat_over_responses_deployment_folds_message_and_two_function_calls_int "role": "assistant", "content": "Checking both.", "tool_calls": _TWO_CALLS_CHAT_SHAPE, + "provider_specific_fields": {"openai_responses_replay": ANY}, }, } ], response.text @@ -1353,7 +1362,11 @@ def test_chat_over_responses_deployment_with_json_schema_returns_first_of_two_fi { "finish_reason": "stop", "index": 0, - "message": {"role": "assistant", "content": '{"city":"Paris"}'}, + "message": { + "role": "assistant", + "content": '{"city":"Paris"}', + "provider_specific_fields": {"phase": "final_answer"}, + }, } ], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/integration/sdk/test_tool_loop_responses_history_replay_sdk.py b/tests/integration/sdk/test_tool_loop_responses_history_replay_sdk.py new file mode 100644 index 00000000000..0c58545182b --- /dev/null +++ b/tests/integration/sdk/test_tool_loop_responses_history_replay_sdk.py @@ -0,0 +1,37 @@ +import uuid +from typing import Final, cast + +from integration._support.responses_replay import ( + ANSWER, + MODEL, + WEATHER_REPLAY, + WEATHER_TOOL, + WEATHER_TURN, + posted_inputs, + responses_upstream, +) +from integration._support.wire import wire_server + +import litellm +from litellm.types.llms.openai import ChatCompletionToolMessage, ChatCompletionToolParam +from litellm.types.utils import ChatCompletionMessageToolCall + + +def _rain(call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + return {"role": "tool", "tool_call_id": call.id, "content": "Rain, 18C"} + + +def test_run_tool_loop_replays_responses_items_in_the_order_openai_produced() -> None: + prompt: Final = f"Pick a city and check whether it is raining there {uuid.uuid4().hex}" + with wire_server(responses_upstream(WEATHER_TURN)) as wire: + content: Final = litellm.run_tool_loop( + model=f"openai/responses/{MODEL}", + messages=[{"role": "user", "content": prompt}], + tools=[cast(ChatCompletionToolParam, WEATHER_TOOL)], + execute_tool=_rain, + api_base=wire.url, + api_key="synthetic-openai-key", + ) + inputs: Final = posted_inputs(wire) + assert content == ANSWER + assert inputs[1] == (("message", "user", None, prompt), *WEATHER_REPLAY), inputs diff --git a/tests/integration/translation/chat_completions/bases/openai_responses.py b/tests/integration/translation/chat_completions/bases/openai_responses.py index d086a56a221..f0fb0af6131 100644 --- a/tests/integration/translation/chat_completions/bases/openai_responses.py +++ b/tests/integration/translation/chat_completions/bases/openai_responses.py @@ -97,7 +97,17 @@ GPT_5_6_SOL_TEST_CASE: Final = TranslationTestCase( "created": ANY, "model": "openai/responses/gpt-5.6-sol", "object": "chat.completion", - "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "Hello!", "role": "assistant"}}], + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello!", + "role": "assistant", + "provider_specific_fields": {"phase": "final_answer", "openai_responses_replay": ANY}, + }, + } + ], "usage": { "completion_tokens": 6, "prompt_tokens": 19, @@ -252,6 +262,7 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( "summary": [], } ], + "provider_specific_fields": {"phase": "commentary", "openai_responses_replay": ANY}, }, } ], diff --git a/tests/integration/translation/response_scenarios.py b/tests/integration/translation/response_scenarios.py index b49d2b71441..7168ad0ade5 100644 --- a/tests/integration/translation/response_scenarios.py +++ b/tests/integration/translation/response_scenarios.py @@ -57,7 +57,10 @@ def two_messages_before_tools(case: TranslationTestCase) -> TranslationTestCase: def _messages_then_tools_response(case: TranslationTestCase) -> Mapping[str, JsonValue]: expected: Final = case.expected_litellm_response if case.litellm_endpoint == "/v1/chat/completions": - choice: Final = items(expected["choices"])[0] + choices: Final = expected["choices"] + assert isinstance(choices, list) and isinstance(choices[0], dict), choices + message: Final = choices[0]["message"] + assert isinstance(message, dict), message return { **expected, "choices": [ @@ -67,7 +70,12 @@ def _messages_then_tools_response(case: TranslationTestCase) -> Mapping[str, Jso "message": { "role": "assistant", "content": "".join(COMMENTARY), - "tool_calls": object_value(choice["message"])["tool_calls"], + "tool_calls": message["tool_calls"], + **( + {"provider_specific_fields": message["provider_specific_fields"]} + if "provider_specific_fields" in message + else {} + ), }, } ], diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_history.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_history.py new file mode 100644 index 00000000000..a2efbfef5dc --- /dev/null +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_history.py @@ -0,0 +1,627 @@ +import json +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.completion_extras.litellm_responses_transformation import history +from litellm.exceptions import MidStreamFallbackError + +OPENAI_RESPONSES_URL: Final = "https://api.openai.com/v1/responses" +AZURE_API_BASE: Final = "https://litellm-history-test.openai.azure.com" +MODEL: Final = "gpt-5.6" +TOOLS: Final = ( + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + }, +) +USER: Final = {"role": "user", "content": "Weather in Paris? Narrate as you go."} + +COMMENTARY: Final = { + "type": "message", + "id": "msg_commentary", + "status": "completed", + "role": "assistant", + "phase": "commentary", + "content": [{"type": "output_text", "text": "Checking the weather.", "annotations": []}], +} +FIRST_REASONING: Final = { + "type": "reasoning", + "id": "rs_first", + "summary": [{"type": "summary_text", "text": "Paris is likely wet."}], + "encrypted_content": "enc_first", +} +SECOND_REASONING: Final = {"type": "reasoning", "id": "rs_second", "summary": [], "encrypted_content": "enc_second"} +WEATHER_CALL: Final = { + "type": "function_call", + "id": "fc_weather", + "call_id": "call_weather_paris", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", +} +FINAL_ANSWER: Final = { + "type": "message", + "id": "msg_final", + "status": "completed", + "role": "assistant", + "phase": "final_answer", + "content": [ + {"type": "output_text", "text": "Paris: ", "annotations": []}, + {"type": "output_text", "text": "light rain.", "annotations": []}, + ], +} +OUTPUT: Final = (COMMENTARY, FIRST_REASONING, SECOND_REASONING, WEATHER_CALL, FINAL_ANSWER) +REPLAYED: Final = ( + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking the weather."}], + "phase": "commentary", + }, + { + "type": "reasoning", + "id": "rs_first", + "summary": [{"type": "summary_text", "text": "Paris is likely wet."}], + "encrypted_content": "enc_first", + }, + {"type": "reasoning", "id": "rs_second", "summary": [], "encrypted_content": "enc_second"}, + {"type": "function_call", "call_id": "call_weather_paris", "name": "get_weather", "arguments": '{"city":"Paris"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Paris: "}, {"type": "output_text", "text": "light rain."}], + "phase": "final_answer", + }, +) +TOOL_RESULT: Final = {"role": "tool", "tool_call_id": "call_weather_paris", "content": "light rain"} +FOLLOW_UP: Final = { + "type": "message", + "id": "msg_follow_up", + "status": "completed", + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], +} + + +def _response(output: Sequence[Mapping[str, object]], status: str = "completed") -> dict[str, object]: + return { + "id": "resp_history", + "object": "response", + "created_at": 1789788253, + "status": status, + "model": MODEL, + "output": list(output), + "incomplete_details": {"reason": "max_output_tokens"} if status == "incomplete" else None, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +def _added_item(item: Mapping[str, object]) -> Mapping[str, object]: + match item.get("type"): + case "message": + return {**item, "status": "in_progress", "content": []} + case "function_call": + return {**item, "status": "in_progress", "arguments": ""} + case "custom_tool_call": + return {**item, "status": "in_progress", "input": ""} + case _: + return item + + +def _delta_events(index: int, item: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + match item.get("type"): + case "message": + parts: Final = item["content"] + assert isinstance(parts, list) + return tuple( + { + "type": "response.output_text.delta", + "item_id": item["id"], + "output_index": index, + "content_index": part_index, + "delta": part["text"], + } + for part_index, part in enumerate(parts) + ) + case "function_call": + return ( + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": index, + "delta": item["arguments"], + }, + ) + case "custom_tool_call": + return ( + { + "type": "response.custom_tool_call_input.delta", + "item_id": item["id"], + "output_index": index, + "delta": item["input"], + }, + ) + case _: + return () + + +def _item_events(index: int, item: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"type": "response.output_item.added", "output_index": index, "item": _added_item(item)}, + *_delta_events(index, item), + {"type": "response.output_item.done", "output_index": index, "item": item}, + ) + + +def _sse(output: Sequence[Mapping[str, object]], terminal_event: str | None = "response.completed") -> httpx.Response: + item_events: Final = tuple(event for index, item in enumerate(output) for event in _item_events(index, item)) + terminal: Final = ( + () + if terminal_event is None + else ( + { + "type": terminal_event, + "response": _response(output, "completed" if terminal_event == "response.completed" else "incomplete"), + }, + ) + ) + events: Final = ( + {"type": "response.created", "response": {**_response((), "in_progress"), "usage": None}}, + *item_events, + *terminal, + ) + body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + return httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + + +def _assistant_turn( + output: Sequence[Mapping[str, object]], + *, + stream: bool, + model: str = f"openai/responses/{MODEL}", + terminal_event: str | None = "response.completed", + status: str = "completed", + response_format: Mapping[str, object] | None = None, +) -> dict[str, object]: + reply: Final = httpx.Response( + 200, content=json.dumps(_response(output, status)), headers={"content-type": "application/json"} + ) + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith=OPENAI_RESPONSES_URL if model.startswith("openai/") else AZURE_API_BASE).mock( + return_value=_sse(output, terminal_event) if stream else reply + ) + response: Final = litellm.completion( + model=model, + messages=[USER], + tools=list(TOOLS), + api_key="sk-test", + api_base=None if model.startswith("openai/") else AZURE_API_BASE, + api_version=None if model.startswith("openai/") else "2025-04-01-preview", + stream=stream, + response_format=response_format, + ) + message: Final = ( + (litellm.stream_chunk_builder(list(response), messages=[USER]) if stream else response).choices[0].message + ) + return json.loads(json.dumps(message.model_dump())) + + +def _next_turn_input( + assistant: Mapping[str, object], + model: str = f"openai/responses/{MODEL}", + tools: Sequence[Mapping[str, object]] = TOOLS, + drop_params: bool = False, +) -> list[Mapping[str, object]]: + with respx.mock(assert_all_called=True) as router: + upstream: Final = router.post( + url__startswith=OPENAI_RESPONSES_URL if model.startswith("openai/") else AZURE_API_BASE + ).mock(return_value=httpx.Response(200, json=_response((FOLLOW_UP,)))) + litellm.completion( + model=model, + messages=[USER, dict(assistant), TOOL_RESULT], + tools=list(tools), + api_key="sk-test", + api_base=None if model.startswith("openai/") else AZURE_API_BASE, + api_version=None if model.startswith("openai/") else "2025-04-01-preview", + drop_params=drop_params, + ) + return json.loads(upstream.calls[0].request.content)["input"] + + +def _replayed_fields(message: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = message.get("provider_specific_fields") + return fields if isinstance(fields, Mapping) else {} + + +def _types(items: Sequence[Mapping[str, object]]) -> list[object]: + return [item.get("type") for item in items] + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_next_turn_replays_every_item_phase_and_boundary_once_in_order(stream: bool) -> None: + assistant: Final = _assistant_turn(OUTPUT, stream=stream) + + sent: Final = _next_turn_input(assistant) + + assert assistant["content"] == "Checking the weather.Paris: light rain." + assert _replayed_fields(assistant)["phase"] == "final_answer" + assert sent[1:-1] == list(REPLAYED) + assert sent[-1]["type"] == "function_call_output" + assert sent[-1]["call_id"] == "call_weather_paris" + + +def _edit_content(message: Mapping[str, object]) -> Mapping[str, object]: + return {**message, "content": "Checking the weather. Paris: heavy rain."} + + +def _edit_phase(message: Mapping[str, object]) -> Mapping[str, object]: + return {**message, "provider_specific_fields": {**_replayed_fields(message), "phase": "commentary"}} + + +def _edit_arguments(message: Mapping[str, object]) -> Mapping[str, object]: + tool_calls: Final = message["tool_calls"] + assert isinstance(tool_calls, list) + edited_call: Final = {**tool_calls[0], "function": {**tool_calls[0]["function"], "arguments": '{"city":"Lyon"}'}} + return {**message, "tool_calls": [edited_call]} + + +@pytest.mark.parametrize( + ("edit", "expected_text", "expected_phase", "expected_arguments"), + [ + (_edit_content, "Checking the weather. Paris: heavy rain.", "final_answer", '{"city":"Paris"}'), + (_edit_phase, "Checking the weather.Paris: light rain.", "commentary", '{"city":"Paris"}'), + (_edit_arguments, "Checking the weather.Paris: light rain.", "final_answer", '{"city":"Lyon"}'), + ], + ids=["content", "phase", "arguments"], +) +def test_an_edited_chat_field_wins_over_the_recorded_items( + edit: Callable[[Mapping[str, object]], Mapping[str, object]], + expected_text: str, + expected_phase: str, + expected_arguments: str, +) -> None: + sent: Final = _next_turn_input(edit(_assistant_turn(OUTPUT, stream=False))) + + assert _types(sent[1:-1]) == ["reasoning", "reasoning", "message", "function_call"] + assert sent[3]["content"] == [{"type": "output_text", "text": expected_text}] + assert sent[3]["phase"] == expected_phase + assert sent[4]["arguments"] == expected_arguments + + +def _without_replay(fields: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in fields.items() if key != "openai_responses_replay"} + + +def _with_replay(fields: Mapping[str, object], replay: object) -> Mapping[str, object]: + return {**fields, "openai_responses_replay": replay} + + +def _replay_of(fields: Mapping[str, object]) -> Mapping[str, object]: + replay: Final = fields["openai_responses_replay"] + assert isinstance(replay, Mapping) + return replay + + +@pytest.mark.parametrize( + "corrupt", + [ + _without_replay, + lambda fields: _with_replay(fields, {**_replay_of(fields), "checksum": "0" * 64}), + lambda fields: _with_replay(fields, {**_replay_of(fields), "version": 2}), + lambda fields: _with_replay(fields, {**_replay_of(fields), "items": [{"type": "message", "parts": [-1]}]}), + lambda fields: _with_replay(fields, "not a replay"), + ], + ids=["missing", "checksum", "version", "items", "not-a-mapping"], +) +def test_missing_or_corrupt_replay_metadata_falls_back_to_the_chat_fields( + corrupt: Callable[[Mapping[str, object]], Mapping[str, object]], +) -> None: + assistant: Final = _assistant_turn(OUTPUT, stream=False) + + sent: Final = _next_turn_input({**assistant, "provider_specific_fields": corrupt(_replayed_fields(assistant))}) + + assert _types(sent[1:-1]) == ["reasoning", "reasoning", "message", "function_call"] + assert sent[3]["content"] == [{"type": "output_text", "text": "Checking the weather.Paris: light rain."}] + assert sent[3]["phase"] == "final_answer" + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_an_unphased_answer_after_commentary_is_never_labeled_commentary(stream: bool) -> None: + unphased_answer: Final = {key: value for key, value in FINAL_ANSWER.items() if key != "phase"} + assistant: Final = _assistant_turn((COMMENTARY, WEATHER_CALL, unphased_answer), stream=stream) + + sent: Final = _next_turn_input(_edit_content(assistant)) + + assert "phase" not in _replayed_fields(assistant) + assert _types(sent[1:-1]) == ["message", "function_call"] + assert "phase" not in sent[1] + + +def _with_text(item: Mapping[str, object], text: str) -> Mapping[str, object]: + return {**item, "content": [{"type": "output_text", "text": text, "annotations": []}]} + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_a_structured_output_turn_carries_the_phase_of_the_message_it_returns(stream: bool) -> None: + output: Final = (_with_text(COMMENTARY, '{"city":"Paris"}'), _with_text(FINAL_ANSWER, '{"city":"Lyon"}')) + + assistant: Final = _assistant_turn(output, stream=stream, response_format={"type": "json_object"}) + + assert assistant["content"] == '{"city":"Paris"}' + assert _replayed_fields(assistant)["phase"] == "commentary" + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_a_lone_surrogate_in_the_reply_still_gets_a_record(stream: bool) -> None: + assistant: Final = _assistant_turn((_with_text(COMMENTARY, "Checking \ud83d"), WEATHER_CALL), stream=stream) + + assert "openai_responses_replay" in _replayed_fields(assistant) + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_a_last_message_the_bridge_cannot_read_leaves_the_turn_without_a_phase(stream: bool) -> None: + unreadable: Final = {**_with_text(FINAL_ANSWER, "It is raining in Paris."), "phase": "summary"} + assistant: Final = _assistant_turn((COMMENTARY, WEATHER_CALL, unreadable), stream=stream) + + assert "phase" not in _replayed_fields(assistant) + + +def test_a_phase_openai_does_not_define_is_dropped_under_drop_params() -> None: + assistant: Final = _assistant_turn(OUTPUT, stream=False) + + sent: Final = _next_turn_input({**assistant, "provider_specific_fields": {"phase": "analysis"}}, drop_params=True) + + assert _types(sent[1:-1]) == ["reasoning", "reasoning", "message", "function_call"] + assert "phase" not in sent[3] + + +def test_a_phase_openai_does_not_define_is_rejected_without_drop_params() -> None: + assistant: Final = _assistant_turn(OUTPUT, stream=False) + + with respx.mock(assert_all_called=False) as router: + upstream: Final = router.post(url__startswith=OPENAI_RESPONSES_URL).mock( + return_value=httpx.Response(200, json=_response((FOLLOW_UP,))) + ) + with pytest.raises(litellm.BadRequestError, match="'analysis'"): + litellm.completion( + model=f"openai/responses/{MODEL}", + messages=[USER, {**assistant, "phase": "analysis"}, TOOL_RESULT], + tools=list(TOOLS), + api_key="sk-test", + ) + + assert not upstream.called + + +@pytest.mark.parametrize("fields", [None, ["not", "a", "mapping"]], ids=["no-fields", "malformed-fields"]) +def test_an_assistant_message_the_app_wrote_replays_without_a_phase(fields: object) -> None: + sent: Final = _next_turn_input( + { + "role": "assistant", + "content": "Checking Paris.", + "tool_calls": [ + { + "id": "call_weather_paris", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + "provider_specific_fields": fields, + } + ) + + assert _types(sent[1:-1]) == ["message", "function_call"] + assert "phase" not in sent[1] + + +SEARCH_CALL: Final = { + "type": "tool_search_call", + "id": "tsc_weather", + "arguments": {"paths": ["get_weather"]}, + "status": "completed", +} +SEARCH_OUTPUT: Final = { + "type": "tool_search_output", + "id": "tso_weather", + "status": "completed", + "tools": [{"name": "get_weather"}], +} + + +@pytest.mark.parametrize( + ("tools", "expected"), + [ + ((*TOOLS, {"type": "tool_search"}), ["tool_search_call", "tool_search_output", *_types(REPLAYED)]), + (TOOLS, ["reasoning", "reasoning", "message", "function_call"]), + ], + ids=["searching", "not-searching"], +) +def test_tool_search_items_replay_only_when_the_next_request_searches( + tools: Sequence[Mapping[str, object]], expected: Sequence[str] +) -> None: + assistant: Final = _assistant_turn((SEARCH_CALL, SEARCH_OUTPUT, *OUTPUT), stream=False) + + sent: Final = _next_turn_input(assistant, tools=tools) + + assert _types(sent[1:-1]) == list(expected) + + +WEB_SEARCH: Final = {"type": "web_search_call", "id": "ws_1", "status": "completed", "action": {"type": "search"}} + + +@pytest.mark.parametrize( + ("output", "stream", "terminal_event", "status"), + [ + (OUTPUT, False, "response.completed", "incomplete"), + (OUTPUT, True, "response.incomplete", "incomplete"), + ((WEB_SEARCH, *OUTPUT), False, "response.completed", "completed"), + ((WEB_SEARCH, *OUTPUT), True, "response.completed", "completed"), + ], + ids=["incomplete-json", "incomplete-stream", "unsupported-item-json", "unsupported-item-stream"], +) +def test_an_incomplete_or_unsupported_output_records_no_replay( + output: Sequence[Mapping[str, object]], stream: bool, terminal_event: str, status: str +) -> None: + assistant: Final = _assistant_turn(output, stream=stream, terminal_event=terminal_event, status=status) + + sent: Final = _next_turn_input(assistant) + + assert "openai_responses_replay" not in _replayed_fields(assistant) + assert _types(sent[1:-1]) == ["reasoning", "reasoning", "message", "function_call"] + + +def _chunks_before_error(stream: Iterator[object]) -> Iterator[object]: + try: + yield from stream + except MidStreamFallbackError: + return + + +def test_an_interrupted_stream_records_no_replay() -> None: + with respx.mock(assert_all_called=True) as router: + router.post(OPENAI_RESPONSES_URL).mock(return_value=_sse(OUTPUT, terminal_event=None)) + stream: Final = litellm.completion( + model=f"openai/responses/{MODEL}", messages=[USER], tools=list(TOOLS), api_key="sk-test", stream=True + ) + received: Final = tuple(_chunks_before_error(iter(stream))) + assistant: Final = litellm.stream_chunk_builder(list(received), messages=[USER]).choices[0].message.model_dump() + + assert assistant["content"] == "Checking the weather.Paris: light rain." + assert _replayed_fields(assistant) == {} + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_a_non_openai_route_records_no_phase_or_replay(stream: bool) -> None: + assistant: Final = _assistant_turn(OUTPUT, stream=stream, model=f"azure/responses/{MODEL}") + + assert _replayed_fields(assistant) == {} + + +def test_a_non_openai_route_ignores_a_recorded_replay() -> None: + sent: Final = _next_turn_input(_assistant_turn(OUTPUT, stream=False), model=f"azure/responses/{MODEL}") + + assert _types(sent[1:-1]) == ["reasoning", "reasoning", "message", "function_call"] + assert all("phase" not in item for item in sent) + + +SECOND_THOUGHT: Final = {**SECOND_REASONING, "summary": [{"type": "summary_text", "text": "Lyon is likely dry."}]} +LYON_CALL: Final = {**WEATHER_CALL, "id": "fc_lyon", "call_id": "call_weather_lyon", "arguments": '{"city":"Lyon"}'} + + +@pytest.mark.parametrize( + ("model", "expected_ids", "expected_text"), + [ + (f"openai/responses/{MODEL}", ["rs_first", "rs_second"], "Paris is likely wet. Lyon is likely dry."), + (f"azure/responses/{MODEL}", ["rs_second"], "Lyon is likely dry."), + ], + ids=["openai", "azure"], +) +def test_only_the_openai_route_keeps_every_reasoning_item_of_a_tool_call_turn( + model: str, expected_ids: list[str], expected_text: str +) -> None: + assistant: Final = _assistant_turn( + (FIRST_REASONING, WEATHER_CALL, SECOND_THOUGHT, LYON_CALL), stream=False, model=model + ) + + reasoning_items: Final = assistant["reasoning_items"] + assert isinstance(reasoning_items, list) + assert [item["id"] for item in reasoning_items] == expected_ids + assert assistant["reasoning_content"] == expected_text + + +@pytest.mark.parametrize( + ("model", "expected_content"), + [(f"openai/responses/{MODEL}", None), (f"azure/responses/{MODEL}", "")], + ids=["openai", "azure"], +) +def test_only_the_openai_route_ends_a_streamed_tool_call_turn_with_null_content( + model: str, expected_content: str | None +) -> None: + assistant: Final = _assistant_turn((FIRST_REASONING, WEATHER_CALL), stream=True, model=model) + + assert assistant["content"] == expected_content + + +def _resealed(message: Mapping[str, object], items: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + chat: Final = history._chat_message(message) + assert chat is not None + entries: Final = history._Replay.model_validate({"version": 1, "items": items, "checksum": ""}).items + replay: Final = {"version": 1, "items": list(items), "checksum": history._checksum(chat, entries)} + return {**message, "provider_specific_fields": _with_replay(_replayed_fields(message), replay)} + + +@pytest.mark.parametrize( + "items", + [ + ( + {"type": "message", "parts": [21]}, + {"type": "function_call", "index": 1}, + {"type": "message", "parts": [7, 11]}, + ), + ( + {"type": "message", "parts": [5]}, + {"type": "function_call", "index": 0}, + {"type": "message", "parts": [7, 11]}, + ), + ], + ids=["tool-call-index", "text-lengths"], +) +def test_a_resealed_record_that_does_not_fit_the_chat_fields_falls_back(items: Sequence[Mapping[str, object]]) -> None: + assistant: Final = _assistant_turn((COMMENTARY, WEATHER_CALL, FINAL_ANSWER), stream=False) + + sent: Final = _next_turn_input(_resealed(assistant, items)) + + assert _types(sent[1:-1]) == ["message", "function_call"] + assert sent[1]["content"] == [{"type": "output_text", "text": "Checking the weather.Paris: light rain."}] + + +PATCH_CALL: Final = { + "type": "custom_tool_call", + "id": "ctc_patch", + "call_id": "call_patch_readme", + "name": "apply_patch", + "input": "*** Begin Patch\n*** End Patch", + "status": "completed", +} + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +def test_a_custom_tool_call_replays_with_its_custom_output(stream: bool) -> None: + assistant: Final = _assistant_turn((COMMENTARY, PATCH_CALL), stream=stream) + + with respx.mock(assert_all_called=True) as router: + upstream: Final = router.post(OPENAI_RESPONSES_URL).mock( + return_value=httpx.Response(200, json=_response((FOLLOW_UP,))) + ) + litellm.completion( + model=f"openai/responses/{MODEL}", + messages=[USER, assistant, {"role": "tool", "tool_call_id": "call_patch_readme", "content": "applied"}], + tools=[{"type": "custom", "custom": {"name": "apply_patch"}}], + api_key="sk-test", + ) + sent: Final = json.loads(upstream.calls[0].request.content)["input"] + + assert sent[1:] == [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking the weather."}], + "phase": "commentary", + }, + { + "type": "custom_tool_call", + "call_id": "call_patch_readme", + "name": "apply_patch", + "input": PATCH_CALL["input"], + }, + {"type": "custom_tool_call_output", "call_id": "call_patch_readme", "output": "applied"}, + ] diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 46f1dbf2e7a..d59f59703ad 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -670,6 +670,41 @@ def test_streamed_tool_search_turn_replays_with_its_reasoning_item(): assert request["input"][5]["namespace"] == "get_weather" +@pytest.mark.parametrize( + "answer_item", + [ + ResponseFunctionToolCall( + type="function_call", id="fc_1", call_id="call_weather", name="get_weather", arguments="{}" + ), + ResponseOutputMessage( + type="message", + id="msg_1", + role="assistant", + status="completed", + content=[ResponseOutputText(type="output_text", text="Rain.", annotations=[])], + ), + ], + ids=["tool-call", "message"], +) +def test_every_reasoning_item_before_an_answer_reaches_an_openai_chat_message(answer_item: object) -> None: + reasoning: Final = tuple( + ResponseReasoningItem( + type="reasoning", + id=f"rs_{name}", + summary=[Summary(type="summary_text", text=f"{name.title()}.")], + encrypted_content=f"enc_{name}", + ) + for name in ("first", "second") + ) + + message: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + [*reasoning, answer_item], replay_openai_history=True + )[0].message + + assert [item["id"] for item in message.reasoning_items] == ["rs_first", "rs_second"] + assert message.reasoning_content == "First. Second." + + def test_transform_response_with_reasoning_and_output(): """Test transform_response handles ResponsesAPIResponse with reasoning items and output messages.""" from unittest.mock import Mock @@ -1270,9 +1305,7 @@ def test_transform_request_forwards_defer_loading_only_beside_tool_search( [{"type": "tool_search", "execution": "client"}, {"type": "tool_search"}], ids=["client-only", "hosted-and-client"], ) -def test_transform_request_drops_defer_loading_beside_client_tool_search( - deferred_tool, drop_params, tool_search_tool -): +def test_transform_request_drops_defer_loading_beside_client_tool_search(deferred_tool, drop_params, tool_search_tool): tool_search_tools: Final = ( [tool_search_tool] if tool_search_tool.get("execution") == "client" @@ -4396,7 +4429,11 @@ def test_tool_search_items_replay_only_to_a_request_that_carries_tool_search(too "provider_specific_fields": { "tool_search_items": [ {"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}}, - {"type": "tool_search_output", "id": "tso_1", "tools": [{"type": "function", "name": "get_weather"}]}, + { + "type": "tool_search_output", + "id": "tso_1", + "tools": [{"type": "function", "name": "get_weather"}], + }, ] }, }, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b580f526a0f..f4502937fce 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7820,6 +7820,38 @@ async def test_openai_stream_relays_the_served_service_tier_on_every_chunk_inclu assert relayed[-1]["usage"]["total_tokens"] == 11 +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_relays_the_served_service_tier_on_the_finish_chunk_it_adds(logging_obj: Logging, sync_mode: bool): + from litellm.utils import ModelResponseListIterator + + def _chunk(content: str, finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-tier", + created=1742056047, + model="gpt-4.1-mini", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content))], + service_tier="priority", + ) + + logging_obj.update_environment_variables( + model="gpt-4.1-mini", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + wrapper = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=[_chunk("Hi", None), _chunk("!", "stop")]), + model="gpt-4.1-mini", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + relayed = ( + [chunk.model_dump() for chunk in wrapper] if sync_mode else [chunk.model_dump() async for chunk in wrapper] + ) + + assert [chunk["choices"][0]["finish_reason"] for chunk in relayed] == [None, None, "stop"], relayed + assert [chunk.get("service_tier") for chunk in relayed] == ["priority"] * len(relayed), relayed + + def _last_chunk_carries_finish_reason_wrapper( logging_obj: Logging, finish_reason: str, sync_stream: bool ) -> CustomStreamWrapper: diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index e3ebfe7258f..de8f62803cc 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -606,3 +606,69 @@ def test_client_tool_turn_is_replayed_without_its_provider_specific_fields(respx replayed_assistant: Final = _request_bodies(respx_mock)[1]["messages"][1] assert replayed_assistant["tool_calls"][0]["id"] == "call_1" assert "provider_specific_fields" not in replayed_assistant + + +def test_responses_bridge_replays_commentary_reasoning_and_tool_call_in_order(respx_mock: respx.MockRouter) -> None: + round_one: Final = _responses_payload( + "resp_1", + [ + { + "type": "message", + "id": "msg_commentary", + "status": "completed", + "role": "assistant", + "phase": "commentary", + "content": [{"type": "output_text", "text": "Checking Paris.", "annotations": []}], + }, + {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_weather_paris", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + ) + round_two: Final = _responses_payload( + "resp_2", + [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "Paris is 72F", "annotations": []}], + } + ], + ) + respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[httpx.Response(200, json=round_one), httpx.Response(200, json=round_two)] + ) + + answer: Final = litellm.run_tool_loop( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + api_key="sk-test", + ) + + assert answer == "Paris is 72F" + assert _request_bodies(respx_mock)[1]["input"][1:] == [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking Paris."}], + "phase": "commentary", + }, + {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"}, + {"type": "function_call", "call_id": "call_weather_paris", "name": "get_weather", "arguments": '{"city": "Paris"}'}, + { + "type": "function_call_output", + "call_id": "call_weather_paris", + "output": [{"type": "input_text", "text": '{"temp": "72F"}'}], + }, + ] From 032342d2b7695c6db7a22f3b78cf6740090c640e Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:31:22 -0700 Subject: [PATCH 11/58] fix(responses): resume background response streams with stream and starting_after (#48819) * fix(responses): resume background response streams with stream and starting_after Retrieving a response with stream=true and starting_after=N now returns the replayed event stream from OpenAI and Azure instead of one JSON object, on the SDK (get_responses / aget_responses) and on the proxy's GET /v1/responses/{id} Providers that cannot replay a stored stream drop the params under drop_params and raise UnsupportedParamsError without it. Malformed values answer 400 with the param named. A replay that reaches its terminal event settles the background response's bill once and passes the provider's usage through without estimating one. An empty replay (cursor at the end) ends normally * fix(responses): encrypt replayed stream ids and resume Azure AI Foundry streams A replayed background stream now gets the same response id encryption as a created one, so the key that made the response can reuse the id. Azure AI Foundry deployments resume streams like Azure ones, and the chat-bridged Responses iterator closes its wrapped stream instead of raising. * fix(responses): drop malformed resume params under drop_params and keep JSON answers to a streamed retrieve A bad stream flag drops to a plain retrieve and a bad cursor to a full replay when drop_params is on, and a cursor without stream is ignored even when malformed. A streamed retrieve that the upstream answers with JSON returns the response object instead of an empty stream. get_responses forwards only the include keys from extra_query. * fix(responses): give proxy stream resumes the deployment stream_timeout and document retrieve query params * fix(responses): keep deployment stream off retrieves and reject oversized cursors * fix(responses): keep router default stream off plain retrieves * test(responses): black box coverage for resume refusals, drop_params, stream timeout and billing * test(responses): replay a real stream in the billing test and type the fake wrapper chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_utils/check_responses_cost.py | 14 +- .../internal_call_metadata.py | 2 +- .../llms/azure/responses/transformation.py | 4 +- .../llms/base_llm/responses/transformation.py | 5 +- litellm/llms/custom_httpx/http_handler.py | 84 ++- litellm/llms/custom_httpx/llm_http_handler.py | 74 ++- .../llms/manus/responses/transformation.py | 4 +- .../llms/openai/responses/transformation.py | 4 +- .../volcengine/responses/transformation.py | 4 +- litellm/proxy/_lazy_openapi_snapshot.json | 48 +- litellm/proxy/common_request_processing.py | 2 +- litellm/proxy/hooks/responses_id_security.py | 9 +- .../proxy/response_api_endpoints/endpoints.py | 118 +++- .../streaming_iterator.py | 3 + litellm/responses/main.py | 98 +++- litellm/responses/retrieve_stream_query.py | 83 +++ litellm/responses/streaming_iterator.py | 18 +- litellm/router.py | 12 +- litellm/utils.py | 11 +- .../test_responses_lifecycle_wire.py | 251 ++++++++- .../test_background_response_billed_once.py | 43 +- .../test_internal_call_metadata.py | 14 + .../common_utils/test_check_responses_cost.py | 7 +- .../response_api_endpoints/test_endpoints.py | 209 ++++++- .../test_streaming_iterator_transformation.py | 31 +- .../responses/test_responses_api_lifecycle.py | 515 ++++++++++++++++++ .../test_responses_streaming_iterator.py | 28 + tests/unit/test_responses_id_security.py | 27 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 45 +- 29 files changed, 1672 insertions(+), 95 deletions(-) create mode 100644 litellm/responses/retrieve_stream_query.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 07048cf2c2f..7baef40e86b 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -30,7 +30,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -290,12 +290,14 @@ class CheckResponsesCost: sees provider env vars, so it fails for every deployment whose credentials live in the config; the row then never leaves ``queued``. """ - if not via_router: - return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) - router_response = await self.llm_router.aget_responses( - response_id=response_id, litellm_metadata=litellm_metadata + retrieved: Final[object] = ( + await self.llm_router.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) + if via_router + else await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) ) - return cast(ResponsesAPIResponse, router_response) + if not isinstance(retrieved, ResponsesAPIResponse): + raise TypeError(f"Retrieving {response_id} did not return a response object") + return retrieved async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse: return await self._get_response( diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 1200d799018..f3dbd612bee 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -120,7 +120,7 @@ def is_unbilled_non_inference_call( origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) if metadata is not None else None if origin == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: return False - if not is_background_response(response): + if not is_background_response(_finished_response(response)): return True return origin == BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN or _is_managed_background_call(metadata) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index eb6d103efb9..09777da8669 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -335,7 +335,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data @@ -343,7 +343,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - GET /v1/responses/{response_id} """ get_url: Final = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id) - data: Final[dict] = {} + data: Final[dict[str, str]] = {} verbose_logger.debug("get response url=%s", get_url) return get_url, data diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 7f45c1eba28..792b14090b8 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -69,6 +69,9 @@ class BaseResponsesAPIConfig(ABC): def supports_encrypted_agent_messages(self) -> bool: return False + def supports_streamed_retrieval(self) -> bool: + return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI) + def sign_request( self, headers: dict, @@ -209,7 +212,7 @@ class BaseResponsesAPIConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: pass @abstractmethod diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 0b64885df45..54adfb5c692 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -764,18 +764,23 @@ class AsyncHTTPHandler: follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, max_response_bytes: int | None = None, + stream: bool = False, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT + query: Final = httpx.QueryParams(params).merge(HTTPHandler.extract_query_params(url)) + request_headers: Final = httpx.Headers(headers) - params = params or {} - params.update(HTTPHandler.extract_query_params(url)) + if stream: + return await self._send_streaming_get( + url, params=query, headers=request_headers, follow_redirects=follow_redirects, timeout=timeout + ) if max_response_bytes is not None: return await self._get_with_response_limit( url, - params=httpx.QueryParams(params), - headers=httpx.Headers(headers), + params=query, + headers=request_headers, max_bytes=max_response_bytes, follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects, timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout), @@ -783,13 +788,41 @@ class AsyncHTTPHandler: response: Final = await self.client.get( url, - params=params, - headers=headers, + params=query, + headers=request_headers, follow_redirects=_follow_redirects, timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT, ) return response + async def _send_streaming_get( + self, + url: str, + *, + params: httpx.QueryParams, + headers: httpx.Headers, + follow_redirects: bool | None, + timeout: float | httpx.Timeout | None, + ) -> httpx.Response: + start_time: Final = time.time() + request: Final = self.client.build_request( + "GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout + ) + try: + response: Final = await self.client.send( + request, + stream=True, + follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, + ) + _anchor_handler_to(response, self) + response.raise_for_status() + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) + except httpx.HTTPStatusError as e: + await _raise_masked_async_error(e, True) + else: + return response + async def _get_with_response_limit( self, url: str, @@ -1468,15 +1501,20 @@ class HTTPHandler: headers: Mapping[str, Any] | None = None, follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, + stream: bool = False, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT - params = params or {} - params.update(self.extract_query_params(url)) + query: Final = httpx.QueryParams(params).merge(self.extract_query_params(url)) + + if stream: + return self._send_streaming_get( + url, params=query, headers=headers, follow_redirects=follow_redirects, timeout=timeout + ) response: Final = self.client.get( url, - params=params, + params=query, headers=headers, follow_redirects=_follow_redirects, timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT, @@ -1484,6 +1522,34 @@ class HTTPHandler: return response + def _send_streaming_get( + self, + url: str, + *, + params: httpx.QueryParams, + headers: Mapping[str, str] | None, + follow_redirects: bool | None, + timeout: float | httpx.Timeout | None, + ) -> httpx.Response: + start_time: Final = time.time() + request: Final = self.client.build_request( + "GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout + ) + try: + response: Final = self.client.send( + request, + stream=True, + follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, + ) + _anchor_handler_to(response, self) + response.raise_for_status() + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) + except httpx.HTTPStatusError as e: + _raise_masked_sync_error(e, True) + else: + return response + @staticmethod def extract_query_params(url: str) -> dict[str, str]: """ diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dc10eff65b8..9700465ebc3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -386,6 +386,18 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) - return {"json": data} +def _answered_json(response: httpx.Response) -> bool: + return "application/json" in response.headers.get("content-type", "").lower() + + +def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]: + if not stream: + return {} + if starting_after is None: + return {"stream": "true"} + return {"stream": "true", "starting_after": starting_after} + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -3524,12 +3536,20 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, object] | None = None, + stream: bool = False, + starting_after: int | None = None, + litellm_metadata: Mapping[str, object] = MappingProxyType({}), timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: + ) -> ( + ResponsesAPIResponse + | SyncResponsesAPIStreamingIterator + | Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator] + ): """ Get a response by ID Uses GET /v1/responses/{response_id} endpoint in the responses API @@ -3542,7 +3562,11 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, + extra_query=extra_query, extra_body=extra_body, + stream=stream, + starting_after=starting_after, + litellm_metadata=litellm_metadata, timeout=timeout, client=client, shared_session=shared_session, @@ -3571,20 +3595,23 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)} ## LOGGING logging_obj.pre_call( input="", api_key="", additional_args={ - "complete_input_dict": data, + "complete_input_dict": query, "api_base": api_base, "headers": headers, }, ) try: - response: Final = sync_httpx_client.get(url=url, headers=headers, params=data) + response: Final = sync_httpx_client.get( + url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None + ) response.raise_for_status() except Exception as e: raise self._handle_error( @@ -3592,6 +3619,20 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if stream and not _answered_json(response): + return SyncResponsesAPIStreamingIterator( + response=response, + model="", + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=dict(litellm_metadata), + custom_llm_provider=custom_llm_provider, + request_data={"litellm_params": dict(litellm_params)}, + replays_stored_stream=True, + ) + + if stream: + response.read() return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, @@ -3605,11 +3646,15 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, object] | None = None, + stream: bool = False, + starting_after: int | None = None, + litellm_metadata: Mapping[str, object] = MappingProxyType({}), timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, - ) -> ResponsesAPIResponse: + ) -> ResponsesAPIResponse | ResponsesAPIStreamingIterator: """ Async version of get_responses """ @@ -3647,20 +3692,23 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)} ## LOGGING logging_obj.pre_call( input="", api_key="", additional_args={ - "complete_input_dict": data, + "complete_input_dict": query, "api_base": api_base, "headers": headers, }, ) try: - response: Final = await async_httpx_client.get(url=url, headers=headers, params=data) + response: Final = await async_httpx_client.get( + url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None + ) response.raise_for_status() except Exception as e: verbose_logger.debug("Error retrieving response: %s", e) @@ -3669,6 +3717,20 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if stream and not _answered_json(response): + return ResponsesAPIStreamingIterator( + response=response, + model="", + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=dict(litellm_metadata), + custom_llm_provider=custom_llm_provider, + request_data={"litellm_params": dict(litellm_params)}, + replays_stored_stream=True, + ) + + if stream: + await response.aread() return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index ed06391a260..91832899141 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -241,7 +241,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data. @@ -252,7 +252,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 45fc1fc3ff7..be40220f94d 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1094,7 +1094,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data @@ -1103,7 +1103,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 4076ed06fc2..7f2425d0909 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -313,10 +313,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 39adb99a0d0..01aa69ba5d2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -33639,7 +33639,7 @@ ] }, "get": { - "description": "Get a response by ID.\n\nSupports both:\n- Polling IDs (litellm_poll_*): Returns cumulative cached content from background responses\n- Provider response IDs: Passes through to provider API\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/get\n\n```bash\n# Get polling response\ncurl -X GET http://localhost:4000/v1/responses/litellm_poll_abc123 -H \"Authorization: Bearer $LITELLM_MASTER_KEY\"\n\n# Get provider response\ncurl -X GET http://localhost:4000/v1/responses/resp_abc123 -H \"Authorization: Bearer $LITELLM_MASTER_KEY\"\n```", + "description": "Get a response by ID.\n\nSupports both:\n- Polling IDs (litellm_poll_*): Returns cumulative cached content from background responses\n- Provider response IDs: Passes through to provider API\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/get\n\n```bash\n# Get polling response\ncurl -X GET http://localhost:4000/v1/responses/litellm_poll_abc123 -H \"Authorization: Bearer $LITELLM_MASTER_KEY\"\n\n# Get provider response\ncurl -X GET http://localhost:4000/v1/responses/resp_abc123 -H \"Authorization: Bearer $LITELLM_MASTER_KEY\"\n\n# Resume a background response's event stream after sequence number 7\ncurl -N \"http://localhost:4000/v1/responses/resp_abc123?stream=true&starting_after=7\" -H \"Authorization: Bearer $LITELLM_MASTER_KEY\"\n```", "operationId": "get_response_openai_v1_responses__response_id__get", "parameters": [ { @@ -33650,6 +33650,46 @@ "title": "Response Id", "type": "string" } + }, + { + "description": "Replay the response's events as server-sent events (OpenAI, Azure, and Azure AI Foundry)", + "in": "query", + "name": "stream", + "required": false, + "schema": { + "type": "boolean" + } + }, + { + "description": "With stream=true, replay only the events after this sequence number", + "in": "query", + "name": "starting_after", + "required": false, + "schema": { + "minimum": 0, + "type": "integer" + } + }, + { + "description": "Additional output data to include, forwarded to the provider", + "in": "query", + "name": "include", + "required": false, + "schema": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + { + "description": "With stream=true, whether delta events carry obfuscation padding", + "in": "query", + "name": "include_obfuscation", + "required": false, + "schema": { + "type": "boolean" + } } ], "responses": { @@ -33659,6 +33699,12 @@ "schema": { "$ref": "#/components/schemas/ResponsesAPIResponse" } + }, + "text/event-stream": { + "schema": { + "description": "Server sent events when stream=true", + "type": "string" + } } }, "description": "Successful Response" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6877d6c0118..fc29a6e3a08 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3206,7 +3206,7 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, request=request, ) - if route_type == "aresponses": + if route_type in ("aresponses", "aget_responses"): selected_data_generator = ( ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( original_stream_response=response, diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 27b5924c540..771d68a0ea0 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" -_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_RESPONSES_API_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingProxyType( { "aresponses": "previous_response_id", @@ -104,7 +104,7 @@ def _rewrite_advertised_id( return event -def _is_responses_api_create_route(request_route: str | None) -> bool: +def _streams_response_ids(request_route: str | None) -> bool: if request_route is None: return False canonical: Final = ( @@ -112,7 +112,8 @@ def _is_responses_api_create_route(request_route: str | None) -> bool: if request_route.startswith(_RESPONSES_API_PROVIDER_PREFIX + "/") else request_route ) - return canonical in _RESPONSES_API_CREATE_ROUTES + parent, _, response_id = canonical.rpartition("/") + return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "") class ResponsesIDSecurity(CustomLogger): @@ -339,7 +340,7 @@ class ResponsesIDSecurity(CustomLogger): async for chunk in response: if ( isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and _is_responses_api_create_route(user_api_key_dict.request_route) + and _streams_response_ids(user_api_key_dict.request_route) and not general_settings.get("disable_responses_id_security", False) ): chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5737421d4d4..9e642a02d49 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -86,6 +86,7 @@ from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD +from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel @@ -106,11 +107,43 @@ router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) _DECODED_JSON: Final = TypeAdapter(object) _PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_RETRIEVE_QUERY_PARAMS: Final = tuple(sorted(RETRIEVE_INCLUDE_QUERY_KEYS | RETRIEVE_STREAM_QUERY_KEYS)) +_RETRIEVE_QUERY_OPENAPI: Final = { + "parameters": [ + { + "name": "stream", + "in": "query", + "required": False, + "schema": {"type": "boolean"}, + "description": "Replay the response's events as server-sent events (OpenAI, Azure, and Azure AI Foundry)", + }, + { + "name": "starting_after", + "in": "query", + "required": False, + "schema": {"type": "integer", "minimum": 0}, + "description": "With stream=true, replay only the events after this sequence number", + }, + { + "name": "include", + "in": "query", + "required": False, + "schema": {"type": "array", "items": {"type": "string"}}, + "description": "Additional output data to include, forwarded to the provider", + }, + { + "name": "include_obfuscation", + "in": "query", + "required": False, + "schema": {"type": "boolean"}, + "description": "With stream=true, whether delta events carry obfuscation padding", + }, + ] +} _ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg -RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {200: {"model": ResponsesAPIResponse}} -RESPONSES_API_CREATE_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { +RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { 200: { "model": ResponsesAPIResponse, "content": { @@ -254,19 +287,19 @@ async def _resolve_cursor_model_variant_before_auth(request: Request) -> None: "/v1/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) @router.post( "/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) @router.post( "/openai/v1/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) async def responses_api( request: Request, @@ -749,18 +782,21 @@ async def cursor_chat_completions( dependencies=[Depends(user_api_key_auth)], tags=["responses"], responses=RESPONSES_API_RESPONSE_SCHEMAS, + openapi_extra=_RETRIEVE_QUERY_OPENAPI, ) @router.get( "/responses/{response_id}", dependencies=[Depends(user_api_key_auth)], tags=["responses"], responses=RESPONSES_API_RESPONSE_SCHEMAS, + openapi_extra=_RETRIEVE_QUERY_OPENAPI, ) @router.get( "/openai/v1/responses/{response_id}", dependencies=[Depends(user_api_key_auth)], tags=["responses"], responses=RESPONSES_API_RESPONSE_SCHEMAS, + openapi_extra=_RETRIEVE_QUERY_OPENAPI, ) async def get_response( response_id: str, @@ -785,6 +821,10 @@ async def get_response( # Get provider response curl -X GET http://localhost:4000/v1/responses/resp_abc123 \ -H "Authorization: Bearer $LITELLM_MASTER_KEY" + + # Resume a background response's event stream after sequence number 7 + curl -N "http://localhost:4000/v1/responses/resp_abc123?stream=true&starting_after=7" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" ``` """ from litellm.proxy.proxy_server import ( @@ -832,6 +872,14 @@ async def get_response( data: Final = await read_request_body(request=request) read_settles_billing: Final = not carries_client_side_overrides(data) data["response_id"] = response_id + retrieve_query: Final = { + key: values[0] if len(values) == 1 else values + for key in _RETRIEVE_QUERY_PARAMS + if (values := request.query_params.getlist(key)) + } + if retrieve_query: + body_extra_query: Final[object] = data.get("extra_query") + data["extra_query"] = {**(body_extra_query if isinstance(body_extra_query, dict) else {}), **retrieve_query} processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: response: Final[object] = await processor.base_process_llm_request( @@ -843,7 +891,12 @@ async def get_response( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=select_data_generator, + select_data_generator=partial( + _streamed_retrieval_data_generator, + partial(select_data_generator, responses_stream_errors=True), + response_id, + read_settles_billing, + ), model=None, user_model=user_model, user_temperature=user_temperature, @@ -917,12 +970,24 @@ async def _bill_finished_background_response( fastapi_response: Response, read_settles_billing: bool, ) -> None: + _report_charge( + await _settle_finished_background_read(response, response_id, user_api_key_dict, read_settles_billing), + fastapi_response, + ) + + +async def _settle_finished_background_read( + response: object, + response_id: str, + user_api_key_dict: UserAPIKeyAuth, + read_settles_billing: bool, +) -> "_BilledCharge | None": from litellm.proxy.proxy_server import proxy_logging_obj if not isinstance(response, ResponsesAPIResponse) or not is_background_response(response): - return + return None if response.status not in TERMINAL_RESPONSE_STATUSES: - return + return None managed_files_hook: Final = _managed_files_hook(proxy_logging_obj) track_untracked: Final = ( partial( @@ -933,12 +998,11 @@ async def _bill_finished_background_response( if managed_files_hook is not None else None ) - charge: Final = await _settle_background_billing( + return await _settle_background_billing( (lambda cost_check: cost_check.bill_if_finished(response_id, response, track_untracked)) if read_settles_billing else (lambda cost_check: cost_check.bill_from_deployment(response_id, track_untracked)) ) - _report_charge(charge, fastapi_response) class _BilledCharge(Protocol): @@ -1088,27 +1152,53 @@ def _streamed_background_data_generator( user_api_key_dict=user_api_key_dict, request_tags=_background_response_spend_tags(request_data), ) - return _bill_streamed_background_response(response, chunks, track_untracked) + return _bill_streamed_background_response( + response, chunks, partial(_bill_finished_stream, track_untracked=track_untracked) + ) + + +def _streamed_retrieval_data_generator( + select_data_generator: Callable[..., AsyncIterator[str | bytes]], + response_id: str, + read_settles_billing: bool, + response: object, + user_api_key_dict: UserAPIKeyAuth, + request_data: Mapping[str, object], + request: Request | None = None, +) -> AsyncIterator[str | bytes]: + chunks: Final = select_data_generator( + response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request + ) + return _bill_streamed_background_response( + response, + chunks, + partial( + _settle_finished_background_read, + response_id=response_id, + user_api_key_dict=user_api_key_dict, + read_settles_billing=read_settles_billing, + ), + ) async def _bill_streamed_background_response( stream: object, chunks: AsyncIterator[str | bytes], - track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]], + bill: Callable[[ResponsesAPIResponse], Awaitable[object]], ) -> AsyncIterator[str | bytes]: async for chunk in chunks: completed: object = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) if not isinstance(completed, ResponsesAPIResponse): yield chunk continue - await _bill_finished_stream(completed, track_untracked) + await bill(completed) yield chunk async for remaining_chunk in chunks: yield remaining_chunk return late_completed: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) if isinstance(late_completed, ResponsesAPIResponse): - await _bill_finished_stream(late_completed, track_untracked) + await bill(late_completed) async def _bill_finished_stream( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 9f0058a5958..892da541935 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -196,6 +196,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_search = tool_search_routing(request_tools) self._tool_search_executions: Mapping[str, ToolSearchExecution] = MappingProxyType({}) + async def aclose(self) -> None: + await self.litellm_custom_stream_wrapper.aclose() + def _tool_search_execution(self, call_id: str, fn_name: str) -> ToolSearchExecution | None: known: Final = self._tool_search_executions.get(call_id) if known is not None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index d474ed3991f..5bed6037fa4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -106,7 +106,17 @@ if TYPE_CHECKING: else: MCPTool = Any -from .streaming_iterator import BaseResponsesAPIStreamingIterator +from .retrieve_stream_query import ( + RETRIEVE_INCLUDE_QUERY_KEYS, + InvalidRetrieveStreamQuery, + RetrieveStreamQuery, + parse_retrieve_stream_query, +) +from .streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) __all__ = ( "acancel_responses", @@ -192,6 +202,7 @@ def mock_responses_api_response( _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_STR_DICT: Final = TypeAdapter(dict[str, object]) _FORWARDED_VALUE: Final = TypeAdapter(object) @@ -2013,6 +2024,53 @@ def delete_responses( ) +def _requested_retrieve_stream_query( + stream: bool | None, + starting_after: int | None, + extra_query: Mapping[str, object] | None, + custom_llm_provider: str | None, + drop_malformed: bool, +) -> RetrieveStreamQuery: + explicit: Final = { + key: value for key, value in (("stream", stream), ("starting_after", starting_after)) if value is not None + } + requested: Final = {**explicit, **(extra_query or {})} + parsed: Final = parse_retrieve_stream_query( + requested.get("stream"), requested.get("starting_after"), drop_malformed=drop_malformed + ) + if isinstance(parsed, InvalidRetrieveStreamQuery): + raise litellm.BadRequestError( + message=parsed.message, + model=None, + llm_provider=custom_llm_provider, + body={"message": parsed.message, "type": "invalid_request_error", "param": parsed.param}, + ) + return parsed + + +def _supported_retrieve_stream_query( + requested: RetrieveStreamQuery, + responses_api_provider_config: BaseResponsesAPIConfig, + custom_llm_provider: str, + drop_params: object, +) -> RetrieveStreamQuery: + if not requested.stream: + return RetrieveStreamQuery(stream=False, starting_after=None) + if responses_api_provider_config.supports_streamed_retrieval(): + return requested + if drops_unsupported_params(drop_params): + return RetrieveStreamQuery(stream=False, starting_after=None) + raise UnsupportedParamsError( + status_code=400, + llm_provider=custom_llm_provider, + message=( + f"{custom_llm_provider} does not support streaming a stored response (stream=true on " + "GET /v1/responses/{response_id}). To drop it, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ), + ) + + @client async def aget_responses( response_id: str, @@ -2024,8 +2082,10 @@ async def aget_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + stream: bool | None = None, + starting_after: int | None = None, **kwargs, -) -> ResponsesAPIResponse: +) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: """ Async: Fetch a response by its ID. @@ -2058,6 +2118,8 @@ async def aget_responses( extra_query=extra_query, extra_body=extra_body, timeout=timeout, + stream=stream, + starting_after=starting_after, **kwargs, ) @@ -2071,7 +2133,7 @@ async def aget_responses( response = init_response # Update the responses_api_response_id with the model_id - if not isinstance(response, ResponsesAPIResponse): # pyright: ignore[reportUnnecessaryIsInstance] # handlers can return non-ResponsesAPIResponse objects at runtime + if not isinstance(response, ResponsesAPIResponse): return response return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -2101,8 +2163,14 @@ def get_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + stream: bool | None = None, + starting_after: int | None = None, **kwargs, -) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: +) -> ( + ResponsesAPIResponse + | SyncResponsesAPIStreamingIterator + | Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator] +): """ Fetch a response by its ID. @@ -2146,6 +2214,20 @@ def get_responses( if responses_api_provider_config is None: raise ValueError(f"GET responses is not supported for {custom_llm_provider}") + stream_query: Final = _supported_retrieve_stream_query( + requested=_requested_retrieve_stream_query( + stream=stream, + starting_after=starting_after, + extra_query=extra_query, + custom_llm_provider=custom_llm_provider, + drop_malformed=drops_unsupported_params(litellm_params.drop_params), + ), + responses_api_provider_config=responses_api_provider_config, + custom_llm_provider=custom_llm_provider, + drop_params=litellm_params.drop_params, + ) + litellm_metadata: Final = _STR_DICT.validate_python(kwargs.get("litellm_metadata") or {}) + local_vars.update(kwargs) # Pre Call logging @@ -2169,7 +2251,13 @@ def get_responses( litellm_params=litellm_params, logging_obj=litellm_logging_obj, extra_headers=extra_headers, + extra_query={ + key: value for key, value in (extra_query or {}).items() if key in RETRIEVE_INCLUDE_QUERY_KEYS + }, extra_body=extra_body, + stream=stream_query.stream, + starting_after=stream_query.starting_after, + litellm_metadata=litellm_metadata, timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -2180,7 +2268,7 @@ def get_responses( if isinstance(response, ResponsesAPIResponse): response = ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, - litellm_metadata=kwargs.get("litellm_metadata", {}), + litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/responses/retrieve_stream_query.py b/litellm/responses/retrieve_stream_query.py new file mode 100644 index 00000000000..086aa39efc6 --- /dev/null +++ b/litellm/responses/retrieve_stream_query.py @@ -0,0 +1,83 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"}) +RETRIEVE_INCLUDE_QUERY_KEYS: Final = frozenset({"include", "include[]", "include_obfuscation"}) +_QUERY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_MAX_STARTING_AFTER_DIGITS: Final = len(str(2**63 - 1)) + + +@dataclass(frozen=True, slots=True) +class RetrieveStreamQuery: + stream: bool + starting_after: int | None + + +@dataclass(frozen=True, slots=True) +class InvalidRetrieveStreamQuery: + param: Literal["stream", "starting_after"] + message: str + + +_NO_STREAM: Final = RetrieveStreamQuery(stream=False, starting_after=None) +_FULL_REPLAY: Final = RetrieveStreamQuery(stream=True, starting_after=None) + + +def _parsed_stream(value: object) -> bool | InvalidRetrieveStreamQuery: + match value: + case None: + return False + case bool(): + return value + case str() if value.lower() in ("true", "false"): + return value.lower() == "true" + case _: + return InvalidRetrieveStreamQuery("stream", f"Invalid 'stream': expected true or false, got {value!r}") + + +def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQuery: + match value: + case None: + return None + case int() if not isinstance(value, bool) and value >= 0: + return value + case str() if value.isascii() and value.isdigit() and len(value) <= _MAX_STARTING_AFTER_DIGITS: + return int(value) + case _: + return InvalidRetrieveStreamQuery( + "starting_after", f"Invalid 'starting_after': expected an integer >= 0, got {value!r}" + ) + + +def parse_retrieve_stream_query( + stream: object, starting_after: object, drop_malformed: bool +) -> RetrieveStreamQuery | InvalidRetrieveStreamQuery: + parsed_stream: Final = _parsed_stream(stream) + if isinstance(parsed_stream, InvalidRetrieveStreamQuery): + return _NO_STREAM if drop_malformed else parsed_stream + if not parsed_stream: + return _NO_STREAM + parsed_starting_after: Final = _parsed_starting_after(starting_after) + if isinstance(parsed_starting_after, InvalidRetrieveStreamQuery): + return _FULL_REPLAY if drop_malformed else parsed_starting_after + return RetrieveStreamQuery(stream=True, starting_after=parsed_starting_after) + + +def requests_streamed_retrieval(request: Mapping[str, object]) -> bool: + extra_query: Final = request.get("extra_query") + return ( + "response_id" in request + and isinstance(extra_query, Mapping) + and _parsed_stream(_QUERY_ADAPTER.validate_python(extra_query).get("stream")) is True + ) + + +def without_retrieve_stream_defaults( + deployment_params: Mapping[str, object], request: Mapping[str, object] +) -> Mapping[str, object]: + if "response_id" not in request: + return deployment_params + return {key: value for key, value in deployment_params.items() if key not in RETRIEVE_STREAM_QUERY_KEYS} diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6c57bfd2a3f..8959c32fff2 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -376,6 +376,7 @@ class BaseResponsesAPIStreamingIterator: _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () _started_tool_call: bool = False _earlier_rounds_spend: BilledRoundsSpend | None = None + _replays_stored_stream: bool = False def __init__( self, @@ -574,7 +575,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None) - _estimate_wanted: Final[bool] = _chunk_type in ( + _estimate_wanted: Final[bool] = not self._replays_stored_stream and _chunk_type in ( openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, ) @@ -995,6 +996,7 @@ class BaseResponsesAPIStreamingIterator: generated_text: Final = self._generated_content + self._generated_tool_arguments if ( not generated_text + or self._replays_stored_stream or requests_background(self.request_data) or is_background_response(self._created_response) or self._completed_response_logged @@ -1110,7 +1112,7 @@ class BaseResponsesAPIStreamingIterator: """ created: Final = self._created_response generated: Final = self._generated_content + self._generated_tool_arguments - if created is None or not generated or self._billed_without_disconnect(): + if self._replays_stored_stream or created is None or not generated or self._billed_without_disconnect(): return None usage: Final = _estimate_usage_safely( self.model or "", @@ -1162,6 +1164,8 @@ class BaseResponsesAPIStreamingIterator: def _raise_if_ended_without_terminal_event(self) -> None: if self.completed_response is not None: return + if self._replays_stored_stream and not self._yielded_first_chunk: + return error: Final = self._stream_ended_early_error() self._handle_failure(error) if self._output_started: @@ -1209,6 +1213,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, + replays_stored_stream: bool = False, ): super().__init__( response, @@ -1220,6 +1225,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) + self._replays_stored_stream = replays_stored_stream self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes()) def __aiter__(self): @@ -1232,6 +1238,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): await self._await_pending_logging() raise + async def aclose(self) -> None: + await self.response.aclose() + async def _next_event(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() @@ -1304,6 +1313,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, + replays_stored_stream: bool = False, ): super().__init__( response, @@ -1315,11 +1325,15 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) + self._replays_stored_stream = replays_stored_stream self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes()) def __iter__(self): return self + def close(self) -> None: + self.response.close() + def __next__(self): try: self._check_max_streaming_duration() diff --git a/litellm/router.py b/litellm/router.py index 02ad6275b82..3809a5d6555 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -184,6 +184,7 @@ from litellm.llms.openai_like.model_info import ( get_openai_compatible_model_info, ) from litellm.responses.bridge_check import keep_registered_mode +from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults from litellm.router_strategy.auto_router.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, STRATEGY_ROUTER_PARAM_FIELDS, @@ -4860,7 +4861,7 @@ class Router: Handles inserting this as either "metadata" or "litellm_metadata" depending on the metadata_variable_name """ # 1) copy your defaults and pull out metadata - defaults: Final = self.default_litellm_params.copy() + defaults: Final = dict(without_retrieve_stream_defaults(self.default_litellm_params, kwargs)) metadata_defaults: Final = defaults.pop("metadata", {}) or {} # 2) add any non-metadata defaults that aren't already in kwargs @@ -5101,12 +5102,13 @@ class Router: router_timeout=_router_timeout, router_stream_timeout=_router_stream_timeout, ) + timeout_request: Final = {**kwargs, "stream": True} if requests_streamed_retrieval(kwargs) else kwargs kwargs.update( CompletionTimeout.stream_request_kwargs( - passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=kwargs) + passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=timeout_request) ) - if splits_stream_timeout and kwargs.get("stream") - else {"timeout": passthrough_timeout(kwargs=kwargs)} + if splits_stream_timeout and timeout_request.get("stream") + else {"timeout": passthrough_timeout(kwargs=timeout_request)} ) elif splits_stream_timeout: request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) @@ -6440,7 +6442,7 @@ class Router: generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") response_kwargs: Final = { - **data, + **without_retrieve_stream_defaults(data, kwargs), "caching": self.cache_responses, **kwargs, "model": model_name, diff --git a/litellm/utils.py b/litellm/utils.py index 6a9beeec916..26ce38d63b3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2032,16 +2032,15 @@ def client(original_function): # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() - if _is_streaming_request( - kwargs=kwargs, - call_type=call_type, - ): + if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result): if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True # RETURN RESULT update_response_metadata: _ResponseMetadataUpdater = litellm_utils.update_response_metadata update_response_metadata( @@ -2578,6 +2577,8 @@ _NEVER_STREAMING_CALL_TYPES: Final = frozenset( } ) +_RESULT_DECIDES_STREAMING_CALL_TYPES: Final = frozenset({"get_responses", "aget_responses"}) + def _is_streaming_request( kwargs: dict[str, object], @@ -2589,7 +2590,7 @@ def _is_streaming_request( - if "stream=True" in kwargs (litellm chat completion, litellm text completion, litellm messages) - if call_type is generate_content_stream or agenerate_content_stream (litellm google genai) """ - if call_type in _NEVER_STREAMING_CALL_TYPES: + if call_type in _NEVER_STREAMING_CALL_TYPES or call_type in _RESULT_DECIDES_STREAMING_CALL_TYPES: return False if "stream" in kwargs and kwargs["stream"] is True: return True diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 6509c8ced96..02269885b44 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -2,17 +2,21 @@ from __future__ import annotations import json import os +import threading import uuid +from collections.abc import Callable +from types import MappingProxyType from typing import Final, Literal from urllib.parse import parse_qsl, urlsplit +import httpx import openai import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness -from openai.types.responses import Response, ResponseStreamEvent +from openai.types.responses import Response, ResponseCreatedEvent, ResponseStreamEvent from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter _MODEL: Final = "gpt-5" @@ -87,11 +91,14 @@ def _stream(identity: str) -> tuple[bytes, ...]: ) + (b"data: [DONE]\n\n",) -def _deployment(scenario: Scenario, wire: Wire, api_key: str = _API_KEY) -> str: +def _deployment( + scenario: Scenario, wire: Wire, api_key: str = _API_KEY, provider: str = "openai", **parameters: JsonValue +) -> str: model: Final = scenario.model( - model=f"openai/{_MODEL}", + model=f"{provider}/{_MODEL}", api_key=api_key, api_base=f"{wire.url}/v1", + **parameters, ) _wait_for_workers(scenario.gateway, model, 1) return model @@ -297,9 +304,6 @@ def test_streamed_response_ids_are_stable_and_retrievable_on_each_alias( @pytest.mark.parametrize("prefix", _ALIASES) def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> None: - pytest.skip( - "BUG: GET /v1/responses/{id} drops include and starting_after; streamed GET returns 500 for provider SSE" - ) normal_raw_id: Final = f"resp_query_normal_{uuid.uuid4().hex}" stream_raw_id: Final = f"resp_query_stream_{uuid.uuid4().hex}" normal_include: Final = ["message.output_text.logprobs", "reasoning.encrypted_content"] @@ -413,6 +417,241 @@ def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> N assert actual_queries == expected_queries, requests +def _resumed_delta() -> bytes: + event: Final = { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_resumed", + "output_index": 0, + "content_index": 0, + "delta": "resumed", + "logprobs": [], + } + return f"data: {json.dumps(event)}\n\n".encode() + + +def _sorted_query(request: Request) -> tuple[tuple[str, str], ...]: + return tuple(sorted(parse_qsl(urlsplit(request.target).query, keep_blank_values=True))) + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_past_the_last_event_ends_the_stream_without_events(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_resume_end_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return Reply(content_type="text/event-stream", chunks=()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="resume past the end") + events: Final = tuple(client.responses.retrieve(created.id, stream=True, starting_after=7)) + requests: Final = drain_contract_requests(wire) + + assert events == () + assert [(request.method, urlsplit(request.target).path) for request in requests] == [ + ("POST", "/v1/responses"), + ("GET", f"/v1/responses/{raw_id}"), + ], requests + assert _sorted_query(requests[1]) == (("starting_after", "7"), ("stream", "true")), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_client_leaving_a_resumed_stream_releases_the_provider_connection(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_resume_leave_{uuid.uuid4().hex}" + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return Reply( + content_type="text/event-stream", + chunks=(_resumed_delta(), b"data: [DONE]\n\n"), + gate_after_first=release, + heartbeat=b": keepalive\n\n", + pause_between_chunks=0.05, + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="leave a resumed stream") + stream: Final = client.responses.retrieve(created.id, stream=True, starting_after=3) + try: + first: Final = next(iter(stream)) + stream.close() + left_from: Final = wire.disconnected.get(timeout=30) + finally: + release.set() + + assert first.type == "response.output_text.delta" + assert urlsplit(left_from).path == f"/v1/responses/{raw_id}" + + +def _created_then(raw_id: str, read: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return read + + return respond + + +def _event_types(answer: httpx.Response) -> list[JsonValue]: + return [ + _JSON_OBJECT.validate_json(line.removeprefix("data: "))["type"] + for line in answer.text.splitlines() + if line.startswith("data: {") + ] + + +_MALFORMED_RESUMES: Final = MappingProxyType( + { + "stream": {"stream": "maybe"}, + "starting_after": {"stream": "true", "starting_after": "-1"}, + } +) + + +@pytest.mark.parametrize("param", tuple(_MALFORMED_RESUMES)) +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_malformed_resume_query_is_refused_without_calling_the_provider( + gateway: Gateway, prefix: str, param: str +) -> None: + raw_id: Final = f"resp_malformed_{uuid.uuid4().hex}" + replay: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + + with wire_server(_created_then(raw_id, replay)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="malformed resume") + refused: Final = gateway.request("GET", f"{prefix}/responses/{created.id}", params=_MALFORMED_RESUMES[param]) + requests: Final = drain_contract_requests(wire) + + assert refused.status_code == 400, refused.text + assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["param"] == param, refused.text + assert [request.method for request in requests] == ["POST"], requests + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_malformed_resume_cursor_replays_the_whole_stream_when_drop_params_is_on( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_malformed_dropped_{uuid.uuid4().hex}" + replay: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + + with wire_server(_created_then(raw_id, replay)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, drop_params=True) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="dropped cursor") + replayed: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params=_MALFORMED_RESUMES["starting_after"] + ) + requests: Final = drain_contract_requests(wire) + + assert replayed.status_code == 200, replayed.text + assert _event_types(replayed) == ["response.created", "response.completed"], replayed.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (("stream", "true"),), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_on_a_provider_that_cannot_replay_is_refused_without_calling_it(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_no_replay_{uuid.uuid4().hex}" + + with wire_server(_created_then(raw_id, Reply(body=_response(raw_id)))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, provider="litellm_proxy") + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="cannot replay") + refused: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params={"stream": "true", "starting_after": "3"} + ) + requests: Final = drain_contract_requests(wire) + + assert refused.status_code == 400, refused.text + message: Final = string_value(object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["message"]) + assert "litellm_proxy does not support streaming a stored response" in message, refused.text + assert [request.method for request in requests] == ["POST"], requests + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_on_a_provider_that_cannot_replay_answers_the_stored_response_when_drop_params_is_on( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_no_replay_dropped_{uuid.uuid4().hex}" + + with wire_server(_created_then(raw_id, Reply(body=_response(raw_id)))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, provider="litellm_proxy", drop_params=True) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="dropped replay") + answered: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params={"stream": "true", "starting_after": "3"} + ) + requests: Final = drain_contract_requests(wire) + + assert answered.status_code == 200, answered.text + body: Final = _JSON_OBJECT.validate_json(answered.content) + assert (body["object"], body["status"], same_response(string_value(body["id"]), created.id)) == ( + "response", + "completed", + True, + ), answered.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_resumed_stream_waits_out_the_stream_timeout_not_the_request_timeout(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_slow_replay_{uuid.uuid4().hex}" + slow: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id), pause_between_chunks=2) + + with wire_server(_created_then(raw_id, slow)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, timeout=1, stream_timeout=30) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="slow replay") + events: Final = tuple(client.responses.retrieve(created.id, stream=True, starting_after=3)) + + assert [event.type for event in events] == ["response.created", "response.completed"], events + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_plain_retrieve_on_a_deployment_set_to_stream_answers_the_response_object( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_streaming_deployment_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + return Reply(body=_response(raw_id)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, stream=True) + events: Final = tuple( + _client(gateway, gateway.key, prefix).responses.create(model=model, input="forced stream", stream=True) + ) + created_id: Final = next(event.response.id for event in events if isinstance(event, ResponseCreatedEvent)) + retrieved: Final = gateway.request("GET", f"{prefix}/responses/{created_id}") + requests: Final = drain_contract_requests(wire) + + assert retrieved.status_code == 200, retrieved.text + assert retrieved.headers["content-type"].startswith("application/json"), retrieved.headers + body: Final = _JSON_OBJECT.validate_json(retrieved.content) + assert (body["object"], body["status"], same_response(string_value(body["id"]), created_id)) == ( + "response", + "completed", + True, + ), retrieved.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (), requests[1] + + @pytest.mark.parametrize("prefix", _ALIASES) def test_delete_forwards_the_provider_id_once_and_reports_deletion( gateway: Gateway, diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 5f948fac3ab..00ac917a92e 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor from itertools import chain from typing import Final +from urllib.parse import parse_qsl, urlsplit import httpx import openai @@ -24,7 +25,7 @@ from integration._support.client import ( from integration._support.database import read_rows, write_rows from integration._support.otlp_sink import Span, decode_spans from integration._support.upstream import delete_scenario, register_scenario -from integration._support.wire import Request, wire_server +from integration._support.wire import Reply, Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store from integration.security._callback_traffic import langfuse_sink from integration.security._sinks import Recorder @@ -67,6 +68,7 @@ from integration.spend._background_responses import ( scripted_routes, scripted_upstream, settling_reads, + sse_frame, streamed_reply, submissions, tag_spend, @@ -350,6 +352,45 @@ def test_a_streamed_background_response_read_later_bills_its_stream_alone(gatewa assert_charged_exactly(lambda: key_spend(key), PRICE) +def _finishes_only_when_resumed(raw_id: str) -> Callable[[Request], Reply]: + finished: Final = (sse_frame("response.completed", 9, response_body("completed", id=raw_id)) + "\n\n").encode() + queued: Final = json.dumps(response_body("queued", id=raw_id)).encode() + + def respond(request: Request) -> Reply: + if urlsplit(request.target).path == "/v1/models": + return Reply(body=b'{"object": "list", "data": []}') + if request.method == "GET" and ("stream", "true") in parse_qsl(urlsplit(request.target).query): + return Reply(content_type="text/event-stream", chunks=(finished, b"data: [DONE]\n\n")) + return Reply(body=queued) + + return respond + + +@pytest.mark.timeout(240) +def test_a_background_response_first_seen_finished_on_a_resumed_stream_bills_once(gateway: Gateway) -> None: + raw_id: Final = f"resp_resumed_{uuid.uuid4().hex}" + with wire_server(_finishes_only_when_resumed(raw_id)) as wire, gateway.scenario() as scenario: + key: Final = scenario.key() + model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_key="sk-scripted-provider", + api_base=f"{wire.url}/v1", + input_cost_per_token=INPUT_COST_PER_TOKEN, + output_cost_per_token=OUTPUT_COST_PER_TOKEN, + ) + response_id: Final = create_background_response(scenario, key, model) + resumed: Final = tuple( + gateway.request( + "GET", f"/v1/responses/{response_id}", key=key, params={"stream": "true", "starting_after": "0"} + ) + for _ in range(READS) + ) + assert [(answer.status_code, "response.completed" in answer.text) for answer in resumed] == [ + (200, True) + ] * READS, [answer.text for answer in resumed] + assert_charged_exactly(lambda: key_spend(key), PRICE) + + def _finished(scenario: Scenario, key: str) -> tuple[Upstream, str]: upstream: Final = scripted_upstream(scenario, background_routes("completed")) return upstream, create_background_response(scenario, key, upstream.model) diff --git a/tests/unit/litellm_core_utils/test_internal_call_metadata.py b/tests/unit/litellm_core_utils/test_internal_call_metadata.py index dfb232e96ce..8a279df8d71 100644 --- a/tests/unit/litellm_core_utils/test_internal_call_metadata.py +++ b/tests/unit/litellm_core_utils/test_internal_call_metadata.py @@ -181,6 +181,20 @@ class TestNonInferenceCallBilling: assert is_unbilled_non_inference_call_from_params("aresponses", {}, self.BACKGROUND) is False assert is_unbilled_non_inference_call_from_params("acompletion", {"metadata": self.MANAGED_READ}, None) is False + def test_a_streamed_read_prices_like_a_plain_read_of_the_same_response(self): + from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse + + for stored in (self.BACKGROUND, self.FOREGROUND): + finished = ResponsesAPIResponse(**{**stored, "created_at": 1, "output": []}) + completed_event = ResponseCompletedEvent(type="response.completed", response=finished) + for metadata in (None, self.MANAGED_READ, self.STATUS_POLL): + plain_read = is_unbilled_non_inference_call("aget_responses", metadata, finished) + assert is_unbilled_non_inference_call("aget_responses", metadata, completed_event) is plain_read + assert ( + is_unbilled_non_inference_call("aget_responses", metadata, completed_event.model_dump()) + is plain_read + ) + def test_a_tracked_streamed_background_create_is_free_in_every_shape_its_stream_logs(self): from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index b8d5ddef3d2..19ebc763413 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -1207,11 +1207,8 @@ class TestCheckResponsesCost: return_value=1 ) - mock_response = MagicMock() - mock_response.status = "completed" - with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: - mock_aget.return_value = mock_response + mock_aget.return_value = _terminal_response() await check_responses_cost_instance.check_responses_cost() assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False] @@ -1252,7 +1249,7 @@ class TestCheckResponsesCost: ) with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: - mock_aget.return_value = MagicMock(status="completed") + mock_aget.return_value = _terminal_response() await check_responses_cost_instance.check_responses_cost() (billed,) = billed_logger.billed diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index f1aac24dd60..e3d514b2ea9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -5,6 +5,7 @@ Test for response_api_endpoints/endpoints.py import asyncio import contextlib import dataclasses +import datetime import json import unittest from collections.abc import Callable, Mapping @@ -22,6 +23,7 @@ import litellm from litellm.constants import EMPTY_MAPPING from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors from tests._master_key import MASTER_KEY @@ -975,7 +977,13 @@ def _request_for(method: str, path: str, body: bytes): return {"type": "http.request", "body": body, "more_body": False} return Request( - {"type": "http", "method": method, "path": path, "headers": [(b"content-type", b"application/json")]}, + { + "type": "http", + "method": method, + "path": path, + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, receive, ) @@ -1294,6 +1302,205 @@ async def test_a_streamed_background_create_bills_its_finished_response_once_und assert billed["request_tags"] == ["nightly"] +@pytest.mark.asyncio +async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_sees_it_finish( + monkeypatch: pytest.MonkeyPatch, +): + """Reading a finished background job as a stream settles its bill like a plain read does, before + the terminal event reaches the client, while the replayed stream's own log carries no cost and a + later read of the same job bills nothing more.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import ( + _bill_finished_background_response, + _streamed_retrieval_data_generator, + ) + + spend_events: Final = _capture_spend_events(monkeypatch) + stored: Final = _one_pending_background_response() + finished: Final = _background_response(stored.response_id, "completed", 100) + stream: Final = _replayed_stream(stored.response_id) + + async def replayed_events(**_: object): + async for event in stream: + yield f"data: {event.model_dump_json()}\n\n" + + status_before_each_event: Final[list[str]] = [] + with _proxy_over(stored.prisma_client, stored.proxy_logging, _configured_deployment_router(finished)): + async for _ in _streamed_retrieval_data_generator( + replayed_events, stored.response_id, True, stream, UserAPIKeyAuth(api_key="sk-reader"), {} + ): + status_before_each_event.extend(row.status for row in stored.table.rows) + await _bill_finished_background_response( + response=finished, + response_id=stored.response_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + fastapi_response=FastAPIResponse(), + read_settles_billing=True, + ) + await _drain_spend_events() + + assert status_before_each_event == ["queued", "completed"] + (billed,) = tuple(payload for payload in spend_events.payloads if payload["response_cost"]) + assert billed["prompt_tokens"] == 100 + + +_RESUMED_UPSTREAM: Final = "https://resume.example/v1" + + +def _resumed_events(response_id: str) -> tuple[Mapping[str, object], ...]: + return ( + { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + "logprobs": [], + }, + { + "type": "response.completed", + "sequence_number": 5, + "response": _background_response(response_id, "completed", 100).model_dump(mode="json"), + }, + ) + + +def _replayed_stream(response_id: str) -> ResponsesAPIStreamingIterator: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + wire: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events(response_id) + ) + return ResponsesAPIStreamingIterator( + response=httpx.Response(200, content=wire, headers={"content-type": "text/event-stream"}), + model="gpt-5.6", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="aresponses", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + litellm_call_id="resumed-read", + function_id="resumed-read", + ), + custom_llm_provider="openai", + request_data={"model": "gpt-5.6", "input": "hi", "stream": True}, + replays_stored_stream=True, + ) + + +def _sse_events(text: str) -> tuple[Mapping[str, object], ...]: + frames: Final = (frame for frame in text.split("\n\n") if "data: " in frame and "data: [DONE]" not in frame) + return tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) for frame in frames + ) + + +async def _read_through_proxy( + monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], upstream: httpx.Response +) -> tuple[httpx.Response, tuple[httpx.Request, ...]]: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": {"model": "openai/gpt-5.6", "api_base": _RESUMED_UPSTREAM, "api_key": "fixture-key"}, + "model_info": {"id": "deployment-read"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock(assert_all_called=False) as transport: + route: Final = transport.get(_RESUMED_UPSTREAM + "/responses/resp_read").mock(return_value=upstream) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.get(f"/v1/responses/{_proxy_issued_id('resp_read')}", params=query) + return result, tuple(call.request for call in route.calls) + + +@pytest.mark.asyncio +async def test_a_streamed_read_through_the_proxy_replays_the_events_after_the_cursor( + monkeypatch: pytest.MonkeyPatch, +): + wire: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events("resp_read") + ) + result, upstream_requests = await _read_through_proxy( + monkeypatch, + {"stream": "true", "starting_after": "3"}, + httpx.Response(200, content=wire, headers={"Content-Type": "text/event-stream"}), + ) + events: Final = _sse_events(result.text) + + assert result.status_code == 200, result.text + assert result.headers["content-type"].startswith("text/event-stream") + assert [event["type"] for event in events] == ["response.output_text.delta", "response.completed"] + assert events[-1]["response"]["id"] == _proxy_issued_id("resp_read") + assert [dict(request.url.params) for request in upstream_requests] == [{"stream": "true", "starting_after": "3"}] + + +@pytest.mark.asyncio +async def test_a_streamed_read_cut_before_its_terminal_event_ends_in_a_responses_failure_event( + monkeypatch: pytest.MonkeyPatch, +): + delta: Final = _resumed_events("resp_read")[0] + result, _ = await _read_through_proxy( + monkeypatch, + {"stream": "true", "starting_after": "3"}, + httpx.Response( + 200, + content=f"event: {delta['type']}\ndata: {json.dumps(delta)}\n\n", + headers={"Content-Type": "text/event-stream"}, + ), + ) + events: Final = _sse_events(result.text) + + assert result.status_code == 200, result.text + assert [event["type"] for event in events] == ["response.output_text.delta", "response.failed"] + assert "before any terminal event" in events[-1]["response"]["error"]["message"] + + +@pytest.mark.asyncio +async def test_a_plain_read_through_the_proxy_still_answers_json(monkeypatch: pytest.MonkeyPatch): + result, upstream_requests = await _read_through_proxy( + monkeypatch, + {"starting_after": "3"}, + httpx.Response(200, json=_background_response("resp_read", "in_progress", 100).model_dump(mode="json")), + ) + + assert result.status_code == 200, result.text + assert result.json()["id"] == _proxy_issued_id("resp_read") + assert [dict(request.url.params) for request in upstream_requests] == [{}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query, bad_param", + [ + ({"stream": "maybe"}, "stream"), + ({"stream": "true", "starting_after": "-1"}, "starting_after"), + ({"stream": "true", "starting_after": "abc"}, "starting_after"), + ], +) +async def test_a_malformed_resume_query_through_the_proxy_is_a_400( + monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], bad_param: str +): + result, upstream_requests = await _read_through_proxy(monkeypatch, query, httpx.Response(500)) + + assert result.status_code == 400, result.text + assert bad_param in result.json()["error"]["message"] + assert upstream_requests == () + + @pytest.mark.asyncio async def test_a_read_that_bills_a_response_before_its_stream_ends_leaves_the_stream_nothing_to_bill( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 5fd0512fbab..86773dc7ddf 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -12,7 +12,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve import itertools import json -from collections.abc import Callable +from collections.abc import Callable, Iterable from itertools import accumulate, chain, islice, product from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -60,7 +60,7 @@ def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStrea class _FakeStreamWrapper: - def __init__(self, chunks): + def __init__(self, chunks: Iterable[ModelResponseStream]) -> None: self._chunks = list(chunks) self.logging_obj = MagicMock() @@ -705,6 +705,33 @@ async def test_streaming_events_share_the_chat_completion_response_id(): assert decoded["custom_llm_provider"] == "anthropic" +class _ClosableStreamWrapper(_FakeStreamWrapper): + def __init__(self, chunks: Iterable[ModelResponseStream]) -> None: + super().__init__(chunks) + self.closed = False + + async def aclose(self): + self.closed = True + + +@pytest.mark.asyncio +async def test_aclose_closes_the_bridged_chat_completion_stream(): + wrapper: Final = _ClosableStreamWrapper([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=wrapper, + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + await iterator.__anext__() + await iterator.aclose() + + assert wrapper.closed is True + + def test_sync_streaming_events_share_the_chat_completion_response_id(): iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) diff --git a/tests/unit/responses/test_responses_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 80da62f884a..a7c605f1fe7 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -1,3 +1,8 @@ +import json +import queue +import threading +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass, field from typing import Final, Literal, TypeAlias import httpx @@ -6,6 +11,8 @@ import pytest import respx import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.conftest import httpx_transport @@ -116,6 +123,514 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url( assert str(requests[0].request.url) == expected_url +StreamProvider: TypeAlias = Literal["openai", "azure", "azure_ai"] +_STORED_RESPONSE_ID: Final = "resp_background" +_STREAM_HEADERS: Final = {"content-type": "text/event-stream"} + + +def _stored_response(status: str) -> dict[str, object]: + return { + "id": _STORED_RESPONSE_ID, + "object": "response", + "created_at": 1761600000, + "status": status, + "background": True, + "model": "gpt-5.5", + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + } + + +def _sse(event: dict[str, object]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _delta_event(sequence_number: int) -> dict[str, object]: + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + "logprobs": [], + } + + +def _replay_body(*events: dict[str, object]) -> bytes: + return b"".join(_sse(event) for event in events) + + +_REPLAY_AFTER_3: Final = _replay_body( + _delta_event(4), + {"type": "response.completed", "sequence_number": 5, "response": _stored_response("completed")}, +) + + +def _retrieve_route(mock_router: respx.MockRouter, provider: StreamProvider) -> respx.Route: + match provider: + case "openai": + return mock_router.get(f"https://api.openai.com/v1/responses/{_STORED_RESPONSE_ID}") + case "azure": + return mock_router.get(f"https://example-resource.openai.azure.com/openai/responses/{_STORED_RESPONSE_ID}") + case "azure_ai": + return mock_router.get( + f"https://example-resource.services.ai.azure.com/openai/v1/responses/{_STORED_RESPONSE_ID}" + ) + + +def _retrieve_kwargs(provider: StreamProvider) -> dict[str, object]: + match provider: + case "openai": + return {"custom_llm_provider": "openai", "api_key": "sk-test", "api_base": "https://api.openai.com/v1"} + case "azure": + return { + "custom_llm_provider": "azure", + "api_key": "sk-test", + "api_base": "https://example-resource.openai.azure.com", + "api_version": "2025-03-01-preview", + } + case "azure_ai": + return { + "custom_llm_provider": "azure_ai", + "api_key": "sk-test", + "api_base": "https://example-resource.services.ai.azure.com", + } + + +async def _retrieve(sync_mode: bool, **kwargs: object) -> object: + if sync_mode: + return litellm.get_responses(response_id=_STORED_RESPONSE_ID, **kwargs) + return await litellm.aget_responses(response_id=_STORED_RESPONSE_ID, **kwargs) + + +async def _event_types(stream: object) -> tuple[str, ...]: + if isinstance(stream, SyncResponsesAPIStreamingIterator): + return tuple(event.type for event in stream) + assert isinstance(stream, ResponsesAPIStreamingIterator) + return tuple([event.type async for event in stream]) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", ("openai", "azure", "azure_ai")) +@pytest.mark.parametrize( + ("resume_kwargs", "expected_query"), + ( + ({"stream": True, "starting_after": 3}, {"stream": "true", "starting_after": "3"}), + ({"stream": True, "starting_after": 0}, {"stream": "true", "starting_after": "0"}), + ({"stream": True}, {"stream": "true"}), + ({"extra_query": {"stream": "true", "starting_after": "3"}}, {"stream": "true", "starting_after": "3"}), + ( + {"stream": True, "starting_after": 1, "extra_query": {"starting_after": 3, "include[]": "usage"}}, + {"stream": "true", "starting_after": "3", "include[]": "usage"}, + ), + ( + {"stream": True, "extra_query": {"include": "usage", "limit": "1", "api-version": "2020-01-01"}}, + {"stream": "true", "include": "usage"}, + ), + ), +) +@pytest.mark.asyncio +async def test_streamed_retrieve_replays_the_stored_events_after_the_cursor( + provider: StreamProvider, + sync_mode: bool, + resume_kwargs: dict[str, object], + expected_query: dict[str, str], +) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, provider).mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = await _retrieve(sync_mode, **_retrieve_kwargs(provider), **resume_kwargs) + event_types: Final = await _event_types(stream) + (call,) = route.calls + + sent_query: Final = dict(call.request.url.params) + api_version: Final = {"api-version": "2025-03-01-preview"} if provider == "azure" else {} + assert sent_query == {**expected_query, **api_version} + assert event_types == ("response.output_text.delta", "response.completed") + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("starting_after", (3, -1, "abc")) +@pytest.mark.asyncio +async def test_retrieve_without_stream_still_returns_the_response_object( + sync_mode: bool, starting_after: object +) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve(sync_mode, starting_after=starting_after, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert retrieved.status == "completed" + assert dict(call.request.url.params) == {} + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_streamed_retrieve_answered_with_json_returns_the_response_object(sync_mode: bool) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert retrieved.status == "completed" + assert dict(call.request.url.params) == {"stream": "true", "starting_after": "3"} + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_resuming_at_the_end_of_a_finished_stream_ends_without_events(sync_mode: bool) -> None: + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock(return_value=httpx.Response(200, headers=_STREAM_HEADERS)) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=5, **_retrieve_kwargs("openai")) + event_types: Final = await _event_types(stream) + + assert event_types == () + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_resumed_stream_cut_before_its_terminal_event_raises(sync_mode: bool) -> None: + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_replay_body(_delta_event(4))) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + with pytest.raises(litellm.APIConnectionError, match="before any terminal event"): + await _event_types(stream) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_resumed_stream_relays_the_stored_response_without_inventing_usage(sync_mode: bool) -> None: + stored_without_usage: Final = {key: value for key, value in _stored_response("completed").items() if key != "usage"} + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response( + 200, + headers=_STREAM_HEADERS, + content=_replay_body( + _delta_event(4), + {"type": "response.completed", "sequence_number": 5, "response": stored_without_usage}, + ), + ) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + if isinstance(stream, SyncResponsesAPIStreamingIterator): + events: Final = tuple(stream) + else: + assert isinstance(stream, ResponsesAPIStreamingIterator) + events = tuple([event async for event in stream]) + + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.usage is None + + +@dataclass(frozen=True, slots=True) +class _ClosableReplay(httpx.SyncByteStream, httpx.AsyncByteStream): + body: bytes + closed: threading.Event = field(default_factory=threading.Event) + + def __iter__(self) -> Iterator[bytes]: + yield self.body + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.body + + def close(self) -> None: + self.closed.set() + + async def aclose(self) -> None: + self.closed.set() + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_closing_a_resumed_stream_early_releases_the_upstream_connection(sync_mode: bool) -> None: + upstream: Final = _ClosableReplay(_REPLAY_AFTER_3) + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, stream=upstream) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + if isinstance(stream, SyncResponsesAPIStreamingIterator): + first_sync: Final = next(stream) + stream.close() + first_type: Final = first_sync.type + else: + assert isinstance(stream, ResponsesAPIStreamingIterator) + first_async: Final = await anext(stream) + await stream.aclose() + first_type = first_async.type + + assert first_type == "response.output_text.delta" + assert upstream.closed.is_set() + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_rejected_resume_keeps_the_provider_error(sync_mode: bool) -> None: + error_message: Final = "Response cannot be streamed, it is more than 5 minutes old." + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response( + 400, json={"error": {"message": error_message, "type": "invalid_request_error", "param": None}} + ) + ) + with pytest.raises(openai.BadRequestError, match=error_message) as exc_info: + await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert exc_info.value.status_code == 400 + assert dict(call.request.url.params) == {"stream": "true", "starting_after": "3"} + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_resume_that_times_out_raises_a_timeout(sync_mode: bool) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock(side_effect=httpx.ReadTimeout("slow")) + with pytest.raises(litellm.Timeout): + await _retrieve(sync_mode, stream=True, timeout=1.0, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert dict(call.request.url.params) == {"stream": "true"} + + +@pytest.mark.parametrize( + "resume_kwargs", + ( + {"stream": "maybe"}, + {"stream": True, "starting_after": -1}, + {"extra_query": {"stream": "true", "starting_after": "abc"}}, + {"extra_query": {"stream": "1"}}, + {"extra_query": {"stream": "true", "starting_after": "9" * 5000}}, + ), +) +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call( + sync_mode: bool, resume_kwargs: dict[str, object] +) -> None: + with respx.mock(assert_all_called=False) as mock_router: + route: Final = _retrieve_route(mock_router, "openai") + with pytest.raises(litellm.BadRequestError) as exc_info: + await _retrieve(sync_mode, **_retrieve_kwargs("openai"), **resume_kwargs) + + assert exc_info.value.status_code == 400 + assert route.call_count == 0 + + +@pytest.mark.parametrize( + ("resume_kwargs", "expected_query"), + ( + ({"stream": "maybe", "starting_after": 3}, {}), + ({"extra_query": {"stream": "1"}}, {}), + ({"stream": True, "starting_after": -1}, {"stream": "true"}), + ({"extra_query": {"stream": "true", "starting_after": "abc"}}, {"stream": "true"}), + ({"extra_query": {"stream": "true", "starting_after": "9" * 5000}}, {"stream": "true"}), + ), +) +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_malformed_resume_query_is_dropped_under_drop_params( + sync_mode: bool, resume_kwargs: dict[str, object], expected_query: dict[str, str] +) -> None: + upstream: Final = ( + httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + if expected_query + else httpx.Response(200, json=_stored_response("completed")) + ) + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock(return_value=upstream) + retrieved: Final = await _retrieve(sync_mode, drop_params=True, **_retrieve_kwargs("openai"), **resume_kwargs) + streamed_types: Final = None if isinstance(retrieved, ResponsesAPIResponse) else await _event_types(retrieved) + (call,) = route.calls + + assert dict(call.request.url.params) == expected_query + assert (streamed_types is not None) == bool(expected_query) + + +_UNSUPPORTED_PROVIDERS: Final = ("xai", "manus", "volcengine") + + +def _unsupported_provider_kwargs(provider: str) -> dict[str, object]: + return { + "custom_llm_provider": provider, + "api_key": f"{provider}-test", + "api_base": f"https://{provider}.example/v1", + } + + +def _unsupported_provider_route(mock_router: respx.MockRouter, provider: str) -> respx.Route: + return mock_router.get(url__regex=rf"^https://{provider}\.example/.*/responses/{_STORED_RESPONSE_ID}$") + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS) +@pytest.mark.asyncio +async def test_streamed_retrieve_on_an_unsupported_provider_is_a_400_without_drop_params( + sync_mode: bool, provider: str +) -> None: + with respx.mock(assert_all_called=False) as mock_router: + route: Final = _unsupported_provider_route(mock_router, provider) + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + await _retrieve(sync_mode, stream=True, starting_after=3, **_unsupported_provider_kwargs(provider)) + + assert exc_info.value.status_code == 400 + assert route.call_count == 0 + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS) +@pytest.mark.asyncio +async def test_streamed_retrieve_on_an_unsupported_provider_drops_the_stream_under_drop_params( + sync_mode: bool, provider: str +) -> None: + with respx.mock() as mock_router: + route: Final = _unsupported_provider_route(mock_router, provider).mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve( + sync_mode, stream=True, starting_after=3, drop_params=True, **_unsupported_provider_kwargs(provider) + ) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert dict(call.request.url.params) == {} + + +@pytest.mark.asyncio +async def test_a_routed_resume_requested_in_extra_query_gets_the_deployment_stream_timeout() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "background-model", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": "https://api.openai.com/v1", + "timeout": 2, + "stream_timeout": 60, + }, + } + ] + ) + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = await router.aget_responses( + model="background-model", + response_id=_STORED_RESPONSE_ID, + extra_query={"stream": "true", "starting_after": "3"}, + ) + event_types: Final = await _event_types(stream) + (call,) = route.calls + + assert event_types == ("response.output_text.delta", "response.completed") + assert call.request.extensions["timeout"]["read"] == 60 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_stream", "router_defaults"), + [({"stream": True}, None), ({}, {"stream": True})], + ids=["deployment_forces_stream", "router_default_forces_stream"], +) +async def test_a_routed_plain_retrieve_where_config_forces_stream_answers_the_response_object( + deployment_stream: dict[str, bool], router_defaults: dict[str, bool] | None +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "background-model", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": "https://api.openai.com/v1", + **deployment_stream, + }, + } + ], + default_litellm_params=router_defaults, + ) + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await router.aget_responses(model="background-model", response_id=_STORED_RESPONSE_ID) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert dict(call.request.url.params) == {} + + +@pytest.mark.asyncio +async def test_a_routed_create_on_a_deployment_that_forces_stream_still_streams() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "background-model", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": "https://api.openai.com/v1", + "stream": True, + }, + } + ] + ) + with respx.mock() as mock_router: + route: Final = mock_router.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = await router.aresponses(model="background-model", input="hi", background=True) + event_types: Final = await _event_types(stream) + (call,) = route.calls + + assert event_types == ("response.output_text.delta", "response.completed") + assert json.loads(call.request.content)["stream"] is True + +class _SuccessLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: queue.Queue[object] = queue.Queue() + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.logged.put(response_obj) + + +def test_a_sync_streamed_retrieve_logs_success_with_the_finished_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + success_log: Final = _SuccessLog() + monkeypatch.setattr(litellm, "callbacks", [success_log]) + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = litellm.get_responses( + response_id=_STORED_RESPONSE_ID, stream=True, starting_after=3, **_retrieve_kwargs("openai") + ) + assert isinstance(stream, SyncResponsesAPIStreamingIterator) + consumed: Final = tuple(stream) + + logged: Final = success_log.logged.get(timeout=10) + + assert len(consumed) == 2 + assert isinstance(logged, ResponsesAPIResponse) + assert logged.status == "completed" + + @pytest.mark.parametrize("sync_mode", (True, False)) @pytest.mark.asyncio async def test_cancel_responses_reply_carries_the_id_retrieve_gives_for_the_same_run(sync_mode: bool) -> None: diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index ae688ff9173..6c89cfbfb16 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -277,6 +277,7 @@ async def _failure_payload_after_mid_stream_timeout( request_input: str | list[dict[str, object]], upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None, background_fields: Mapping[str, object] | None = None, + replays_stored_stream: bool = False, ) -> tuple[FailureRecorder, StandardLoggingPayload]: upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas() recorder = FailureRecorder() @@ -297,6 +298,7 @@ async def _failure_payload_after_mid_stream_timeout( logging_obj=logging_obj, custom_llm_provider="openai", request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})}, + replays_stored_stream=replays_stored_stream, ) received = _ChunkRecorder() @@ -355,6 +357,17 @@ async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_r assert payload["response_cost"] == 0 +@pytest.mark.asyncio +async def test_mid_stream_failure_of_a_resumed_background_stream_bills_nothing_since_its_create_bills_the_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, replays_stored_stream=True) + + assert payload["status"] == "failure" + assert payload["completion_tokens"] == 0 + assert payload["response_cost"] == 0 + + @pytest.mark.asyncio async def test_mid_stream_failure_usage_estimate_never_fetches_image_dimensions( monkeypatch: pytest.MonkeyPatch, @@ -599,6 +612,21 @@ def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_ser assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens +def test_disconnecting_from_a_resumed_background_stream_bills_no_partial_answer(): + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=_make_logging_obj(), + custom_llm_provider="openai", + replays_stored_stream=True, + ) + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)): + iterator._process_chunk(json.dumps(event)) + + assert iterator.partial_response_on_disconnect() is None + + @pytest.mark.parametrize( "events", [ diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index c7fd7fe247f..7e74bfd0f17 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -11,7 +11,7 @@ from fastapi import HTTPException from litellm.proxy.hooks.responses_id_security import ( ResponsesIDSecurity, - _is_responses_api_create_route, + _streams_response_ids, ) from litellm.types.llms.openai import ( GenericEvent, @@ -527,7 +527,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() -class TestIsResponsesApiCreateRoute: +class TestStreamsResponseIds: """Test the route gate that decides whether a streamed response id is encrypted.""" @pytest.mark.parametrize( @@ -536,10 +536,13 @@ class TestIsResponsesApiCreateRoute: "/v1/responses", "/responses", "/openai/v1/responses", + "/v1/responses/resp_abc", + "/responses/resp_abc", + "/openai/v1/responses/resp_abc", ], ) - def test_create_routes_match(self, route): - assert _is_responses_api_create_route(route) is True + def test_create_and_retrieve_routes_match(self, route): + assert _streams_response_ids(route) is True @pytest.mark.parametrize( "route", @@ -547,14 +550,14 @@ class TestIsResponsesApiCreateRoute: None, "/chat/completions", "/openai/v1/chat/completions", - "/v1/responses/{response_id}", - "/openai/v1/responses/{response_id}", + "/v1/responses/", + "/v1/responses/resp_abc/input_items", "/v1/responsesX", "/responsesX", ], ) - def test_non_create_routes_do_not_match(self, route): - assert _is_responses_api_create_route(route) is False + def test_other_routes_do_not_match(self, route): + assert _streams_response_ids(route) is False class TestAsyncPostCallStreamingIteratorHook: @@ -609,7 +612,13 @@ class TestAsyncPostCallStreamingIteratorHook: @pytest.mark.asyncio @pytest.mark.parametrize( "route", - ["/v1/responses", "/responses", "/openai/v1/responses"], + [ + "/v1/responses", + "/responses", + "/openai/v1/responses", + "/v1/responses/resp_rawprovider123", + "/openai/v1/responses/resp_rawprovider123", + ], ) async def test_streamed_id_encrypted_on_all_responses_routes( self, responses_id_security, route, monkeypatch diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 51348fa9c75..2147afc4931 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -11211,6 +11211,9 @@ export interface paths { * * # Get provider response * curl -X GET http://localhost:4000/v1/responses/resp_abc123 -H "Authorization: Bearer $LITELLM_MASTER_KEY" + * + * # Resume a background response's event stream after sequence number 7 + * curl -N "http://localhost:4000/v1/responses/resp_abc123?stream=true&starting_after=7" -H "Authorization: Bearer $LITELLM_MASTER_KEY" * ``` */ get: operations["get_response_openai_v1_responses__response_id__get"]; @@ -14269,6 +14272,9 @@ export interface paths { * * # Get provider response * curl -X GET http://localhost:4000/v1/responses/resp_abc123 -H "Authorization: Bearer $LITELLM_MASTER_KEY" + * + * # Resume a background response's event stream after sequence number 7 + * curl -N "http://localhost:4000/v1/responses/resp_abc123?stream=true&starting_after=7" -H "Authorization: Bearer $LITELLM_MASTER_KEY" * ``` */ get: operations["get_response_responses__response_id__get"]; @@ -21891,6 +21897,9 @@ export interface paths { * * # Get provider response * curl -X GET http://localhost:4000/v1/responses/resp_abc123 -H "Authorization: Bearer $LITELLM_MASTER_KEY" + * + * # Resume a background response's event stream after sequence number 7 + * curl -N "http://localhost:4000/v1/responses/resp_abc123?stream=true&starting_after=7" -H "Authorization: Bearer $LITELLM_MASTER_KEY" * ``` */ get: operations["get_response_v1_responses__response_id__get"]; @@ -67129,7 +67138,16 @@ export interface operations { }; get_response_openai_v1_responses__response_id__get: { parameters: { - query?: never; + query?: { + /** @description Replay the response's events as server-sent events (OpenAI, Azure, and Azure AI Foundry) */ + stream?: boolean; + /** @description With stream=true, replay only the events after this sequence number */ + starting_after?: number; + /** @description Additional output data to include, forwarded to the provider */ + include?: string[]; + /** @description With stream=true, whether delta events carry obfuscation padding */ + include_obfuscation?: boolean; + }; header?: never; path: { response_id: string; @@ -67145,6 +67163,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */ @@ -70594,7 +70613,16 @@ export interface operations { }; get_response_responses__response_id__get: { parameters: { - query?: never; + query?: { + /** @description Replay the response's events as server-sent events (OpenAI, Azure, and Azure AI Foundry) */ + stream?: boolean; + /** @description With stream=true, replay only the events after this sequence number */ + starting_after?: number; + /** @description Additional output data to include, forwarded to the provider */ + include?: string[]; + /** @description With stream=true, whether delta events carry obfuscation padding */ + include_obfuscation?: boolean; + }; header?: never; path: { response_id: string; @@ -70610,6 +70638,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */ @@ -81397,7 +81426,16 @@ export interface operations { }; get_response_v1_responses__response_id__get: { parameters: { - query?: never; + query?: { + /** @description Replay the response's events as server-sent events (OpenAI, Azure, and Azure AI Foundry) */ + stream?: boolean; + /** @description With stream=true, replay only the events after this sequence number */ + starting_after?: number; + /** @description Additional output data to include, forwarded to the provider */ + include?: string[]; + /** @description With stream=true, whether delta events carry obfuscation padding */ + include_obfuscation?: boolean; + }; header?: never; path: { response_id: string; @@ -81413,6 +81451,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */ From 42d29092105f3f1b618a2375e15743a703b3d077 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:28:39 -0700 Subject: [PATCH 12/58] test(integration): scope the signoz, otel conversation id, and series cap proxy fixtures to their module (#48965) * test(integration): scope the signoz, otel conversation id, and series cap proxy fixtures to their module Session scope booted one 2-worker proxy per xdist worker and kept it until the worker's session ended, so a CircleCI node holding two of these files piled up proxies until the VM stopped responding. Module scope lets the harness group each file's rig tests onto one worker and tears the proxies down after them. * docs(tests): keep proxy fixtures module scoped in integration tests * docs(tests): say how the harness groups proxy fixtures in the integration AGENTS.md rule --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- tests/integration/AGENTS.md | 11 +++++++++++ tests/integration/observability/conftest.py | 2 +- .../observability/test_otel_conversation_id.py | 6 +++--- .../integration/observability/test_signoz_delivery.py | 4 ++-- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index a727b20596c..5bdc538b2d9 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -37,3 +37,14 @@ together with its dedicated job, and list a new dedicated GitHub workflow with t `DEDICATED_INTEGRATION_WORKFLOWS` in `.github/scripts/assert_ci_coverage.py`. The `ci-coverage` check fails a file here that a GitHub workflow runs unless it is in `GITHUB_FILES` and that workflow owns its directory (#46780) + +## Fixtures that boot a proxy + +Scope them to the module or class, or pin a session one to a single worker with an `xdist_group` mark +as `_llm_gateway` in `mcp/test_mcp_llm_endpoints.py` does, and request them as arguments, never through +`request.getfixturevalue`. The harness puts tests that request the same set of module or class fixtures +on one xdist worker, so the proxy they share boots there once and stops when that worker moves to another +file. A test that requests one more of those fixtures than its neighbors gets its own group and can boot +the proxy again on another worker, so keep that set the same across the tests that share a proxy. An +unpinned session fixture boots a copy on every worker that runs one of its tests and keeps them all until +the node ends, and those copies ran CircleCI integration-extensions nodes out of memory diff --git a/tests/integration/observability/conftest.py b/tests/integration/observability/conftest.py index 101c2958b64..eb5c12a1cea 100644 --- a/tests/integration/observability/conftest.py +++ b/tests/integration/observability/conftest.py @@ -34,7 +34,7 @@ def langfuse_vars(audit_sinks: SpanSinks) -> dict[str, JsonValue]: return langfuse_callback_variables(audit_sinks) -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def capped(tmp_path_factory: pytest.TempPathFactory) -> Iterator[CapRig]: """A two-worker proxy capped at three series per metric, with three keys already holding a series each.""" with series_cap_rig( diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index af893abab48..d7f4de2961b 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -305,12 +305,12 @@ class RigFactory: yield Rig(owned.gateway, owned, model, self.provider, self.sink) -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel"), {}, 2).start() -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def generating_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: factory: Final = RigFactory( provider, collector, tmp_path_factory.mktemp("otel-generate"), {"missing_session_id": "generate"}, 2 @@ -318,7 +318,7 @@ def generating_rig(provider: Wire, collector: Collector, tmp_path_factory: pytes yield from factory.start() -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def two_worker_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel-workers"), {}, 2).start() diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index 6e50a9bedf5..e25126e4fa8 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -424,7 +424,7 @@ class RigFactory: yield Rig(owned.gateway, owned, model, self.provider, self.sink, self.tenant_sink) -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def rig( provider: Wire, operator_sink: Collector, tenant_sink: Collector, tmp_path_factory: pytest.TempPathFactory ) -> Iterator[Rig]: @@ -434,7 +434,7 @@ def rig( yield from factory.start() -@pytest.fixture(scope="session") +@pytest.fixture(scope="module") def v2_rig( provider: Wire, operator_sink: Collector, tenant_sink: Collector, tmp_path_factory: pytest.TempPathFactory ) -> Iterator[Rig]: From 310c96162b45eed093442ea97d130dae2c9a7c98 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:36:20 -0700 Subject: [PATCH 13/58] fix(anthropic): drop the non-spec custom object from /v1/messages tool definitions (#48893) * fix(anthropic): drop the non-spec custom object from /v1/messages tool definitions Anthropic and Vertex AI reject a tool carrying a custom object with tools.N.custom.custom: Extra inputs are not permitted. Drop it on the shared Anthropic request path, keep the hoist of a boolean custom.defer_loading to top-level defer_loading, reuse the helper on Bedrock, and stop the chat adapter from merging custom into the tool's JSON schema * fix(messages): hoist custom.defer_loading only on Claude routes and strip custom on count_tokens Non-Claude Messages providers such as DeepSeek now get custom dropped without a defer_loading hoist, since they have no tool search. /v1/messages/count_tokens on Anthropic, Azure AI, and Vertex AI drops custom too, and the local token counter no longer counts it as schema. * fix(anthropic): drop tool custom object only under drop_params and hoist defer_loading for Claude models Without drop_params the tool is forwarded as sent, so the provider's 400 stands. A boolean custom.defer_loading still becomes the top-level flag for Claude models and on Bedrock Invoke and Mantle. Bedrock count_tokens now normalizes invokeModel tools the way Bedrock Messages sends them. * fix(bedrock): drop tool custom object on the Mantle count_tokens fallback Mantle Messages drops custom unconditionally, so its count_tokens call now counts the tools the same way. * fix(messages): keep custom-tool Messages calls on Python and match Mantle count_tokens to Mantle Messages * test(vertex): annotate count_tokens tool test locals with Final * fix(messages): pass tools without custom through untouched so Bedrock fallbacks keep filled tool names * fix(messages): keep Bedrock's in-place tool cleanup so fallbacks still work, and stop hoisting defer_loading on the OpenAI-like passthrough --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/token_counter.py | 2 +- litellm/llms/anthropic/common_utils.py | 48 ++++++++ .../llms/anthropic/count_tokens/handler.py | 2 + .../anthropic/count_tokens/token_counter.py | 2 + .../anthropic/count_tokens/transformation.py | 21 +++- .../messages/adapters/transformation.py | 1 + .../messages/native/transformation.py | 25 +++- .../anthropic/count_tokens/handler.py | 2 + litellm/llms/bedrock/common_utils.py | 37 ++++-- .../bedrock/count_tokens/mantle_handler.py | 9 +- .../bedrock/count_tokens/transformation.py | 11 +- .../anthropic_claude_transformation.py | 3 + .../openai_like/messages/transformation.py | 3 + .../count_tokens/handler.py | 21 +++- litellm/main.py | 14 +-- litellm/messages/dispatch.py | 1 + litellm/rust_bridge/catalog.py | 3 + litellm/rust_bridge/messages/detectors.py | 5 + .../messages/test_dispatch.py | 6 + .../litellm_core_utils/test_token_counter.py | 1 + ...hropic_messages_adapters_transformation.py | 15 +++ .../test_anthropic_messages_transformation.py | 82 ++++++++++++++ .../anthropic/test_anthropic_common_utils.py | 20 ++++ ...t_anthropic_count_tokens_transformation.py | 48 ++++++++ ...est_bedrock_count_tokens_transformation.py | 21 ++++ ...est_bedrock_mantle_count_tokens_handler.py | 28 +++++ .../test_anthropic_claude_transformation.py | 107 ++++++++++++++++++ ..._bedrock_mantle_messages_transformation.py | 14 +++ .../test_count_tokens_location.py | 21 ++++ tests/unit/rust_bridge/test_detectors.py | 3 + tests/unit/test_main.py | 19 ---- 31 files changed, 546 insertions(+), 49 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 2fe840dc863..591d77bdb16 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -611,7 +611,7 @@ T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") _COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") _ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE: Final = frozenset( - ("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming") + ("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming", "custom") ) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 1f5e4e7de39..e591f3a3fde 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -121,6 +121,54 @@ _DECODED_JSON: Final = TypeAdapter(object) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") +def is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: + if custom_llm_provider == "anthropic": + return True + model_lower: Final = model.lower() + if custom_llm_provider == "bedrock": + return "claude" in model_lower or ("arn:" in model_lower and ":bedrock:" in model_lower) + if custom_llm_provider == "azure_ai" or custom_llm_provider == "vertex_ai": + return "claude" in model_lower + return False + + +def is_claude_model_target(custom_llm_provider: str | None, model: str) -> bool: + return "claude" in model.lower() and is_claude_tool_target(custom_llm_provider, model) + + +def _hoistable_defer_loading(tool_dict: Mapping[object, object], custom: object) -> bool | None: + if not isinstance(custom, dict) or "defer_loading" in tool_dict: + return None + defer_loading: Final = OBJECT_MAPPING_ADAPTER.validate_python(custom).get("defer_loading") + return defer_loading if isinstance(defer_loading, bool) else None + + +def _tool_without_custom_field(tool: object, hoist_defer_loading: bool, drop_custom: bool) -> object: + if not isinstance(tool, dict): + return tool + + if "custom" not in tool: + return tool + + tool_dict: Final = OBJECT_MAPPING_ADAPTER.validate_python(tool) + + custom: Final = tool_dict["custom"] + without_custom: Final = {key: value for key, value in tool_dict.items() if key != "custom"} + defer_loading: Final = _hoistable_defer_loading(tool_dict, custom) if hoist_defer_loading else None + if defer_loading is None: + return without_custom if drop_custom else tool_dict + + unhoisted: Final = { + key: value for key, value in OBJECT_MAPPING_ADAPTER.validate_python(custom).items() if key != "defer_loading" + } + kept_custom: Final = {"custom": unhoisted} if unhoisted and not drop_custom else {} + return {**without_custom, "defer_loading": defer_loading, **kept_custom} + + +def tools_without_custom_field(tools: Sequence[object], hoist_defer_loading: bool, drop_custom: bool) -> list[object]: + return [_tool_without_custom_field(tool, hoist_defer_loading, drop_custom) for tool in tools] + + def is_anthropic_messages_url(url: str) -> bool: """Check whether a URL addresses Anthropic's Messages API.""" parsed_url: Final = urlparse(url) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index a230294ad3f..f8bf7232d9c 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -39,6 +39,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, optional_params: Mapping[str, JsonValue] | None = None, + drop_params: bool = False, ) -> dict[str, JsonValue]: """ Handle a CountTokens request using httpx. @@ -69,6 +70,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): tools=tools, system=system, optional_params=optional_params, + drop_params=drop_params, ) verbose_logger.debug("Transformed request: %s", request_body) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 368bb929eeb..91fc77773d1 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -11,6 +11,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountFields, @@ -131,6 +132,7 @@ class AnthropicTokenCounter(BaseTokenCounter): tools=tools, system=system, optional_params=count_options, + drop_params=drops_unsupported_params(litellm_params.get("drop_params")), ) if result is not None: diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 0967c22f93a..e60ffe67c8d 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -26,7 +26,11 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, ) from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers +from litellm.llms.anthropic.common_utils import ( + is_claude_model_target, + merge_anthropic_beta_headers, + tools_without_custom_field, +) from litellm.llms.anthropic.wif import resolve_anthropic_base from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse @@ -372,6 +376,7 @@ class AnthropicCountTokensConfig: tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, optional_params: Mapping[str, JsonValue] | None = None, + drop_params: bool = False, ) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries """ Transform request to Anthropic CountTokens format. @@ -391,7 +396,19 @@ class AnthropicCountTokensConfig: **MappingProxyType( { key: value - for key, value in (("system", counted_system), ("tools", tools)) + for key, value in ( + ("system", counted_system), + ( + "tools", + None + if tools is None + else tools_without_custom_field( + tools, + hoist_defer_loading=is_claude_model_target("anthropic", model), + drop_custom=drop_params, + ), + ), + ) if value is not None } ), diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 1ed552d46e3..5fc7a56600c 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -817,6 +817,7 @@ class LiteLLMAnthropicMessagesAdapter: "name", "input_schema", "description", + "custom", "cache_control", "strict", "type", diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index cc53858f842..6c0e463bd23 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -29,13 +29,16 @@ from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( + OBJECT_LIST_ADAPTER, AnthropicError, AnthropicModelInfo, + is_claude_model_target, merge_anthropic_beta_headers, optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks_from_messages, strip_encrypted_reasoning_blocks_from_anthropic_messages, + tools_without_custom_field, ) from .mid_conversation_system import ( as_system_content_blocks, @@ -181,6 +184,16 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): """ return False + def hoists_wrapped_defer_loading(self, model: str) -> bool: + return is_claude_model_target(self._resolved_provider, model) + + def tools_for_upstream(self, tools: Sequence[object], model: str, drop_params: object) -> Sequence[object]: + return tools_without_custom_field( + tools, + hoist_defer_loading=self.hoists_wrapped_defer_loading(model), + drop_custom=drops_unsupported_params(drop_params), + ) + @staticmethod def _filter_billing_headers_from_system(system_param): """ @@ -700,7 +713,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Auto-strip advisor blocks from history if advisor tool is absent. # Prevents Anthropic 400: advisor_tool_result in history requires advisor tool. - _tools: Final = anthropic_messages_optional_request_params.get("tools") or [] + tools: Final = anthropic_messages_optional_request_params.get("tools") + _tools: Final = tools or [] _has_advisor: Final = any(isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for t in _tools) if not _has_advisor: messages = strip_advisor_blocks_from_messages(messages) @@ -713,7 +727,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): model=model, **anthropic_messages_optional_request_params, ) - return dict(anthropic_messages_request) + if not isinstance(tools, list): + return dict(anthropic_messages_request) + return { + **anthropic_messages_request, + "tools": self.tools_for_upstream( + OBJECT_LIST_ADAPTER.validate_python(tools), model, litellm_params.drop_params + ), + } def transform_anthropic_messages_response( self, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 0f8afaae681..4b3f263201a 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -13,6 +13,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -69,6 +70,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): tools=tools, system=system, optional_params=optional_params, + drop_params=drops_unsupported_params((litellm_params or {}).get("drop_params")), ) verbose_logger.debug("Transformed request: %s", request_body) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index fde995c7907..dd721e059a9 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -15,7 +15,11 @@ from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ( + ReadOnly, + TypedDict, + TypeIs, # noqa: TID251 # narrows the caller's tools without copying them +) if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage @@ -573,6 +577,26 @@ def strip_unsupported_bedrock_invoke_output_config_keys( request_body["output_config"] = {"format": preserved_format} # rebind-ok: out-param +_TOOL_ITEMS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) + + +def _is_str_keyed_dict( + value: object, +) -> TypeIs[dict[str, object]]: # guard-ok: JSON tool objects # mutable-ok: fallbacks reuse tools edited in place + return isinstance(value, dict) + + +def _normalize_custom_field_on_tool(tool: object) -> None: + if not _is_str_keyed_dict(tool): + return + custom: Final = tool.pop("custom", None) + if not _is_str_keyed_dict(custom) or "defer_loading" in tool: + return + deferred: Final = custom.get("defer_loading") + if isinstance(deferred, bool): + tool["defer_loading"] = deferred # rebind-ok: a fallback reuses the caller's tool, as on main + + def normalize_custom_field_on_tools(request_body: dict) -> None: """ Drop the ``custom`` field from each tool, first hoisting a boolean @@ -590,15 +614,8 @@ def normalize_custom_field_on_tools(request_body: dict) -> None: tools: Final = request_body.get("tools") if not tools or not isinstance(tools, list): return - for tool in tools: - if not isinstance(tool, dict): - continue - custom: dict[str, object] | None = tool.pop("custom", None) - if not isinstance(custom, dict) or "defer_loading" in tool: - continue - deferred: object = custom.get("defer_loading") - if isinstance(deferred, bool): - tool["defer_loading"] = deferred + for tool in _TOOL_ITEMS_ADAPTER.validate_python(tools): + _normalize_custom_field_on_tool(tool) _TOOL_DICTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 34e622b5e8d..5cd26235ba2 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -8,6 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.anthropic.common_utils import tools_without_custom_field from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -26,6 +27,7 @@ class MantleCountTokensRequest(TypedDict): _COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) _COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) _COUNT_OPTIONS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue]) +_TOOLS: Final = TypeAdapter(list[dict[str, JsonValue]]) class BedrockMantleCountTokensHandler(BaseAWSLLM): @@ -59,10 +61,15 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None ) endpoint_url: Final = await run_aws_signing(self.get_mantle_count_tokens_endpoint, aws_region_name) + tools: Final = request.get("tools") body: Final = self._anthropic_config.transform_request_to_count_tokens( model=resolved_model, messages=request["messages"], - tools=request.get("tools"), + tools=None + if tools is None + else _TOOLS.validate_python( + tools_without_custom_field(tools, hoist_defer_loading=True, drop_custom=True) + ), system=request.get("system"), optional_params=_COUNT_OPTIONS.validate_python( {name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES} diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index bf9c119f6f2..df5a3827c1f 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -12,6 +12,7 @@ from typing import Final, Literal from pydantic import JsonValue from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.anthropic.common_utils import tools_without_custom_field from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -194,7 +195,15 @@ class BedrockCountTokensConfig(BaseAWSLLM): # For InvokeModel, we need to provide the raw body that would be sent to the model # Remove the 'model' field from the body as it's not part of the model input - body_data: Final = {k: v for k, v in request_data.items() if k != "model"} + tools: Final = request_data.get("tools") + body_data: Final = { + **{k: v for k, v in request_data.items() if k != "model"}, + **( + {"tools": tools_without_custom_field(tools, hoist_defer_loading=True, drop_custom=True)} + if isinstance(tools, list) + else {} + ), + } if "messages" in body_data: # Bedrock validates the body against the model's InvokeModel schema; diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py index e1e15861a90..60ce73abd82 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py @@ -102,6 +102,9 @@ class AmazonAnthropicClaudeMessagesConfig( def should_strip_billing_metadata(self) -> bool: return True + def tools_for_upstream(self, tools: Sequence[object], model: str, drop_params: object) -> Sequence[object]: + return tools + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: return True diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 8daae2a8e6a..a5d825a39d7 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -37,6 +37,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def is_native_messages_route(self) -> bool: return False + def hoists_wrapped_defer_loading(self, model: str) -> bool: + return False + def validate_anthropic_messages_environment( self, headers: dict[str, str], diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 4be41e26326..180761f15d9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -6,9 +6,12 @@ Unlike Gemini models which use Google's token counting API, partner models use their respective publisher-specific count-tokens endpoints. """ +from collections.abc import Mapping from typing import Any, Final +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.anthropic.common_utils import OBJECT_LIST_ADAPTER, is_claude_model_target, tools_without_custom_field from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -20,6 +23,22 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( ) +def count_tokens_request_body( + request_data: Mapping[str, object], model: str, drop_params: object +) -> Mapping[str, object]: + tools: Final = request_data.get("tools") + if not isinstance(tools, list): + return request_data + return { + **request_data, + "tools": tools_without_custom_field( + OBJECT_LIST_ADAPTER.validate_python(tools), + hoist_defer_loading=is_claude_model_target("vertex_ai", model), + drop_custom=drops_unsupported_params(drop_params), + ), + } + + def count_tokens_publisher(model: str) -> str | None: return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) @@ -169,7 +188,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): response: Final = await async_client.post( endpoint_url, headers=headers, - json=request_data, + json=dict(count_tokens_request_body(request_data, model, litellm_params.get("drop_params"))), timeout=30.0, ) diff --git a/litellm/main.py b/litellm/main.py index b8ad28c26fe..996e93589fe 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -234,6 +234,7 @@ from .litellm_core_utils.streaming_chunk_builder_utils import ( ChunkProcessor as ChunkProcessor, # noqa: PLC0414 # public re-export ) from .llms.anthropic.chat import AnthropicChatCompletion +from .llms.anthropic.common_utils import is_claude_tool_target from .llms.azure.audio_transcriptions import AzureAudioTranscription from .llms.azure.azure import AzureChatCompletion, check_dynamic_azure_params from .llms.azure.chat.o_series_handler import AzureOpenAIO1ChatCompletion @@ -1135,17 +1136,6 @@ from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the o _ANTHROPIC_ONLY_TOOL_KEYS: Final = frozenset({"input_examples", "eager_input_streaming"}) -def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: - if custom_llm_provider == "anthropic": - return True - model_lower: Final = model.lower() - if custom_llm_provider == "bedrock": - return "claude" in model_lower or ("arn:" in model_lower and ":bedrock:" in model_lower) - if custom_llm_provider == "azure_ai" or custom_llm_provider == "vertex_ai": - return "claude" in model_lower - return False - - def _history_without_unanswered_tool_search_calls( messages: Sequence[object], custom_llm_provider: str | None, provider_config: BaseConfig ) -> tuple[object, ...] | None: @@ -5752,7 +5742,7 @@ def completion( } ) - if not _is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): + if not is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_anthropic_only_tool_keys(tools=tools) if provider_specific_header is not None: diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 0e2fc2e3dcb..8e6becc60b4 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -111,6 +111,7 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry, + Feature.TOOL_CUSTOM_FIELD: route_detectors.tool_custom_field, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 74cd28161a4..fe52f8416d0 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -96,6 +96,8 @@ class Feature(str, Enum): """litellm.include_cost_in_streaming_usage, where Python adds cost to streamed usage.""" THINKING_SIGNATURE_RETRY = "thinking_signature_retry" """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" + TOOL_CUSTOM_FIELD = "tool_custom_field" + """Tools carrying a non-spec custom object, which Python drops or hoists for the provider.""" CLIENT_HEADERS = "client_headers" """Proxy-forwarded headers that Python merges into extra_headers.""" DEPLOYMENT_HOOK = "deployment_hook" @@ -176,6 +178,7 @@ MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.CACHE_CONTROL_INJECTION, Feature.STREAM_COST_INJECTION, Feature.THINKING_SIGNATURE_RETRY, + Feature.TOOL_CUSTOM_FIELD, } RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 5f75858d350..52917126b3b 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -43,6 +43,11 @@ def thinking_signature_retry(request: NativeCall) -> bool: return any(_replays_thinking(message) for message in optional_sequence(request.resolved.get("messages")) or ()) +def tool_custom_field(request: NativeCall) -> bool: + tools: Final = optional_sequence(request.resolved.get("tools")) or () + return any("custom" in (optional_mapping(tool) or {}) for tool in tools) + + def stream_cost_injection(request: NativeCall) -> bool: """Python adds response cost to streaming usage when the global setting is truthy.""" return streaming(request) and bool(litellm.include_cost_in_streaming_usage) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 9430e01ea8a..15f58afb4e5 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -473,6 +473,12 @@ ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", " {"messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]}, {}, ), + ( + Feature.TOOL_CUSTOM_FIELD, + "anthropic/claude-sonnet-4-5", + {"tools": [{"name": "a", "input_schema": {"type": "object"}, "custom": {"eager_input_streaming": True}}]}, + {}, + ), ( Feature.STREAM_COST_INJECTION, "anthropic/claude-sonnet-4-5", diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 9b0f6607708..000e3993bc2 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -700,6 +700,7 @@ def test_token_counter_counts_anthropic_tool_choices_like_openai_choices( _anthropic_tool(_WEATHER_FUNCTION), {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True}, {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]}, + {**_anthropic_tool(_WEATHER_FUNCTION), "custom": {"eager_input_streaming": True}}, _anthropic_tool(_WEB_SEARCH_FUNCTION), ], ) diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index 4e4b89b764b..51ec71aaea0 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -4238,6 +4238,21 @@ def test_translate_anthropic_tools_to_openai_preserves_parameters_type(): assert new_tools[0]["type"] == "function" +def test_translate_anthropic_custom_field_does_not_change_openai_parameters(): + input_schema: Final = {"type": "object", "properties": {"city": {"type": "string"}}} + tools: Final = [ + { + "name": "get_weather", + "input_schema": input_schema, + "custom": {"eager_input_streaming": True}, + } + ] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert new_tools[0]["function"]["parameters"] == input_schema + + def test_translate_anthropic_tools_to_openai_maps_strict_onto_function_not_parameters(): """A tool-level `strict` lands on the OpenAI function, leaving the caller's `input_schema` untouched.""" adapter = LiteLLMAnthropicMessagesAdapter() diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index aa6625561ba..13729ecb6cb 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -14,9 +14,13 @@ from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.wif import get_anthropic_wif_token +from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, +) from litellm.types.llms.anthropic import AnthropicFinishReason from litellm.types.router import GenericLiteLLMParams from tests.unit.llms.anthropic.test_anthropic_wif import ( @@ -79,6 +83,84 @@ def test_native_response_defaults_stop_sequence_without_changing_other_fields( } +_EMPTY_SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} + + +@pytest.mark.parametrize( + "config_class", + [AnthropicMessagesConfig, VertexAIPartnerModelsAnthropicMessagesConfig], +) +@pytest.mark.parametrize( + ("drop_params", "expected_eager_tool"), + [ + (True, {"name": "A", "input_schema": _EMPTY_SCHEMA}), + (False, {"name": "A", "input_schema": _EMPTY_SCHEMA, "custom": {"eager_input_streaming": True}}), + ], +) +def test_messages_request_hoists_wrapped_defer_loading_and_drops_custom_only_under_drop_params( + config_class: type[AnthropicMessagesConfig], drop_params: bool, expected_eager_tool: dict[str, object] +) -> None: + tools: Final[list[dict[str, object]]] = [ + {"name": "A", "input_schema": _EMPTY_SCHEMA, "custom": {"eager_input_streaming": True}}, + {"name": "B", "input_schema": _EMPTY_SCHEMA, "custom": {"defer_loading": True}}, + {"name": "C", "input_schema": _EMPTY_SCHEMA, "custom": {"defer_loading": True, "eager_input_streaming": True}}, + {"name": "D", "input_schema": _EMPTY_SCHEMA}, + ] + optional_params: Final[dict[str, object]] = {"max_tokens": 50, "tools": tools} + + request: Final = config_class().transform_anthropic_messages_request( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "Hi"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + assert request["tools"] == [ + expected_eager_tool, + {"name": "B", "input_schema": _EMPTY_SCHEMA, "defer_loading": True}, + { + "name": "C", + "input_schema": _EMPTY_SCHEMA, + "defer_loading": True, + **({} if drop_params else {"custom": {"eager_input_streaming": True}}), + }, + {"name": "D", "input_schema": _EMPTY_SCHEMA}, + ] + assert tools[0]["custom"] == {"eager_input_streaming": True} + assert optional_params["tools"] is tools + + +@pytest.mark.parametrize( + ("config", "model"), + [ + (DeepSeekAnthropicMessagesConfig(), "deepseek-chat"), + (MinimaxMessagesConfig(), "MiniMax-M2"), + (AnthropicMessagesConfig(), "deepseek-chat"), + (OpenAILikeAnthropicMessagesConfig(), "anthropic/claude-haiku-4.5"), + ], +) +@pytest.mark.parametrize("drop_params", [True, False]) +def test_messages_request_never_defers_tools_off_claude( + config: AnthropicMessagesConfig, model: str, drop_params: bool +) -> None: + wrapped_tool: Final[dict[str, object]] = { + "name": "A", + "input_schema": _EMPTY_SCHEMA, + "custom": {"defer_loading": True}, + } + + request: Final = config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Hi"}], + anthropic_messages_optional_request_params={"max_tokens": 50, "tools": [wrapped_tool]}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + assert request["tools"] == [{"name": "A", "input_schema": _EMPTY_SCHEMA} if drop_params else wrapped_tool] + + def test_workload_identity_allowed_for_anthropic() -> None: assert AnthropicMessagesConfig()._allows_workload_identity is True diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index eb5c1598c17..d408433e59c 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -25,6 +25,7 @@ import pytest import respx import litellm +from litellm.llms.anthropic.common_utils import is_claude_tool_target from litellm.llms.custom_httpx.http_handler import HTTPHandler sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) @@ -4522,3 +4523,22 @@ def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_dro ) assert optional_params == {"thinking": {"type": "adaptive"}} + + +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 89352ce8702..a3e723eaca2 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -171,6 +171,54 @@ def test_transform_includes_tools(): assert result["tools"] == tools +_PATH_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"path": {"type": "string"}}} + + +@pytest.mark.parametrize( + ("model", "drop_params", "expected_tools"), + [ + ( + "claude-haiku-5-5", + True, + [ + {"name": "read_file", "input_schema": _PATH_SCHEMA}, + {"name": "list_files", "input_schema": _PATH_SCHEMA, "defer_loading": True}, + ], + ), + ( + "claude-haiku-5-5", + False, + [ + {"name": "read_file", "input_schema": _PATH_SCHEMA, "custom": {"eager_input_streaming": True}}, + {"name": "list_files", "input_schema": _PATH_SCHEMA, "defer_loading": True}, + ], + ), + ( + "deepseek-chat", + True, + [ + {"name": "read_file", "input_schema": _PATH_SCHEMA}, + {"name": "list_files", "input_schema": _PATH_SCHEMA}, + ], + ), + ], +) +def test_transform_counts_tools_the_way_messages_sends_them( + model: str, drop_params: bool, expected_tools: list[dict[str, JsonValue]] +): + result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( + model=model, + messages=[{"role": "user", "content": "Hello"}], + tools=[ + {"name": "read_file", "input_schema": _PATH_SCHEMA, "custom": {"eager_input_streaming": True}}, + {"name": "list_files", "input_schema": _PATH_SCHEMA, "custom": {"defer_loading": True}}, + ], + drop_params=drop_params, + ) + + assert result["tools"] == expected_tools + + def test_transform_includes_system_and_tools(): """Test that both system and tools are included together.""" config = AnthropicCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 052baf601f9..8ac35caba9d 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -73,6 +73,27 @@ def test_transform_to_invoke_model_format_base64_encodes_body(): assert body["max_tokens"] == DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS +def test_transform_to_invoke_model_format_counts_tools_the_way_bedrock_messages_sends_them(): + schema = {"type": "object", "properties": {}} + + result = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens( + { + "model": "anthropic.claude-haiku-5-5", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + "tools": [ + {"name": "read", "input_schema": schema, "custom": {"defer_loading": True}}, + {"name": "write", "input_schema": schema, "custom": {"eager_input_streaming": True}}, + ], + } + ) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body["tools"] == [ + {"name": "read", "input_schema": schema, "defer_loading": True}, + {"name": "write", "input_schema": schema}, + ] + + def test_transform_to_invoke_model_format_raw_body_unchanged(): """Non-messages bodies (e.g. Titan inputText) must not get Anthropic fields.""" config = BedrockCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py index a6b309c1b83..930a290098f 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -74,6 +74,34 @@ async def test_counts_on_mantle_with_the_base_model_and_the_deployment_credentia } +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resolved_model", + ["anthropic.claude-opus-4-8", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123"], +) +async def test_counts_tools_the_way_mantle_messages_sends_them(resolved_model: str) -> None: + mantle: Final = _MantleEndpoint(200, {"input_tokens": 11}) + schema: Final = {"type": "object", "properties": {}} + + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data={ + **REQUEST, + "tools": [ + {"name": "eager", "input_schema": schema, "custom": {"eager_input_streaming": True}}, + {"name": "deferred", "input_schema": schema, "custom": {"defer_loading": True}}, + ], + }, + litellm_params=dict(LITELLM_PARAMS), + resolved_model=resolved_model, + client=_client(mantle), + ) + + assert json.loads(mantle.only_request().content)["tools"] == [ + {"name": "eager", "input_schema": schema}, + {"name": "deferred", "input_schema": schema, "defer_loading": True}, + ] + + @pytest.mark.asyncio async def test_bedrock_mantle_api_base_env_names_the_host(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws") diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py index 3c0bbadc133..c096e2126ae 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py @@ -4198,3 +4198,110 @@ def test_bedrock_messages_without_guardrail_config_sends_no_guardrail_headers() headers={}, model="anthropic.claude-sonnet-5-v1:0", messages=[], optional_params={}, litellm_params={} ) assert not any(name.startswith("X-Amzn-Bedrock-Guardrail") for name in headers) + + +@dataclass(slots=True) +class BedrockDownAnthropicUp: + anthropic_bodies: tuple[dict[str, object], ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.host != "api.anthropic.com": + return httpx.Response(500, json={"message": "Bedrock is unavailable"}) + self.anthropic_bodies = (*self.anthropic_bodies, REASONING_BODY_ADAPTER.validate_json(request.content)) + return httpx.Response(200, json=_ANTHROPIC_MESSAGE) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("unnamed_tool_extras", "sent_extras"), + [ + ({}, {}), + ({"custom": {"eager_input_streaming": True}}, {}), + ({"custom": {"defer_loading": False}}, {"defer_loading": False}), + ], +) +async def test_bedrock_messages_fallback_to_anthropic_sends_the_tool_names_bedrock_filled_in( + respx_mock: respx.Router, + monkeypatch: pytest.MonkeyPatch, + unnamed_tool_extras: dict[str, object], + sent_extras: dict[str, object], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = BedrockDownAnthropicUp() + respx_mock.post(url__regex=r"https://.*").mock(side_effect=upstream) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/invoke/us.anthropic.claude-opus-4-8", + "aws_access_key_id": "test-access-key", + "aws_secret_access_key": "test-secret-key", + "aws_region_name": "us-east-1", + }, + }, + {"model_name": "anthropic-claude", "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "test"}}, + ], + fallbacks=[{"bedrock-claude": ["anthropic-claude"]}], + num_retries=0, + ) + schema: Final[Mapping[str, object]] = {"type": "object", "properties": {}} + + await router.aanthropic_messages( + model="bedrock-claude", + messages=[{"role": "user", "content": "What is the weather?"}], + max_tokens=100, + tools=[{"input_schema": schema, **unnamed_tool_extras}, {"name": "get_weather", "input_schema": schema}], + ) + + (body,) = upstream.anthropic_bodies + assert body["tools"] == [ + {"name": "litellm_unnamed_tool_0", "input_schema": schema, **sent_extras}, + {"name": "get_weather", "input_schema": schema}, + ] + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("custom", "sent_extras"), + [({"eager_input_streaming": True}, {}), ({"defer_loading": False}, {"defer_loading": False})], +) +async def test_bedrock_chat_fallback_to_anthropic_sends_the_tool_bedrock_cleaned( + respx_mock: respx.Router, + monkeypatch: pytest.MonkeyPatch, + custom: dict[str, object], + sent_extras: dict[str, object], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = BedrockDownAnthropicUp() + respx_mock.post(url__regex=r"https://.*").mock(side_effect=upstream) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/invoke/us.anthropic.claude-opus-4-8", + "aws_access_key_id": "test-access-key", + "aws_secret_access_key": "test-secret-key", + "aws_region_name": "us-east-1", + }, + }, + { + "model_name": "anthropic-claude", + "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "test", "drop_params": True}, + }, + ], + fallbacks=[{"bedrock-claude": ["anthropic-claude"]}], + num_retries=0, + ) + schema: Final[Mapping[str, object]] = {"type": "object", "properties": {}} + + await router.acompletion( + model="bedrock-claude", + messages=[{"role": "user", "content": "What is the weather?"}], + max_tokens=100, + tools=[{"name": "get_weather", "input_schema": schema, "custom": custom}], + ) + + (body,) = upstream.anthropic_bodies + assert body["tools"] == [{"name": "get_weather", "input_schema": schema, **sent_extras}] diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 3236efbbfba..432d0912248 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -224,6 +224,20 @@ class TestRequestBody: ) assert "stream" not in body + def test_body_hoists_wrapped_defer_loading_and_drops_custom(self): + schema: Final = {"type": "object", "properties": {}} + body = BedrockMantleAnthropicMessagesConfig().transform_anthropic_messages_request( + model="anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + anthropic_messages_optional_request_params={ + "max_tokens": 8, + "tools": [{"name": "Read", "input_schema": schema, "custom": {"defer_loading": True}}], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["tools"] == [{"name": "Read", "input_schema": schema, "defer_loading": True}] + class TestAuth: def test_bearer_from_api_key_skips_aws_credentials(self): diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py index e2fb81bc240..9474d1c83e1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py @@ -4,10 +4,13 @@ Tests for Vertex AI partner models count_tokens location resolution. Ref: https://github.com/BerriAI/litellm/issues/23872 """ +from typing import Final + import pytest from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( VertexAIPartnerModelsTokenCounter, + count_tokens_request_body, ) @@ -294,3 +297,21 @@ class TestCountTokensVersionSuffixStripping: # The model name sent to the API must NOT have the @default suffix assert captured_json["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize("drop_params", [True, False]) +def test_count_tokens_counts_tools_the_way_vertex_messages_sends_them(drop_params): + schema: Final = {"type": "object", "properties": {}} + eager_tool: Final = {"name": "write", "input_schema": schema, "custom": {"eager_input_streaming": True}} + request_data: Final = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "read", "input_schema": schema, "custom": {"defer_loading": True}}, eager_tool], + } + + body: Final = count_tokens_request_body(request_data, "claude-sonnet-4-6", drop_params) + + assert body["messages"] == request_data["messages"] + assert body["tools"] == [ + {"name": "read", "input_schema": schema, "defer_loading": True}, + {"name": "write", "input_schema": schema} if drop_params else eager_tool, + ] diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 78722458cf3..c02d2c12dbf 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -189,6 +189,9 @@ def _wrap_nested_level(value: object, depth: int) -> object: {"messages": [{"role": "user", "content": [{"type": "text"}]}]}, False, ), + (messages_detectors.tool_custom_field, {"tools": [{"name": "a"}, {"name": "b", "custom": {}}]}, True), + (messages_detectors.tool_custom_field, {"tools": [{"name": "a", "defer_loading": True}]}, False), + (messages_detectors.tool_custom_field, {}, False), (detectors.drop_params, {"drop_params": True}, True), (detectors.drop_params, {"drop_params": False}, False), (detectors.drop_params, {"drop_params": " True "}, True), diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f38bbcf4404..e119fe133d5 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -503,25 +503,6 @@ def test_bedrock_latency_optimized_inference(): assert json_data["performanceConfig"]["latency"] == "optimized" -@pytest.mark.parametrize( - ("custom_llm_provider", "model", "expected"), - [ - ("anthropic", "claude-sonnet-5", True), - ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), - ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), - ("bedrock", "us.amazon.nova-2-lite-v1:0", False), - ("vertex_ai", "claude-sonnet-5", True), - ("vertex_ai", "gemini-3.8-flash", False), - ("azure_ai", "claude-sonnet-4-6", True), - ("azure_ai", "gpt-5.6", False), - ("openai", "gpt-5.6", False), - ("gemini", "gemini-3.8-flash", False), - ], -) -def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): - assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected - - @pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): tools = [ From b93d26bc854c0519d60d8fb1643689cbb36589b8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:37:54 -0700 Subject: [PATCH 14/58] fix(responses): drop a malformed prompt_cache_breakpoint under drop_params on the direct /v1/responses path (#45330) * fix(responses): drop a malformed prompt_cache_breakpoint under drop_params on the direct /v1/responses path * refactor(responses): reuse the shared prompt_cache_breakpoint_for_wire helper in the chat bridge * fix(responses): read a non-flag deployment drop_params as unset in the marker walk, like the param mapper * fix(responses): drop every prompt_cache_breakpoint under drop_params when the model has no breakpoints Also drop a marker set on the message item itself, sanitize before the cache_control conversion so a client typo no longer erases that breakpoint, and skip items that carry no marker so unmarked history costs no extra walk * fix(responses): keep a valid prompt_cache_breakpoint unless the model is known to reject it A deployment alias the cost map does not know (an Azure deployment name, a fine-tune, a litellm_proxy alias) kept its valid markers before this PR and must keep them under drop_params. base_model is read from litellm_params first, then model_info, as chat completions does. * fix(responses): drop a valid marker on known non-breakpoint models and keep the tool schema gate unchanged * fix(responses): judge breakpoint support by model name so cost-map aliases keep valid markers * fix(responses): keep the caller's request stripped of cache_control as before when a marker is copied * fix(responses): keep valid prompt cache markers in place and let cache_control fill a dropped one Under drop_params the marker step now hands back the caller's own objects when every marker is already wire-safe, so a request with valid markers is no longer copied and walked a second time. A deployment that turns cache_control into markers keeps the caller's own markers too, whatever its base_model says, so one request never gets two answers about breakpoint support. On /v1/messages a malformed raw marker next to cache_control is dropped before the conversion, so the block gets the cache hint cache_control alone would give, as on /v1/responses. * test(responses): keep a null marker and the original key order on a model that reads breakpoints --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../transformation.py | 25 +- .../anthropic_cache_control_hook.py | 20 +- .../prompt_templates/common_utils.py | 15 +- .../messages/responses_adapters/handler.py | 24 +- .../llms/openai/responses/transformation.py | 164 ++++++++-- ...ses_bridge_prompt_cache_breakpoint_wire.py | 44 ++- ...ore_utils_prompt_templates_common_utils.py | 47 +++ .../test_responses_adapters_handler.py | 22 ++ .../test_openai_responses_transformation.py | 300 ++++++++++++++++++ 9 files changed, 605 insertions(+), 56 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 6173794a210..580e55273f6 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -36,6 +36,7 @@ from litellm.completion_extras.litellm_responses_transformation.history import ( from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( + prompt_cache_breakpoint_for_wire, responses_reasoning_items_from_thinking_blocks, with_prompt_cache_breakpoint, ) @@ -61,7 +62,6 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, - PromptCacheBreakpoint, Reasoning, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, @@ -631,19 +631,6 @@ def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Li return "length" -_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint) - - -def _prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: - if marker is None or not drop_params: - return marker - try: - return _PROMPT_CACHE_BREAKPOINT.validate_python(marker) - except ValidationError: - verbose_logger.debug("Chat provider: dropping malformed prompt_cache_breakpoint %r under drop_params", marker) - return None - - def _input_file_from_file_value(file_value: object) -> dict[str, object]: if not isinstance(file_value, dict): return {"type": "input_file"} @@ -1679,7 +1666,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if original_type == "text": converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(item.get("text", ""), role), - _prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), ) result.append(converted) verbose_logger.debug("Chat provider: text -> %s", converted) @@ -1692,14 +1679,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): cast(ChatCompletionImageObject, item), role ), ), - _prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), ) result.append(converted) verbose_logger.debug("Chat provider: image_url -> %s", converted) elif original_type == "input_audio": converted = with_prompt_cache_breakpoint( {"type": "input_audio", "input_audio": item.get("input_audio")}, - _prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), ) result.append(converted) verbose_logger.debug("Chat provider: input_audio -> %s", converted) @@ -1715,7 +1702,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): _input_file_from_file_value( cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked ), - _prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), ) result.append(converted) verbose_logger.debug("Chat provider: file -> %s", converted) @@ -1739,7 +1726,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Default to input_text for unknown types converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(str(item.get("text", item)), role), - _prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), ) result.append(converted) verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index b0deec24c1a..fc1f1830129 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -72,6 +72,8 @@ MAX_CACHE_CONTROL_BLOCKS: Final = 4 CACHE_BREAKPOINT_KEYS: Final = ("cache_control", "prompt_cache_breakpoint") OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION: Final = (5, 6) _GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+))?") +_FINE_TUNED_BASE_PATTERN: Final = re.compile(r"^ft:([^:]+)") +_O_SERIES_PATTERN: Final = re.compile(r"^o\d+(?:-|$)") OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( {"text", "image", "image_url", "file", "input_audio", "input_text", "input_image", "input_file"} ) @@ -125,6 +127,10 @@ def _served_provider_for_model(model: str) -> str | None: def supports_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None = None) -> bool: + return openai_prompt_cache_breakpoint_support(model, custom_llm_provider) is True + + +def openai_prompt_cache_breakpoint_support(model: str, custom_llm_provider: str | None = None) -> bool | None: hosted_flag: Final = ( None if custom_llm_provider is None @@ -135,9 +141,19 @@ def supports_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model) if model_map_flag is not None: return model_map_flag - version_match: Final = _GPT_VERSION_PATTERN.match(model.rsplit("/", 1)[-1].lower()) - if version_match is None: + fine_tuned_base: Final = _FINE_TUNED_BASE_PATTERN.match(model.rsplit("/", 1)[-1]) + if fine_tuned_base is not None: + return False if _named_model_breakpoint_support(fine_tuned_base.group(1)) is False else None + return _named_model_breakpoint_support(model) + + +def _named_model_breakpoint_support(model: str) -> bool | None: + bare_model: Final = model.rsplit("/", 1)[-1].lower() + if _O_SERIES_PATTERN.match(bare_model) is not None: return False + version_match: Final = _GPT_VERSION_PATTERN.match(bare_model) + if version_match is None: + return None version: Final = (int(version_match.group(1)), int(version_match.group(2) or 0)) return version >= OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 101f069c0f8..9bc53791b36 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -17,7 +17,7 @@ from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger @@ -2073,6 +2073,19 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key +_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint) + + +def prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: + if marker is None or not drop_params: + return marker + try: + return _PROMPT_CACHE_BREAKPOINT.validate_python(marker) + except ValidationError: + verbose_logger.debug("Dropping malformed prompt_cache_breakpoint %r under drop_params", marker) + return None + + def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: """Anthropic's ``cache_control`` mark rewritten as OpenAI's ``prompt_cache_breakpoint`` on the same block.""" if block.get("cache_control") is None: diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index 99a07a1e9bf..9b22170c9db 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -15,7 +15,11 @@ from litellm.integrations.anthropic_cache_control_hook import ( supports_openai_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.core_helpers import drops_unsupported_params -from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + cache_control_as_prompt_cache_breakpoint, + prompt_cache_breakpoint_for_wire, + with_prompt_cache_breakpoint, +) from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( @@ -90,14 +94,28 @@ def _without_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str return {key: value for key, value in block.items() if key != "prompt_cache_breakpoint"} +def _cache_control_as_wire_safe_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str, object]: + if "prompt_cache_breakpoint" not in block: + return cache_control_as_prompt_cache_breakpoint(block) + wire_marker: Final = prompt_cache_breakpoint_for_wire(block["prompt_cache_breakpoint"], drop_params=True) + return cache_control_as_prompt_cache_breakpoint( + with_prompt_cache_breakpoint(_without_prompt_cache_breakpoint(block), wire_marker) + ) + + def _prompt_cache_block_rewrite(model: str, forwarded_kwargs: Mapping[str, object]) -> BlockRewrite | None: """Client cache marks become breakpoints for a deployment that reads them, and a model that takes no breakpoints drops raw ones under drop_params.""" + drops_params: Final = drops_unsupported_params(forwarded_kwargs.get("drop_params")) if _reads_prompt_cache_breakpoint(model, forwarded_kwargs): - return cache_control_as_prompt_cache_breakpoint + return ( + _cache_control_as_wire_safe_prompt_cache_breakpoint + if drops_params + else cache_control_as_prompt_cache_breakpoint + ) if _model_takes_prompt_cache_breakpoint(model, forwarded_kwargs): return None - if drops_unsupported_params(forwarded_kwargs.get("drop_params")): + if drops_params: return _without_prompt_cache_breakpoint return None diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index be40220f94d..903e6177e7b 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Collection, Mapping, Sequence +from collections.abc import Callable, Collection, Iterable, Mapping, Sequence from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints @@ -10,7 +10,11 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook, targets_openai_api +from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + openai_prompt_cache_breakpoint_support, + targets_openai_api, +) from litellm.litellm_core_utils.core_helpers import normalize_drop_params, process_response_headers from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -20,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( cache_control_as_prompt_cache_breakpoint, drop_non_python_regex_patterns, flatten_combinators_and_drop_non_python_regex_patterns, + prompt_cache_breakpoint_for_wire, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -66,6 +71,7 @@ _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, _PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS _ROLES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"system", "developer", "user"}) _PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"input_text", "input_image", "input_file"}) +_INPUT_PART_LIST_KEYS: Final = frozenset({"content", "output"}) _COMPACT_TRACKING_PARAMS: Final = frozenset({"metadata", "user"}) _COMPACT_REJECTED_PARAMS: Final = frozenset( { @@ -93,6 +99,13 @@ _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, _CONTENT_PARTS: Final[TypeAdapter[Sequence[Mapping[str, object]]]] = TypeAdapter(Sequence[Mapping[str, object]]) +class _DeclaredModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + base_model: object = None + model_info: object = None + + class _DeploymentHost(BaseModel): """`responses()` builds the params from raw kwargs, where callers may spell the host `base_url`.""" @@ -216,6 +229,63 @@ def _with_breakpoint_on_last_reading_part( return (*parts[:last], marked_last, *parts[last + 1 :]) +def _json_object_or_none(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + +def _json_array_or_empty(value: object) -> Sequence[object]: + return value if isinstance(value, (list, tuple)) else () + + +def _item_with_wire_safe_prompt_cache_breakpoints(item: object, keeps_markers: bool) -> object: + fields: Final = _json_object_or_none(item) + if fields is None or not _carries_prompt_cache_breakpoint(fields): + return item + wire_safe_fields: Final = { + key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, keeps_markers) + for key, value in fields.items() + if key != "prompt_cache_breakpoint" + } + unchanged: Final = "prompt_cache_breakpoint" not in fields and _same_objects( + wire_safe_fields.values(), fields.values() + ) + return item if unchanged else wire_safe_fields + + +def _same_objects(rebuilt: Iterable[object], originals: Iterable[object]) -> bool: + return all(new is old for new, old in zip(rebuilt, originals, strict=True)) + + +def _carries_prompt_cache_breakpoint(fields: Mapping[str, object]) -> bool: + return "prompt_cache_breakpoint" in fields or any( + _parts_carry_prompt_cache_breakpoint(fields.get(key)) for key in _INPUT_PART_LIST_KEYS + ) + + +def _parts_carry_prompt_cache_breakpoint(parts: object) -> bool: + return any(isinstance(part, Mapping) and "prompt_cache_breakpoint" in part for part in _json_array_or_empty(parts)) + + +def _parts_with_wire_safe_prompt_cache_breakpoints(key: str, value: object, keeps_markers: bool) -> object: + if key not in _INPUT_PART_LIST_KEYS or not _parts_carry_prompt_cache_breakpoint(value): + return value + parts: Final = _json_array_or_empty(value) + wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, keeps_markers) for part in parts) + return value if _same_objects(wire_safe_parts, parts) else list(wire_safe_parts) + + +def _part_with_wire_safe_prompt_cache_breakpoint(part: object, keeps_markers: bool) -> object: + fields: Final = _json_object_or_none(part) + if fields is None or "prompt_cache_breakpoint" not in fields: + return part + marker: Final = fields["prompt_cache_breakpoint"] + wire_marker: Final = prompt_cache_breakpoint_for_wire(marker, drop_params=True) if keeps_markers else None + if keeps_markers and wire_marker == marker: + return part + unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"} + return with_prompt_cache_breakpoint(unmarked, wire_marker) + + class _ReasoningSupportEntry(LiteLLMBaseModel): litellm_provider: str | None = None supports_reasoning: bool | None = None @@ -544,12 +614,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): validated_input: Final = self._validate_input_param( self._drop_bridge_minted_reasoning_items(input), drops_set_created_by=drops_params(litellm_params) ) - marked_input: Final = ( - self._input_with_prompt_cache_breakpoints(validated_input) - if self._reads_prompt_cache_breakpoints( - model=model, litellm_params=litellm_params, prompt_cache_options=prompt_cache_options - ) - else validated_input + reads_breakpoints: Final = self._reads_prompt_cache_breakpoints( + model=model, litellm_params=litellm_params, prompt_cache_options=prompt_cache_options + ) + marker_safe_input: Final = self._input_with_wire_safe_prompt_cache_breakpoints( + model=model, input=validated_input, litellm_params=litellm_params, reads_breakpoints=reads_breakpoints ) maps_anthropic_tool_search: Final = self._targets_openai_api(litellm_params) and any( is_anthropic_tool_search_chat_tool(tool) for tool in tools or () @@ -559,20 +628,71 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): if maps_anthropic_tool_search else tools ) - replayable_input: Final = ( - input_without_unanswered_anthropic_tool_search_calls(marked_input) - if maps_anthropic_tool_search - else marked_input - ) stripped_input, stripped_tools = self.remove_cache_control_flag_from_input_and_tools( - model=model, input=replayable_input, tools=mapped_tools + model=model, + input=self._replayable_input(marker_safe_input, reads_breakpoints, maps_anthropic_tool_search), + tools=mapped_tools, ) + if marker_safe_input is not validated_input: + self._strip_caller_input_as_if_unmarked( + model=model, + caller_input=validated_input, + reads_breakpoints=reads_breakpoints, + maps_anthropic_tool_search=maps_anthropic_tool_search, + ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( model=model, tools=object_schema_tools, litellm_params=litellm_params ) return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools + def _replayable_input( + self, input: str | ResponseInputParam, reads_breakpoints: bool, maps_anthropic_tool_search: bool + ) -> str | ResponseInputParam: + marked_input: Final = self._input_with_prompt_cache_breakpoints(input) if reads_breakpoints else input + if not maps_anthropic_tool_search: + return marked_input + return input_without_unanswered_anthropic_tool_search_calls(marked_input) + + def _strip_caller_input_as_if_unmarked( + self, + model: str, + caller_input: str | ResponseInputParam, + reads_breakpoints: bool, + maps_anthropic_tool_search: bool, + ) -> None: + self.remove_cache_control_flag_from_input_and_tools( + model=model, input=self._replayable_input(caller_input, reads_breakpoints, maps_anthropic_tool_search) + ) + + def _input_with_wire_safe_prompt_cache_breakpoints( + self, + model: str, + input: str | ResponseInputParam, + litellm_params: GenericLiteLLMParams, + reads_breakpoints: bool, + ) -> str | ResponseInputParam: + if not isinstance(input, list) or not drops_params(litellm_params): + return input + keeps_markers: Final = reads_breakpoints or self._keeps_prompt_cache_breakpoints( + model=model, litellm_params=litellm_params + ) + wire_safe_items: Final = tuple( + _item_with_wire_safe_prompt_cache_breakpoints(item, keeps_markers) for item in input + ) + if _same_objects(wire_safe_items, input): + return input + wire_safe_input: Final = list(wire_safe_items) + return cast("ResponseInputParam", wire_safe_input) # cast-ok: items keep their shape, minus a rejected marker + + def _keeps_prompt_cache_breakpoints(self, model: str, litellm_params: GenericLiteLLMParams) -> bool: + provider: Final = _provider_name(self.custom_llm_provider) + base_model: Final = self._declared_base_model(model=model, litellm_params=litellm_params) + known_support: Final = ( + openai_prompt_cache_breakpoint_support(name, provider) for name in dict.fromkeys((base_model, model)) + ) + return next((support for support in known_support if support is not None), None) is not False + def _targets_openai_api(self, litellm_params: GenericLiteLLMParams) -> bool: host: Final = _DeploymentHost.model_validate(litellm_params) return self.custom_llm_provider == LlmProviders.OPENAI and targets_openai_api(host.api_base or host.base_url) @@ -688,14 +808,15 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @staticmethod def _without_foreign_tool_call_item_id(item: object) -> object: - if not isinstance(item, dict): + fields: Final = _json_object_or_none(item) + if fields is None: return item - item_type: Final = item.get("type") - item_id: Final = item.get("id") + item_type: Final = fields.get("type") + item_id: Final = fields.get("id") genuine_prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type) if isinstance(item_type, str) else None if genuine_prefix is None or not isinstance(item_id, str) or item_id.startswith(genuine_prefix): return item - return {key: value for key, value in item.items() if key != "id"} + return {key: value for key, value in fields.items() if key != "id"} def _sanitized_tool_schemas_for_openai( self, @@ -739,6 +860,13 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): base_model: Final[object] = model_info.get("base_model") if isinstance(model_info, dict) else None return base_model if isinstance(base_model, str) and base_model else model + @staticmethod + def _declared_base_model(model: str, litellm_params: GenericLiteLLMParams) -> str: + declared: Final = _DeclaredModel.model_validate(litellm_params) + model_info: Final = _json_object_or_none(declared.model_info) + candidates: Final = (declared.base_model, None if model_info is None else model_info.get("base_model")) + return next((name for name in candidates if isinstance(name, str) and name), model) + @staticmethod def _sanitized_tool_entry( entry: Mapping[str, object], diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py index 12a7bfe68e9..33338cae95d 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py @@ -39,14 +39,6 @@ _CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = ( ) _CASE_IDS: Final = tuple(case[0] for case in _CASES) _CASE_VALUES: Final = tuple(case[1:] for case in _CASES) -_ADAPTER_CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = ( - ("valid-on", "on", pcb.EXPLICIT, pcb.EXPLICIT), - ("valid-off", "off", pcb.EXPLICIT, pcb.EXPLICIT), - ("malformed-on", "on", "yes", "yes"), - ("malformed-off", "off", "yes", "yes"), -) -_ADAPTER_CASE_IDS: Final = tuple(case[0] for case in _ADAPTER_CASES) -_ADAPTER_CASE_VALUES: Final = tuple(case[1:] for case in _ADAPTER_CASES) @dataclass(frozen=True, slots=True) @@ -331,7 +323,7 @@ def _anthropic_text(message: anthropic.types.Message, marker: str) -> None: assert content.type == "text" and content.text == rv.answer(marker), message -@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _ADAPTER_CASE_VALUES, ids=_ADAPTER_CASE_IDS) +@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _CASE_VALUES, ids=_CASE_IDS) def test_anthropic_sdk_marker_on_user_text( bridge: _Bridge, mode: Mode, breakpoint: JsonValue, expected: JsonValue ) -> None: @@ -540,14 +532,17 @@ def test_anthropic_sdk_cache_control_on_an_assistant_turn_is_dropped(bridge: _Br ], "OpenAI accepts a breakpoint on output_text but writes no cache for it (live gpt-5.6, 2026-10-12)" -def test_native_responses_request_never_enters_the_bridge(bridge: _Bridge) -> None: +@pytest.mark.parametrize( + ("mode", "expected"), (("on", pcb.EXPLICIT), ("off", _UNKNOWN_KEY)), ids=("normalized-on", "verbatim-off") +) +def test_native_responses_marker_with_an_unknown_key(bridge: _Bridge, mode: Mode, expected: JsonValue) -> None: marker: Final = uuid.uuid4().hex block: Final[dict[str, JsonValue]] = {"type": "input_text", "text": pcb.prompt(marker)} response: Final = bridge.gateway.request( "POST", "/v1/responses", { - "model": bridge.on, + "model": bridge.model(mode), "input": [{"type": "message", "role": "user", "content": [pcb.marked(block, _UNKNOWN_KEY)]}], **pcb.NO_CACHE, }, @@ -559,8 +554,31 @@ def test_native_responses_request_never_enters_the_bridge(bridge: _Bridge) -> No request: Final = pcb.posted(bridge.wire, marker) _wire_body(request) on_wire: Final = pcb.single_block(pcb.input_items(request), "user") - assert on_wire == pcb.marked(block, _UNKNOWN_KEY), on_wire - bridge.spend.landed(bridge.on, response.headers["x-litellm-call-id"], marker) + assert on_wire == pcb.marked(block, expected), on_wire + bridge.spend.landed(bridge.model(mode), response.headers["x-litellm-call-id"], marker) + + +@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _CASE_VALUES, ids=_CASE_IDS) +def test_openai_sdk_native_responses_marker_on_user_text( + bridge: _Bridge, mode: Mode, breakpoint: JsonValue, expected: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + block: Final[dict[str, JsonValue]] = {"type": "input_text", "text": pcb.prompt(marker)} + with openai.OpenAI(api_key=bridge.gateway.key, base_url=_v1(bridge.gateway), max_retries=0) as client: + raw: Final = client.responses.with_raw_response.create( + model=bridge.model(mode), + input=[{"type": "message", "role": "user", "content": [pcb.marked(block, breakpoint)]}], + extra_body=dict(pcb.NO_CACHE), + ) + response: Final = raw.parse() + assert pcb.answers(response.id, marker), response + assert response.output_text == rv.answer(marker), response + request: Final = pcb.posted(bridge.wire, marker) + _wire_body(request) + on_wire: Final = pcb.single_block(pcb.input_items(request), "user") + assert on_wire["type"] == "input_text" and on_wire["text"] == pcb.prompt(marker), on_wire + pcb.assert_marker(on_wire, expected) + bridge.spend.landed(bridge.model(mode), raw.headers["x-litellm-call-id"], marker) @pytest.mark.parametrize("breakpoint", _MALFORMED_VALUES, ids=_MALFORMED_IDS) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index b9675f6fd5c..72e72714cd4 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( merge_consecutive_system_messages, native_redacted_thinking_data, parse_tool_call_arguments, + prompt_cache_breakpoint_for_wire, responses_reasoning_items_from_thinking_blocks, salvage_concatenated_tool_arguments, schema_grammar_branch_count, @@ -2328,6 +2329,52 @@ def test_get_semantic_cache_prompt_from_messages_matches_get_str_from_messages_w assert get_semantic_cache_prompt_from_messages(messages) == get_str_from_messages(messages) +_MALFORMED_PROMPT_CACHE_BREAKPOINTS: Final = ( + "yes", + 1, + ["explicit"], + "", + {}, + {"mode": "bogus"}, + {"mode": "explicit", "ttl": "1h"}, +) +_MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = ( + "string", + "int", + "list", + "empty-string", + "empty-object", + "bogus-mode", + "bad-ttl", +) +_VALID_PROMPT_CACHE_BREAKPOINTS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) + + +@pytest.mark.parametrize("drop_params", (True, False), ids=("drop", "keep")) +def test_prompt_cache_breakpoint_for_wire_passes_a_missing_marker_through(drop_params: bool) -> None: + assert prompt_cache_breakpoint_for_wire(None, drop_params) is None + + +@pytest.mark.parametrize("marker", _MALFORMED_PROMPT_CACHE_BREAKPOINTS, ids=_MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS) +def test_prompt_cache_breakpoint_for_wire_drops_a_malformed_marker_under_drop_params(marker: object) -> None: + assert prompt_cache_breakpoint_for_wire(marker, True) is None + + +@pytest.mark.parametrize("marker", _MALFORMED_PROMPT_CACHE_BREAKPOINTS, ids=_MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS) +def test_prompt_cache_breakpoint_for_wire_forwards_a_malformed_marker_without_drop_params(marker: object) -> None: + assert prompt_cache_breakpoint_for_wire(marker, False) is marker + + +@pytest.mark.parametrize("drop_params", (True, False), ids=("drop", "keep")) +@pytest.mark.parametrize("marker", _VALID_PROMPT_CACHE_BREAKPOINTS, ids=("explicit", "explicit-30m")) +def test_prompt_cache_breakpoint_for_wire_keeps_a_valid_marker(marker: dict[str, str], drop_params: bool) -> None: + assert prompt_cache_breakpoint_for_wire(marker, drop_params) == marker + + +def test_prompt_cache_breakpoint_for_wire_strips_an_unknown_key_under_drop_params() -> None: + assert prompt_cache_breakpoint_for_wire({"mode": "explicit", "note": "kept"}, True) == {"mode": "explicit"} + + def test_cache_control_is_replaced_by_an_explicit_prompt_cache_breakpoint_on_the_same_block() -> None: marked: Final = {"type": "text", "text": "rules", "cache_control": {"type": "ephemeral", "ttl": "1h"}} diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 7de722e1171..bba5413fbdf 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -556,6 +556,28 @@ def test_raw_prompt_cache_breakpoints_reach_a_deployment_unless_drop_params_drop ] +def _gpt_5_6_input_for_a_cached_block(drop_params: bool, raw_mark: Mapping[str, object]) -> object: + responses_kwargs: Final = _build_responses_kwargs( + model="openai/gpt-5.6", + extra_kwargs={**OPENAI_API, "drop_params": drop_params}, + max_tokens=1024, + messages=[{"role": "user", "content": [{"type": "text", "text": "look", "cache_control": EPHEMERAL, **raw_mark}]}], + ) + return responses_kwargs["input"] + + +def test_a_malformed_raw_marker_beside_cache_control_gives_way_to_it_only_under_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + malformed: Final = {"prompt_cache_breakpoint": "yes"} + + assert _gpt_5_6_input_for_a_cached_block(True, malformed) == _gpt_5_6_input_for_a_cached_block(True, {}) + assert _gpt_5_6_input_for_a_cached_block(False, malformed) == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "look", **malformed}]} + ] + + @pytest.mark.parametrize( ("model", "extra_kwargs"), [ diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 48bb0955f34..a185cdaf718 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -25,6 +25,7 @@ from litellm.types.llms.openai import ( OutputTextDeltaEvent, ResponseAPIUsage, ResponseCompletedEvent, + ResponseInputParam, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) @@ -4246,6 +4247,305 @@ async def test_extra_body_merges_with_request_data(extra_body_mock_response_data assert request_body["custom_field"] == "custom_value" +_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "explicit", "ttl": "1h"}) +_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bad-ttl") +_VALID_MARKERS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) +_VALID_MARKER_IDS: Final = ("explicit", "explicit-30m") +_PARTS: Final[Mapping[str, Mapping[str, object]]] = { + "input_text": {"type": "input_text", "text": "Say ok"}, + "input_image": {"type": "input_image", "image_url": "https://example.com/a.png"}, + "input_file": {"type": "input_file", "file_id": "file-abc"}, +} +_CONFIGS: Final = (OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig()) +_CONFIG_IDS: Final = ("openai", "azure") + + +def _marked_message(part_type: str, marker: object) -> list[dict[str, object]]: + return [{"type": "message", "role": "user", "content": [{**_PARTS[part_type], "prompt_cache_breakpoint": marker}]}] + + +def _transformed_input( + config: BaseResponsesAPIConfig, input: str | list[dict[str, object]], litellm_params: GenericLiteLLMParams +) -> object: + request: Final = config.transform_responses_api_request( + model="gpt-6.1-sol", + input=cast("ResponseInputParam", input), # cast-ok: plain dicts shaped like the SDK's input items + response_api_optional_request_params={}, + litellm_params=litellm_params, + headers={}, + ) + return request["input"] + + +def _single_part(transformed_input: object) -> Mapping[str, object]: + (item,) = cast(list[Mapping[str, object]], transformed_input) + (part,) = cast(list[Mapping[str, object]], item["content"]) + return part + + +@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS) +@pytest.mark.parametrize("part_type", tuple(_PARTS)) +@pytest.mark.parametrize("marker", _MALFORMED_MARKERS, ids=_MALFORMED_MARKER_IDS) +def test_malformed_prompt_cache_breakpoint_is_dropped_under_deployment_drop_params( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, part_type: str, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + transformed: Final = _transformed_input( + config, _marked_message(part_type, marker), GenericLiteLLMParams(drop_params=True) + ) + assert _single_part(transformed) == _PARTS[part_type] + + +@pytest.mark.parametrize("marker", _MALFORMED_MARKERS, ids=_MALFORMED_MARKER_IDS) +def test_malformed_prompt_cache_breakpoint_is_dropped_under_global_drop_params( + monkeypatch: pytest.MonkeyPatch, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + transformed: Final = _transformed_input(_CONFIGS[0], _marked_message("input_text", marker), GenericLiteLLMParams()) + assert _single_part(transformed) == _PARTS["input_text"] + + +@pytest.mark.parametrize("drop_params", (False, None), ids=("false", "null")) +@pytest.mark.parametrize("marker", _MALFORMED_MARKERS, ids=_MALFORMED_MARKER_IDS) +def test_malformed_prompt_cache_breakpoint_is_forwarded_verbatim_without_drop_params( + monkeypatch: pytest.MonkeyPatch, marker: object, drop_params: bool | None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=drop_params)) == input + + +@pytest.mark.parametrize("marker", _MALFORMED_MARKERS, ids=_MALFORMED_MARKER_IDS) +def test_a_non_flag_deployment_drop_params_keeps_the_marker_like_the_param_mapper_does( + monkeypatch: pytest.MonkeyPatch, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params="maybe")) == input + + +@pytest.mark.parametrize("drop_params", (True, False), ids=("drop", "keep")) +@pytest.mark.parametrize("marker", _VALID_MARKERS, ids=_VALID_MARKER_IDS) +def test_valid_prompt_cache_breakpoint_reaches_the_wire_under_both_settings( + monkeypatch: pytest.MonkeyPatch, marker: Mapping[str, str], drop_params: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=drop_params)) == input + + +def test_prompt_cache_breakpoint_unknown_key_is_trimmed_under_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", {"mode": "explicit", "note": "kept"}) + dropped: Final = _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(dropped)["prompt_cache_breakpoint"] == {"mode": "explicit"} + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=False)) == input + + +def test_malformed_prompt_cache_breakpoint_on_a_tool_output_part_is_dropped_under_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + marked_output: Final = [{"type": "input_text", "text": "42", "prompt_cache_breakpoint": "yes"}] + input: Final[list[dict[str, object]]] = [{"type": "function_call_output", "call_id": "call_1", "output": marked_output}] + (item,) = cast( + list[Mapping[str, object]], _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=True)) + ) + assert item["output"] == [{"type": "input_text", "text": "42"}] + + +@pytest.mark.parametrize("input", ("hi", [{"type": "message", "role": "user", "content": "hi"}]), ids=("string", "unmarked")) +def test_input_without_a_marker_is_untouched_under_drop_params( + monkeypatch: pytest.MonkeyPatch, input: str | list[dict[str, object]] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=True)) == input + + +def _transformed_input_for( + model: str, + input: list[dict[str, object]], + litellm_params: GenericLiteLLMParams, + config: BaseResponsesAPIConfig = _CONFIGS[0], +) -> object: + request: Final = config.transform_responses_api_request( + model=model, + input=cast("ResponseInputParam", input), # cast-ok: plain dicts shaped like the SDK's input items + response_api_optional_request_params={}, + litellm_params=litellm_params, + headers={}, + ) + return request["input"] + + +@pytest.mark.parametrize( + "marker", (*_VALID_MARKERS, *_MALFORMED_MARKERS), ids=(*_VALID_MARKER_IDS, *_MALFORMED_MARKER_IDS) +) +def test_every_prompt_cache_breakpoint_is_dropped_under_drop_params_for_a_model_without_breakpoints( + monkeypatch: pytest.MonkeyPatch, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + dropped: Final = _transformed_input_for("gpt-5.4-mini", input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(dropped) == _PARTS["input_text"] + assert _transformed_input_for("gpt-5.4-mini", input, GenericLiteLLMParams(drop_params=False)) == input + + +@pytest.mark.parametrize( + "model", + ("o4-mini", "ft:gpt-4o-mini-2024-07-18:my-org::abc123", "ft:o4-mini-2025-04-16:my-org::abc123"), + ids=("o-series", "fine-tune-of-older-gpt", "fine-tune-of-o-series"), +) +def test_a_valid_marker_is_dropped_under_drop_params_on_a_model_known_to_lack_breakpoints( + monkeypatch: pytest.MonkeyPatch, model: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", {"mode": "explicit"}) + dropped: Final = _transformed_input_for(model, input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(dropped) == _PARTS["input_text"] + + +@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS) +def test_a_cost_map_entry_registered_under_the_deployment_name_keeps_a_valid_marker_under_drop_params( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + registered: Final = {"id": "litellm-qa-sol-alias", "supports_vision": True, "input_cost_per_token": 1e-6} + monkeypatch.setitem(litellm.model_cost, "litellm-qa-sol-alias", registered) + valid: Final = _marked_message("input_text", {"mode": "explicit"}) + assert _transformed_input_for("litellm-qa-sol-alias", valid, GenericLiteLLMParams(drop_params=True), config) == valid + + +@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS) +def test_a_valid_marker_survives_drop_params_on_a_deployment_name_with_unknown_breakpoint_support( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + litellm_params: Final = GenericLiteLLMParams(drop_params=True) + valid: Final = _marked_message("input_text", {"mode": "explicit"}) + malformed: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "1h"}) + assert _transformed_input_for("litellm-qa-sol-alias", valid, litellm_params, config) == valid + dropped: Final = _transformed_input_for("litellm-qa-sol-alias", malformed, litellm_params, config) + assert _single_part(dropped) == _PARTS["input_text"] + + +def _declaring_base_model(where: str, base_model: str) -> GenericLiteLLMParams: + if where == "litellm_params": + return GenericLiteLLMParams(drop_params=True, base_model=base_model) + return GenericLiteLLMParams(drop_params=True, model_info={"base_model": base_model}) + + +@pytest.mark.parametrize("where", ("litellm_params", "model_info")) +def test_a_declared_base_model_decides_whether_a_valid_marker_survives_drop_params( + monkeypatch: pytest.MonkeyPatch, where: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", {"mode": "explicit"}) + kept: Final = _transformed_input_for("gpt-5.4-mini-prod", input, _declaring_base_model(where, "gpt-6.1-sol")) + assert kept == input + dropped: Final = _transformed_input_for("prod-deployment", input, _declaring_base_model(where, "gpt-5.4-mini")) + assert _single_part(dropped) == _PARTS["input_text"] + + +def test_a_deployment_that_marks_cache_control_keeps_the_callers_own_marker_whatever_its_base_model_says( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + cached: Final = {"type": "input_text", "text": "a", "cache_control": {"type": "ephemeral"}} + marked: Final = {"type": "input_text", "text": "b", "prompt_cache_breakpoint": {"mode": "explicit"}} + input: Final = [{"type": "message", "role": "user", "content": [cached, marked]}] + transformed: Final = _transformed_input_for("gpt-6.1-sol", input, _declaring_base_model("model_info", "gpt-5.5")) + (item,) = cast(list[Mapping[str, object]], transformed) + parts: Final = cast(list[Mapping[str, object]], item["content"]) + assert [part.get("prompt_cache_breakpoint") for part in parts] == [{"mode": "explicit"}, {"mode": "explicit"}] + + +def test_a_model_that_takes_markers_sends_kept_parts_byte_for_byte_and_drops_only_the_rejected_marker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + kept: Final = [ + {"type": "input_text", "text": "a", "prompt_cache_breakpoint": None}, + {"prompt_cache_breakpoint": {"mode": "explicit"}, "type": "input_text", "text": "b"}, + ] + rejected: Final = {"type": "input_text", "text": "c"} + bad_ttl: Final = {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}} + input: Final = [ + {"type": "message", "role": "user", "content": kept}, + {"type": "message", "role": "user", "content": [{**bad_ttl, **rejected}]}, + ] + transformed: Final = _transformed_input_for("gpt-6.1-sol", input, GenericLiteLLMParams(drop_params=True)) + assert json.dumps(transformed) == json.dumps( + [{"type": "message", "role": "user", "content": kept}, {"type": "message", "role": "user", "content": [rejected]}] + ) + + +@pytest.mark.parametrize("content", ("Say ok", [{"type": "input_text", "text": "Say ok"}]), ids=("string", "parts")) +@pytest.mark.parametrize("marker", (_VALID_MARKERS[0], "yes"), ids=("explicit", "string")) +def test_a_prompt_cache_breakpoint_on_the_message_item_is_dropped_under_drop_params( + monkeypatch: pytest.MonkeyPatch, content: object, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + unmarked: Final[dict[str, object]] = {"type": "message", "role": "user", "content": content} + input: Final = [{**unmarked, "prompt_cache_breakpoint": marker}] + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=True)) == [unmarked] + assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=False)) == input + + +def test_cache_control_still_marks_a_part_whose_own_marker_is_malformed_under_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + part: Final = {"type": "input_text", "text": "Say ok", "cache_control": {"type": "ephemeral"}} + input: Final = [{"type": "message", "role": "user", "content": [{**part, "prompt_cache_breakpoint": "yes"}]}] + transformed: Final = _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(transformed) == { + "type": "input_text", + "text": "Say ok", + "prompt_cache_breakpoint": {"mode": "explicit"}, + } + + +def _caller_input_after_transform(model: str, marker: object) -> str: + hint: Final = {"type": "ephemeral"} + marked: Final = {} if marker is None else {"prompt_cache_breakpoint": marker} + input: Final = [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Say ok", "cache_control": hint, **marked}, + {"type": "input_text", "text": "Thanks", "cache_control": hint}, + ], + "cache_control": hint, + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "cache_control": hint, **marked}], + }, + ] + _CONFIGS[0].transform_responses_api_request( + model=model, + input=cast("ResponseInputParam", input), # cast-ok: plain dicts shaped like the SDK's input items + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(drop_params=True), + headers={}, + ) + return json.dumps(input) + + +@pytest.mark.parametrize("model", ("gpt-4.1", "gpt-6.1-sol")) +@pytest.mark.parametrize("marker", (_VALID_MARKERS[0], "yes"), ids=("explicit", "malformed")) +def test_a_marker_leaves_the_caller_request_that_fallbacks_and_logs_read_as_an_unmarked_one_would( + monkeypatch: pytest.MonkeyPatch, model: str, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + marker_json: Final = f', "prompt_cache_breakpoint": {json.dumps(marker)}' + unmarked: Final = _caller_input_after_transform(model, None) + assert _caller_input_after_transform(model, marker).replace(marker_json, "") == unmarked + + @pytest.mark.asyncio @pytest.mark.usefixtures(httpx_transport.__name__) async def test_aresponses_forwards_previous_response_id_to_openai() -> None: From e85cb4052e3b99f2847ef55633fb43a7c6ed210d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:46:32 -0700 Subject: [PATCH 15/58] test: fix flaky spend-log cleanup budget, login lockout, and wedged Redis tests (#47108) * test: stabilize cleanup, lockout, and post-kill timing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): bound the wedged Redis burst p95 by the socket timeout the test sets The burst p95 under a fixed 0.75s cap tracked how fast the box ran the 30 concurrent writes, so a loaded runner failed it at unchanged code. A handler that waits on the wedged publish takes at least the coordination Redis socket timeout, so the test now sets that timeout and asserts the p95 stays under it. * test: pin the lockout clock and bound the wedged Redis burst by the proxy's 2s wait cap The lockout tests now set the login throttle's clock instead of reading the real one, so Retry-After is pinned at 77 and a round-down fails on every run The wedged Redis burst p95 is now bounded by the 2s cap the proxy puts on its own Redis waits, so a burst wait capped under the 5s socket timeout still fails the test The fresh post-kill clients keep the harness's 2s keep-alive expiry, here and in the two tests #47080 changed * test: tell ceil from round in the lockout tests and bound the burst by the imported eviction cap The refused attempt now runs with 76.4s left on the 77s block, where ceil gives 77 but round and int give 76, so both lockout tests fail if Retry-After stops rounding up. The burst bound imports _BULK_EVICTION_WAIT_SECONDS instead of copying its 2s value, so lowering the cap moves the bound with it. * test: bound the wedged Redis burst by the shorter of the eviction cap and the socket timeout Any wait the proxy puts on the stopped Redis lasts at least the shorter of the two, so the bound now follows whichever one a later change lowers. * test: drive the tool index cleanup test with the manual clock and keep the post-kill block as on main The post-kill reset is already fixed on main by making kill() wait for the worker's last thread, so the fresh-client and keep-alive edits go back to main's version. The tool index cleanup test now advances the injected clock instead of timing the real one, and the lockout clock helper gets a type hint. * test: advance the spend log cleanup tests' manual clock through its own method * test: freeze the throttle clock in the login_utils lockout tests and drive the run budget tests with the manual clock --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 60 +++++++------- .../coordination_redis_proxy_config.yaml | 1 + ..._user_updates_wedged_coordination_redis.py | 9 ++- tests/unit/proxy/auth/test_login_utils.py | 15 ++++ .../proxy_server/test_routes_login_sso.py | 22 +++--- tests/unit/proxy/test_spend_log_cleanup.py | 79 +++++++++++++------ 6 files changed, 123 insertions(+), 63 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 7f5efcdaec3..170eb450339 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -421,7 +421,10 @@ class SpendLogCleanup: general_settings=None, redis_cache: RedisCache | None = None, partition_manager: SpendLogsPartitionManager | None = None, + *, + clock: Callable[[], float] = time.monotonic, ): + self._clock = clock self.retention_seconds: int | None = None self.partition_manager = partition_manager or SpendLogsPartitionManager() from litellm.proxy.proxy_server import general_settings as default_settings @@ -586,7 +589,7 @@ class SpendLogCleanup: async def _probe_storage(self, prisma_client: PrismaClient, deadline: float) -> StorageProbe | None: """Read what the spend-log tables weigh from the catalog, or None when the probe fails or has no budget.""" - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -624,7 +627,7 @@ class SpendLogCleanup: stale.allocated_bytes, _ANALYZE_SQL, ) - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -653,7 +656,7 @@ class SpendLogCleanup: async def _load_size_state(self, prisma_client: PrismaClient, deadline: float) -> SizeCleanupState | None: """The reclaimable count saved by the last check, or None on the first check or when it cannot be read.""" - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -686,7 +689,7 @@ class SpendLogCleanup: the tool-index rows that expire with their spend logs. When fewer rows than that remain, the step is everything up to now. """ - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) async with db_span("size_step_cutoff", "LiteLLM_SpendLogs"), bounded_tx(prisma_client, timeout_ms) as tx: @@ -718,11 +721,11 @@ class SpendLogCleanup: Interpolating this into SQL is safe by construction: an int cannot carry SQL, and SET does not accept a bind parameter. """ - remaining_ms: Final = int((deadline - time.monotonic()) * 1000) + remaining_ms: Final = int((deadline - self._clock()) * 1000) return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms)) @staticmethod - def _group_deadline(overall_deadline: float, groups_remaining: int) -> float: + def _group_deadline(overall_deadline: float, groups_remaining: int, now: float) -> float: """ Give each pending cleanup group an equal share of the time left. @@ -731,10 +734,9 @@ class SpendLogCleanup: """ if groups_remaining == 1: return overall_deadline - current_time: Final = time.monotonic() - if current_time >= overall_deadline: + if now >= overall_deadline: return overall_deadline - return current_time + (overall_deadline - current_time) / groups_remaining + return now + (overall_deadline - now) / groups_remaining def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs: """ @@ -748,7 +750,7 @@ class SpendLogCleanup: """ def remaining() -> int | None: - return None if time.monotonic() >= deadline else self._timeout_ms(deadline) + return None if self._clock() >= deadline else self._timeout_ms(deadline) return remaining @@ -841,7 +843,7 @@ class SpendLogCleanup: run_count = 0 consecutive_failures = 0 while True: - if time.monotonic() >= deadline: + if self._clock() >= deadline: verbose_proxy_logger.info( "Run budget exhausted during %s cleanup after %d rows; the next run resumes from here", table_name, @@ -858,13 +860,13 @@ class SpendLogCleanup: prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline ) # Find rows and delete them in one go without fetching to application - batch_started_at = time.monotonic() + batch_started_at = self._clock() try: batch_result = await self._execute_delete_batch( prisma_client, delete_sql, cutoff_date, table_name, deadline ) except Exception as batch_exc: - if time.monotonic() >= deadline: + if self._clock() >= deadline: # The statement timeout was clamped to the budget that was # left, so this batch was cancelled by the deadline itself. # That is the bound working, not a database fault, and @@ -920,7 +922,7 @@ class SpendLogCleanup: consecutive_failures = 0 deleted_count = batch_result - SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at) + SpendLogCleanupMetrics.record_batch(table_name, deleted_count, self._clock() - batch_started_at) verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name) if deleted_count == 0: @@ -957,7 +959,7 @@ class SpendLogCleanup: operator the backlog was not drained; the gauge simply keeps its value from the last run that finished inside its budget. """ - if time.monotonic() >= deadline: + if self._clock() >= deadline: return TableCleanupResult(table_name=table_name, rows_deleted=rows_deleted, stop_reason=stop_reason) remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline) if remaining is not None: @@ -1070,7 +1072,7 @@ class SpendLogCleanup: # behind any long-running reader for as long as that reader lives, which # is the one way this job could still outlast its budget without bound. remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) - if time.monotonic() >= deadline: + if self._clock() >= deadline: verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run") elif self.general_settings.get( "use_spend_logs_partitioning", False @@ -1082,7 +1084,7 @@ class SpendLogCleanup: verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped) logs_result: Final = await self._delete_old_logs( - prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2) + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock()) ) verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted) @@ -1210,7 +1212,7 @@ class SpendLogCleanup: self, prisma_client: PrismaClient, deadline: float, still_over_target: Callable[[], Awaitable[bool]] ) -> tuple[str, ...]: remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) - if time.monotonic() >= deadline or not self._setting("use_spend_logs_partitioning"): + if self._clock() >= deadline or not self._setting("use_spend_logs_partitioning"): return () if not await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms): return () @@ -1257,7 +1259,7 @@ class SpendLogCleanup: needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted) if needed <= 0: return state.stopped("target_reached") - if time.monotonic() >= deadline: + if self._clock() >= deadline: return state.stopped("budget_exhausted") if state.batches_left <= 0: return state.stopped("batch_cap_reached") @@ -1318,11 +1320,11 @@ class SpendLogCleanup: except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job verbose_proxy_logger.warning("Auto-router baseline retention remains pending") sessions_result: Final = await self._delete_old_autorouter_session_rows( - prisma_client, session_cutoff, self._group_deadline(deadline, 3) + prisma_client, session_cutoff, self._group_deadline(deadline, 3, self._clock()) ) verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) user_sessions_result: Final = await self._delete_old_autorouter_user_session_rows( - prisma_client, session_cutoff, self._group_deadline(deadline, 2) + prisma_client, session_cutoff, self._group_deadline(deadline, 2, self._clock()) ) verbose_proxy_logger.info( "Deleted %s expired auto-router user session rollup rows", user_sessions_result.rows_deleted @@ -1443,7 +1445,7 @@ class SpendLogCleanup: async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None: lock_acquired = False - run_started_at: Final = time.monotonic() + run_started_at: Final = self._clock() progress: Final = _RunProgress() progress_token: Final = _run_progress.set(progress) try: @@ -1493,7 +1495,7 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("skipped_locked") return - deadline: Final = time.monotonic() + self.run_budget_seconds + deadline: Final = self._clock() + self.run_budget_seconds size_group: Final = int(active_size_rule is not None) configured_group_count: Final = ( int(delete_spend_logs and self.retention_seconds is not None) @@ -1506,7 +1508,7 @@ class SpendLogCleanup: spend_log_results: Final = ( await self._clean_spend_log_tables( prisma_client, - self._group_deadline(deadline, configured_group_count), + self._group_deadline(deadline, configured_group_count, self._clock()), ) if delete_spend_logs and self.retention_seconds is not None else () @@ -1526,7 +1528,7 @@ class SpendLogCleanup: await self._clean_session_rollup( prisma_client, autorouter_retention_seconds, - self._group_deadline(deadline, remaining_groups_after_spend_logs), + self._group_deadline(deadline, remaining_groups_after_spend_logs, self._clock()), ) if autorouter_retention_seconds is not None else () @@ -1535,7 +1537,7 @@ class SpendLogCleanup: await self._clean_health_checks( prisma_client, health_check_retention_seconds, - self._group_deadline(deadline, remaining_groups_after_sessions), + self._group_deadline(deadline, remaining_groups_after_sessions, self._clock()), ) if health_check_retention_seconds is not None else () @@ -1544,7 +1546,7 @@ class SpendLogCleanup: await self._clean_daily_tag_spend( prisma_client, daily_tag_spend_retention_seconds, - self._group_deadline(deadline, 1 + size_group), + self._group_deadline(deadline, 1 + size_group, self._clock()), ) if daily_tag_spend_retention_seconds is not None else () @@ -1561,12 +1563,12 @@ class SpendLogCleanup: ) outcome: Final = self._run_outcome(results) SpendLogCleanupMetrics.record_run(outcome) - self._log_run_summary(outcome, results, time.monotonic() - run_started_at, size_report) + self._log_run_summary(outcome, results, self._clock() - run_started_at, size_report) except asyncio.CancelledError: verbose_proxy_logger.error( "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", - time.monotonic() - run_started_at, + self._clock() - run_started_at, progress.rows_deleted, progress.batches, ) diff --git a/tests/integration/coordination_redis_proxy_config.yaml b/tests/integration/coordination_redis_proxy_config.yaml index e28e0f5f2ad..49a7a24e0ac 100644 --- a/tests/integration/coordination_redis_proxy_config.yaml +++ b/tests/integration/coordination_redis_proxy_config.yaml @@ -9,5 +9,6 @@ general_settings: coordination_redis: host: os.environ/REDIS_HOST port: os.environ/REDIS_PORT + socket_timeout: os.environ/COORDINATION_REDIS_SOCKET_TIMEOUT router_settings: disable_cooldowns: true diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 7934e927d93..2b05d35d553 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -16,6 +16,7 @@ from pydantic import JsonValue from redis import Redis from redis.client import PubSub +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import _BULK_EVICTION_WAIT_SECONDS from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from tests.integration._support.database import read_rows @@ -27,6 +28,8 @@ _BURST: Final = 30 _WEDGED_PUBLISHES_PER_BURST_USER: Final = 2 _HANDLER_BUDGET_SECONDS: Final = 0.75 _BULK_BUDGET_SECONDS: Final = 2.0 +_COORDINATION_SOCKET_TIMEOUT_SECONDS: Final = 5.0 +_SHORTEST_REDIS_WAIT_SECONDS: Final = min(_BULK_EVICTION_WAIT_SECONDS, _COORDINATION_SOCKET_TIMEOUT_SECONDS) _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" @@ -102,6 +105,7 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( "DATABASE_URL": database_url, "REDIS_HOST": coordination.host, "REDIS_PORT": str(coordination.port), + "COORDINATION_REDIS_SOCKET_TIMEOUT": str(_COORDINATION_SOCKET_TIMEOUT_SECONDS), }, config=Path("tests/integration/coordination_redis_proxy_config.yaml"), remove_environment=("DATABASE_URL_READ_REPLICA",), @@ -235,8 +239,9 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200 ) timings["c1_burst_p95"] = elapsed_sorted[int(len(elapsed_sorted) * 0.95) - 1] - assert timings["c1_burst_p95"] < _HANDLER_BUDGET_SECONDS, ( - f"burst p95 {timings['c1_burst_p95']:.3f}s" + assert timings["c1_burst_p95"] < _SHORTEST_REDIS_WAIT_SECONDS, ( + f"burst p95 {timings['c1_burst_p95']:.3f}s reached the {_SHORTEST_REDIS_WAIT_SECONDS}s " + "shortest wait the proxy puts on Redis, so burst updates waited on the wedged Redis" ) eventually( lambda: candidate.request("GET", "/health/liveliness").status_code, diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index 37ea05f8ba9..5148e38aa61 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -10,6 +10,7 @@ import os from collections.abc import Mapping from contextlib import ExitStack from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -678,6 +679,12 @@ def _throttle( ) +def _set_throttle_clock(monkeypatch: pytest.MonkeyPatch, now: float) -> None: + from litellm.proxy.auth import login_throttle + + monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: now)) + + def _stores(): from litellm.caching.in_memory_cache import InMemoryCache @@ -756,12 +763,14 @@ async def test_too_many_failures_for_one_username_block_that_pair_and_carry_retr monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=2, block_seconds=77) keys = throttle._keys("admin") assert [await _fail(throttle) for _ in range(3)] == ["401", "401", "401"], "the limit itself is a plain 401" assert throttle._local_block_ttl(keys.pair_block) == 77 + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as blocked: await _guess(throttle) assert blocked.value.code == "429" @@ -814,10 +823,12 @@ async def test_a_correct_password_is_refused_while_its_pair_is_blocked(monkeypat monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") monkeypatch.setenv("DATABASE_URL", "postgresql://stub") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=1, block_seconds=90) assert [await _fail(throttle, username="user@corp.com") for _ in range(3)] == ["401", "401", "429"] + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as refused: await _db_login(throttle, "user@corp.com", "right", correct=True) assert refused.value.code == "429" @@ -1328,12 +1339,14 @@ async def test_a_source_block_outranks_a_pair_block_in_the_retry_after(monkeypat monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=1, source_limit=3, block_seconds=120, client_ip="203.0.113.45") assert [await _fail(throttle) for _ in range(2)] == ["401", "401"], "the admin pair is now blocked" throttle.blocks.set_cache(throttle._keys("admin").pair_block, 1, ttl=30) assert [await _fail(throttle, username=f"spray-{i}@corp.com") for i in range(3)] == ["401"] * 3 assert throttle._local_block_ttl(throttle._keys("admin").source_block) == 120, "the source is now blocked too" + _set_throttle_clock(monkeypatch, 1_000.6) for name in ("admin", "spray-0@corp.com", "never-seen@corp.com"): with pytest.raises(ProxyException) as refused: await _guess(throttle, username=name) @@ -1457,6 +1470,7 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=2, block_seconds=300, redis_cache=_DownRedis()) records: list[logging.LogRecord] = [] @@ -1465,6 +1479,7 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch verbose_proxy_logger.addHandler(handler) try: assert [await _fail(throttle) for _ in range(3)] == ["401"] * 3 + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as blocked: await _guess(throttle) finally: diff --git a/tests/unit/proxy/proxy_server/test_routes_login_sso.py b/tests/unit/proxy/proxy_server/test_routes_login_sso.py index 43629b926ad..a6de5c7b6a0 100644 --- a/tests/unit/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/unit/proxy/proxy_server/test_routes_login_sso.py @@ -11,6 +11,7 @@ Routes covered: from __future__ import annotations from datetime import datetime, timezone +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -504,6 +505,12 @@ def test_login_form_ignores_open_redirect_return_to(client, monkeypatch): # --------------------------------------------------------------------------- +def _set_throttle_clock(monkeypatch: pytest.MonkeyPatch, now: float) -> None: + from litellm.proxy.auth import login_throttle + + monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: now)) + + def _install_real_auth(monkeypatch, **settings): """Run the real authenticate_user so the throttle inside it is exercised. @@ -583,11 +590,12 @@ def test_budget_is_shared_across_username_casing(client, monkeypatch, reset_logi def test_a_refused_attempt_carries_retry_after(client, monkeypatch, reset_login_throttle): - """The 429 tells the caller how long the block has left.""" + """The 429 tells the caller how long the block has left, rounded up to whole seconds.""" + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=77) - assert [_json_login(client, "/v2/login") for _ in range(2)] == [401, 401] + _set_throttle_clock(monkeypatch, 1_000.6) refused = client.post("/v2/login", json={"username": "admin", "password": "wrong"}) assert refused.status_code == 429 assert refused.headers.get("retry-after") == "77" @@ -595,10 +603,11 @@ def test_a_refused_attempt_carries_retry_after(client, monkeypatch, reset_login_ def test_the_form_returns_a_human_readable_lockout_page(client, monkeypatch, reset_login_throttle): """The no-JavaScript form must render a wait page when its POST is throttled.""" + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=77) - assert [_form_login(client) for _ in range(2)] == [401, 401] + _set_throttle_clock(monkeypatch, 1_000.6) refused = client.post("/login", data={"username": "admin", "password": "wrong"}) assert refused.status_code == 429 assert refused.headers.get("content-type", "").startswith("text/html") @@ -697,12 +706,7 @@ def test_the_master_key_as_a_bearer_token_still_works_while_the_ui_password_is_b def test_a_database_users_correct_password_is_refused_while_blocked(client, monkeypatch, reset_login_throttle): """The block is hard: while it lasts, nothing from that source signs in as that user, right password or not, and the block is not extended by the refused attempts.""" - from types import SimpleNamespace - - from litellm.proxy.auth import login_throttle - - frozen_now = login_throttle.time.time() - monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: frozen_now)) + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=64) _db_user(monkeypatch, "user@corp.com") diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 876ffe96584..114c970e163 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -8,6 +8,7 @@ import math import time from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,17 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( from tests.unit.proxy.db.fake_prisma_engine import engine_call +@dataclass(slots=True) +class _ManualClock: + now: float = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + def _far_deadline() -> float: """A run deadline far enough out that only the other bounds can stop a batch loop.""" return time.monotonic() + 3600 @@ -908,26 +920,31 @@ async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( general_settings={ "maximum_spend_logs_retention_period": "7d", # Comfortably more batches than a sub-second budget can reach (each - # batch sleeps 0.1s), but small enough that a broken deadline fails - # this test in seconds instead of hanging it + # batch advances the clock 0.1s), but small enough that a broken + # deadline still ends the loop "maximum_spend_logs_cleanup_max_batches": 50, - } + }, + clock=clock, ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - started_at = time.monotonic() - result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 1.0) - elapsed = time.monotonic() - started_at + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, clock() + 1.0) assert result.stop_reason == "budget_exhausted" - assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s" + assert clock.now < 1.15, f"budgeted run overran its deadline: {clock.now}s" assert mock_db.execute_raw.call_count < 50 assert result.rows_deleted > 0 @@ -943,7 +960,13 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -951,20 +974,19 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): "maximum_spend_logs_retention_period": "7d", "maximum_autorouter_session_retention_period": "365d", # Comfortably more batches than a sub-second budget can reach (each - # batch sleeps 0.1s), but small enough that a broken deadline fails - # this test in seconds instead of hanging it + # batch advances the clock 0.1s), but small enough that a broken + # deadline still ends the loop "maximum_spend_logs_cleanup_max_batches": 50, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None - started_at = time.monotonic() await cleaner.cleanup_old_spend_logs(mock_prisma_client) - elapsed = time.monotonic() - started_at # three tables are eligible; a per-table budget would push this past 3s - assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s" + assert clock.now < 1.15, f"budget was granted per table, not per run: {clock.now}s" tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list} assert "LiteLLM_SpendLogs" in tables_touched @@ -972,10 +994,15 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): @pytest.mark.asyncio async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete(): mock_prisma_client = MagicMock() - _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -984,7 +1011,8 @@ async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete() "maximum_health_check_retention_period": "30d", "maximum_spend_logs_cleanup_max_batches": 50, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None @@ -1598,7 +1626,13 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): mock_prisma_client = MagicMock() mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -1606,18 +1640,17 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): "maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_max_batches": 500, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None - started_at = time.monotonic() await cleaner.cleanup_old_spend_logs(mock_prisma_client) - elapsed = time.monotonic() - started_at tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list] assert tables.count("LiteLLM_SpendLogs") > 0 assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog" - assert elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s" + assert clock.now < 1.15, f"splitting the budget must not extend the run: {clock.now}s" @pytest.mark.asyncio From 184ec7b89449394d00cb7d7c50a7afee96d4e927 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:46:35 -0700 Subject: [PATCH 16/58] fix(vertex_ai): build the cachedContents collection URL under a custom api_base (#44783) * fix(gemini): use collection endpoint for custom context cache * fix(vertex_ai): build the cachedContents collection URL under a custom api_base * fix(vertex_ai): keep a gateway prefix that contains /publishers/ * test(integration): expect the cachedContents collection URL on a custom Gemini api_base * fix(vertex_ai): cut cache collection URLs at the base's location, graft a mount prefix * fix(vertex_ai): cache under a publisher model api_base in that model's location A Vertex AI api_base that names an endpoint or no model keeps its old cache URL, so no cache is created that the chat call cannot use. The cache body takes its project and location from the api_base model, matching the collection URL. * fix(vertex_ai): cache for the api_base's model and skip caching on a Gemini v1 base A Vertex api_base naming a publisher model now gives the cache its model as well as its project and location, so the cache matches the model generateContent calls. A Gemini api_base ending in /v1 skips context caching, since the v1 API has no cachedContents and rejects the cachedContent field. * fix(vertex_ai): check the cache token minimum against the api_base's model The cache is created for the model a publisher-model api_base names, so the minimum prompt size now comes from that model too. A prompt below its minimum is sent uncached instead of failing at cache create. * fix(vertex_ai): use the stricter cache minimum and send a Gemini base's 404 uncached --------- Co-authored-by: paulo037 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../vertex_ai_context_caching.py | 238 +++--- litellm/llms/vertex_ai/vertex_llm_base.py | 51 +- litellm/types/llms/vertex_ai.py | 2 +- .../cost_calculation/cases/gemini.py | 4 +- ...test_gemini_messages_cache_control_wire.py | 8 +- ...test_cache_control_tool_call_marks_wire.py | 4 +- ...ni_context_caching_custom_api_base_wire.py | 130 +++ .../test_vertex_ai_context_caching.py | 802 ++++++++++-------- 8 files changed, 767 insertions(+), 472 deletions(-) create mode 100644 tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 5bff43828d6..884c82fe331 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,4 +1,5 @@ from collections.abc import Sequence +from dataclasses import dataclass from typing import Final, Literal import httpx @@ -8,7 +9,6 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching.caching import Cache, LiteLLMCacheType -from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -16,10 +16,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, ) from litellm.llms.openai.openai import AllMessageValues -from litellm.utils import is_prompt_caching_valid_prompt +from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url -from ..vertex_llm_base import VertexBase +from ..vertex_llm_base import VertexBase, publisher_model_location from .known_cached_contents import KnownCachedContents, known_cached_contents from .transformation import ( cached_messages_end_on_supported_turn, @@ -65,6 +66,19 @@ class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) +_TOOL_PARAMS: Final = ("tools", "tool_choice") + + +def _sync_cache_client(timeout: float | httpx.Timeout | None) -> HTTPHandler: + if timeout is None: + return HTTPHandler() + return HTTPHandler(timeout=httpx.Timeout(timeout) if isinstance(timeout, (float, int)) else timeout) + + +class _CachedContentsNotServedError(VertexAIError): + pass + + def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedContent) -> None: tokens: Final = (created.get("usageMetadata") or {}).get("totalTokenCount") or 0 logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count @@ -78,6 +92,61 @@ def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None: return timeout +@dataclass(frozen=True, slots=True) +class _CacheTarget: + model: str + project: str | None + location: str | None + + +def _cache_target( + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + api_base: str | None, + model: str, + vertex_project: str | None, + vertex_location: str | None, +) -> _CacheTarget: + model_location: Final = ( + None if custom_llm_provider == "gemini" or api_base is None else publisher_model_location(api_base) + ) + if model_location is None: + return _CacheTarget(model=model, project=vertex_project, location=vertex_location) + return _CacheTarget(model=model_location.model, project=model_location.project, location=model_location.location) + + +def _prompt_cache_min_tokens(models: Sequence[str], custom_llm_provider: str) -> int: + return max( + get_prompt_cache_min_tokens( + model if model.startswith(custom_llm_provider) else f"{custom_llm_provider}/{model}" + ) + for model in models + ) + + +def _context_caching_skip_reason( + cached_messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, + deployment_model: str, + cache_model: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], +) -> str | None: + if not cached_messages_end_on_supported_turn(cached_messages): + return ( + "Vertex AI context caching: cached message block ends on a model turn once " + "system messages are extracted, which the cachedContents API rejects." + ) + min_tokens: Final = _prompt_cache_min_tokens((deployment_model, cache_model), custom_llm_provider) + if not is_prompt_caching_valid_prompt( + model=cache_model, + messages=list(cached_messages), + tools=None if tools is None else list(tools), + custom_llm_provider=custom_llm_provider, + min_token_count=min_tokens, + ): + return f"Vertex AI context caching: cached content is below minimum token count ({min_tokens})." + return None + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -135,6 +204,7 @@ class ContextCachingEndpoints(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_api_version=("v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1"), + collection_endpoint=True, ) def check_cache( @@ -206,6 +276,8 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None + if e.response.status_code == 404 and custom_llm_provider == "gemini": + raise _CachedContentsNotServedError(status_code=404, message=e.response.text) raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -305,6 +377,8 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None + if e.response.status_code == 404 and custom_llm_provider == "gemini": + raise _CachedContentsNotServedError(status_code=404, message=e.response.text) raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -370,41 +444,25 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - if not cached_messages_end_on_supported_turn(cached_messages): - verbose_logger.debug( - "Vertex AI context caching: cached message block ends on a model turn once " - "system messages are extracted, which the cachedContents API rejects. " - "Skipping context caching." - ) + target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location) + tools: Final = optional_params.get("tools") + tool_choice: Final = optional_params.get("tool_choice") + skip_reason: Final = _context_caching_skip_reason( + cached_messages, tools, model, target.model, custom_llm_provider + ) + if skip_reason is not None: + verbose_logger.debug("%s Skipping context caching.", skip_reason) return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. - if not is_prompt_caching_valid_prompt( - model=model, - messages=cached_messages, - tools=optional_params.get("tools"), - custom_llm_provider=custom_llm_provider, - ): - verbose_logger.debug( - "Vertex AI context caching: cached content is below minimum token " - "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, - ) - return messages, optional_params, None - - tools: Final = optional_params.pop("tools", None) - tool_choice: Final = optional_params.pop("tool_choice", None) - ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) headers: Final = { @@ -417,44 +475,42 @@ class ContextCachingEndpoints(VertexBase): if extra_headers is not None: headers.update(extra_headers) - if client is None or not isinstance(client, HTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - timeout = httpx.Timeout(timeout) - _params["timeout"] = timeout - client = HTTPHandler(**_params) - else: - client = client + http_client: Final = client if isinstance(client, HTTPHandler) else _sync_cache_client(timeout) ## CHECK IF CACHED ALREADY generated_cache_key: Final = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model - ) - google_cache_name: Final = self.check_cache( - cache_key=generated_cache_key, - client=client, - headers=headers, - api_key=api_key, - api_base=api_base, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header, - model=model, + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model ) + try: + google_cache_name: Final = self.check_cache( + cache_key=generated_cache_key, + client=http_client, + headers=headers, + api_key=api_key, + api_base=api_base, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=target.project, + vertex_location=target.location, + vertex_auth_header=vertex_auth_header, + model=target.model, + ) + except _CachedContentsNotServedError: + verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url) + return messages, optional_params, None + for tool_param in _TOOL_PARAMS: + optional_params.pop(tool_param, None) if google_cache_name: return non_cached_messages, optional_params, google_cache_name ## TRANSFORM REQUEST cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, + model=target.model, messages=cached_messages, cache_key=generated_cache_key, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, ) cached_content_request_body["tools"] = tools @@ -473,7 +529,7 @@ class ContextCachingEndpoints(VertexBase): ) try: - response: Final = client.post( + response: Final = http_client.post( url=url, headers=headers, json=cached_content_request_body, @@ -489,7 +545,7 @@ class ContextCachingEndpoints(VertexBase): _record_cache_creation_tokens(logging_obj, raw_response_cached) self._known_caches.remember( KnownCachedContents.scope( - url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location ), generated_cache_key, raw_response_cached.get("name"), @@ -536,41 +592,25 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - if not cached_messages_end_on_supported_turn(cached_messages): - verbose_logger.debug( - "Vertex AI context caching: cached message block ends on a model turn once " - "system messages are extracted, which the cachedContents API rejects. " - "Skipping context caching." - ) + target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location) + tools: Final = optional_params.get("tools") + tool_choice: Final = optional_params.get("tool_choice") + skip_reason: Final = _context_caching_skip_reason( + cached_messages, tools, model, target.model, custom_llm_provider + ) + if skip_reason is not None: + verbose_logger.debug("%s Skipping context caching.", skip_reason) return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. - if not is_prompt_caching_valid_prompt( - model=model, - messages=cached_messages, - tools=optional_params.get("tools"), - custom_llm_provider=custom_llm_provider, - ): - verbose_logger.debug( - "Vertex AI context caching: cached content is below minimum token " - "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, - ) - return messages, optional_params, None - - tools: Final = optional_params.pop("tools", None) - tool_choice: Final = optional_params.pop("tool_choice", None) - ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) headers: Final = { @@ -590,10 +630,10 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key: Final = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model ) scope: Final = KnownCachedContents.scope( - url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location ) async def lookup_or_create() -> str | None: @@ -605,10 +645,10 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) if google_cache_name: @@ -616,12 +656,12 @@ class ContextCachingEndpoints(VertexBase): ## TRANSFORM REQUEST cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, + model=target.model, messages=cached_messages, cache_key=generated_cache_key, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, ) cached_content_request_body["tools"] = tools @@ -659,9 +699,15 @@ class ContextCachingEndpoints(VertexBase): ) return raw_response_cached.get("name") - cache_name: Final = await self._known_caches.single_flight( - scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) - ) + try: + cache_name: Final = await self._known_caches.single_flight( + scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) + ) + except _CachedContentsNotServedError: + verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url) + return messages, optional_params, None + for tool_param in _TOOL_PARAMS: + optional_params.pop(tool_param, None) return non_cached_messages, optional_params, cache_name def get_cache(self): diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 0cdd92bb071..7dd98212f76 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -6,9 +6,11 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json +import re import threading from collections import OrderedDict from collections.abc import Mapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse @@ -46,6 +48,31 @@ def graft_default_vertex_path(api_base: str, default_url: str) -> str: _graft_default_vertex_path = graft_default_vertex_path + +_PUBLISHER_MODEL_LOCATION: Final = re.compile(r".*/projects/[^/]+/locations/[^/]+(?=/publishers/[^/]+/models/[^/]+/?$)") + + +@dataclass(frozen=True, slots=True) +class PublisherModelLocation: + parent_url: str + project: str + location: str + model: str + + +def publisher_model_location(api_base: str) -> PublisherModelLocation | None: + match: Final = _PUBLISHER_MODEL_LOCATION.match(api_base) + if match is None: + return None + _, _, project, _, location = match.group().rsplit("/", 4) + return PublisherModelLocation( + parent_url=match.group(), + project=project, + location=location, + model=api_base.rstrip("/").rsplit("/", 1)[-1], + ) + + GOOGLE_IMPORT_ERROR_MESSAGE: Final = ( "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform" ) @@ -637,15 +664,18 @@ class VertexBase: vertex_location: str | None = None, vertex_api_version: Literal["v1", "v1beta1"] | None = None, use_psc_endpoint_format: bool = False, + collection_endpoint: bool = False, ) -> tuple[str | None, str]: """ for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317 Handles custom api_base for: - 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint} + 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}, or {api_base}/{endpoint} + for a collection endpoint 2. Vertex AI with standard proxies - grafts the default vertex URL path onto the api_base when its path is empty or only an API version (/v1, /v1beta1); - otherwise constructs {api_base}:{endpoint} + otherwise constructs {api_base}:{endpoint}, or, for a collection endpoint under an api_base + naming a publisher model, the collection under that model's projects/{project}/locations/{location} 3. Vertex AI with PSC endpoints - constructs full path structure {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} (only when use_psc_endpoint_format=True) @@ -653,22 +683,22 @@ class VertexBase: Args: use_psc_endpoint_format: If True, constructs PSC endpoint URL format. If False (default), uses api_base as-is and appends :{endpoint} + collection_endpoint: True when `endpoint` names a collection resource such as + cachedContents rather than a model action such as generateContent ## Returns - (auth_header, url) - Tuple[Optional[str], str] """ if api_base: if custom_llm_provider == "gemini": - # For Gemini (Google AI Studio), construct the full path like other providers - if model is None: + if model is None and not collection_endpoint: raise ValueError("Model parameter is required for Gemini custom API base URLs") - url = join_api_base(api_base, f"models/{model}:{endpoint}") + url = join_api_base(api_base, endpoint if collection_endpoint else f"models/{model}:{endpoint}") if gemini_api_key is None: raise ValueError( "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." ) - if gemini_api_key is not None: - auth_header = {"x-goog-api-key": gemini_api_key} + auth_header = {"x-goog-api-key": gemini_api_key} else: # For Vertex AI if use_psc_endpoint_format: @@ -692,7 +722,12 @@ class VertexBase: elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path: url = graft_default_vertex_path(api_base=api_base, default_url=url) else: - url = f"{api_base.rstrip('/')}:{endpoint}" + model_location: Final = publisher_model_location(api_base) if collection_endpoint else None + url = ( + f"{api_base.rstrip('/')}:{endpoint}" + if model_location is None + else join_api_base(model_location.parent_url, endpoint) + ) if stream is True: parsed_stream_url: Final = urlparse(url) stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse" diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index f1176f777c7..7cb5432d009 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -376,7 +376,7 @@ class RequestBody(TypedDict, total=False): class CachedContentRequestBody(TypedDict, total=False): contents: Required[list[ContentType]] system_instruction: SystemInstructions - tools: Tools + tools: Tools | None toolConfig: ToolConfig model: Required[str] # Format: models/{model} ttl: str # ending in 's' - Example: "3.5s". diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index 86a0d3b9067..9fe4b1c50d5 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -3348,10 +3348,10 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( response=RoutedResponse( content_type="application/x-routed", routes={ - "GET /models/gemini-3.1-pro:cachedContents": JsonResponse( + "GET /cachedContents": JsonResponse( content_type="application/json", body={"cachedContents": []} ), - "POST /models/gemini-3.1-pro:cachedContents": JsonResponse( + "POST /cachedContents": JsonResponse( content_type="application/json", body={ "name": "cachedContents/$REQUEST_ID", diff --git a/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py b/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py index 71073b3dd5e..cd53b54e601 100644 --- a/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py +++ b/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py @@ -39,11 +39,11 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from def respond(request: Request) -> Reply: assert request.headers["x-goog-api-key"] == _API_KEY, request.headers if request.method == "GET": - assert request.target == f"/models/{_BACKEND}:cachedContents", request.target + assert request.target == "/cachedContents", request.target return Reply(body=b"{}") assert request.method == "POST", request.method body: Final = _JSON_OBJECT.validate_json(request.body) - if request.target == f"/models/{_BACKEND}:cachedContents": + if request.target == "/cachedContents": assert isinstance(body["displayName"], str) and body["displayName"], body assert body == { "contents": [{"role": "user", "parts": [{"text": "."}]}], @@ -80,7 +80,7 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["content"] == [{"type": "text", "text": "The policy applies."}], response.text assert [(request.method, request.target) for request in wire.drain()] == [ - ("GET", f"/models/{_BACKEND}:cachedContents"), - ("POST", f"/models/{_BACKEND}:cachedContents"), + ("GET", "/cachedContents"), + ("POST", "/cachedContents"), ("POST", f"/models/{_BACKEND}:generateContent"), ] diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py index 6389d90597e..a18224a9f36 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py @@ -406,8 +406,8 @@ def _gemini_peer(request: Request) -> Reply: _GEMINI: Final = "gemini/gemini-3.8-flash" _GEMINI_CACHE_WRITE: Final = [ - ("GET", "/models/gemini-3.8-flash:cachedContents"), - ("POST", "/models/gemini-3.8-flash:cachedContents"), + ("GET", "/cachedContents"), + ("POST", "/cachedContents"), ("POST", "/models/gemini-3.8-flash:generateContent"), ] diff --git a/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py b/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py new file mode 100644 index 00000000000..b76f0cd162b --- /dev/null +++ b/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py @@ -0,0 +1,130 @@ +import json +from collections.abc import Callable +from typing import Final + +import pytest + +from integration._support.client import Gateway +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-3.8-flash" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_MODEL_LOCATION: Final = "europe-west4" +_VERTEX_MODEL_RESOURCE: Final = f"projects/{_PROJECT}/locations/{_MODEL_LOCATION}/publishers/google/models/{_BACKEND}" +_VERTEX_COLLECTION: Final = f"/v1/projects/{_PROJECT}/locations/{_MODEL_LOCATION}/cachedContents" +_GEMINI_COLLECTION: Final = "/v1beta/cachedContents" +_GEMINI_MODEL_PATH: Final = f"/v1beta/models/{_BACKEND}" +_GEMINI_KEY: Final = "scripted-gemini-key" +_CACHE_NAME: Final = "cachedContents/scripted-cache" +_CACHED_TOKENS: Final = 5120 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CACHED_TEXT: Final = " ".join( + f"Policy clause {index}: the gateway keeps this clause in the cache." for index in range(700) +) +_MESSAGES: Final = ( + { + "role": "user", + "content": [{"type": "text", "text": _CACHED_TEXT, "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "How many clauses are there?"}, +) + + +def _generate_content_reply() -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "700"}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": _CACHED_TOKENS + 9, + "cachedContentTokenCount": _CACHED_TOKENS, + "candidatesTokenCount": 2, + "totalTokenCount": _CACHED_TOKENS + 11, + }, + "modelVersion": _BACKEND, + } + ).encode() + ) + + +def _google_responder(collection: str, model_path: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == collection and request.method == "GET": + return Reply(body=b'{"cachedContents": []}') + if request.target == collection and request.method == "POST": + return Reply(body=json.dumps({"name": _CACHE_NAME, "model": f"models/{_BACKEND}"}).encode()) + assert request.target == f"{model_path}:generateContent", f"{request.method} {request.target}" + return _generate_content_reply() + + return respond + + +def _chat(gateway: Gateway, model: str) -> dict[str, JsonValue]: + response: Final = gateway.client.post( + "/v1/chat/completions", + json={"model": model, "messages": list(_MESSAGES), "max_tokens": 32}, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=30, + ) + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _assert_cache_round_trip(requests: tuple[Request, ...], collection: str, model_path: str) -> None: + assert tuple((request.method, request.target) for request in requests) == ( + ("GET", collection), + ("POST", collection), + ("POST", f"{model_path}:generateContent"), + ), [(request.method, request.target) for request in requests] + created: Final = _JSON_OBJECT.validate_json(requests[1].body) + assert "contents" in created, created + generated: Final = _JSON_OBJECT.validate_json(requests[2].body) + assert generated["cachedContent"] == _CACHE_NAME, generated + + +def _cached_tokens_of(response: dict[str, JsonValue]) -> JsonValue: + usage: Final = response["usage"] + assert isinstance(usage, dict) + details: Final = usage["prompt_tokens_details"] + assert isinstance(details, dict) + return details["cached_tokens"] + + +def test_gemini_custom_api_base_lists_and_creates_the_cached_contents_collection(gateway: Gateway) -> None: + with ( + wire_server(_google_responder(_GEMINI_COLLECTION, _GEMINI_MODEL_PATH)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=f"{wire.url}/v1beta", api_key=_GEMINI_KEY) + response: Final = _chat(gateway, model) + requests: Final = wire.drain() + _assert_cache_round_trip(requests, _GEMINI_COLLECTION, _GEMINI_MODEL_PATH) + assert all(request.headers.get("x-goog-api-key") == _GEMINI_KEY for request in requests) + assert _cached_tokens_of(response) == _CACHED_TOKENS + + +@pytest.mark.parametrize("gateway_prefix", ["", "/vertex-proxy"], ids=["model-path", "prefixed-model-path"]) +def test_vertex_custom_api_base_creates_the_cache_where_its_model_lives(gateway: Gateway, gateway_prefix: str) -> None: + api_base_path: Final = f"{gateway_prefix}/v1/{_VERTEX_MODEL_RESOURCE}" + collection: Final = f"{gateway_prefix}{_VERTEX_COLLECTION}" + with ( + wire_server(_google_responder(collection, api_base_path)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{api_base_path}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + response: Final = _chat(gateway, model) + requests: Final = wire.drain() + _assert_cache_round_trip(requests, collection, api_base_path) + assert _JSON_OBJECT.validate_json(requests[1].body)["model"] == _VERTEX_MODEL_RESOURCE + assert all(request.headers.get("authorization") == "Bearer scripted-token" for request in requests) + assert _cached_tokens_of(response) == _CACHED_TOKENS diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 19e7e2dee52..78b9e930001 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1,4 +1,4 @@ -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.context_caching.known_cached_contents import KnownCachedContents from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ContextCachingEndpoints, ) @@ -84,18 +85,10 @@ class TestContextCachingEndpoints: """Teardown for each test method""" self._token_check_patcher.stop() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) - def test_check_and_create_cache_with_cached_content( - self, mock_cache_obj, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + def test_check_and_create_cache_with_cached_content(self, mock_cache_obj, mock_separate, custom_llm_provider): """Test check_and_create_cache when cached_content is provided""" # Setup cached_content = "cached_content_123" @@ -130,15 +123,9 @@ class TestContextCachingEndpoints: mock_separate.assert_not_called() mock_cache_obj.get_cache_key.assert_not_called() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) # No cached messages @@ -168,15 +155,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache is None - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_existing_cache_found( self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -224,15 +205,181 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("vertex_location", ["us-central1", None]) + def test_check_and_create_cache_creates_the_cache_where_the_api_base_model_lives(self, vertex_location: str | None): + context_caching = ContextCachingEndpoints(known_caches=KnownCachedContents()) + listed = MagicMock() + listed.json.return_value = {"cachedContents": []} + created = MagicMock() + created.json.return_value = {"name": "cachedContents/base-cache"} + self.mock_client.get.return_value = listed + self.mock_client.post.return_value = created + base_model_location = "https://gateway.example.com/v1/projects/base-project/locations/europe-west4" + + _, _, cache_name = context_caching.check_and_create_cache( + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Summarize it"}, + ], + optional_params={}, + api_key="unused", + api_base=f"{base_model_location}/publishers/google/models/gemini-2.5-pro", + model="gemini-2.5-flash", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location=vertex_location, + vertex_auth_header="vertex-token", + ) + + assert cache_name == "cachedContents/base-cache" + assert self.mock_client.get.call_args.kwargs["url"] == f"{base_model_location}/cachedContents" + assert self.mock_client.post.call_args.kwargs["url"] == f"{base_model_location}/cachedContents" + assert ( + self.mock_client.post.call_args.kwargs["json"]["model"] + == "projects/base-project/locations/europe-west4/publishers/google/models/gemini-2.5-pro" + ) + + def test_check_and_create_cache_keeps_one_cache_per_api_base_model(self): + context_caching = ContextCachingEndpoints(known_caches=KnownCachedContents()) + created_bodies: Final = [] + + def create(**kwargs): + created_bodies.append(kwargs["json"]) + response = MagicMock() + response.json.return_value = {"name": f"cachedContents/cache-{len(created_bodies)}"} + return response + + def list_created(**kwargs): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cachedContents/cache-{index}", "displayName": body["displayName"]} + for index, body in enumerate(created_bodies, start=1) + ] + } + return response + + self.mock_client.get.side_effect = list_created + self.mock_client.post.side_effect = create + base_model_location = "https://gateway.example.com/v1/projects/base-project/locations/europe-west4" + + def cache_for(base_model: str) -> str | None: + _, _, cache_name = context_caching.check_and_create_cache( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}} + ], + }, + {"role": "user", "content": "Summarize it"}, + ], + optional_params={}, + api_key="unused", + api_base=f"{base_model_location}/publishers/google/models/{base_model}", + model="gemini-2.5-flash", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + ) + return cache_name + + pro_cache: Final = cache_for("gemini-2.5-pro") + lite_cache: Final = cache_for("gemini-2.5-flash-lite") + + assert pro_cache != lite_cache + assert [body["model"].rsplit("/", 1)[-1] for body in created_bodies] == [ + "gemini-2.5-pro", + "gemini-2.5-flash-lite", + ] + assert cache_for("gemini-2.5-pro") == pro_cache + + def _gemini_v1_gateway_call_args(self, list_response: MagicMock): + self.mock_client.get.return_value = list_response + self.mock_async_client.get = AsyncMock(return_value=list_response) + created: Final = MagicMock() + created.json.return_value = {"name": "cachedContents/gateway-cache"} + self.mock_client.post.return_value = created + self.mock_async_client.post = AsyncMock(return_value=created) + messages: Final = [ + { + "role": "user", + "content": [{"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Summarize it"}, + ] + call_args: Final = dict( + messages=messages, + optional_params={"tools": self.sample_tools}, + api_key="unused", + api_base="https://gateway.example.com/gemini/v1", + model="gemini-2.5-flash", + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + return call_args + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_check_and_create_cache_sends_uncached_when_the_gemini_api_base_serves_no_cached_contents( + self, is_async: bool + ): + not_served: Final = MagicMock() + not_served.raise_for_status.side_effect = httpx.HTTPStatusError( + "Not Found", request=MagicMock(), response=MagicMock(status_code=404, text="") + ) + + call_args: Final = self._gemini_v1_gateway_call_args(not_served) + + result: Final = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + + assert result == (call_args["messages"], {"tools": self.sample_tools}, None) + self.mock_client.post.assert_not_called() + self.mock_async_client.post.assert_not_called() + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_check_and_create_cache_caches_through_a_gemini_v1_gateway_that_serves_cached_contents( + self, is_async: bool + ): + empty_list: Final = MagicMock() + empty_list.json.return_value = {} + + call_args: Final = self._gemini_v1_gateway_call_args(empty_list) + + _, params, cache_name = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + client: Final = self.mock_async_client if is_async else self.mock_client + + assert params == {} + assert cache_name == "cachedContents/gateway-cache" + assert client.get.call_args.kwargs["url"] == "https://gateway.example.com/gemini/v1/cachedContents" + assert client.post.call_args.kwargs["url"].startswith("https://gateway.example.com/gemini/v1/cachedContents") + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -299,15 +446,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert call_args.kwargs["json"]["tools"] == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") def test_check_and_create_cache_http_error( @@ -332,9 +473,7 @@ class TestContextCachingEndpoints: mock_response = MagicMock() mock_response.status_code = 400 mock_response.text = "Bad Request" - http_error = httpx.HTTPStatusError( - "Error", request=MagicMock(), response=mock_response - ) + http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response) self.mock_client.post.side_effect = http_error optional_params = self.sample_optional_params.copy() @@ -362,15 +501,9 @@ class TestContextCachingEndpoints: assert "Bad Request" in str(exc_info.value.message) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") async def test_async_check_and_create_cache_with_cached_content( self, mock_cache_obj, mock_separate, custom_llm_provider ): @@ -405,15 +538,9 @@ class TestContextCachingEndpoints: assert returned_cache == cached_content @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - async def test_async_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test async_check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) @@ -444,15 +571,9 @@ class TestContextCachingEndpoints: assert returned_cache is None @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") async def test_async_check_and_create_cache_existing_cache_found( self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -501,23 +622,15 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_create_new_cache( self, mock_get_client, @@ -581,20 +694,12 @@ class TestContextCachingEndpoints: assert call_args.kwargs["json"]["tools"] == self.sample_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_timeout_error( self, mock_get_client, @@ -615,9 +720,7 @@ class TestContextCachingEndpoints: mock_get_token_url.return_value = ("token", "https://test-url.com") # Mock timeout error - self.mock_async_client.post = AsyncMock( - side_effect=httpx.TimeoutException("Timeout") - ) + self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) optional_params = self.sample_optional_params.copy() test_project = "test_project" @@ -643,20 +746,14 @@ class TestContextCachingEndpoints: assert exc_info.value.status_code == 408 assert "Timeout error occurred" in str(exc_info.value.message) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -666,9 +763,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): # Execute result = self.context_caching.check_and_create_cache( messages=self.sample_messages, @@ -691,12 +786,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -732,12 +823,8 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params in async version when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -773,20 +860,14 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params in async version when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -796,9 +877,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): # Execute result = await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, @@ -821,12 +900,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """tool_choice is popped from optional_params when cached messages exist.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -838,9 +913,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): self.context_caching.check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -858,12 +931,8 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider): """tool_choice is NOT popped when there are no cached messages (early return).""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -892,12 +961,8 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -909,9 +974,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -930,9 +993,7 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( self, custom_llm_provider ): @@ -963,15 +1024,9 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1033,15 +1088,9 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1101,15 +1150,9 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1161,15 +1204,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert "toolConfig" not in call_args.kwargs["json"] - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1227,15 +1264,9 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == function_pin - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1278,9 +1309,7 @@ class TestContextCachingEndpoints: } self.mock_client.post.return_value = mock_response - tool_choice = ToolConfig( - functionCallingConfig=FunctionCallingConfig(mode="ANY") - ) + tool_choice = ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY")) optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = tool_choice @@ -1301,9 +1330,7 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == tool_choice - assert call_args.kwargs["json"]["toolConfig"] == { - "functionCallingConfig": {"mode": "ANY"} - } + assert call_args.kwargs["json"]["toolConfig"] == {"functionCallingConfig": {"mode": "ANY"}} mock_cache_obj.get_cache_key.assert_called_once_with( messages=cached_messages, tools=self.sample_tools, @@ -1311,12 +1338,8 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( self, @@ -1361,15 +1384,9 @@ class TestContextCachingEndpoints: second_cache_key = check_cache_calls[1].kwargs["cache_key"] assert first_cache_key != second_cache_key - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that context caching is skipped when cached content is below 1024 tokens. Gemini requires a minimum of 1024 tokens for context caching. If the cached @@ -1416,9 +1433,7 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() - def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum( - self, local_model_cost_map - ): + def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum(self, local_model_cost_map): model = "gemini-2.5-pro" self._token_check_patcher.stop() @@ -1454,15 +1469,48 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ("api_base_model", "deployment_model"), + [("gemini-2.5-pro", "gemini-2.5-flash"), ("gemini-not-in-the-cost-map", "gemini-2.5-pro")], ) + @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider + async def test_check_and_create_cache_applies_the_stricter_model_minimum( + self, local_model_cost_map, is_async: bool, api_base_model: str, deployment_model: str ): + self._token_check_patcher.stop() + messages: Final = [ + {"role": "system", "content": " ".join(["word"] * 1500), "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Hello"}, + ] + call_args: Final = dict( + messages=messages, + optional_params={}, + api_key="unused", + api_base=f"https://gateway.example.com/v1/projects/base-project/locations/global/publishers/google/models/{api_base_model}", + model=deployment_model, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + ) + + result: Final = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + + self._token_check_patcher.start() + assert result == (messages, {}, None) + assert not self.mock_client.mock_calls + assert not self.mock_async_client.mock_calls + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @pytest.mark.asyncio + async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that async context caching is skipped when cached content is below 1024 tokens.""" # Stop the default mock so the real token count check runs self._token_check_patcher.stop() @@ -1505,13 +1553,9 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize("is_async", [False, True]) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) @pytest.mark.asyncio - async def test_check_and_create_cache_considers_tools_for_min_tokens( - self, custom_llm_provider, is_async - ): + async def test_check_and_create_cache_considers_tools_for_min_tokens(self, custom_llm_provider, is_async): """Test that context caching accounts for tools when validating minimum token count. Fixes #42804: When messages alone are below the threshold, but tools push the total @@ -1536,11 +1580,15 @@ class TestContextCachingEndpoints: "type": "function", "function": { "name": f"synthetic_tool_{i}", - "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " * 8, + "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " + * 8, "parameters": { "type": "object", "properties": { - f"arg_{j}": {"type": "string", "description": "Argument description for caching verification " * 4} + f"arg_{j}": { + "type": "string", + "description": "Argument description for caching verification " * 4, + } for j in range(10) }, "required": [f"arg_{j}" for j in range(5)], @@ -1569,19 +1617,23 @@ class TestContextCachingEndpoints: self.mock_client.post.return_value = mock_response self.mock_async_client.post = AsyncMock(return_value=mock_response) - with patch.object( - self.context_caching, - "_get_token_and_url_context_caching", - return_value=("fake_token", "https://fake.url/cachedContents"), - ), patch.object( - self.context_caching, - "check_cache", - return_value=None, - ), patch.object( - self.context_caching, - "async_check_cache", - new_callable=AsyncMock, - return_value=None, + with ( + patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("fake_token", "https://fake.url/cachedContents"), + ), + patch.object( + self.context_caching, + "check_cache", + return_value=None, + ), + patch.object( + self.context_caching, + "async_check_cache", + new_callable=AsyncMock, + return_value=None, + ), ): if is_async: result = await self.context_caching.async_check_and_create_cache( @@ -1625,13 +1677,10 @@ class TestContextCachingEndpoints: post_mock.assert_called_once() call_kwargs = post_mock.call_args.kwargs assert call_kwargs["json"]["tools"] == large_tools - assert call_kwargs["json"]["contents"] == [ - {"role": "user", "parts": [{"text": "Short system instruction."}]} - ] + assert call_kwargs["json"]["contents"] == [{"role": "user", "parts": [{"text": "Short system instruction."}]}] self._token_check_patcher.start() - def _model_turn_final_messages(self, final_cached_role): tool_call = { "id": "call_abc123", @@ -1678,9 +1727,7 @@ class TestContextCachingEndpoints: ] @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) - def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """The cachedContents API rejects contents ending on an assistant or tool turn with HTTP 400 "Requests ending with a model turn are not supported", so the request must proceed uncached instead of failing. @@ -1713,9 +1760,7 @@ class TestContextCachingEndpoints: @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """Async variant: an unsupported terminal turn skips caching instead of failing.""" all_messages = self._model_turn_final_messages(final_cached_role) optional_params = self.sample_optional_params.copy() @@ -1773,14 +1818,8 @@ def test_cached_messages_end_on_supported_turn(): ) is True ) - assert ( - cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) - is False - ) - assert ( - cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) - is False - ) + assert cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) is False + assert cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) is False assert cached_messages_end_on_supported_turn([]) is False @@ -1795,13 +1834,9 @@ class TestCheckCachePagination: self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_finds_cache_on_second_page( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_finds_cache_on_second_page(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly handles pagination and finds cache on second page""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1850,13 +1885,9 @@ class TestCheckCachePagination: second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] assert "pageToken=token_page_2" in second_call_url - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1891,13 +1922,9 @@ class TestCheckCachePagination: assert result is None assert self.mock_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_multiple_pages( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_multiple_pages(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly iterates through multiple pages""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1942,9 +1969,7 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 3 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") async def test_async_check_cache_pagination_finds_cache_on_second_page( self, mock_get_token_url, custom_llm_provider @@ -1974,9 +1999,7 @@ class TestCheckCachePagination: } # Setup mock async client to return different responses - self.mock_async_client.get = AsyncMock( - side_effect=[first_page_response, second_page_response] - ) + self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response]) # Execute result = await self.context_caching.async_check_cache( @@ -2000,13 +2023,9 @@ class TestCheckCachePagination: assert "pageToken=token_page_2" in second_call_url @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that async_check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -2041,29 +2060,21 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_client.get.side_effect = [ - create_page_response(i) for i in range(100) - ] + self.mock_client.get.side_effect = [create_page_response(i) for i in range(100)] result = self.context_caching.check_cache( cache_key=cache_key_to_find, @@ -2082,29 +2093,21 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 100 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_async_client.get = AsyncMock( - side_effect=[create_page_response(i) for i in range(100)] - ) + self.mock_async_client.get = AsyncMock(side_effect=[create_page_response(i) for i in range(100)]) result = await self.context_caching.async_check_cache( cache_key=cache_key_to_find, @@ -2123,8 +2126,6 @@ class TestCheckCachePagination: assert self.mock_async_client.get.call_count == 100 - - class TestVertexAIGlobalLocation: """Test global location handling in context caching.""" @@ -2150,9 +2151,7 @@ class TestVertexAIGlobalLocation: # Assert correct URL format for global expected_url = "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/cachedContents" assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" def test_regional_location_url_construction_v1(self): """Test that regional location uses correct URL (with location prefix) for v1 API.""" @@ -2195,32 +2194,117 @@ class TestVertexAIGlobalLocation: ) # Assert correct URL format for global with beta API - expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + expected_url = ( + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + ) assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" - def test_gemini_context_caching_with_custom_api_base_passes_model(self): - """Gemini context caching with custom api_base must pass model to _check_custom_proxy. - - Regression test for https://github.com/BerriAI/litellm/issues/23846 - Previously model was hardcoded to None, causing ValueError when api_base was set. - """ + @pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://my-proxy.example.com", "https://my-proxy.example.com/cachedContents"), + ("https://my-proxy.example.com/v1beta/", "https://my-proxy.example.com/v1beta/cachedContents"), + ( + "https://gateway.example.com/publishers/relay/v1beta", + "https://gateway.example.com/publishers/relay/v1beta/cachedContents", + ), + ], + ) + def test_gemini_context_caching_with_custom_api_base_uses_collection_endpoint( + self, api_base: str, expected_url: str + ): caching = ContextCachingEndpoints() auth_header, url = caching._get_token_and_url_context_caching( gemini_api_key="test-key", custom_llm_provider="gemini", - api_base="https://my-proxy.example.com", + api_base=api_base, vertex_project=None, vertex_location=None, vertex_auth_header=None, model="gemini-1.5-pro", ) - assert "models/gemini-1.5-pro" in url - assert url.startswith("https://my-proxy.example.com/") + assert auth_header == {"x-goog-api-key": "test-key"} + assert url == expected_url + + @pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "expected_url"), + [ + ( + "vertex_ai", + "https://gateway.example.com", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/v1", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/publishers/relay/v1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro/", + "https://gateway.example.com/publishers/relay/v1/projects/my-project/locations/us-central1" + "/cachedContents", + ), + ( + "vertex_ai_beta", + "https://gateway.example.com/v1beta1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro", + "https://gateway.example.com/v1beta1/projects/my-project/locations/us-central1/cachedContents", + ), + ], + ) + def test_vertex_context_caching_with_custom_api_base_uses_collection_endpoint( + self, custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta"], api_base: str, expected_url: str + ): + caching = ContextCachingEndpoints() + + auth_header, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + vertex_project="my-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + model="gemini-2.5-pro", + ) + + assert auth_header == "vertex-token" + assert url == expected_url + + @pytest.mark.parametrize( + "api_base", + [ + "https://gateway.example.com/v1/projects/tuned-project/locations/europe-west4/endpoints/123", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1", + "https://gateway.example.com/vertex-proxy", + ], + ) + def test_vertex_context_caching_without_a_publisher_model_in_api_base_keeps_the_api_base_resource( + self, api_base: str + ): + caching = ContextCachingEndpoints() + + _, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider="vertex_ai", + api_base=api_base, + vertex_project="my-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + model="gemini-2.5-pro", + ) + + assert url == f"{api_base}:cachedContents" def test_gemini_context_caching_without_api_base_ignores_model(self): """Without custom api_base, model param is not needed (default URL is used).""" @@ -2305,9 +2389,7 @@ def _cacheable_messages() -> list: return [ { "role": "system", - "content": [ - {"type": "text", "text": "cacheable prefix " * 1200, "cache_control": {"type": "ephemeral"}} - ], + "content": [{"type": "text", "text": "cacheable prefix " * 1200, "cache_control": {"type": "ephemeral"}}], }, {"role": "user", "content": "question"}, ] @@ -2404,7 +2486,9 @@ def test_check_and_create_cache_clamps_negative_creation_tokens() -> None: with ( patch.object(ContextCachingEndpoints, "check_cache", return_value=None), - patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u")), + patch.object( + ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u") + ), ): _, _, cache_name = ContextCachingEndpoints().check_and_create_cache( messages=_cacheable_messages(), From bee9f679dcd2581b223d99c1198842a1cd13625f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:46:44 -0700 Subject: [PATCH 17/58] fix(github_copilot): treat a malformed api-key.json like a missing one (#44723) * fix(github_copilot): treat a malformed api-key.json like a missing one * fix(github_copilot): read the key and the endpoints of api-key.json as separate strict views * fix(github_copilot): ignore a non-UTF-8 api-key.json like any other malformed file Read the cached key file as bytes so a UTF-16 file, which Windows PowerShell 5.1 writes with >, fails JSON validation instead of raising the decoder's error, and the request gets the refresh or the login hint * fix(github_copilot): decode UTF-16 key files and refresh over a malformed endpoints field * fix(github_copilot): decode a UTF-16 access-token file and refuse an undecodable one like a missing one * fix(github_copilot): use the refreshed key when api-key.json cannot be written * fix(github_copilot): keep a refreshed key that cannot be saved in memory * fix(github_copilot): serve a kept key and its endpoint together after the key file is repaired A key held in memory is now checked before api-key.json, so the token and the endpoint come from the same source, and the map holds one entry per key file with the access token beside it * test(github_copilot): type the authenticator test fixture and the tests this branch adds --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/github_copilot/authenticator.py | 158 ++++++++---- tests/integration/_support/device_login.py | 12 +- .../test_device_code_login_guard_wire.py | 60 +++-- .../test_github_copilot_authenticator.py | 230 +++++++++++++++++- 4 files changed, 377 insertions(+), 83 deletions(-) diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 59e1804f453..ee0e7af7e9b 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -1,11 +1,12 @@ import json import os import time +from dataclasses import dataclass from datetime import datetime -from typing import Any, Final +from typing import Any, Final, TypeVar import httpx -from pydantic import ConfigDict, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -13,7 +14,6 @@ from litellm.litellm_core_utils.asyncify import can_block_current_thread from litellm.llms.custom_httpx.http_handler import get_httpx_client from .common_utils import ( - APIKeyExpiredError, GetAccessTokenError, GetAPIKeyError, GetDeviceCodeError, @@ -27,6 +27,44 @@ DEFAULT_GITHUB_ACCESS_TOKEN_URL: Final = "https://github.com/login/oauth/access_ DEFAULT_GITHUB_API_KEY_URL: Final = "https://api.github.com/copilot_internal/v2/token" +class CopilotEndpoints(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) + + api: str | None = None + + +class CopilotCachedKey(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) + + token: str + expires_at: float + endpoints: CopilotEndpoints | None = None + + +class CopilotCachedEndpoints(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) + + endpoints: CopilotEndpoints | None = None + + +ApiKeyFileView = TypeVar("ApiKeyFileView", CopilotCachedKey, CopilotCachedEndpoints) + + +@dataclass(frozen=True, slots=True) +class _UnsavedApiKey: + access_token: str + key: CopilotCachedKey + + +_UNSAVED_API_KEYS: Final[dict[str, _UnsavedApiKey]] = {} # mutable-ok: refreshed keys api-key.json could not hold + + +def _read_token_file(path: str) -> str: + with open(path, "rb") as f: + raw: Final = f.read() + return raw.decode(json.detect_encoding(raw)) + + @with_config(ConfigDict(extra="allow", strict=True)) class _DeviceCode(TypedDict): device_code: ReadOnly[str] @@ -87,12 +125,13 @@ class Authenticator: GetAccessTokenError: If unable to obtain an access token after retries. """ try: - with open(self.access_token_file, "r") as f: - access_token = f.read().strip() - if access_token: - return access_token + saved_token: Final = _read_token_file(self.access_token_file).strip() + if saved_token: + return saved_token except OSError: verbose_logger.warning("No existing access token found or error reading file") + except ValueError as e: + verbose_logger.warning("Ignoring %s, not a text access token: %s", self.access_token_file, e) if not can_block_current_thread(): raise GetAccessTokenError( @@ -135,64 +174,83 @@ class Authenticator: Raises: GetAPIKeyError: If unable to obtain an API key. """ - try: - with open(self.api_key_file, "r") as f: - api_key_info = json.load(f) - if api_key_info.get("expires_at", 0) > datetime.now().timestamp(): - return api_key_info.get("token") - else: - verbose_logger.warning("API key expired, refreshing") - raise APIKeyExpiredError( - message="API key expired", - status_code=401, - ) - except OSError: - verbose_logger.warning("No API key file found or error opening file") - except (json.JSONDecodeError, KeyError) as e: - verbose_logger.warning("Error reading API key from file: %s", e) - except APIKeyExpiredError: - pass # Already logged in the try block + unsaved: Final = self._unsaved_api_key() + if unsaved is not None: + return unsaved.token + cached: Final = self._read_api_key_file(CopilotCachedKey) + if cached is not None and cached.expires_at > datetime.now().timestamp(): + return cached.token + if cached is not None: + verbose_logger.warning("API key expired, refreshing") try: - api_key_info = self._refresh_api_key() - with open(self.api_key_file, "w") as f: - json.dump(api_key_info, f) - token: Final = api_key_info.get("token") - if token: - return token - else: - raise GetAPIKeyError( - message="API key response missing token", - status_code=401, - ) - except OSError as e: - verbose_logger.error("Error saving API key to file: %s", e) - raise GetAPIKeyError( - message=f"Failed to save API key: {e}", - status_code=500, - ) + api_key_info: Final = self._refresh_api_key() except RefreshAPIKeyError as e: raise GetAPIKeyError( message=f"Failed to refresh API key: {e}", status_code=401, + ) from e + serialized: Final = json.dumps(api_key_info) + try: + with open(self.api_key_file, "w") as f: + f.write(serialized) + except OSError as e: + verbose_logger.warning( + "Keeping the refreshed API key in memory, it could not be saved to %s: %s", self.api_key_file, e ) + self._keep_unsaved_api_key(serialized) + token: Final = api_key_info.get("token") + if not token: + raise GetAPIKeyError( + message="API key response missing token", + status_code=401, + ) + return token def get_api_base(self) -> str | None: """ - Get the API endpoint from the api-key.json file. + Get the API endpoint of the current API key. Returns: Optional[str]: The GitHub Copilot API endpoint, or None if not found. """ - try: - with open(self.api_key_file, "r") as f: - api_key_info: Final = json.load(f) - endpoints: Final = api_key_info.get("endpoints", {}) - api_endpoint: Final = endpoints.get("api") - return api_endpoint - except (OSError, json.JSONDecodeError, KeyError) as e: - verbose_logger.warning("Error reading API endpoint from file: %s", e) + cached: Final = self._unsaved_api_key() or self._read_api_key_file(CopilotCachedEndpoints) + if cached is None or cached.endpoints is None: return None + return cached.endpoints.api + + def _read_api_key_file(self, view: type[ApiKeyFileView]) -> ApiKeyFileView | None: + try: + return view.model_validate_json(_read_token_file(self.api_key_file)) + except OSError: + verbose_logger.warning("No API key file found or error opening file") + return None + except ValueError as e: + verbose_logger.warning("Ignoring %s, not a Copilot API key file: %s", self.api_key_file, e) + return None + + def _stored_access_token(self) -> str | None: + try: + access_token: Final = _read_token_file(self.access_token_file).strip() + except (OSError, ValueError): + return None + return access_token or None + + def _unsaved_api_key(self) -> CopilotCachedKey | None: + unsaved: Final = _UNSAVED_API_KEYS.get(self.api_key_file) + if unsaved is None or unsaved.key.expires_at <= time.time(): + return None + return unsaved.key if unsaved.access_token == self._stored_access_token() else None + + def _keep_unsaved_api_key(self, serialized: str) -> None: + access_token: Final = self._stored_access_token() + if access_token is None: + return + try: + key: Final = CopilotCachedKey.model_validate_json(serialized) + except ValueError: + return + _UNSAVED_API_KEYS[self.api_key_file] = _UnsavedApiKey(access_token=access_token, key=key) def _refresh_api_key(self) -> dict[str, Any]: """ diff --git a/tests/integration/_support/device_login.py b/tests/integration/_support/device_login.py index 9b990fd2167..55519678f40 100644 --- a/tests/integration/_support/device_login.py +++ b/tests/integration/_support/device_login.py @@ -73,6 +73,7 @@ COPILOT_USER_CODE: Final = "COPI-LEAK" COPILOT_ACCESS: Final = "device-login-github-access" COPILOT_REJECTED_ACCESS: Final = "device-login-github-access-rejected" COPILOT_FIRST_LOGIN_ACCESS: Final = "device-login-github-access-first-login" +COPILOT_UNSAVED_KEY_ACCESS: Final = "device-login-github-access-unsaved-key" COPILOT_STORED_KEY: Final = "tid=device-login;stored" COPILOT_MINTED_KEY: Final = "tid=device-login;minted" COPILOT_REJECTED_KEY: Final = "tid=device-login;rejected" @@ -132,7 +133,11 @@ def _oauth_token(request: Request, grant: threading.Event) -> Reply: def _copilot_api_key(request: Request, api_url: str) -> Reply: - if request.headers.get("authorization") not in (f"token {COPILOT_ACCESS}", f"token {COPILOT_FIRST_LOGIN_ACCESS}"): + if request.headers.get("authorization") not in ( + f"token {COPILOT_ACCESS}", + f"token {COPILOT_FIRST_LOGIN_ACCESS}", + f"token {COPILOT_UNSAVED_KEY_ACCESS}", + ): return _json(401, {"message": "Bad credentials"}) return _json( 200, @@ -380,4 +385,7 @@ def write_copilot_access(directory: Path, access_token: str) -> None: def clear_tokens(*directories: Path) -> None: for directory in directories: for path in directory.iterdir(): - path.unlink() + if path.is_dir(): + path.rmdir() + else: + path.unlink() diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 4f01d79d43b..b4044674ed5 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -595,6 +595,27 @@ def test_copilot_api_key_is_minted_once_from_the_stored_access_token(scene: _Sce assert json.loads((scene.rig.copilot_dir / "api-key.json").read_text())["token"] == dl.COPILOT_MINTED_KEY +def test_copilot_api_key_that_cannot_be_saved_is_minted_at_most_once_per_worker(scene: _Scene) -> None: + dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_UNSAVED_KEY_ACCESS) + unwritable_key_file: Final = scene.rig.copilot_dir / "api-key.json" + unwritable_key_file.mkdir() + first: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) + first_served: Final = _serve(first) + assert first_served.status == 200, first_served.text + (minted,) = scene.rig.peers.auth.drain() + assert minted.headers.get("authorization") == f"token {dl.COPILOT_UNSAVED_KEY_ACCESS}" + later: Final = tuple(scene.call(_Cell("copilot", "chat", False, "httpx")) for _ in range(3)) + later_served: Final = tuple(_serve(call) for call in later) + assert [served.status for served in later_served] == [200, 200, 200], [served.text for served in later_served] + assert len(scene.rig.peers.auth.drain()) <= 1 + received: Final = scene.rig.peers.api.drain() + forwarded: Final = tuple(_carrying(received, call.marker) for call in (first, *later)) + assert [tuple(dl.bearer(request) for request in requests) for requests in forwarded] == [ + (dl.COPILOT_MINTED_KEY,) + ] * 4, _describe(received) + assert unwritable_key_file.is_dir() + + def test_copilot_access_token_github_rejects_fails_the_request_without_a_device_login(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) served: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) @@ -655,44 +676,39 @@ def test_chatgpt_token_with_a_text_expiry_falls_back_to_the_expiry_inside_the_to assert scene.rig.peers.auth_connections() == () -_PROVIDER_RESOLUTION_ERROR: Final = "GetLLMProvider Exception" - - -@dataclass(frozen=True, slots=True) -class _UnusableCopilotFiles: - files: Mapping[str, str] - answer: str - - -_UNUSABLE_COPILOT_FILES: Final[Mapping[str, _UnusableCopilotFiles]] = MappingProxyType( +_UNUSABLE_COPILOT_FILES: Final[Mapping[str, Mapping[str, bytes]]] = MappingProxyType( { - "empty-access-token": _UnusableCopilotFiles(MappingProxyType({"access-token": ""}), dl.COPILOT_REFUSAL), - "blank-access-token": _UnusableCopilotFiles(MappingProxyType({"access-token": " \n"}), dl.COPILOT_REFUSAL), - "invalid-api-key-json": _UnusableCopilotFiles( - MappingProxyType({"api-key.json": "{not json"}), dl.COPILOT_REFUSAL + "empty-access-token": MappingProxyType({"access-token": b""}), + "blank-access-token": MappingProxyType({"access-token": b" \n"}), + "access-token-truncated-utf-16": MappingProxyType({"access-token": b"\xff\xfe\x00\xd8"}), + "invalid-api-key-json": MappingProxyType({"api-key.json": b"{not json"}), + "api-key-list": MappingProxyType({"api-key.json": b"[]"}), + "api-key-expiry-text": MappingProxyType( + {"api-key.json": json.dumps({"token": "copilot-text-expiry", "expires_at": "soon"}).encode()} ), - "api-key-list": _UnusableCopilotFiles(MappingProxyType({"api-key.json": "[]"}), _PROVIDER_RESOLUTION_ERROR), - "api-key-expiry-text": _UnusableCopilotFiles( - MappingProxyType({"api-key.json": json.dumps({"token": "copilot-text-expiry", "expires_at": "soon"})}), - _PROVIDER_RESOLUTION_ERROR, + "api-key-expired-utf-16": MappingProxyType( + {"api-key.json": json.dumps({"token": "copilot-utf-16", "expires_at": 1}).encode("utf-16")} ), + "api-key-truncated-utf-16": MappingProxyType({"api-key.json": b"\xff\xfe\x00\xd8"}), } ) @pytest.mark.parametrize("name", tuple(_UNUSABLE_COPILOT_FILES)) def test_unusable_copilot_token_files_are_refused_like_missing_ones_each_time(scene: _Scene, name: str) -> None: - variant: Final = _UNUSABLE_COPILOT_FILES[name] - for file_name, content in variant.files.items(): - (scene.rig.copilot_dir / file_name).write_text(content) + files: Final = _UNUSABLE_COPILOT_FILES[name] + for file_name, content in files.items(): + (scene.rig.copilot_dir / file_name).write_bytes(content) first: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) second: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) for served in (first, second): assert served.status == 400, served.text - assert variant.answer in _error_message(served.text), served.text + assert dl.COPILOT_REFUSAL in _error_message(served.text), served.text assert dl.COPILOT_USER_CODE not in served.text assert scene.rig.peers.auth_connections() == () assert _describe(scene.rig.peers.api.drain()) == () + for file_name, content in files.items(): + assert (scene.rig.copilot_dir / file_name).read_bytes() == content @pytest.mark.parametrize("provider", _PROVIDERS) diff --git a/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py index 21d4901cf81..55fd02bc4dd 100644 --- a/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py +++ b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py @@ -1,11 +1,16 @@ import json +import logging import os import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta +from pathlib import Path +from types import MappingProxyType +from typing import Final from unittest.mock import MagicMock, mock_open, patch import pytest +import respx from litellm.llms.github_copilot.authenticator import Authenticator from litellm.llms.github_copilot.common_utils import ( @@ -16,6 +21,20 @@ from litellm.llms.github_copilot.common_utils import ( RefreshAPIKeyError, ) +API_KEY_URL: Final = "https://copilot-key.test/copilot_internal/v2/token" +REFRESHED_KEY: Final = MappingProxyType({"token": "new-api-key", "expires_at": 4102444800}) +TENANT_API: Final = "https://api.tenant.githubcopilot.com" +REPAIRED_API: Final = "https://api.repaired.githubcopilot.com" +TRUNCATED_UTF_16: Final = b"\xff\xfe\x00\xd8" + + +@pytest.fixture +def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", API_KEY_URL) + (tmp_path / "access-token").write_text("gho_logged_in") + return tmp_path + class TestGitHubCopilotAuthenticator: @pytest.fixture @@ -69,13 +88,9 @@ class TestGitHubCopilotAuthenticator: headers_with_token = authenticator._get_github_headers("test-token") assert headers_with_token["authorization"] == "token test-token" - def test_get_access_token_from_file(self, authenticator): + def test_get_access_token_from_file(self, logged_in_token_dir: Path) -> None: """Test retrieving an access token from a file.""" - mock_token = "mock-access-token" - - with patch("builtins.open", mock_open(read_data=mock_token)): - token = authenticator.get_access_token() - assert token == mock_token + assert Authenticator().get_access_token() == "gho_logged_in" def test_get_access_token_login(self, authenticator): """Test logging in to get an access token.""" @@ -118,6 +133,30 @@ class TestGitHubCopilotAuthenticator: assert authenticator.access_token_file not in str(exc.value) mock_login.assert_not_called() + @pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"]) + def test_saved_access_token_in_another_encoding_is_sent_on_refresh( + self, logged_in_token_dir: Path, encoding: str + ) -> None: + """An access-token file a Windows editor saved as UTF-16 or with a BOM still refreshes the key.""" + (logged_in_token_dir / "access-token").write_bytes("gho_logged_in\n".encode(encoding)) + + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) + assert Authenticator().get_api_key() == "new-api-key" + + assert route.calls.last.request.headers["authorization"] == "token gho_logged_in" + + def test_undecodable_access_token_is_refused_like_a_missing_one(self, logged_in_token_dir: Path) -> None: + """A proxy worker gets the login hint for an access-token file it cannot decode, and the file is kept.""" + access_token_file: Final = logged_in_token_dir / "access-token" + access_token_file.write_bytes(TRUNCATED_UTF_16) + + with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc: + pool.submit(Authenticator().get_access_token).result() + + assert "worker thread" in str(exc.value) + assert access_token_file.read_bytes() == TRUNCATED_UTF_16 + def test_get_access_token_failure(self, authenticator): """Test that an exception is raised after multiple login failures.""" with ( @@ -137,7 +176,7 @@ class TestGitHubCopilotAuthenticator: future_time = (datetime.now() + timedelta(hours=1)).timestamp() mock_api_key_data = json.dumps( {"token": "mock-api-key", "expires_at": future_time} - ) + ).encode() with patch("builtins.open", mock_open(read_data=mock_api_key_data)): api_key = authenticator.get_api_key() @@ -148,7 +187,7 @@ class TestGitHubCopilotAuthenticator: past_time = (datetime.now() - timedelta(hours=1)).timestamp() mock_expired_data = json.dumps( {"token": "expired-api-key", "expires_at": past_time} - ) + ).encode() mock_new_data = { "token": "new-api-key", "expires_at": (datetime.now() + timedelta(hours=1)).timestamp(), @@ -163,6 +202,179 @@ class TestGitHubCopilotAuthenticator: assert api_key == "new-api-key" authenticator._refresh_api_key.assert_called_once() + @pytest.mark.parametrize( + "malformed", + [ + b"[]", + json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(), + TRUNCATED_UTF_16, + json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(), + ], + ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"], + ) + def test_get_api_key_malformed_file_refreshes_like_missing( + self, logged_in_token_dir: Path, malformed: bytes + ) -> None: + """A file that is not a Copilot key object is ignored, refreshed over, never answered with a parse error.""" + api_key_file: Final = logged_in_token_dir / "api-key.json" + api_key_file.write_bytes(malformed) + + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) + assert Authenticator().get_api_key() == "new-api-key" + + assert route.call_count == 1 + assert route.calls.last.request.headers["authorization"] == "token gho_logged_in" + assert json.loads(api_key_file.read_text()) == REFRESHED_KEY + + @pytest.mark.parametrize( + "malformed", + [ + b"[]", + json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(), + TRUNCATED_UTF_16, + json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(), + ], + ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"], + ) + def test_get_api_key_malformed_file_is_kept_when_refresh_fails( + self, logged_in_token_dir: Path, malformed: bytes + ) -> None: + """Ignoring the file never deletes it, so a failed refresh leaves the operator's file for inspection.""" + api_key_file: Final = logged_in_token_dir / "api-key.json" + api_key_file.write_bytes(malformed) + + with respx.mock(assert_all_called=True) as router: + router.get(API_KEY_URL).respond(401) + with pytest.raises(GetAPIKeyError): + Authenticator().get_api_key() + + assert api_key_file.read_bytes() == malformed + + def test_refreshed_key_that_cannot_be_saved_is_reused_with_its_endpoint(self, logged_in_token_dir: Path) -> None: + unwritable_key_file: Final = logged_in_token_dir / "api-key.json" + unwritable_key_file.mkdir() + first: Final = Authenticator() + second: Final = Authenticator() + + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}}) + served: Final = (first.get_api_key(), second.get_api_key(), first.get_api_key()) + + assert served == ("new-api-key",) * 3 + assert (first.get_api_base(), second.get_api_base()) == (TENANT_API, TENANT_API) + assert route.call_count == 1 + assert unwritable_key_file.is_dir() + + def test_refreshed_key_that_cannot_be_saved_and_has_expired_is_refreshed_again( + self, logged_in_token_dir: Path + ) -> None: + (logged_in_token_dir / "api-key.json").mkdir() + auth: Final = Authenticator() + + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(API_KEY_URL).respond(json={"token": "expired-api-key", "expires_at": 1}) + served: Final = (auth.get_api_key(), auth.get_api_key()) + + assert served == ("expired-api-key",) * 2 + assert route.call_count == 2 + + def test_kept_key_and_its_endpoint_stay_paired_after_the_key_file_is_repaired( + self, logged_in_token_dir: Path + ) -> None: + key_file: Final = logged_in_token_dir / "api-key.json" + key_file.mkdir() + auth: Final = Authenticator() + + with respx.mock(assert_all_called=True) as router: + router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}}) + auth.get_api_key() + key_file.rmdir() + key_file.write_text( + json.dumps({"token": "repaired-api-key", "expires_at": 9999999999, "endpoints": {"api": REPAIRED_API}}) + ) + + assert (auth.get_api_key(), auth.get_api_base()) in ( + ("new-api-key", TENANT_API), + ("repaired-api-key", REPAIRED_API), + ) + + def test_unsaved_key_is_not_reused_after_the_access_token_changes(self, logged_in_token_dir: Path) -> None: + (logged_in_token_dir / "api-key.json").mkdir() + + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) + Authenticator().get_api_key() + (logged_in_token_dir / "access-token").write_text("gho_logged_in_again") + Authenticator().get_api_key() + + assert [call.request.headers["authorization"] for call in route.calls] == [ + "token gho_logged_in", + "token gho_logged_in_again", + ] + + @pytest.mark.parametrize( + "malformed", + [ + b"[]", + json.dumps({"endpoints": {"api": 5}}).encode(), + json.dumps({"token": "t", "expires_at": 1, "endpoints": []}).encode(), + TRUNCATED_UTF_16, + ], + ids=["json-list", "number-api", "list-endpoints", "truncated-utf-16"], + ) + def test_get_api_base_malformed_file_returns_none( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, malformed: bytes + ) -> None: + """get_api_base runs before get_api_key on every request, so it must survive malformed shapes too.""" + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_bytes(malformed) + + assert Authenticator().get_api_base() is None + + def test_get_api_base_keeps_endpoints_when_the_key_is_malformed(self, logged_in_token_dir: Path) -> None: + """The endpoint is read on its own, so a bad key never sends the refreshed key to the default host.""" + (logged_in_token_dir / "api-key.json").write_text( + json.dumps({"token": "tid=abc", "expires_at": "soon", "endpoints": {"api": TENANT_API}}) + ) + auth: Final = Authenticator() + + assert auth.get_api_base() == TENANT_API + with respx.mock(assert_all_called=True) as router: + router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) + assert auth.get_api_key() == "new-api-key" + + @pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"]) + def test_live_key_file_in_another_json_encoding_is_used_without_a_refresh( + self, logged_in_token_dir: Path, encoding: str + ) -> None: + """A key file a Windows editor saved as UTF-16 or with a BOM still holds a live key and tenant.""" + live_key: Final = json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": {"api": TENANT_API}}) + api_key_file: Final = logged_in_token_dir / "api-key.json" + api_key_file.write_bytes(live_key.encode(encoding)) + auth: Final = Authenticator() + + with respx.mock: + assert auth.get_api_base() == TENANT_API + assert auth.get_api_key() == "tid=abc" + + assert api_key_file.read_bytes() == live_key.encode(encoding) + + def test_ignored_key_file_is_logged_without_its_contents( + self, logged_in_token_dir: Path, caplog: pytest.LogCaptureFixture + ) -> None: + """The warning names the file but never echoes a token saved in the wrong shape.""" + (logged_in_token_dir / "api-key.json").write_text( + json.dumps({"token": ["tid=secret-value"], "expires_at": 4102444800}) + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"), respx.mock(assert_all_called=True) as router: + router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) + assert Authenticator().get_api_key() == "new-api-key" + + assert "not a Copilot API key file" in caplog.text + assert "tid=secret-value" not in caplog.text + def test_refresh_api_key(self, authenticator, mock_http_client): """Test refreshing an API key.""" mock_client, mock_response = mock_http_client @@ -270,7 +482,7 @@ class TestGitHubCopilotAuthenticator: "expires_at": (datetime.now() + timedelta(hours=1)).timestamp(), "endpoints": {"api": "https://api.enterprise.githubcopilot.com"}, } - ) + ).encode() with patch("builtins.open", mock_open(read_data=mock_api_key_data)): api_base = authenticator.get_api_base() assert api_base == "https://api.enterprise.githubcopilot.com" From c489aa7117c864d2b70dfd23681a831ae9f4e805 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:47:04 -0700 Subject: [PATCH 18/58] feat(proxy)!: add suppress_startup_output to turn off startup prints (#47708) * fix(proxy): keep every startup line JSON when json_logs is on Boot-time messages the proxy printed straight to stdout (banner, feedback box, model list, cache and callback notices, policy table, migration notices, CLI launch notes) and the first records of uvicorn- or granian-spawned workers came out as plain text while json_logs was on, breaking log pipelines that parse every line as JSON. - print_startup_message() prints unchanged in plain mode and, with json_logs on, emits one INFO record on the proxy logger (ANSI stripped) that skips the logger and handler level checks like the print it replaces; show_banner() routes to it in JSON mode only - the CLI exports JSON_LOGS=true when the config sets json_logs so spawned workers log JSON from their first record - the --enforce_prisma_migration_check deprecation notice prints after the config is read so a config-only json_logs covers it Resolves LIT-1624 * fix(proxy): keep show_banner() SDK-free, pick the JSON banner in initialize() Since #46803 the `lite` shell calls show_banner() against an uninitialized litellm package, and reading litellm.json_logs inside it initialized the whole SDK at the banner (0.14 s to 0.83 s for a bare `lite` here, the seconds that PR removed on small VMs). show_banner() is back to its main version; the proxy's initialize() now emits the JSON banner record itself and calls show_banner() only in plain mode. The banner unit tests move to the initialize() call site. * fix(proxy): keep the JSON banner when an operator --log_config names the proxy logger uvicorn applies --log_config after proxy_cli has switched JSON logging on and before initialize() switches it on again through load_config(). A logging config that names "LiteLLM Proxy" removes that logger's handlers in the window between the two, which is where initialize() emits the banner and the feedback box, so at one worker both went nowhere; two-worker children import litellm after uvicorn's logging setup and were not affected. print_startup_message() now falls back to a one-off JSON handler, on the stream the JSON handler uses, when the proxy logger has no handler of its own. Live, json_logs on plus a log_config naming the proxy logger, one worker: main prints the banner and the feedback box as plaintext, e2303a3925 prints neither, this head prints both as one JSON record each. New unit test fails on e2303a3925; new integration cell test_json_logs_with_operator_log_config_keeps_every_line_json passes at one and two workers. * fix(proxy): keep a long startup listing whole in its JSON record print_startup_message ran its record through the proxy logger's filters, so in JSON mode the required-migration SQL list and the policy list were middle-truncated at MAX_STRING_LENGTH_STDOUT_LOG (4096 by default) with the litellm_truncated marker, where the prints they replace showed them whole. That cap bounds echoed request payloads; a startup message is operator text, so the record now gets the uncapped processing of a DEBUG line: secret redaction and base64 collapse still apply, the text cap does not. Red on 41f92b745a, green here: - tests/unit/test_logging.py: a listing longer than the cap comes out whole; a credential in a startup message is still redacted - tests/integration/configuration/test_json_logs_startup.py: a 60-policy config boots with json_logs on and the policy listing record (9762 chars) carries every policy with no truncation marker (4096 chars, marker, on the unfixed tree) * fix(proxy): keep the --enforce_prisma_migration_check notice on --health and --test This PR moved the deprecation notice after the config is read, so that a json_logs set only in the config applies to it; --health and --test return before that point and stopped printing it, a visible change main did not have. The notice is now also emitted before those two returns, through the same helper, so a deployment running `litellm --health --enforce_prisma_migration_check` still learns the flag is dead, as on main. A normal server start emits it once, after the config, as in the previous commit. Red on 2d7a2d6844, green here: tests/unit/proxy/test_proxy_cli.py asserts the notice on --health and on --test. * fix(proxy): keep the JSON startup records on stdout under an operator handler on the proxy logger An operator --log_config that gives "LiteLLM Proxy" a handler of its own, a NullHandler that silences it or a plain WARNING-level handler on stderr, is in force in the window where the banner, the feedback box and the cache notice are emitted: uvicorn applies it after the CLI process switched JSON on and before initialize() re-installs the JSON handler through load_config(). 41f92b745a sent the startup records through the proxy logger's handlers and fell back to a JSON stdout handler only when there were none, so under such a config the records landed in the operator's handler, as plaintext on stderr, or vanished, where the prints they replace reached stdout whatever the logging config said. print_startup_message() now writes every startup record through the proxy's own JSON handler on stdout, filters included, and never through the proxy logger's handlers; the operator's handler gets nothing from a startup message, as it got nothing from the print. The normal JSON handler is the same class, formatter and filters, so the records a pipeline reads are unchanged when no operator config is in play. Tests: tests/unit/test_logging.py gets the two handler shapes (red on f6f5957d16: the record went to stderr as plaintext or nowhere); tests/integration/configuration/ test_json_logs_startup.py gets a silencing operator --log_config cell at 1 and 2 workers (red on f6f5957d16: banner, feedback box and cache notice missing; green here). * fix(proxy): print the --enforce_prisma_migration_check notice when the config fails to load f6f5957d16 printed the deprecation notice after the config is read, so a json_logs set only in the config makes it a JSON record. A boot whose config file is missing or does not parse then exited before that point and never printed it, where main printed it before the error (10-20 gauntlet finding). The notice is printed where main printed it, unless a config file can still switch json_logs on: then it is held until the config is read and printed right after the json_logs switch, and a click close callback prints it on the way out if the boot fails first. functools.cache keeps it to one print either way. Tests: tests/unit/proxy/test_proxy_cli.py covers a failed config read (red on 75043aed83), a plain config printing it once and a config-only json_logs making it one JSON record; tests/integration/configuration/test_json_logs_startup.py adds the JSON cell and a missing and an unparseable config cell (red on 75043aed83, green here), through refused_boot_log, which now takes extra CLI arguments. * feat(proxy): add suppress_startup_output to turn off startup prints * fix(proxy): print startup messages once and without color codes off a terminal * test(policy-engine): check the policy table keeps colors on a terminal and drops them on a pipe * fix(proxy): keep startup warnings when suppress_startup_output is on and apply it from CONFIG_FILE_PATH and before other settings Warnings that name something the operator must fix now print through print_startup_warning, which skips the suppress gate. The config key is read with the CLI's json_logs handling for --config and CONFIG_FILE_PATH, applied before the rest of litellm_settings in load_config, and can only turn the setting on. The Prometheus config panel obeys the setting, and startup prints use the plain log format's color rule. * fix(proxy): read CONFIG_FILE_PATH in the CLI only when workers would load it Workers load CONFIG_FILE_PATH only when it names an existing YAML file, so the CLI's early read of json_logs and suppress_startup_output skips any other value instead of failing a boot that used to start. * fix(proxy): read only suppress_startup_output from CONFIG_FILE_PATH in the CLI, without secrets The CLI now reads CONFIG_FILE_PATH as a plain file with its includes, skips it when a config bucket is set, and applies only suppress_startup_output, also when --config is passed. json_logs there stays with the workers as before. Fatal boot errors drop their color codes off a terminal like the startup prints. * fix(logging): emit startup warnings as WARNING records in JSON mode * fix(proxy): keep one JSON record per search tool in the startup listing * fix(proxy): keep booting with an ASCII stdout or an empty litellm_settings Startup messages now go through click.echo, which writes UTF-8 when stdout is ASCII, as the banner did before. An empty litellm_settings section reads as no settings in the CLI, as the workers already read it * test(proxy): boot the CLI tests on the real ProxyConfig config readers * test(proxy): cover the startup warnings that print with startup output suppressed * test(proxy): boot the deprecated resolver flag test on the real ProxyConfig * test(proxy): drop the startup output contract's narrating docstrings * test(proxy): cover the pgbouncer start failure message off a terminal --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/_logging.py | 94 +++-- litellm/integrations/prometheus.py | 4 +- litellm/proxy/db/check_migration.py | 8 +- litellm/proxy/policy_engine/init_policies.py | 4 +- litellm/proxy/proxy_cli.py | 112 +++-- litellm/proxy/proxy_server.py | 32 +- .../configuration/test_startup_output.py | 133 ++++++ .../integrations/test_prometheus.py | 20 + tests/unit/proxy/db/test_check_migration.py | 4 +- .../proxy/policy_engine/test_init_policies.py | 29 +- .../proxy/proxy_server/test_proxy_config.py | 114 +++++ tests/unit/proxy/test_max_budget_env_var.py | 10 +- tests/unit/proxy/test_proxy_cli.py | 394 ++++++++++++++++-- ...test_proxy_server_endpoints_and_startup.py | 4 +- tests/unit/test_logging.py | 148 +++++-- 16 files changed, 949 insertions(+), 162 deletions(-) create mode 100644 tests/integration/configuration/test_startup_output.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 651997321a3..5e8f4246cef 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -39,6 +39,7 @@ from litellm._logging import ( turn_on_json, log_level, ) +from litellm._logging import suppress_startup_output as suppress_startup_output import re from litellm.constants import ( DEFAULT_BATCH_SIZE, diff --git a/litellm/_logging.py b/litellm/_logging.py index 8974f85f517..5279cb233b4 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -10,10 +10,12 @@ import sys from collections.abc import Iterator, Mapping from datetime import datetime from logging import Formatter +from types import FrameType from typing import Final, TextIO from urllib.parse import unquote from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -598,6 +600,7 @@ _LOG_FORMAT_PREFIX: Final = "%(asctime)s - %(name)s:%(levelname)s" _LOG_FORMAT_SUFFIX: Final = ": %(filename)s:%(lineno)s - %(message)s" _PLAIN_LOG_FORMAT: Final = _LOG_FORMAT_PREFIX + _LOG_FORMAT_SUFFIX _COLOR_LOG_FORMAT: Final = f"\033[92m{_LOG_FORMAT_PREFIX}\033[0m{_LOG_FORMAT_SUFFIX}" +_ANSI_ESCAPE: Final = re.compile(r"\x1b\[[0-9;]*m") def _stream_is_tty(stream: TextIO | None) -> bool: @@ -612,15 +615,28 @@ def _stream_is_tty(stream: TextIO | None) -> bool: return False -def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str: - """The plain-text log format, colorized only when both streams are an interactive terminal. +def _color_enabled(stream: TextIO | None) -> bool: + """True when ANSI colors written to the stream reach an interactive terminal. Honors the NO_COLOR convention from no-color.org: color is disabled when NO_COLOR is present with a non-empty value. """ - if os.environ.get("NO_COLOR"): - return _PLAIN_LOG_FORMAT - return _COLOR_LOG_FORMAT if _stream_is_tty(stdout) and _stream_is_tty(stderr) else _PLAIN_LOG_FORMAT + return not os.environ.get("NO_COLOR") and _stream_is_tty(stream) + + +def _terminal_colors(stdout: TextIO | None, stderr: TextIO | None) -> bool: + """True when both streams are an interactive color terminal, the one rule for log lines and startup prints.""" + return _color_enabled(stdout) and _color_enabled(stderr) + + +def terminal_text(message: str) -> str: + """The message with its ANSI colors kept only on a color terminal, by the rule the plain log format uses.""" + return message if _terminal_colors(sys.stdout, sys.stderr) else _ANSI_ESCAPE.sub("", message) + + +def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str: + """The plain-text log format, colorized only when both streams are an interactive terminal.""" + return _COLOR_LOG_FORMAT if _terminal_colors(stdout, stderr) else _PLAIN_LOG_FORMAT class LevelRoutingStreamHandler(logging.StreamHandler): @@ -642,8 +658,8 @@ class LevelRoutingStreamHandler(logging.StreamHandler): super().emit(record) -def _parse_json_logs_env(value: str | None) -> bool: - """Strict opt-in parse for the JSON_LOGS env var: only "true" (any case) enables JSON logs. +def _parse_env_flag(value: str | None) -> bool: + """Strict opt-in parse for a boolean env var such as JSON_LOGS: only "true" (any case) enables it. Matches the reader in litellm-proxy-extras/_logging.py. The previous bool(os.getenv(...)) treated any non-empty value, including "false" and "0", @@ -656,7 +672,18 @@ def resolve_log_level(log_level: str) -> int: return getattr(logging, log_level.upper()) -json_logs: Final = _parse_json_logs_env(os.getenv("JSON_LOGS")) +json_logs: Final = _parse_env_flag(os.getenv("JSON_LOGS")) +suppress_startup_output: Final = _parse_env_flag(os.getenv("LITELLM_SUPPRESS_STARTUP_OUTPUT")) + + +class StartupOutputSettings(TypedDict, total=False): + json_logs: ReadOnly[object] + suppress_startup_output: ReadOnly[object] + + +STARTUP_OUTPUT_SETTINGS: Final = TypeAdapter(StartupOutputSettings) + + # Create a handler for the logger (you may need to adapt this based on your needs) log_level: Final = os.getenv("LITELLM_LOG", "DEBUG") numeric_level: Final[int] = resolve_log_level(log_level) @@ -743,7 +770,7 @@ class JsonFormatter(Formatter): return dt.isoformat() def format(self, record): - message_str: Final = record.getMessage() + message_str: Final = _ANSI_ESCAPE.sub("", record.getMessage()) json_record: Final[dict[str, object]] = { "message": message_str, "level": record.levelname, @@ -799,12 +826,14 @@ class CorrelationPlainFormatter(logging.Formatter): """Appends trace_id/session_id to plain-text log lines stamped by CorrelationContextFilter. Mirrors JsonFormatter's handling of these two fields so request_correlation_in_logs - behaves the same whether or not json_logs is enabled. + behaves the same whether or not json_logs is enabled. Unless it was built with the color format, + it also drops the ANSI color codes a message carries, which a piped log would store as text. """ def format(self, record: logging.LogRecord) -> str: rendered: Final = super().format(record) - formatted: Final = rendered if _is_redacted(record) else redact_string(rendered) + uncolored: Final = rendered if self._fmt == _COLOR_LOG_FORMAT else _ANSI_ESCAPE.sub("", rendered) + formatted: Final = uncolored if _is_redacted(record) else redact_string(uncolored) trace_id: Final = getattr(record, "trace_id", None) session_id: Final = getattr(record, "session_id", None) if not trace_id and not session_id: @@ -898,29 +927,46 @@ verbose_proxy_logger.addFilter(_diagnostic_filter) verbose_proxy_stdout_logger.addFilter(_diagnostic_filter) verbose_logger.addFilter(_diagnostic_filter) -_ANSI_ESCAPE: Final = re.compile(r"\x1b\[[0-9;]*m") - def print_startup_message(message: str) -> None: """Write a boot-time message meant for the operator's terminal. - Plain mode prints it unchanged, colors and spacing included. With JSON logs on, every line of - stdout is read by a log pipeline, so the message goes out as one INFO record of the proxy logger - instead, on stdout like the print whatever handlers the logging config gave that logger, with the - ANSI color codes and surrounding blank lines removed, located at the caller and emitted at any log - level, whole and unredacted as the print showed it; a blank spacer produces no record. + LITELLM_SUPPRESS_STARTUP_OUTPUT or litellm_settings.suppress_startup_output turns every one of them + off; a warning goes through print_startup_warning instead. Plain mode prints it with its spacing, + keeping the ANSI colors only on a color terminal by the same check as the plain log format, so a + container log never stores escape codes. + With JSON logs on, every line of stdout is read by a log pipeline, so the message goes out as one + INFO record of the proxy logger instead, on stdout like the print whatever handlers the logging + config gave that logger, with the ANSI color codes and surrounding blank lines removed, located at + the caller and emitted at any log level, whole and unredacted as the print showed it; a blank + spacer produces no record. """ + if litellm.suppress_startup_output is True: + return + frame: Final = inspect.currentframe() + _write_startup_message(message, frame.f_back if frame is not None else None, logging.INFO) + + +def print_startup_warning(message: str) -> None: + """Write a boot-time warning as print_startup_message writes a message, even with suppress_startup_output + on, since it names something the operator has to fix. With JSON logs on it is a WARNING record, on + stdout where the print put it, so a collector that filters on the level field keeps it.""" + frame: Final = inspect.currentframe() + _write_startup_message(message, frame.f_back if frame is not None else None, logging.WARNING) + + +def _write_startup_message(message: str, caller: FrameType | None, level: int) -> None: if not litellm.json_logs: - print(message) # noqa: T201 # plain mode: this is the operator-terminal print it replaces + import click + + click.echo(terminal_text(message)) return text: Final = _ANSI_ESCAPE.sub("", message).strip("\n") if not text: return - frame: Final = inspect.currentframe() - caller: Final = frame.f_back if frame is not None else None record: Final = verbose_proxy_logger.makeRecord( verbose_proxy_logger.name, - logging.INFO, + level, caller.f_code.co_filename if caller is not None else __file__, caller.f_lineno if caller is not None else 0, text, @@ -934,8 +980,8 @@ def print_startup_message(message: str) -> None: def _startup_json_handler() -> logging.Handler: - """The proxy's JSON handler on stdout, filters included, for one startup message.""" - handler: Final = LevelRoutingStreamHandler() + """The proxy's JSON handler on stdout, filters included, for one startup message or warning.""" + handler: Final = logging.StreamHandler(sys.stdout) handler.setFormatter(JsonFormatter()) handler.addFilter(_secret_filter) handler.addFilter(_correlation_filter) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index bc8344fa686..d15956c0e95 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1173,8 +1173,8 @@ class PrometheusLogger(CustomLogger): # Build label filters from valid configurations label_filters: Final = self._build_label_filters(parsed_configs) - # Pretty print the processed configuration - self._pretty_print_prometheus_config(label_filters) + if litellm.suppress_startup_output is not True: + self._pretty_print_prometheus_config(label_filters) return label_filters def _parse_exclude_config(self) -> tuple[frozenset[str], frozenset[str]]: diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 83cc4501466..d126f112f90 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -4,7 +4,7 @@ import os import subprocess from typing import Final -from litellm._logging import print_startup_message, verbose_logger +from litellm._logging import print_startup_message, print_startup_warning, verbose_logger def extract_sql_commands(diff_output: str) -> list[str]: @@ -63,7 +63,7 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: run_prisma, ) except ImportError as e: - print_startup_message( + print_startup_warning( f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}", ) return False, [] @@ -96,14 +96,14 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: else: return False, [] except subprocess.TimeoutExpired: - print_startup_message( + print_startup_warning( f"Timed out after {timeout}s generating the migration diff. " f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer.", ) return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" - print_startup_message(error_message) + print_startup_warning(error_message) return False, [] diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index 59d15b0ba0a..2d0e9d25368 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -11,7 +11,7 @@ from collections.abc import Mapping, Sequence from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional -from litellm._logging import print_startup_message, verbose_proxy_logger +from litellm._logging import print_startup_message, print_startup_warning, verbose_proxy_logger from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_validator import PolicyValidator @@ -93,7 +93,7 @@ def _print_policies_on_startup( ) else: print_startup_message("\n".join((header, *policy_lines))) - print_startup_message( + print_startup_warning( f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}\n", ) sys.stdout.flush() diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 909f6de40cd..0d8ab98a0b1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -10,16 +10,17 @@ import sys import urllib.parse as urlparse from collections.abc import Callable, Iterable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Annotated, Any, Final import click import httpx from click.core import ParameterSource from dotenv import load_dotenv -from pydantic import ConfigDict +from pydantic import BeforeValidator, ConfigDict, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm -from litellm._logging import print_startup_message +from litellm._logging import StartupOutputSettings, print_startup_message, print_startup_warning, terminal_text from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY from litellm.proxy.db.pgbouncer import ( PgBouncerError, @@ -205,7 +206,7 @@ def deprecated_v2_flag_passed_on_cli() -> bool: def _print_enforce_prisma_migration_check_notice() -> None: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM Proxy: --enforce_prisma_migration_check is " "deprecated and has no effect, because the proxy always exits " "when database setup fails at startup. You can safely remove it.\033[0m", @@ -227,6 +228,36 @@ def _print_or_hold_enforce_prisma_migration_check_notice( return print_once +def _empty_as_none(value: object) -> object: + return value or None + + +class _ConfigStartupOutputSettings(TypedDict, total=False): + litellm_settings: ReadOnly[Annotated[StartupOutputSettings | None, BeforeValidator(_empty_as_none)]] + + +_CONFIG_STARTUP_OUTPUT_SETTINGS: Final = TypeAdapter(_ConfigStartupOutputSettings) + + +def _apply_config_suppress_startup_output(litellm_settings: StartupOutputSettings | None) -> None: + if litellm_settings is None or litellm_settings.get("suppress_startup_output") is not True: + return + litellm.suppress_startup_output = True + os.environ["LITELLM_SUPPRESS_STARTUP_OUTPUT"] = "true" + + +def _apply_config_logging_settings(litellm_settings: StartupOutputSettings | None) -> None: + """Turn on the config's json_logs and suppress_startup_output before any worker starts. Workers that + uvicorn or granian spawn import litellm afresh and read both from the env at import, so they log JSON + and stay quiet from their first line instead of from when load_config reaches the key.""" + _apply_config_suppress_startup_output(litellm_settings) + if litellm_settings is None or litellm_settings.get("json_logs") is not True: + return + litellm.json_logs = True + os.environ["JSON_LOGS"] = "true" + litellm.turn_on_json() + + class ProxyInitializationHelpers: @staticmethod def echo_litellm_version(): @@ -373,7 +404,7 @@ class ProxyInitializationHelpers: if "timeout_worker_healthcheck" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["timeout_worker_healthcheck"] = timeout_worker_healthcheck else: - print_startup_message( + print_startup_warning( f"\033[1;33mLiteLLM Proxy: --timeout_worker_healthcheck " f"requires uvicorn>=0.37.0, but installed uvicorn=={uvicorn.__version__}. " f"Ignoring the flag.\033[0m", @@ -396,7 +427,7 @@ class ProxyInitializationHelpers: import uvicorn if max_requests_before_restart is None: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " "has no effect without --max_requests_before_restart\033[0m\n", ) @@ -404,7 +435,7 @@ class ProxyInitializationHelpers: if "limit_max_requests_jitter" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["limit_max_requests_jitter"] = jitter else: - print_startup_message( + print_startup_warning( f"\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " f"requires uvicorn>=0.41.0, but installed uvicorn=={uvicorn.__version__}. " f"Ignoring the flag.\033[0m", @@ -571,12 +602,12 @@ class ProxyInitializationHelpers: print_startup_message(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") if max_requests_before_restart is not None: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM: --max_requests_before_restart is not supported by Granian " "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n", ) if ciphers is not None: - print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + print_startup_warning("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") kwargs: Final[dict[str, object]] = { "target": "litellm.proxy.proxy_server:app", @@ -698,7 +729,7 @@ class ProxyInitializationHelpers: gunicorn_options["max_requests"] = max_requests_before_restart if max_requests_before_restart_jitter is not None: if max_requests_before_restart is None: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " "has no effect without --max_requests_before_restart\033[0m\n", ) @@ -739,7 +770,7 @@ class ProxyInitializationHelpers: with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - print_startup_message(f""" + print_startup_warning(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` """) @@ -1325,6 +1356,19 @@ def run_server( os.environ[k] = v litellm_settings = None + env_config_file: Final = os.environ.get("CONFIG_FILE_PATH") + if ( + env_config_file + and os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is None + and ProxyConfig().is_yaml(config_file_path=env_config_file) + ): + import asyncio + + _apply_config_suppress_startup_output( + _CONFIG_STARTUP_OUTPUT_SETTINGS.validate_python( + asyncio.run(ProxyConfig().read_config_file(config_file_path=env_config_file)) + ).get("litellm_settings") + ) if config is not None: """ Allow user to pass in db url via config @@ -1342,17 +1386,9 @@ def run_server( ### LITELLM SETTINGS ### litellm_settings = _config.get("litellm_settings", None) - if ( - litellm_settings is not None - and "json_logs" in litellm_settings - and litellm_settings["json_logs"] is True - ): - import litellm - - litellm.json_logs = True - os.environ["JSON_LOGS"] = "true" - - litellm.turn_on_json() + _apply_config_logging_settings( + _CONFIG_STARTUP_OUTPUT_SETTINGS.validate_python(_config).get("litellm_settings") + ) if held_prisma_notice is not None: held_prisma_notice() ### GENERAL SETTINGS ### @@ -1450,7 +1486,9 @@ def run_server( _bad_scheme = unsupported_db_scheme(_candidate_url) if _bad_scheme is not None: print( - f"\033[1;31mLiteLLM Proxy: {unsupported_db_scheme_message(_db_env, _bad_scheme)}\033[0m", + terminal_text( + f"\033[1;31mLiteLLM Proxy: {unsupported_db_scheme_message(_db_env, _bad_scheme)}\033[0m" + ), file=sys.stderr, flush=True, ) @@ -1553,7 +1591,7 @@ def run_server( env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"), ) if deprecated_v2_flag_passed_on_cli() and use_v2_resolver: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM Proxy: --use_v2_migration_resolver is " "deprecated and has no effect, because the v2 migration " "resolver is now the default. You can safely remove it. To " @@ -1561,7 +1599,7 @@ def run_server( "--use_legacy_migration_resolver.\033[0m", ) if not use_v2_resolver: - print_startup_message( + print_startup_warning( "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration " "resolver. It performs the diff-and-force recovery that can " "cause schema thrashing during rolling deploys where two " @@ -1583,22 +1621,26 @@ def run_server( # issues, and any `prisma db push` against a # partitioned LiteLLM_SpendLogs. print( - f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m", + terminal_text(f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m"), file=sys.stderr, flush=True, ) sys.exit(2) if not setup_ok: print( - "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " - "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" + terminal_text( + "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " + "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" + ) ) sys.exit(1) else: print( - "\033[1;31mLiteLLM Proxy: a database URL is set but the prisma CLI is neither on PATH nor importable " - "as a package, so the database cannot be set up. Install it with `pip install 'litellm[extra_proxy]'` " - "or run a shipped LiteLLM image.\033[0m" + terminal_text( + "\033[1;31mLiteLLM Proxy: a database URL is set but the prisma CLI is neither on PATH nor importable " + "as a package, so the database cannot be set up. Install it with `pip install 'litellm[extra_proxy]'` " + "or run a shipped LiteLLM image.\033[0m" + ) ) sys.exit(1) pgbouncer_settings: Final = PgBouncerSettings() @@ -1609,8 +1651,10 @@ def run_server( ) if isinstance(pooled_database_url, PgBouncerError): print( - f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer " - f"could not start: {pooled_database_url.reason}\033[0m", + terminal_text( + f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer " + f"could not start: {pooled_database_url.reason}\033[0m" + ), file=sys.stderr, flush=True, ) @@ -1675,7 +1719,7 @@ def run_server( run_gunicorn=run_gunicorn is True, ) if single_worker_warning is not None: - print_startup_message(f"\033[1;33m{single_worker_warning}\033[0m\n") + print_startup_warning(f"\033[1;33m{single_worker_warning}\033[0m\n") if run_gunicorn is False and run_hypercorn is False and run_granian is False: if limit_concurrency is not None: uvicorn_args["limit_concurrency"] = limit_concurrency diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5e37ff1ba65..5d034a72934 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -310,6 +310,7 @@ import litellm._redis from litellm import Router from litellm._internal_context import service_target, with_service_target from litellm._logging import ( + STARTUP_OUTPUT_SETTINGS, print_startup_message, redact_string, verbose_proxy_logger, @@ -5805,6 +5806,12 @@ class ProxyConfig: # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config + async def read_config_file(self, config_file_path: str) -> Mapping[str, object]: + """The config file with its includes merged, without the bucket, database and secret steps of get_config.""" + return await self._process_includes( + config=self._load_yaml_file(config_file_path), config_file_path=os.path.abspath(config_file_path) + ) + async def _process_includes(self, config: dict, config_file_path: str) -> dict: """ Process includes by appending their contents to the main config @@ -6429,13 +6436,8 @@ class ProxyConfig: search_tools_parsed: Final[list[SearchToolTypedDict]] = [] - print_startup_message("\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m") - for search_tool in search_tools_raw: - # Display loaded search tool search_tool_name = search_tool.get("search_tool_name", "") - search_provider = search_tool.get("litellm_params", {}).get("search_provider", "") - print_startup_message(f"\033[32m {search_tool_name} ({search_provider})\033[0m") # Handle os.environ/ variables in litellm_params litellm_params = search_tool.get("litellm_params", {}) @@ -6605,13 +6607,15 @@ class ProxyConfig: from litellm.types.integrations.prometheus import validate_caller_identity_settings validate_caller_identity_settings(litellm_settings) + if STARTUP_OUTPUT_SETTINGS.validate_python(litellm_settings).get("suppress_startup_output") is True: + litellm.suppress_startup_output = True # ANSI escape code for blue text blue_color_code: Final = "\033[94m" reset_color_code: Final = "\033[0m" for key, value in litellm_settings.items(): if key == "cache" and value is True: - print_startup_message(f"{blue_color_code}\nSetting Cache on Proxy{reset_color_code}") + print_startup_message(f"\n{blue_color_code}Setting Cache on Proxy{reset_color_code}") from litellm.caching.caching import Cache cache_params = {} @@ -6926,6 +6930,8 @@ class ProxyConfig: setattr(litellm, key, value) elif key == "drop_params": litellm.drop_params = drop_params_flag(value, "litellm_settings.drop_params", verbose_proxy_logger) + elif key == "suppress_startup_output": + continue else: verbose_proxy_logger.debug( "%s setting litellm.%s=%s%s", @@ -7281,6 +7287,13 @@ class ProxyConfig: ## SEARCH TOOLS SETTINGS search_tools: Final[list[SearchToolTypedDict] | None] = self.parse_search_tools(config) + if search_tools: + print_startup_message("\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m") + for tool in search_tools: + print_startup_message( + f"\033[32m {tool.get('search_tool_name', '')} " + f"({tool.get('litellm_params', {}).get('search_provider', '')})\033[0m" + ) ## SANDBOX TOOLS SETTINGS from litellm.sandbox.sandbox_tools import register_sandbox_tools @@ -9683,12 +9696,9 @@ async def initialize( master_key, \ user_custom_auth, \ prisma_client - from litellm.proxy.common_utils.banner import LITELLM_BANNER, show_banner + from litellm.proxy.common_utils.banner import LITELLM_BANNER - if litellm.json_logs: - print_startup_message(LITELLM_BANNER) - else: - show_banner() + print_startup_message(f"\n{LITELLM_BANNER}\n") if os.getenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "").lower() != "true": generate_feedback_box() user_model = model diff --git a/tests/integration/configuration/test_startup_output.py b/tests/integration/configuration/test_startup_output.py new file mode 100644 index 00000000000..fe9a5031778 --- /dev/null +++ b/tests/integration/configuration/test_startup_output.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +import yaml +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.process import owned_proxy_process + +STARTED_WORKER: Final = "Started server process" +RELOAD_LINE: Final = "search tool(s) into router" +RELOADS_PER_WORKER: Final = 3 +RELOAD_WAIT_SECONDS: Final = 60 +SEARCH_TOOLS_HEADER: Final = "Proxy initialized with Search Tools:" +STARTUP_MESSAGES: Final = ( + "██╗", + "Give Feedback / Get Help", + "Proxy initialized with Config, Set models:", + "Setting Cache on Proxy", + SEARCH_TOOLS_HEADER, + "Postgres connection budget", +) +ANSI_ESCAPE: Final = "\x1b[" +WORKER_COUNTS: Final = (1, 2) +CACHE_SETTINGS: Final[Mapping[str, JsonValue]] = {"cache": True, "cache_params": {"type": "redis"}} + + +def _lines_with(log: str, text: str) -> int: + return sum(text in line for line in log.splitlines()) + + +def _write_config( + directory: Path, upstream_url: str, model: str, search_tool: str, litellm_settings: Mapping[str, JsonValue] +) -> Path: + config: Final = directory / f"startup_output_{uuid.uuid4().hex}.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": model, + "litellm_params": { + "model": f"openai/{model}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + } + ], + "search_tools": [ + { + "search_tool_name": search_tool, + "litellm_params": {"search_provider": "tavily", "api_key": "integration-search-key"}, + } + ], + "litellm_settings": dict(litellm_settings), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "proxy_config_reload_interval_seconds": 1, + }, + }, + sort_keys=False, + ) + ) + return config + + +def _boot_reload_and_stop( + gateway: Gateway, + directory: Path, + litellm_settings: Mapping[str, JsonValue], + overrides: Mapping[str, str], + workers: int, +) -> tuple[str, str]: + stamp: Final = uuid.uuid4().hex + model: Final = f"startup-output-{stamp}" + search_tool: Final = f"startup-output-search-{stamp}" + config: Final = _write_config(directory, gateway.upstream_url, model, search_tool, litellm_settings) + with owned_proxy_process( + gateway, + directory, + {"LITELLM_LOG": "INFO", **overrides}, + config=config, + workers=workers, + ) as owned: + reply: Final = owned.gateway.chat(model, text=f"startup output {stamp}") + assert reply["model"] == model, reply + reloads: Final = eventually( + lambda: _lines_with(owned.log.read_text(), RELOAD_LINE), + lambda count: count >= RELOADS_PER_WORKER * workers, + seconds=RELOAD_WAIT_SECONDS, + return_last_on_timeout=True, + ) + assert reloads >= RELOADS_PER_WORKER * workers, reloads + log: Final = owned.log + return log.read_text(), search_tool + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_startup_messages_reach_a_piped_stdout_once_and_without_color_codes( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + log, search_tool = _boot_reload_and_stop(gateway, tmp_path, CACHE_SETTINGS, {}, workers) + + for message in STARTUP_MESSAGES: + assert message in log, (message, log) + assert _lines_with(log, SEARCH_TOOLS_HEADER) == workers, log + assert _lines_with(log, f"{search_tool} (tavily)") == workers, log + colored: Final = tuple(line for line in log.splitlines() if ANSI_ESCAPE in line) + assert colored == (), "lines with ANSI color codes on a piped stdout:\n" + "\n".join(colored) + + +@pytest.mark.parametrize( + ("litellm_settings", "overrides"), + [ + pytest.param(CACHE_SETTINGS, {"LITELLM_SUPPRESS_STARTUP_OUTPUT": "true"}, id="env-plain"), + pytest.param({**CACHE_SETTINGS, "json_logs": True, "suppress_startup_output": True}, {}, id="config-json"), + ], +) +def test_suppress_startup_output_leaves_out_every_startup_message( + gateway: Gateway, tmp_path: Path, litellm_settings: Mapping[str, JsonValue], overrides: Mapping[str, str] +) -> None: + workers: Final = 2 + log, search_tool = _boot_reload_and_stop(gateway, tmp_path, litellm_settings, overrides, workers) + + assert _lines_with(log, STARTED_WORKER) == workers, log + printed: Final = tuple(message for message in (*STARTUP_MESSAGES, search_tool) if message in log) + assert printed == (), (printed, log) diff --git a/tests/unit/enterprise/integrations/test_prometheus.py b/tests/unit/enterprise/integrations/test_prometheus.py index 987cfa03a5e..fc4a729ec98 100644 --- a/tests/unit/enterprise/integrations/test_prometheus.py +++ b/tests/unit/enterprise/integrations/test_prometheus.py @@ -217,6 +217,26 @@ def test_future_metrics_with_end_user_are_filtered(): litellm.enable_end_user_cost_tracking_prometheus_only = original_setting +@pytest.mark.parametrize("suppressed", [False, True]) +def test_prometheus_config_panel_obeys_suppress_startup_output( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str], suppressed: bool +): + """The processed-config panel is a startup print like the banner, so suppress_startup_output turns it off + while the label filters still load.""" + clear_prometheus_registry() + monkeypatch.setattr(litellm, "suppress_startup_output", suppressed) + monkeypatch.setattr( + litellm, + "prometheus_metrics_config", + [{"group": "service_metrics", "metrics": ["litellm_deployment_total_requests"], "include_labels": ["team"]}], + ) + + label_filters = PrometheusLogger()._parse_prometheus_config() + + assert label_filters["litellm_deployment_total_requests"] == ["team"] + assert ("Prometheus Configuration Processed" in capfd.readouterr().out) is not suppressed + + def test_prometheus_config_parsing(): """Test that prometheus metrics configuration is parsed correctly""" # Clear registry before test diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index b5685676941..3cc3d8c6713 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/tests/unit/proxy/db/test_check_migration.py @@ -81,7 +81,7 @@ def test_a_skipped_migration_diff_is_one_json_record_in_json_mode( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: """ - With json_logs on, the notice that the migration diff was skipped is one INFO record on stdout, + With json_logs on, the notice that the migration diff was skipped is one WARNING record on stdout, also while the loggers sit at WARNING, as they do at boot without LITELLM_LOG. """ from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper @@ -101,7 +101,7 @@ def test_a_skipped_migration_diff_is_one_json_record_in_json_mode( assert len(lines) == 1, lines record: Final = json.loads(lines[0]) assert (record["level"], record["message"]) == ( - "INFO", + "WARNING", "Skipping the migration diff: litellm-proxy-extras has no Prisma runner. " "Error: import of litellm_proxy_extras.prisma_toolchain halted; None in sys.modules", ) diff --git a/tests/unit/proxy/policy_engine/test_init_policies.py b/tests/unit/proxy/policy_engine/test_init_policies.py index b9b5f278252..ce0c48096d0 100644 --- a/tests/unit/proxy/policy_engine/test_init_policies.py +++ b/tests/unit/proxy/policy_engine/test_init_policies.py @@ -1,8 +1,10 @@ """The policy table printed at boot keeps its terminal form in plain mode and is JSON when json_logs is on.""" import json +import re import sys from collections.abc import Mapping, Sequence +from io import StringIO from typing import Final import pytest @@ -40,6 +42,15 @@ NO_ATTACHMENTS_WARNING: Final = ( ) +class _Stream(StringIO): + def __init__(self, terminal: bool) -> None: + super().__init__() + self._terminal = terminal + + def isatty(self) -> bool: + return self._terminal + + @pytest.mark.parametrize( ("attachments", "expected"), [ @@ -53,18 +64,24 @@ NO_ATTACHMENTS_WARNING: Final = ( ], ids=["with-attachments", "without-attachments"], ) +@pytest.mark.parametrize("terminal", [True, False], ids=["terminal", "piped"]) def test_policy_table_prints_unchanged_in_plain_mode( monkeypatch: pytest.MonkeyPatch, - capfd: pytest.CaptureFixture[str], attachments: Sequence[Mapping[str, object]] | None, expected: str, + terminal: bool, ) -> None: - """Operators without JSON logs keep the colored policy table byte for byte.""" + """Operators without JSON logs keep the colored policy table byte for byte on a terminal, and the same text + without color codes on a piped stdout.""" monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.delenv("NO_COLOR", raising=False) + stdout: Final = _Stream(terminal) + monkeypatch.setattr(sys, "stdout", stdout) + monkeypatch.setattr(sys, "stderr", _Stream(terminal)) _print_policies_on_startup(POLICIES, attachments) - assert capfd.readouterr() == (expected, "") + assert stdout.getvalue() == (expected if terminal else re.sub(r"\x1b\[\d+m", "", expected)) def test_policy_table_is_one_json_record_in_json_mode( @@ -98,8 +115,8 @@ def test_policy_table_is_one_json_record_in_json_mode( def test_missing_attachments_warning_is_a_second_json_record( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: - """Without attachments the table and the yellow warning are two INFO records on stdout, where the two - prints put them.""" + """Without attachments the table and the yellow warning are an INFO and a WARNING record on stdout, + where the two prints put them.""" monkeypatch.setattr(litellm, "json_logs", True) monkeypatch.setattr(sys, "excepthook", sys.excepthook) turn_on_json() @@ -113,4 +130,4 @@ def test_missing_attachments_warning_is_a_second_json_record( "\n".join(("LiteLLM Policy Engine: Loaded 1 policies", "", *POLICY_ROWS)), NO_ATTACHMENTS_WARNING, ] - assert {record["level"] for record in records} == {"INFO"} + assert [record["level"] for record in records] == ["INFO", "WARNING"] diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index e59dd19271c..20a4ca87d42 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -13,6 +13,7 @@ import json import logging import os import re +import sys from collections.abc import Awaitable, Callable, Mapping from contextlib import nullcontext from dataclasses import dataclass @@ -28,6 +29,7 @@ from prisma.errors import DataError from pydantic import JsonValue, TypeAdapter, ValidationError import litellm +from litellm._logging import turn_on_json from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.types.router import Deployment, LiteLLM_Params @@ -2088,6 +2090,118 @@ def test_ProxyConfig_parse_search_tools_missing_returns_none(): assert pc.parse_search_tools({}) is None +@pytest.mark.asyncio +async def test_ProxyConfig_lists_search_tools_once_at_startup_not_on_every_reparse( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +): + """The search tools reload job and add_deployment re-parse search_tools every config reload interval, so + the startup listing belongs to load_config alone, or a container log repeats it every 30 seconds.""" + config_file: Final = tmp_path / "search.yaml" + config_file.write_text( + "model_list: []\n" + "search_tools:\n" + " - search_tool_name: web-search\n" + " litellm_params:\n" + " search_provider: tavily\n" + ) + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(litellm, "suppress_startup_output", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + proxy_config: Final = ProxyConfig() + + await proxy_config.load_config(router=None, config_file_path=str(config_file)) + for _ in range(3): + assert proxy_config.parse_search_tools(proxy_config.get_config_state()) is not None + + out: Final = capfd.readouterr().out + assert out.count("Proxy initialized with Search Tools:") == 1, out + assert out.count("web-search (tavily)") == 1, out + + +@pytest.mark.asyncio +async def test_ProxyConfig_lists_each_search_tool_as_its_own_json_record( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +): + """With json_logs on, the startup listing is a header record and one record per tool, so a pipeline that + reads one tool per record keeps working.""" + config_file: Final = tmp_path / "search.yaml" + config_file.write_text( + "model_list: []\n" + "search_tools:\n" + " - search_tool_name: web-search\n" + " litellm_params:\n" + " search_provider: tavily\n" + " - search_tool_name: news-search\n" + " litellm_params:\n" + " search_provider: exa_ai\n" + ) + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(litellm, "suppress_startup_output", False) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + turn_on_json() + + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + messages: Final = [json.loads(raw)["message"] for raw in capfd.readouterr().out.splitlines() if raw.strip()] + listing: Final = [message for message in messages if "Search Tools" in message or "-search (" in message] + assert listing == [ + "LiteLLM: Proxy initialized with Search Tools:", + " web-search (tavily)", + " news-search (exa_ai)", + ], messages + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("already_on", "config_value"), + [ + pytest.param(True, False, id="env-var-on-and-config-false"), + pytest.param(False, True, id="config-key-below-cache"), + ], +) +async def test_ProxyConfig_load_config_keeps_startup_output_off_once_either_setting_turns_it_on( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], + already_on: bool, + config_value: bool, +): + """Only `true` turns suppress_startup_output on, from the env var or the config key, and a config `false` + never turns the env var back off. The key applies before any other litellm_settings key, so a `cache: true` + above it prints nothing either.""" + config_file: Final = tmp_path / "suppress.yaml" + config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " cache: true\n" + " cache_params:\n" + " type: local\n" + f" suppress_startup_output: {str(config_value).lower()}\n" + "search_tools:\n" + " - search_tool_name: web-search\n" + " litellm_params:\n" + " search_provider: tavily\n" + ) + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(litellm, "suppress_startup_output", already_on) + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + out: Final = capfd.readouterr().out + assert litellm.suppress_startup_output is True + assert "Setting Cache on Proxy" not in out, out + assert "Proxy initialized with Search Tools:" not in out, out + + + def test_ProxyConfig_merge_config_and_db_search_tools_returns_superset(): config_tools = [ { diff --git a/tests/unit/proxy/test_max_budget_env_var.py b/tests/unit/proxy/test_max_budget_env_var.py index ec71f70fa8f..95074ed1edd 100644 --- a/tests/unit/proxy/test_max_budget_env_var.py +++ b/tests/unit/proxy/test_max_budget_env_var.py @@ -18,10 +18,7 @@ async def test_max_budget_string_converted_to_float(): string. initialize() should convert it to float so the comparison `litellm.max_budget > 0` doesn't raise TypeError. """ - with ( - patch("litellm.proxy.common_utils.banner.show_banner"), - patch("litellm.proxy.proxy_server.generate_feedback_box"), - ): + with patch("litellm.proxy.proxy_server.generate_feedback_box"): from litellm.proxy.proxy_server import initialize original = litellm.max_budget @@ -36,10 +33,7 @@ async def test_max_budget_string_converted_to_float(): @pytest.mark.asyncio async def test_max_budget_float_stays_float(): """max_budget as float should still work.""" - with ( - patch("litellm.proxy.common_utils.banner.show_banner"), - patch("litellm.proxy.proxy_server.generate_feedback_box"), - ): + with patch("litellm.proxy.proxy_server.generate_feedback_box"): from litellm.proxy.proxy_server import initialize original = litellm.max_budget diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index d77731aa83b..afefae0dbb1 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -215,6 +215,23 @@ class TestProxyInitializationHelpers: ) assert args["timeout_worker_healthcheck"] == 30 + def test_an_old_uvicorn_drops_worker_healthcheck_and_warns_with_startup_output_suppressed( + self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + class _UvicornConfigWithoutWorkerHealthcheck: + def __init__(self, app=None): + pass + + monkeypatch.setattr(litellm, "suppress_startup_output", True) + monkeypatch.setattr(litellm, "json_logs", False) + with patch("uvicorn.Config", _UvicornConfigWithoutWorkerHealthcheck): + args: Final = ProxyInitializationHelpers._get_default_unvicorn_init_args( + "localhost", 8000, timeout_worker_healthcheck=15 + ) + + assert "timeout_worker_healthcheck" not in args + assert "--timeout_worker_healthcheck requires uvicorn>=0.37.0" in capsys.readouterr().out + def test_get_reload_options_no_config_still_watches_env(self): opts = ProxyInitializationHelpers._get_reload_options(None) assert opts["reload"] is True @@ -1835,10 +1852,9 @@ class TestProxyInitializationHelpers: assert "uvicorn exits" not in warning @patch("uvicorn.run") - @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") def test_single_worker_max_requests_warning_printed_on_uvicorn( - self, mock_setup_db, mock_print, mock_uvicorn_run + self, mock_setup_db, mock_uvicorn_run ): """uvicorn single-worker startup prints the does-not-recycle warning""" from click.testing import CliRunner @@ -1881,11 +1897,8 @@ class TestProxyInitializationHelpers: assert ( result.exit_code == 0 ), f"exit_code={result.exit_code}, output={result.output}" - assert any( - "--max_requests_before_restart 123 with --num_workers 1" in str(c) - and "uvicorn exits the process" in str(c) - for c in mock_print.call_args_list - ) + assert "--max_requests_before_restart 123 with --num_workers 1" in result.output + assert "uvicorn exits the process" in result.output @patch("litellm.proxy.proxy_cli.start_query_engine_reaper") @patch("uvicorn.run") @@ -1942,10 +1955,9 @@ class TestProxyInitializationHelpers: @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server") @patch("uvicorn.run") - @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") def test_single_worker_max_requests_warning_printed_on_gunicorn( - self, mock_setup_db, mock_print, mock_uvicorn_run, mock_run_gunicorn + self, mock_setup_db, mock_uvicorn_run, mock_run_gunicorn ): """--run_gunicorn single-worker startup prints the gunicorn warning variant""" from click.testing import CliRunner @@ -1989,11 +2001,8 @@ class TestProxyInitializationHelpers: assert ( result.exit_code == 0 ), f"exit_code={result.exit_code}, output={result.output}" - assert any( - "--max_requests_before_restart 123 with --num_workers 1" in str(c) - and "gunicorn accepts no connections" in str(c) - for c in mock_print.call_args_list - ) + assert "--max_requests_before_restart 123 with --num_workers 1" in result.output + assert "gunicorn accepts no connections" in result.output @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.init_hypercorn_server") @patch("uvicorn.run") @@ -2149,7 +2158,7 @@ class TestProxyInitializationHelpers: assert [call.args for call in reserved_before_run] == [("the gunicorn master",)] @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") - def test_gunicorn_jitter_without_base_warns(self): + def test_gunicorn_jitter_without_base_warns(self, capsys: pytest.CaptureFixture[str]): """gunicorn path warns when jitter is set without --max_requests_before_restart""" pytest.importorskip("gunicorn") @@ -2160,7 +2169,6 @@ class TestProxyInitializationHelpers: with ( patch("gunicorn.app.base.BaseApplication.run", capture_run), - patch("builtins.print") as mock_print, patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership "litellm.proxy.proxy_cli.start_query_engine_reaper" ), @@ -2178,7 +2186,7 @@ class TestProxyInitializationHelpers: assert "max_requests" not in captured["options"] assert "max_requests_jitter" not in captured["options"] - assert any("has no effect" in str(c) for c in mock_print.call_args_list) + assert "has no effect" in capsys.readouterr().out def test_apply_uvicorn_jitter_sets_arg_when_supported(self): class _NewUvicornConfig: @@ -2194,16 +2202,13 @@ class TestProxyInitializationHelpers: ) assert uvicorn_args["limit_max_requests_jitter"] == 50 - def test_apply_uvicorn_jitter_skipped_on_old_uvicorn(self): + def test_apply_uvicorn_jitter_skipped_on_old_uvicorn(self, capsys: pytest.CaptureFixture[str]): class _FakeUvicornConfig: def __init__(self, limit_max_requests=None): pass uvicorn_args: dict = {} - with ( - patch("uvicorn.Config", _FakeUvicornConfig), - patch("builtins.print") as mock_print, - ): + with patch("uvicorn.Config", _FakeUvicornConfig): ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter( uvicorn_args=uvicorn_args, max_requests_before_restart=1000, @@ -2211,19 +2216,18 @@ class TestProxyInitializationHelpers: ) assert "limit_max_requests_jitter" not in uvicorn_args - assert any("0.41.0" in str(c) for c in mock_print.call_args_list) + assert "0.41.0" in capsys.readouterr().out - def test_apply_uvicorn_jitter_without_base_warns(self): + def test_apply_uvicorn_jitter_without_base_warns(self, capsys: pytest.CaptureFixture[str]): uvicorn_args: dict = {} - with patch("builtins.print") as mock_print: - ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter( - uvicorn_args=uvicorn_args, - max_requests_before_restart=None, - jitter=50, - ) + ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter( + uvicorn_args=uvicorn_args, + max_requests_before_restart=None, + jitter=50, + ) assert "limit_max_requests_jitter" not in uvicorn_args - assert any("has no effect" in str(c) for c in mock_print.call_args_list) + assert "has no effect" in capsys.readouterr().out @patch.dict(os.environ, {}, clear=True) def test_construct_database_url_from_env_vars(self): @@ -2648,10 +2652,56 @@ class TestRunServerDbSetup: assert exit_info.value.code == 1 assert "a database URL is set but the prisma CLI is neither on PATH nor importable" in out assert "pip install 'litellm[extra_proxy]'" in out + assert "\x1b[" not in out else: assert "prisma CLI" not in out assert "Setup complete" in out + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True) + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=True) + def test_a_pgbouncer_that_cannot_start_stops_the_cli_with_an_uncolored_error_off_a_terminal( + self, + mock_build_indexes, + mock_setup_database, + mock_atexit_register, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig + + mock_proxy_module: Final = MagicMock( + app=MagicMock(), + ProxyConfig=ProxyConfig, + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + clean_env: Final = { + k: v + for k, v in os.environ.items() + if not k.startswith(("DATABASE_URL", "DIRECT_URL", "LITELLM_PGBOUNCER_", "DISABLE_SCHEMA_UPDATE")) + } + pgbouncer_env: Final = { + "DATABASE_URL": "postgresql://test:test@localhost:5432/test", + "LITELLM_PGBOUNCER_ENABLED": "true", + "LITELLM_PGBOUNCER_BINARY": str(tmp_path / "missing-pgbouncer"), + } + + with ( + patch.dict(os.environ, {**clean_env, **pgbouncer_env}, clear=True), + patch.dict( + "sys.modules", + {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, + ), + pytest.raises(SystemExit) as exit_info, + ): + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) + + error: Final = capsys.readouterr().err + assert exit_info.value.code == 1 + assert "LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer could not start" in error + assert "\x1b[" not in error + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") @@ -2726,6 +2776,7 @@ class TestRunServerDbSetup: mock_setup_database, mock_atexit_register, mock_subprocess_run, + capsys, ): """A sqlite DATABASE_URL must exit immediately, before any prisma call, instead of stalling on a migration against the postgresql-only schema.""" @@ -2764,6 +2815,9 @@ class TestRunServerDbSetup: ) assert exc_info.value.code == 1 mock_setup_database.assert_not_called() + err: Final = capsys.readouterr().err + assert "sqlite" in err + assert "\x1b[" not in err @patch("subprocess.run") @patch("atexit.register") @@ -2825,6 +2879,51 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out + @patch("subprocess.run", return_value=MagicMock(returncode=0)) + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True) + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") + def test_the_deprecated_v2_resolver_flag_warns_with_startup_output_suppressed( + self, + mock_build_indexes, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig + + mock_proxy_module: Final = MagicMock( + app=MagicMock(), + ProxyConfig=ProxyConfig, + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + clean_env: Final = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER", "DISABLE_SCHEMA_UPDATE") + } + monkeypatch.setattr(litellm, "suppress_startup_output", True) + monkeypatch.setattr(litellm, "json_logs", False) + + with ( + patch.dict( + os.environ, {**clean_env, "DATABASE_URL": "postgresql://test:test@localhost:5432/test"}, clear=True + ), + patch.dict( + "sys.modules", + {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, + ), + ): + run_server.main( + ["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False + ) + + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) + assert "--use_v2_migration_resolver is deprecated and has no effect" in capsys.readouterr().out + @pytest.mark.parametrize( ("arguments", "environment", "warned"), ( @@ -3968,11 +4067,104 @@ def _boot_from_config_file(config_path: Path, *flags: str) -> str | None: return os.environ.get("JSON_LOGS") +COUNTING_SECRET_MANAGER_MODULE: Final = """ +import os + +from litellm.integrations.custom_secret_manager import CustomSecretManager + + +class CountingSecretManager(CustomSecretManager): + def __init__(self): + super().__init__() + self._log("constructed") + + @staticmethod + def _log(event): + with open(os.environ["SECRET_READ_LOG"], "a") as f: + f.write(event + "\\n") + + def sync_read_secret(self, secret_name, optional_params=None, timeout=None, **kwargs): + self._log(secret_name) + return "from-secret-manager" + + async def async_read_secret(self, secret_name, optional_params=None, timeout=None, **kwargs): + self._log(secret_name) + return "from-secret-manager" +""" + +SECRET_MANAGER_BACKED_CONFIG: Final = """ +model_list: + - model_name: my-model + litellm_params: + model: openai/gpt-5.4-mini + api_key: os.environ/MY_PROVIDER_KEY + +litellm_settings: + suppress_startup_output: true + +general_settings: + key_management_system: custom + key_management_settings: + custom_secret_manager: counting_secret_manager.CountingSecretManager + hosted_keys: + - MY_PROVIDER_KEY +""" + + class TestJsonLogsConfigReachesSpawnedWorkers: """`litellm_settings.json_logs: true` turns JSON on inside the CLI process, but the workers that uvicorn or granian spawn are fresh interpreters that pick plain or JSON from the JSON_LOGS environment variable at import. The CLI has to export it, or every worker logs plain text until - its own load_config reaches the key.""" + its own load_config reaches the key. `suppress_startup_output` has the same problem with the banner + each worker prints before its load_config.""" + + @staticmethod + def _write_config(path: Path, litellm_settings: dict[str, object]) -> Path: + path.write_text(yaml.dump({"model_list": [], "litellm_settings": litellm_settings})) + return path + + @staticmethod + def _env_after_boot( + name: str, + config: Path | None = None, + config_file_path_env: Path | None = None, + extra_env: dict[str, str] | None = None, + proxy_config_class: object = None, + ) -> str | None: + from litellm.proxy.proxy_server import ProxyConfig + + mock_proxy_module: Final = MagicMock( + app=MagicMock(), + ProxyConfig=proxy_config_class or ProxyConfig, + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + clean_env: Final = { + k: v + for k, v in os.environ.items() + if k + not in ( + "DATABASE_URL", + "DIRECT_URL", + "JSON_LOGS", + "LITELLM_SUPPRESS_STARTUP_OUTPUT", + "CONFIG_FILE_PATH", + "LITELLM_CONFIG_BUCKET_NAME", + ) + } + config_env: Final = {} if config_file_path_env is None else {"CONFIG_FILE_PATH": str(config_file_path_env)} + config_args: Final = [] if config is None else ["--config", str(config)] + with ( + patch.dict(os.environ, {**clean_env, **config_env, **(extra_env or {})}, clear=True), + patch.dict( + "sys.modules", + {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, + ), + patch("atexit.register"), + patch("litellm.turn_on_json"), # test-quality-ok: it swaps every logger handler in this pytest worker + ): + run_server.main([*config_args, "--local", "--skip_server_startup"], standalone_mode=False) + return os.environ.get(name) def test_json_logs_true_exports_json_logs_to_spawned_workers( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -3991,11 +4183,126 @@ class TestJsonLogsConfigReachesSpawnedWorkers: assert _boot_from_config_file(config_path) is None + @pytest.mark.parametrize("litellm_settings", [[], "", 0, False]) + @pytest.mark.parametrize("via_config_file_path_env", [False, True]) + def test_an_empty_litellm_settings_still_boots( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, litellm_settings: object, via_config_file_path_env: bool + ) -> None: + monkeypatch.setattr(litellm, "suppress_startup_output", False) + config_file: Final = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"model_list": [], "litellm_settings": litellm_settings})) + exported: Final = ( + self._env_after_boot("LITELLM_SUPPRESS_STARTUP_OUTPUT", config_file_path_env=config_file) + if via_config_file_path_env + else self._env_after_boot("LITELLM_SUPPRESS_STARTUP_OUTPUT", config=config_file) + ) + assert exported is None + + @pytest.mark.parametrize(("setting", "exported"), [(True, "true"), ("true", None), (False, None)]) + def test_only_a_boolean_suppress_startup_output_reaches_spawned_workers( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: object, exported: str | None + ) -> None: + monkeypatch.setattr(litellm, "suppress_startup_output", False) + config: Final = self._write_config(tmp_path / "config.yaml", {"suppress_startup_output": setting}) + assert self._env_after_boot("LITELLM_SUPPRESS_STARTUP_OUTPUT", config=config) == exported + + @pytest.mark.parametrize(("name", "exported"), [("LITELLM_SUPPRESS_STARTUP_OUTPUT", "true"), ("JSON_LOGS", None)]) + def test_a_config_named_by_CONFIG_FILE_PATH_passes_only_suppress_startup_output_to_spawned_workers( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, name: str, exported: str | None + ) -> None: + """Its json_logs stays the workers' own to read, as before, so uvicorn's lines keep their plain format.""" + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(litellm, "suppress_startup_output", False) + settings: Final[dict[str, object]] = {"json_logs": True, "suppress_startup_output": True} + config_file: Final = self._write_config(tmp_path / "config.yaml", settings) + assert self._env_after_boot(name, config_file_path_env=config_file) == exported + + def test_a_CONFIG_FILE_PATH_quiet_config_wins_over_a_cli_config_too( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Workers load CONFIG_FILE_PATH ahead of --config, so its key quiets the CLI when both are set.""" + monkeypatch.setattr(litellm, "suppress_startup_output", False) + cli_config: Final = self._write_config(tmp_path / "cli.yaml", {}) + env_config: Final = self._write_config(tmp_path / "env.yaml", {"suppress_startup_output": True}) + exported: Final = self._env_after_boot( + "LITELLM_SUPPRESS_STARTUP_OUTPUT", config=cli_config, config_file_path_env=env_config + ) + assert exported == "true" + + @pytest.mark.parametrize( + ("config_file", "env_file", "extra_env"), + [ + ("config.yaml", "missing.yaml", {}), + ("config.json", "config.json", {}), + ("config.yaml", "config.yaml", {"LITELLM_CONFIG_BUCKET_NAME": "litellm-configs"}), + ], + ids=["missing-file", "not-yaml", "bucket-config"], + ) + def test_a_CONFIG_FILE_PATH_workers_skip_is_skipped_by_the_cli_too( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + config_file: str, + env_file: str, + extra_env: dict[str, str], + ) -> None: + """Workers load CONFIG_FILE_PATH only when it names a YAML file and no config bucket is set, and a boot + with any other value starts, so the CLI reads nothing from it either instead of failing the boot.""" + monkeypatch.setattr(litellm, "suppress_startup_output", False) + self._write_config(tmp_path / config_file, {"suppress_startup_output": True}) + exported: Final = self._env_after_boot( + "LITELLM_SUPPRESS_STARTUP_OUTPUT", config_file_path_env=tmp_path / env_file, extra_env=extra_env + ) + assert exported is None + + def test_a_CONFIG_FILE_PATH_boot_reads_no_secret_before_the_workers_do( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The CLI takes only the suppress key from CONFIG_FILE_PATH, so a config backed by a secret manager is + read from it once per worker, as before, not once more by the CLI.""" + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "suppress_startup_output", False) + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_system", None) + monkeypatch.setattr(litellm, "_key_management_settings", litellm._key_management_settings) + reads: Final = tmp_path / "secret_reads.log" + (tmp_path / "counting_secret_manager.py").write_text(COUNTING_SECRET_MANAGER_MODULE) + config_file: Final = tmp_path / "config.yaml" + config_file.write_text(SECRET_MANAGER_BACKED_CONFIG) + + exported: Final = self._env_after_boot( + "LITELLM_SUPPRESS_STARTUP_OUTPUT", + config_file_path_env=config_file, + extra_env={"SECRET_READ_LOG": str(reads)}, + proxy_config_class=ProxyConfig, + ) + + assert exported == "true" + assert not reads.exists(), reads.read_text() + + def test_a_CONFIG_FILE_PATH_key_set_in_an_included_file_reaches_spawned_workers( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "suppress_startup_output", False) + self._write_config(tmp_path / "quiet.yaml", {"suppress_startup_output": True}) + config_file: Final = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"include": ["quiet.yaml"], "model_list": []})) + + exported: Final = self._env_after_boot( + "LITELLM_SUPPRESS_STARTUP_OUTPUT", config_file_path_env=config_file, proxy_config_class=ProxyConfig + ) + + assert exported == "true" + class TestEnforcePrismaMigrationCheckNoticeAroundTheConfigRead: """With a config file, the `--enforce_prisma_migration_check` deprecation notice waits until the config - is read, so a `json_logs: true` set only there makes it a JSON record. It is printed once either way, - and a boot whose config read fails still prints it before the error, as it did before the wait.""" + is read, so a `json_logs: true` set only there makes it a JSON record, while `suppress_startup_output: true` + leaves it, since it is a warning. Otherwise it is printed once, and a boot whose config read fails still + prints it before the error, as it did before the wait.""" NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect" FLAG: Final = "--enforce_prisma_migration_check" @@ -4035,13 +4342,26 @@ class TestEnforcePrismaMigrationCheckNoticeAroundTheConfigRead: assert len(notice_lines) == 1, notice_lines assert self.NOTICE in json.loads(notice_lines[0])["message"] + @pytest.mark.parametrize("json_logs", [False, True]) + def test_a_config_only_suppress_startup_output_keeps_the_notice( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], json_logs: bool + ) -> None: + monkeypatch.setattr(litellm, "json_logs", json_logs) + monkeypatch.setattr(litellm, "suppress_startup_output", False) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"suppress_startup_output": True}})) + + _boot_from_config_file(config_path, self.FLAG) + + assert capsys.readouterr().out.count(self.NOTICE) == 1 + def test_granian_flag_notices_are_json_records_in_json_mode( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: """The yellow notices for the flags Granian ignores (`--max_requests_before_restart`, `--ciphers`) are - printed while the launcher starts; with json_logs on they are INFO records on stdout like every other - line, not plaintext between the records.""" + printed while the launcher starts; with json_logs on they are WARNING records on stdout after the INFO + start line, not plaintext between the records.""" pytest.importorskip("granian") monkeypatch.setattr(litellm, "json_logs", True) monkeypatch.setattr(sys, "excepthook", sys.excepthook) @@ -4072,4 +4392,4 @@ def test_granian_flag_notices_are_json_records_in_json_mode( "(Granian uses workers_lifetime in seconds, not a per-request limit).", "LiteLLM: --ciphers is not applied when using --run_granian.", ] - assert {record["level"] for record in records} == {"INFO"} + assert [record["level"] for record in records] == ["INFO", "WARNING", "WARNING"] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 55d1587669a..0c83d6c6c37 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15615,8 +15615,8 @@ def test_feedback_box_is_one_json_record_in_json_mode( async def test_initialize_prints_the_banner_between_blank_lines_in_plain_mode( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: - """Operators without JSON logs see the banner exactly as before, through show_banner()'s flushing - click.echo.""" + """Operators without JSON logs see the banner exactly as before, flushed at once as show_banner()'s + click.echo did.""" monkeypatch.setattr(litellm, "json_logs", False) monkeypatch.setenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "true") diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 195d5afe000..fcd835e93dd 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -7,7 +7,7 @@ import logging import re import sys import time -from io import StringIO +from io import BytesIO, StringIO, TextIOWrapper from pathlib import Path from typing import Final, List @@ -31,10 +31,11 @@ from litellm._logging import ( StdoutLogTruncationFilter, _get_uvicorn_json_log_config, _initialize_loggers_with_handler, - _parse_json_logs_env, + _parse_env_flag, _plain_log_format, _stdout_truncation_marker, print_startup_message, + print_startup_warning, turn_on_json, format_base64_size, session_id_var, @@ -46,6 +47,7 @@ from litellm._logging import ( verbose_router_logger, ) from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD +from litellm.proxy.common_utils.banner import LITELLM_BANNER from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import secret_redaction from litellm.types.utils import StandardLoggingPayload @@ -1402,6 +1404,32 @@ def test_plain_log_format_survives_hostile_streams(stdout, stderr): assert _plain_log_format(stdout, stderr) == _PLAIN_LOG_FORMAT +def _router_success_record() -> logging.LogRecord: + return logging.LogRecord( + "LiteLLM Router", logging.INFO, "router.py", 1, "litellm.acompletion(model=%s)\033[32m 200 OK\033[0m", ("gpt-5.6",), None + ) + + +@pytest.mark.parametrize( + ("log_format", "want_suffix"), + [ + (_COLOR_LOG_FORMAT, "litellm.acompletion(model=gpt-5.6)\033[32m 200 OK\033[0m"), + (_PLAIN_LOG_FORMAT, "router.py:1 - litellm.acompletion(model=gpt-5.6) 200 OK"), + ], +) +def test_plain_formatter_keeps_message_colors_only_with_the_color_format(log_format: str, want_suffix: str): + """A message that embeds color codes keeps them on a color terminal. Piped, or under NO_COLOR, the plain + format is chosen and the codes go too, so a container log never stores them as text.""" + assert CorrelationPlainFormatter(log_format).format(_router_success_record()).endswith(want_suffix) + + +def test_json_formatter_drops_color_codes_from_the_message(): + """A JSON log line is read by a log pipeline, never a terminal, so a colored message arrives as plain text.""" + assert json.loads(JsonFormatter().format(_router_success_record()))["message"] == ( + "litellm.acompletion(model=gpt-5.6) 200 OK" + ) + + def test_level_routing_handler_falls_back_to_stderr_when_stdout_is_unusable(monkeypatch, capsys): logger = logging.getLogger("test_level_routing_fallback") logger.handlers.clear() @@ -1436,9 +1464,9 @@ def test_level_routing_handler_falls_back_to_stderr_when_stdout_is_unusable(monk (None, False), ], ) -def test_parse_json_logs_env_enables_only_on_true(value, want): +def test_parse_env_flag_enables_only_on_true(value: str | None, want: bool): """JSON_LOGS=false / 0 must not enable JSON logs (LIT-5558).""" - assert _parse_json_logs_env(value) is want + assert _parse_env_flag(value) is want def test_plain_log_format_survives_none_streams(): @@ -1813,17 +1841,98 @@ def test_print_verbose_formats_args_only_when_set_verbose_is_on(monkeypatch, cap assert capsys.readouterr().out == "key k1 value cached-vector\n" -def test_print_startup_message_prints_plain_mode_text_unchanged( - monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] -) -> None: - """Operators without JSON logs keep the colored, spaced terminal output exactly as before.""" +class _CapturedStream(StringIO): + def __init__(self, tty: bool) -> None: + super().__init__() + self._tty = tty + + def isatty(self) -> bool: + return self._tty + + +@pytest.mark.parametrize( + "stdout_tty, stderr_tty, no_color, want_color", + [(True, True, None, True), (False, True, None, False), (True, False, None, False), (True, True, "1", False)], +) +def test_print_startup_message_keeps_colors_only_on_a_color_terminal( + monkeypatch: pytest.MonkeyPatch, stdout_tty: bool, stderr_tty: bool, no_color: str | None, want_color: bool +): + """Operators without JSON logs keep the colored, spaced output on a terminal. A container runtime reads + stdout through a pipe and stores the escape codes as text, so there, and under NO_COLOR, the colors go. + The rule is the plain log format's, so a `2>file` run drops the colors from prints and log lines alike.""" monkeypatch.setattr(litellm, "json_logs", False) + if no_color is None: + monkeypatch.delenv("NO_COLOR", raising=False) + else: + monkeypatch.setenv("NO_COLOR", no_color) + stdout: Final = _CapturedStream(stdout_tty) + monkeypatch.setattr(sys, "stdout", stdout) + monkeypatch.setattr(sys, "stderr", _CapturedStream(stderr_tty)) + assert (_plain_log_format(sys.stdout, sys.stderr) == _COLOR_LOG_FORMAT) is want_color print_startup_message("\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n") + assert stdout.getvalue() == ( + "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\n" + if want_color + else "\nLiteLLM: Proxy initialized with Config, Set models:\n\n" + ) + + +def test_print_startup_message_writes_the_banner_as_utf8_to_an_ascii_stdout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(litellm, "suppress_startup_output", False) + raw: Final = BytesIO() + monkeypatch.setattr(sys, "stdout", TextIOWrapper(raw, encoding="ascii")) + + print_startup_message(f"\n{LITELLM_BANNER}\n") + + assert raw.getvalue().decode("utf-8") == f"\n{LITELLM_BANNER}\n\n" + + +@pytest.mark.parametrize("json_logs", [False, True]) +def test_print_startup_message_writes_nothing_when_suppressed( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str], json_logs: bool +): + """suppress_startup_output turns the startup messages off in plain and JSON mode alike, for operators + whose log pipeline wants only the logger's own lines.""" + monkeypatch.setattr(litellm, "json_logs", json_logs) + monkeypatch.setattr(litellm, "suppress_startup_output", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + if json_logs: + turn_on_json() + + print_startup_message("\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m\n\033[32m web-search (tavily)\033[0m") + + assert capfd.readouterr() == ("", "") + + +@pytest.mark.parametrize("json_logs", [False, True]) +def test_print_startup_warning_still_prints_when_suppressed( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str], json_logs: bool +): + """A warning names something the operator has to fix, like an ignored flag or a failed migration diff, + so suppress_startup_output leaves it in place. In JSON mode it is a WARNING record at the caller's + location, on stdout where the print put it, so a collector's WARNING filter keeps it.""" + monkeypatch.setattr(litellm, "json_logs", json_logs) + monkeypatch.setattr(litellm, "suppress_startup_output", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + if json_logs: + turn_on_json() + + print_startup_warning("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + out, err = capfd.readouterr() - assert out == "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\n" assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + if json_logs: + record: Final = json.loads(lines[0]) + assert record["message"] == "LiteLLM: --ciphers is not applied when using --run_granian." + assert record["level"] == "WARNING" + assert str(record["logger"]).startswith("test_logging.py:"), record["logger"] + else: + assert lines[0] == "LiteLLM: --ciphers is not applied when using --run_granian." def test_print_startup_message_is_one_json_record_in_json_mode( @@ -1852,27 +1961,6 @@ def test_print_startup_message_is_one_json_record_in_json_mode( assert str(record["logger"]).startswith("test_logging.py:"), record["logger"] -def test_print_startup_message_yellow_notice_is_an_info_record_on_stdout( - monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] -) -> None: - """An operator notice printed in yellow (an ignored flag, a deprecation) stays on stdout in JSON mode, - where the print put it, as one INFO record with the color codes removed.""" - monkeypatch.setattr(litellm, "json_logs", True) - monkeypatch.setattr(sys, "excepthook", sys.excepthook) - turn_on_json() - verbose_proxy_logger.setLevel(logging.WARNING) - - print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") - - out, err = capfd.readouterr() - assert err == "" - lines: Final = [raw for raw in out.splitlines() if raw.strip()] - assert len(lines) == 1, lines - record: Final = json.loads(lines[0]) - assert record["message"] == "LiteLLM: --ciphers is not applied when using --run_granian." - assert record["level"] == "INFO" - - def test_print_startup_message_shows_above_the_json_handler_level( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: From e8f0735d8a935cf4bf6ec97c5ef6b8d846667b95 Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:47:27 -0700 Subject: [PATCH 19/58] feat(proxy): let a Claude Code gateway user choose the session team (#49001) * feat(proxy): let a Claude Code gateway user choose the session team The gateway's browser verification page now lists the developer's teams, and the device code grant mints the session for the picked team. A team the developer is not on is refused with a 400, and a form without a team keeps the first-team default. * fix(proxy): refuse a stale gateway team pick and require an explicit choice * fix(gateway): list the browser-picked team first so an older replica mid-deploy issues it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_endpoints/gateway_endpoints.py | 17 ++- .../html_forms/native_client_consent.py | 9 +- litellm/proxy/management_endpoints/ui_sso.py | 68 ++++++++- .../test_claude_code_gateway_oauth.py | 91 ++++++++++++ .../test_gateway_endpoints.py | 55 ++++++++ .../proxy/management_endpoints/test_ui_sso.py | 132 ++++++++++++++++++ 6 files changed, 364 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 1911cb574ef..6c58f2edc02 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -45,7 +45,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_set_request_parsed_body, ) from litellm.proxy.management_endpoints.sso_helper_utils import CLI_SSO_SESSIONS_TARGET -from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail +from litellm.proxy.management_endpoints.ui_sso import ( + CLI_SSO_BROWSER_TEAM_CHOICE, + CLI_SSO_CHOSEN_TEAM_ID, + CliSsoTeamDetail, +) from litellm.proxy.mcp_server.gateway_dcr_flow import ( PROXY_API_AUDIENCE, MintedProxyCredential, @@ -235,6 +239,7 @@ async def device_authorization(request: Request) -> JSONResponse: "sso_complete": False, "user_code_verified": False, "session_data": None, + CLI_SSO_BROWSER_TEAM_CHOICE: True, } _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) @@ -267,7 +272,15 @@ def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError: status_code=400, error="invalid_grant", description="The login session is malformed; sign in again" ) - team_id: Final = session_data.teams[0] if session_data.teams else None + chosen: Final = flow.get(CLI_SSO_CHOSEN_TEAM_ID) + chosen_team_id: Final = chosen if isinstance(chosen, str) else None + if chosen_team_id is not None and chosen_team_id not in session_data.teams: + return _OAuthError( + status_code=400, + error="invalid_grant", + description="The team picked at sign-in is no longer one of your teams; sign in again", + ) + team_id: Final = chosen_team_id or next(iter(session_data.teams), None) selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id) if selected_team is None: return _OAuthError( diff --git a/litellm/proxy/common_utils/html_forms/native_client_consent.py b/litellm/proxy/common_utils/html_forms/native_client_consent.py index dac92c4e787..aeadc93221c 100644 --- a/litellm/proxy/common_utils/html_forms/native_client_consent.py +++ b/litellm/proxy/common_utils/html_forms/native_client_consent.py @@ -62,7 +62,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p

Approving issues it a personal credential that expires within {CLI_JWT_EXPIRATION_HOURS} hours. lite logout stops it from being renewed. Only approve if you started this sign-in yourself.

-{_team_field(teams)} +{team_select_field(teams)}
@@ -74,7 +74,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p """ -def _team_field(teams: Sequence[tuple[str, str]]) -> str: +def team_select_field(teams: Sequence[tuple[str, str]], *, require_choice: bool = False) -> str: if not teams: return "" if len(teams) == 1: @@ -83,9 +83,12 @@ def _team_field(teams: Sequence[tuple[str, str]]) -> str: f'' f"

Requests are attributed to team {escape(team_label)}.

" ) + placeholder: Final = '' if require_choice else "" options: Final = "".join( f'' for team_id, team_label in teams ) + required: Final = " required" if require_choice else "" return ( - f'' + '' + f'' ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 22c4f27d94a..a980b9c33ff 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,6 +16,7 @@ import json import os import re import secrets +from collections import Counter from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from html import escape @@ -42,7 +43,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator +from pydantic import AliasPath, BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator import litellm from litellm._internal_context import with_service_target @@ -106,6 +107,7 @@ from litellm.proxy.common_utils.html_forms.default_credentials_hint import shoul from litellm.proxy.common_utils.html_forms.jwt_display_template import ( jwt_display_template, ) +from litellm.proxy.common_utils.html_forms.native_client_consent import team_select_field from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -184,6 +186,8 @@ _CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS: Final = 60 _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS: Final = 30 _CLI_SSO_USER_CODE_ALPHABET: Final = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" _CLI_SSO_LOGIN_ID_RE: Final = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") +CLI_SSO_BROWSER_TEAM_CHOICE: Final = "browser_team_choice" +CLI_SSO_CHOSEN_TEAM_ID: Final = "chosen_team_id" _CLI_SSO_USER_CODE_RE = re.compile(rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$") _CLI_SSO_SCALAR_TYPES: Final = (str, int, float, bool) _CLI_SSO_DEST_KEY_RE: Final = re.compile(r"^[A-Za-z0-9_.-]+$") @@ -670,6 +674,7 @@ def _render_cli_sso_verification_page( verify_url: str, browser_complete_token: str, prefill_user_code: str | None = None, + team_choices: Sequence[tuple[str, str]] = (), ) -> str: escaped_verify_url: Final = escape(verify_url, quote=True) escaped_browser_complete_token: Final = escape(browser_complete_token, quote=True) @@ -716,6 +721,15 @@ def _render_cli_sso_verification_page( letter-spacing: 0.08em; text-transform: uppercase; }} + select {{ + box-sizing: border-box; + width: 100%; + padding: 12px; + margin-bottom: 18px; + border: 1px solid #cbd5e1; + border-radius: 6px; + font-size: 16px; + }} button {{ margin-top: 16px; width: 100%; @@ -735,6 +749,7 @@ def _render_cli_sso_verification_page(

{instructions}

+ {team_select_field(team_choices, require_choice=True)} @@ -823,13 +838,46 @@ async def cli_sso_complete(request: Request, login_id: str): ): raise HTTPException(status_code=400, detail="Invalid verification code") - flow["user_code_verified"] = True - _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) + team_choice_flow: Final = _CliSsoTeamChoiceFlow.model_validate(flow) + chosen_team_id: Final = _browser_chosen_team_id(team_choice_flow, form_values) + if chosen_team_id is not None and chosen_team_id not in team_choice_flow.teams: + raise HTTPException(status_code=400, detail="Choose one of the teams you belong to") + + _set_cli_sso_flow( + login_id=login_id, + cache=cli_sso_session_cache, + flow={ + **flow, + "user_code_verified": True, + CLI_SSO_CHOSEN_TEAM_ID: chosen_team_id, + "session_data": ( + flow["session_data"] + if chosen_team_id is None + else { + **flow["session_data"], + "teams": [chosen_team_id, *(team for team in team_choice_flow.teams if team != chosen_team_id)], + } + ), + }, + ) html_content: Final = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) +class _CliSsoTeamChoiceFlow(BaseModel): + model_config = ConfigDict(frozen=True) + + browser_team_choice: bool = Field(default=False, validation_alias=CLI_SSO_BROWSER_TEAM_CHOICE) + teams: tuple[str, ...] = Field(default=(), validation_alias=AliasPath("session_data", "teams")) + + +def _browser_chosen_team_id(flow: _CliSsoTeamChoiceFlow, form_values: Mapping[str, Sequence[str]]) -> str | None: + if not flow.browser_team_choice: + return None + return (form_values.get("team_id") or [""])[0] or None + + def normalize_email(email: str | None) -> str | None: """ Normalize email address to lowercase for consistent storage and comparison. @@ -2416,6 +2464,18 @@ def _cli_sso_session_teams(team_details: Sequence[CliSsoTeamDetail]) -> list[str return [detail.team_id for detail in team_details if detail.team_id is not None] +def _browser_team_choices( + flow: _CliSsoTeamChoiceFlow, team_details: Sequence[CliSsoTeamDetail] +) -> tuple[tuple[str, str], ...]: + if not flow.browser_team_choice: + return () + labeled: Final = tuple( + (detail.team_id, detail.team_alias or detail.team_id) for detail in team_details if detail.team_id is not None + ) + label_counts: Final = Counter(label for _, label in labeled) + return tuple((team_id, label if label_counts[label] == 1 else f"{label} ({team_id})") for team_id, label in labeled) + + def selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> CliSsoTeamDetail | None: """``None`` means the team's grants are unknown. An empty grant is a real value meaning unrestricted, so an unknown one must not be minted as empty.""" @@ -2484,6 +2544,7 @@ async def _complete_cli_sso_callback_session( request.scope.get("litellm_microsoft_interactive_subject"), resolved_user_id, prisma_client ) resolved_teams: Final = _cli_sso_session_teams(team_details) + team_choices: Final = _browser_team_choices(_CliSsoTeamChoiceFlow.model_validate(flow), team_details) attribution_metadata: Final = build_cli_sso_attribution_metadata(result=result) if attribution_metadata: await _persist_cli_sso_user_metadata( @@ -2521,6 +2582,7 @@ async def _complete_cli_sso_callback_session( verify_url=verify_url, browser_complete_token=browser_complete_token, prefill_user_code=prefill_user_code, + team_choices=team_choices, ), status_code=200, ) diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index a80eeacad46..4dedbdcd8e4 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -545,3 +545,94 @@ def test_device_authorization_polls_complete_sso_and_mints_a_model_restricted_be assert set(token_form) == {"code", "code_verifier", "grant_type", "redirect_uri"}, token_form assert userinfo_request.method == "GET" and urlsplit(userinfo_request.target).path == "/userinfo" assert user_id == _USER_ID + + +_TEAM_OPTION: Final = re.compile(r'