import json import uuid from hashlib import sha256 from typing import Final import httpx import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from openai import OpenAI from pydantic import BaseModel, JsonValue, TypeAdapter _TEXT_BACKEND: Final = "gpt-3.5-turbo-instruct" _CHAT_BACKEND: Final = "gpt-5.4-mini" _API_KEY: Final = "synthetic-openai-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _SPEND_ROW: Final = 'SELECT request_tags, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s' class _TextChoice(BaseModel): index: int text: str finish_reason: str class _TextCompletion(BaseModel): id: str object: str model: str choices: list[_TextChoice] def _text_completion(identity: str, texts: tuple[str, ...]) -> bytes: return json.dumps( { "id": identity, "object": "text_completion", "created": 1, "model": _TEXT_BACKEND, "choices": [ {"index": index, "text": text, "finish_reason": "stop", "logprobs": None} for index, text in enumerate(texts) ], "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, } ).encode() def _chat_completion(identity: str, content: str) -> bytes: return json.dumps( { "id": identity, "object": "chat.completion", "created": 1, "model": _CHAT_BACKEND, "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, } ).encode() def _openai(gateway: Gateway, key: str) -> OpenAI: return OpenAI( api_key=key, base_url=str(gateway.client.base_url).rstrip("/") + "/v1", max_retries=0, http_client=httpx.Client(timeout=15, trust_env=False), ) def _form_fields(fields: dict[str, str]) -> list[tuple[str, tuple[None, str]]]: return [(name, (None, value)) for name, value in fields.items()] def _assert_text_completion_request(request: Request, prompt: str) -> None: assert (request.method, request.target) == ("POST", "/completions"), request assert request.headers["authorization"] == f"Bearer {_API_KEY}" assert request.headers["content-type"] == "application/json" assert _JSON_OBJECT.validate_json(request.body) == {"model": _TEXT_BACKEND, "prompt": prompt} def _assert_spend_row_carries_client_tag_and_key_identity( request_id: str, key: str, user_id: str, team_id: str ) -> None: rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (request_id,)), lambda values: len(values) == 1, seconds=70) tags: Final = rows[0]["request_tags"] metadata: Final = rows[0]["metadata"] assert isinstance(tags, list) and isinstance(metadata, dict), rows assert "t1" in tags, rows assert ( metadata["user_api_key"], metadata["user_api_key_user_id"], metadata["user_api_key_team_id"], ) == (sha256(key.encode()).hexdigest(), user_id, team_id), metadata def test_text_completion_json_string_metadata_reaches_upstream_and_tags_spend(gateway: Gateway) -> None: pytest.skip("BUG: /v1/completions with metadata sent as a JSON string returns 500 before dispatch") identity: Final = f"cmpl-{uuid.uuid4().hex}" prompt: Final = f"tagged prompt {identity}" def respond(request: Request) -> Reply: _assert_text_completion_request(request, prompt) return Reply(body=_text_completion(identity, ("tagged",))) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY) team: Final = scenario.team() user_id: Final = f"integration-{uuid.uuid4().hex}" key: Final = scenario.key(models=[model], team_id=team, user_id=user_id) with _openai(gateway, key) as client: completion: Final = client.completions.create( model=model, prompt=prompt, extra_body={"metadata": json.dumps({"tags": ["t1"]})} ) assert [(choice.index, choice.text) for choice in completion.choices] == [(0, "tagged")], completion assert completion.id == identity, completion assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] _assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team) def test_text_completion_form_metadata_reaches_upstream_and_tags_spend(gateway: Gateway) -> None: identity: Final = f"cmpl-{uuid.uuid4().hex}" prompt: Final = f"tagged prompt {identity}" def respond(request: Request) -> Reply: _assert_text_completion_request(request, prompt) return Reply(body=_text_completion(identity, ("tagged",))) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY) team: Final = scenario.team() user_id: Final = f"integration-{uuid.uuid4().hex}" key: Final = scenario.key(models=[model], team_id=team, user_id=user_id) response: Final = gateway.client.post( "/v1/completions", files=_form_fields({"model": model, "prompt": prompt, "metadata": json.dumps({"tags": ["t1"]})}), headers={"Authorization": f"Bearer {key}"}, ) assert response.status_code == 200, response.text assert response.request.headers["content-type"].startswith("multipart/form-data"), response.request.headers payload: Final = _TextCompletion.model_validate_json(response.content) assert (payload.id, payload.object, payload.model) == (identity, "text_completion", model), response.text assert [(choice.index, choice.text) for choice in payload.choices] == [(0, "tagged")], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] _assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team) def test_text_completion_typed_form_fields_reach_upstream_as_json_types(gateway: Gateway) -> None: pytest.skip( "BUG: multipart /v1/completions forwards temperature, max_tokens, n and echo upstream as strings, " "so echo='false' is truthy" ) identity: Final = f"cmpl-{uuid.uuid4().hex}" prompt: Final = f"typed form prompt {identity}" def respond(request: Request) -> Reply: assert (request.method, request.target) == ("POST", "/completions"), request assert request.headers["authorization"] == f"Bearer {_API_KEY}" assert request.headers["content-type"] == "application/json" assert _JSON_OBJECT.validate_json(request.body) == { "model": _TEXT_BACKEND, "prompt": prompt, "temperature": 0.2, "max_tokens": 5, "n": 1, "echo": False, }, request.body return Reply(body=_text_completion(identity, ("tagged",))) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY) team: Final = scenario.team() user_id: Final = f"integration-{uuid.uuid4().hex}" key: Final = scenario.key(models=[model], team_id=team, user_id=user_id) response: Final = gateway.client.post( "/v1/completions", files=_form_fields( { "model": model, "prompt": prompt, "temperature": "0.2", "max_tokens": "5", "n": "1", "echo": "false", "metadata": json.dumps({"tags": ["t1"]}), } ), headers={"Authorization": f"Bearer {key}"}, ) assert response.status_code == 200, response.text assert response.request.headers["content-type"].startswith("multipart/form-data"), response.request.headers payload: Final = _TextCompletion.model_validate_json(response.content) assert (payload.id, payload.object, payload.model) == (identity, "text_completion", model), response.text assert [(choice.index, choice.text) for choice in payload.choices] == [(0, "tagged")], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] _assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team) def test_text_completion_sdk_parameter_set_reaches_upstream_whole(gateway: Gateway) -> None: identity: Final = f"cmpl-full-params-{uuid.uuid4().hex}" prompt: Final = f"full parameter prompt {identity}" expected: Final = { "model": _TEXT_BACKEND, "prompt": prompt, "suffix": " tail", "echo": True, "logprobs": 2, "best_of": 2, "n": 1, "logit_bias": {"50256": -100}, "seed": 7, "user": "integration-user", "stop": ["\n", "END"], "presence_penalty": 0.5, "frequency_penalty": 0.25, "temperature": 0.3, "top_p": 0.9, "max_tokens": 12, } def respond(request: Request) -> Reply: assert (request.method, request.target) == ("POST", "/completions"), request assert request.headers["authorization"] == f"Bearer {_API_KEY}" assert request.headers["content-type"] == "application/json" assert _JSON_OBJECT.validate_json(request.body) == expected, request.body return Reply(body=_text_completion(identity, ("one",))) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY) key: Final = scenario.key(models=[model]) with _openai(gateway, key) as client: completion: Final = client.completions.create( model=model, prompt=prompt, suffix=" tail", echo=True, logprobs=2, best_of=2, n=1, logit_bias={"50256": -100}, seed=7, user="integration-user", stop=["\n", "END"], presence_penalty=0.5, frequency_penalty=0.25, temperature=0.3, top_p=0.9, max_tokens=12, ) assert (completion.id, completion.object) == (identity, "text_completion"), completion assert [(choice.index, choice.text) for choice in completion.choices] == [(0, "one")], completion assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] @pytest.mark.parametrize( "prompt_shape", ( pytest.param("string-list", id="string-list"), pytest.param("token-ids", id="token-ids"), pytest.param("token-id-batch", id="token-id-batch"), ), ) def test_text_completion_prompt_spellings_reach_upstream_unchanged(gateway: Gateway, prompt_shape: str) -> None: identity: Final = f"cmpl-{uuid.uuid4().hex}" prompt_identity: Final = uuid.uuid4() token_ids: Final = tuple( int.from_bytes(prompt_identity.bytes[offset : offset + 4], "big") % 50_000 + 50_000 * (offset // 4) for offset in range(0, 16, 4) ) prompt: Final[JsonValue] = ( [f"first {prompt_identity.hex}", f"second {prompt_identity.hex}"] if prompt_shape == "string-list" else list(token_ids[:3]) if prompt_shape == "token-ids" else [list(token_ids[:2]), list(token_ids[2:])] ) texts: Final = ("first", "second") if isinstance(prompt, list) and isinstance(prompt[0], (str, list)) else ("one",) def respond(request: Request) -> Reply: assert (request.method, request.target) == ("POST", "/completions"), request assert request.headers["authorization"] == f"Bearer {_API_KEY}" assert _JSON_OBJECT.validate_json(request.body) == {"model": _TEXT_BACKEND, "prompt": prompt} return Reply(body=_text_completion(identity, texts)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY) key: Final = scenario.key(models=[model]) with _openai(gateway, key) as client: completion: Final = client.completions.create(model=model, prompt=prompt) assert completion.id == identity, completion assert completion.object == "text_completion", completion assert [(choice.index, choice.text) for choice in completion.choices] == list(enumerate(texts)), completion assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] def test_text_completion_string_list_prompt_becomes_user_messages_for_chat_only_deployment(gateway: Gateway) -> None: identity: Final = f"chatcmpl-{uuid.uuid4().hex}" prompts: Final = [f"first {identity}", f"second {identity}"] def respond(request: Request) -> Reply: assert (request.method, request.target) == ("POST", "/chat/completions"), request assert request.headers["authorization"] == f"Bearer {_API_KEY}" assert _JSON_OBJECT.validate_json(request.body) == { "model": _CHAT_BACKEND, "messages": [{"role": "user", "content": prompts[0]}, {"role": "user", "content": prompts[1]}], } return Reply(body=_chat_completion(identity, "answered")) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"openai/{_CHAT_BACKEND}", api_base=wire.url, api_key=_API_KEY) key: Final = scenario.key(models=[model]) with _openai(gateway, key) as client: completion: Final = client.completions.create(model=model, prompt=prompts) assert (completion.id, completion.object) == (identity, "text_completion"), completion assert [(choice.index, choice.text, choice.finish_reason) for choice in completion.choices] == [ (0, "answered", "stop") ], completion assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]