mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
import io
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
|
|
import pytest
|
|
from prometheus_client import REGISTRY
|
|
|
|
import litellm
|
|
from litellm import completion
|
|
from litellm._logging import verbose_logger
|
|
from litellm.integrations.prometheus import PrometheusLogger
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
|
|
verbose_logger.setLevel(logging.DEBUG)
|
|
|
|
litellm.set_verbose = True
|
|
import time
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_async_prometheus_success_logging():
|
|
run_id = str(uuid.uuid4())
|
|
litellm.set_verbose = True
|
|
litellm.success_callback = ["prometheus"]
|
|
litellm.failure_callback = ["prometheus"]
|
|
|
|
response = await litellm.acompletion(
|
|
model="claude-instant-1.2",
|
|
messages=[{"role": "user", "content": "what llm are u"}],
|
|
max_tokens=10,
|
|
mock_response="hi",
|
|
temperature=0.2,
|
|
metadata={
|
|
"id": run_id,
|
|
"tags": ["tag1", "tag2"],
|
|
"user_api_key": "6eb81e014497d89f3cc1aa9da7c2b37bda6b7fea68e4b710d33d94201e68970c",
|
|
"user_api_key_alias": "ishaans-prometheus-key",
|
|
"user_api_end_user_max_budget": None,
|
|
"litellm_api_version": "1.40.19",
|
|
"global_max_parallel_requests": None,
|
|
"user_api_key_user_id": "admin",
|
|
"user_api_key_org_id": None,
|
|
"user_api_key_team_id": "dbe2f686-a686-4896-864a-4c3924458709",
|
|
"user_api_key_team_alias": "testing-team",
|
|
},
|
|
)
|
|
print(response)
|
|
await asyncio.sleep(3)
|
|
|
|
# get prometheus logger
|
|
from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers
|
|
|
|
for callback in _in_memory_loggers:
|
|
if isinstance(callback, PrometheusLogger):
|
|
test_prometheus_logger = callback
|
|
|
|
print("done with success request")
|
|
|
|
print(
|
|
"vars of test_prometheus_logger",
|
|
vars(test_prometheus_logger.litellm_requests_metric),
|
|
)
|
|
|
|
# Get the metrics
|
|
metrics = {}
|
|
for metric in REGISTRY.collect():
|
|
for sample in metric.samples:
|
|
metrics[sample.name] = sample.value
|
|
|
|
print("metrics from prometheus", metrics)
|
|
assert metrics["litellm_requests_metric_total"] == 1.0
|
|
assert metrics["litellm_total_tokens_total"] == 30.0
|
|
assert metrics["llm_deployment_success_responses_total"] == 1.0
|
|
assert metrics["llm_deployment_total_requests_total"] == 1.0
|
|
assert metrics["llm_deployment_latency_per_output_token_bucket"] == 1.0
|