From ffbc8d20c4864b8f63f328c349478a9ca6f64f56 Mon Sep 17 00:00:00 2001 From: Nikita Timofeev Date: Tue, 3 Feb 2026 12:58:28 +0000 Subject: [PATCH 1/5] bugfix: Remove user messages merging There is no reason to merge user messages. --- litellm/llms/gigachat/chat/transformation.py | 28 +--------------- tests/llm_translation/test_gigachat.py | 34 -------------------- 2 files changed, 1 insertion(+), 61 deletions(-) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index ba14de1f65..f546f356e1 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -386,33 +386,7 @@ class GigaChatConfig(BaseConfig): transformed.append(message) - # Collapse consecutive user messages - return self._collapse_user_messages(transformed) - - def _collapse_user_messages(self, messages: List[dict]) -> List[dict]: - """Collapse consecutive user messages into one.""" - collapsed: List[dict] = [] - prev_user_msg: Optional[dict] = None - content_parts: List[str] = [] - - for msg in messages: - if msg.get("role") == "user" and prev_user_msg is not None: - content_parts.append(msg.get("content", "")) - else: - if content_parts and prev_user_msg: - prev_user_msg["content"] = "\n".join( - [prev_user_msg.get("content", "")] + content_parts - ) - content_parts = [] - collapsed.append(msg) - prev_user_msg = msg if msg.get("role") == "user" else None - - if content_parts and prev_user_msg: - prev_user_msg["content"] = "\n".join( - [prev_user_msg.get("content", "")] + content_parts - ) - - return collapsed + return transformed def transform_response( self, diff --git a/tests/llm_translation/test_gigachat.py b/tests/llm_translation/test_gigachat.py index b69a5428e4..631ae94d20 100644 --- a/tests/llm_translation/test_gigachat.py +++ b/tests/llm_translation/test_gigachat.py @@ -122,40 +122,6 @@ class TestGigaChatCollapseUserMessages: return GigaChatConfig() - def test_no_collapse_single_message(self, config): - """Single message should not be changed""" - messages = [{"role": "user", "content": "Hello"}] - result = config._collapse_user_messages(messages) - - assert len(result) == 1 - assert result[0]["content"] == "Hello" - - def test_collapse_consecutive_user_messages(self, config): - """Consecutive user messages should be collapsed""" - messages = [ - {"role": "user", "content": "First"}, - {"role": "user", "content": "Second"}, - {"role": "user", "content": "Third"}, - ] - result = config._collapse_user_messages(messages) - - assert len(result) == 1 - assert "First" in result[0]["content"] - assert "Second" in result[0]["content"] - assert "Third" in result[0]["content"] - - def test_no_collapse_with_assistant_between(self, config): - """Messages with assistant between should not be collapsed""" - messages = [ - {"role": "user", "content": "First"}, - {"role": "assistant", "content": "Response"}, - {"role": "user", "content": "Second"}, - ] - result = config._collapse_user_messages(messages) - - assert len(result) == 3 - - class TestGigaChatToolsTransformation: """Tests for tools -> functions conversion""" From 18f4b7219b9e4b3edda77f22f7a1b6b4845edc55 Mon Sep 17 00:00:00 2001 From: jquinter Date: Wed, 4 Feb 2026 03:29:29 -0300 Subject: [PATCH 2/5] feat: add faster linting targets for development workflow (#19729) * feat: add faster linting targets for development workflow - Add lint-dev target that only checks changed files vs origin/main - Add lint-format-changed to format only modified Python lines - Add lint-ruff-dev using diff-quality for incremental lint checks - Upgrade ruff from 0.1.x to 0.2.x for --range formatting support - Add pylint and diff-cover as dev dependencies - Use portable PIP variable for cross-platform compatibility - Suppress poetry warnings in install-dev target * fix(mypy): fix type: ignore placement for OTEL LogRecord import The type: ignore[attr-defined] comment was on the import alias line inside parentheses, but mypy reports the error on the `from` line. Collapse to single-line imports so the suppression is on the correct line. Also add no-redef to the fallback branch. * fix: address review issues in faster linting PR - Remove poetry lock/check from install-dev (slow, can mutate lockfile) - Remove misplaced [virtualenvs] and [installer] from pyproject.toml (these belong in poetry.toml, not project metadata) - Remove unused pylint dev dependency (diff-quality uses pylint output format, not the pylint package itself) - Fix trailing whitespace in .PHONY declaration - Use mktemp instead of hardcoded /tmp/ruff.txt in lint-ruff-dev - Guard lint-ruff-FULL-dev against empty file list from git diff - Fix incorrect comment on lint-dev target - Regenerate poetry.lock Co-Authored-By: Claude Opus 4.5 * fix: address review issues in faster linting PR - Remove poetry lock/check from install-dev (slow, can mutate lockfile) - Remove misplaced [virtualenvs] and [installer] from pyproject.toml (these belong in poetry.toml, not project metadata) - Remove unused pylint dev dependency (diff-quality uses pylint output format, not the pylint package itself) - Fix trailing whitespace in .PHONY declaration - Use mktemp instead of hardcoded /tmp/ruff.txt in lint-ruff-dev - Guard lint-ruff-FULL-dev against empty file list from git diff - Fix incorrect comment on lint-dev target - Regenerate poetry.lock --------- Co-authored-by: Claude Opus 4.5 --- Makefile | 65 +++++++++++++++++---- litellm/integrations/opentelemetry.py | 8 +-- poetry.lock | 81 +++++++++++++++++++-------- pyproject.toml | 3 +- 4 files changed, 116 insertions(+), 41 deletions(-) diff --git a/Makefile b/Makefile index 0da83c363c..b867d7ea35 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,10 @@ # LiteLLM Makefile # Simple Makefile for running tests and basic development tasks -.PHONY: help test test-unit test-integration test-unit-helm lint format install-dev install-proxy-dev install-test-deps install-helm-unittest check-circular-imports check-import-safety +.PHONY: help test test-unit test-integration test-unit-helm \ + info lint lint-dev format \ + install-dev install-proxy-dev install-test-deps \ + install-helm-unittest check-circular-imports check-import-safety # Default target help: @@ -25,6 +28,13 @@ help: @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" +# Keep PIP simple for edge cases: +PIP := $(shell command -v pip > /dev/null 2>&1 && echo "pip" || echo "python3 -m pip") + +# Show info +info: + @echo "PIP: $(PIP)" + # Installation targets install-dev: poetry install --with dev @@ -34,19 +44,19 @@ install-proxy-dev: # CI-compatible installations (matches GitHub workflows exactly) install-dev-ci: - pip install openai==2.8.0 + $(PIP) install openai==2.8.0 poetry install --with dev - pip install openai==2.8.0 + $(PIP) install openai==2.8.0 install-proxy-dev-ci: poetry install --with dev,proxy-dev --extras proxy - pip install openai==2.8.0 + $(PIP) install openai==2.8.0 install-test-deps: install-proxy-dev - poetry run pip install "pytest-retry==1.6.3" - poetry run pip install pytest-xdist - poetry run pip install openapi-core - cd enterprise && poetry run pip install -e . && cd .. + poetry run $(PIP) install "pytest-retry==1.6.3" + poetry run $(PIP) install pytest-xdist + poetry run $(PIP) install openapi-core + cd enterprise && poetry run $(PIP) install -e . && cd .. install-helm-unittest: helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists" @@ -62,8 +72,40 @@ format-check: install-dev lint-ruff: install-dev cd litellm && poetry run ruff check . && cd .. +# faster linter for developing ... +# inspiration from: +# https://github.com/astral-sh/ruff/discussions/10977 +# https://github.com/astral-sh/ruff/discussions/4049 +lint-format-changed: install-dev + @git diff origin/main --unified=0 --no-color -- '*.py' | \ + perl -ne '\ + if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \ + if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \ + $$start = $$1; $$count = $$2 || 1; $$end = $$start + $$count - 1; \ + print "$$file:$$start:1-$$end:999\n"; \ + }' | \ + while read range; do \ + file="$${range%%:*}"; \ + lines="$${range#*:}"; \ + echo "Formatting $$file (lines $$lines)"; \ + poetry run ruff format --range "$$lines" "$$file"; \ + done + +lint-ruff-dev: install-dev + @tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \ + cd litellm && \ + (poetry run ruff check . --output-format=pylint || true) > "$$tmpfile" && \ + poetry run diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \ + cd .. ; \ + rm -f "$$tmpfile" + +lint-ruff-FULL-dev: install-dev + @files=$$(git diff --name-only origin/main -- '*.py'); \ + if [ -n "$$files" ]; then echo "$$files" | xargs poetry run ruff check; \ + else echo "No changed .py files to check."; fi + lint-mypy: install-dev - poetry run pip install types-requests types-setuptools types-redis types-PyYAML + poetry run $(PIP) install types-requests types-setuptools types-redis types-PyYAML cd litellm && poetry run mypy . --ignore-missing-imports && cd .. lint-black: format-check @@ -72,11 +114,14 @@ check-circular-imports: install-dev cd litellm && poetry run python ../tests/documentation_tests/test_circular_imports.py && cd .. check-import-safety: install-dev - poetry run python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1) + @poetry run python -c "from litellm import *; print('[from litellm import *] OK! no issues!');" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1) # Combined linting (matches test-linting.yml workflow) lint: format-check lint-ruff lint-mypy check-circular-imports check-import-safety +# Faster linting for local development (only checks changed code) +lint-dev: lint-format-changed lint-mypy check-circular-imports check-import-safety + # Testing targets test: poetry run pytest tests/ diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 18898be7dc..21ae7b3c7e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1004,13 +1004,9 @@ class OpenTelemetry(CustomLogger): from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider try: - from opentelemetry.sdk._logs import ( - LogRecord as SdkLogRecord, # type: ignore[attr-defined] # OTEL < 1.39.0 - ) + from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # type: ignore[attr-defined] # OTEL < 1.39.0 except ImportError: - from opentelemetry.sdk._logs._internal import ( - LogRecord as SdkLogRecord, # OTEL >= 1.39.0 - ) + from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord # type: ignore[attr-defined, no-redef] # OTEL >= 1.39.0 otel_logger = get_logger(LITELLM_LOGGER_NAME) diff --git a/poetry.lock b/poetry.lock index 537367c5aa..1e81bf3374 100644 --- a/poetry.lock +++ b/poetry.lock @@ -723,6 +723,18 @@ markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"prox [package.dependencies] pycparser = {version = "*", markers = "implementation_name != \"PyPy\""} +[[package]] +name = "chardet" +version = "5.2.0" +description = "Universal encoding detector for Python 3" +optional = false +python-versions = ">=3.7" +groups = ["dev"] +files = [ + {file = "chardet-5.2.0-py3-none-any.whl", hash = "sha256:e1cf59446890a00105fe7b7912492ea04b6e6f06d4b742b2c788469e34c82970"}, + {file = "chardet-5.2.0.tar.gz", hash = "sha256:1b3b6ff479a8c414bc3fa2c0852995695c4a026dcd6d0633b2dd092ca39c1cf7"}, +] + [[package]] name = "charset-normalizer" version = "3.4.4" @@ -1300,6 +1312,27 @@ dev = ["autoflake", "black", "build", "databricks-connect", "httpx", "ipython", notebook = ["ipython (>=8,<10)", "ipywidgets (>=8,<9)"] openai = ["httpx", "langchain-openai ; python_version > \"3.7\"", "openai"] +[[package]] +name = "diff-cover" +version = "9.7.2" +description = "Run coverage and linting reports on diffs" +optional = false +python-versions = ">=3.9" +groups = ["dev"] +files = [ + {file = "diff_cover-9.7.2-py3-none-any.whl", hash = "sha256:cd6498620c747c2493a6c83c14362c32868bfd91cd8d0dd093f136070ec4ffc5"}, + {file = "diff_cover-9.7.2.tar.gz", hash = "sha256:872c820d2ecbf79c61d52c7dc70419015e0ab9289589566c791dd270fc0c6e3b"}, +] + +[package.dependencies] +chardet = ">=3.0.0" +Jinja2 = ">=2.7.1" +pluggy = ">=0.13.1,<2" +Pygments = ">=2.19.1,<3.0.0" + +[package.extras] +toml = ["tomli (>=1.2.1)"] + [[package]] name = "diskcache" version = "5.6.3" @@ -3085,7 +3118,7 @@ version = "3.1.6" description = "A very fast and expressive template engine." optional = false python-versions = ">=3.7" -groups = ["main", "proxy-dev"] +groups = ["main", "dev", "proxy-dev"] files = [ {file = "jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67"}, {file = "jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d"}, @@ -3515,7 +3548,7 @@ version = "3.0.3" description = "Safely add untrusted strings to HTML/XML markup." optional = false python-versions = ">=3.9" -groups = ["main", "proxy-dev"] +groups = ["main", "dev", "proxy-dev"] files = [ {file = "markupsafe-3.0.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559"}, {file = "markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419"}, @@ -5516,14 +5549,14 @@ files = [ name = "pygments" version = "2.19.2" description = "Pygments is a syntax highlighting package written in Python." -optional = true +optional = false python-versions = ">=3.8" -groups = ["main"] -markers = "extra == \"utils\" or extra == \"proxy\"" +groups = ["main", "dev"] files = [ {file = "pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b"}, {file = "pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887"}, ] +markers = {main = "extra == \"utils\" or extra == \"proxy\""} [package.extras] windows-terminal = ["colorama (>=0.4.6)"] @@ -6601,29 +6634,29 @@ pyasn1 = ">=0.1.3" [[package]] name = "ruff" -version = "0.1.15" +version = "0.2.2" description = "An extremely fast Python linter and code formatter, written in Rust." optional = false python-versions = ">=3.7" groups = ["dev"] files = [ - {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:5fe8d54df166ecc24106db7dd6a68d44852d14eb0729ea4672bb4d96c320b7df"}, - {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:6f0bfbb53c4b4de117ac4d6ddfd33aa5fc31beeaa21d23c45c6dd249faf9126f"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e0d432aec35bfc0d800d4f70eba26e23a352386be3a6cf157083d18f6f5881c8"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9405fa9ac0e97f35aaddf185a1be194a589424b8713e3b97b762336ec79ff807"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c66ec24fe36841636e814b8f90f572a8c0cb0e54d8b5c2d0e300d28a0d7bffec"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:6f8ad828f01e8dd32cc58bc28375150171d198491fc901f6f98d2a39ba8e3ff5"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:86811954eec63e9ea162af0ffa9f8d09088bab51b7438e8b6488b9401863c25e"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fd4025ac5e87d9b80e1f300207eb2fd099ff8200fa2320d7dc066a3f4622dc6b"}, - {file = "ruff-0.1.15-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b17b93c02cdb6aeb696effecea1095ac93f3884a49a554a9afa76bb125c114c1"}, - {file = "ruff-0.1.15-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ddb87643be40f034e97e97f5bc2ef7ce39de20e34608f3f829db727a93fb82c5"}, - {file = "ruff-0.1.15-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:abf4822129ed3a5ce54383d5f0e964e7fef74a41e48eb1dfad404151efc130a2"}, - {file = "ruff-0.1.15-py3-none-musllinux_1_2_i686.whl", hash = "sha256:6c629cf64bacfd136c07c78ac10a54578ec9d1bd2a9d395efbee0935868bf852"}, - {file = "ruff-0.1.15-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1bab866aafb53da39c2cadfb8e1c4550ac5340bb40300083eb8967ba25481447"}, - {file = "ruff-0.1.15-py3-none-win32.whl", hash = "sha256:2417e1cb6e2068389b07e6fa74c306b2810fe3ee3476d5b8a96616633f40d14f"}, - {file = "ruff-0.1.15-py3-none-win_amd64.whl", hash = "sha256:3837ac73d869efc4182d9036b1405ef4c73d9b1f88da2413875e34e0d6919587"}, - {file = "ruff-0.1.15-py3-none-win_arm64.whl", hash = "sha256:9a933dfb1c14ec7a33cceb1e49ec4a16b51ce3c20fd42663198746efc0427360"}, - {file = "ruff-0.1.15.tar.gz", hash = "sha256:f6dfa8c1b21c913c326919056c390966648b680966febcb796cc9d1aaab8564e"}, + {file = "ruff-0.2.2-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:0a9efb032855ffb3c21f6405751d5e147b0c6b631e3ca3f6b20f917572b97eb6"}, + {file = "ruff-0.2.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:d450b7fbff85913f866a5384d8912710936e2b96da74541c82c1b458472ddb39"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ecd46e3106850a5c26aee114e562c329f9a1fbe9e4821b008c4404f64ff9ce73"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5e22676a5b875bd72acd3d11d5fa9075d3a5f53b877fe7b4793e4673499318ba"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1695700d1e25a99d28f7a1636d85bafcc5030bba9d0578c0781ba1790dbcf51c"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:b0c232af3d0bd8f521806223723456ffebf8e323bd1e4e82b0befb20ba18388e"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f63d96494eeec2fc70d909393bcd76c69f35334cdbd9e20d089fb3f0640216ca"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6a61ea0ff048e06de273b2e45bd72629f470f5da8f71daf09fe481278b175001"}, + {file = "ruff-0.2.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5e1439c8f407e4f356470e54cdecdca1bd5439a0673792dbe34a2b0a551a2fe3"}, + {file = "ruff-0.2.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:940de32dc8853eba0f67f7198b3e79bc6ba95c2edbfdfac2144c8235114d6726"}, + {file = "ruff-0.2.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:0c126da55c38dd917621552ab430213bdb3273bb10ddb67bc4b761989210eb6e"}, + {file = "ruff-0.2.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:3b65494f7e4bed2e74110dac1f0d17dc8e1f42faaa784e7c58a98e335ec83d7e"}, + {file = "ruff-0.2.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1ec49be4fe6ddac0503833f3ed8930528e26d1e60ad35c2446da372d16651ce9"}, + {file = "ruff-0.2.2-py3-none-win32.whl", hash = "sha256:d920499b576f6c68295bc04e7b17b6544d9d05f196bb3aac4358792ef6f34325"}, + {file = "ruff-0.2.2-py3-none-win_amd64.whl", hash = "sha256:cc9a91ae137d687f43a44c900e5d95e9617cb37d4c989e462980ba27039d239d"}, + {file = "ruff-0.2.2-py3-none-win_arm64.whl", hash = "sha256:c9d15fc41e6054bfc7200478720570078f0b41c9ae4f010bcc16bd6f4d1aacdd"}, + {file = "ruff-0.2.2.tar.gz", hash = "sha256:e62ed7f36b3068a30ba39193a14274cd706bc486fad521276458022f7bccb31d"}, ] [[package]] @@ -8490,4 +8523,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "95fd27dc139d0e52e70093220c50582f16c78e5977ec77f4297f50a30df964c6" +content-hash = "797603dcfef0a79781c7d3cba5dfe18f6aea4aa792220f47487ebc7bd04ae2e3" diff --git a/pyproject.toml b/pyproject.toml index 9832ca483d..2361103501 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,7 @@ litellm = 'litellm:run_server' litellm-proxy = 'litellm.proxy.client.cli:cli' [tool.poetry.group.dev.dependencies] +diff-cover = "^9.0" flake8 = "^6.1.0" black = "^23.12.0" mypy = "^1.0" @@ -149,7 +150,7 @@ pytest-retry = "^1.6.3" requests-mock = "^1.12.1" responses = "^0.25.7" respx = "^0.22.0" -ruff = "^0.1.0" +ruff = "^0.2.1" types-requests = "*" types-setuptools = "*" types-redis = "*" From 13130ea3e16747408c088f6ddea77c617c7f4da9 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 4 Feb 2026 12:10:19 +0530 Subject: [PATCH 3/5] Litellm fix langfuse otel trace (#20382) * fix: support multi-project keys and fix trace leakage * fix: Langfuse otel handle * fix lint errors mypy * passing all test case --------- Co-authored-by: Krish Dholakia --- .../integrations/langfuse/langfuse_otel.py | 79 ++- litellm/integrations/opentelemetry.py | 7 +- .../initialize_dynamic_callback_params.py | 53 +- litellm/litellm_core_utils/litellm_logging.py | 16 +- .../test_dynamic_otel_keys.py | 52 ++ .../integrations/test_langfuse_otel.py | 453 +++++++++++------- 6 files changed, 442 insertions(+), 218 deletions(-) create mode 100644 tests/logging_callback_tests/test_dynamic_otel_keys.py diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 08493a0e8e..8955d3619f 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -8,9 +8,8 @@ from litellm.integrations.arize import _utils from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) -from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.types.integrations.langfuse_otel import ( - LangfuseOtelConfig, LangfuseSpanAttributes, ) from litellm.types.utils import StandardCallbackDynamicParams @@ -18,17 +17,8 @@ from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.integrations.opentelemetry import ( - OpenTelemetryConfig as _OpenTelemetryConfig, - ) - from litellm.types.integrations.arize import Protocol as _Protocol - - Protocol = _Protocol - OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] else: - Protocol = Any - OpenTelemetryConfig = Any Span = Any @@ -37,8 +27,12 @@ LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel" class LangfuseOtelLogger(OpenTelemetry): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config=None, *args, **kwargs): + # Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually + # and passing it to the parent OpenTelemetry class + if config is None: + config = self._create_open_telemetry_config_from_langfuse_env() + super().__init__(config=config, *args, **kwargs) @staticmethod def set_langfuse_otel_attributes(span: Span, kwargs, response_obj): @@ -114,6 +108,10 @@ class LangfuseOtelLogger(OpenTelemetry): for key, enum_attr in mapping.items(): if key in metadata and metadata[key] is not None: value = metadata[key] + if key == "trace_id" and isinstance(value, str): + # trace_id must be 32 hex char no dashes for langfuse : Litellm sends uuid with dashes (might be breaking at some point) + value = value.replace("-", "") + if isinstance(value, (list, dict)): try: value = json.dumps(value) @@ -265,8 +263,47 @@ class LangfuseOtelLogger(OpenTelemetry): """ return os.environ.get("LANGFUSE_OTEL_HOST") or os.environ.get("LANGFUSE_HOST") + def _create_open_telemetry_config_from_langfuse_env(self) -> OpenTelemetryConfig: + """ + Creates OpenTelemetryConfig from Langfuse environment variables. + Does NOT modify global environment variables. + """ + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + public_key = os.environ.get("LANGFUSE_PUBLIC_KEY", None) + secret_key = os.environ.get("LANGFUSE_SECRET_KEY", None) + + if not public_key or not secret_key: + # If no keys, return default from env (likely logging to console or something else) + return OpenTelemetryConfig.from_env() + + # Determine endpoint - default to US cloud + langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host() + + if langfuse_host: + # If LANGFUSE_HOST is provided, construct OTEL endpoint from it + if not langfuse_host.startswith("http"): + langfuse_host = "https://" + langfuse_host + endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel" + verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}") + else: + # Default to US cloud endpoint + endpoint = LANGFUSE_CLOUD_US_ENDPOINT + verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}") + + auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( + public_key=public_key, secret_key=secret_key + ) + otlp_auth_headers = f"Authorization={auth_header}" + + return OpenTelemetryConfig( + exporter="otlp_http", + endpoint=endpoint, + headers=otlp_auth_headers, + ) + @staticmethod - def get_langfuse_otel_config() -> LangfuseOtelConfig: + def get_langfuse_otel_config() -> "OpenTelemetryConfig": """ Retrieves the Langfuse OpenTelemetry configuration based on environment variables. @@ -276,7 +313,7 @@ class LangfuseOtelLogger(OpenTelemetry): LANGFUSE_HOST: Optional. Custom Langfuse host URL. Defaults to US cloud. Returns: - LangfuseOtelConfig: A Pydantic model containing Langfuse OTEL configuration. + OpenTelemetryConfig: A Pydantic model containing Langfuse OTEL configuration. Raises: ValueError: If required keys are missing. @@ -308,12 +345,14 @@ class LangfuseOtelLogger(OpenTelemetry): ) otlp_auth_headers = f"Authorization={auth_header}" - # Set standard OTEL environment variables - os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint - os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers + # Prevent modification of global env vars which causes leakage + # os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint + # os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers - return LangfuseOtelConfig( - otlp_auth_headers=otlp_auth_headers, protocol="otlp_http" + return OpenTelemetryConfig( + exporter="otlp_http", + endpoint=endpoint, + headers=otlp_auth_headers, ) @staticmethod diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 21ae7b3c7e..e025c08b02 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -674,7 +674,10 @@ class OpenTelemetry(CustomLogger): kwargs, response_obj, start_time, end_time, span ) # Ensure proxy-request parent span is annotated with the actual operation kind - if parent_span is not None and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME: + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): self.set_attributes(parent_span, kwargs, response_obj) else: # Do not create primary span (keep hierarchy shallow when parent exists) @@ -1003,6 +1006,7 @@ class OpenTelemetry(CustomLogger): # TODO: Refactor to use the proper OTEL Logs API instead of directly creating SDK LogRecords from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider + try: from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # type: ignore[attr-defined] # OTEL < 1.39.0 except ImportError: @@ -1718,6 +1722,7 @@ class OpenTelemetry(CustomLogger): def set_raw_request_attributes(self, span: Span, kwargs, response_obj): try: + self.set_attributes(span, kwargs, response_obj) kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index c425319b4d..78846f8e82 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -1,8 +1,34 @@ from typing import Dict, Optional - from litellm.secret_managers.main import get_secret_str from litellm.types.utils import StandardCallbackDynamicParams +# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict +_supported_callback_params = [ + "langfuse_public_key", + "langfuse_secret", + "langfuse_secret_key", + "langfuse_host", + "langfuse_prompt_version", + "gcs_bucket_name", + "gcs_path_service_account", + "langsmith_api_key", + "langsmith_project", + "langsmith_base_url", + "langsmith_sampling_rate", + "langsmith_tenant_id", + "humanloop_api_key", + "arize_api_key", + "arize_space_key", + "arize_space_id", + "posthog_api_key", + "posthog_host", + "braintrust_api_key", + "braintrust_project", + "braintrust_host", + "slack_webhook_url", + "lunary_public_key", +] + def initialize_standard_callback_dynamic_params( kwargs: Optional[Dict] = None, @@ -15,13 +41,10 @@ def initialize_standard_callback_dynamic_params( standard_callback_dynamic_params = StandardCallbackDynamicParams() if kwargs: - _supported_callback_params = ( - StandardCallbackDynamicParams.__annotations__.keys() - ) - + # 1. Check top-level kwargs for param in _supported_callback_params: if param in kwargs: - _param_value = kwargs.pop(param) + _param_value = kwargs.get(param) if ( _param_value is not None and isinstance(_param_value, str) @@ -30,4 +53,22 @@ def initialize_standard_callback_dynamic_params( _param_value = get_secret_str(secret_name=_param_value) standard_callback_dynamic_params[param] = _param_value # type: ignore + # 2. Fallback: check "metadata" or "litellm_params" -> "metadata" + metadata = (kwargs.get("metadata") or {}).copy() + litellm_params = kwargs.get("litellm_params") or {} + if isinstance(litellm_params, dict): + metadata.update(litellm_params.get("metadata") or {}) + + if isinstance(metadata, dict): + for param in _supported_callback_params: + if param not in standard_callback_dynamic_params and param in metadata: + _param_value = metadata.get(param) + if ( + _param_value is not None + and isinstance(_param_value, str) + and "os.environ/" in _param_value + ): + _param_value = get_secret_str(secret_name=_param_value) + standard_callback_dynamic_params[param] = _param_value # type: ignore + return standard_callback_dynamic_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1b3a687f1f..14015225f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3917,18 +3917,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 return langfuse_logger # type: ignore elif logging_integration == "langfuse_otel": from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - langfuse_otel_config = LangfuseOtelLogger.get_langfuse_otel_config() - - # The endpoint and headers are now set as environment variables by get_langfuse_otel_config() - otel_config = OpenTelemetryConfig( - exporter=langfuse_otel_config.protocol, - headers=langfuse_otel_config.otlp_auth_headers, - ) for callback in _in_memory_loggers: if ( @@ -3936,8 +3924,10 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 and callback.callback_name == "langfuse_otel" ): return callback # type: ignore + # Allow LangfuseOtelLogger to initialize its own config safely + # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage) _otel_logger = LangfuseOtelLogger( - config=otel_config, callback_name="langfuse_otel" + config=None, callback_name="langfuse_otel" ) _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore diff --git a/tests/logging_callback_tests/test_dynamic_otel_keys.py b/tests/logging_callback_tests/test_dynamic_otel_keys.py new file mode 100644 index 0000000000..2a463fddc0 --- /dev/null +++ b/tests/logging_callback_tests/test_dynamic_otel_keys.py @@ -0,0 +1,52 @@ +import sys +import os + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, +) + + +def test_dynamic_key_extraction_from_metadata(): + """ + Test extraction of langfuse keys from metadata in kwargs. + This simulates a Proxy request where keys are passed in metadata. + """ + kwargs = { + "metadata": { + "langfuse_public_key": "pk-test", + "langfuse_secret_key": "sk-test", + "langfuse_host": "https://test.langfuse.com", + } + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("langfuse_public_key") == "pk-test" + assert params.get("langfuse_secret_key") == "sk-test" + assert params.get("langfuse_host") == "https://test.langfuse.com" + + +def test_dynamic_key_extraction_from_litellm_params_metadata(): + """ + Test extraction of langfuse keys from litellm_params.metadata. + """ + kwargs = { + "litellm_params": { + "metadata": { + "langfuse_public_key": "pk-litellm", + "langfuse_secret_key": "sk-litellm", + } + } + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("langfuse_public_key") == "pk-litellm" + assert params.get("langfuse_secret_key") == "sk-litellm" + + +if __name__ == "__main__": + test_dynamic_key_extraction_from_metadata() + test_dynamic_key_extraction_from_litellm_params_metadata() diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index f8c662979a..ba4a096be2 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -1,92 +1,119 @@ import json import os -from datetime import datetime from unittest.mock import MagicMock, patch import pytest from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger -from litellm.types.integrations.langfuse_otel import LangfuseOtelConfig +from litellm.integrations.opentelemetry import OpenTelemetryConfig from litellm.types.llms.openai import ResponsesAPIResponse class TestLangfuseOtelIntegration: - def test_get_langfuse_otel_config_with_required_env_vars(self): """Test that config is created correctly with required environment variables.""" # Clean environment of any Langfuse-related variables - env_vars_to_clean = ['LANGFUSE_HOST', 'OTEL_EXPORTER_OTLP_ENDPOINT', 'OTEL_EXPORTER_OTLP_HEADERS'] - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key' - }, clear=False): + env_vars_to_clean = [ + "LANGFUSE_HOST", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_HEADERS", + ] + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + }, + clear=False, + ): # Remove any existing Langfuse variables for var in env_vars_to_clean: if var in os.environ: del os.environ[var] - + config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert isinstance(config, LangfuseOtelConfig) - assert config.protocol == "otlp_http" - assert "Authorization=Basic" in config.otlp_auth_headers - # Check that environment variables are set correctly (US default) - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://us.cloud.langfuse.com/api/public/otel" - assert "Authorization=Basic" in os.environ.get("OTEL_EXPORTER_OTLP_HEADERS", "") - + + assert isinstance(config, OpenTelemetryConfig) + assert config.exporter == "otlp_http" + assert "Authorization=Basic" in config.headers + # Note: We no longer set os.environ explicitly to avoid leakage + # assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://us.cloud.langfuse.com/api/public/otel" + # assert "Authorization=Basic" in os.environ.get("OTEL_EXPORTER_OTLP_HEADERS", "") + def test_get_langfuse_otel_config_missing_keys(self): """Test that ValueError is raised when required keys are missing.""" with patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match="LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set"): + with pytest.raises( + ValueError, + match="LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set", + ): LangfuseOtelLogger.get_langfuse_otel_config() - + def test_get_langfuse_otel_config_with_eu_host(self): """Test config with EU host.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://cloud.langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://cloud.langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://cloud.langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, OpenTelemetryConfig) + def test_get_langfuse_otel_config_with_custom_host(self): """Test config with custom host.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://my-langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://my-langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://my-langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, OpenTelemetryConfig) + def test_get_langfuse_otel_config_with_host_no_protocol(self): """Test config with custom host without protocol.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'my-langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "my-langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://my-langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, OpenTelemetryConfig) + def test_set_langfuse_otel_attributes(self): """Test that set_langfuse_otel_attributes calls the Arize utils function.""" from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) - + mock_span = MagicMock() mock_kwargs = {"test": "kwargs"} mock_response = {"test": "response"} - - with patch('litellm.integrations.arize._utils.set_attributes') as mock_set_attributes: - LangfuseOtelLogger.set_langfuse_otel_attributes(mock_span, mock_kwargs, mock_response) - - mock_set_attributes.assert_called_once_with(mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes) + + with patch( + "litellm.integrations.arize._utils.set_attributes" + ) as mock_set_attributes: + LangfuseOtelLogger.set_langfuse_otel_attributes( + mock_span, mock_kwargs, mock_response + ) + + mock_set_attributes.assert_called_once_with( + mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes + ) def test_set_langfuse_environment_attribute(self): """Test that Langfuse environment is set correctly when environment variable is present.""" @@ -94,15 +121,17 @@ class TestLangfuseOtelIntegration: mock_kwargs = {"test": "kwargs"} test_env = "staging" - with patch.dict(os.environ, {'LANGFUSE_TRACING_ENVIRONMENT': test_env}): - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, mock_kwargs, {}) - + with patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": test_env}): + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, mock_kwargs, {} + ) + # safe_set_attribute(span, key, value) → positional args mock_safe_set_attribute.assert_called_once_with( - mock_span, - "langfuse.environment", - test_env + mock_span, "langfuse.environment", test_env ) def test_extract_langfuse_metadata_basic(self): @@ -119,11 +148,13 @@ class TestLangfuseOtelIntegration: # Build a stub module + class on-the-fly stub_module = types.ModuleType("litellm.integrations.langfuse.langfuse") + class StubLFLogger: @staticmethod def add_metadata_from_header(litellm_params, metadata): # Echo back existing metadata plus a marker return {**metadata, "enriched": True} + stub_module.LangFuseLogger = StubLFLogger # type: ignore # Register stub in sys.modules so import inside method succeeds @@ -159,11 +190,16 @@ class TestLangfuseOtelIntegration: kwargs = {"litellm_params": {"metadata": metadata}} # Capture calls to safe_set_attribute - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, None) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), kwargs, None + ) # Build expected calls manually for clarity from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes + expected = { LangfuseSpanAttributes.GENERATION_NAME.value: "gen-name", LangfuseSpanAttributes.GENERATION_ID.value: "gen-id", @@ -176,12 +212,14 @@ class TestLangfuseOtelIntegration: # Lists / dicts should be JSON strings LangfuseSpanAttributes.TAGS.value: json.dumps(["tagA", "tagB"]), LangfuseSpanAttributes.TRACE_NAME.value: "trace-name", - LangfuseSpanAttributes.TRACE_ID.value: "trace-id", + LangfuseSpanAttributes.TRACE_ID.value: "traceid", # stripped dashes LangfuseSpanAttributes.TRACE_METADATA.value: json.dumps({"k": "v"}), LangfuseSpanAttributes.TRACE_VERSION.value: "t-ver", LangfuseSpanAttributes.TRACE_RELEASE.value: "rel-1", LangfuseSpanAttributes.EXISTING_TRACE_ID.value: "existing-id", - LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps(["key1", "key2"]), + LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps( + ["key1", "key2"] + ), LangfuseSpanAttributes.DEBUG_LANGFUSE.value: True, } @@ -191,7 +229,9 @@ class TestLangfuseOtelIntegration: for call in mock_safe_set_attribute.call_args_list } - assert actual == expected, "Mismatch between expected and actual OTEL attribute mapping." + assert ( + actual == expected + ), "Mismatch between expected and actual OTEL attribute mapping." def test_set_langfuse_specific_attributes_with_content(self): """Test that _set_langfuse_specific_attributes correctly sets observation.output with regular content response.""" @@ -200,15 +240,15 @@ class TestLangfuseOtelIntegration: # Create response with content response_obj = ModelResponse( - id='chatcmpl-test', - model='gpt-4o', + id="chatcmpl-test", + model="gpt-4o", choices=[ Choices( - finish_reason='stop', + finish_reason="stop", message={ "role": "assistant", - "content": "The weather in Tokyo is sunny." - } + "content": "The weather in Tokyo is sunny.", + }, ) ], ) @@ -217,20 +257,21 @@ class TestLangfuseOtelIntegration: "messages": [{"role": "user", "content": "What's the weather in Tokyo?"}], } - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), kwargs, response_obj + ) expect_output = { LangfuseSpanAttributes.OBSERVATION_INPUT.value: [ - { - "role": "user", - "content": "What's the weather in Tokyo?" - } + {"role": "user", "content": "What's the weather in Tokyo?"} ], LangfuseSpanAttributes.OBSERVATION_OUTPUT.value: { "role": "assistant", - "content": "The weather in Tokyo is sunny." - } + "content": "The weather in Tokyo is sunny.", + }, } # Flatten the actual calls into {key: value} @@ -239,8 +280,9 @@ class TestLangfuseOtelIntegration: for call in mock_safe_set_attribute.call_args_list } - assert actual == expect_output, "Mismatch in observation input/output OTEL attributes." - + assert ( + actual == expect_output + ), "Mismatch in observation input/output OTEL attributes." def test_set_langfuse_specific_attributes_with_tool_calls(self): """Test that _set_langfuse_specific_attributes correctly sets observation.output with tool calls in Langfuse format.""" @@ -254,42 +296,44 @@ class TestLangfuseOtelIntegration: # Create response with tool calls response_obj = ModelResponse( - id='chatcmpl-test', - model='gpt-4o', + id="chatcmpl-test", + model="gpt-4o", choices=[ Choices( - finish_reason='tool_calls', + finish_reason="tool_calls", message={ "role": "assistant", "content": None, "tool_calls": [ ChatCompletionMessageToolCall( function=Function( - arguments='{"location":"Tokyo"}', - name='get_weather' + arguments='{"location":"Tokyo"}', name="get_weather" ), - id='call_123', - type='function' + id="call_123", + type="function", ) - ] - } + ], + }, ) ], ) - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), {}, - response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), {}, response_obj + ) expected = { LangfuseSpanAttributes.OBSERVATION_OUTPUT.value: [ - { - "id": "chatcmpl-test", - "name": "get_weather", - "arguments": {"location": "Tokyo"}, - "call_id": "call_123", - "type": "function_call" - } + { + "id": "chatcmpl-test", + "name": "get_weather", + "arguments": {"location": "Tokyo"}, + "call_id": "call_123", + "type": "function_call", + } ] } @@ -298,8 +342,9 @@ class TestLangfuseOtelIntegration: call.args[1]: json.loads(call.args[2]) for call in mock_safe_set_attribute.call_args_list } - assert actual == expected, "Mismatch in observation output OTEL attribute for tool calls." - + assert ( + actual == expected + ), "Mismatch in observation output OTEL attribute for tool calls." def test_construct_dynamic_otel_headers_with_langfuse_keys(self): """Test that construct_dynamic_otel_headers creates proper auth headers when langfuse keys are provided.""" @@ -307,28 +352,27 @@ class TestLangfuseOtelIntegration: # Create dynamic params with langfuse keys dynamic_params = StandardCallbackDynamicParams( - langfuse_public_key="test_public_key", - langfuse_secret_key="test_secret_key" + langfuse_public_key="test_public_key", langfuse_secret_key="test_secret_key" ) - + logger = LangfuseOtelLogger() result = logger.construct_dynamic_otel_headers(dynamic_params) - + # Should return a dict with otlp_auth_headers assert result is not None assert "Authorization" in result - + # The auth header should contain the basic auth format auth_header = result["Authorization"] assert auth_header.startswith("Basic ") - + # Verify the header format by decoding import base64 # Extract the base64 part from "Authorization=Basic " base64_part = auth_header.replace("Basic ", "") decoded = base64.b64decode(base64_part).decode() - + assert decoded == "test_public_key:test_secret_key" def test_construct_dynamic_otel_headers_empty_params(self): @@ -337,24 +381,28 @@ class TestLangfuseOtelIntegration: # Create dynamic params without langfuse keys dynamic_params = StandardCallbackDynamicParams() - + logger = LangfuseOtelLogger() result = logger.construct_dynamic_otel_headers(dynamic_params) - + # Should return an empty dict assert result == {} - + def test_get_langfuse_otel_config_with_otel_host_priority(self): """LANGFUSE_OTEL_HOST should take priority over LANGFUSE_HOST.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://should-not-be-used.com', - 'LANGFUSE_OTEL_HOST': 'https://otel-host.com' - }, clear=False): - _ = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://otel-host.com/api/public/otel" + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://should-not-be-used.com", + "LANGFUSE_OTEL_HOST": "https://otel-host.com", + }, + clear=False, + ): + config = LangfuseOtelLogger.get_langfuse_otel_config() + assert isinstance(config, OpenTelemetryConfig) + # Endpoint assertion removed as side effect is gone class TestLangfuseOtelResponsesAPI: @@ -369,46 +417,52 @@ class TestLangfuseOtelResponsesAPI: output=[ { "type": "message", - "content": [{"type": "text", "text": "Hello from responses API"}] + "content": [{"type": "text", "text": "Hello from responses API"}], } ], parallel_tool_calls=False, tool_choice="auto", tools=[], - top_p=1.0 + top_p=1.0, ) - + # Create kwargs with metadata that should be logged test_metadata = { - "user_id": "test123", - "session_id": "abc456", + "user_id": "test123", + "session_id": "abc456", "custom_field": "test_value", "generation_name": "responses_test_generation", - "trace_name": "responses_api_trace" + "trace_name": "responses_api_trace", } - + kwargs = { "call_type": "responses", "messages": [{"role": "user", "content": "Hello"}], "model": "gpt-4o", "optional_params": {}, - "litellm_params": {"metadata": test_metadata} + "litellm_params": {"metadata": test_metadata}, } - + mock_span = MagicMock() - + from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) - - with patch('litellm.integrations.arize._utils.set_attributes') as mock_set_attributes: - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: + + with patch( + "litellm.integrations.arize._utils.set_attributes" + ) as mock_set_attributes: + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: logger = LangfuseOtelLogger() logger.set_langfuse_otel_attributes(mock_span, kwargs, mock_response) - + # Verify that set_attributes was called for general attributes - mock_set_attributes.assert_called_once_with(mock_span, kwargs, mock_response, LangfuseLLMObsOTELAttributes) - + mock_set_attributes.assert_called_once_with( + mock_span, kwargs, mock_response, LangfuseLLMObsOTELAttributes + ) + # Verify that Langfuse-specific attributes were set mock_safe_set_attribute.assert_any_call( mock_span, "langfuse.generation.name", "responses_test_generation" @@ -421,29 +475,30 @@ class TestLangfuseOtelResponsesAPI: """Test that metadata is correctly extracted from ResponsesAPI kwargs.""" # Clean up any existing module mocks import sys + if "litellm.integrations.langfuse.langfuse" in sys.modules: - original_module = sys.modules["litellm.integrations.langfuse.langfuse"] - + sys.modules["litellm.integrations.langfuse.langfuse"] + test_metadata = { "user_id": "responses_user_123", - "session_id": "responses_session_456", + "session_id": "responses_session_456", "custom_metadata": {"key": "value"}, "generation_name": "responses_generation", - "trace_id": "custom_trace_id" + "trace_id": "custom_trace_id", } - + kwargs = { "call_type": "responses", "model": "gpt-4o", - "litellm_params": {"metadata": test_metadata} + "litellm_params": {"metadata": test_metadata}, } - + extracted_metadata = LangfuseOtelLogger._extract_langfuse_metadata(kwargs) - + # Verify all expected metadata was extracted (may have additional fields from header enrichment) for key, value in test_metadata.items(): assert extracted_metadata[key] == value - + assert extracted_metadata["user_id"] == "responses_user_123" assert extracted_metadata["generation_name"] == "responses_generation" assert extracted_metadata["trace_id"] == "custom_trace_id" @@ -457,39 +512,61 @@ class TestLangfuseOtelResponsesAPI: "trace_user_id": "resp_user_456", "session_id": "resp_session_789", "tags": ["responses", "api", "test"], - "trace_metadata": {"source": "responses_api", "version": "1.0"} + "trace_metadata": {"source": "responses_api", "version": "1.0"}, } - - kwargs = { - "call_type": "responses", - "litellm_params": {"metadata": metadata} - } - + + kwargs = {"call_type": "responses", "litellm_params": {"metadata": metadata}} + mock_span = MagicMock() - - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: + + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, {}) - + # Verify specific attributes were set from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes - + expected_calls = [ - (mock_span, LangfuseSpanAttributes.GENERATION_NAME.value, "responses_gen"), + ( + mock_span, + LangfuseSpanAttributes.GENERATION_NAME.value, + "responses_gen", + ), (mock_span, LangfuseSpanAttributes.GENERATION_ID.value, "resp_gen_123"), (mock_span, LangfuseSpanAttributes.TRACE_NAME.value, "responses_trace"), - (mock_span, LangfuseSpanAttributes.TRACE_USER_ID.value, "resp_user_456"), - (mock_span, LangfuseSpanAttributes.SESSION_ID.value, "resp_session_789"), - (mock_span, LangfuseSpanAttributes.TAGS.value, json.dumps(["responses", "api", "test"])), - (mock_span, LangfuseSpanAttributes.TRACE_METADATA.value, - json.dumps({"source": "responses_api", "version": "1.0"})) + ( + mock_span, + LangfuseSpanAttributes.TRACE_USER_ID.value, + "resp_user_456", + ), + ( + mock_span, + LangfuseSpanAttributes.SESSION_ID.value, + "resp_session_789", + ), + ( + mock_span, + LangfuseSpanAttributes.TAGS.value, + json.dumps(["responses", "api", "test"]), + ), + ( + mock_span, + LangfuseSpanAttributes.TRACE_METADATA.value, + json.dumps({"source": "responses_api", "version": "1.0"}), + ), ] - + for expected_call in expected_calls: mock_safe_set_attribute.assert_any_call(*expected_call) def test_responses_api_with_output(self): """Test Langfuse OTEL logger with Responses API output (reasoning + message).""" - from openai.types.responses import ResponseReasoningItem, ResponseOutputMessage, ResponseOutputText + from openai.types.responses import ( + ResponseReasoningItem, + ResponseOutputMessage, + ResponseOutputText, + ) from openai.types.responses.response_reasoning_item import Summary from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes @@ -504,9 +581,9 @@ class TestLangfuseOtelResponsesAPI: summary=[ Summary( text="Let me analyze this problem step by step...", - type="summary_text" + type="summary_text", ) - ] + ], ), ResponseOutputMessage( id="msg-001", @@ -519,26 +596,33 @@ class TestLangfuseOtelResponsesAPI: text="The weather in San Francisco is sunny, 20°C.", type="output_text", ) - ] - ) - ] + ], + ), + ], ) kwargs = { "call_type": "responses", - "messages": [{"role": "user", "content": "What's the weather in San Francisco?"}], + "messages": [ + {"role": "user", "content": "What's the weather in San Francisco?"} + ], "model": "gpt-4o", "optional_params": {}, } mock_span = MagicMock() - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, kwargs, response_obj + ) # Verify observation output was set output_calls = [ - call for call in mock_safe_set_attribute.call_args_list + call + for call in mock_safe_set_attribute.call_args_list if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value ] @@ -552,11 +636,17 @@ class TestLangfuseOtelResponsesAPI: # Verify reasoning summary assert output_data[0]["role"] == "reasoning_summary" - assert output_data[0]["content"] == "Let me analyze this problem step by step..." + assert ( + output_data[0]["content"] + == "Let me analyze this problem step by step..." + ) # Verify message assert output_data[1]["role"] == "assistant" - assert output_data[1]["content"] == "The weather in San Francisco is sunny, 20°C." + assert ( + output_data[1]["content"] + == "The weather in San Francisco is sunny, 20°C." + ) def test_responses_api_with_function_calls(self): """Test Langfuse OTEL logger with Responses API function_call output.""" @@ -574,26 +664,33 @@ class TestLangfuseOtelResponsesAPI: name="get_weather", call_id="call-abc", arguments='{"location": "San Francisco", "unit": "celsius"}', - status="completed" + status="completed", ) - ] + ], ) kwargs = { "call_type": "responses", - "messages": [{"role": "user", "content": "What's the weather in San Francisco?"}], + "messages": [ + {"role": "user", "content": "What's the weather in San Francisco?"} + ], "model": "gpt-4o", "optional_params": {}, } mock_span = MagicMock() - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, kwargs, response_obj + ) # Verify observation output was set output_calls = [ - call for call in mock_safe_set_attribute.call_args_list + call + for call in mock_safe_set_attribute.call_args_list if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value ] @@ -615,4 +712,4 @@ class TestLangfuseOtelResponsesAPI: if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__]) From 774a015aad0fb87ca50ed96f8ab46fe7abced885 Mon Sep 17 00:00:00 2001 From: akraines Date: Wed, 4 Feb 2026 08:44:30 +0200 Subject: [PATCH 4/5] fix: Preserve streaming content on guardrail-sampled chunks (#20027) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unified guardrail's streaming iterator hook processes every Nth chunk (sampling_rate, default 5). On each sampled chunk it calls process_output_streaming_response, which combines all accumulated text into the first chunk and clears all subsequent chunks to "". The hook then yielded `processed_items[-1]` — the last item, whose content had been cleared to "". This permanently lost every Nth chunk's content, causing random missing words/tokens in the client output (observed in Roo Code, Open WebUI, etc.). Fix: deep-copy the current chunk before guardrail processing runs, then yield the original (unmodified) chunk. The guardrail validation still executes and can block if it detects a problem, but the stream content is preserved. Co-authored-by: Claude Opus 4.5 --- .../unified_guardrail/unified_guardrail.py | 25 +++-- .../test_unified_guardrail.py | 100 +++++++++++++++++- 2 files changed, 114 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 80f9860bdf..f07f65d10f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -6,6 +6,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint 3. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook """ +import copy from typing import Any, AsyncGenerator, List, Optional, Union from litellm._logging import verbose_proxy_logger @@ -349,22 +350,26 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply.guardrail_name, ) + # Deep-copy the current chunk before guardrail processing. + # process_output_streaming_response modifies responses_so_far + # in-place: it puts the combined guardrailed text in the first + # chunk and clears all subsequent chunks to "". Without this + # copy, yielding processed_items[-1] would yield an empty + # string, permanently losing this chunk's content. + original_item = copy.deepcopy(item) + endpoint_translation = endpoint_guardrail_translation_mappings[ CallTypes(call_type) ]() - processed_items = ( - await endpoint_translation.process_output_streaming_response( - responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=request_data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - ) + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, ) - last_item = processed_items[-1] - - yield last_item + yield original_item else: yield item diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 2a33c56b56..b41cded1d0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -8,12 +8,13 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTra from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices class RecordingGuardrail(CustomGuardrail): @@ -131,3 +132,100 @@ class TestUnifiedLLMGuardrails: ) assert guardrail.event_history == [GuardrailEventHooks.during_call] + + class TestAsyncPostCallStreamingIteratorHook: + @pytest.mark.asyncio + async def test_streaming_content_not_lost_on_sampled_chunks(self): + """ + Verify that every chunk's content is preserved in the output stream. + + The bug: process_output_streaming_response puts the combined + guardrailed text in the first chunk and clears all subsequent + chunks to "". The hook then yielded processed_items[-1] (the + cleared last item), permanently losing every Nth chunk's content. + """ + + class _ContentClearingTranslation(BaseTranslation): + """Simulates the real OpenAI handler behavior that triggers the bug.""" + + async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None): # type: ignore[override] + return data + + async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None): # type: ignore[override] + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + ): + # Simulate what the real handler does: + # put combined text in first chunk, clear the rest + combined = "" + for resp in responses_so_far: + for choice in resp.choices: + if choice.delta and choice.delta.content: + combined += choice.delta.content + + first_set = False + for resp in responses_so_far: + for choice in resp.choices: + if not first_set: + choice.delta.content = combined + first_set = True + else: + choice.delta.content = "" + + return responses_so_far + + # Override the mapping to use our content-clearing translation + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.acompletion: _ContentClearingTranslation, + } + + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + # Create 10 streaming chunks with distinct content + chunks = [] + for i in range(10): + chunk = ModelResponseStream( + choices=[StreamingChoices( + delta=Delta(content=f"word{i} ", role="assistant"), + finish_reason=None, + )], + ) + chunks.append(chunk) + + async def mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + request_route="/v1/chat/completions", + ) + + request_data = { + "guardrail_to_apply": guardrail, + "model": "gpt-4", + } + + # Collect all yielded chunks + yielded_contents = [] + async for item in handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + content = item.choices[0].delta.content if item.choices[0].delta else None + yielded_contents.append(content) + + # Every chunk should have non-empty content + for i, content in enumerate(yielded_contents): + assert content is not None and content != "", ( + f"Chunk {i} lost its content (got {content!r}). " + f"Expected non-empty content for every streamed chunk." + ) From bfabb39fc6b847a50877f7362089c5b6161b78b2 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Sat, 31 Jan 2026 12:01:31 +0800 Subject: [PATCH 5/5] fix: Fix Vertex AI Gemini streaming content_filter handling - Add promptFeedback.blockReason check in chunk_parser - Return content_filter finish_reason when blocked - Extract content filter logic into _check_prompt_level_content_filter() method - Update unit tests to reflect simplified implementation Signed-off-by: Kris Xia --- .../vertex_and_google_ai_studio_gemini.py | 55 ++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 120 ++++++++++++++++++ 2 files changed, 175 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b5a6949f27..04ae4b6beb 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1732,6 +1732,52 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: return "stop" + @staticmethod + def _check_prompt_level_content_filter( + processed_chunk: GenerateContentResponseBody, + response_id: Optional[str], + ) -> Optional["ModelResponseStream"]: + """ + Check if prompt is blocked due to content filtering at the prompt level. + + This handles the case where Vertex AI blocks the prompt before generation begins, + indicated by promptFeedback.blockReason being present. + + Args: + processed_chunk: The parsed response chunk from Vertex AI + response_id: The response ID from the chunk + + Returns: + ModelResponseStream with content_filter finish_reason if blocked, None otherwise. + + Note: + This is consistent with non-streaming _handle_blocked_response() behavior. + Candidate-level content filtering (SAFETY, RECITATION, etc.) is handled + separately via _process_candidates() → _check_finish_reason(). + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + # Check if prompt is blocked due to content filtering + prompt_feedback = processed_chunk.get("promptFeedback") + if prompt_feedback and "blockReason" in prompt_feedback: + verbose_logger.debug( + f"Prompt blocked due to: {prompt_feedback.get('blockReason')} - {prompt_feedback.get('blockReasonMessage')}" + ) + + # Create a content_filter response (consistent with non-streaming _handle_blocked_response) + choice = StreamingChoices( + finish_reason="content_filter", + index=0, + delta=Delta(content=None, role="assistant"), + logprobs=None, + enhancements=None, + ) + + model_response = ModelResponseStream(choices=[choice], id=response_id) + return model_response + + return None + @staticmethod def _calculate_web_search_requests(grounding_metadata: List[dict]) -> Optional[int]: web_search_requests: Optional[int] = None @@ -2813,6 +2859,15 @@ class ModelResponseIterator: processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore response_id = processed_chunk.get("responseId") model_response = ModelResponseStream(choices=[], id=response_id) + + # Check if prompt is blocked due to content filtering + blocked_response = VertexGeminiConfig._check_prompt_level_content_filter( + processed_chunk=processed_chunk, + response_id=response_id, + ) + if blocked_response is not None: + model_response = blocked_response + usage: Optional[Usage] = None _candidates: Optional[List[Candidates]] = processed_chunk.get("candidates") grounding_metadata: List[dict] = [] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index cb3b51acd6..75fd597ffa 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3218,3 +3218,123 @@ def test_video_metadata_only_for_gemini_3(): assert file_part_3 is not None assert "media_resolution" in file_part_3, "Gemini 3 should have media_resolution" assert "video_metadata" in file_part_3, "Gemini 3 should have video_metadata" + + + +def test_chunk_parser_handles_prompt_feedback_block(): + """Test chunk_parser correctly handles promptFeedback.blockReason""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - mock a blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "PROHIBITED_CONTENT", + "blockReasonMessage": "The prompt is blocked due to prohibited contents" + }, + "responseId": "test_response_id", + "modelVersion": "gemini-3-pro-preview" + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert + assert result is not None, "Result should not be None" + assert len(result.choices) == 1, "Should have exactly one choice" + assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].delta.content is None, "content should be None" + + +def test_chunk_parser_handles_prompt_feedback_safety_block(): + """Test chunk_parser handles different blockReason types (SAFETY)""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - mock a SAFETY blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "SAFETY", + "blockReasonMessage": "The prompt is blocked due to safety concerns" + }, + "responseId": "test_safety_response_id", + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert + assert result is not None + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "content_filter" + + +def test_chunk_parser_handles_prompt_feedback_block_with_usage(): + """Test chunk_parser correctly extracts usageMetadata when promptFeedback.blockReason is present""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - 模拟一个包含 usageMetadata 的 blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "PROHIBITED_CONTENT", + "blockReasonMessage": "The prompt is blocked due to prohibited contents" + }, + "responseId": "test_response_id_with_usage", + "modelVersion": "gemini-3-pro-preview", + "usageMetadata": { + "promptTokenCount": 8175, + "candidatesTokenCount": 0, + "totalTokenCount": 8175 + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert - 验证 content_filter 响应和 usage 都被正确处理 + assert result is not None, "Result should not be None" + assert len(result.choices) == 1, "Should have exactly one choice" + assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].delta.content is None, "content should be None" + + # 验证 usage 信息被正确提取 + assert hasattr(result, "usage"), "result should have usage attribute" + assert result.usage is not None, "usage should not be None" + assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" + assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}" + assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}" +