litellm/tests/integration/streaming/faked_responses_case.py

import json
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Final

import httpx
from anthropic.types import ToolParam as AnthropicToolParam
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, Wire
from openai.types.chat import ChatCompletionFunctionToolParam
from openai.types.responses import FunctionToolParam
from pydantic import JsonValue, TypeAdapter

API_KEY: Final = "synthetic-faked-stream-key"
AZURE_API_VERSION: Final = "2025-03-01-preview"
CALL_ID: Final = "x-litellm-call-id"
NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}
FAKED_STREAM: Final[dict[str, JsonValue]] = {"supports_native_streaming": False}
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
QUESTION: Final = "What is the weather in Paris?"
ANSWER: Final = "Paris is 22 degrees Celsius with clear skies."
ARGUMENTS: Final = json.dumps({"city": "Paris", "unit": "celsius"})
SECOND_ARGUMENTS: Final = json.dumps({"city": "Lyon", "days": 3})
SQL: Final = "SELECT city FROM weather WHERE celsius > 20"
REFUSAL: Final = "I cannot share the forecast for that place."
SUMMARIES: Final = ("First, find the city.", "Then, look up its weather.")
RESPONSES_TOOLS: Final[list[JsonValue]] = [
    {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}},
    {"type": "function", "name": "get_forecast", "parameters": {"type": "object", "properties": {}}},
]
CHAT_TOOLS: Final[list[JsonValue]] = [
    {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}},
    {"type": "function", "function": {"name": "get_forecast", "parameters": {"type": "object", "properties": {}}}},
]
ANTHROPIC_TOOLS: Final[list[JsonValue]] = [
    {"name": "get_weather", "input_schema": {"type": "object", "properties": {}}},
    {"name": "get_forecast", "input_schema": {"type": "object", "properties": {}}},
]

_TOOL_NAMES: Final = ("get_weather", "get_forecast")
RESPONSES_TOOL_PARAMS: Final[list[FunctionToolParam]] = [
    {"type": "function", "name": name, "parameters": {"type": "object", "properties": {}}, "strict": False}
    for name in _TOOL_NAMES
]
CHAT_TOOL_PARAMS: Final[list[ChatCompletionFunctionToolParam]] = [
    {"type": "function", "function": {"name": name, "parameters": {"type": "object", "properties": {}}}}
    for name in _TOOL_NAMES
]
ANTHROPIC_TOOL_PARAMS: Final[list[AnthropicToolParam]] = [
    {"name": name, "input_schema": {"type": "object", "properties": {}}} for name in _TOOL_NAMES
]


def reasoning_item(identity: str, summaries: Sequence[str] = SUMMARIES) -> dict[str, JsonValue]:
    return {
        "type": "reasoning",
        "id": f"rs_{identity}",
        "summary": [{"type": "summary_text", "text": text} for text in summaries],
    }


def message_item(identity: str, text: str = ANSWER) -> dict[str, JsonValue]:
    return {
        "type": "message",
        "id": f"msg_{identity}",
        "role": "assistant",
        "status": "completed",
        "content": [{"type": "output_text", "text": text, "annotations": []}],
    }


def refusal_item(identity: str, refusal: str = REFUSAL) -> dict[str, JsonValue]:
    return {
        "type": "message",
        "id": f"msg_refusal_{identity}",
        "role": "assistant",
        "status": "completed",
        "content": [{"type": "refusal", "refusal": refusal}],
    }


def function_call_item(
    identity: str, arguments: JsonValue = ARGUMENTS, name: str = "get_weather"
) -> dict[str, JsonValue]:
    return {
        "type": "function_call",
        "id": f"fc_{name}_{identity}",
        "call_id": f"call_{name}_{identity}",
        "name": name,
        "arguments": arguments,
        "status": "completed",
    }


def custom_tool_call_item(identity: str, tool_input: str = SQL) -> dict[str, JsonValue]:
    return {
        "type": "custom_tool_call",
        "id": f"ctc_{identity}",
        "call_id": f"call_sql_{identity}",
        "name": "run_sql",
        "input": tool_input,
    }


def every_item_kind(identity: str) -> tuple[dict[str, JsonValue], ...]:
    return (
        reasoning_item(identity),
        message_item(identity),
        function_call_item(identity),
        custom_tool_call_item(identity),
        refusal_item(identity),
    )


