diff --git a/tests/e2e/llm_translation/test_ocr_matrix_e2e.py b/tests/e2e/llm_translation/test_ocr_matrix_e2e.py index 5c917bfb3a2..91ec7649953 100644 --- a/tests/e2e/llm_translation/test_ocr_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_matrix_e2e.py @@ -251,17 +251,11 @@ class Case: async def run(self, credentials: Mapping[str, object]) -> OCRResponse: kwargs: Final = {**self.provider.params, **credentials} document: Final = self.document.build() - try: - response: Final = ( - await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **kwargs) - if self.call == "async" - else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **kwargs) - ) - except Exception as exc: - reason: Final = str(exc).casefold() - if any(value in reason for value in ("quota", "rate limit", "region")): - pytest.skip(f"{self.provider.id} rejected the account due to quota or region") - raise + response: Final = ( + await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **kwargs) + if self.call == "async" + else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **kwargs) + ) assert isinstance(response, OCRResponse) return response diff --git a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json index 6da1fffd329..749796e0d04 100644 --- a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_router.json @@ -21,17 +21,10 @@ "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": 20, - "total": 30, + "output": 10, + "total": 20, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0 }, diff --git a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json index ccedcfe1e9d..83f4e648b8d 100644 --- a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_tags_stream.json @@ -3,7 +3,7 @@ "parent_span_id": null, "attributes": { "langfuse.observation.cost_details": { - "total": 3.5e-05 + "total": 1.8e-05 }, "langfuse.observation.input": { "messages": [ @@ -28,9 +28,9 @@ }, "langfuse.observation.type": "generation", "langfuse.observation.usage_details": { - "input": 10, - "output": 20, - "total": 30, + "input": 9, + "output": 9, + "total": 18, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0 }, diff --git a/tests/unit/integrations/langfuse/test_langfuse.py b/tests/unit/integrations/langfuse/test_langfuse.py index a23e3f479c6..9ef6c9eda85 100644 --- a/tests/unit/integrations/langfuse/test_langfuse.py +++ b/tests/unit/integrations/langfuse/test_langfuse.py @@ -32,7 +32,6 @@ _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", @@ -41,7 +40,6 @@ _VOLATILE_ATTRIBUTES: Final = frozenset( "langfuse.observation.metadata.requester_metadata", "langfuse.observation.metadata.response_id", "langfuse.observation.metadata.usage_object", - "langfuse.observation.usage_details", } ) @@ -112,14 +110,6 @@ def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) - 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 @@ -449,6 +439,13 @@ async def test_langfuse_logging_with_router(langfuse_export: tuple[respx.Route, await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello!"}], + 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(), metadata={"trace_id": trace_id}, ) await _flush_langfuse() diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 56049c53736..9ba09085ce8 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -53,9 +53,8 @@ from litellm.types.services import ServiceLoggerPayload, ServiceTypes from collections.abc import AsyncIterator from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.types.utils import ModelResponse +from litellm.types.utils import ChatCompletionMessageToolCall, Function from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from tests.logging_callback_tests.base_test import BaseLoggingCallbackTest exporter = InMemorySpanExporter() @@ -7140,9 +7139,20 @@ async def test_arize_phoenix_creates_nested_spans_on_dedicated_provider(): @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") -class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): - def test_parallel_tool_calls(self, mock_response_obj: ModelResponse): - tool_calls = mock_response_obj.choices[0].message.tool_calls +class TestOpentelemetryUnitTests: + def test_parallel_tool_calls(self): + tool_calls: Final = [ + ChatCompletionMessageToolCall( + function=Function(arguments='{"city": "New York"}', name="get_weather"), + id="call_weather", + type="function", + ), + ChatCompletionMessageToolCall( + function=Function(arguments='{"city": "New York"}', name="get_news"), + id="call_news", + type="function", + ), + ] from litellm.integrations.opentelemetry import OpenTelemetry from litellm.proxy._types import SpanAttributes diff --git a/tests/unit/proxy/image_endpoints/test_azure_routes.py b/tests/unit/proxy/image_endpoints/test_azure_routes.py index 87cb08e9176..b6160266716 100644 --- a/tests/unit/proxy/image_endpoints/test_azure_routes.py +++ b/tests/unit/proxy/image_endpoints/test_azure_routes.py @@ -1,5 +1,4 @@ import asyncio -import os from pathlib import Path from unittest import mock @@ -102,16 +101,9 @@ def test_azure_image_generation_route(client_no_auth): def test_azure_image_edit_route(client_no_auth): litellm.turn_on_debug() client, _, mock_aimage_edit = client_no_auth - image_path = os.path.join( - os.path.dirname(__file__), - "../../../image_gen_tests/test_image.png", - ) - with open(image_path, "rb") as f: - files = {"image": ("test_image.png", f, "image/png")} - data = {"prompt": "A cute baby sea otter"} - response = client.post( - "/openai/deployments/dall-e-3/images/edits", files=files, data=data - ) + files = {"image": ("test_image.png", b"\x89PNG\r\n\x1a\nimage-bytes", "image/png")} + data = {"prompt": "A cute baby sea otter"} + response = client.post("/openai/deployments/dall-e-3/images/edits", files=files, data=data) mock_aimage_edit.assert_called_once() called_kwargs = mock_aimage_edit.call_args.kwargs diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 21ff5534e03..1fee3d0ab4a 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4216,6 +4216,53 @@ class TestAddMCPServerAtomicity: assert result.server_id == "created-1" assert result.rpm == 5 + @pytest.mark.asyncio + async def test_create_response_hides_stored_credentials(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + payload: Final = NewMCPServerRequest( + alias="echo", + url="https://echo.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret"}, + ) + admin: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + created_server: Final = generate_mock_mcp_server_db_record( + server_id="created-1", alias="echo", auth_type="api_key" + ).model_copy(update={"credentials": {"auth_value": "secret"}}) + + mock_manager: Final = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server_if_identifier_free", + AsyncMock(return_value=created_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result: Final = await add_mcp_server(payload=payload, user_api_key_dict=admin) + + assert result.server_id == "created-1" + assert result.credentials is None + assert created_server.credentials == {"auth_value": "secret"} + mock_manager.add_server.assert_awaited_once_with(created_server) + @pytest.mark.asyncio async def test_create_500s_and_skips_registry_when_db_write_fails(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import (