mirror of
https://github.com/hanzoai/insights.git
synced 2026-10-30 16:15:31 +00:00
zeekay
1b367e2b62
The tree was a debranded snapshot of upstream taken around 2026-02-20, and upstream has moved about five months since. This replaces it with upstream master, mechanically re-derived: bin/debrand does the rename, and the pipeline and squashed migrations are restored from the shipped tree. Not a merge. The previous import replaced the tree wholesale in a single commit, so there is no shared content history to merge against. Reproducing the rename on the new tree is what actually converges. What survives: .hanzo/workflows, the Dockerfiles and bin/docker* entrypoints, compose files, LLM.md, and insights/migrations. The squashed initial migration is load-bearing — the live database records it by name, so upstream's ~918 individual migrations cannot come back. Nine products are dropped as unreferenced. Twenty-one more that we do not run are kept anyway because the core app still imports them; unpicking those is a refactor and does not belong in a rebase. This base does not build yet, and the reasons are not in the rename: upstream now wants Python 3.13.13 against our 3.12.12 and Django 5.2 against our 4.2, and the dependency manifests still name upstream packages that we republish under our own names. Those are decisions, not renames.
211 lines
9.2 KiB
Python
211 lines
9.2 KiB
Python
"""
|
|
Integration tests for LiteLLM callbacks.
|
|
|
|
Verifies that callbacks are properly invoked during actual LLM requests.
|
|
Skipped unless ANTHROPIC_API_KEY is set.
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from anthropic import Anthropic
|
|
from openai import OpenAI
|
|
|
|
from llm_gateway.metrics.prometheus import CALLBACK_SUCCESS, TOKENS_INPUT, TOKENS_OUTPUT
|
|
|
|
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
|
|
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
|
|
ANTHROPIC_TEST_MODEL = "claude-haiku-4-5-20251001"
|
|
|
|
pytestmark = pytest.mark.xfail(strict=False, reason="Upstream provider may be rate-limited or temporarily unavailable")
|
|
|
|
|
|
class TestCallbacksFireOnAnthropicRequest:
|
|
pytestmark = pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY not set")
|
|
|
|
def test_callbacks_fire_on_non_streaming_request(self, anthropic_client: Anthropic) -> None:
|
|
initial_rate_limit = CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get()
|
|
initial_prometheus = CALLBACK_SUCCESS.labels(callback="prometheus")._value.get()
|
|
|
|
anthropic_client.messages.create(
|
|
model=ANTHROPIC_TEST_MODEL,
|
|
messages=[{"role": "user", "content": "Say 'test'"}],
|
|
max_tokens=5,
|
|
)
|
|
|
|
assert CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get() > initial_rate_limit
|
|
assert CALLBACK_SUCCESS.labels(callback="prometheus")._value.get() > initial_prometheus
|
|
|
|
def test_callbacks_fire_on_streaming_request(self, anthropic_client: Anthropic) -> None:
|
|
initial_rate_limit = CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get()
|
|
initial_prometheus = CALLBACK_SUCCESS.labels(callback="prometheus")._value.get()
|
|
|
|
with anthropic_client.messages.stream(
|
|
model=ANTHROPIC_TEST_MODEL,
|
|
messages=[{"role": "user", "content": "Say 'hi'"}],
|
|
max_tokens=5,
|
|
) as stream:
|
|
stream.get_final_text()
|
|
|
|
assert CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get() > initial_rate_limit
|
|
assert CALLBACK_SUCCESS.labels(callback="prometheus")._value.get() > initial_prometheus
|
|
|
|
def test_prometheus_callback_records_tokens(self, anthropic_client: Anthropic) -> None:
|
|
initial_input = TOKENS_INPUT.labels(
|
|
provider="anthropic", model=ANTHROPIC_TEST_MODEL, product="llm_gateway"
|
|
)._value.get()
|
|
initial_output = TOKENS_OUTPUT.labels(
|
|
provider="anthropic", model=ANTHROPIC_TEST_MODEL, product="llm_gateway"
|
|
)._value.get()
|
|
|
|
anthropic_client.messages.create(
|
|
model=ANTHROPIC_TEST_MODEL,
|
|
messages=[{"role": "user", "content": "Say 'a'"}],
|
|
max_tokens=5,
|
|
)
|
|
|
|
assert (
|
|
TOKENS_INPUT.labels(provider="anthropic", model=ANTHROPIC_TEST_MODEL, product="llm_gateway")._value.get()
|
|
> initial_input
|
|
)
|
|
assert (
|
|
TOKENS_OUTPUT.labels(provider="anthropic", model=ANTHROPIC_TEST_MODEL, product="llm_gateway")._value.get()
|
|
> initial_output
|
|
)
|
|
|
|
|
|
class TestCallbacksFireOnOpenAIRequest:
|
|
pytestmark = pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY not set")
|
|
|
|
def test_callbacks_fire_on_non_streaming_request(self, openai_client: OpenAI) -> None:
|
|
initial_rate_limit = CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get()
|
|
initial_prometheus = CALLBACK_SUCCESS.labels(callback="prometheus")._value.get()
|
|
|
|
openai_client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "Say 'test'"}],
|
|
max_tokens=5,
|
|
)
|
|
|
|
assert CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get() > initial_rate_limit
|
|
assert CALLBACK_SUCCESS.labels(callback="prometheus")._value.get() > initial_prometheus
|
|
|
|
def test_callbacks_fire_on_streaming_request(self, openai_client: OpenAI) -> None:
|
|
initial_rate_limit = CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get()
|
|
initial_prometheus = CALLBACK_SUCCESS.labels(callback="prometheus")._value.get()
|
|
|
|
stream = openai_client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "Say 'hi'"}],
|
|
max_tokens=5,
|
|
stream=True,
|
|
)
|
|
for _ in stream:
|
|
pass
|
|
|
|
assert CALLBACK_SUCCESS.labels(callback="rate_limit")._value.get() > initial_rate_limit
|
|
assert CALLBACK_SUCCESS.labels(callback="prometheus")._value.get() > initial_prometheus
|
|
|
|
|
|
class TestCallbacksReceiveCorrectData:
|
|
pytestmark = pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY not set")
|
|
|
|
def test_callback_receives_token_counts(self, anthropic_client: Anthropic) -> None:
|
|
received_data = {}
|
|
|
|
async def capture_on_success(self, kwargs, response_obj, start_time, end_time, end_user_id):
|
|
standard_logging_object = kwargs.get("standard_logging_object", {})
|
|
received_data["prompt_tokens"] = standard_logging_object.get("prompt_tokens")
|
|
received_data["completion_tokens"] = standard_logging_object.get("completion_tokens")
|
|
received_data["model"] = standard_logging_object.get("model")
|
|
received_data["provider"] = standard_logging_object.get("custom_llm_provider")
|
|
|
|
from llm_gateway.callbacks.prometheus import PrometheusCallback
|
|
|
|
with patch.object(PrometheusCallback, "_on_success", capture_on_success):
|
|
anthropic_client.messages.create(
|
|
model=ANTHROPIC_TEST_MODEL,
|
|
messages=[{"role": "user", "content": "Say 'x'"}],
|
|
max_tokens=5,
|
|
)
|
|
|
|
prompt_tokens = received_data.get("prompt_tokens")
|
|
completion_tokens = received_data.get("completion_tokens")
|
|
assert prompt_tokens is not None
|
|
assert prompt_tokens > 0
|
|
assert completion_tokens is not None
|
|
assert completion_tokens > 0
|
|
assert "claude" in received_data.get("model", "").lower()
|
|
assert received_data.get("provider") == "anthropic"
|
|
|
|
def test_callback_receives_streaming_token_counts(self, anthropic_client: Anthropic) -> None:
|
|
received_data = {}
|
|
|
|
async def capture_on_success(self, kwargs, response_obj, start_time, end_time, end_user_id):
|
|
standard_logging_object = kwargs.get("standard_logging_object", {})
|
|
received_data["prompt_tokens"] = standard_logging_object.get("prompt_tokens")
|
|
received_data["completion_tokens"] = standard_logging_object.get("completion_tokens")
|
|
|
|
from llm_gateway.callbacks.prometheus import PrometheusCallback
|
|
|
|
with patch.object(PrometheusCallback, "_on_success", capture_on_success):
|
|
with anthropic_client.messages.stream(
|
|
model=ANTHROPIC_TEST_MODEL,
|
|
messages=[{"role": "user", "content": "Say 'y'"}],
|
|
max_tokens=5,
|
|
) as stream:
|
|
stream.get_final_text()
|
|
|
|
assert received_data.get("prompt_tokens") is not None
|
|
assert received_data.get("completion_tokens") is not None
|
|
|
|
|
|
class TestEndUserIdExtraction:
|
|
pytestmark = pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY not set")
|
|
|
|
def test_openai_user_param_is_not_used_as_end_user(self, openai_client: OpenAI) -> None:
|
|
"""Verify that the 'user' parameter in OpenAI requests is NOT used as end_user_id.
|
|
|
|
The end_user_id is always the authenticated user's ID to prevent rate limit poisoning.
|
|
"""
|
|
received_data: dict[str, object] = {}
|
|
|
|
async def capture_on_success(self, kwargs, response_obj, start_time, end_time, end_user_id):
|
|
received_data["end_user_id"] = end_user_id
|
|
received_data["user_in_kwargs"] = kwargs.get("user")
|
|
|
|
from llm_gateway.callbacks.rate_limiting import RateLimitCallback
|
|
|
|
with patch.object(RateLimitCallback, "_on_success", capture_on_success):
|
|
openai_client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "Say 'test'"}],
|
|
max_tokens=5,
|
|
user="test-end-user-123",
|
|
)
|
|
|
|
# end_user_id is always the authenticated user's ID (123 from conftest fixture),
|
|
# never the client-provided 'user' param
|
|
assert received_data.get("end_user_id") == "123"
|
|
assert received_data.get("user_in_kwargs") == "test-end-user-123"
|
|
|
|
def test_openai_request_without_user_param(self, openai_client: OpenAI) -> None:
|
|
"""Verify that requests without 'user' parameter still have end_user_id from auth."""
|
|
received_data: dict[str, str | None] = {}
|
|
|
|
async def capture_on_success(self, kwargs, response_obj, start_time, end_time, end_user_id):
|
|
received_data["end_user_id"] = end_user_id
|
|
|
|
from llm_gateway.callbacks.rate_limiting import RateLimitCallback
|
|
|
|
with patch.object(RateLimitCallback, "_on_success", capture_on_success):
|
|
openai_client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "Say 'test'"}],
|
|
max_tokens=5,
|
|
)
|
|
|
|
# end_user_id is always the authenticated user's ID (123 from conftest fixture)
|
|
assert received_data.get("end_user_id") == "123"
|