def response_body(identity: str, output: Sequence[Mapping[str, JsonValue]]) -> dict[str, JsonValue]:
    return {
        "id": f"resp_{identity}",
        "object": "response",
        "created_at": 1,
        "status": "completed",
        "model": "faked-stream-backend",
        "output": [dict(item) for item in output],
        "usage": {"input_tokens": 12, "output_tokens": 30, "total_tokens": 42},
    }


def answering(output: Callable[[str], Sequence[Mapping[str, JsonValue]]]) -> Callable[[Request], Reply]:
    def respond(request: Request) -> Reply:
        identity: Final = uuid.uuid4().hex[:16]
        return Reply(body=json.dumps(response_body(identity, output(identity))).encode())

    return respond


def faked_model(scenario: Scenario, wire: Wire, model: str = "openai/faked-stream-backend") -> str:
    return scenario.model(model=model, api_base=wire.url + "/v1", api_key=API_KEY, model_info=FAKED_STREAM)


def faked_azure_model(scenario: Scenario, wire: Wire) -> str:
    return scenario.model(
        model="azure/responses/faked-stream-deployment",
        api_base=wire.url,
        api_key=API_KEY,
        api_version=AZURE_API_VERSION,
        model_info=FAKED_STREAM,
    )


def posted_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
    return tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST")


def assert_one_unstreamed_upstream_call(wire: Wire, target: str) -> dict[str, JsonValue]:
    received: Final = tuple(request for request in wire.drain() if request.method == "POST")
    assert [request.target for request in received] == [target], [request.target for request in received]
    body: Final = JSON_OBJECT.validate_json(received[0].body)
    assert body.get("stream") is not True, body
    return body


def data_events(lines: Iterator[str]) -> tuple[dict[str, JsonValue], ...]:
    return tuple(JSON_OBJECT.validate_json(line.removeprefix("data: ")) for line in lines if line.startswith("data: {"))


def streamed(
    gateway: Gateway, path: str, body: Mapping[str, JsonValue], key: str | None = None
) -> tuple[httpx.Headers, tuple[dict[str, JsonValue], ...]]:
    with gateway.client.stream(
        "POST",
        path,
        json={**body, "stream": True},
        headers={"Authorization": f"Bearer {gateway.key if key is None else key}", "anthropic-version": "2023-06-01"},
        timeout=120,
    ) as response:
        assert response.status_code == 200, response.read()
        return response.headers, data_events(response.iter_lines())


def stream_events(
    gateway: Gateway, path: str, body: Mapping[str, JsonValue], key: str | None = None
) -> tuple[str, tuple[dict[str, JsonValue], ...]]:
    headers, events = streamed(gateway, path, body, key)
    return headers[CALL_ID], events


def responses_request(model: str, **extra: JsonValue) -> dict[str, JsonValue]:
    return {"model": model, "input": f"{QUESTION} {uuid.uuid4().hex}", "tools": RESPONSES_TOOLS, **extra}


def joined(events: Sequence[Mapping[str, JsonValue]], kind: str, output_index: int, **match: JsonValue) -> str:
    return "".join(string_value(event["delta"]) for event in events if _matches(event, kind, output_index, match))


def _matches(event: Mapping[str, JsonValue], kind: str, output_index: int, match: Mapping[str, JsonValue]) -> bool:
    return (
        event["type"] == kind
        and event.get("output_index") == output_index
        and all(event.get(name) == value for name, value in match.items())
    )


def _of_kind(events: Sequence[Mapping[str, JsonValue]], kind: str) -> Iterator[dict[str, JsonValue]]:
    for event in events:
        if event["type"] == kind:
            yield dict(event)


def items_by_index(events: Sequence[Mapping[str, JsonValue]], kind: str) -> dict[int, dict[str, JsonValue]]:
    return {int(str(event["output_index"])): object_value(event["item"]) for event in _of_kind(events, kind)}


def _announced_parts(
    events: Sequence[Mapping[str, JsonValue]], kind: str, output_index: int
) -> tuple[dict[str, JsonValue], ...]:
    return tuple(
        object_value(event["part"]) for event in _of_kind(events, kind) if event["output_index"] == output_index
    )


