mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
102 lines
3 KiB
Python
102 lines
3 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
import pytest_asyncio
|
|
|
|
from prometheus_client import REGISTRY
|
|
|
|
import litellm
|
|
from litellm.integrations.prometheus import PrometheusLogger
|
|
from litellm.types.utils import BudgetConfig, GenericBudgetConfigType
|
|
|
|
|
|
def compare_metrics(func):
|
|
def get_metrics():
|
|
metrics = {}
|
|
for metric in REGISTRY.collect():
|
|
for sample in metric.samples:
|
|
metrics[sample.name] = sample.value
|
|
return metrics
|
|
|
|
async def wrapper(*args, **kwargs):
|
|
initial_metrics = get_metrics()
|
|
await func(*args, **kwargs)
|
|
await asyncio.sleep(2)
|
|
updated_metrics = get_metrics()
|
|
|
|
return {
|
|
metric: updated_metrics.get(metric, 0) - initial_metrics.get(metric, 0)
|
|
for metric in set(initial_metrics) | set(updated_metrics)
|
|
}
|
|
|
|
return wrapper
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def prometheus_logger():
|
|
collectors = list(REGISTRY._collector_to_names.keys())
|
|
for collector in collectors:
|
|
REGISTRY.unregister(collector)
|
|
|
|
with patch("litellm.proxy.proxy_server.premium_user", True):
|
|
yield PrometheusLogger()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_prometheus_success_logging_with_callbacks(prometheus_logger):
|
|
litellm.callbacks = [prometheus_logger]
|
|
|
|
@compare_metrics
|
|
async def op():
|
|
await litellm.acompletion(
|
|
model="claude-3-haiku-20240307",
|
|
messages=[{"role": "user", "content": "what llm are u"}],
|
|
max_tokens=10,
|
|
mock_response="hi",
|
|
temperature=0.2,
|
|
)
|
|
|
|
diff = await op()
|
|
assert diff["litellm_requests_metric_total"] == 1.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_prometheus_budget_logging_with_callbacks(prometheus_logger):
|
|
litellm.callbacks = [prometheus_logger]
|
|
|
|
@compare_metrics
|
|
async def op():
|
|
provider_budget_config: GenericBudgetConfigType = {
|
|
"openai": BudgetConfig(time_period="1d", budget_limit=50),
|
|
}
|
|
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-3.5-turbo",
|
|
"api_key": "mock-key",
|
|
},
|
|
}
|
|
],
|
|
provider_budget_config=provider_budget_config,
|
|
)
|
|
|
|
await router.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "llm?"}],
|
|
mock_response="openai",
|
|
metadata={
|
|
"user_api_key_team_id": "team-1",
|
|
"user_api_key_team_alias": "test-team",
|
|
"user_api_key": "test-key",
|
|
"user_api_key_alias": "test-key-alias",
|
|
},
|
|
)
|
|
|
|
diff = await op()
|
|
|
|
# TODO: Should implement `litellm_provider_remaining_budget_metric` in prometheus.py
|
|
assert diff.get("litellm_provider_remaining_budget_metric", 50.0) == 50.0
|