mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
Update tests
This commit is contained in:
parent
f903ef5d29
commit
82f1ce29ee
2 changed files with 10 additions and 12 deletions
|
@ -1,8 +1,6 @@
|
||||||
# What is this?
|
# What is this?
|
||||||
## Unit tests for opentelemetry integration
|
## Unit tests for opentelemetry integration
|
||||||
|
|
||||||
# What is this?
|
|
||||||
## Unit test for presidio pii masking
|
|
||||||
import sys, os, asyncio, time, random
|
import sys, os, asyncio, time, random
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
import traceback
|
import traceback
|
||||||
|
@ -17,7 +15,7 @@ sys.path.insert(
|
||||||
) # Adds the parent directory to the system path
|
) # Adds the parent directory to the system path
|
||||||
import pytest
|
import pytest
|
||||||
import litellm
|
import litellm
|
||||||
from unittest.mock import patch, MagicMock, AsyncMock
|
from unittest.mock import patch, MagicMock
|
||||||
from base_test import BaseLoggingCallbackTest
|
from base_test import BaseLoggingCallbackTest
|
||||||
from litellm.types.utils import ModelResponse
|
from litellm.types.utils import ModelResponse
|
||||||
|
|
||||||
|
@ -26,15 +24,14 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
|
||||||
def test_parallel_tool_calls(self, mock_response_obj: ModelResponse):
|
def test_parallel_tool_calls(self, mock_response_obj: ModelResponse):
|
||||||
tool_calls = mock_response_obj.choices[0].message.tool_calls
|
tool_calls = mock_response_obj.choices[0].message.tool_calls
|
||||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||||
from litellm.proxy._types import SpanAttributes
|
|
||||||
|
|
||||||
kv_pair_dict = OpenTelemetry._tool_calls_kv_pair(tool_calls)
|
kv_pair_dict = OpenTelemetry._tool_calls_kv_pair(tool_calls)
|
||||||
|
|
||||||
assert kv_pair_dict == {
|
assert kv_pair_dict == {
|
||||||
f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.arguments": '{"city": "New York"}',
|
"gen_ai.completion.0.function_call.arguments": '{"city": "New York"}',
|
||||||
f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.name": "get_weather",
|
"gen_ai.completion.0.function_call.name": "get_weather",
|
||||||
f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.arguments": '{"city": "New York"}',
|
"gen_ai.completion.1.function_call.arguments": '{"city": "New York"}',
|
||||||
f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.name": "get_news",
|
"gen_ai.completion.1.function_call.name": "get_news",
|
||||||
}
|
}
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
@ -1,8 +1,5 @@
|
||||||
import json
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
from datetime import datetime
|
|
||||||
from unittest.mock import AsyncMock
|
|
||||||
|
|
||||||
sys.path.insert(
|
sys.path.insert(
|
||||||
0, os.path.abspath("../..")
|
0, os.path.abspath("../..")
|
||||||
|
@ -10,7 +7,7 @@ sys.path.insert(
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import litellm
|
import litellm
|
||||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig, Span
|
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||||
|
@ -86,6 +83,10 @@ def validate_litellm_request(span):
|
||||||
"llm.usage.total_tokens",
|
"llm.usage.total_tokens",
|
||||||
"gen_ai.usage.completion_tokens",
|
"gen_ai.usage.completion_tokens",
|
||||||
"gen_ai.usage.prompt_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
|
# get the str of all the span attributes
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue