litellm/tests/e2e/logging/logging_client.py

"""Client for the logging e2e suite: team/key/org-scoped Langfuse OTEL callbacks,
chat (including tools), Prometheus scrape, and Langfuse observation read-back.

Holds the shared ProxyClient so the ``resources`` fixture cleans up keys, teams,
users, orgs, and models it creates. External Langfuse reads go through
``e2e_http`` (the only module allowed to call ``requests.*``).

Uses the ``langfuse_otel`` callback (OTLP to ``{host}/api/public/otel``), not
the classic ``langfuse`` SDK callback. OTEL generations land as name
``litellm_request``; correlate by unique prompt marker and ``user_api_key_alias``
in metadata. Spend is on ``calculatedTotalCost`` (StandardLogging response_cost).
"""

from __future__ import annotations

import base64
import json
import os
import time
from dataclasses import dataclass
from typing import Callable, Literal

import pytest
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError

from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation
from e2e_metadata import step
from proxy_client import ProxyClient
from e2e_http import (
    URL,
    AuthHeaders,
    require_successful_call,
    NoBody,
    StreamingResponse,
    Success,
    get,
    unwrap,
)
from models import (
    AnthropicMessagesBody,
    ChatBody,
    ChatMessage,
    ChatResponse,
    ChatTool,
    ChatToolFunction,
    KeyGenerateBody,
    KeyLoggingCallback,
    KeyLoggingCallbackVars,
    KeyMetadata,
    LiteLLMParamsBody,
    OrgDeleteBody,
    OrgNewBody,
    OrgNewResponse,
    SpendLogRow,
    TeamDeleteBody,
    TeamNewBody,
    TeamNewResponse,
    UserDeleteBody,
    UserNewBody,
    UserNewResponse,
)

# Deliberately invalid *upstream provider* key for failure-path tests.
# Not a LiteLLM virtual key; OpenAI must reject it after the proxy accepts the call.
INVALID_UPSTREAM_API_KEY = "sk-upstream-invalid-for-langfuse-e2e-only"

