mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #20398 from BerriAI/litellm_oss_staging_02_04_2026
Litellm oss staging 02 04 2026
This commit is contained in:
Sameer Kankute
GitHub
commit
b74afec2be
21 changed files with 1069 additions and 445 deletions
65
Makefile
65
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/
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -599,9 +599,9 @@ class OpenTelemetry(CustomLogger):
|
|||
|
||||
def _get_dynamic_otel_headers_from_kwargs(self, kwargs) -> Optional[dict]:
|
||||
"""Extract dynamic headers from kwargs if available."""
|
||||
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
|
||||
kwargs.get("standard_callback_dynamic_params")
|
||||
)
|
||||
standard_callback_dynamic_params: Optional[
|
||||
StandardCallbackDynamicParams
|
||||
] = kwargs.get("standard_callback_dynamic_params")
|
||||
|
||||
if not standard_callback_dynamic_params:
|
||||
return None
|
||||
|
|
@ -619,7 +619,9 @@ class OpenTelemetry(CustomLogger):
|
|||
# Prevents thread exhaustion by reusing providers for the same credential sets (e.g. per-team keys)
|
||||
cache_key = str(sorted(dynamic_headers.items()))
|
||||
if cache_key in self._tracer_provider_cache:
|
||||
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
|
||||
return self._tracer_provider_cache[cache_key].get_tracer(
|
||||
LITELLM_TRACER_NAME
|
||||
)
|
||||
|
||||
# Create a temporary tracer provider with dynamic headers
|
||||
temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config))
|
||||
|
|
@ -674,7 +676,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,14 +1008,11 @@ 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
|
||||
)
|
||||
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)
|
||||
|
||||
|
|
@ -1618,7 +1620,6 @@ class OpenTelemetry(CustomLogger):
|
|||
|
||||
for idx, choice in enumerate(response_obj.get("choices")):
|
||||
if choice.get("finish_reason"):
|
||||
|
||||
message = choice.get("message")
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
|
|
@ -1631,7 +1632,9 @@ class OpenTelemetry(CustomLogger):
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
self.handle_callback_failure(callback_name=self.callback_name or "opentelemetry")
|
||||
self.handle_callback_failure(
|
||||
callback_name=self.callback_name or "opentelemetry"
|
||||
)
|
||||
verbose_logger.exception(
|
||||
"OpenTelemetry logging error in set_attributes %s", str(e)
|
||||
)
|
||||
|
|
@ -1722,6 +1725,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")
|
||||
|
|
|
|||
|
|
@ -1,8 +1,35 @@
|
|||
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",
|
||||
"turn_off_message_logging",
|
||||
]
|
||||
|
||||
|
||||
def initialize_standard_callback_dynamic_params(
|
||||
kwargs: Optional[Dict] = None,
|
||||
|
|
@ -15,13 +42,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 +54,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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -130,6 +130,11 @@ def perform_redaction(model_call_details: dict, result):
|
|||
def should_redact_message_logging(model_call_details: dict) -> bool:
|
||||
"""
|
||||
Determine if message logging should be redacted.
|
||||
|
||||
Priority order:
|
||||
1. Dynamic parameter (turn_off_message_logging in request)
|
||||
2. Headers (litellm-disable-message-redaction / litellm-enable-message-redaction)
|
||||
3. Global setting (litellm.turn_off_message_logging)
|
||||
"""
|
||||
litellm_params = model_call_details.get("litellm_params", {})
|
||||
|
||||
|
|
@ -139,36 +144,36 @@ def should_redact_message_logging(model_call_details: dict) -> bool:
|
|||
# Get headers from the metadata
|
||||
request_headers = metadata.get("headers", {}) if isinstance(metadata, dict) else {}
|
||||
|
||||
possible_request_headers = [
|
||||
# Check for headers that explicitly control redaction
|
||||
if request_headers and bool(
|
||||
request_headers.get("litellm-disable-message-redaction", False)
|
||||
):
|
||||
# User explicitly disabled redaction via header
|
||||
return False
|
||||
|
||||
possible_enable_headers = [
|
||||
"litellm-enable-message-redaction", # old header. maintain backwards compatibility
|
||||
"x-litellm-enable-message-redaction", # new header
|
||||
]
|
||||
|
||||
is_redaction_enabled_via_header = False
|
||||
for header in possible_request_headers:
|
||||
for header in possible_enable_headers:
|
||||
if bool(request_headers.get(header, False)):
|
||||
is_redaction_enabled_via_header = True
|
||||
break
|
||||
|
||||
# check if user opted out of logging message/response to callbacks
|
||||
if (
|
||||
litellm.turn_off_message_logging is not True
|
||||
and is_redaction_enabled_via_header is not True
|
||||
and _get_turn_off_message_logging_from_dynamic_params(model_call_details)
|
||||
is not True
|
||||
):
|
||||
return False
|
||||
|
||||
if request_headers and bool(
|
||||
request_headers.get("litellm-disable-message-redaction", False)
|
||||
):
|
||||
return False
|
||||
|
||||
# user has OPTED OUT of message redaction
|
||||
if _get_turn_off_message_logging_from_dynamic_params(model_call_details) is False:
|
||||
return False
|
||||
|
||||
return True
|
||||
# Priority 1: Check dynamic parameter first (if explicitly set)
|
||||
dynamic_turn_off = _get_turn_off_message_logging_from_dynamic_params(model_call_details)
|
||||
if dynamic_turn_off is not None:
|
||||
# Dynamic parameter is explicitly set, use it
|
||||
return dynamic_turn_off
|
||||
|
||||
# Priority 2: Check if header explicitly enables redaction
|
||||
if is_redaction_enabled_via_header:
|
||||
return True
|
||||
|
||||
# Priority 3: Fall back to global setting
|
||||
return litellm.turn_off_message_logging is True
|
||||
|
||||
|
||||
def redact_message_input_output_from_logging(
|
||||
|
|
|
|||
|
|
@ -30,6 +30,8 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.types.llms.bedrock import *
|
||||
|
||||
from ..common_utils import is_claude_4_5_on_bedrock
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantMessage,
|
||||
|
|
@ -306,9 +308,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
return "nova-2-lite" in model_without_region
|
||||
|
||||
def _map_web_search_options(
|
||||
self,
|
||||
web_search_options: dict,
|
||||
model: str
|
||||
self, web_search_options: dict, model: str
|
||||
) -> Optional[BedrockToolBlock]:
|
||||
"""
|
||||
Map web_search_options to Nova grounding systemTool.
|
||||
|
|
@ -634,7 +634,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
Filtered list of beta headers
|
||||
"""
|
||||
filtered_betas = []
|
||||
|
||||
|
||||
# 1. Filter out beta headers that are universally unsupported on Bedrock Converse
|
||||
for beta in beta_list:
|
||||
should_keep = True
|
||||
|
|
@ -642,10 +642,10 @@ class AmazonConverseConfig(BaseConfig):
|
|||
if unsupported_pattern in beta.lower():
|
||||
should_keep = False
|
||||
break
|
||||
|
||||
|
||||
if should_keep:
|
||||
filtered_betas.append(beta)
|
||||
|
||||
|
||||
return filtered_betas
|
||||
|
||||
def _separate_computer_use_tools(
|
||||
|
|
@ -808,11 +808,11 @@ class AmazonConverseConfig(BaseConfig):
|
|||
if param == "web_search_options" and isinstance(value, dict):
|
||||
# Note: we use `isinstance(value, dict)` instead of `value and isinstance(value, dict)`
|
||||
# because empty dict {} is falsy but is a valid way to enable Nova grounding
|
||||
grounding_tool = self._map_web_search_options(value, model)
|
||||
if grounding_tool is not None:
|
||||
optional_params = self._add_tools_to_optional_params(
|
||||
optional_params=optional_params, tools=[grounding_tool]
|
||||
)
|
||||
grounding_tool = self._map_web_search_options(value, model)
|
||||
if grounding_tool is not None:
|
||||
optional_params = self._add_tools_to_optional_params(
|
||||
optional_params=optional_params, tools=[grounding_tool]
|
||||
)
|
||||
|
||||
# Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models
|
||||
# Nova Lite 2 handles token budgeting differently through reasoningConfig
|
||||
|
|
@ -926,6 +926,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
ChatCompletionAssistantMessage,
|
||||
],
|
||||
block_type: Literal["system"],
|
||||
model: Optional[str] = None,
|
||||
) -> Optional[SystemContentBlock]:
|
||||
pass
|
||||
|
||||
|
|
@ -939,6 +940,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
ChatCompletionAssistantMessage,
|
||||
],
|
||||
block_type: Literal["content_block"],
|
||||
model: Optional[str] = None,
|
||||
) -> Optional[ContentBlock]:
|
||||
pass
|
||||
|
||||
|
|
@ -951,16 +953,26 @@ class AmazonConverseConfig(BaseConfig):
|
|||
ChatCompletionAssistantMessage,
|
||||
],
|
||||
block_type: Literal["system", "content_block"],
|
||||
model: Optional[str] = None,
|
||||
) -> Optional[Union[SystemContentBlock, ContentBlock]]:
|
||||
if message_block.get("cache_control", None) is None:
|
||||
cache_control = message_block.get("cache_control", None)
|
||||
if cache_control is None:
|
||||
return None
|
||||
|
||||
cache_point = CachePointBlock(type="default")
|
||||
if isinstance(cache_control, dict) and "ttl" in cache_control:
|
||||
ttl = cache_control["ttl"]
|
||||
if ttl in ["5m", "1h"] and model is not None:
|
||||
if is_claude_4_5_on_bedrock(model):
|
||||
cache_point["ttl"] = ttl
|
||||
|
||||
if block_type == "system":
|
||||
return SystemContentBlock(cachePoint=CachePointBlock(type="default"))
|
||||
return SystemContentBlock(cachePoint=cache_point)
|
||||
else:
|
||||
return ContentBlock(cachePoint=CachePointBlock(type="default"))
|
||||
return ContentBlock(cachePoint=cache_point)
|
||||
|
||||
def _transform_system_message(
|
||||
self, messages: List[AllMessageValues]
|
||||
self, messages: List[AllMessageValues], model: Optional[str] = None
|
||||
) -> Tuple[List[AllMessageValues], List[SystemContentBlock]]:
|
||||
system_prompt_indices = []
|
||||
system_content_blocks: List[SystemContentBlock] = []
|
||||
|
|
@ -972,7 +984,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
SystemContentBlock(text=message["content"])
|
||||
)
|
||||
cache_block = self._get_cache_point_block(
|
||||
message, block_type="system"
|
||||
message, block_type="system", model=model
|
||||
)
|
||||
if cache_block:
|
||||
system_content_blocks.append(cache_block)
|
||||
|
|
@ -983,7 +995,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
SystemContentBlock(text=m["text"])
|
||||
)
|
||||
cache_block = self._get_cache_point_block(
|
||||
m, block_type="system"
|
||||
m, block_type="system", model=model
|
||||
)
|
||||
if cache_block:
|
||||
system_content_blocks.append(cache_block)
|
||||
|
|
@ -1137,13 +1149,13 @@ class AmazonConverseConfig(BaseConfig):
|
|||
if beta not in seen:
|
||||
unique_betas.append(beta)
|
||||
seen.add(beta)
|
||||
|
||||
|
||||
# Filter out unsupported beta headers for Bedrock Converse API
|
||||
filtered_betas = self._filter_unsupported_beta_headers_for_bedrock(
|
||||
model=model,
|
||||
beta_list=unique_betas,
|
||||
)
|
||||
|
||||
|
||||
additional_request_params["anthropic_beta"] = filtered_betas
|
||||
|
||||
return bedrock_tools, anthropic_beta_list
|
||||
|
|
@ -1196,9 +1208,11 @@ class AmazonConverseConfig(BaseConfig):
|
|||
)
|
||||
|
||||
# Prepare and separate parameters
|
||||
inference_params, additional_request_params, request_metadata = self._prepare_request_params(
|
||||
optional_params, model
|
||||
)
|
||||
(
|
||||
inference_params,
|
||||
additional_request_params,
|
||||
request_metadata,
|
||||
) = self._prepare_request_params(optional_params, model)
|
||||
|
||||
original_tools = inference_params.pop("tools", [])
|
||||
|
||||
|
|
@ -1250,7 +1264,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: Optional[dict] = None,
|
||||
) -> RequestObject:
|
||||
messages, system_content_blocks = self._transform_system_message(messages)
|
||||
messages, system_content_blocks = self._transform_system_message(
|
||||
messages, model=model
|
||||
)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
|
|
@ -1306,7 +1322,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: Optional[dict] = None,
|
||||
) -> RequestObject:
|
||||
messages, system_content_blocks = self._transform_system_message(messages)
|
||||
messages, system_content_blocks = self._transform_system_message(
|
||||
messages, model=model
|
||||
)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
|
|
@ -1484,7 +1502,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
return message, returned_finish_reason
|
||||
|
||||
def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[
|
||||
def _translate_message_content(
|
||||
self, content_blocks: List[ContentBlock]
|
||||
) -> Tuple[
|
||||
str,
|
||||
List[ChatCompletionToolCallChunk],
|
||||
Optional[List[BedrockConverseReasoningContentBlock]],
|
||||
|
|
@ -1501,9 +1521,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
"""
|
||||
content_str = ""
|
||||
tools: List[ChatCompletionToolCallChunk] = []
|
||||
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
|
||||
None
|
||||
)
|
||||
reasoningContentBlocks: Optional[
|
||||
List[BedrockConverseReasoningContentBlock]
|
||||
] = None
|
||||
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
|
||||
for idx, content in enumerate(content_blocks):
|
||||
"""
|
||||
|
|
@ -1557,7 +1577,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
return content_str, tools, reasoningContentBlocks, citationsContentBlocks
|
||||
|
||||
def _transform_response( # noqa: PLR0915
|
||||
def _transform_response( # noqa: PLR0915
|
||||
self,
|
||||
model: str,
|
||||
response: httpx.Response,
|
||||
|
|
@ -1630,9 +1650,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}
|
||||
content_str = ""
|
||||
tools: List[ChatCompletionToolCallChunk] = []
|
||||
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
|
||||
None
|
||||
)
|
||||
reasoningContentBlocks: Optional[
|
||||
List[BedrockConverseReasoningContentBlock]
|
||||
] = None
|
||||
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
|
||||
|
||||
if message is not None:
|
||||
|
|
@ -1651,15 +1671,17 @@ class AmazonConverseConfig(BaseConfig):
|
|||
provider_specific_fields["citationsContent"] = citationsContentBlocks
|
||||
|
||||
if provider_specific_fields:
|
||||
chat_completion_message["provider_specific_fields"] = provider_specific_fields
|
||||
chat_completion_message[
|
||||
"provider_specific_fields"
|
||||
] = provider_specific_fields
|
||||
|
||||
if reasoningContentBlocks is not None:
|
||||
chat_completion_message["reasoning_content"] = (
|
||||
self._transform_reasoning_content(reasoningContentBlocks)
|
||||
)
|
||||
chat_completion_message["thinking_blocks"] = (
|
||||
self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
)
|
||||
chat_completion_message[
|
||||
"reasoning_content"
|
||||
] = self._transform_reasoning_content(reasoningContentBlocks)
|
||||
chat_completion_message[
|
||||
"thinking_blocks"
|
||||
] = self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
chat_completion_message["content"] = content_str
|
||||
if (
|
||||
json_mode is True
|
||||
|
|
|
|||
|
|
@ -446,6 +446,29 @@ def get_bedrock_base_model(model: str) -> str:
|
|||
return model
|
||||
|
||||
|
||||
def is_claude_4_5_on_bedrock(model: str) -> bool:
|
||||
"""
|
||||
Check if the model is a Claude 4.5 model on Bedrock.
|
||||
Claude 4.5 models support prompt caching with '5m' and '1h' TTL on Bedrock.
|
||||
"""
|
||||
model_lower = model.lower()
|
||||
claude_4_5_patterns = [
|
||||
"sonnet-4.5",
|
||||
"sonnet_4.5",
|
||||
"sonnet-4-5",
|
||||
"sonnet_4_5",
|
||||
"haiku-4.5",
|
||||
"haiku_4.5",
|
||||
"haiku-4-5",
|
||||
"haiku_4_5",
|
||||
"opus-4.5",
|
||||
"opus_4.5",
|
||||
"opus-4-5",
|
||||
"opus_4_5",
|
||||
]
|
||||
return any(pattern in model_lower for pattern in claude_4_5_patterns)
|
||||
|
||||
|
||||
# Import after standalone functions to avoid circular imports
|
||||
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
|
||||
|
||||
|
|
@ -815,21 +838,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]:
|
|||
# If it's already a list, return it
|
||||
if isinstance(anthropic_beta_header, list):
|
||||
return anthropic_beta_header
|
||||
|
||||
|
||||
# Try to parse as JSON array first (e.g., '["interleaved-thinking-2025-05-14", "claude-code-20250219"]')
|
||||
if isinstance(anthropic_beta_header, str):
|
||||
anthropic_beta_header = anthropic_beta_header.strip()
|
||||
if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith("]"):
|
||||
if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith(
|
||||
"]"
|
||||
):
|
||||
try:
|
||||
parsed = json.loads(anthropic_beta_header)
|
||||
if isinstance(parsed, list):
|
||||
return [str(beta).strip() for beta in parsed]
|
||||
except json.JSONDecodeError:
|
||||
pass # Fall through to comma-separated parsing
|
||||
|
||||
|
||||
# Fall back to comma-separated values
|
||||
return [beta.strip() for beta in anthropic_beta_header.split(",")]
|
||||
|
||||
|
||||
return []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,7 +23,10 @@ from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
|
|||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
AmazonInvokeConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import get_anthropic_beta_from_headers
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
get_anthropic_beta_from_headers,
|
||||
is_claude_4_5_on_bedrock,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -54,7 +57,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
# These will be filtered out to prevent 400 "invalid beta flag" errors
|
||||
UNSUPPORTED_BEDROCK_INVOKE_BETA_PATTERNS = [
|
||||
"advanced-tool-use", # Bedrock Invoke doesn't support advanced-tool-use beta headers
|
||||
"prompt-caching-scope"
|
||||
"prompt-caching-scope",
|
||||
]
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
|
|
@ -116,15 +119,22 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
)
|
||||
|
||||
def _remove_ttl_from_cache_control(
|
||||
self, anthropic_messages_request: Dict
|
||||
self, anthropic_messages_request: Dict, model: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
Remove `ttl` field from cache_control in messages.
|
||||
Bedrock doesn't support the ttl field in cache_control.
|
||||
|
||||
Update: Bedock supports `5m` and `1h` for Claude 4.5 models.
|
||||
|
||||
Args:
|
||||
anthropic_messages_request: The request dictionary to modify in-place
|
||||
model: The model name to check if it supports ttl
|
||||
"""
|
||||
is_claude_4_5 = False
|
||||
if model:
|
||||
is_claude_4_5 = self._is_claude_4_5_on_bedrock(model)
|
||||
|
||||
if "messages" in anthropic_messages_request:
|
||||
for message in anthropic_messages_request["messages"]:
|
||||
if isinstance(message, dict) and "content" in message:
|
||||
|
|
@ -133,7 +143,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
for item in content:
|
||||
if isinstance(item, dict) and "cache_control" in item:
|
||||
cache_control = item["cache_control"]
|
||||
if isinstance(cache_control, dict) and "ttl" in cache_control:
|
||||
if (
|
||||
isinstance(cache_control, dict)
|
||||
and "ttl" in cache_control
|
||||
):
|
||||
ttl = cache_control["ttl"]
|
||||
if is_claude_4_5 and ttl in ["5m", "1h"]:
|
||||
continue
|
||||
|
||||
cache_control.pop("ttl", None)
|
||||
|
||||
def _supports_extended_thinking_on_bedrock(self, model: str) -> bool:
|
||||
|
|
@ -155,10 +172,18 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
# Supported models on Bedrock for extended thinking
|
||||
supported_patterns = [
|
||||
"opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", # Opus 4.5
|
||||
"opus-4.1", "opus_4.1", "opus-4-1", "opus_4_1", # Opus 4.1
|
||||
"opus-4", "opus_4", # Opus 4
|
||||
"sonnet-4", "sonnet_4", # Sonnet 4
|
||||
"opus-4.5",
|
||||
"opus_4.5",
|
||||
"opus-4-5",
|
||||
"opus_4_5", # Opus 4.5
|
||||
"opus-4.1",
|
||||
"opus_4.1",
|
||||
"opus-4-1",
|
||||
"opus_4_1", # Opus 4.1
|
||||
"opus-4",
|
||||
"opus_4", # Opus 4
|
||||
"sonnet-4",
|
||||
"sonnet_4", # Sonnet 4
|
||||
]
|
||||
|
||||
return any(pattern in model_lower for pattern in supported_patterns)
|
||||
|
|
@ -175,10 +200,27 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
"""
|
||||
model_lower = model.lower()
|
||||
opus_4_5_patterns = [
|
||||
"opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5",
|
||||
"opus-4.5",
|
||||
"opus_4.5",
|
||||
"opus-4-5",
|
||||
"opus_4_5",
|
||||
]
|
||||
return any(pattern in model_lower for pattern in opus_4_5_patterns)
|
||||
|
||||
def _is_claude_4_5_on_bedrock(self, model: str) -> bool:
|
||||
"""
|
||||
Check if the model is Claude 4.5 on Bedrock.
|
||||
|
||||
Claude Sonnet 4.5, Haiku 4.5, and Opus 4.5 support 1-hour prompt caching.
|
||||
|
||||
Args:
|
||||
model: The model name
|
||||
|
||||
Returns:
|
||||
True if the model is Claude 4.5
|
||||
"""
|
||||
return is_claude_4_5_on_bedrock(model)
|
||||
|
||||
def _supports_tool_search_on_bedrock(self, model: str) -> bool:
|
||||
"""
|
||||
Check if the model supports tool search on Bedrock.
|
||||
|
|
@ -199,9 +241,15 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
# Supported models for tool search on Bedrock
|
||||
supported_patterns = [
|
||||
# Opus 4.5
|
||||
"opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5",
|
||||
"opus-4.5",
|
||||
"opus_4.5",
|
||||
"opus-4-5",
|
||||
"opus_4_5",
|
||||
# Sonnet 4.5
|
||||
"sonnet-4.5", "sonnet_4.5", "sonnet-4-5", "sonnet_4_5",
|
||||
"sonnet-4.5",
|
||||
"sonnet_4.5",
|
||||
"sonnet-4-5",
|
||||
"sonnet_4_5",
|
||||
]
|
||||
|
||||
return any(pattern in model_lower for pattern in supported_patterns)
|
||||
|
|
@ -238,8 +286,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
beta_headers_to_remove.add(beta)
|
||||
has_advanced_tool_use = True
|
||||
break
|
||||
|
||||
|
||||
|
||||
# 2. Filter out extended thinking headers for models that don't support them
|
||||
extended_thinking_patterns = [
|
||||
"extended-thinking",
|
||||
|
|
@ -263,7 +310,6 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
beta_set.add("tool-search-tool-2025-10-19")
|
||||
beta_set.add("tool-examples-2025-10-29")
|
||||
|
||||
|
||||
def _get_tool_search_beta_header_for_bedrock(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -290,7 +336,9 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
input_examples_used: Whether input examples are used
|
||||
beta_set: The set of beta headers to modify in-place
|
||||
"""
|
||||
if tool_search_used and not (programmatic_tool_calling_used or input_examples_used):
|
||||
if tool_search_used and not (
|
||||
programmatic_tool_calling_used or input_examples_used
|
||||
):
|
||||
beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER)
|
||||
if "opus-4" in model.lower() or "opus_4" in model.lower():
|
||||
beta_set.add("tool-search-tool-2025-10-19")
|
||||
|
|
@ -302,13 +350,13 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
) -> None:
|
||||
"""
|
||||
Convert Anthropic output_format to inline schema in message content.
|
||||
|
||||
|
||||
Bedrock Invoke doesn't support the output_format parameter, so we embed
|
||||
the schema directly into the user message content as text instructions.
|
||||
|
||||
|
||||
This approach adds the schema to the last user message, instructing the model
|
||||
to respond in the specified JSON format.
|
||||
|
||||
|
||||
Args:
|
||||
output_format: The output_format dict with 'type' and 'schema'
|
||||
anthropic_messages_request: The request dict to modify in-place
|
||||
|
|
@ -321,35 +369,32 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
schema = output_format.get("schema")
|
||||
if not schema:
|
||||
return
|
||||
|
||||
|
||||
# Get messages from the request
|
||||
messages = anthropic_messages_request.get("messages", [])
|
||||
if not messages:
|
||||
return
|
||||
|
||||
|
||||
# Find the last user message
|
||||
last_user_message_idx = None
|
||||
for idx in range(len(messages) - 1, -1, -1):
|
||||
if messages[idx].get("role") == "user":
|
||||
last_user_message_idx = idx
|
||||
break
|
||||
|
||||
|
||||
if last_user_message_idx is None:
|
||||
return
|
||||
|
||||
|
||||
last_user_message = messages[last_user_message_idx]
|
||||
content = last_user_message.get("content", [])
|
||||
|
||||
|
||||
# Ensure content is a list
|
||||
if isinstance(content, str):
|
||||
content = [{"type": "text", "text": content}]
|
||||
last_user_message["content"] = content
|
||||
|
||||
|
||||
# Add schema as text content to the message
|
||||
schema_text = {
|
||||
"type": "text",
|
||||
"text": json.dumps(schema)
|
||||
}
|
||||
schema_text = {"type": "text", "text": json.dumps(schema)}
|
||||
content.append(schema_text)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
|
|
@ -374,9 +419,9 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
# 1. anthropic_version is required for all claude models
|
||||
if "anthropic_version" not in anthropic_messages_request:
|
||||
anthropic_messages_request["anthropic_version"] = (
|
||||
self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION
|
||||
)
|
||||
anthropic_messages_request[
|
||||
"anthropic_version"
|
||||
] = self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION
|
||||
|
||||
# 2. `stream` is not allowed in request body for bedrock invoke
|
||||
if "stream" in anthropic_messages_request:
|
||||
|
|
@ -386,8 +431,10 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
if "model" in anthropic_messages_request:
|
||||
anthropic_messages_request.pop("model", None)
|
||||
|
||||
# 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it)
|
||||
self._remove_ttl_from_cache_control(anthropic_messages_request)
|
||||
# 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models)
|
||||
self._remove_ttl_from_cache_control(
|
||||
anthropic_messages_request=anthropic_messages_request, model=model
|
||||
)
|
||||
|
||||
# 5. Convert `output_format` to inline schema (Bedrock invoke doesn't support output_format)
|
||||
output_format = anthropic_messages_request.pop("output_format", None)
|
||||
|
|
@ -396,14 +443,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
output_format=output_format,
|
||||
anthropic_messages_request=anthropic_messages_request,
|
||||
)
|
||||
|
||||
|
||||
# 6. AUTO-INJECT beta headers based on features used
|
||||
anthropic_model_info = AnthropicModelInfo()
|
||||
tools = anthropic_messages_optional_request_params.get("tools")
|
||||
messages_typed = cast(List[AllMessageValues], messages)
|
||||
tool_search_used = anthropic_model_info.is_tool_search_used(tools)
|
||||
programmatic_tool_calling_used = anthropic_model_info.is_programmatic_tool_calling_used(
|
||||
tools
|
||||
programmatic_tool_calling_used = (
|
||||
anthropic_model_info.is_programmatic_tool_calling_used(tools)
|
||||
)
|
||||
input_examples_used = anthropic_model_info.is_input_examples_used(tools)
|
||||
|
||||
|
|
@ -436,8 +483,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
if beta_set:
|
||||
anthropic_messages_request["anthropic_beta"] = list(beta_set)
|
||||
|
||||
|
||||
|
||||
return anthropic_messages_request
|
||||
|
||||
def get_async_streaming_response_iterator(
|
||||
|
|
@ -455,7 +501,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
)
|
||||
# Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients.
|
||||
return self.bedrock_sse_wrapper(
|
||||
completion_stream=completion_stream,
|
||||
completion_stream=completion_stream,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
request_body=request_body,
|
||||
)
|
||||
|
|
@ -474,14 +520,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
BaseAnthropicMessagesStreamingIterator,
|
||||
)
|
||||
|
||||
handler = BaseAnthropicMessagesStreamingIterator(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
request_body=request_body,
|
||||
)
|
||||
|
||||
|
||||
async for chunk in handler.async_sse_wrapper(completion_stream):
|
||||
yield chunk
|
||||
|
||||
|
||||
|
||||
class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import List, Optional, Tuple
|
||||
|
||||
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.openai.openai import OpenAIConfig
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -29,9 +30,7 @@ class GithubCopilotConfig(OpenAIConfig):
|
|||
api_key: Optional[str],
|
||||
custom_llm_provider: str,
|
||||
) -> Tuple[Optional[str], Optional[str], str]:
|
||||
dynamic_api_base = (
|
||||
self.authenticator.get_api_base() or GITHUB_COPILOT_API_BASE
|
||||
)
|
||||
dynamic_api_base = self.authenticator.get_api_base() or GITHUB_COPILOT_API_BASE
|
||||
try:
|
||||
dynamic_api_key = self.authenticator.get_api_key()
|
||||
except GetAPIKeyError as e:
|
||||
|
|
@ -140,7 +139,7 @@ class GithubCopilotConfig(OpenAIConfig):
|
|||
"""
|
||||
Check if any message contains vision content (images).
|
||||
Returns True if any message has content with vision-related types, otherwise False.
|
||||
|
||||
|
||||
Checks for:
|
||||
- image_url content type (OpenAI format)
|
||||
- Content items with type 'image_url'
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from .openai import ChatCompletionToolCallChunk
|
|||
|
||||
class CachePointBlock(TypedDict, total=False):
|
||||
type: Literal["default"]
|
||||
ttl: str
|
||||
|
||||
|
||||
class SystemContentBlock(TypedDict, total=False):
|
||||
|
|
@ -961,6 +962,7 @@ class BedrockGetBatchResponse(TypedDict, total=False):
|
|||
timeoutDurationInHours: Optional[int]
|
||||
clientRequestToken: Optional[str]
|
||||
|
||||
|
||||
class BedrockToolBlock(TypedDict, total=False):
|
||||
toolSpec: Optional[ToolSpecBlock]
|
||||
systemTool: Optional[SystemToolBlock] # For Nova grounding
|
||||
|
|
|
|||
83
poetry.lock
generated
83
poetry.lock
generated
|
|
@ -725,6 +725,18 @@ markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"pro
|
|||
[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"
|
||||
|
|
@ -1302,6 +1314,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"
|
||||
|
|
@ -3087,7 +3120,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"},
|
||||
|
|
@ -3517,7 +3550,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"},
|
||||
|
|
@ -5522,14 +5555,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)"]
|
||||
|
|
@ -6608,29 +6641,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]]
|
||||
|
|
@ -8498,4 +8531,8 @@ utils = ["numpydoc"]
|
|||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = ">=3.9,<4.0"
|
||||
<<<<<<< litellm_oss_staging_02_04_2026
|
||||
content-hash = "797603dcfef0a79781c7d3cba5dfe18f6aea4aa792220f47487ebc7bd04ae2e3"
|
||||
=======
|
||||
content-hash = "e5447e14dd37e324ac07a8fc6286d27e9a0d355ed93ebb24fc11e3f5df12fd3e"
|
||||
>>>>>>> main
|
||||
|
|
|
|||
|
|
@ -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 = "*"
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
52
tests/logging_callback_tests/test_dynamic_otel_keys.py
Normal file
52
tests/logging_callback_tests/test_dynamic_otel_keys.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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>"
|
||||
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__])
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue