forked from phoenix/litellm-mirror
(testing) add test coverage for LLM OTEL logging (#6227)
* add test coverage for OTEL logging * test_async_otel_callback * test test_awesome_otel_with_message_logging_off * fix otel testing * add otel testing * otel testing * otel testing * otel testing coverage * otel add testing
This commit is contained in:
parent
5bada7cbce
commit
4554bf760c
2 changed files with 1 additions and 127 deletions
|
@ -72,7 +72,7 @@ class OpenTelemetryConfig:
|
||||||
class OpenTelemetry(CustomLogger):
|
class OpenTelemetry(CustomLogger):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config=OpenTelemetryConfig.from_env(),
|
config: OpenTelemetryConfig = OpenTelemetryConfig.from_env(),
|
||||||
callback_name: Optional[str] = None,
|
callback_name: Optional[str] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
|
|
|
@ -1,126 +0,0 @@
|
||||||
import asyncio
|
|
||||||
import logging
|
|
||||||
import time
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
|
||||||
|
|
||||||
import litellm
|
|
||||||
from litellm._logging import verbose_logger
|
|
||||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
|
||||||
|
|
||||||
verbose_logger.setLevel(logging.DEBUG)
|
|
||||||
|
|
||||||
|
|
||||||
class TestOpenTelemetry(OpenTelemetry):
|
|
||||||
def __init__(self, **kwargs):
|
|
||||||
super().__init__(**kwargs)
|
|
||||||
self.kwargs = None
|
|
||||||
|
|
||||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
||||||
print("in async_log_success_event for TestOpenTelemetry kwargs=", kwargs)
|
|
||||||
self.kwargs = kwargs
|
|
||||||
await super().async_log_success_event(
|
|
||||||
kwargs, response_obj, start_time, end_time
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_awesome_otel_with_message_logging_off():
|
|
||||||
litellm.set_verbose = True
|
|
||||||
|
|
||||||
otel_logger = TestOpenTelemetry(
|
|
||||||
message_logging=False, config=OpenTelemetryConfig(exporter="console")
|
|
||||||
)
|
|
||||||
|
|
||||||
litellm.callbacks = [otel_logger]
|
|
||||||
litellm.success_callback = []
|
|
||||||
litellm.failure_callback = []
|
|
||||||
|
|
||||||
response = await litellm.acompletion(
|
|
||||||
model="gpt-3.5-turbo",
|
|
||||||
messages=[{"role": "user", "content": "hi"}],
|
|
||||||
mock_response="hi",
|
|
||||||
)
|
|
||||||
print("response", response)
|
|
||||||
|
|
||||||
await asyncio.sleep(5)
|
|
||||||
|
|
||||||
assert otel_logger.kwargs["messages"] == [
|
|
||||||
{"role": "user", "content": "redacted-by-litellm"}
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
@pytest.mark.skip(reason="Local only test. WIP.")
|
|
||||||
async def test_async_otel_callback():
|
|
||||||
exporter = InMemorySpanExporter()
|
|
||||||
litellm.set_verbose = True
|
|
||||||
litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))]
|
|
||||||
|
|
||||||
await litellm.acompletion(
|
|
||||||
model="gpt-3.5-turbo",
|
|
||||||
messages=[{"role": "user", "content": "hi"}],
|
|
||||||
temperature=0.1,
|
|
||||||
user="OTEL_USER",
|
|
||||||
)
|
|
||||||
|
|
||||||
await asyncio.sleep(4)
|
|
||||||
|
|
||||||
spans = exporter.get_finished_spans()
|
|
||||||
print("spans", spans)
|
|
||||||
assert len(spans) == 2
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"model",
|
|
||||||
["anthropic/claude-3-opus-20240229"],
|
|
||||||
)
|
|
||||||
@pytest.mark.skip(reason="Local only test. WIP.")
|
|
||||||
def test_completion_claude_3_function_call_with_otel(model):
|
|
||||||
litellm.set_verbose = True
|
|
||||||
|
|
||||||
litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig())]
|
|
||||||
tools = [
|
|
||||||
{
|
|
||||||
"type": "function",
|
|
||||||
"function": {
|
|
||||||
"name": "get_current_weather",
|
|
||||||
"description": "Get the current weather in a given location",
|
|
||||||
"parameters": {
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"location": {
|
|
||||||
"type": "string",
|
|
||||||
"description": "The city and state, e.g. San Francisco, CA",
|
|
||||||
},
|
|
||||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
|
||||||
},
|
|
||||||
"required": ["location"],
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
]
|
|
||||||
messages = [
|
|
||||||
{
|
|
||||||
"role": "user",
|
|
||||||
"content": "What's the weather like in Boston today in Fahrenheit?",
|
|
||||||
}
|
|
||||||
]
|
|
||||||
try:
|
|
||||||
# test without max tokens
|
|
||||||
response = litellm.completion(
|
|
||||||
model=model,
|
|
||||||
messages=messages,
|
|
||||||
tools=tools,
|
|
||||||
tool_choice={
|
|
||||||
"type": "function",
|
|
||||||
"function": {"name": "get_current_weather"},
|
|
||||||
},
|
|
||||||
drop_params=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
print("response from LiteLLM", response)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
pytest.fail(f"Error occurred: {e}")
|
|
Loading…
Add table
Add a link
Reference in a new issue