WEATHER_TOOL = ChatTool(
    type="function",
    function=ChatToolFunction(
        name="get_weather",
        description="Get the current weather for a city",
        parameters={
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    ),
)


class ResponsesRequestBody(BaseModel):
    """OpenAI Responses API /v1/responses request."""

    model: str
    input: str
    max_output_tokens: int
    stream: bool | None = None


class TeamCallbackBody(BaseModel):
    callback_name: Literal["langfuse_otel", "langfuse", "langsmith", "gcs"]
    callback_type: Literal["success", "failure", "success_and_failure"]
    callback_vars: dict[str, str]


class TeamCallbackResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")

    status: str


class GuardrailLitellmParams(BaseModel):
    guardrail: str
    mode: str
    default_on: bool = False
    rules: list[dict[str, object]] | None = None
    default_action: str | None = None
    on_disallowed_action: str | None = None


class GuardrailSpec(BaseModel):
    guardrail_name: str
    litellm_params: GuardrailLitellmParams


class CreateGuardrailBody(BaseModel):
    guardrail: GuardrailSpec


class CreateGuardrailResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")

    guardrail_id: str | None = None
    guardrail_name: str | None = None


class LangfuseObservation(BaseModel):
    model_config = ConfigDict(extra="ignore", populate_by_name=True)

    id: str
    trace_id: str | None = Field(default=None, alias="traceId")
    name: str | None = None
    type: str | None = None
    calculated_total_cost: float | None = Field(default=None, alias="calculatedTotalCost")
    level: str | None = None
    input: object | None = None
    output: object | None = None
    metadata: object | None = None
    usage: object | None = None
    usage_details: object | None = Field(default=None, alias="usageDetails")
    model: str | None = None


class LangfuseObservationList(BaseModel):
    model_config = ConfigDict(extra="ignore")

    data: list[LangfuseObservation] = []


class LangfuseOtelMetadata(BaseModel):
    """Langfuse stores every OTel span attribute under metadata.attributes."""

    model_config = ConfigDict(extra="ignore")

    attributes: dict[str, str] = {}


def otel_attributes(obs: LangfuseObservation) -> dict[str, str]:
    try:
        return LangfuseOtelMetadata.model_validate(obs.metadata).attributes
    except ValidationError:
        return {}


def is_otel_v2_generation(obs: LangfuseObservation, *, key_alias: str) -> bool:
    attributes = otel_attributes(obs)
    return (
        attributes.get("langfuse.observation.type") == "generation"
        and attributes.get("litellm.metadata.user_api_key_alias") == key_alias
    )


class LangfuseListParams(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    limit: int = 100
    trace_id: str | None = Field(default=None, alias="traceId")
    name: str | None = None
    from_start_time: str | None = Field(default=None, alias="fromStartTime")


@dataclass(frozen=True, slots=True)
class LangfuseCreds:
    public_key: str
    secret_key: str
    host: str

    @property
    def auth_headers(self) -> AuthHeaders:
        token = base64.b64encode(f"{self.public_key}:{self.secret_key}".encode()).decode()
        return AuthHeaders(authorization=f"Basic {token}")

    def callback_vars(self) -> dict[str, str]:
        return {
            "langfuse_public_key": self.public_key,
            "langfuse_secret_key": self.secret_key,
            "langfuse_host": self.host,
        }

    def key_logging_metadata(self) -> KeyMetadata:
        return KeyMetadata(
            logging=[
                KeyLoggingCallback(
                    callback_name="langfuse_otel",
                    callback_type="success_and_failure",
                    callback_vars=KeyLoggingCallbackVars(
                        langfuse_public_key=self.public_key,
                        langfuse_secret_key=self.secret_key,
                        langfuse_host=self.host,
                    ),
                )
            ]
        )


@dataclass(frozen=True, slots=True)
class WeaveCreds:
    """Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback.

    The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment:
    the weave_otel logger is constructed from those before the per-key vars are
    applied, so a key-scoped callback on a proxy without them never initializes.
    The per-key vars are what direct THIS key's spans at this project.
    """

    api_key: str
    project_id: str

    def key_logging_metadata(self) -> KeyMetadata:
        return KeyMetadata(
            logging=[
                KeyLoggingCallback(
                    callback_name="weave_otel",
                    callback_type="success_and_failure",
                    callback_vars=KeyLoggingCallbackVars(
                        wandb_api_key=self.api_key,
                        weave_project_id=self.project_id,
                    ),
                )
            ]
        )


def load_weave_creds() -> WeaveCreds:
    api_key = os.getenv("WANDB_API_KEY")
    project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip()
    if not (api_key and project_id):
        pytest.fail(
            "Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
            "format <entity>/<project>); missing credentials is a hard failure, not a skip"
        )
    return WeaveCreds(api_key=api_key, project_id=project_id)


def load_langfuse_creds() -> LangfuseCreds:
    public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
    secret_key = os.getenv("LANGFUSE_SECRET_KEY")
    host = (os.getenv("LANGFUSE_BASE_URL") or os.getenv("LANGFUSE_HOST") or "").rstrip("/")
    if not (public_key and secret_key and host):
        pytest.fail(
            "Langfuse e2e requires LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, and "
            "LANGFUSE_BASE_URL (or LANGFUSE_HOST); missing credentials is a hard failure, not a skip"
        )
    return LangfuseCreds(public_key=public_key, secret_key=secret_key, host=host)


def observation_spend(obs: LangfuseObservation) -> float | None:
    """Langfuse calculatedTotalCost is populated from StandardLogging response_cost."""
    return obs.calculated_total_cost


def costs_agree(expected: float, actual: float, *, rel_tol: float = 0.05) -> bool:
    """Costs agree within 5% relative (or 1e-9 absolute for near-zero)."""
    return abs(expected - actual) <= max(1e-9, abs(expected) * rel_tol)


_COMPLETION_BODY_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue])


def completion_response_id(body: str) -> str | None:
    """SpendLogs.request_id is the chat completion body id, not x-litellm-call-id."""
    if not body or body == "<streamed>":
        return None
    try:
        parsed = _COMPLETION_BODY_ADAPTER.validate_json(body)
    except ValidationError:
        return None
    raw = parsed.get("id")
    return raw if isinstance(raw, str) and raw else None


def _matches_run(obs: LangfuseObservation, *, key_alias: str, prompt_marker: str) -> bool:
    """Match a Langfuse generation for this run.

    langfuse_otel names generations ``litellm_request`` (not ``litellm:{alias}``).
    Prefer the unique prompt marker in input; fall back to key alias in metadata
    (user_api_key_alias) or the classic SDK generation name.
    """
    if prompt_marker and prompt_marker in json.dumps(obs.input, default=str):
        return True
    meta_blob = json.dumps(obs.metadata, default=str) if obs.metadata is not None else ""
    if key_alias and key_alias in meta_blob:
        return True
    if obs.name == f"litellm:{key_alias}":
        return True
    return False


def observation_mentions_tool(obs: LangfuseObservation, tool_name: str) -> bool:
    blob = json.dumps(
        {"input": obs.input, "output": obs.output, "metadata": obs.metadata},
        default=str,
    )
    return tool_name in blob


