forked from phoenix/litellm-mirror
f59cb46e71
* fix(__init__.py): add 'watsonx_text' as mapped llm api route Fixes https://github.com/BerriAI/litellm/issues/6663 * fix(opentelemetry.py): fix passing parallel tool calls to otel Fixes https://github.com/BerriAI/litellm/issues/6677 * refactor(test_opentelemetry_unit_tests.py): create a base set of unit tests for all logging integrations - test for parallel tool call handling reduces bugs in repo * fix(__init__.py): update provider-model mapping to include all known provider-model mappings Fixes https://github.com/BerriAI/litellm/issues/6669 * feat(anthropic): support passing document in llm api call * docs(anthropic.md): add pdf anthropic call to docs + expose new 'supports_pdf_input' function * fix(factory.py): fix linting error
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
# 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}.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",
|
|
}
|
|
|
|
@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()
|