mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
test(integration): port langfuse callbacks-in-db coverage to the local harness (#43282)
* test(integration): port langfuse callbacks-in-db coverage to the local harness Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the langfuse db rows after the exported span Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
4eb340abe0
commit
e11c3f5815
2 changed files with 121 additions and 120 deletions
|
|
@ -2,24 +2,28 @@ import base64
|
|||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
|
||||
from opentelemetry.proto.common.v1.common_pb2 import KeyValue
|
||||
from opentelemetry.proto.trace.v1.trace_pb2 import Span
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
PUBLIC_KEY: Final = "pk-lf-integration"
|
||||
SECRET_KEY: Final = "sk-lf-integration"
|
||||
PROJECTS_PATH: Final = "/api/public/projects"
|
||||
TRACES_PATH: Final = "/api/public/otel/v1/traces"
|
||||
PROMPTS_PATH: Final = "/api/public/v2/prompts/"
|
||||
STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml")
|
||||
CONFIG_SECTIONS: Final = ("litellm_settings", "environment_variables")
|
||||
LANGFUSE_ENVIRONMENT: Final = ("LANGFUSE_HOST", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY")
|
||||
_PROXY_CONFIG: Final = TypeAdapter(dict[str, object])
|
||||
_SETTINGS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
|
@ -64,9 +68,7 @@ def _text_prompt(name: str) -> Reply:
|
|||
|
||||
|
||||
def _langfuse_config(tmp_path: Path) -> Path:
|
||||
config: Final = _PROXY_CONFIG.validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
config: Final = _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text()))
|
||||
settings: Final = {
|
||||
**_SETTINGS.validate_python(config["litellm_settings"]),
|
||||
"success_callback": ["langfuse"],
|
||||
|
|
@ -86,6 +88,32 @@ def _langfuse_environment(langfuse: Wire) -> dict[str, str]:
|
|||
}
|
||||
|
||||
|
||||
def _config_rows() -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT param_name, param_value FROM "LiteLLM_Config" WHERE param_name IN (%s, %s) ORDER BY param_name',
|
||||
CONFIG_SECTIONS,
|
||||
)
|
||||
|
||||
|
||||
def _restore_config_rows(snapshot: Sequence[Mapping[str, JsonValue]]) -> None:
|
||||
saved: Final = {string_value(row["param_name"]): row["param_value"] for row in snapshot}
|
||||
for section in CONFIG_SECTIONS:
|
||||
if section not in saved:
|
||||
write_rows('DELETE FROM "LiteLLM_Config" WHERE param_name = %s', (section,))
|
||||
elif saved[section] is None:
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, NULL) '
|
||||
"ON CONFLICT (param_name) DO UPDATE SET param_value = NULL",
|
||||
(section,),
|
||||
)
|
||||
else:
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s::jsonb) '
|
||||
"ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value",
|
||||
(section, json.dumps(saved[section])),
|
||||
)
|
||||
|
||||
|
||||
def _attribute(entries: Sequence[KeyValue], key: str) -> str | list[str] | None:
|
||||
for entry in entries:
|
||||
if entry.key != key:
|
||||
|
|
@ -192,6 +220,93 @@ def test_langfuse_callback_delivers_the_generation_over_otlp_v4_with_the_caller_
|
|||
)
|
||||
|
||||
|
||||
def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_generation_over_otlp_v4(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
marker: Final = "langfusedb" + uuid.uuid4().hex
|
||||
provider_secret: Final = "synthetic-provider-secret-" + marker
|
||||
public_key: Final = "pk-lf-db-" + marker
|
||||
secret_key: Final = "sk-lf-db-" + marker
|
||||
assert "langfuse" not in STOCK_CONFIG.read_text()
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
assert request.headers["authorization"] == f"Bearer {provider_secret}"
|
||||
return _completion(marker + "-answer")
|
||||
|
||||
def langfuse(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target.startswith(PROJECTS_PATH):
|
||||
return _projects()
|
||||
return Reply(body=b"", content_type="application/x-protobuf")
|
||||
|
||||
snapshot: Final = _config_rows()
|
||||
try:
|
||||
with (
|
||||
wire_server(upstream) as provider,
|
||||
wire_server(langfuse) as destination,
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"LANGFUSE_FLUSH_INTERVAL": "1"},
|
||||
remove_environment=LANGFUSE_ENVIRONMENT,
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
candidate.post(
|
||||
"/config/update",
|
||||
{
|
||||
"litellm_settings": {"success_callback": ["langfuse"]},
|
||||
"environment_variables": {
|
||||
"LANGFUSE_HOST": destination.url,
|
||||
"LANGFUSE_PUBLIC_KEY": public_key,
|
||||
"LANGFUSE_SECRET_KEY": secret_key,
|
||||
},
|
||||
},
|
||||
)
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret)
|
||||
body: Final = candidate.post(
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": marker + "-question"}],
|
||||
"metadata": {"generation_name": marker},
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
)
|
||||
received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones
|
||||
|
||||
def exported() -> tuple[Span, ...]:
|
||||
received.extend(destination.drain())
|
||||
return tuple(span for span in _spans(received) if span.name == marker)
|
||||
|
||||
spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20)
|
||||
posts: Final = tuple(request for request in received if request.method == "POST")
|
||||
assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received]
|
||||
basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode()
|
||||
for request in posts:
|
||||
assert request.headers["authorization"] == basic
|
||||
assert request.headers["content-type"] == "application/x-protobuf"
|
||||
assert request.headers["x-langfuse-ingestion-version"] == "4"
|
||||
assert provider_secret.encode() not in request.body
|
||||
assert candidate.key.encode() not in request.body
|
||||
|
||||
attributes: Final = spans[0].attributes
|
||||
assert _attribute(attributes, "langfuse.observation.type") == "generation"
|
||||
assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"])
|
||||
assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input"))
|
||||
assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output"))
|
||||
|
||||
stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows()}
|
||||
callbacks: Final = TypeAdapter(list[str]).validate_python(
|
||||
object_value(stored["litellm_settings"]).get("success_callback") or []
|
||||
)
|
||||
assert "langfuse" in callbacks, stored
|
||||
assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored
|
||||
assert secret_key not in json.dumps(stored["environment_variables"]), stored
|
||||
finally:
|
||||
_restore_config_rows(snapshot)
|
||||
assert _config_rows() == snapshot
|
||||
|
||||
|
||||
def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_headers_off_the_client(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -1,114 +0,0 @@
|
|||
"""
|
||||
PROD TEST - DO NOT Delete this Test
|
||||
|
||||
e2e test for langfuse callback in DB
|
||||
- Add langfuse callback to DB - with /config/update
|
||||
- wait 20 seconds for the callback to be loaded into the instance
|
||||
- Make a /chat/completions request to the proxy
|
||||
- Check if the request is logged in Langfuse
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import os
|
||||
import dotenv
|
||||
from dotenv import load_dotenv
|
||||
from openai import AsyncOpenAI, APIConnectionError
|
||||
from openai.types.chat import ChatCompletion
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# used for testing
|
||||
LANGFUSE_BASE_URL = "https://exampleopenaiendpoint-production-c715.up.railway.app"
|
||||
PROXY_BASE_URL = "http://127.0.0.1:4000"
|
||||
|
||||
|
||||
async def wait_for_proxy_ready(session, timeout: int = 60):
|
||||
for _ in range(timeout):
|
||||
try:
|
||||
async with session.get(f"{PROXY_BASE_URL}/health/liveliness") as response:
|
||||
if response.status == 200:
|
||||
return
|
||||
except aiohttp.ClientError:
|
||||
pass
|
||||
await asyncio.sleep(1)
|
||||
raise RuntimeError(f"Proxy at {PROXY_BASE_URL} not ready after {timeout}s")
|
||||
|
||||
|
||||
async def config_update(session, routing_strategy=None):
|
||||
url = f"{PROXY_BASE_URL}/config/update"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
print("routing_strategy: ", routing_strategy)
|
||||
data = {
|
||||
"litellm_settings": {"success_callback": ["langfuse"]},
|
||||
"environment_variables": {
|
||||
"LANGFUSE_PUBLIC_KEY": "any-public-key",
|
||||
"LANGFUSE_SECRET_KEY": "any-secret-key",
|
||||
"LANGFUSE_HOST": LANGFUSE_BASE_URL,
|
||||
},
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print("status: ", status)
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def check_langfuse_request(response_id: str):
|
||||
async with aiohttp.ClientSession() as session:
|
||||
url = f"{LANGFUSE_BASE_URL}/langfuse/trace/{response_id}"
|
||||
async with session.get(url) as response:
|
||||
response_json = await response.json()
|
||||
assert response.status == 200, f"Expected status 200, got {response.status}"
|
||||
assert (
|
||||
response_json["exists"] == True
|
||||
), f"Request {response_id} not found in Langfuse traces"
|
||||
assert response_json["request_id"] == response_id, f"Request ID mismatch"
|
||||
|
||||
|
||||
async def make_chat_completions_request() -> ChatCompletion:
|
||||
client = AsyncOpenAI(api_key="sk-1234", base_url=PROXY_BASE_URL)
|
||||
last_error = None
|
||||
for _ in range(10):
|
||||
try:
|
||||
response = await client.chat.completions.create(
|
||||
model="fake-openai-endpoint",
|
||||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||||
)
|
||||
print(response)
|
||||
return response
|
||||
except APIConnectionError as e:
|
||||
last_error = e
|
||||
await asyncio.sleep(2)
|
||||
raise AssertionError(
|
||||
f"Proxy at {PROXY_BASE_URL} unreachable after retries: {last_error!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_langfuse_callbacks_in_db():
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# add langfuse callback to DB
|
||||
await config_update(session)
|
||||
|
||||
# wait 20 seconds for the callback to be loaded into the instance
|
||||
await asyncio.sleep(20)
|
||||
await wait_for_proxy_ready(session)
|
||||
|
||||
# make a /chat/completions request to the proxy
|
||||
response = await make_chat_completions_request()
|
||||
print(response)
|
||||
response_id = response.id
|
||||
print("response_id: ", response_id)
|
||||
|
||||
await asyncio.sleep(20)
|
||||
# check if the request is logged in Langfuse
|
||||
await check_langfuse_request(response_id)
|
||||
Loading…
Add table
Reference in a new issue