diff --git a/tests/e2e/llm_translation/test_anthropic_passthrough_migrations_e2e.py b/tests/e2e/llm_translation/test_anthropic_passthrough_migrations_e2e.py new file mode 100644 index 00000000000..06c2d128272 --- /dev/null +++ b/tests/e2e/llm_translation/test_anthropic_passthrough_migrations_e2e.py @@ -0,0 +1,297 @@ +from typing import Final, cast + +import pytest +from anthropic.types import Message, MessageParam, TextBlock, ToolUnionParam +from e2e_config import unique_marker +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients + +BEDROCK_MODEL_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +BEDROCK_BACKENDS: Final = ( + pytest.param(f"bedrock/converse/{BEDROCK_MODEL_ID}", id="bedrock-converse"), + pytest.param(f"bedrock/invoke/{BEDROCK_MODEL_ID}", id="bedrock-invoke"), +) +ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" +TOOL_SEARCH_BACKENDS: Final = ( + pytest.param(ANTHROPIC_MODEL, id="anthropic"), + *BEDROCK_BACKENDS, +) +def _tool_search_tools() -> list[ToolUnionParam]: + return [ + cast( + ToolUnionParam, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + ), + cast( + ToolUnionParam, + { + "name": "get_weather", + "description": "Get the current weather for a location", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + "defer_loading": True, + }, + ), + cast( + ToolUnionParam, + { + "name": "get_stock_price", + "description": "Get the current stock price for a ticker symbol", + "input_schema": { + "type": "object", + "properties": {"ticker": {"type": "string"}}, + "required": ["ticker"], + }, + "defer_loading": True, + }, + ), + ] +CACHED_DOCUMENT: Final = "This agreement describes payment terms, renewal dates, and service obligations. " * 120 + +pytestmark = pytest.mark.e2e + + +def _register_deployment( + proxy: ProxyClient, + resources: ResourceManager, + params: LiteLLMParamsBody, + prefix: str, +) -> tuple[str, str]: + model_name: Final = f"{prefix}-{unique_marker()}" + model_id: Final = proxy.create_model(model_name, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model_name, resources.key() + + +def _deployment_params(backend: str) -> LiteLLMParamsBody: + if not backend.startswith("bedrock/"): + return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY") + return LiteLLMParamsBody( + model=backend, + aws_access_key_id="os.environ/AWS_BEDROCK_TEST_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_BEDROCK_TEST_SECRET_ACCESS_KEY", + aws_region_name="us-east-1", + ) + + +def _cached_messages(marker: str) -> list[MessageParam]: + return [ + cast( + MessageParam, + { + "role": "user", + "content": [ + { + "type": "text", + "text": f"{marker} {CACHED_DOCUMENT}", + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": "What are the payment terms?"}, + ], + }, + ) + ] + + +def _message_text(message: Message) -> str: + return "".join(block.text for block in message.content if isinstance(block, TextBlock)) + + +@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) +@pytest.mark.covers( + "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", + "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", +) +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) +) +def test_bedrock_messages_prompt_caching_creates_cache( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str +) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-create") + client = sdk.anthropic(key) + message = client.messages.create( + model=model, + max_tokens=100, + messages=_cached_messages(unique_marker()), + extra_body=NO_PROXY_CACHE, + ) + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert message.usage.cache_creation_input_tokens > 0 or message.usage.cache_read_input_tokens > 0, ( + f"prompt cache was not created or read: {message.usage!r}" + ) + + +@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) +@pytest.mark.covers( + "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", + "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", +) +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) +) +def test_bedrock_messages_prompt_caching_reads_cache_on_second_call( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str +) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-read") + client = sdk.anthropic(key) + messages: Final = _cached_messages(unique_marker()) + + client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + message = client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + + assert message.usage.cache_read_input_tokens > 0, f"cache read was not reported: {message.usage!r}" + + +@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) +@pytest.mark.covers( + "llm.messages.bedrock_converse.basic.stream.works", + "llm.messages.bedrock_invoke.basic.stream.works", +) +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.STREAM, + ) +) +def test_bedrock_messages_streaming_prompt_caching_reads_cache( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str +) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-stream") + client = sdk.anthropic(key) + messages: Final = _cached_messages(unique_marker()) + client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + + with client.messages.stream( + model=model, + max_tokens=100, + messages=messages, + extra_body=NO_PROXY_CACHE, + ) as stream: + message: Final = stream.get_final_message() + + assert message.usage.cache_read_input_tokens > 0, f"streaming cache read was not reported: {message.usage!r}" + assert _message_text(message).strip(), f"streaming response contained no text: {message.content!r}" + + +@pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) +@pytest.mark.covers( + "llm.messages.anthropic.tool_search.nonstream.works", + "llm.messages.bedrock_converse.tool_search.nonstream.works", + "llm.messages.bedrock_invoke.tool_search.nonstream.works", +) +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.ANTHROPIC, Provider.BEDROCK), + models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), + mode=Mode.NONSTREAM, + ) +) +def test_anthropic_tool_search_returns_a_message( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str +) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search") + message = sdk.anthropic(key).messages.create( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": f"Find the weather tool for San Francisco. {unique_marker()}"}], + tools=_tool_search_tools(), + extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, + extra_body=NO_PROXY_CACHE, + ) + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert message.content, "tool search response contained no content" + assert message.usage.input_tokens > 0, f"tool search response had no input usage: {message.usage!r}" + + +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.ANTHROPIC, Provider.BEDROCK), + models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), + mode=Mode.STREAM, + ) +) +@pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) +@pytest.mark.covers( + "llm.messages.anthropic.tool_search.nonstream.works", + "llm.messages.bedrock_converse.tool_search.nonstream.works", + "llm.messages.bedrock_invoke.tool_search.nonstream.works", +) +def test_anthropic_tool_search_streaming_returns_a_message( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str +) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search-stream") + client = sdk.anthropic(key) + + with client.messages.stream( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": f"Find the weather tool for Tokyo. {unique_marker()}"}], + tools=_tool_search_tools(), + extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, + extra_body=NO_PROXY_CACHE, + ) as stream: + message: Final = stream.get_final_message() + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert message.content, "streaming tool search response contained no content" + assert message.usage.input_tokens > 0, f"streaming tool search response had no input usage: {message.usage!r}" + + +@pytest.mark.covers("llm.messages.bedrock_invoke.basic.nonstream.works") +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) +) +def test_bedrock_invoke_messages_accepts_forwarded_headers( + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients +) -> None: + params: Final = _deployment_params(f"bedrock/invoke/{BEDROCK_MODEL_ID}") + model, key = _register_deployment(proxy, resources, params, "e2e-bedrock-forwarded-headers") + message = sdk.anthropic(key).messages.create( + model=model, + max_tokens=32, + messages=[{"role": "user", "content": f"Reply with one word. {unique_marker()}"}], + extra_headers={ + "x-forwarded-for": "10.11.232.194", + "x-forwarded-port": "443", + "x-forwarded-proto": "https", + "x-app": "cli", + }, + extra_body=NO_PROXY_CACHE, + ) + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert _message_text(message).strip(), f"forwarded-header response contained no text: {message.content!r}" diff --git a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py index f860fecf47d..a8b0de8f978 100644 --- a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py +++ b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py @@ -1,65 +1,24 @@ import hashlib -import json -import uuid -from typing import Final, Protocol, cast +from typing import Final import pytest from e2e_metadata import Domain, Subject, meta -from pydantic import JsonValue, TypeAdapter +from pydantic import TypeAdapter -import litellm from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.secret_managers.main import get_secret pytestmark = pytest.mark.e2e _AWS_FIXTURE_MASTER_KEY_SHA256: Final = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _STRING_MAP: Final = TypeAdapter(dict[str, str]) -_STRING: Final = TypeAdapter(str) -class _AWSSecretManagerOperations(Protocol): - async def async_write_secret(self, secret_name: str, secret_value: str) -> object: ... - - async def async_delete_secret( - self, secret_name: str, recovery_window_in_days: int | None = 7 - ) -> object: ... - - -@pytest.mark.asyncio @meta(Subject(domain=Domain.DEPLOY_OPS)) -async def test_aws_secret_manager() -> None: +def test_aws_secret_manager() -> None: AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True) - - fixture_value: Final[str | None] = cast( - str | None, TypeAdapter(str | None).validate_python(get_secret("litellm_master_key")) + secret_value: Final = _STRING_MAP.validate_json(get_secret("litellm_master_key")) + assert ( + hashlib.sha256(secret_value["litellm_master_key"].encode()).hexdigest() + == _AWS_FIXTURE_MASTER_KEY_SHA256 ) - if fixture_value is not None: - fixture_data: Final = _STRING_MAP.validate_json(_STRING.validate_python(fixture_value)) - assert hashlib.sha256(fixture_data["litellm_master_key"].encode()).hexdigest() == ( - _AWS_FIXTURE_MASTER_KEY_SHA256 - ) - return - - secret_name: Final = f"litellm-e2e-secret-manager-{uuid.uuid4().hex}" - secret_value: Final = f"sk-litellm-e2e-{uuid.uuid4().hex}" - manager: Final = litellm.secret_manager_client - assert isinstance(manager, AWSSecretsManagerV2) - secret_manager: Final = cast(_AWSSecretManagerOperations, manager) - try: - creation: Final = _JSON_OBJECT.validate_python( - await secret_manager.async_write_secret( - secret_name, - json.dumps({"litellm_master_key": secret_value}), - ) - ) - assert creation.get("Name") == secret_name - created_value: Final = _STRING.validate_python(get_secret(secret_name)) - created_data: Final = _STRING_MAP.validate_json(created_value) - assert created_data["litellm_master_key"] == secret_value - finally: - deletion: Final = _JSON_OBJECT.validate_python( - await secret_manager.async_delete_secret(secret_name, recovery_window_in_days=7) - ) - assert deletion.get("Name") == secret_name diff --git a/tests/logging_callback_tests/logging_worker_drain_canary.py b/tests/integration/observability/logging_worker_drain_canary.py similarity index 89% rename from tests/logging_callback_tests/logging_worker_drain_canary.py rename to tests/integration/observability/logging_worker_drain_canary.py index bff29129d7d..fbdbf3a9744 100644 --- a/tests/logging_callback_tests/logging_worker_drain_canary.py +++ b/tests/integration/observability/logging_worker_drain_canary.py @@ -19,5 +19,5 @@ async def test_2_never_inherits_the_pending_event() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) queued_on, ran_on = RUNS.get_nowait() assert RUNS.empty() - assert ran_on is queued_on - assert ran_on is not asyncio.get_running_loop() + assert queued_on.is_closed() + assert ran_on is asyncio.get_running_loop() diff --git a/tests/logging_callback_tests/test_logging_worker_drain.py b/tests/integration/observability/test_logging_worker_drain.py similarity index 90% rename from tests/logging_callback_tests/test_logging_worker_drain.py rename to tests/integration/observability/test_logging_worker_drain.py index e6fef9880a1..1d6a141a197 100644 --- a/tests/logging_callback_tests/test_logging_worker_drain.py +++ b/tests/integration/observability/test_logging_worker_drain.py @@ -6,6 +6,8 @@ from tests.test_litellm_rust.support.child_interpreter import run_child_interpre CANARY_MODULE: Final = Path(__file__).with_name("logging_worker_drain_canary.py") CANARY_RUN: Final = ( + "import sys\n" + f"sys.path.insert(0, {str(Path(__file__).parents[2])!r})\n" "import pytest\n" f"raise SystemExit(pytest.main([{str(CANARY_MODULE)!r}, '-p', 'no:xdist', '-p', 'no:cacheprovider', '-q']))\n" ) diff --git a/tests/logging_callback_tests/base_test.py b/tests/logging_callback_tests/base_test.py deleted file mode 100644 index 11d89996474..00000000000 --- a/tests/logging_callback_tests/base_test.py +++ /dev/null @@ -1,85 +0,0 @@ -import asyncio -import httpx -import json -import pytest -from typing import Any, Dict, List -from unittest.mock import MagicMock, Mock, patch - -import litellm -from litellm.exceptions import BadRequestError -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.utils import CustomStreamWrapper - -# test_example.py -from abc import ABC - - -class BaseLoggingCallbackTest(ABC): - """ - Abstract base test class that enforces a common test across all test classes. - """ - - @pytest.fixture - def mock_response_obj(self): - from litellm.types.utils import ( - ModelResponse, - Choices, - Message, - ChatCompletionMessageToolCall, - Function, - Usage, - CompletionTokensDetailsWrapper, - PromptTokensDetailsWrapper, - ) - - # Create a mock response object with the structure you need - return ModelResponse( - id="chatcmpl-ASId3YJWagBpBskWfoNEMPFSkmrEw", - created=1731308157, - model="gpt-4o-mini-2024-07-18", - object="chat.completion", - system_fingerprint="fp_0ba0d124f1", - choices=[ - Choices( - finish_reason="tool_calls", - index=0, - message=Message( - content=None, - role="assistant", - tool_calls=[ - ChatCompletionMessageToolCall( - function=Function( - arguments='{"city": "New York"}', name="get_weather" - ), - id="call_PngsQS5YGmIZKnswhnUOnOVb", - type="function", - ), - ChatCompletionMessageToolCall( - function=Function( - arguments='{"city": "New York"}', name="get_news" - ), - id="call_1zsDThBu0VSK7KuY7eCcJBnq", - type="function", - ), - ], - function_call=None, - ), - ) - ], - usage=Usage( - completion_tokens=46, - prompt_tokens=86, - total_tokens=132, - completion_tokens_details=CompletionTokensDetailsWrapper( - accepted_prediction_tokens=0, - audio_tokens=0, - reasoning_tokens=0, - rejected_prediction_tokens=0, - text_tokens=None, - ), - prompt_tokens_details=PromptTokensDetailsWrapper( - audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None - ), - ), - service_tier=None, - ) diff --git a/tests/logging_callback_tests/conftest.py b/tests/logging_callback_tests/conftest.py deleted file mode 100644 index 066afdf5c15..00000000000 --- a/tests/logging_callback_tests/conftest.py +++ /dev/null @@ -1,239 +0,0 @@ -# conftest.py -# -# xdist-compatible test isolation for logging callback tests. -# -# Key design: capture litellm's true default values at conftest import time -# (BEFORE test modules are imported) so we can reset to clean defaults before -# each test. This is necessary because some test modules set module-level -# globals like `litellm.num_retries = 3` which pollute state for all tests -# in the same xdist worker. - -import asyncio -import importlib -import os -from collections.abc import AsyncIterator -from typing import Final - -import pytest -import pytest_asyncio - -import litellm -from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - -from tests._vcr_conftest_common import ( # noqa: E402,F401 - VerboseReporterState, - _pin_multipart_boundary, - apply_vcr_auto_marker_to_items, - emit_cassette_cache_session_banner, - emit_vcr_classification_summary, - emit_vcr_diagnostic_log, - install_live_call_probe, - record_vcr_outcome, - register_persister_if_enabled, - reset_vcr_diag_dir, - vcr_config_dict, -) - -# vcrpy and respx both patch the httpx transport ā€” applying both makes one -# silently win, so respx-using files opt out of the auto-marker. -_RESPX_CONFLICTING_FILES = frozenset( - { - "test_assemble_streaming_responses.py", - "test_langfuse_unit_tests.py", - } -) - -_VCR_INCOMPATIBLE_FILES = frozenset() - -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () - - -_verbose_state = VerboseReporterState() - - -@pytest.fixture(scope="module") -def vcr_config(): - return vcr_config_dict() - - -def pytest_recording_configure(config, vcr): - register_persister_if_enabled(vcr) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_makereport(item, call): - outcome = yield - rep = outcome.get_result() - setattr(item, f"rep_{rep.when}", rep) - - -@pytest.fixture(autouse=True) -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -def pytest_configure(config): - _verbose_state.remember_pluginmanager(config) - reset_vcr_diag_dir() - - -def pytest_runtest_logreport(report): - _verbose_state.maybe_emit_verdict(report) - - -_LIST_ATTRS = ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - "service_callback", - "pre_call_rules", - "post_call_rules", -) - -_SCALAR_ATTRS = ( - "set_verbose", - "cache", - "num_retries", - "num_retries_per_request", - "turn_off_message_logging", - "redact_messages_in_exceptions", - "redact_user_api_key_info", - "s3_callback_params", - "s3_audit_callback_params", - "datadog_params", - "vector_store_registry", -) - -# ---- Capture true defaults at conftest import time ---- -# This runs BEFORE any test modules are imported, so values are clean. -_DEFAULTS: dict = {} -for _attr in _LIST_ATTRS: - if hasattr(litellm, _attr): - _val = getattr(litellm, _attr) - _DEFAULTS[_attr] = _val.copy() if isinstance(_val, list) else _val -for _attr in _SCALAR_ATTRS: - if hasattr(litellm, _attr): - _DEFAULTS[_attr] = getattr(litellm, _attr) - - -@pytest.fixture(scope="function", autouse=True) -def isolate_litellm_state(): - """ - Per-function isolation fixture. - - Resets litellm state to the true defaults captured at conftest import time, - then restores after the test. This prevents module-level mutations (e.g. - `litellm.num_retries = 3` at the top of test_langfuse_e2e_test.py) from - leaking across tests within the same xdist worker. - """ - from litellm.litellm_core_utils import litellm_logging as ll_logging - from litellm.proxy.management_helpers import audit_logs as ll_audit_logs - - # Flush cache and clear internal logger instances before test - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - - # Clear cached logger instances (LangsmithLogger, SlackAlerting, etc.) - ll_logging._in_memory_loggers.clear() - ll_audit_logs._audit_log_callback_cache.clear() - - # Reset ALL attrs to their true defaults before the test runs. - # This undoes any module-level mutations from test file imports. - for attr in _LIST_ATTRS: - if attr in _DEFAULTS: - default = _DEFAULTS[attr] - setattr( - litellm, attr, default.copy() if isinstance(default, list) else default - ) - - for attr in _SCALAR_ATTRS: - if attr in _DEFAULTS: - setattr(litellm, attr, _DEFAULTS[attr]) - - yield - - # Teardown: reset back to defaults again (belt-and-suspenders) - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - - ll_logging._in_memory_loggers.clear() - ll_audit_logs._audit_log_callback_cache.clear() - - for attr in _LIST_ATTRS: - if attr in _DEFAULTS: - default = _DEFAULTS[attr] - setattr( - litellm, attr, default.copy() if isinstance(default, list) else default - ) - - for attr in _SCALAR_ATTRS: - if attr in _DEFAULTS: - setattr(litellm, attr, _DEFAULTS[attr]) - - -LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0 - - -@pytest_asyncio.fixture(loop_scope="function", autouse=True) -async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]: - yield - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) - - -@pytest.fixture(scope="module", autouse=True) -def setup_and_teardown(): - """ - Module-scoped setup. Reloads litellm only in single-process mode - (skipped under xdist to avoid cross-worker interference). - """ - - import litellm - - worker_id = os.environ.get("PYTEST_XDIST_WORKER", None) - if worker_id is None: - importlib.reload(litellm) - - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - import litellm.proxy.proxy_server - - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - - yield - - -def pytest_collection_modifyitems(config, items): - apply_vcr_auto_marker_to_items( - items, - skip_files=_RESPX_CONFLICTING_FILES | _VCR_INCOMPATIBLE_FILES, - skip_nodeid_suffixes=_VCR_INCOMPATIBLE_NODEID_SUFFIXES, - ) - - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests - custom_logger_tests = [ - item for item in items if "custom_logger" in item.parent.name - ] - other_tests = [item for item in items if "custom_logger" not in item.parent.name] - - # Sort tests based on their names - custom_logger_tests.sort(key=lambda x: x.name) - other_tests.sort(key=lambda x: x.name) - - # Reorder the items list - items[:] = custom_logger_tests + other_tests - - -def pytest_terminal_summary(terminalreporter, exitstatus, config): - emit_cassette_cache_session_banner(terminalreporter) - emit_vcr_classification_summary(terminalreporter) - emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/logging_callback_tests/create_mock_standard_logging_payload.py b/tests/logging_callback_tests/create_mock_standard_logging_payload.py deleted file mode 100644 index 096c8ff8c60..00000000000 --- a/tests/logging_callback_tests/create_mock_standard_logging_payload.py +++ /dev/null @@ -1,128 +0,0 @@ -import io - - - -import asyncio -import gzip -import json -import logging -import time -from unittest.mock import AsyncMock, patch - -import pytest - -import litellm -from litellm import completion -from litellm._logging import verbose_logger -from litellm.integrations.datadog.datadog import * -from datetime import datetime, timedelta -from litellm.types.utils import ( - StandardLoggingPayload, - StandardLoggingModelInformation, - StandardLoggingMetadata, - StandardLoggingHiddenParams, -) - -verbose_logger.setLevel(logging.DEBUG) - - -def create_standard_logging_payload() -> StandardLoggingPayload: - return StandardLoggingPayload( - id="test_id", - call_type="completion", - response_cost=0.1, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=20, - completion_tokens=10, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-5-mini", model_map_value=None - ), - model="gpt-5-mini", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.1", - additional_headers=None, - ), - ) - - -def create_standard_logging_payload_with_long_content() -> StandardLoggingPayload: - return StandardLoggingPayload( - id="test_id", - call_type="completion", - response_cost=0.1, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=20, - completion_tokens=10, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-5-mini", model_map_value=None - ), - model="gpt-5-mini", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!" * 80000}], - response={"choices": [{"message": {"content": "Hi there!" * 80000}}]}, - error_str="error_str" * 80000, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.1", - additional_headers=None, - ), - ) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json deleted file mode 100644 index 82173b57cc2..00000000000 --- a/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json +++ /dev/null @@ -1,179 +0,0 @@ -{ - "id": "chatcmpl-2299b6a2-82a3-465a-b47c-04e685a2227f", - "trace_id": null, - "call_type": "acompletion", - "cache_hit": null, - "stream": true, - "status": "success", - "custom_llm_provider": "openai", - "saved_cache_cost": 0.0, - "startTime": "2025-01-24 09:20:46.847371", - "endTime": "2025-01-24 09:20:46.851954", - "completionStartTime": "2025-01-24 09:20:46.851954", - "response_time": 0.007394075393676758, - "model": "gpt-4o", - "metadata": { - "user_api_key_hash": null, - "user_api_key_alias": null, - "user_api_key_team_id": null, - "user_api_key_org_id": null, - "user_api_key_user_id": null, - "user_api_key_team_alias": null, - "user_api_key_user_email": null, - "spend_logs_metadata": null, - "requester_ip_address": null, - "requester_metadata": null, - "user_api_key_end_user_id": null, - "prompt_management_metadata": null, - "applied_guardrails": [], - "mcp_tool_call_metadata": null, - "vector_store_request_metadata": null, - "guardrail_information": null - }, - "cache_key": null, - "response_cost": 0.00022500000000000002, - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - "request_tags": [], - "end_user": "", - "api_base": "", - "model_group": "", - "model_id": "", - "requester_ip_address": null, - "messages": [ - { - "role": "user", - "content": "Hello, world!" - } - ], - "response": { - "id": "chatcmpl-2299b6a2-82a3-465a-b47c-04e685a2227f", - "created": 1742855151, - "model": "gpt-4o", - "object": "chat.completion", - "system_fingerprint": null, - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "message": { - "content": "hi", - "role": "assistant", - "tool_calls": null, - "function_call": null, - "provider_specific_fields": null - } - } - ], - "usage": { - "completion_tokens": 20, - "prompt_tokens": 10, - "total_tokens": 30, - "completion_tokens_details": null, - "prompt_tokens_details": null - } - }, - "model_parameters": {}, - "hidden_params": { - "model_id": null, - "cache_key": null, - "api_base": "https://api.openai.com", - "response_cost": 0.00022500000000000002, - "additional_headers": {}, - "litellm_overhead_time_ms": null, - "batch_models": null, - "litellm_model_name": "gpt-4o" - }, - "model_map_information": { - "model_map_key": "gpt-4o", - "model_map_value": { - "key": "gpt-4o", - "max_tokens": 16384, - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "input_cost_per_token": 2.5e-06, - "cache_creation_input_token_cost": null, - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_character": null, - "input_cost_per_token_above_128k_tokens": null, - "input_cost_per_query": null, - "input_cost_per_second": null, - "input_cost_per_audio_token": null, - "input_cost_per_token_batches": 1.25e-06, - "output_cost_per_token_batches": 5e-06, - "output_cost_per_token": 1e-05, - "output_cost_per_audio_token": null, - "output_cost_per_character": null, - "output_cost_per_token_above_128k_tokens": null, - "output_cost_per_character_above_128k_tokens": null, - "output_cost_per_second": null, - "output_cost_per_image": null, - "output_vector_size": null, - "litellm_provider": "openai", - "mode": "chat", - "supports_system_messages": true, - "supports_response_schema": true, - "supports_vision": true, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_assistant_prefill": false, - "supports_prompt_caching": true, - "supports_audio_input": false, - "supports_audio_output": false, - "supports_pdf_input": false, - "supports_embedding_image_input": false, - "supports_native_streaming": null, - "supports_web_search": true, - "search_context_cost_per_query": { - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.05 - }, - "tpm": null, - "rpm": null, - "supported_openai_params": [ - "frequency_penalty", - "logit_bias", - "logprobs", - "top_logprobs", - "max_tokens", - "max_completion_tokens", - "modalities", - "prediction", - "n", - "presence_penalty", - "seed", - "stop", - "stream", - "stream_options", - "temperature", - "top_p", - "tools", - "tool_choice", - "function_call", - "functions", - "max_retries", - "extra_headers", - "parallel_tool_calls", - "audio", - "response_format", - "user" - ] - } - }, - "error_str": null, - "error_information": { - "error_code": "", - "error_class": "", - "llm_provider": "", - "traceback": "", - "error_message": "" - }, - "response_cost_failure_debug_info": null, - "guardrail_information": null, - "standard_built_in_tools_params": { - "web_search_options": null, - "file_search": null - } - } \ No newline at end of file diff --git a/tests/logging_callback_tests/gettysburg.wav b/tests/logging_callback_tests/gettysburg.wav deleted file mode 100644 index 9690f521e84..00000000000 Binary files a/tests/logging_callback_tests/gettysburg.wav and /dev/null differ diff --git a/tests/logging_callback_tests/log.txt b/tests/logging_callback_tests/log.txt deleted file mode 100644 index 8e170fa500f..00000000000 --- a/tests/logging_callback_tests/log.txt +++ /dev/null @@ -1,21 +0,0 @@ -============================= test session starts ============================== -platform darwin -- Python 3.13.1, pytest-8.3.5, pluggy-1.5.0 -rootdir: /Users/krrishdholakia/Documents/litellm -configfile: pyproject.toml -plugins: respx-0.22.0, postgresql-7.0.1, anyio-4.4.0, asyncio-0.26.0, mock-3.14.0, ddtrace-2.19.0rc1 -asyncio: mode=Mode.STRICT, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function -collected 1 item - -test_gcs_pub_sub.py . [100%] - -=============================== warnings summary =============================== -../../myenv/lib/python3.13/site-packages/pydantic/_internal/_config.py:295 - /Users/krrishdholakia/Documents/litellm/myenv/lib/python3.13/site-packages/pydantic/_internal/_config.py:295: PydanticDeprecatedSince20: Support for class-based `config` is deprecated, use ConfigDict instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.10/migration/ - warnings.warn(DEPRECATION_MESSAGE, DeprecationWarning) - -../../litellm/caching/llm_caching_handler.py:17 - /Users/krrishdholakia/Documents/litellm/litellm/caching/llm_caching_handler.py:17: DeprecationWarning: There is no current event loop - event_loop = asyncio.get_event_loop() - --- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html -======================== 1 passed, 2 warnings in 4.10s ========================= diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py deleted file mode 100644 index 8c17ee132d8..00000000000 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ /dev/null @@ -1,228 +0,0 @@ -import os - - - -import litellm -import litellm.vector_stores.main -import json -from typing import Optional - -import pytest - -import litellm -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import ( - StandardLoggingPayload, -) - - -class MockCustomLogger(CustomLogger): - def __init__(self): - self.standard_logging_payload: Optional[StandardLoggingPayload] = None - self.completion_logging_payload: Optional[StandardLoggingPayload] = None - super().__init__() - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - payload = kwargs.get("standard_logging_object") - # Store the payload - completion calls have call_type='acompletion' - if payload and payload.get("call_type") == "acompletion": - self.completion_logging_payload = payload - self.standard_logging_payload = payload - pass - - -@pytest.fixture(autouse=True) -def add_aws_region_to_env(monkeypatch): - monkeypatch.setenv("AWS_REGION", "us-west-2") - - -@pytest.fixture -def setup_vector_store_registry(): - from litellm.vector_stores.vector_store_registry import ( - VectorStoreRegistry, - LiteLLM_ManagedVectorStore, - ) - - # Init vector store registry - litellm.vector_store_registry = VectorStoreRegistry( - vector_stores=[ - LiteLLM_ManagedVectorStore( - vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock" - ) - ] - ) - - -@pytest.mark.asyncio -async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( - setup_vector_store_registry, -): - """ - Test that the Bedrock Knowledge Base Hook works when making a real llm api call and returns citations. - """ - - # Init client - litellm.turn_on_debug() - async_client = AsyncHTTPHandler() - response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], - client=async_client, - ) - print("OPENAI RESPONSE:", json.dumps(dict(response), indent=4, default=str)) - assert response is not None - - # Check that search_results are present in provider_specific_fields - assert hasattr(response.choices[0].message, "provider_specific_fields") - provider_fields = response.choices[0].message.provider_specific_fields - assert provider_fields is not None - assert "search_results" in provider_fields - search_results = provider_fields["search_results"] - assert search_results is not None - assert len(search_results) > 0 - - # Check search result structure (OpenAI-compatible format) - first_search_result = search_results[0] - assert "object" in first_search_result - assert first_search_result["object"] == "vector_store.search_results.page" - assert "data" in first_search_result - assert len(first_search_result["data"]) > 0 - - # Check individual result structure - first_result = first_search_result["data"][0] - assert "score" in first_result - assert "content" in first_result - print(f"Search results returned: {len(search_results)}") - print(f"First search result has {len(first_search_result['data'])} items") - - -@pytest.mark.asyncio -async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( - setup_vector_store_registry, -): - """ - Test that the Bedrock Knowledge Base Hook works when making a real llm api call - """ - - # Init client - litellm.turn_on_debug() - response = await litellm.acompletion( - model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", - messages=[{"role": "user", "content": "what is litellm?"}], - max_tokens=10, - tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], - ) - assert response is not None - - -@pytest.mark.asyncio -async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_and_filters( - setup_vector_store_registry, -): - """ - Test that filters from file_search tools are properly passed through to vector store search. - This test verifies the entire flow: tool parsing -> filter extraction -> vector store API call. - - In this case we filter for a non-existent user_id, which should return no results. - """ - litellm.turn_on_debug() - - response = await litellm.acompletion( - model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", - messages=[{"role": "user", "content": "what is litellm?"}], - max_tokens=10, - tools=[ - { - "type": "file_search", - "vector_store_ids": ["T37J8R4WTM"], - "filters": { - "key": "user_id", - "value": "fake-user-id", - "operator": "eq", - }, - } - ], - ) - - # Verify response is not None - assert response is not None - - # Verify search results were added to the response (this proves the search was called) - assert hasattr(response.choices[0].message, "provider_specific_fields") - provider_fields = response.choices[0].message.provider_specific_fields - assert provider_fields is not None - assert ( - "search_results" in provider_fields - ), "search_results not in provider_specific_fields" - - search_results = provider_fields["search_results"] - assert ( - search_results is not None and len(search_results) > 0 - ), "No search results found" - - # The search was performed - this confirms filters were passed through - # The logs above show: litellm.asearch(... filters={'key': 'user_id', 'value': 'fake-user-id', 'operator': 'eq'}) - # And the Bedrock API request contains: {'filter': {'equals': {'key': 'user_id', 'value': 'fake-user-id'}}} - - print("āœ… Filters were successfully passed through to vector store search") - print(f" Search was performed and {len(search_results)} result(s) returned") - - -# @pytest.mark.asyncio -# async def test_logging_with_knowledge_base_hook(setup_vector_store_registry): -# """ -# Test that the knowledge base request was logged in standard logging payload -# """ -# test_custom_logger = MockCustomLogger() -# litellm.set_verbose = True -# await litellm.acompletion( -# model="gpt-5.5", -# messages=[{"role": "user", "content": "what is litellm?"}], -# vector_store_ids = [ -# "T37J8R4WTM" -# ], -# ) - -# # sleep for 1 second to allow the logging callback to run -# await asyncio.sleep(1) - -# # assert that the knowledge base request was logged in the standard logging payload -# standard_logging_payload: Optional[StandardLoggingPayload] = test_custom_logger.standard_logging_payload -# assert standard_logging_payload is not None - - -# metadata = standard_logging_payload["metadata"] -# standard_logging_vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = metadata["vector_store_request_metadata"] - -# print("standard_logging_vector_store_request_metadata:", json.dumps(standard_logging_vector_store_request_metadata, indent=4, default=str)) - -# # 1 vector store request was made, expect 1 vector store request metadata object -# assert len(standard_logging_vector_store_request_metadata) == 1 - -# # expect the vector store request metadata object to have the correct values -# vector_store_request_metadata = standard_logging_vector_store_request_metadata[0] -# assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM" -# assert vector_store_request_metadata.get("query") == "what is litellm?" -# assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock" - - -# vector_store_search_response: VectorStoreSearchResponse = vector_store_request_metadata.get("vector_store_search_response") -# assert vector_store_search_response is not None -# assert vector_store_search_response.get("search_query") == "what is litellm?" -# assert len(vector_store_search_response.get("data", [])) >=0 -# for item in vector_store_search_response.get("data", []): -# assert item.get("score") is not None -# assert item.get("content") is not None -# assert len(item.get("content", [])) >= 0 -# for content_item in item.get("content", []): -# text_content = content_item.get("text") -# assert text_content is not None -# assert len(text_content) > 0 - - - - - - diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py deleted file mode 100644 index 1f1ca8960f6..00000000000 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ /dev/null @@ -1,297 +0,0 @@ -import io -import os - - - -import asyncio -import litellm -import gzip -import json -import logging -import time -from unittest.mock import AsyncMock, patch - -import pytest - -from litellm import completion -from litellm._logging import verbose_logger -from litellm.integrations.gcs_pubsub.pub_sub import * -from datetime import datetime, timedelta -from litellm.types.utils import ( - StandardLoggingPayload, - StandardLoggingModelInformation, - StandardLoggingMetadata, - StandardLoggingHiddenParams, -) - -verbose_logger.setLevel(logging.DEBUG) - -ignored_keys = [ - "request_id", - "litellm_call_id", - "metadata.litellm_call_id", - "session_id", - "startTime", - "endTime", - "completionStartTime", - "endTime", - "request_duration_ms", - "metadata.model_map_information", - "metadata.usage_object", - "metadata.cold_storage_object_key", - "metadata.litellm_overhead_time_ms", - "metadata.cost_breakdown", - "metadata.autorouter_savings", - "metadata.eval_information", -] - - -def _compare_nested_dicts( - actual: dict, expected: dict, path: str = "", ignore_keys: list[str] = [] -) -> list[str]: - """Compare nested dictionaries and return a list of differences in a human-friendly format.""" - differences = [] - - # Check if current path should be ignored - if path in ignore_keys: - return differences - - # Check for keys in actual but not in expected - for key in actual.keys(): - current_path = f"{path}.{key}" if path else key - if current_path not in ignore_keys and key not in expected: - differences.append(f"Extra key in actual: {current_path}") - - for key, expected_value in expected.items(): - current_path = f"{path}.{key}" if path else key - if current_path in ignore_keys: - continue - if key not in actual: - differences.append(f"Missing key: {current_path}") - continue - - actual_value = actual[key] - - # Try to parse JSON strings - if isinstance(expected_value, str): - try: - expected_value = json.loads(expected_value) - except json.JSONDecodeError: - pass - if isinstance(actual_value, str): - try: - actual_value = json.loads(actual_value) - except json.JSONDecodeError: - pass - - if isinstance(expected_value, dict) and isinstance(actual_value, dict): - differences.extend( - _compare_nested_dicts( - actual_value, expected_value, current_path, ignore_keys - ) - ) - elif isinstance(expected_value, dict) or isinstance(actual_value, dict): - differences.append( - f"Type mismatch at {current_path}: expected dict, got {type(actual_value).__name__}" - ) - else: - # For non-dict values, only report if they're different - if actual_value != expected_value: - # Format the values to be more readable - actual_str = str(actual_value) - expected_str = str(expected_value) - if len(actual_str) > 50 or len(expected_str) > 50: - actual_str = f"{actual_str[:50]}..." - expected_str = f"{expected_str[:50]}..." - differences.append( - f"Value mismatch at {current_path}:\n expected: {expected_str}\n got: {actual_str}" - ) - return differences - - -def assert_gcs_pubsub_request_matches_expected( - actual_request_body: dict, - expected_file_name: str, -): - """ - Helper function to compare actual GCS PubSub request body with expected JSON file. - - Args: - actual_request_body (dict): The actual request body received from the API call - expected_file_name (str): Name of the JSON file containing expected request body - """ - # Get the current directory and read the expected request body - pwd = os.path.dirname(os.path.realpath(__file__)) - expected_body_path = os.path.join(pwd, "gcs_pub_sub_body", expected_file_name) - - with open(expected_body_path, "r") as f: - expected_request_body = json.load(f) - - # Replace dynamic values in actual request body - differences = _compare_nested_dicts( - actual_request_body, expected_request_body, ignore_keys=ignored_keys - ) - if differences: - pytest.fail(f"Dictionary mismatch: {differences}") - - -def assert_gcs_pubsub_request_matches_expected_standard_logging_payload( - actual_request_body: dict, - expected_file_name: str, -): - """ - Helper function to compare actual GCS PubSub request body with expected JSON file. - - Args: - actual_request_body (dict): The actual request body received from the API call - expected_file_name (str): Name of the JSON file containing expected request body - """ - # Get the current directory and read the expected request body - pwd = os.path.dirname(os.path.realpath(__file__)) - expected_body_path = os.path.join(pwd, "gcs_pub_sub_body", expected_file_name) - - with open(expected_body_path, "r") as f: - expected_request_body = json.load(f) - - # Replace dynamic values in actual request body - FIELDS_TO_VALIDATE = [ - "custom_llm_provider", - "hidden_params", - "messages", - "response", - "model", - "status", - "stream", - ] - - actual_request_body["response"]["id"] = expected_request_body["response"]["id"] - actual_request_body["response"]["created"] = expected_request_body["response"][ - "created" - ] - - for field in FIELDS_TO_VALIDATE: - assert field in actual_request_body - - FIELDS_EXISTENCE_CHECKS = [ - "response_cost", - "response_time", - "completion_tokens", - "prompt_tokens", - "total_tokens", - ] - - for field in FIELDS_EXISTENCE_CHECKS: - assert field in actual_request_body - - -@pytest.mark.asyncio -async def test_async_gcs_pub_sub(): - # Create a mock for the async_httpx_client's post method - mock_post = AsyncMock() - mock_post.return_value.status_code = 202 - mock_post.return_value.text = "Accepted" - - # Initialize the GcsPubSubLogger and set the mock - gcs_pub_sub_logger = GcsPubSubLogger( - project_id="STUBBED_PROJECT_ID", topic_id="STUBBED_TOPIC_ID", flush_interval=1 - ) - gcs_pub_sub_logger.async_httpx_client.post = mock_post - - mock_construct_request_headers = AsyncMock() - mock_construct_request_headers.return_value = {"Authorization": "Bearer mock_token"} - gcs_pub_sub_logger.construct_request_headers = mock_construct_request_headers - litellm.callbacks = [gcs_pub_sub_logger] - - # Make the completion call - response = await litellm.acompletion( - model="gpt-4o", - messages=[{"role": "user", "content": "Hello, world!"}], - mock_response="hi", - ) - - await asyncio.sleep(3) # Wait for async flush - - # Assert httpx post was called - mock_post.assert_called_once() - - # Get the actual request body from the mock - actual_url = mock_post.call_args[1]["url"] - print("sent to url", actual_url) - assert ( - actual_url - == "https://pubsub.googleapis.com/v1/projects/STUBBED_PROJECT_ID/topics/STUBBED_TOPIC_ID:publish" - ) - actual_request = mock_post.call_args[1]["json"] - - # Extract and decode the base64 encoded message - encoded_message = actual_request["messages"][0]["data"] - import base64 - - decoded_message = base64.b64decode(encoded_message).decode("utf-8") - - # Parse the JSON string into a dictionary - actual_request = json.loads(decoded_message) - print("##########\n") - print(json.dumps(actual_request, indent=4)) - print("##########\n") - # Verify the request body matches expected format - assert_gcs_pubsub_request_matches_expected_standard_logging_payload( - actual_request, "standard_logging_payload.json" - ) - - -@pytest.mark.asyncio -async def test_async_gcs_pub_sub_v1(): - # Create a mock for the async_httpx_client's post method - litellm.gcs_pub_sub_use_v1 = True - mock_post = AsyncMock() - mock_post.return_value.status_code = 202 - mock_post.return_value.text = "Accepted" - - # Initialize the GcsPubSubLogger and set the mock - gcs_pub_sub_logger = GcsPubSubLogger( - project_id="STUBBED_PROJECT_ID", topic_id="STUBBED_TOPIC_ID", flush_interval=1 - ) - gcs_pub_sub_logger.async_httpx_client.post = mock_post - - mock_construct_request_headers = AsyncMock() - mock_construct_request_headers.return_value = {"Authorization": "Bearer mock_token"} - gcs_pub_sub_logger.construct_request_headers = mock_construct_request_headers - litellm.callbacks = [gcs_pub_sub_logger] - - # Make the completion call - response = await litellm.acompletion( - model="gpt-4o", - messages=[{"role": "user", "content": "Hello, world!"}], - mock_response="hi", - ) - - await asyncio.sleep(3) # Wait for async flush - - # Assert httpx post was called - mock_post.assert_called_once() - - # Get the actual request body from the mock - actual_url = mock_post.call_args[1]["url"] - print("sent to url", actual_url) - assert ( - actual_url - == "https://pubsub.googleapis.com/v1/projects/STUBBED_PROJECT_ID/topics/STUBBED_TOPIC_ID:publish" - ) - actual_request = mock_post.call_args[1]["json"] - - # Extract and decode the base64 encoded message - encoded_message = actual_request["messages"][0]["data"] - import base64 - - decoded_message = base64.b64decode(encoded_message).decode("utf-8") - - # Parse the JSON string into a dictionary - actual_request = json.loads(decoded_message) - print("##########\n") - print(json.dumps(actual_request, indent=4)) - print("##########\n") - # Verify the request body matches expected format - assert_gcs_pubsub_request_matches_expected( - actual_request, "spend_logs_payload.json" - ) diff --git a/tests/logging_callback_tests/test_langfuse_e2e_test.py b/tests/logging_callback_tests/test_langfuse_e2e_test.py deleted file mode 100644 index 8330a18490d..00000000000 --- a/tests/logging_callback_tests/test_langfuse_e2e_test.py +++ /dev/null @@ -1,545 +0,0 @@ -import asyncio -import json -import logging -import os -import threading -from collections.abc import Mapping -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from opentelemetry.proto.common.v1.common_pb2 import AnyValue - -logging.basicConfig(level=logging.DEBUG) - -import litellm -from litellm.integrations.langfuse.langfuse_sdk import resolve_observation_id, resolve_trace_id -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - -litellm.num_retries = 3 -litellm.success_callback = ["langfuse"] -os.environ["LANGFUSE_DEBUG"] = "True" - -import pytest -import pytest_asyncio - -LANGFUSE_EXPORT_POST: Final = "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" -LANGFUSE_EXPORT_PATH: Final = "/api/public/otel/v1/traces" - -_PER_RUN_ATTRIBUTES: Final = frozenset( - { - "langfuse.observation.completion_start_time", - "langfuse.observation.metadata.applied_guardrails", - "langfuse.observation.metadata.cache_hit", - "langfuse.observation.metadata.hidden_params", - "langfuse.observation.metadata.litellm_call_id", - "langfuse.observation.metadata.litellm_response_cost", - "langfuse.observation.metadata.requester_metadata", - "langfuse.observation.metadata.response_id", - "langfuse.observation.metadata.usage_object", - } -) - - -def _decode_attribute(value: AnyValue) -> object: - match value.WhichOneof("value"): - case "string_value": - try: - return json.loads(value.string_value) - except json.JSONDecodeError: - return value.string_value - case "bool_value": - return value.bool_value - case "int_value": - return value.int_value - case "double_value": - return value.double_value - case "array_value": - return [_decode_attribute(item) for item in value.array_value.values] - case _: - return None - - -def _exported_spans(mock_post: MagicMock) -> list[dict[str, object]]: - spans: list[dict[str, object]] = [] - for call in mock_post.call_args_list: - url: str = call.args[0] if call.args else call.kwargs["url"] - assert url.endswith(LANGFUSE_EXPORT_PATH), url - request = ExportTraceServiceRequest.FromString(call.kwargs["data"]) - for resource_spans in request.resource_spans: - for scope_spans in resource_spans.scope_spans: - for span in scope_spans.spans: - spans.append( - { - "name": span.name, - "trace_id": span.trace_id.hex(), - "span_id": span.span_id.hex(), - "parent_span_id": span.parent_span_id.hex() or None, - "attributes": { - attribute.key: _decode_attribute(attribute.value) for attribute in span.attributes - }, - } - ) - return spans - - -def _comparable(span: Mapping[str, object]) -> dict[str, object]: - attributes = span["attributes"] - assert isinstance(attributes, dict) - return { - "name": span["name"], - "parent_span_id": span["parent_span_id"], - "attributes": {key: value for key, value in sorted(attributes.items()) if key not in _PER_RUN_ATTRIBUTES}, - } - - -def assert_langfuse_request_matches_expected( - spans: list[dict[str, object]], - expected_file_name: str, - trace_id: str, -): - """Compare the generation langfuse exported for ``trace_id`` with the expected JSON file.""" - pwd = os.path.dirname(os.path.realpath(__file__)) - expected_body_path = os.path.join(pwd, "langfuse_expected_request_body", expected_file_name) - with open(expected_body_path, "r") as f: - expected_generation = json.load(f) - - otel_trace_id: Final = resolve_trace_id(trace_id) - generations: Final = [ - span - for span in spans - if span["trace_id"] == otel_trace_id and span["attributes"]["langfuse.observation.type"] == "generation" # pyright: ignore[reportIndexIssue] # built as dict in _exported_spans - ] - assert len(generations) == 1, ( - f"Expected exactly one generation for trace_id={trace_id} ({otel_trace_id}), " - f"got {len(generations)}. Spans: {json.dumps(spans, indent=2)}" - ) - - actual_generation: Final = _comparable(generations[0]) - assert actual_generation == expected_generation, ( - f"Difference in exported generation: {json.dumps(actual_generation, indent=2)} " - f"!= {json.dumps(expected_generation, indent=2)}" - ) - - -class TestLangfuseLogging: - @pytest_asyncio.fixture - async def mock_setup(self): - """Common setup for Langfuse logging tests""" - from litellm._uuid import uuid - - mock_post = MagicMock(return_value=MagicMock(ok=True, status_code=200)) - - litellm.set_verbose = True - litellm.success_callback = ["langfuse"] - - return {"trace_id": f"litellm-test-{uuid.uuid4()!s}", "mock_post": mock_post} - - async def _verify_langfuse_call( - self, - mock_post, - expected_file_name: str, - trace_id: str, - ): - """Wait for the batch processor to export, then compare the generation it shipped.""" - otel_trace_id: Final = resolve_trace_id(trace_id) - for _ in range(100): - if any(span["trace_id"] == otel_trace_id for span in _exported_spans(mock_post)): - break - await asyncio.sleep(0.1) - - assert mock_post.call_count >= 1, "langfuse exported nothing" - assert_langfuse_request_matches_expected( - _exported_spans(mock_post), - expected_file_name, - trace_id, - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion(self, mock_setup): - """Test Langfuse logging for chat completion""" - setup = mock_setup - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata={"trace_id": setup["trace_id"]}, - ) - await self._verify_langfuse_call(setup["mock_post"], "completion.json", setup["trace_id"]) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_tags(self, mock_setup): - """Test Langfuse logging for chat completion with tags""" - setup = mock_setup - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata={ - "trace_id": setup["trace_id"], - "tags": ["test_tag", "test_tag_2"], - }, - ) - await self._verify_langfuse_call(setup["mock_post"], "completion_with_tags.json", setup["trace_id"]) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_tags_stream(self, mock_setup): - """Test Langfuse logging for chat completion with tags""" - setup = mock_setup - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata={ - "trace_id": setup["trace_id"], - "tags": ["test_tag_stream", "test_tag_2_stream"], - }, - ) - await self._verify_langfuse_call( - setup["mock_post"], - "completion_with_tags_stream.json", - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_generation_id_metadata_names_the_exported_observation(self, mock_setup): - """v2 let callers pick the generation id; v4 only has span ids, so the requested id must become one.""" - setup = mock_setup - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata={"trace_id": setup["trace_id"], "generation_id": "my-generation"}, - ) - await self._verify_langfuse_call(setup["mock_post"], "completion.json", setup["trace_id"]) - - generation: Final = next( - span - for span in _exported_spans(setup["mock_post"]) - if span["trace_id"] == resolve_trace_id(setup["trace_id"]) - ) - assert generation["span_id"] == resolve_observation_id("my-generation") - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_langfuse_metadata(self, mock_setup): - """Test Langfuse logging for chat completion with metadata for langfuse""" - setup = mock_setup - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata={ - "trace_id": setup["trace_id"], - "tags": ["test_tag", "test_tag_2"], - "generation_name": "test_generation_name", - "parent_observation_id": "test_parent_observation_id", - "version": "test_version", - "trace_user_id": "test_user_id", - "session_id": "test_session_id", - "trace_name": "test_trace_name", - "trace_metadata": {"test_key": "test_value"}, - "trace_version": "test_trace_version", - "trace_release": "test_trace_release", - }, - ) - await self._verify_langfuse_call( - setup["mock_post"], - "completion_with_langfuse_metadata.json", - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_with_non_serializable_metadata(self, mock_setup): - """Test Langfuse logging with metadata that requires preparation (Pydantic models, sets, etc)""" - import datetime - - from pydantic import BaseModel - - class UserPreferences(BaseModel): - favorite_colors: set[str] - last_login: datetime.datetime - settings: dict - - setup = mock_setup - - test_metadata = { - "user_prefs": UserPreferences( - favorite_colors={"red", "blue"}, - last_login=datetime.datetime.now(), - settings={"theme": "dark", "notifications": True}, - ), - "nested_set": { - "inner_set": {1, 2, 3}, - "inner_pydantic": UserPreferences( - favorite_colors={"green", "yellow"}, - last_login=datetime.datetime.now(), - settings={"theme": "light"}, - ), - }, - "trace_id": setup["trace_id"], - } - - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata=test_metadata, - ) - - await self._verify_langfuse_call( - setup["mock_post"], - "completion_with_complex_metadata.json", - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "test_metadata, response_json_file", - [ - ({"a": 1, "b": 2, "c": 3}, "simple_metadata.json"), - ( - {"a": {"nested_a": 1}, "b": {"nested_b": 2}}, - "nested_metadata.json", - ), - ({"a": [1, 2, 3], "b": {4, 5, 6}}, "simple_metadata2.json"), - ( - {"a": (1, 2), "b": frozenset([3, 4]), "c": {"d": [5, 6]}}, - "simple_metadata3.json", - ), - ({"lock": threading.Lock()}, "metadata_with_lock.json"), - ({"func": lambda x: x + 1}, "metadata_with_function.json"), - ( - { - "int": 42, - "str": "hello", - "list": [1, 2, 3], - "set": {4, 5}, - "dict": {"nested": "value"}, - "non_copyable": threading.Lock(), - "function": print, - }, - "complex_metadata.json", - ), - ( - {"list": ["list", "not", "a", "dict"]}, - "complex_metadata_2.json", - ), - ({}, "empty_metadata.json"), - ], - ) - @pytest.mark.flaky(retries=6, delay=1) - async def test_langfuse_logging_with_various_metadata_types(self, mock_setup, test_metadata, response_json_file): - """Test Langfuse logging with various metadata types including non-serializable objects""" - setup = mock_setup - - if test_metadata is not None: - test_metadata["trace_id"] = setup["trace_id"] - - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response="Hello! How can I assist you today?", - metadata=test_metadata, - ) - - await self._verify_langfuse_call( - setup["mock_post"], - response_json_file, - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_malformed_llm_response(self, mock_setup): - """Test Langfuse logging for chat completion with malformed LLM response""" - setup = mock_setup - litellm.turn_on_debug() - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - mock_response = litellm.ModelResponse( - choices=[], - usage=litellm.Usage( - prompt_tokens=10, - completion_tokens=10, - total_tokens=20, - ), - model="gpt-3.5-turbo", - object="chat.completion", - created=1723081200, - ).model_dump() - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response=mock_response, - metadata={"trace_id": setup["trace_id"]}, - ) - await self._verify_langfuse_call(setup["mock_post"], "completion_with_no_choices.json", setup["trace_id"]) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_bedrock_llm_response(self, mock_setup): - """Test Langfuse logging for chat completion with malformed LLM response""" - setup = mock_setup - litellm.turn_on_debug() - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - mock_response = litellm.ModelResponse( - choices=[], - usage=litellm.Usage( - prompt_tokens=10, - completion_tokens=10, - total_tokens=20, - ), - model="anthropic.claude-haiku-4-5-20251001-v1:0", - object="chat.completion", - created=1723081200, - ).model_dump() - await litellm.acompletion( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[{"role": "user", "content": "Hello!"}], - mock_response=mock_response, - metadata={"trace_id": setup["trace_id"]}, - aws_access_key_id="fake-key", - aws_secret_access_key="fake-key", - aws_region="us-east-1", - ) - await self._verify_langfuse_call( - setup["mock_post"], - "completion_with_bedrock_call.json", - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_completion_with_vertex_llm_response(self, mock_setup): - """Test Langfuse logging for chat completion with malformed LLM response""" - setup = mock_setup - litellm.turn_on_debug() - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - mock_response = litellm.ModelResponse( - choices=[], - usage=litellm.Usage( - prompt_tokens=10, - completion_tokens=10, - total_tokens=20, - ), - model="vertex/gemini-3-flash-preview", - object="chat.completion", - created=1723081200, - ).model_dump() - await litellm.acompletion( - model="vertex_ai/gemini-3-flash-preview", - messages=[{"role": "user", "content": "Hello!"}], - mock_response=mock_response, - metadata={"trace_id": setup["trace_id"]}, - vertex_credentials="my-mock-credentials", - api_key="my-mock-credentials-2", - ) - await self._verify_langfuse_call( - setup["mock_post"], - "completion_with_vertex_call.json", - setup["trace_id"], - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_vllm_embedding(self, mock_setup): - """ - Test that the request sent to the vllm embedding endpoint is correct. - - Verifies the request body matches the expected JSON fixture, - including that the hosted_vllm/ prefix is stripped from the model name - and that no unexpected fields (e.g. encoding_format) are included. - """ - setup = mock_setup - - vllm_response_data = { - "object": "list", - "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], - "model": "BAAI/bge-small-en-v1.5", - "usage": {"prompt_tokens": 10, "total_tokens": 10}, - } - mock_vllm_response = httpx.Response( - status_code=200, - json=vllm_response_data, - ) - - mock_async_client = AsyncHTTPHandler() - mock_async_client.post = AsyncMock(return_value=mock_vllm_response) - - with patch(LANGFUSE_EXPORT_POST, setup["mock_post"]): - await litellm.aembedding( - model="hosted_vllm/BAAI/bge-small-en-v1.5", - input=["Hello from litellm!"], - api_base="http://my-fake-vllm.com/v1", - metadata={"trace_id": setup["trace_id"]}, - client=mock_async_client, - ) - - # Verify the request sent to vllm matches the expected JSON fixture - assert mock_async_client.post.call_count == 1 - actual_vllm_request = mock_async_client.post.call_args.kwargs["json"] - - pwd = os.path.dirname(os.path.realpath(__file__)) - expected_body_path = os.path.join(pwd, "langfuse_expected_request_body", "embedding_with_vllm.json") - with open(expected_body_path, "r") as f: - expected_vllm_request = json.load(f) - - assert actual_vllm_request == expected_vllm_request, ( - f"vllm request body mismatch:\n" - f"actual: {json.dumps(actual_vllm_request, indent=2)}\n" - f"expected: {json.dumps(expected_vllm_request, indent=2)}" - ) - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=1) - async def test_langfuse_logging_with_router(self, mock_setup): - """Test Langfuse logging with router""" - litellm.turn_on_debug() - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "mock_response": "Hello! How can I assist you today?", - "api_key": "test_api_key", - }, - } - ] - ) - with patch(LANGFUSE_EXPORT_POST, mock_setup["mock_post"]): - mock_response = litellm.ModelResponse( - choices=[], - usage=litellm.Usage( - prompt_tokens=10, - completion_tokens=10, - total_tokens=20, - ), - model="gpt-3.5-turbo", - object="chat.completion", - created=1723081200, - ).model_dump() - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello!"}], - mock_response=mock_response, - metadata={"trace_id": mock_setup["trace_id"]}, - ) - await self._verify_langfuse_call( - mock_setup["mock_post"], - "completion_with_router.json", - mock_setup["trace_id"], - ) diff --git a/tests/logging_callback_tests/test_langfuse_unit_tests.py b/tests/logging_callback_tests/test_langfuse_unit_tests.py deleted file mode 100644 index 1a834a37c46..00000000000 --- a/tests/logging_callback_tests/test_langfuse_unit_tests.py +++ /dev/null @@ -1,61 +0,0 @@ -from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache - - -def test_langfuse_e2e_sync(monkeypatch): - """A sync completion must reach langfuse over the wire, not just build a span. - - v4 exports OTLP over ``requests`` rather than the v2 ingestion endpoint over - httpx, so this stands up a real receiver and asserts langfuse posted to it. - """ - import threading - import time - from http.server import BaseHTTPRequestHandler, HTTPServer - - import litellm - from litellm import completion - from litellm.integrations.langfuse.langfuse import LangFuseLogger - from litellm.integrations.langfuse.langfuse_prompt_management import langfuse_client_init - from litellm.litellm_core_utils import litellm_logging - - received_paths = [] - - class _Receiver(BaseHTTPRequestHandler): - def do_POST(self): - received_paths.append(self.path) - self.rfile.read(int(self.headers.get("Content-Length") or 0)) - self.send_response(200) - self.send_header("Content-Length", "0") - self.end_headers() - - def log_message(self, *args): - pass - - server = HTTPServer(("127.0.0.1", 0), _Receiver) - threading.Thread(target=server.serve_forever, daemon=True).start() - monkeypatch.setenv("LANGFUSE_HOST", f"http://127.0.0.1:{server.server_port}") - monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-e2e-sync") - monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-e2e-sync") - monkeypatch.setattr(litellm, "success_callback", ["langfuse"]) - monkeypatch.setattr(litellm_logging, "langFuseLogger", None) - monkeypatch.setattr(litellm_logging, "in_memory_dynamic_logger_cache", DynamicLoggingCache()) - monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) - langfuse_client_init.cache_clear() - - try: - completion( - model="openai/my-fake-endpoint", - messages=[{"role": "user", "content": "hello from litellm"}], - stream=False, - mock_response="Hello from litellm 2", - ) - for logger in litellm.logging_callback_manager.get_all_callbacks(): - if isinstance(logger, LangFuseLogger): - logger.flush() - deadline = time.time() + 10 - while not received_paths and time.time() < deadline: - time.sleep(0.1) - finally: - server.shutdown() - - assert received_paths, "langfuse exported nothing" - assert all(path.endswith("/api/public/otel/v1/traces") for path in received_paths) diff --git a/tests/logging_callback_tests/test_log_db_redis_services.py b/tests/logging_callback_tests/test_log_db_redis_services.py deleted file mode 100644 index aad404a5525..00000000000 --- a/tests/logging_callback_tests/test_log_db_redis_services.py +++ /dev/null @@ -1,33 +0,0 @@ -import io - - -import gzip -import json -import logging -from unittest.mock import AsyncMock, patch - -import pytest - -import litellm -from litellm import completion -from litellm._logging import verbose_logger -from litellm.proxy.utils import ServiceTypes - - -@pytest.mark.asyncio -async def test_dd_log_db_spend_failure_metrics(): - from litellm._service_logger import ServiceLogging - from litellm.integrations.datadog.datadog import DataDogLogger - - dd_logger = DataDogLogger() - with patch.object(dd_logger, "async_service_failure_hook", new_callable=AsyncMock): - service_logging_obj = ServiceLogging() - - litellm.service_callback = [dd_logger] - - await service_logging_obj.async_service_failure_hook( - service=ServiceTypes.DB, - call_type="test_call_type", - error="test_error", - duration=1.0, - ) diff --git a/tests/logging_callback_tests/test_pagerduty_alerting.py b/tests/logging_callback_tests/test_pagerduty_alerting.py deleted file mode 100644 index 1426dc32081..00000000000 --- a/tests/logging_callback_tests/test_pagerduty_alerting.py +++ /dev/null @@ -1,97 +0,0 @@ -import asyncio -import random -from datetime import datetime, timedelta -from typing import Optional - -import pytest -import litellm - -from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( - PagerDutyAlerting, - AlertingConfig, -) -from litellm.proxy._types import UserAPIKeyAuth - - -@pytest.mark.asyncio -async def test_pagerduty_alerting(): - pagerduty = PagerDutyAlerting( - alerting_args=AlertingConfig( - failure_threshold=1, failure_threshold_window_seconds=10 - ) - ) - litellm.callbacks = [pagerduty] - - try: - await litellm.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - mock_response="litellm.RateLimitError", - ) - except litellm.RateLimitError: - pass - - await asyncio.sleep(2) - - -@pytest.mark.asyncio -async def test_pagerduty_alerting_high_failure_rate(): - pagerduty = PagerDutyAlerting( - alerting_args=AlertingConfig( - failure_threshold=3, failure_threshold_window_seconds=600 - ) - ) - litellm.callbacks = [pagerduty] - - try: - await litellm.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - mock_response="litellm.RateLimitError", - ) - except litellm.RateLimitError: - pass - - await asyncio.sleep(2) - - # make 3 more fails - for _ in range(3): - try: - await litellm.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - mock_response="litellm.RateLimitError", - ) - except litellm.RateLimitError: - pass - - await asyncio.sleep(2) - - -@pytest.mark.asyncio -async def test_pagerduty_hanging_request_alerting(): - pagerduty = PagerDutyAlerting( - alerting_args=AlertingConfig(hanging_threshold_seconds=0.0000001) - ) - litellm.callbacks = [pagerduty] - - await pagerduty.async_pre_call_hook( - cache=None, - user_api_key_dict=UserAPIKeyAuth( - api_key="test", - key_alias="test-pagerduty", - team_alias="test-team", - org_id="test-org", - user_id="test-user", - end_user_id="test-end-user", - ), - data={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, - call_type="completion", - ) - - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "hi"}], - ) - - await asyncio.sleep(1) diff --git a/tests/logging_callback_tests/test_sqs_logger.py b/tests/logging_callback_tests/test_sqs_logger.py deleted file mode 100644 index f7b8b5e692f..00000000000 --- a/tests/logging_callback_tests/test_sqs_logger.py +++ /dev/null @@ -1,181 +0,0 @@ -import asyncio -import base64 -import json -import os -from copy import deepcopy -from unittest.mock import AsyncMock, MagicMock, patch -from urllib.parse import unquote - -import litellm -import pytest - -from litellm.integrations.sqs import SQSLogger -from litellm.types.utils import StandardLoggingPayload - -from litellm.litellm_core_utils.app_crypto import AppCrypto - - -@pytest.mark.asyncio -async def test_async_sqs_logger_flush(): - expected_queue_url = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" - expected_region = "us-east-1" - - sqs_logger = SQSLogger( - sqs_queue_url=expected_queue_url, - sqs_region_name=expected_region, - sqs_flush_interval=1, - ) - - # Mock the httpx client - mock_response = MagicMock() - mock_response.raise_for_status = MagicMock() - sqs_logger.async_httpx_client.post = AsyncMock(return_value=mock_response) - - litellm.callbacks = [sqs_logger] - - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "hello"}], - mock_response="hi", - ) - - await asyncio.sleep(2) - - # Verify that httpx post was called - sqs_logger.async_httpx_client.post.assert_called() - - # Get the call arguments - call_args = sqs_logger.async_httpx_client.post.call_args - - # Verify the URL is correct - called_url = call_args[0][0] # First positional argument - assert ( - called_url == expected_queue_url - ), f"Expected URL {expected_queue_url}, got {called_url}" - - # Verify the payload contains StandardLoggingPayload data - called_data = call_args.kwargs["data"] - - # Extract the MessageBody from the URL-encoded data - # Format: "Action=SendMessage&Version=2012-11-05&MessageBody=" - assert "Action=SendMessage" in called_data - assert "Version=2012-11-05" in called_data - assert "MessageBody=" in called_data - - # Extract and decode the message body - message_body_start = called_data.find("MessageBody=") + len("MessageBody=") - message_body_encoded = called_data[message_body_start:] - message_body_json = unquote(message_body_encoded) - - # Parse the JSON to verify it's a StandardLoggingPayload - payload_data = json.loads(message_body_json) - - # Verify it has the expected StandardLoggingPayload structure - assert "model" in payload_data - assert "messages" in payload_data - assert "response" in payload_data - assert payload_data["model"] == "gpt-5.5" - assert len(payload_data["messages"]) == 1 - assert payload_data["messages"][0]["role"] == "user" - assert payload_data["messages"][0]["content"] == "hello" - - -@pytest.mark.asyncio -async def test_async_sqs_logger_error_flush(): - expected_queue_url = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" - expected_region = "us-east-1" - - sqs_logger = SQSLogger( - sqs_queue_url=expected_queue_url, - sqs_region_name=expected_region, - sqs_flush_interval=1, - ) - - # Mock the httpx client - mock_response = MagicMock() - mock_response.raise_for_status = Exception("Something went wrong") - sqs_logger.async_httpx_client.post = AsyncMock(return_value=mock_response) - - litellm.callbacks = [sqs_logger] - - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "hello"}], - mock_response="Error occurred", - ) - - await asyncio.sleep(2) - - # Verify that httpx post was called - sqs_logger.async_httpx_client.post.assert_called() - - # Get the call arguments - call_args = sqs_logger.async_httpx_client.post.call_args - - # Verify the URL is correct - called_url = call_args[0][0] # First positional argument - assert ( - called_url == expected_queue_url - ), f"Expected URL {expected_queue_url}, got {called_url}" - - # Verify the payload contains StandardLoggingPayload data - called_data = call_args.kwargs["data"] - - # Extract the MessageBody from the URL-encoded data - # Format: "Action=SendMessage&Version=2012-11-05&MessageBody=" - assert "Action=SendMessage" in called_data - assert "Version=2012-11-05" in called_data - assert "MessageBody=" in called_data - - # Extract and decode the message body - message_body_start = called_data.find("MessageBody=") + len("MessageBody=") - message_body_encoded = called_data[message_body_start:] - message_body_json = unquote(message_body_encoded) - - # Parse the JSON to verify it's a StandardLoggingPayload - payload_data = json.loads(message_body_json) - - # Verify it has the expected StandardLoggingPayload structure - assert "model" in payload_data - assert "messages" in payload_data - assert "response" in payload_data - assert payload_data["model"] == "gpt-5.5" - assert len(payload_data["messages"]) == 1 - assert payload_data["messages"][0]["role"] == "user" - assert payload_data["messages"][0]["content"] == "hello" - - -@pytest.mark.asyncio -async def test_async_send_batch_triggers_tasks(monkeypatch): - monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) - monkeypatch.setattr(asyncio, "create_task", MagicMock()) - logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2") - - logger.async_send_message = AsyncMock() - logger.log_queue = [{"log": 1}, {"log": 2}] - - await logger.async_send_batch() - # It uses asyncio.create_task() so direct await count = 0 is expected - asyncio.create_task.assert_called() - - -@pytest.mark.asyncio -async def test_strip_base64_keeps_non_file_content(): - logger = SQSLogger(sqs_strip_base64_files=True) - - payload = { - "messages": [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Just text"}, - {"type": "text", "text": "Another message"}, - ], - } - ] - } - - stripped = await logger._strip_base64_from_messages(payload) - - # Should not modify normal text messages - assert stripped["messages"][0]["content"] == payload["messages"][0]["content"] diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py deleted file mode 100644 index 41b351631f6..00000000000 --- a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py +++ /dev/null @@ -1,390 +0,0 @@ -""" -Base test class for Anthropic Messages API prompt caching E2E tests. - -Tests that prompt caching works correctly via litellm.anthropic.messages interface -by making actual API calls and validating usage metrics. - -Per AWS docs (https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html): -- Converse API uses: cachePoint: { type: "default" } -- InvokeModel API uses: cache_control: { type: "ephemeral" } -- Claude 3.7 Sonnet: GA, 1024 min tokens -- Claude 3.5 Haiku: GA, 2048 min tokens -""" - -import json -import os -import sys -from abc import ABC, abstractmethod -from typing import Any, Dict, List - -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) - -import pytest -import litellm - -from tests._live_test_helpers import _skip_live_prompt_caching_test - -# Large document for caching tests (needs 1024+ tokens for Claude models) -LARGE_DOCUMENT_FOR_CACHING = """ -This is a comprehensive legal agreement between Party A and Party B. - -ARTICLE 1: DEFINITIONS -1.1 "Agreement" means this document and all attachments. -1.2 "Confidential Information" means any non-public information. -1.3 "Effective Date" means the date of last signature. -1.4 "Term" means the period during which this Agreement is in effect. - -ARTICLE 2: SCOPE OF SERVICES -2.1 Party A agrees to provide the following services... -2.2 Party B agrees to compensate Party A for services rendered... -2.3 All services shall be performed in a professional manner... - -ARTICLE 3: PAYMENT TERMS -3.1 Payment shall be made within 30 days of invoice receipt. -3.2 Late payments shall accrue interest at 1.5% per month. -3.3 All fees are non-refundable unless otherwise specified. - -ARTICLE 4: INTELLECTUAL PROPERTY -4.1 All pre-existing IP remains with the original owner. -4.2 Work product created under this Agreement shall be owned by Party B. -4.3 Party A grants a license to use any tools or methodologies. - -ARTICLE 5: CONFIDENTIALITY -5.1 Both parties agree to maintain confidentiality of all shared information. -5.2 Confidential information shall not be disclosed to third parties. -5.3 This obligation survives termination of the Agreement. - -ARTICLE 6: TERMINATION -6.1 Either party may terminate with 30 days written notice. -6.2 Immediate termination is permitted for material breach. -6.3 Upon termination, all confidential information must be returned. - -ARTICLE 7: LIMITATION OF LIABILITY -7.1 Neither party shall be liable for consequential damages. -7.2 Total liability shall not exceed fees paid in the prior 12 months. -7.3 This limitation does not apply to willful misconduct. - -ARTICLE 8: DISPUTE RESOLUTION -8.1 Disputes shall first be addressed through good faith negotiation. -8.2 If negotiation fails, disputes shall be submitted to arbitration. -8.3 Arbitration shall be conducted under AAA rules. - -ARTICLE 9: GENERAL PROVISIONS -9.1 This Agreement constitutes the entire understanding between parties. -9.2 Amendments must be in writing and signed by both parties. -9.3 This Agreement shall be governed by the laws of Delaware. -9.4 Neither party may assign this Agreement without consent. -9.5 Waiver of any provision shall not constitute ongoing waiver. - -IN WITNESS WHEREOF, the parties have executed this Agreement. -""" * 8 # Repeat to ensure we have enough tokens (need 1024+ for Claude models) - - -class BaseAnthropicMessagesPromptCachingTest(ABC): - """ - Base test class for prompt caching E2E tests across different providers. - - Subclasses must implement: - - get_model(): Returns the model string to use for tests - """ - - @abstractmethod - def get_model(self) -> str: - """ - Returns the model string to use for tests. - - Examples: - - "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" - - "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0" - """ - pass - - def get_messages_with_cache_control(self) -> List[Dict[str, Any]]: - """ - Returns test messages with cache_control set on content blocks. - """ - return [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": LARGE_DOCUMENT_FOR_CACHING, - "cache_control": {"type": "ephemeral"}, - }, - { - "type": "text", - "text": "What are the payment terms in this agreement?", - }, - ], - }, - ] - - @pytest.mark.asyncio - async def test_prompt_caching_returns_cache_creation_tokens(self): - """ - E2E test: First call should return cache_creation_input_tokens > 0. - - This validates that the cache_control field is being passed through - correctly and the provider is creating a cache. - """ - _skip_live_prompt_caching_test() - litellm.turn_on_debug() - - messages = self.get_messages_with_cache_control() - - response = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - ) - - print(f"Response: {json.dumps(response, indent=2, default=str)}") - - # Validate response structure - assert "usage" in response, "Response should contain usage" - usage = response["usage"] - - # Check for cache tokens in usage - cache_creation = usage.get("cache_creation_input_tokens", 0) - cache_read = usage.get("cache_read_input_tokens", 0) - - print(f"cache_creation_input_tokens: {cache_creation}") - print(f"cache_read_input_tokens: {cache_read}") - - # First call should create cache (cache_creation > 0) OR read from existing cache - assert cache_creation > 0 or cache_read > 0, ( - f"Expected cache_creation_input_tokens > 0 or cache_read_input_tokens > 0, " - f"but got cache_creation={cache_creation}, cache_read={cache_read}. " - f"This indicates cache_control is not being passed through correctly." - ) - - @pytest.mark.asyncio - async def test_prompt_caching_returns_cache_read_tokens_on_second_call(self): - """ - E2E test: Second call with same content should return cache_read_input_tokens > 0. - - This validates that caching is working end-to-end. - """ - _skip_live_prompt_caching_test() - litellm.turn_on_debug() - - messages = self.get_messages_with_cache_control() - - # First call - creates cache - response1 = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - ) - - print( - f"First response usage: {json.dumps(response1.get('usage', {}), indent=2)}" - ) - - # Second call - should read from cache - response2 = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - ) - - print( - f"Second response usage: {json.dumps(response2.get('usage', {}), indent=2)}" - ) - - usage = response2.get("usage", {}) - cache_read = usage.get("cache_read_input_tokens", 0) - - # Second call should read from cache - assert cache_read > 0, ( - f"Expected cache_read_input_tokens > 0 on second call, " - f"but got {cache_read}. Full usage: {usage}" - ) - - def _parse_sse_chunks(self, chunk: bytes) -> list: - """ - Parse SSE format chunks and return list of JSON objects. - """ - results = [] - chunk_str = chunk.decode("utf-8") - for line in chunk_str.split("\n"): - if line.startswith("data: "): - try: - json_data = json.loads(line[6:]) # Skip the 'data: ' prefix - results.append(json_data) - except json.JSONDecodeError: - pass - return results - - @pytest.mark.asyncio - async def test_prompt_caching_streaming_returns_cache_tokens(self): - """ - E2E test: Streaming response should include cache tokens in usage. - - This validates that cache_creation_input_tokens and cache_read_input_tokens - are correctly returned in the streaming response's message_delta event. - """ - _skip_live_prompt_caching_test() - litellm.turn_on_debug() - - messages = self.get_messages_with_cache_control() - - response = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - stream=True, - ) - - # Collect all chunks and find the message_delta with usage - cache_creation = 0 - cache_read = 0 - found_usage = False - - async for chunk in response: - # Handle SSE format chunks (bytes) - if isinstance(chunk, bytes): - json_chunks = self._parse_sse_chunks(chunk) - for json_data in json_chunks: - print( - f"Parsed chunk: {json.dumps(json_data, indent=2, default=str)}" - ) - - # Look for message_delta with usage (final chunk) - if json_data.get("type") == "message_delta": - usage = json_data.get("usage", {}) - if usage: - found_usage = True - cache_creation = max( - cache_creation, - usage.get("cache_creation_input_tokens", 0), - ) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - print( - f"Found usage in message_delta: cache_creation={cache_creation}, cache_read={cache_read}" - ) - - # Also check message_start for usage (Anthropic includes it there too) - if json_data.get("type") == "message_start": - message = json_data.get("message", {}) - usage = message.get("usage", {}) - if usage: - found_usage = True - cache_creation = max( - cache_creation, - usage.get("cache_creation_input_tokens", 0), - ) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - print( - f"Found usage in message_start: cache_creation={cache_creation}, cache_read={cache_read}" - ) - elif isinstance(chunk, dict): - print(f"Dict chunk: {json.dumps(chunk, indent=2, default=str)}") - # Handle dict chunks directly - if chunk.get("type") == "message_delta": - usage = chunk.get("usage", {}) - if usage: - found_usage = True - cache_creation = max( - cache_creation, usage.get("cache_creation_input_tokens", 0) - ) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - - if chunk.get("type") == "message_start": - message = chunk.get("message", {}) - usage = message.get("usage", {}) - if usage: - found_usage = True - cache_creation = max( - cache_creation, usage.get("cache_creation_input_tokens", 0) - ) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - - assert found_usage, "Expected to find usage in streaming response" - - # Should have cache tokens (either creation or read) - assert cache_creation > 0 or cache_read > 0, ( - f"Expected cache_creation_input_tokens > 0 or cache_read_input_tokens > 0 in streaming response, " - f"but got cache_creation={cache_creation}, cache_read={cache_read}. " - f"This indicates cache tokens are not being passed through in streaming mode." - ) - - @pytest.mark.asyncio - async def test_prompt_caching_streaming_second_call_returns_cache_read(self): - """ - E2E test: Second streaming call should return cache_read_input_tokens > 0. - """ - _skip_live_prompt_caching_test() - litellm.turn_on_debug() - - messages = self.get_messages_with_cache_control() - - # First call - creates cache - response1 = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - stream=True, - ) - - # Consume the first stream - async for chunk in response1: - pass - - # Second call - should read from cache - response2 = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - max_tokens=100, - stream=True, - ) - - cache_read = 0 - async for chunk in response2: - # Handle SSE format chunks (bytes) - if isinstance(chunk, bytes): - json_chunks = self._parse_sse_chunks(chunk) - for json_data in json_chunks: - print( - f"Second call parsed chunk: {json.dumps(json_data, indent=2, default=str)}" - ) - - if json_data.get("type") == "message_delta": - usage = json_data.get("usage", {}) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - - if json_data.get("type") == "message_start": - message = json_data.get("message", {}) - usage = message.get("usage", {}) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - elif isinstance(chunk, dict): - if chunk.get("type") == "message_delta": - usage = chunk.get("usage", {}) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - - if chunk.get("type") == "message_start": - message = chunk.get("message", {}) - usage = message.get("usage", {}) - cache_read = max( - cache_read, usage.get("cache_read_input_tokens", 0) - ) - - assert cache_read > 0, ( - f"Expected cache_read_input_tokens > 0 on second streaming call, " - f"but got {cache_read}" - ) diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py deleted file mode 100644 index 966048e609e..00000000000 --- a/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py +++ /dev/null @@ -1,194 +0,0 @@ -""" -Base test class for Anthropic Messages API tool search E2E tests. - -Tests that tool search works correctly via litellm.anthropic.messages interface -by making actual API calls and validating that tool search discovers deferred tools. - -Reference: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool -""" - -import json -from abc import ABC, abstractmethod -from typing import Any, Dict, List - - -import pytest -import litellm - - -# Sample tools for tool search testing -def get_deferred_tools() -> List[Dict[str, Any]]: - """ - Returns a list of tools with defer_loading: true. - These tools should only be discovered via tool search. - """ - return [ - { - "name": "get_weather", - "description": "Get the current weather for a location", - "input_schema": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - "defer_loading": True, - }, - { - "name": "get_stock_price", - "description": "Get the current stock price for a ticker symbol", - "input_schema": { - "type": "object", - "properties": { - "ticker": { - "type": "string", - "description": "The stock ticker symbol, e.g. AAPL", - } - }, - "required": ["ticker"], - }, - "defer_loading": True, - }, - { - "name": "search_web", - "description": "Search the web for information", - "input_schema": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "The search query"} - }, - "required": ["query"], - }, - "defer_loading": True, - }, - ] - - -def get_tool_search_tool_regex() -> Dict[str, Any]: - """Returns the tool search tool using regex variant.""" - return {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"} - - -def get_tool_search_tool_bm25() -> Dict[str, Any]: - """Returns the tool search tool using BM25 variant.""" - return {"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"} - - -class BaseAnthropicMessagesToolSearchTest(ABC): - """ - Base test class for tool search E2E tests across different providers. - - Subclasses must implement: - - get_model(): Returns the model string to use for tests - - Tests pass the anthropic-beta header via extra_headers to validate - that the header is correctly forwarded to downstream providers. - """ - - @abstractmethod - def get_model(self) -> str: - """ - Returns the model string to use for tests. - - Examples: - - "anthropic/claude-sonnet-4-5-20250929" - - "vertex_ai/claude-sonnet-4-5@20250929" - - "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0" - """ - pass - - def get_extra_headers(self) -> Dict[str, str]: - """ - Returns extra headers to pass with the request. - Includes the anthropic-beta header for tool search. - - This is what claude code forwards, simulate the same behavior here. - """ - return {"anthropic-beta": "advanced-tool-use-2025-11-20"} - - def get_tools_with_tool_search(self) -> List[Dict[str, Any]]: - """ - Returns tools list with tool search tool and deferred tools. - """ - return [get_tool_search_tool_regex()] + get_deferred_tools() - - @pytest.mark.asyncio - async def test_tool_search_basic_request(self): - """ - E2E test: Basic tool search request should succeed. - - This validates that the tool search beta header is being passed via - extra_headers and forwarded correctly to the downstream provider. - """ - litellm.turn_on_debug() - - tools = self.get_tools_with_tool_search() - messages = [{"role": "user", "content": "What's the weather in San Francisco?"}] - - response = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - tools=tools, - max_tokens=1024, - extra_headers=self.get_extra_headers(), - ) - - print(f"Response: {json.dumps(response, indent=2, default=str)}") - - # Validate response structure - assert "content" in response, "Response should contain content" - assert "usage" in response, "Response should contain usage" - - # The model should either respond with text or use a tool - content = response.get("content", []) - assert len(content) > 0, "Response should have content" - - @pytest.mark.asyncio - @pytest.mark.flaky(retries=3, delay=5) - async def test_tool_search_streaming(self): - """ - E2E test: Tool search should work with streaming responses. - """ - litellm.turn_on_debug() - - tools = self.get_tools_with_tool_search() - messages = [{"role": "user", "content": "What's the weather like in Tokyo?"}] - - response = await litellm.anthropic.messages.acreate( - model=self.get_model(), - messages=messages, - tools=tools, - max_tokens=1024, - stream=True, - extra_headers=self.get_extra_headers(), - ) - - # Collect all chunks - chunks = [] - async for chunk in response: - if isinstance(chunk, bytes): - chunk_str = chunk.decode("utf-8") - for line in chunk_str.split("\n"): - if line.startswith("data: "): - try: - json_data = json.loads(line[6:]) - chunks.append(json_data) - print( - f"Chunk: {json.dumps(json_data, indent=2, default=str)}" - ) - except json.JSONDecodeError: - pass - elif isinstance(chunk, dict): - chunks.append(chunk) - print(f"Chunk: {json.dumps(chunk, indent=2, default=str)}") - - # Should have received chunks - assert len(chunks) > 0, "Expected to receive streaming chunks" - - # Should have message_start - message_starts = [c for c in chunks if c.get("type") == "message_start"] - assert len(message_starts) > 0, "Expected message_start in streaming response" diff --git a/tests/pass_through_unit_tests/base_anthropic_unified_messages_test.py b/tests/pass_through_unit_tests/base_anthropic_unified_messages_test.py deleted file mode 100644 index 858a6713f7b..00000000000 --- a/tests/pass_through_unit_tests/base_anthropic_unified_messages_test.py +++ /dev/null @@ -1,64 +0,0 @@ -import json -from datetime import datetime -from typing import AsyncIterator, Dict, Any -import asyncio -import unittest.mock -from unittest.mock import AsyncMock, MagicMock - -import litellm -import pytest -from dotenv import load_dotenv -from litellm.llms.anthropic.pass_through.messages.handler import ( - anthropic_messages, -) - -from typing import Optional -from litellm.types.utils import StandardLoggingPayload -from litellm.integrations.custom_logger import CustomLogger -from litellm.router import Router -import importlib -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - -class TestCustomLogger(CustomLogger): - def __init__(self): - super().__init__() - self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print("inside async_log_success_event") - self.logged_standard_logging_payload = kwargs.get("standard_logging_object") - - pass - - -class BaseAnthropicMessagesTest: - """Base class for anthropic messages tests to reduce code duplication""" - - @property - def model_config(self) -> Dict[str, Any]: - """Override in subclasses to provide model-specific configuration""" - raise NotImplementedError("Subclasses must implement model_config") - - @property - def expected_model_name_in_logging(self) -> str: - """ - This is the model name that is expected to be in the logging payload - """ - raise NotImplementedError( - "Subclasses must implement expected_model_name_in_logging" - ) - - def _validate_response(self, response: Any): - """Validate non-streaming response structure""" - # Handle type checking - response should be a dict for non-streaming - if isinstance(response, AsyncIterator): - pytest.fail("Expected non-streaming response but got AsyncIterator") - - assert isinstance( - response, dict - ), f"Expected dict response, got {type(response)}" - assert "id" in response - assert "content" in response - assert "model" in response - assert response.get("role") == "assistant" diff --git a/tests/pass_through_unit_tests/conftest.py b/tests/pass_through_unit_tests/conftest.py deleted file mode 100644 index df8196f1785..00000000000 --- a/tests/pass_through_unit_tests/conftest.py +++ /dev/null @@ -1,83 +0,0 @@ -import asyncio - -import pytest - -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - -from tests._vcr_conftest_common import ( # noqa: E402,F401 - VerboseReporterState, - _pin_multipart_boundary, - apply_vcr_auto_marker_to_items, - emit_cassette_cache_session_banner, - emit_vcr_classification_summary, - emit_vcr_diagnostic_log, - install_live_call_probe, - record_vcr_outcome, - register_persister_if_enabled, - reset_vcr_diag_dir, - vcr_config_dict, -) - -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () - - -_verbose_state = VerboseReporterState() - - -@pytest.fixture(scope="module") -def vcr_config(): - return vcr_config_dict() - - -def pytest_recording_configure(config, vcr): - register_persister_if_enabled(vcr) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_makereport(item, call): - outcome = yield - rep = outcome.get_result() - setattr(item, f"rep_{rep.when}", rep) - - -@pytest.fixture(autouse=True) -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -@pytest.fixture(autouse=True) -async def _drain_logging_worker(): - """ - The logging queue is bound to the running loop, so anything left queued when a test's loop - goes away is carried onto the next loop and fires against that test's callbacks. - """ - GLOBAL_LOGGING_WORKER.start() - try: - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) - except asyncio.TimeoutError: - pass - await GLOBAL_LOGGING_WORKER.stop() - yield - - -def pytest_configure(config): - _verbose_state.remember_pluginmanager(config) - reset_vcr_diag_dir() - - -def pytest_runtest_logreport(report): - _verbose_state.maybe_emit_verdict(report) - - -def pytest_collection_modifyitems(config, items): - apply_vcr_auto_marker_to_items( - items, skip_nodeid_suffixes=_VCR_INCOMPATIBLE_NODEID_SUFFIXES - ) - - -def pytest_terminal_summary(terminalreporter, exitstatus, config): - emit_cassette_cache_session_banner(terminalreporter) - emit_vcr_classification_summary(terminalreporter) - emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py deleted file mode 100644 index f0d03ca77d6..00000000000 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ /dev/null @@ -1,80 +0,0 @@ -import os -from datetime import datetime -from typing import Dict, Any -import unittest.mock -from unittest.mock import MagicMock - -from dotenv import load_dotenv -from litellm.llms.anthropic.pass_through.messages.handler import ( - anthropic_messages, -) - -from typing import Optional -from litellm.types.utils import StandardLoggingPayload -from litellm.integrations.custom_logger import CustomLogger -from base_anthropic_unified_messages_test import BaseAnthropicMessagesTest - -# Load environment variables -load_dotenv() - - -def _validate_anthropic_response(response: Dict[str, Any]): - assert "id" in response - assert "content" in response - assert "model" in response - assert response["role"] == "assistant" - - -class TestAnthropicDirectAPI(BaseAnthropicMessagesTest): - """Tests for direct Anthropic API calls""" - - test_non_streaming_base = None - - @property - def model_config(self) -> Dict[str, Any]: - return { - "model": "claude-haiku-4-5-20251001", - "api_key": os.getenv("ANTHROPIC_API_KEY"), - } - - @property - def expected_model_name_in_logging(self) -> str: - """ - This is the model name that is expected to be in the logging payload - """ - return "claude-haiku-4-5-20251001" - - -class TestAnthropicBedrockAPI(BaseAnthropicMessagesTest): - """Tests for Anthropic via Bedrock""" - - @property - def model_config(self) -> Dict[str, Any]: - return { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - } - - @property - def expected_model_name_in_logging(self) -> str: - """ - This is the model name that is expected to be in the logging payload - """ - return "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" - - -class TestAnthropicOpenAIAPI(BaseAnthropicMessagesTest): - """Tests for OpenAI via Anthropic messages interface""" - - @property - def model_config(self) -> Dict[str, Any]: - return { - "model": "openai/gpt-4.1-mini", - "client": None, - } - - @property - def expected_model_name_in_logging(self) -> str: - """ - This is the model name that is expected to be in the logging payload - """ - return "gpt-4.1-mini" diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py deleted file mode 100644 index e64218b677f..00000000000 --- a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py +++ /dev/null @@ -1,43 +0,0 @@ -""" -E2E Test suite for Anthropic Messages API prompt caching across different providers. - -Tests that prompt caching works correctly via litellm.anthropic.messages interface -by making actual API calls and validating usage metrics. - -Per AWS docs (https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html): -- Converse API uses: cachePoint: { type: "default" } -- InvokeModel API uses: cache_control: { type: "ephemeral" } -- Claude 3.7 Sonnet: GA, 1024 min tokens -- Claude 3.5 Haiku: GA, 2048 min tokens -""" - - - -import pytest -from base_anthropic_messages_prompt_caching_test import ( - BaseAnthropicMessagesPromptCachingTest, -) - - -class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest): - """ - E2E tests for prompt caching with Bedrock Converse API. - - Uses the bedrock/converse/ prefix which routes through litellm.completion() - and the AmazonConverseConfig transformation. - """ - - def get_model(self) -> str: - return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0" - - -class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): - """ - E2E tests for prompt caching with Bedrock Invoke API. - - Uses the bedrock/invoke/ prefix which routes through the native - Anthropic Messages API format. - """ - - def get_model(self) -> str: - return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0" diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_tool_search.py b/tests/pass_through_unit_tests/test_anthropic_messages_tool_search.py deleted file mode 100644 index 9006356ff2a..00000000000 --- a/tests/pass_through_unit_tests/test_anthropic_messages_tool_search.py +++ /dev/null @@ -1,80 +0,0 @@ -""" -E2E Test suite for Anthropic Messages API tool search across different providers. - -Tests that tool search works correctly via litellm.anthropic.messages interface -by making actual API calls. - -Supported providers: -- Anthropic API: advanced-tool-use-2025-11-20 -- Azure Anthropic: advanced-tool-use-2025-11-20 -- Vertex AI: tool-search-tool-2025-10-19 -- Bedrock Invoke: tool-search-tool-2025-10-19 - -Reference: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool -""" - - - -import pytest -from base_anthropic_messages_tool_search_test import ( - BaseAnthropicMessagesToolSearchTest, -) - - -class TestAnthropicAPIToolSearch(BaseAnthropicMessagesToolSearchTest): - """ - E2E tests for tool search with Anthropic API directly. - - Uses the anthropic/ prefix which routes through the native - Anthropic Messages API. - - Beta header: advanced-tool-use-2025-11-20 - - Note: Tool search is only supported on Claude Opus 4.5 and Claude Sonnet 4.5. - """ - - def get_model(self) -> str: - return "anthropic/claude-sonnet-4-5-20250929" - - -# class TestAzureAnthropicToolSearch(BaseAnthropicMessagesToolSearchTest): -# """ -# E2E tests for tool search with Azure Anthropic (Microsoft Foundry). - -# Uses the azure/ prefix which routes through Azure's Anthropic endpoint. - -# Beta header: advanced-tool-use-2025-11-20 -# """ - -# def get_model(self) -> str: -# return "azure/claude-sonnet-4-5-20250929" - - -# class TestVertexAIToolSearch(BaseAnthropicMessagesToolSearchTest): -# """ -# E2E tests for tool search with Vertex AI. - -# Uses the vertex_ai/ prefix which routes through Google Cloud's -# Vertex AI Anthropic partner models. - -# Beta header: tool-search-tool-2025-10-19 -# """ - -# def get_model(self) -> str: -# return "vertex_ai/claude-sonnet-4@20250514" - - -# class TestBedrockInvokeToolSearch(BaseAnthropicMessagesToolSearchTest): -# """ -# E2E tests for tool search with Bedrock Invoke API. - -# Uses the bedrock/invoke/ prefix which routes through the native -# Anthropic Messages API format on Bedrock. - -# Beta header: advanced-tool-use-2025-11-20 (passed via extra_headers) - -# Note: Tool search on Bedrock is only supported on Claude Opus 4.5. -# """ - -# def get_model(self) -> str: -# return "bedrock/invoke/us.anthropic.claude-opus-4-5-20251101-v1:0" diff --git a/tests/pass_through_unit_tests/test_bedrock_anthropic_messages_test.py b/tests/pass_through_unit_tests/test_bedrock_anthropic_messages_test.py deleted file mode 100644 index aacf65ea4a3..00000000000 --- a/tests/pass_through_unit_tests/test_bedrock_anthropic_messages_test.py +++ /dev/null @@ -1,86 +0,0 @@ -import json -import os -from datetime import datetime -from typing import AsyncIterator, Dict, Any -import asyncio -import unittest.mock -from unittest.mock import MagicMock -import pytest -from litellm.router import Router - -import litellm -from base_anthropic_unified_messages_test import BaseAnthropicMessagesTest - -INSTANCE_BASE_ANTHROPIC_MESSAGES_TEST = BaseAnthropicMessagesTest() - - -@pytest.mark.asyncio -async def test_anthropic_messages_bedrock_converse_with_thinking(): - """ - Test that bedrock/converse model works with thinking parameter. - Validates the request body from issue where budget_tokens was being lost. - """ - router = Router( - model_list=[ - { - "model_name": "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "litellm_params": { - "model": "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - }, - }, - ] - ) - - messages = [{"role": "user", "content": "What is 2+2?"}] - - response = await router.aanthropic_messages( - messages=messages, - model="bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - max_tokens=1026, - thinking={"type": "enabled", "budget_tokens": 1025}, - ) - print("bedrock response: ", response) - - # Verify response - INSTANCE_BASE_ANTHROPIC_MESSAGES_TEST._validate_response(response) - - -@pytest.mark.asyncio -async def test_should_not_fail_with_forwarded_headers_bedrock_invoke_messages(): - """ - E2E test for Bedrock invoke messages with header forwarding enabled. - This calls the real Bedrock endpoint (no mocks) and should not raise - SigV4 signature mismatch errors when forwarded headers are present. - """ - router = Router( - model_list=[ - { - "model_name": "claude-sonnet-4-5-20250929", - "litellm_params": { - "model": "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "aws_region_name": os.getenv("AWS_REGION_NAME"), - }, - } - ] - ) - - forwarded_headers = { - "x-forwarded-for": "10.11.232.194", - "x-forwarded-port": "443", - "x-forwarded-proto": "https", - "x-app": "cli", - } - - response = await router.aanthropic_messages( - messages=[{"role": "user", "content": "hi"}], - model="claude-sonnet-4-5-20250929", - max_tokens=5, - stream=False, - headers=forwarded_headers, # simulates forward_client_headers_to_llm_api - aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"), - aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), - aws_region_name=os.getenv("AWS_REGION_NAME"), - ) - print("INVOKE API RESPONSE: ", response) - - INSTANCE_BASE_ANTHROPIC_MESSAGES_TEST._validate_response(response) diff --git a/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py b/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py deleted file mode 100644 index 675b8e6b5fb..00000000000 --- a/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py +++ /dev/null @@ -1,65 +0,0 @@ -""" -Simple E2E test for Bedrock with advanced-tool-use beta header. - -Tests that LiteLLM correctly filters out the advanced-tool-use-2025-11-20 beta header -for Bedrock Invoke API, which doesn't support it and returns a 400 "invalid beta flag" error. -""" - -import pytest - - -import litellm - - -@pytest.mark.asyncio -async def test_bedrock_sonnet_4_5_with_advanced_tool_use_beta_header(): - """ - Simple E2E test: Call Bedrock Sonnet 4.5 with advanced-tool-use beta header. - - This should work without throwing "invalid beta flag" error because LiteLLM - filters out the advanced-tool-use beta header for Bedrock Invoke API. - """ - litellm.turn_on_debug() - response = await litellm.anthropic.messages.acreate( - model="bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[{"role": "user", "content": "What is 2+2?"}], - max_tokens=100, - provider_specific_header={ - "custom_llm_provider": "bedrock", - "extra_headers": { - "anthropic-beta": "advanced-tool-use-2025-11-20", - }, - }, - ) - - # Verify response - assert response is not None - assert "content" in response - print(f"āœ… Test passed! Response: {response}") - - -# @pytest.mark.asyncio -# async def test_bedrock_claude_3_5_with_advanced_tool_use_beta_header_filtered(): -# """ -# Simple E2E test: Call Bedrock Claude 3.5 with advanced-tool-use beta header. - -# This should work because the beta header is filtered out by LiteLLM before -# sending the request to Bedrock Invoke API. -# """ - -# response = await litellm.anthropic.messages.acreate( -# model="bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", -# messages=[{"role": "user", "content": "What is 2+2?"}], -# max_tokens=100, -# provider_specific_header={ -# "custom_llm_provider": "bedrock", -# "extra_headers": { -# "anthropic-beta": "advanced-tool-use-2025-11-20", -# }, -# }, -# ) - -# # Verify response -# assert response is not None -# assert "content" in response -# print(f"āœ… Test passed! Claude 3.5 response (beta header filtered): {response}") diff --git a/tests/pass_through_unit_tests/test_custom_logger_passthrough.py b/tests/pass_through_unit_tests/test_custom_logger_passthrough.py deleted file mode 100644 index 70fc8f9ccf2..00000000000 --- a/tests/pass_through_unit_tests/test_custom_logger_passthrough.py +++ /dev/null @@ -1,142 +0,0 @@ -import json -import threading -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Optional -from fastapi import Request -import pytest -import asyncio - - -import litellm -from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - PassthroughStandardLoggingPayload, -) -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - pass_through_request, -) - - -class TestCustomLogger(CustomLogger): - def __init__(self): - self.logged_kwargs: Optional[dict] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print( - "in async log success event kwargs", - json.dumps(kwargs, indent=4, default=str), - ) - self.logged_kwargs = kwargs - - -UPSTREAM_RESPONSE_BODY = { - "id": "modr-abc123", - "model": "omni-moderation-latest", - "results": [ - { - "flagged": False, - "categories": {"violence": False}, - "category_scores": {"violence": 1.2e-06}, - } - ], -} - - -@pytest.fixture -def upstream(): - received: dict = {} - - class Handler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def do_POST(self): - body = self.rfile.read(int(self.headers.get("content-length", 0) or 0)) - received["path"] = self.path - received["body"] = json.loads(body or b"{}") - payload = json.dumps(UPSTREAM_RESPONSE_BODY).encode() - self.send_response(200) - self.send_header("content-type", "application/json") - self.send_header("content-length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - - def log_message(self, *args): - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target=server.serve_forever, daemon=True).start() - try: - yield f"http://127.0.0.1:{server.server_port}", received - finally: - server.shutdown() - server.server_close() - - -@pytest.mark.asyncio -async def test_passthrough_logging_payload_for_a_route_no_provider_handler_claims( - upstream, -): - base_url, upstream_received = upstream - - test_custom_logger = TestCustomLogger() - litellm._async_success_callback = [test_custom_logger] - - TARGET_URL = f"{base_url}/v1/moderations" - REQUEST_BODY = { - "model": "omni-moderation-latest", - "input": "I want to bake a cake for my friend's birthday.", - } - TARGET_METHOD = "POST" - - result = await pass_through_request( - request=Request( - scope={ - "type": "http", - "method": TARGET_METHOD, - "path": "/v1/moderations", - "query_string": b"", - "headers": [ - (b"content-type", b"application/json"), - (b"authorization", b"Bearer sk-test-passthrough"), - ], - }, - ), - target=TARGET_URL, - custom_headers={ - "Content-Type": "application/json", - "Authorization": "Bearer sk-test-passthrough", - }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test", - user_id="test", - team_id="test", - end_user_id="test", - ), - custom_body=REQUEST_BODY, - forward_headers=False, - merge_query_params=False, - ) - - print("got result", result) - print("result status code", result.status_code) - print("result content", result.body) - - assert upstream_received.get("path") == "/v1/moderations" - assert upstream_received.get("body") == REQUEST_BODY - assert result.status_code == 200 - - await asyncio.sleep(1) - - assert test_custom_logger.logged_kwargs is not None - passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( - test_custom_logger.logged_kwargs["passthrough_logging_payload"] - ) - assert passthrough_logging_payload is not None - assert passthrough_logging_payload["url"] == TARGET_URL - assert passthrough_logging_payload["request_body"] == REQUEST_BODY - assert passthrough_logging_payload["request_method"] == TARGET_METHOD - - client_facing_response_body = json.loads(result.body) - assert client_facing_response_body == UPSTREAM_RESPONSE_BODY - assert passthrough_logging_payload["response_body"] == client_facing_response_body diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py deleted file mode 100644 index 665d6f05334..00000000000 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ /dev/null @@ -1,991 +0,0 @@ -""" -Real E2E Tests for WebSearch Interception - -Makes actual calls to test WebSearch interception with Perplexity. -Tests both streaming and non-streaming requests. -""" - -import sys - - -import litellm -from litellm.integrations.websearch_interception import ( - WebSearchInterceptionLogger, -) -from litellm.anthropic_interface import messages -from litellm.types.utils import LlmProviders - - -async def test_websearch_interception_non_streaming(): - """ - Test WebSearch interception with non-streaming request. - Validates that agentic loop executes transparently. - """ - litellm.turn_on_debug() - - print("\n" + "=" * 80) - print("E2E TEST 1: WebSearch Interception (Non-Streaming)") - print("=" * 80) - - # Initialize real router with search_tools configuration - import litellm.proxy.proxy_server as proxy_server - from litellm import Router - - # Create real router with search_tools - router = Router( - search_tools=[ - { - "search_tool_name": "my-perplexity-search", - "litellm_params": {"search_provider": "perplexity"}, - } - ] - ) - proxy_server.llm_router = router - - print("\nāœ… Initialized router with search_tools:") - print(f" - search_tool_name: my-perplexity-search") - print(f" - search_provider: perplexity") - - # Enable WebSearch interception for bedrock - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.BEDROCK], - search_tool_name="my-perplexity-search", - ) - litellm.callbacks = [websearch_logger] - litellm.set_verbose = True - - print("\nāœ… Configured WebSearch interception for Bedrock") - print("āœ… Will use search tool from router") - - try: - # Make request with WebSearch tool (non-streaming) - print("\nšŸ“ž Making litellm.messages.acreate() call...") - print(f" Model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") - print(f" Query: 'What is LiteLLM?'") - print(f" Tools: WebSearch") - print(f" Stream: False") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[ - { - "role": "user", - "content": "What is LiteLLM? Give me a brief overview.", - } - ], - tools=[ - { - "name": "WebSearch", - "description": "Search the web for information", - "input_schema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query", - } - }, - "required": ["query"], - }, - } - ], - max_tokens=1024, - stream=False, - ) - - print("\nāœ… Received response!") - - # Handle both dict and object responses - if isinstance(response, dict): - response_id = response.get("id") - response_model = response.get("model") - response_stop_reason = response.get("stop_reason") - response_content = response.get("content", []) - else: - response_id = response.id - response_model = response.model - response_stop_reason = response.stop_reason - response_content = response.content - - print(f"\nšŸ“„ Response ID: {response_id}") - print(f"šŸ“„ Model: {response_model}") - print(f"šŸ“„ Stop Reason: {response_stop_reason}") - print(f"šŸ“„ Content blocks: {len(response_content)}") - - # Debug: Print all content block types - for i, block in enumerate(response_content): - block_type = block.get("type") if isinstance(block, dict) else block.type - print(f" Block {i}: type={block_type}") - if block_type == "tool_use": - block_name = ( - block.get("name") if isinstance(block, dict) else block.name - ) - print(f" name={block_name}") - - # Validate response - assert response is not None, "Response should not be None" - assert response_content is not None, "Response should have content" - assert ( - len(response_content) > 0 - ), "Response should have at least one content block" - - # Check if response contains tool_use (means interception didn't work) - has_tool_use = any( - (block.get("type") if isinstance(block, dict) else block.type) == "tool_use" - for block in response_content - ) - - # Check if we got a text response - has_text = any( - (block.get("type") if isinstance(block, dict) else block.type) == "text" - for block in response_content - ) - - if has_tool_use: - print("\nāŒ TEST 1 FAILED: Interception did not work") - print(f"āŒ Stop reason: {response_stop_reason}") - print("āŒ Response contains tool_use blocks") - return False - - elif has_text and response_stop_reason != "tool_use": - text_block = next( - block - for block in response_content - if (block.get("type") if isinstance(block, dict) else block.type) - == "text" - ) - text_content = ( - text_block.get("text") - if isinstance(text_block, dict) - else text_block.text - ) - - print(f"\nšŸ“ Response Text:") - print(f" {text_content[:200]}...") - - if "litellm" in text_content.lower(): - print("\n" + "=" * 80) - print("āœ… TEST 1 PASSED!") - print("=" * 80) - print("āœ… User made ONE litellm.messages.acreate() call") - print("āœ… Got back final answer (not tool_use)") - print("āœ… Agentic loop executed transparently") - print("āœ… WebSearch interception working!") - print("=" * 80) - return True - else: - print("\nāš ļø Got text response but doesn't mention LiteLLM") - return False - else: - print("\nāŒ Unexpected response format") - return False - - except Exception as e: - print(f"\nāŒ Test 1 failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -async def test_websearch_interception_streaming(): - """ - Test WebSearch interception with streaming request. - Validates that stream=True is converted to stream=False transparently. - """ - print("\n" + "=" * 80) - print("E2E TEST 2: WebSearch Interception (Streaming)") - print("=" * 80) - - # Router already initialized from test 1 - print("\nāœ… Using existing router configuration") - print("āœ… WebSearch interception already enabled for Bedrock") - print("āœ… Streaming will be converted to non-streaming for WebSearch interception") - - try: - # Make request with WebSearch tool AND stream=True - print("\nšŸ“ž Making litellm.messages.acreate() call with stream=True...") - print(f" Model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") - print(f" Query: 'What is LiteLLM?'") - print(f" Tools: WebSearch") - print(f" Stream: True (will be converted to False)") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[ - { - "role": "user", - "content": "What is LiteLLM? Give me a brief overview.", - } - ], - tools=[ - { - "name": "WebSearch", - "description": "Search the web for information", - "input_schema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query", - } - }, - "required": ["query"], - }, - } - ], - max_tokens=1024, - stream=True, # REQUEST STREAMING - ) - - print("\nāœ… Received response!") - - # Check if response is actually a stream (async generator) - import inspect - - is_stream = inspect.isasyncgen(response) - - if is_stream: - print("\nāš ļø WARNING: Response is a stream (async_generator)") - print("āš ļø This means stream conversion didn't work!") - print("\nšŸ“¦ Consuming stream chunks:") - - chunks = [] - chunk_count = 0 - async for chunk in response: - chunk_count += 1 - print(f"\n--- Chunk {chunk_count} ---") - print(chunk) - chunks.append(chunk) - - print( - f"\nāŒ TEST 2 FAILED: Got {len(chunks)} stream chunks instead of single response" - ) - return False - - # If not a stream, validate as normal response - print("āœ… Response is NOT a stream (conversion worked!)") - - # Handle both dict and object responses - if isinstance(response, dict): - response_id = response.get("id") - response_model = response.get("model") - response_stop_reason = response.get("stop_reason") - response_content = response.get("content", []) - else: - response_id = response.id - response_model = response.model - response_stop_reason = response.stop_reason - response_content = response.content - - print(f"\nšŸ“„ Response ID: {response_id}") - print(f"šŸ“„ Model: {response_model}") - print(f"šŸ“„ Stop Reason: {response_stop_reason}") - print(f"šŸ“„ Content blocks: {len(response_content)}") - - # Debug: Print all content block types - for i, block in enumerate(response_content): - block_type = block.get("type") if isinstance(block, dict) else block.type - print(f" Block {i}: type={block_type}") - - # Validate response - assert response is not None, "Response should not be None" - assert response_content is not None, "Response should have content" - assert ( - len(response_content) > 0 - ), "Response should have at least one content block" - - # Check if response contains tool_use (means interception didn't work) - has_tool_use = any( - (block.get("type") if isinstance(block, dict) else block.type) == "tool_use" - for block in response_content - ) - - # Check if we got a text response - has_text = any( - (block.get("type") if isinstance(block, dict) else block.type) == "text" - for block in response_content - ) - - if has_tool_use: - print("\nāŒ TEST 2 FAILED: Interception did not work") - print("āŒ Response contains tool_use blocks") - return False - - elif has_text and response_stop_reason != "tool_use": - text_block = next( - block - for block in response_content - if (block.get("type") if isinstance(block, dict) else block.type) - == "text" - ) - text_content = ( - text_block.get("text") - if isinstance(text_block, dict) - else text_block.text - ) - - print(f"\nšŸ“ Response Text:") - print(f" {text_content[:200]}...") - - if "litellm" in text_content.lower(): - print("\n" + "=" * 80) - print("āœ… TEST 2 PASSED!") - print("=" * 80) - print( - "āœ… User made ONE litellm.messages.acreate() call with stream=True" - ) - print("āœ… Stream was transparently converted to non-streaming") - print("āœ… Got back final answer (not tool_use)") - print("āœ… Agentic loop executed transparently") - print("āœ… WebSearch interception working with streaming!") - print("=" * 80) - return True - else: - print("\nāš ļø Got text response but doesn't mention LiteLLM") - return False - else: - print("\nāŒ Unexpected response format") - return False - - except Exception as e: - print(f"\nāŒ Test 2 failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -async def test_websearch_interception_no_tool_call_streaming(): - """ - Test WebSearch interception when LLM doesn't make a tool call with streaming. - - This tests the scenario where: - 1. User requests stream=True - 2. WebSearch tool is provided - 3. LLM decides NOT to use the tool (just responds with text) - 4. System should return a fake stream - """ - print("\n" + "=" * 80) - print("E2E TEST 3: WebSearch Interception (No Tool Call, Streaming)") - print("=" * 80) - - # Router already initialized from test 1 - print("\nāœ… Using existing router configuration") - print("āœ… WebSearch interception already enabled for Bedrock") - - try: - # Make request with WebSearch tool AND stream=True - # Use a query that the LLM will answer directly without using the tool - print("\nšŸ“ž Making litellm.messages.acreate() call with stream=True...") - print(f" Model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") - print(f" Query: 'What is 2+2?'") - print(f" Tools: WebSearch") - print(f" Stream: True") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[ - { - "role": "user", - "content": "What is 2+2? Just give me the answer, no need to search.", - } - ], - tools=[ - { - "name": "WebSearch", - "description": "Search the web for information", - "input_schema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query", - } - }, - "required": ["query"], - }, - } - ], - max_tokens=1024, - stream=True, # REQUEST STREAMING - ) - - print("\nāœ… Received response!") - - # Check if response is actually a stream (async generator or async iterator) - import inspect - - is_async_gen = inspect.isasyncgen(response) - is_async_iter = hasattr(response, "__aiter__") and hasattr( - response, "__anext__" - ) - is_stream = is_async_gen or is_async_iter - - if not is_stream: - print("\nāŒ TEST 3 FAILED: Response is NOT a stream") - print(f"āŒ Expected a fake stream when LLM doesn't use the tool") - print(f"āŒ Response type: {type(response)}") - return False - - print( - f"āœ… Response is a stream (async_gen={is_async_gen}, async_iter={is_async_iter})" - ) - print("\nšŸ“¦ Consuming stream chunks:") - - chunks = [] - chunk_count = 0 - async for chunk in response: - chunk_count += 1 - print(f"\n--- Chunk {chunk_count} ---") - print(f" Type: {type(chunk)}") - print( - f" Content: {chunk[:200] if isinstance(chunk, bytes) else str(chunk)[:200]}..." - ) - chunks.append(chunk) - - print(f"\nāœ… Received {len(chunks)} stream chunk(s)") - - if len(chunks) > 0: - print("\n" + "=" * 80) - print("āœ… TEST 3 PASSED!") - print("=" * 80) - print("āœ… User made ONE litellm.messages.acreate() call with stream=True") - print("āœ… LLM didn't use the WebSearch tool") - print("āœ… Got back a fake stream (not a non-streaming response)") - print("āœ… WebSearch interception handles no-tool-call case correctly!") - print("=" * 80) - return True - else: - print("\nāŒ TEST 3 FAILED: No chunks received") - return False - - except Exception as e: - print(f"\nāŒ Test 3 failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -async def test_claude_code_native_websearch(): - """ - Test WebSearch interception with Claude Code's native web_search_20250305 tool. - - This tests the exact request format that Claude Code sends: - - tools: [{'type': 'web_search_20250305', 'name': 'web_search', 'max_uses': 8}] - - Model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 - """ - print("\n" + "=" * 80) - print("E2E TEST: Claude Code Native WebSearch (web_search_20250305)") - print("=" * 80) - - # Router already initialized from test 1 - print("\nāœ… Using existing router configuration") - print("āœ… WebSearch interception already enabled for Bedrock") - - try: - # Make request with Claude Code's exact native web_search tool format - print("\nšŸ“ž Making litellm.messages.acreate() call...") - print(f" Model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0") - print(f" Query: 'Perform a web search for the query: litellm what is it'") - print(f" Tools: Native web_search_20250305") - print(f" Stream: False") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[ - { - "role": "user", - "content": "Perform a web search for the query: litellm what is it", - } - ], - tools=[ - {"type": "web_search_20250305", "name": "web_search", "max_uses": 8} - ], - max_tokens=1024, - stream=False, - ) - - print("\nāœ… Received response!") - - # Handle both dict and object responses - if isinstance(response, dict): - response_id = response.get("id") - response_model = response.get("model") - response_stop_reason = response.get("stop_reason") - response_content = response.get("content", []) - else: - response_id = response.id - response_model = response.model - response_stop_reason = response.stop_reason - response_content = response.content - - print(f"\nšŸ“„ Response ID: {response_id}") - print(f"šŸ“„ Model: {response_model}") - print(f"šŸ“„ Stop Reason: {response_stop_reason}") - print(f"šŸ“„ Content blocks: {len(response_content)}") - - # Debug: Print all content block types - for i, block in enumerate(response_content): - block_type = block.get("type") if isinstance(block, dict) else block.type - print(f" Block {i}: type={block_type}") - if block_type == "tool_use": - block_name = ( - block.get("name") if isinstance(block, dict) else block.name - ) - print(f" name={block_name}") - - # Validate response - assert response is not None, "Response should not be None" - assert response_content is not None, "Response should have content" - assert ( - len(response_content) > 0 - ), "Response should have at least one content block" - - # Check if response contains tool_use (means interception didn't work) - has_tool_use = any( - (block.get("type") if isinstance(block, dict) else block.type) == "tool_use" - for block in response_content - ) - - # Check if we got a text response - has_text = any( - (block.get("type") if isinstance(block, dict) else block.type) == "text" - for block in response_content - ) - - if has_tool_use: - print("\nāŒ TEST FAILED: Interception did not work") - print(f"āŒ Stop reason: {response_stop_reason}") - print("āŒ Response contains tool_use blocks") - return False - - elif has_text and response_stop_reason != "tool_use": - text_block = next( - block - for block in response_content - if (block.get("type") if isinstance(block, dict) else block.type) - == "text" - ) - text_content = ( - text_block.get("text") - if isinstance(text_block, dict) - else text_block.text - ) - - print(f"\nšŸ“ Response Text:") - print(f" {text_content[:200]}...") - - if "litellm" in text_content.lower(): - print("\n" + "=" * 80) - print("āœ… TEST PASSED!") - print("=" * 80) - print( - "āœ… Claude Code's native web_search_20250305 tool was intercepted" - ) - print("āœ… Tool was converted to LiteLLM standard format") - print("āœ… User made ONE litellm.messages.acreate() call") - print("āœ… Got back final answer with search results") - print("āœ… Agentic loop executed transparently") - print("āœ… WebSearch interception working with Claude Code!") - print("=" * 80) - return True - else: - print("\nāš ļø Got text response but doesn't mention LiteLLM") - return False - else: - print("\nāŒ Unexpected response format") - return False - - except Exception as e: - print(f"\nāŒ Test failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -if __name__ == "__main__": - import asyncio - - async def run_all_tests(): - """Run all E2E tests""" - test_results = [] - - # Test 1: Non-streaming - result1 = await test_websearch_interception_non_streaming() - test_results.append(("Non-Streaming", result1)) - - # Test 2: Streaming - result2 = await test_websearch_interception_streaming() - test_results.append(("Streaming", result2)) - - # Test 3: No tool call with streaming - result3 = await test_websearch_interception_no_tool_call_streaming() - test_results.append(("No Tool Call Streaming", result3)) - - # Test 4: Claude Code native web_search - result4 = await test_claude_code_native_websearch() - test_results.append(("Claude Code Native WebSearch", result4)) - - # Print summary - print("\n" + "=" * 80) - print("TEST SUMMARY") - print("=" * 80) - for test_name, result in test_results: - status = "āœ… PASSED" if result else "āŒ FAILED" - print(f"{test_name}: {status}") - print("=" * 80) - - # Return overall result - return all(result for _, result in test_results) - - result = asyncio.run(run_all_tests()) - import sys - - sys.exit(0 if result else 1) - - -async def test_litellm_standard_websearch_tool(): - """ - PRIORITY TEST #1: Test with the canonical litellm_web_search tool format. - - This validates that using get_litellm_web_search_tool() directly - works end-to-end without any conversion needed. - """ - print("\n" + "=" * 80) - print("E2E TEST: LiteLLM Standard WebSearch Tool") - print("=" * 80) - - from litellm.integrations.websearch_interception import get_litellm_web_search_tool - - print("\nāœ… Using existing router configuration") - print("āœ… WebSearch interception already enabled for Bedrock") - - try: - print("\nšŸ“ž Making litellm.messages.acreate() call...") - print(f" Model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0") - print(f" Query: 'What is the latest news about AI?'") - print(f" Tool: litellm_web_search (standard format, no conversion needed)") - print(f" Stream: False") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[ - { - "role": "user", - "content": "What is the latest news about AI? Give me a brief overview.", - } - ], - tools=[get_litellm_web_search_tool()], - max_tokens=1024, - stream=False, - ) - - print("\nāœ… Received response!") - - if isinstance(response, dict): - response_id = response.get("id") - response_stop_reason = response.get("stop_reason") - response_content = response.get("content", []) - else: - response_id = response.id - response_stop_reason = response.stop_reason - response_content = response.content - - print(f"\nšŸ“„ Response ID: {response_id}") - print(f"šŸ“„ Stop Reason: {response_stop_reason}") - print(f"šŸ“„ Content blocks: {len(response_content)}") - - for i, block in enumerate(response_content): - block_type = block.get("type") if isinstance(block, dict) else block.type - print(f" Block {i}: type={block_type}") - - has_tool_use = any( - (block.get("type") if isinstance(block, dict) else block.type) == "tool_use" - for block in response_content - ) - - has_text = any( - (block.get("type") if isinstance(block, dict) else block.type) == "text" - for block in response_content - ) - - if has_tool_use: - print("\nāŒ TEST FAILED: Interception did not work") - return False - - elif has_text and response_stop_reason != "tool_use": - text_block = next( - block - for block in response_content - if (block.get("type") if isinstance(block, dict) else block.type) - == "text" - ) - text_content = ( - text_block.get("text") - if isinstance(text_block, dict) - else text_block.text - ) - - print(f"\nšŸ“ Response Text: {text_content[:200]}...") - - print("\n" + "=" * 80) - print("āœ… TEST PASSED!") - print("=" * 80) - print("āœ… LiteLLM standard tool format works without conversion") - print("āœ… Agentic loop executed transparently") - print("=" * 80) - return True - else: - print("\nāŒ Unexpected response format") - return False - - except Exception as e: - print(f"\nāŒ Test failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -async def test_claude_code_native_websearch_streaming(): - """ - PRIORITY TEST #2: Test Claude Code's native tool WITH stream=True. - - Validates: - - Native tool conversion (web_search_20250305 ā†’ litellm_web_search) - - Stream=True ā†’ Stream=False conversion - - Agentic loop executes with both conversions - """ - print("\n" + "=" * 80) - print("E2E TEST: Claude Code Native WebSearch + Streaming") - print("=" * 80) - - print("\nāœ… Using existing router configuration") - print("āœ… WebSearch interception already enabled for Bedrock") - - try: - print("\nšŸ“ž Making litellm.messages.acreate() call with stream=True...") - print(f" Model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0") - print(f" Tool: Native web_search_20250305") - print(f" Stream: True (will be converted to False)") - - response = await messages.acreate( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[ - {"role": "user", "content": "Search for the latest AI developments."} - ], - tools=[ - {"type": "web_search_20250305", "name": "web_search", "max_uses": 8} - ], - max_tokens=1024, - stream=True, - ) - - print("\nāœ… Received response!") - - import inspect - - is_stream = inspect.isasyncgen(response) - - if is_stream: - print("\nāš ļø Response is a stream (stream conversion didn't work)") - return False - - print("āœ… Response is NOT a stream (conversion worked!)") - - if isinstance(response, dict): - response_stop_reason = response.get("stop_reason") - response_content = response.get("content", []) - else: - response_stop_reason = response.stop_reason - response_content = response.content - - has_tool_use = any( - (block.get("type") if isinstance(block, dict) else block.type) == "tool_use" - for block in response_content - ) - - has_text = any( - (block.get("type") if isinstance(block, dict) else block.type) == "text" - for block in response_content - ) - - if has_tool_use: - print("\nāŒ TEST FAILED: Interception did not work") - return False - - elif has_text and response_stop_reason != "tool_use": - print("\n" + "=" * 80) - print("āœ… TEST PASSED!") - print("=" * 80) - print("āœ… Native tool converted to litellm_web_search") - print("āœ… Stream=True converted to Stream=False") - print("āœ… Both conversions working together!") - print("=" * 80) - return True - else: - print("\nāŒ Unexpected response format") - return False - - except Exception as e: - print(f"\nāŒ Test failed with error: {str(e)}") - import traceback - - traceback.print_exc() - return False - - -async def test_pre_request_hook_modifies_request_body(): - """ - Unit test to verify async_pre_request_hook correctly modifies request body. - - Tests that: - 1. WebSearchInterceptionLogger is active - 2. Native web_search_20250305 tool is converted to litellm_web_search - 3. Stream is converted from True to False - 4. Modified parameters reach the API call - """ - import asyncio - from unittest.mock import AsyncMock, patch, MagicMock - from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME - - litellm.turn_on_debug() - - print("\n" + "=" * 80) - print("UNIT TEST: Pre-Request Hook Modifies Request Body") - print("=" * 80) - - # Initialize WebSearchInterceptionLogger - litellm.callbacks = [ - WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.BEDROCK], - search_tool_name="test-search-tool", - ) - ] - - print("āœ… WebSearchInterceptionLogger initialized") - - mock_router = MagicMock() - mock_router.search_tools = [ - { - "search_tool_name": "test-search-tool", - "litellm_params": {"search_provider": "tavily"}, - } - ] - - # Track what actually gets sent to the API - captured_request = {} - - def mock_anthropic_messages_handler( - max_tokens, - messages, - model, - metadata=None, - stop_sequences=None, - stream=None, - system=None, - temperature=None, - thinking=None, - tool_choice=None, - tools=None, - top_k=None, - top_p=None, - container=None, - api_key=None, - api_base=None, - client=None, - custom_llm_provider=None, - **kwargs, - ): - """Mock handler that captures the actual request parameters""" - # Capture what gets sent to the handler (after hook modifications) - captured_request["tools"] = tools - captured_request["stream"] = stream - captured_request["max_tokens"] = max_tokens - captured_request["model"] = model - - # Return a mock response (non-streaming) - from litellm.types.llms.anthropic_messages.anthropic_response import ( - AnthropicMessagesResponse, - ) - - return AnthropicMessagesResponse( - id="msg_test", - type="message", - role="assistant", - content=[{"type": "text", "text": "Test response"}], - model="claude-sonnet-4-5", - stop_reason="end_turn", - usage={"input_tokens": 10, "output_tokens": 20}, - ) - - # Patch the anthropic_messages_handler function (called after hooks) - with patch( - "litellm.llms.anthropic.pass_through.messages.handler.anthropic_messages_handler", - side_effect=mock_anthropic_messages_handler, - ), patch( # test-quality-ok: the hook imports this process-global router at call time; no injection seam exists to register search_tools - "litellm.proxy.proxy_server.llm_router", - mock_router, - ): - - print( - "\nšŸ“ Making request with native web_search_20250305 tool (stream=True)..." - ) - - # Make the request with native tool format - response = await messages.acreate( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[{"role": "user", "content": "Test query"}], - tools=[ - {"type": "web_search_20250305", "name": "web_search", "max_uses": 8} - ], - max_tokens=100, - stream=True, # Should be converted to False - ) - - print("\nšŸ” Verifying request modifications...") - - # Verify tool was converted - tools = captured_request.get("tools") - print(f"\n Captured tools: {tools}") - - if tools and len(tools) > 0: - tool = tools[0] - tool_name = tool.get("name") - - if tool_name == LITELLM_WEB_SEARCH_TOOL_NAME: - print( - f" āœ… Tool converted: web_search_20250305 ā†’ {LITELLM_WEB_SEARCH_TOOL_NAME}" - ) - else: - print( - f" āŒ Tool NOT converted: expected {LITELLM_WEB_SEARCH_TOOL_NAME}, got {tool_name}" - ) - return False - else: - print(" āŒ No tools captured in request") - return False - - # Verify stream was converted - stream = captured_request.get("stream") - print(f" Captured stream: {stream}") - - if stream is False: - print(" āœ… Stream converted: True ā†’ False") - else: - print(f" āŒ Stream NOT converted: expected False, got {stream}") - return False - - print("\n" + "=" * 80) - print("āœ… PRE-REQUEST HOOK TEST PASSED!") - print("=" * 80) - print("āœ… CustomLogger is active") - print("āœ… async_pre_request_hook modifies request body") - print("āœ… Tool conversion works correctly") - print("āœ… Stream conversion works correctly") - print("=" * 80) - - return True diff --git a/tests/unit/enterprise/enterprise_callbacks/test_pagerduty_alerting.py b/tests/unit/enterprise/enterprise_callbacks/test_pagerduty_alerting.py new file mode 100644 index 00000000000..c26966eeb38 --- /dev/null +++ b/tests/unit/enterprise/enterprise_callbacks/test_pagerduty_alerting.py @@ -0,0 +1,98 @@ +from collections.abc import Iterator +from typing import Final, cast + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.types.utils import StandardLoggingPayload +from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting +from litellm.proxy._types import UserAPIKeyAuth + +PAGERDUTY_URL: Final = "https://events.pagerduty.com/v2/enqueue" + + +@pytest.fixture(autouse=True) +def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _pagerduty_logger(monkeypatch: pytest.MonkeyPatch, failure_threshold: int = 1) -> PagerDutyAlerting: + monkeypatch.setenv("PAGERDUTY_API_KEY", "unit-routing-key") + return PagerDutyAlerting( + alerting_args={ + "failure_threshold": failure_threshold, + "failure_threshold_window_seconds": 600, + "hanging_threshold_seconds": 0, + "hanging_threshold_fails": 1, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("failure_count", "threshold", "request_count"), [(1, 1, 1), (2, 3, 0), (3, 3, 1)]) +async def test_failure_alerts_follow_the_configured_threshold( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + failure_count: int, + threshold: int, + request_count: int, +) -> None: + route: Final = respx_mock.post(PAGERDUTY_URL).mock(return_value=httpx.Response(202)) + logger: Final = _pagerduty_logger(monkeypatch, threshold) + payload: Final = cast( + StandardLoggingPayload, + { + "error_information": { + "error_class": "RateLimitError", + "error_code": "429", + "llm_provider": "openai", + }, + "metadata": {}, + }, + ) + + for _ in range(failure_count): + await logger.async_log_failure_event({"standard_logging_object": payload}, None, None, None) + + assert route.call_count == request_count + if request_count == 1: + body: Final = TypeAdapter(dict[str, object]).validate_json(route.calls.last.request.content) + payload_body: Final = TypeAdapter(dict[str, object]).validate_python(body["payload"]) + details: Final = TypeAdapter(dict[str, object]).validate_python(payload_body["custom_details"]) + recent_errors: Final = TypeAdapter(list[dict[str, object]]).validate_python(details["recent_errors"]) + assert body["routing_key"] == "unit-routing-key" + assert body["event_action"] == "trigger" + assert payload_body["summary"] == f"High LLM API Failure Rate: {failure_count} in the last 600 seconds." + assert recent_errors[0]["error_class"] == "RateLimitError" + assert recent_errors[0]["error_code"] == "429" + assert recent_errors[0]["error_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_hanging_requests_alert_at_the_default_threshold( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + route: Final = respx_mock.post(PAGERDUTY_URL).mock(return_value=httpx.Response(202)) + logger: Final = _pagerduty_logger(monkeypatch) + user: Final = UserAPIKeyAuth(api_key="unit-key", key_alias="unit-alias", user_id="unit-user") + + for _ in range(60): + await logger.hanging_response_handler( + {"request_id": "request-1", "litellm_call_id": "unit-call"}, + user, + ) + + assert route.call_count == 1 + body: Final = TypeAdapter(dict[str, object]).validate_json(route.calls.last.request.content) + payload_body: Final = TypeAdapter(dict[str, object]).validate_python(body["payload"]) + details: Final = TypeAdapter(dict[str, object]).validate_python(payload_body["custom_details"]) + recent_errors: Final = TypeAdapter(list[dict[str, object]]).validate_python(details["recent_errors"]) + assert body["routing_key"] == "unit-routing-key" + assert payload_body["summary"] == "High Number of Hanging LLM Requests: 60 in the last 600 seconds." + assert recent_errors[0]["failure_event_type"] == "hanging_response" diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json similarity index 100% rename from tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json rename to tests/unit/integrations/gcs_pubsub/spend_logs_payload.json diff --git a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py index 2196116c383..0000b778626 100644 --- a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py +++ b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py @@ -1,14 +1,40 @@ +import base64 import datetime import json import os import unittest -from typing import List, Optional, Tuple +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, List, Optional, Tuple, cast from unittest.mock import ANY, MagicMock, Mock, patch import httpx import pytest +import respx import litellm +from pydantic import TypeAdapter +from litellm.types.utils import StandardLoggingPayload + + +@dataclass(frozen=True, slots=True) +class _CachedVertexCredentials: + token: str + quota_project_id: str | None + expired: bool = False + + def refresh(self, request: object) -> None: + return None + + +@pytest.fixture(autouse=True) +def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() @pytest.mark.asyncio @@ -62,3 +88,169 @@ async def test_construct_request_headers_project_id_from_env(monkeypatch): finally: # Clean up environment variable del os.environ["GCS_PUBSUB_PROJECT_ID"] + + +def _authorized_user_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, project_id: str) -> str: + credentials_path: Final = tmp_path / "google-auth.json" + credentials_path.write_text( + json.dumps( + { + "type": "authorized_user", + "client_id": "unit-client", + "client_secret": "unit-secret", + "refresh_token": "unit-refresh-token", + "token": "unit-access-token", + "token_uri": "https://oauth2.googleapis.com/token", + "quota_project_id": "unit-project", + "expiry": "2099-01-01T00:00:00Z", + } + ) + ) + credentials_path_string: Final = str(credentials_path) + credentials: Final = _CachedVertexCredentials(token="unit-access-token", quota_project_id=project_id) + vertex_chat_completion: Final = litellm.vertex_chat_completion + monkeypatch.setitem( + vertex_chat_completion._credentials_project_mapping, + (credentials_path_string, project_id), + (credentials, project_id), + ) + return credentials_path_string + + +@pytest.mark.asyncio +async def test_publish_message_encodes_the_standard_payload( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, tmp_path: Path +) -> None: + from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger + + project_id: Final = "unit-project" + topic_id: Final = "unit-topic" + url: Final = f"https://pubsub.googleapis.com/v1/projects/{project_id}/topics/{topic_id}:publish" + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + route: Final = respx_mock.post(url).mock(return_value=httpx.Response(202, json={"messageIds": ["publish-1"]})) + logger: Final = GcsPubSubLogger( + project_id=project_id, + topic_id=topic_id, + credentials_path=_authorized_user_credentials(monkeypatch, tmp_path, project_id), + ) + payload: Final = cast( + StandardLoggingPayload, + { + "request_id": "request-1", + "model": "vertex_ai/gemini-2.5-flash", + "messages": [{"role": "user", "content": "publish this"}], + "response": {"choices": [{"message": {"content": "published"}}]}, + }, + ) + + response: Final = TypeAdapter(dict[str, list[str]]).validate_python(await logger.publish_message(payload)) + request_json: Final = TypeAdapter(dict[str, list[dict[str, str]]]).validate_json(route.calls.last.request.content) + decoded_message: Final = TypeAdapter(dict[str, object]).validate_json( + base64.b64decode(request_json["messages"][0]["data"]) + ) + + assert route.call_count == 1 + assert route.calls.last.request.url == url + assert route.calls.last.request.headers["Authorization"] == "Bearer unit-access-token" + assert request_json == { + "messages": [ + { + "data": base64.b64encode(json.dumps(payload, default=str).encode("utf-8")).decode("utf-8"), + } + ] + } + assert decoded_message == payload + assert response == {"messageIds": ["publish-1"]} + + +def _without_volatile_gcs_payload_fields(value: object, path: str = "") -> object: + ignored_paths: Final = frozenset( + { + "request_id", + "litellm_call_id", + "metadata.litellm_call_id", + "session_id", + "startTime", + "endTime", + "completionStartTime", + "request_duration_ms", + "metadata.model_map_information", + "metadata.usage_object", + "metadata.cold_storage_object_key", + "metadata.litellm_overhead_time_ms", + "metadata.cost_breakdown", + "metadata.autorouter_savings", + "metadata.eval_information", + } + ) + if isinstance(value, Mapping): + mapping: Final = TypeAdapter(dict[str, object]).validate_python(value) + return MappingProxyType( + { + key: _without_volatile_gcs_payload_fields(item, f"{path}.{key}" if path else key) + for key, item in mapping.items() + if (f"{path}.{key}" if path else key) not in ignored_paths + } + ) + if isinstance(value, list): + items: Final = TypeAdapter(list[object]).validate_python(value) + return tuple( + _without_volatile_gcs_payload_fields(item, f"{path}[{index}]") for index, item in enumerate(items) + ) + return value + + +@pytest.mark.asyncio +async def test_v1_success_event_publishes_the_spend_log_fixture( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, tmp_path: Path +) -> None: + from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger + + project_id: Final = "unit-project" + topic_id: Final = "unit-topic" + url: Final = f"https://pubsub.googleapis.com/v1/projects/{project_id}/topics/{topic_id}:publish" + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr(litellm, "gcs_pub_sub_use_v1", True) + route: Final = respx_mock.post(url).mock(return_value=httpx.Response(202, json={"messageIds": ["publish-2"]})) + logger: Final = GcsPubSubLogger( + project_id=project_id, + topic_id=topic_id, + credentials_path=_authorized_user_credentials(monkeypatch, tmp_path, project_id), + ) + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": "gpt-4o", + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + "metadata": {}, + "hidden_params": {}, + "model_map_information": None, + }, + ) + response_obj: Final = { + "id": "unit-response", + "model": "gpt-4o", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + "choices": [{"message": {"content": "hi"}}], + } + await logger.async_log_success_event( + kwargs={"model": "gpt-4o", "standard_logging_object": standard_logging_payload}, + response_obj=response_obj, + start_time=datetime.datetime(2025, 1, 1), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + await logger.async_send_batch() + + request_json: Final = TypeAdapter(dict[str, list[dict[str, str]]]).validate_json(route.calls.last.request.content) + decoded_payload: Final = TypeAdapter(dict[str, object]).validate_json( + base64.b64decode(request_json["messages"][0]["data"]) + ) + + assert route.call_count == 1 + assert route.calls.last.request.url == url + assert decoded_payload["model"] == "gpt-4o" + assert decoded_payload["prompt_tokens"] == 1 + assert decoded_payload["completion_tokens"] == 1 + assert decoded_payload["total_tokens"] == 2 diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_complex_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_complex_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_langfuse_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_langfuse_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_no_choices.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_no_choices.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json similarity index 75% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json index 749796e0d04..6da1fffd329 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json @@ -21,10 +21,17 @@ "extra_body": "{}" }, "langfuse.observation.type": "generation", + "langfuse.observation.output": { + "content": "Hello! How can I assist you today?", + "role": "assistant", + "tool_calls": null, + "function_call": null, + "provider_specific_fields": null + }, "langfuse.observation.usage_details": { "input": 10, - "output": 10, - "total": 20, + "output": 20, + "total": 30, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0 }, diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json similarity index 95% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json index 5223538f919..ccedcfe1e9d 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json @@ -16,7 +16,8 @@ "langfuse.observation.level": "DEFAULT", "langfuse.observation.model.name": "gpt-3.5-turbo", "langfuse.observation.model.parameters": { - "extra_body": "{}" + "extra_body": "{}", + "stream": true }, "langfuse.observation.output": { "content": "Hello! How can I assist you today?", diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_vertex_call.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_vertex_call.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/complex_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/complex_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/complex_metadata_2.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/complex_metadata_2.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/embedding_with_vllm.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/embedding_with_vllm.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/embedding_with_vllm.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/embedding_with_vllm.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/empty_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/empty_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/metadata_with_function.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/metadata_with_function.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/metadata_with_lock.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/metadata_with_lock.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/nested_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/nested_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata2.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata2.json diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata3.json similarity index 100% rename from tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json rename to tests/unit/integrations/langfuse/langfuse_expected_request_body/simple_metadata3.json diff --git a/tests/unit/integrations/langfuse/test_langfuse_callback_migrations.py b/tests/unit/integrations/langfuse/test_langfuse_callback_migrations.py new file mode 100644 index 00000000000..a23e3f479c6 --- /dev/null +++ b/tests/unit/integrations/langfuse/test_langfuse_callback_migrations.py @@ -0,0 +1,471 @@ +import asyncio +import json +import threading +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +import respx +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.proto.common.v1.common_pb2 import AnyValue +from pydantic import TypeAdapter + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.integrations.langfuse.langfuse import LangFuseLogger +from litellm.integrations.langfuse.langfuse_prompt_management import langfuse_client_init +from litellm.integrations.langfuse.langfuse_sdk import ( + flush_langfuse_tracing, + resolve_observation_id, + resolve_trace_id, +) +from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache as LegacyDynamicLoggingCache +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.thread_pool_executor import MAX_THREADS, executor +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +_EXPORT_URL: Final = "https://langfuse.unit/api/public/otel/v1/traces" +_VOLATILE_ATTRIBUTES: Final = frozenset( + { + "langfuse.observation.completion_start_time", + "langfuse.observation.cost_details", + "langfuse.observation.metadata.applied_guardrails", + "langfuse.observation.metadata.cache_hit", + "langfuse.observation.metadata.hidden_params", + "langfuse.observation.metadata.litellm_call_id", + "langfuse.observation.metadata.litellm_response_cost", + "langfuse.observation.metadata.requester_metadata", + "langfuse.observation.metadata.response_id", + "langfuse.observation.metadata.usage_object", + "langfuse.observation.usage_details", + } +) + + +def _decode_attribute(value: AnyValue) -> object: + match value.WhichOneof("value"): + case "string_value": + try: + return json.loads(value.string_value) + except json.JSONDecodeError: + return value.string_value + case "bool_value": + return value.bool_value + case "int_value": + return value.int_value + case "double_value": + return value.double_value + case "array_value": + return [_decode_attribute(item) for item in value.array_value.values] + case _: + return None + + +def _exported_spans(route: respx.Route) -> tuple[dict[str, object], ...]: + exported: Final = tuple( + ExportTraceServiceRequest.FromString(call.request.content) + for call in route.calls + ) + return tuple( + { + "trace_id": span.trace_id.hex(), + "span_id": span.span_id.hex(), + "parent_span_id": span.parent_span_id.hex() or None, + "name": span.name, + "attributes": { + attribute.key: _decode_attribute(attribute.value) + for attribute in span.attributes + }, + } + for request in exported + for resource in request.resource_spans + for scope in resource.scope_spans + for span in scope.spans + ) + + +def _comparable(span: Mapping[str, object]) -> dict[str, object]: + attributes: Final = TypeAdapter(dict[str, object]).validate_python(span["attributes"]) + return { + "name": span["name"], + "parent_span_id": span["parent_span_id"], + "attributes": { + key: value for key, value in sorted(attributes.items()) if key not in _VOLATILE_ATTRIBUTES + }, + } + + +def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) -> dict[str, object]: + otel_trace_id: Final = resolve_trace_id(trace_id) + generations: Final = tuple( + span + for span in _exported_spans(route) + if span["trace_id"] == otel_trace_id + and TypeAdapter(dict[str, object]).validate_python(span["attributes"]).get("langfuse.observation.type") + == "generation" + ) + assert len(generations) == 1 + expected_path: Final = Path(__file__).with_name("langfuse_expected_request_body") / fixture_name + expected: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text())) + actual: Final = _comparable(generations[0]) + attributes: Final = TypeAdapter(dict[str, object]).validate_python(generations[0]["attributes"]) + cost_details: Final = TypeAdapter(dict[str, float]).validate_python(attributes["langfuse.observation.cost_details"]) + usage_details: Final = TypeAdapter(dict[str, int]).validate_python(attributes["langfuse.observation.usage_details"]) + assert cost_details["total"] >= 0 + assert usage_details["input"] >= 0 + assert usage_details["output"] >= 0 + assert usage_details["total"] >= usage_details["input"] + assert usage_details["total"] >= usage_details["output"] + expected_comparable: Final = _comparable(expected) + assert actual == expected_comparable, json.dumps( + {"actual": actual, "expected": expected_comparable}, indent=2, sort_keys=True + ) + return actual + + +def _drain_logging_executor() -> None: + workers: Final = MAX_THREADS + gate: Final = threading.Barrier(workers + 1) + futures: Final = tuple(executor.submit(gate.wait) for _ in range(workers)) + gate.wait() + for future in futures: + future.result() + + +async def _flush_langfuse() -> None: + loop: Final = asyncio.get_running_loop() + scheduled_callbacks: Final[asyncio.Future[None]] = loop.create_future() + loop.call_soon(scheduled_callbacks.set_result, None) + await scheduled_callbacks + await GLOBAL_LOGGING_WORKER.flush() + _drain_logging_executor() + assert flush_langfuse_tracing() + + +@pytest.fixture +def langfuse_export( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> Iterator[tuple[respx.Route, str]]: + trace_id: Final = f"litellm-migration-{uuid4()}" + route: Final = respx_mock.post(_EXPORT_URL).respond(status_code=200) + monkeypatch.setenv("LANGFUSE_HOST", "https://langfuse.unit") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", f"{trace_id}-public") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", f"{trace_id}-secret") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + client_cache: Final = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache) + http_client: Final = httpx.Client(transport=httpx.MockTransport(respx_mock.handler)) + client_cache.set_cache( + key="httpx_client", + value=HTTPHandler(client=http_client), + litellm_owned_client=False, + ) + monkeypatch.setattr(litellm, "success_callback", ["langfuse"]) + monkeypatch.setattr(litellm_logging, "langFuseLogger", None) + monkeypatch.setattr(litellm_logging, "in_memory_dynamic_logger_cache", LegacyDynamicLoggingCache()) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + langfuse_client_init.cache_clear() + yield route, trace_id + logger: Final = litellm_logging.in_memory_dynamic_logger_cache.get_cache( + credentials={}, service_name="langfuse" + ) + if isinstance(logger, LangFuseLogger): + logger.stop() + client_cache.flush_cache() + http_client.close() + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion(langfuse_export: tuple[respx.Route, str]) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={"trace_id": trace_id}, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_tags(langfuse_export: tuple[respx.Route, str]) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={"trace_id": trace_id, "tags": ["test_tag", "test_tag_2"]}, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_tags.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_tags_stream(langfuse_export: tuple[respx.Route, str]) -> None: + route, trace_id = langfuse_export + response: Final = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={"trace_id": trace_id, "tags": ["test_tag_stream", "test_tag_2_stream"]}, + stream=True, + ) + chunks: Final = tuple([chunk async for chunk in response]) + assert chunks + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_tags_stream.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_generation_id_metadata_names_the_exported_observation( + langfuse_export: tuple[respx.Route, str], +) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={"trace_id": trace_id, "generation_id": "my-generation"}, + ) + await _flush_langfuse() + generation: Final = next(span for span in _exported_spans(route) if span["trace_id"] == resolve_trace_id(trace_id)) + assert generation["span_id"] == resolve_observation_id("my-generation") + _expected_generation(route, "completion.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_langfuse_metadata( + langfuse_export: tuple[respx.Route, str], +) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={ + "trace_id": trace_id, + "tags": ["test_tag", "test_tag_2"], + "generation_name": "test_generation_name", + "parent_observation_id": "test_parent_observation_id", + "version": "test_version", + "trace_user_id": "test_user_id", + "session_id": "test_session_id", + "trace_name": "test_trace_name", + "trace_metadata": {"test_key": "test_value"}, + "trace_version": "test_trace_version", + "trace_release": "test_trace_release", + }, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_langfuse_metadata.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_with_non_serializable_metadata( + langfuse_export: tuple[respx.Route, str], +) -> None: + from pydantic import BaseModel, ConfigDict + + class UserPreferences(BaseModel): + model_config = ConfigDict(frozen=True) + favorite_colors: frozenset[str] + settings: dict[str, object] + + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={ + "user_prefs": UserPreferences(favorite_colors=frozenset({"red", "blue"}), settings={"theme": "dark"}), + "nested_set": {"inner_set": frozenset({1, 2, 3})}, + "trace_id": trace_id, + }, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_complex_metadata.json", trace_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "fixture_name"), + [ + ({"a": 1, "b": 2, "c": 3}, "simple_metadata.json"), + ({"a": {"nested_a": 1}, "b": {"nested_b": 2}}, "nested_metadata.json"), + ({"a": [1, 2, 3], "b": {4, 5, 6}}, "simple_metadata2.json"), + ({"a": (1, 2), "b": frozenset({3, 4}), "c": {"d": [5, 6]}}, "simple_metadata3.json"), + ({"lock": "non-copyable"}, "metadata_with_lock.json"), + ({"func": "callable"}, "metadata_with_function.json"), + ({"list": ["list", "not", "a", "dict"]}, "complex_metadata_2.json"), + ({}, "empty_metadata.json"), + ], +) +async def test_langfuse_logging_with_various_metadata_types( + langfuse_export: tuple[respx.Route, str], metadata: dict[str, object], fixture_name: str +) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response="Hello! How can I assist you today?", + metadata={**metadata, "trace_id": trace_id}, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, fixture_name, trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_malformed_llm_response( + langfuse_export: tuple[respx.Route, str], +) -> None: + route, trace_id = langfuse_export + response: Final = litellm.ModelResponse( + choices=[], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), + model="gpt-3.5-turbo", + object="chat.completion", + created=1723081200, + ).model_dump() + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + mock_response=response, + metadata={"trace_id": trace_id}, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_no_choices.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_bedrock_llm_response( + langfuse_export: tuple[respx.Route, str], +) -> None: + route, trace_id = langfuse_export + response: Final = litellm.ModelResponse( + choices=[], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), + model="anthropic.claude-haiku-4-5-20251001-v1:0", + object="chat.completion", + created=1723081200, + ).model_dump() + await litellm.acompletion( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + mock_response=response, + metadata={"trace_id": trace_id}, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-key", + aws_region="us-east-1", + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_bedrock_call.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_completion_with_vertex_llm_response( + langfuse_export: tuple[respx.Route, str], +) -> None: + route, trace_id = langfuse_export + response: Final = litellm.ModelResponse( + choices=[], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), + model="vertex/gemini-3-flash-preview", + object="chat.completion", + created=1723081200, + ).model_dump() + await litellm.acompletion( + model="vertex_ai/gemini-3-flash-preview", + messages=[{"role": "user", "content": "Hello!"}], + mock_response=response, + metadata={"trace_id": trace_id}, + vertex_credentials="my-mock-credentials", + api_key="my-mock-credentials-2", + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_vertex_call.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_logging_vllm_embedding( + langfuse_export: tuple[respx.Route, str], respx_mock: respx.MockRouter +) -> None: + route, trace_id = langfuse_export + embedding_url: Final = "http://my-fake-vllm.com/v1/embeddings" + upstream: Final = respx_mock.post(embedding_url).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "BAAI/bge-small-en-v1.5", + "usage": {"prompt_tokens": 10, "total_tokens": 10}, + }, + ) + ) + await litellm.aembedding( + model="hosted_vllm/BAAI/bge-small-en-v1.5", + input=["Hello from litellm!"], + api_base="http://my-fake-vllm.com/v1", + metadata={"trace_id": trace_id}, + ) + await _flush_langfuse() + assert upstream.call_count == 1 + expected_path: Final = Path(__file__).with_name("langfuse_expected_request_body") / "embedding_with_vllm.json" + assert TypeAdapter(dict[str, object]).validate_json(upstream.calls.last.request.content) == json.loads( + expected_path.read_text() + ) + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_langfuse_logging_with_router(langfuse_export: tuple[respx.Route, str]) -> None: + route, trace_id = langfuse_export + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "mock_response": "Hello! How can I assist you today?", + "api_key": "test_api_key", + }, + } + ] + ) + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + metadata={"trace_id": trace_id}, + ) + await _flush_langfuse() + assert route.call_count == 1 + _expected_generation(route, "completion_with_router.json", trace_id) + + +@pytest.mark.asyncio +async def test_langfuse_e2e_sync(langfuse_export: tuple[respx.Route, str]) -> None: + route, trace_id = langfuse_export + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "sync callback route"}], + mock_response="Hello from litellm", + metadata={"trace_id": trace_id}, + ) + await _flush_langfuse() + assert route.call_count == 1 + assert str(route.calls[0].request.url).endswith("/api/public/otel/v1/traces") + assert _exported_spans(route) diff --git a/tests/unit/integrations/test_sqs.py b/tests/unit/integrations/test_sqs.py index 5fd923362af..8a4279d5ddf 100644 --- a/tests/unit/integrations/test_sqs.py +++ b/tests/unit/integrations/test_sqs.py @@ -2,37 +2,67 @@ import asyncio import base64 import json import os -from unittest.mock import AsyncMock, MagicMock +from collections.abc import Iterator +from datetime import datetime, timezone +from typing import Final, cast +from urllib.parse import parse_qs +from unittest.mock import MagicMock +import httpx import pytest +import respx +import litellm from litellm.integrations.sqs import SQSLogger +from litellm.types.utils import StandardLoggingPayload + + +@pytest.fixture(autouse=True) +def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _fixed_signing_time() -> datetime: + return datetime(2025, 1, 1, tzinfo=timezone.utc) @pytest.mark.asyncio -async def test_async_send_batch_does_not_await_send_directly(monkeypatch): +async def test_async_send_batch_dispatches_queued_payload( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr("botocore.auth.get_current_datetime", _fixed_signing_time) monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) - spawned = [] - real_create_task = asyncio.create_task + queue_url: Final = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" + sent: Final = asyncio.Event() + payload: Final = cast( + StandardLoggingPayload, + { + "model": "openai/gpt-5.5", + "messages": [{"role": "user", "content": "queued request"}], + }, + ) - def spy_create_task(coro, *args, **kwargs): - task = real_create_task(coro, *args, **kwargs) - spawned.append(task) - return task + async def respond(_request: httpx.Request) -> httpx.Response: + sent.set() + return httpx.Response(200) - monkeypatch.setattr(asyncio, "create_task", spy_create_task) + route: Final = respx_mock.post(queue_url).mock(side_effect=respond) + logger: Final = SQSLogger( + sqs_queue_url=queue_url, + sqs_region_name="us-east-1", + sqs_aws_access_key_id="unit-access-key", + sqs_aws_secret_access_key="unit-secret-key", + ) + logger.log_queue = [payload] - logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2") - logger.async_send_message = AsyncMock() - logger.log_queue = [{"log": 1}, {"log": 2}] + await logger.async_send_batch() + await sent.wait() - try: - await logger.async_send_batch() - assert logger.async_send_message.await_count == 0 - finally: - for task in spawned: - task.cancel() - await asyncio.gather(*spawned, return_exceptions=True) + assert route.call_count == 1 + assert json.loads(parse_qs(route.calls.last.request.content.decode())["MessageBody"][0]) == payload def test_sqs_logger_init_without_encryption(monkeypatch): @@ -139,6 +169,89 @@ async def test_strip_base64_removes_file_and_nontext_entries(): assert content.get("type") == "text" +@pytest.mark.asyncio +async def test_async_send_message_posts_the_queue_url_and_encoded_payload( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr("botocore.auth.get_current_datetime", _fixed_signing_time) + monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) + monkeypatch.setattr(asyncio, "create_task", MagicMock()) + queue_url: Final = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" + payload: Final = cast( + StandardLoggingPayload, + { + "model": "openai/gpt-5.5", + "messages": [{"role": "user", "content": "log this request"}], + "response": {"choices": [{"message": {"content": "logged"}}]}, + }, + ) + route: Final = respx_mock.post(queue_url).mock(return_value=httpx.Response(200)) + logger: Final = SQSLogger( + sqs_queue_url=queue_url, + sqs_region_name="us-east-1", + sqs_aws_access_key_id="unit-access-key", + sqs_aws_secret_access_key="unit-secret-key", + ) + + await logger.async_send_message(payload) + + body: Final = parse_qs(route.calls.last.request.content.decode()) + assert route.call_count == 1 + assert body["Action"] == ["SendMessage"] + assert body["Version"] == ["2012-11-05"] + assert json.loads(body["MessageBody"][0]) == payload + + +@pytest.mark.asyncio +async def test_async_send_message_handles_a_rejected_queue_request( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr("botocore.auth.get_current_datetime", _fixed_signing_time) + monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) + monkeypatch.setattr(asyncio, "create_task", MagicMock()) + queue_url: Final = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" + payload: Final = cast( + StandardLoggingPayload, + { + "model": "openai/gpt-5.5", + "messages": [{"role": "user", "content": "log this request"}], + "response": {"choices": [{"message": {"content": "logged"}}]}, + }, + ) + route: Final = respx_mock.post(queue_url).mock(return_value=httpx.Response(403, text="rejected")) + logger: Final = SQSLogger( + sqs_queue_url=queue_url, + sqs_region_name="us-east-1", + sqs_aws_access_key_id="unit-access-key", + sqs_aws_secret_access_key="unit-secret-key", + ) + + await logger.async_send_message(payload) + + body: Final = parse_qs(route.calls.last.request.content.decode()) + assert route.call_count == 1 + assert json.loads(body["MessageBody"][0]) == payload + + +@pytest.mark.asyncio +async def test_strip_base64_returns_text_only_messages_unchanged() -> None: + logger: Final = SQSLogger(sqs_strip_base64_files=True) + payload: Final = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello world"}, + {"type": "text", "text": "Keep this sentence"}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": "Response"}]}, + ] + } + + assert await logger._strip_base64_from_messages(payload) == payload + + @pytest.mark.asyncio async def test_strip_base64_handles_empty_or_missing_messages(): logger = SQSLogger(sqs_strip_base64_files=True) @@ -211,27 +324,3 @@ async def test_strip_base64_recursive_redaction(): if isinstance(c, dict): s = json.dumps(c).lower() assert "base64," not in s, f"Found real base64 blob in: {s}" - - -@pytest.mark.asyncio -async def test_async_health_check_healthy(monkeypatch): - monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) - monkeypatch.setattr(asyncio, "create_task", MagicMock()) - logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2") - logger.async_send_message = AsyncMock(return_value=None) - - result = await logger.async_health_check() - assert result["status"] == "healthy" - assert result.get("error_message") is None - - -@pytest.mark.asyncio -async def test_async_health_check_unhealthy(monkeypatch): - monkeypatch.setattr("litellm.aws_sqs_callback_params", {}) - monkeypatch.setattr(asyncio, "create_task", MagicMock()) - logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2") - logger.async_send_message = AsyncMock(side_effect=Exception("boom")) - - result = await logger.async_health_check() - assert result["status"] == "unhealthy" - assert "boom" in (result.get("error_message") or "") diff --git a/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py b/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py index 20d3a310bdb..7096c5fa76a 100644 --- a/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py +++ b/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py @@ -109,8 +109,8 @@ async def test_completion_with_vector_store_ids_prepends_the_kb_context_block(re @pytest.mark.asyncio async def test_streaming_completion_carries_the_search_results_on_a_chunk_delta(respx_mock: respx.MockRouter) -> None: - _kb_route(respx_mock) - respx_mock.post(_ANTHROPIC_URL).mock( + kb: Final = _kb_route(respx_mock) + anthropic: Final = respx_mock.post(_ANTHROPIC_URL).mock( return_value=httpx.Response(200, text=_ANTHROPIC_STREAM, headers={"content-type": "text/event-stream"}) ) @@ -128,6 +128,8 @@ async def test_streaming_completion_carries_the_search_results_on_a_chunk_delta( if choice.delta.provider_specific_fields and "search_results" in choice.delta.provider_specific_fields ) + assert kb.call_count == 1 + assert anthropic.call_count == 1 assert len(chunks) > 0 assert len(annotated) >= 1 assert annotated[0][0]["object"] == "vector_store.search_results.page" @@ -137,7 +139,9 @@ async def test_streaming_completion_carries_the_search_results_on_a_chunk_delta( @pytest.mark.asyncio async def test_file_search_filters_reach_the_kb_as_a_bedrock_equals_filter(respx_mock: respx.MockRouter) -> None: kb: Final = _kb_route(respx_mock) - respx_mock.post(_ANTHROPIC_URL).mock(return_value=httpx.Response(200, json=_ANTHROPIC_MESSAGE)) + anthropic: Final = respx_mock.post(_ANTHROPIC_URL).mock( + return_value=httpx.Response(200, json=_ANTHROPIC_MESSAGE) + ) response: Final = await litellm.acompletion( model="anthropic/claude-haiku-4-5-20251001", @@ -156,6 +160,15 @@ async def test_file_search_filters_reach_the_kb_as_a_bedrock_equals_filter(respx _sent_body(kb)["retrievalConfiguration"] ) assert retrieval["vectorSearchConfiguration"]["filter"] == {"equals": {"key": "user_id", "value": "fake-user-id"}} + assert _sent_body(anthropic)["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"{_PREFIX}{_KB_TEXT}\n\n"}, + {"type": "text", "text": "what is litellm?"}, + ], + } + ] assert response.choices[0].message.content == "LiteLLM simplifies LLM access." diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 8406166b772..dc486d953b1 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -4241,6 +4241,32 @@ def test_thinking_with_max_completion_tokens(): assert result["thinking"]["budget_tokens"] == 5000 +def test_anthropic_messages_converse_thinking_budget_is_preserved_in_request_body(): + model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={ + "thinking": {"type": "enabled", "budget_tokens": 1025}, + "max_tokens": 1026, + }, + optional_params={}, + model=model, + drop_params=False, + ) + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "What is 2+2?"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == { + "type": "enabled", + "budget_tokens": 1025, + } + + def test_drop_thinking_param_when_thinking_blocks_missing(): """ Test that thinking param is dropped when modify_params=True and diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index 9cd27a79fe0..2cc03949ff1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -6,8 +6,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import TypeAdapter import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy.pass_through_endpoints.streaming_handler import ( @@ -16,7 +18,10 @@ from litellm.proxy.pass_through_endpoints.streaming_handler import ( from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) -from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + PassthroughStandardLoggingPayload, +) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -203,6 +208,65 @@ async def test_handle_logging_runs_async_handler_for_passthrough(): mock_sync.assert_not_called() +@pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") +@pytest.mark.asyncio +async def test_unhandled_route_passes_passthrough_payload_to_custom_logger(): + class _PayloadLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payload: PassthroughStandardLoggingPayload | None = None + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + self.payload = TypeAdapter(PassthroughStandardLoggingPayload).validate_python( + kwargs["passthrough_logging_payload"] + ) + + start_time: Final = datetime(2025, 1, 2, 3, 4, 5, tzinfo=timezone.utc) + end_time: Final = datetime(2025, 1, 2, 3, 4, 9, tzinfo=timezone.utc) + url: Final = "https://upstream.example/v1/moderations" + request_body: Final = {"model": "omni-moderation-latest", "input": "check this"} + response_body: Final = {"id": "moderation-1", "results": [{"flagged": False}]} + payload: Final = PassthroughStandardLoggingPayload( + url=url, + request_body=request_body, + request_method="POST", + response_body=response_body, + ) + logger: Final = _PayloadLogger() + logging_obj: Final = LiteLLMLoggingObj( + model="omni-moderation-latest", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=start_time, + litellm_call_id="moderation-1", + function_id="moderations", + dynamic_async_success_callbacks=[logger], + ) + response: Final = httpx.Response(200, request=httpx.Request("POST", url), json=response_body) + + await PassThroughEndpointLogging().pass_through_async_success_handler( + httpx_response=response, + response_body=response_body, + logging_obj=logging_obj, + url_route="/v1/moderations", + result="", + start_time=start_time, + end_time=end_time, + cache_hit=False, + request_body=request_body, + passthrough_logging_payload=payload, + ) + + assert logger.payload == payload + + @pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") def test_convert_raw_bytes_to_str_lines(): """