def _rebuilt_part(
    events: Sequence[Mapping[str, JsonValue]], output_index: int, content_index: int, part: Mapping[str, JsonValue]
) -> dict[str, JsonValue]:
    match part.get("type"):
        case "output_text":
            text: Final = joined(events, "response.output_text.delta", output_index, content_index=content_index)
            return {**part, "text": string_value(part["text"]) + text}
        case "refusal":
            refusal: Final = joined(events, "response.refusal.delta", output_index, content_index=content_index)
            return {**part, "refusal": string_value(part["refusal"]) + refusal}
        case _:
            return dict(part)


def _rebuilt_summary(
    events: Sequence[Mapping[str, JsonValue]], output_index: int, summary_index: int, part: Mapping[str, JsonValue]
) -> dict[str, JsonValue]:
    text: Final = joined(events, "response.reasoning_summary_text.delta", output_index, summary_index=summary_index)
    return {**part, "text": string_value(part["text"]) + text}


def _rebuilt_item(
    events: Sequence[Mapping[str, JsonValue]], output_index: int, added: Mapping[str, JsonValue]
) -> dict[str, JsonValue]:
    match added.get("type"):
        case "message":
            parts: Final = (
                *map(object_value, list_value(added["content"])),
                *_announced_parts(events, "response.content_part.added", output_index),
            )
            return {
                **added,
                "content": [_rebuilt_part(events, output_index, index, part) for index, part in enumerate(parts)],
            }
        case "reasoning":
            summaries: Final = (
                *map(object_value, list_value(added["summary"])),
                *_announced_parts(events, "response.reasoning_summary_part.added", output_index),
            )
            return {
                **added,
                "summary": [
                    _rebuilt_summary(events, output_index, index, part) for index, part in enumerate(summaries)
                ],
            }
        case "function_call":
            arguments: Final = joined(events, "response.function_call_arguments.delta", output_index)
            return {**added, "arguments": string_value(added["arguments"]) + arguments}
        case "custom_tool_call":
            tool_input: Final = joined(events, "response.custom_tool_call_input.delta", output_index)
            return {**added, "input": string_value(added["input"]) + tool_input}
        case "mcp_call":
            mcp_arguments: Final = joined(events, "response.mcp_call_arguments.delta", output_index)
            return {**added, "arguments": string_value(added["arguments"]) + mcp_arguments}
        case _:
            return dict(added)


def rebuilt(events: Sequence[Mapping[str, JsonValue]]) -> dict[int, dict[str, JsonValue]]:
    return {
        index: _rebuilt_item(events, index, added)
        for index, added in items_by_index(events, "response.output_item.added").items()
    }


def values(item: Mapping[str, JsonValue]) -> JsonValue:
    match item.get("type"):
        case "message":
            return [
                [string_value(part["type"]), part.get("text", part.get("refusal"))]
                for part in map(object_value, list_value(item["content"]))
            ]
        case "reasoning":
            return [part["text"] for part in map(object_value, list_value(item["summary"]))]
        case "function_call" | "mcp_call":
            return [item["name"], item["arguments"]]
        case "custom_tool_call":
            return [item["name"], item["input"]]
        case _:
            return json.dumps(item, sort_keys=True)


def completed_output(events: Sequence[Mapping[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]:
    completed: Final = tuple(_of_kind(events, "response.completed"))
    assert len(completed) == 1, [event["type"] for event in events]
    return tuple(map(object_value, list_value(object_value(completed[0]["response"])["output"])))


def assert_each_value_arrives_once(
    events: Sequence[Mapping[str, JsonValue]], expected: Sequence[Mapping[str, JsonValue]]
) -> None:
    wanted: Final = [values(item) for item in expected]
    seen_by_client: Final = rebuilt(events)
    done: Final = items_by_index(events, "response.output_item.done")
    assert [values(seen_by_client[index]) for index in sorted(seen_by_client)] == wanted, [
        event["type"] for event in events
    ]
    assert [values(done[index]) for index in sorted(done)] == wanted
    assert [values(item) for item in completed_output(events)] == wanted


def spend_statuses(call_id: str) -> list[JsonValue]:
    rows: Final = eventually(
        lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,)),
        lambda found: len(found) >= 1,
        seconds=70,
    )
    return [row["status"] for row in rows]


def http_client(gateway: Gateway) -> httpx.AsyncClient:
    return httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False)