Update tests

This commit is contained in:
Alex Burgel 2025-03-14 14:53:21 -04:00
parent f903ef5d29
commit 82f1ce29ee
No known key found for this signature in database
2 changed files with 10 additions and 12 deletions

View file

@ -1,8 +1,6 @@
# What is this?
## Unit tests for opentelemetry integration
# What is this?
## Unit test for presidio pii masking
import sys, os, asyncio, time, random
from datetime import datetime
import traceback
@ -17,7 +15,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
import pytest
import litellm
from unittest.mock import patch, MagicMock, AsyncMock
from unittest.mock import patch, MagicMock
from base_test import BaseLoggingCallbackTest
from litellm.types.utils import ModelResponse
@ -26,15 +24,14 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
def test_parallel_tool_calls(self, mock_response_obj: ModelResponse):
tool_calls = mock_response_obj.choices[0].message.tool_calls
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.proxy._types import SpanAttributes
kv_pair_dict = OpenTelemetry._tool_calls_kv_pair(tool_calls)
assert kv_pair_dict == {
f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.arguments": '{"city": "New York"}',
f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.name": "get_weather",
f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.arguments": '{"city": "New York"}',
f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.name": "get_news",
"gen_ai.completion.0.function_call.arguments": '{"city": "New York"}',
"gen_ai.completion.0.function_call.name": "get_weather",
"gen_ai.completion.1.function_call.arguments": '{"city": "New York"}',
"gen_ai.completion.1.function_call.name": "get_news",
}
@pytest.mark.asyncio

View file

@ -1,8 +1,5 @@
import json
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock
sys.path.insert(
0, os.path.abspath("../..")
@ -10,7 +7,7 @@ sys.path.insert(
import pytest
import litellm
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig, Span
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
import asyncio
import logging
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@ -86,6 +83,10 @@ def validate_litellm_request(span):
"llm.usage.total_tokens",
"gen_ai.usage.completion_tokens",
"gen_ai.usage.prompt_tokens",
"gen_ai.prompt.0.role",
"gen_ai.prompt.0.content",
"gen_ai.completion.0.role",
"gen_ai.completion.0.content",
]
# get the str of all the span attributes