Spaces:
Configuration error
Configuration error
File size: 1,952 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# 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
from dotenv import load_dotenv
load_dotenv()
import os
import asyncio
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
from unittest.mock import patch, MagicMock, AsyncMock
from base_test import BaseLoggingCallbackTest
from litellm.types.utils import ModelResponse
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.value}.0.function_call.arguments": '{"city": "New York"}',
f"{SpanAttributes.LLM_COMPLETIONS.value}.0.function_call.name": "get_weather",
f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.arguments": '{"city": "New York"}',
f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.name": "get_news",
}
@pytest.mark.asyncio
async def test_opentelemetry_integration(self):
"""
Unit test to confirm the parent otel span is ended
"""
parent_otel_span = MagicMock()
litellm.callbacks = ["otel"]
await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello, world!"}],
mock_response="Hey!",
metadata={"litellm_parent_otel_span": parent_otel_span},
)
await asyncio.sleep(1)
parent_otel_span.end.assert_called_once()
|