def observation_has_guardrail(obs: LangfuseObservation, *, guardrail_name: str) -> bool:
    blob = json.dumps(obs.metadata, default=str) if obs.metadata is not None else ""
    if guardrail_name in blob or "guardrail" in blob.lower():
        return True
    if obs.name is not None and "guardrail" in obs.name.lower():
        return True
    return False


@dataclass(frozen=True, slots=True)
class LoggingClient:
    proxy: ProxyClient

    @step("Generate a virtual key named {alias} with models: {models}")
    def key_with_alias(
        self,
        alias: str,
        *,
        models: list[str],
        team_id: str | None = None,
        user_id: str | None = None,
        organization_id: str | None = None,
        metadata: KeyMetadata | None = None,
    ) -> str:
        return self.proxy.generate_key(
            KeyGenerateBody(
                key_alias=alias,
                models=models,
                user_id=user_id or f"e2e-{alias}",
                team_id=team_id,
                organization_id=organization_id,
                metadata=metadata,
            )
        )

    @step("Delete the virtual key")
    def delete_key(self, key: str) -> None:
        self.proxy.delete_key(key)

    @step("Create the team {alias} with models: {models}")
    def create_team(
        self,
        alias: str,
        *,
        models: list[str],
        organization_id: str | None = None,
    ) -> str:
        return unwrap(
            self.proxy.transport.post(
                "/team/new",
                headers=self.proxy.transport.master,
                json=TeamNewBody(
                    team_alias=alias,
                    models=models,
                    organization_id=organization_id,
                ),
                response_type=TeamNewResponse,
            )
        ).team_id

    @step("Delete the team")
    def delete_team(self, team_id: str) -> None:
        _ = self.proxy.transport.post(
            "/team/delete",
            headers=self.proxy.transport.master,
            json=TeamDeleteBody(team_ids=[team_id]),
            response_type=NoBody,
        )

    @step("Create the internal user {user_email}")
    def create_user(self, *, user_email: str, user_id: str | None = None) -> str:
        return unwrap(
            self.proxy.transport.post(
                "/user/new",
                headers=self.proxy.transport.master,
                json=UserNewBody(
                    user_email=user_email,
                    user_role="internal_user",
                    user_id=user_id,
                ),
                response_type=UserNewResponse,
            )
        ).user_id

    @step("Delete the internal user")
    def delete_user(self, user_id: str) -> None:
        _ = self.proxy.transport.post(
            "/user/delete",
            headers=self.proxy.transport.master,
            json=UserDeleteBody(user_ids=[user_id]),
            response_type=NoBody,
        )

    @step("Create the organization {alias} with models: {models}")
    def create_org(self, alias: str, *, models: list[str]) -> str:
        return unwrap(
            self.proxy.transport.post(
                "/organization/new",
                headers=self.proxy.transport.master,
                json=OrgNewBody(organization_alias=alias, models=models),
                response_type=OrgNewResponse,
            )
        ).organization_id

    @step("Delete the organization")
    def delete_org(self, organization_id: str) -> None:
        _ = self.proxy.transport.delete(
            "/organization/delete",
            headers=self.proxy.transport.master,
            json=OrgDeleteBody(organization_ids=[organization_id]),
            response_type=NoBody,
        )

    @step("Add a Langfuse OTel logging callback for {callback_type} events to the team")
    def add_team_langfuse_callback(
        self,
        team_id: str,
        creds: LangfuseCreds,
        *,
        callback_type: Literal["success", "failure", "success_and_failure"] = "success_and_failure",
    ) -> None:
        response = unwrap(
            self.proxy.transport.post(
                f"/team/{team_id}/callback",
                headers=self.proxy.transport.master,
                json=TeamCallbackBody(
                    callback_name="langfuse_otel",
                    callback_type=callback_type,
                    callback_vars=creds.callback_vars(),
                ),
                response_type=TeamCallbackResponse,
            )
        )
        assert response.status == "success", (
            f"POST /team/{team_id}/callback must return status=success; got {response.status!r}"
        )

    @step("Create the tool_permission guardrail {name} that allows only the tool {allowed_tool}")
    def create_tool_permission_guardrail(self, name: str, *, allowed_tool: str) -> str:
        """Register a tool_permission guardrail that allows one tool and denies the rest."""
        response = unwrap(
            self.proxy.transport.post(
                "/guardrails",
                headers=self.proxy.transport.master,
                json=CreateGuardrailBody(
                    guardrail=GuardrailSpec(
                        guardrail_name=name,
                        litellm_params=GuardrailLitellmParams(
                            guardrail="tool_permission",
                            mode="post_call",
                            default_on=False,
                            default_action="deny",
                            on_disallowed_action="block",
                            rules=[
                                {
                                    "id": "allow-named-tool",
                                    "tool_name": allowed_tool,
                                    "decision": "allow",
                                }
                            ],
                        ),
                    )
                ),
                response_type=CreateGuardrailResponse,
            )
        )
        guardrail_id = response.guardrail_id
        assert guardrail_id, f"create guardrail returned no id: {response!r}"
        settle_propagation(time.monotonic())
        return guardrail_id

    @step("Delete the guardrail")
    def delete_guardrail(self, guardrail_id: str) -> None:
        _ = self.proxy.transport.delete(
            f"/guardrails/{guardrail_id}",
            headers=self.proxy.transport.master,
            json=NoBody(),
            response_type=NoBody,
        )

    @step("Add a deployment named {model_name} that calls {litellm_params.model}")
    def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
        return self.proxy.create_model(model_name, litellm_params)

    @step("Delete the deployment")
    def delete_model(self, model_id: str) -> None:
        self.proxy.delete_model(model_id)

    @step('Send a /chat/completions request to {model} with the prompt "{text}"')
    def chat(self, key: str, model: str, text: str) -> ChatResponse:
        return unwrap(
            self.proxy.chat(
                key,
                ChatBody(
                    model=model,
                    messages=[ChatMessage(role="user", content=text)],
                    max_tokens=64,
                ),
            )
        )

    @step('Send a /chat/completions request to {model} with stream={stream} and the prompt "{text}"')
    def chat_raw(
        self,
        key: str,
        model: str,
        text: str,
        *,
        stream: bool = False,
        tools: list[ChatTool] | None = None,
        tool_choice: str | None = None,
        guardrails: list[str] | None = None,
        max_tokens: int = 64,
    ) -> StreamingResponse:
        body = ChatBody(
            model=model,
            messages=[ChatMessage(role="user", content=text)],
            max_tokens=max_tokens,
            stream=stream,
            tools=tools,
            tool_choice=tool_choice,
            guardrails=guardrails,
        )
        if stream:
            return self.proxy.chat_stream(key, body)
        return self.proxy.transport.send(
            "/chat/completions",
            headers=self.proxy.transport.bearer(key),
            json=body,
        )

    @step('Send a /v1/messages request to {model} with stream={stream} and the prompt "{text}"')
    def messages_raw(
        self, key: str, model: str, text: str, *, max_tokens: int = 16, stream: bool = False
    ) -> StreamingResponse:
        """POST /v1/messages (Anthropic-native body): raw outcome judged by
        status/body/headers, for tests that need x-litellm-call-id. With
        ``stream=True`` the SSE body is consumed and its events counted."""
        body = AnthropicMessagesBody(
            model=model,
            max_tokens=max_tokens,
            messages=[ChatMessage(role="user", content=text)],
            stream=True if stream else None,
        )
        if stream:
            return self.proxy.transport.stream("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)
        return self.proxy.transport.send("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)

    @step('Send a /v1/responses request to {model} with stream={stream} and the prompt "{text}"')
    def responses_raw(
        self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False
    ) -> StreamingResponse:
        """POST /v1/responses (OpenAI Responses API): raw outcome judged by
        status/body/headers, for tests that need x-litellm-call-id.
        max_output_tokens caps reasoning-model output cost; a capped response is
        still a 200 and still exports the trace. With ``stream=True`` the SSE
        body is consumed and its events counted."""
        body = ResponsesRequestBody(
            model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None
        )
        if stream:
            return self.proxy.transport.stream("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)
        return self.proxy.transport.send("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)

    @step("Scrape the Prometheus metrics from /metrics")
    def scrape_metrics(self) -> str:
        return self.proxy.probe("/metrics", params=NoBody()).body

    @step("Wait for the key's spend log in /spend/logs")
    def poll_proxy_spend_for_key(
        self,
        key: str,
        *,
        response_id: str | None = None,
        require_positive_spend: bool = True,
    ) -> SpendLogRow | None:
        """Poll /spend/logs by virtual key.

        When ``response_id`` is set, only that SpendLogs.request_id may match.
        When unset, any positive-spend row for the key is accepted. Never falls
        back to an unmatched row; missing match returns None.
        """

        def _matches(row: SpendLogRow) -> bool:
            if response_id is not None and row.request_id != response_id:
                return False
            if require_positive_spend and not (row.spend is not None and row.spend > 0):
                return False
            return True

        rows = self.proxy.poll_logs_for_key(key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs))
        for row in rows:
            if _matches(row):
                return row
        return None

    @step("List observations from Langfuse")
    def list_langfuse_observations(
        self,
        creds: LangfuseCreds,
        *,
        trace_id: str | None = None,
        name: str | None = None,
        from_start_time: str | None = None,
    ) -> list[LangfuseObservation]:
        result = get(
            URL(f"{creds.host}/api/public/observations"),
            headers=creds.auth_headers,
            params=LangfuseListParams(
                limit=100,
                traceId=trace_id,
                name=name,
                fromStartTime=from_start_time,
            ),
            response_type=LangfuseObservationList,
            timeout=30.0,
        )
        match result:
            case Success(data=page):
                return page.data
            case _:
                return []

    @step("Look up the Langfuse generation for the key {key_alias}")
    def find_langfuse_observation(
        self,
        creds: LangfuseCreds,
        *,
        key_alias: str,
        prompt_marker: str,
    ) -> LangfuseObservation | None:
        # langfuse_otel generations are named litellm_request; classic SDK used
        # litellm:{key_alias}. Search both, then a recent unfiltered page.
        for name in ("litellm_request", f"litellm:{key_alias}"):
            for obs in self.list_langfuse_observations(creds, name=name):
                if _matches_run(obs, key_alias=key_alias, prompt_marker=prompt_marker):
                    return obs
        for obs in self.list_langfuse_observations(creds):
            if _matches_run(obs, key_alias=key_alias, prompt_marker=prompt_marker):
                return obs
        return None

    @step("Wait for the Langfuse generation for the key {key_alias}")
    def poll_langfuse_observation(
        self,
        creds: LangfuseCreds,
        *,
        key_alias: str,
        prompt_marker: str,
        require_positive_cost: bool = False,
    ) -> LangfuseObservation | None:
        deadline = time.monotonic() + POLL_TIMEOUT
        last: LangfuseObservation | None = None
        while time.monotonic() < deadline:
            last = self.find_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker)
            if last is not None:
                cost = observation_spend(last)
                if not require_positive_cost or (cost is not None and cost > 0):
                    return last
            time.sleep(POLL_INTERVAL)
        return last

    @step("Wait for the OTel v2 Langfuse generation for the key {key_alias}")
    def poll_langfuse_generation(
        self, creds: LangfuseCreds, *, key_alias: str, from_start_time: str
    ) -> LangfuseObservation | None:
        """The OTel v2 generation the proxy exported for one key alias since from_start_time."""
        deadline = time.monotonic() + POLL_TIMEOUT
        while time.monotonic() < deadline:
            for obs in self.list_langfuse_observations(creds, from_start_time=from_start_time):
                if is_otel_v2_generation(obs, key_alias=key_alias):
                    return obs
            time.sleep(POLL_INTERVAL)
        return None

    @step("Wait for the Langfuse trace of the key {key_alias} and every observation in it")
    def poll_langfuse_trace_observations(
        self,
        creds: LangfuseCreds,
        *,
        key_alias: str,
        prompt_marker: str,
    ) -> list[LangfuseObservation]:
        """Generation plus any sibling/child observations (guardrail spans, etc.)."""
        gen = self.poll_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker)
        if gen is None or not gen.trace_id:
            return [] if gen is None else [gen]
        return self.list_langfuse_observations(creds, trace_id=gen.trace_id) or [gen]


@step("Retry the call until the fresh key stops answering 401")
def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> StreamingResponse:
    """First successful call on a fresh key. A fresh key may briefly 401 until
    the data plane's auth cache picks it up, so retry on 401 to a deadline; a
    401 is rejected before the LLM call, so it cannot contaminate delivery or
    trace assertions. Any other failure is behavior under test and fails hard."""
    deadline = time.monotonic() + client.proxy.poll_timeout
    while True:
        outcome = send()
        if outcome.ok:
            return outcome
        if outcome.status_code != 401 or time.monotonic() >= deadline:
            require_successful_call(outcome)
        time.sleep(client.proxy.poll_interval)


def build_logging_client(proxy: ProxyClient) -> LoggingClient:
    return LoggingClient(proxy=proxy)


@step("Read the proxy's callback list from /health/readiness/details")
def readiness_details_body(client: LoggingClient) -> str:
    """/health/readiness/details, tolerating the 503 it serves while the
    ephemeral stack's DB leg blips: the recorded state the logging suites check
    here is the callback list, which the body carries either way."""
    result = client.proxy.probe("/health/readiness/details", params=NoBody())
    db_blip = result.status_code == 503 and '"db":"disconnected"' in result.body
    assert result.status_code == 200 or db_blip, (
        f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
    )
    return result.body