mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
476 lines
13 KiB
Python
476 lines
13 KiB
Python
"""Client for the non-chat inference endpoints (responses, messages, rerank,
|
|
embeddings, audio speech, image generation).
|
|
|
|
Each test registers the deployment it needs through /model/new (deleted on
|
|
teardown), so nothing is hardcoded into the gateway config, then drives the
|
|
endpoint with `send` and parses the provider-native body with a suite-local model
|
|
so the assertion is on real content, not just a 200.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Literal
|
|
|
|
from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS
|
|
from e2e_http import BinaryStream, Result, StreamingResponse
|
|
from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock
|
|
from proxy_client import ProxyClient
|
|
from pydantic import BaseModel
|
|
|
|
__all__ = [
|
|
"CacheControl",
|
|
"ImageEditForm",
|
|
"ImagesResult",
|
|
"RichMessage",
|
|
"TextBlock",
|
|
"TranscriptionForm",
|
|
"TranscriptionResult",
|
|
]
|
|
|
|
|
|
class FunctionParameterProperty(BaseModel):
|
|
type: str
|
|
description: str | None = None
|
|
|
|
|
|
class FunctionParameters(BaseModel):
|
|
type: Literal["object"] = "object"
|
|
properties: dict[str, FunctionParameterProperty]
|
|
required: list[str] = []
|
|
|
|
|
|
class ResponsesFunctionTool(BaseModel):
|
|
type: Literal["function"] = "function"
|
|
name: str
|
|
description: str | None = None
|
|
parameters: FunctionParameters
|
|
|
|
|
|
class ResponsesInputTextPart(BaseModel):
|
|
type: Literal["input_text"] = "input_text"
|
|
text: str
|
|
|
|
|
|
class ResponsesInputImagePart(BaseModel):
|
|
type: Literal["input_image"] = "input_image"
|
|
image_url: str
|
|
|
|
|
|
ResponsesInputContentPart = ResponsesInputTextPart | ResponsesInputImagePart
|
|
|
|
|
|
class ResponsesInputMessage(BaseModel):
|
|
role: Literal["user", "assistant", "system"] = "user"
|
|
content: list[ResponsesInputContentPart]
|
|
|
|
|
|
ResponsesInput = str | list[ResponsesInputMessage]
|
|
|
|
|
|
class ResponsesRequest(BaseModel):
|
|
model: str
|
|
input: ResponsesInput
|
|
instructions: str | None = None
|
|
stream: bool = False
|
|
tools: list[ResponsesFunctionTool] | None = None
|
|
guardrails: list[str] | None = None
|
|
safety_identifier: str | None = None
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class MessagesRequest(BaseModel):
|
|
model: str
|
|
max_tokens: int
|
|
messages: list[ChatMessage]
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class RichMessagesRequest(BaseModel):
|
|
model: str
|
|
max_tokens: int = 64
|
|
system: list[TextBlock]
|
|
messages: list[RichMessage]
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class CompletionsRequest(BaseModel):
|
|
model: str
|
|
prompt: str
|
|
max_tokens: int = 32
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class EmbeddingsRequest(BaseModel):
|
|
model: str
|
|
input: str
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class RerankRequest(BaseModel):
|
|
model: str
|
|
query: str
|
|
documents: list[str]
|
|
top_n: int
|
|
cache: dict[str, bool] | None = {"no-cache": True}
|
|
|
|
|
|
class SpeechRequest(BaseModel):
|
|
model: str
|
|
input: str
|
|
voice: str
|
|
|
|
|
|
class ImageRequest(BaseModel):
|
|
model: str
|
|
prompt: str
|
|
n: int = 1
|
|
size: str = "1024x1024"
|
|
|
|
|
|
class ImageEditForm(BaseModel):
|
|
model: str
|
|
prompt: str
|
|
n: int = 1
|
|
|
|
|
|
class TranscriptionForm(BaseModel):
|
|
model: str
|
|
response_format: str = "json"
|
|
|
|
|
|
class ModerationRequest(BaseModel):
|
|
model: str
|
|
input: str
|
|
|
|
|
|
class GenerateContentPart(BaseModel):
|
|
text: str
|
|
|
|
|
|
class GenerateContentContent(BaseModel):
|
|
role: Literal["user"] = "user"
|
|
parts: tuple[GenerateContentPart, ...]
|
|
|
|
|
|
class GenerateContentBody(BaseModel):
|
|
contents: tuple[GenerateContentContent, ...]
|
|
|
|
|
|
class ResponsesOutputContent(BaseModel):
|
|
type: str | None = None
|
|
text: str | None = None
|
|
|
|
|
|
class ResponsesOutputItem(BaseModel):
|
|
type: str | None = None
|
|
content: list[ResponsesOutputContent] = []
|
|
name: str | None = None
|
|
arguments: str | None = None
|
|
call_id: str | None = None
|
|
|
|
|
|
class ResponsesResult(BaseModel):
|
|
id: str | None = None
|
|
status: str | None = None
|
|
model: str | None = None
|
|
output: list[ResponsesOutputItem] = []
|
|
|
|
@property
|
|
def text(self) -> str:
|
|
return "".join(
|
|
content.text or "" for item in self.output for content in item.content
|
|
)
|
|
|
|
@property
|
|
def function_calls(self) -> tuple[ResponsesOutputItem, ...]:
|
|
return tuple(
|
|
item
|
|
for item in self.output
|
|
if item.type == "function_call"
|
|
and item.name is not None
|
|
and item.arguments is not None
|
|
)
|
|
|
|
|
|
class ResponsesStreamEvent(BaseModel):
|
|
event_id: str | None = None
|
|
|
|
|
|
class ResponsesStreamEventType(BaseModel):
|
|
type: str
|
|
|
|
|
|
class ResponsesOutputTextDeltaEvent(ResponsesStreamEvent):
|
|
type: Literal["response.output_text.delta"]
|
|
delta: str
|
|
|
|
|
|
class AnthropicContentBlock(BaseModel):
|
|
type: str | None = None
|
|
text: str | None = None
|
|
|
|
|
|
class MessagesUsage(BaseModel):
|
|
input_tokens: int = 0
|
|
output_tokens: int = 0
|
|
cache_creation_input_tokens: int = 0
|
|
cache_read_input_tokens: int = 0
|
|
|
|
|
|
class MessagesResult(BaseModel):
|
|
id: str | None = None
|
|
role: str | None = None
|
|
model: str | None = None
|
|
content: list[AnthropicContentBlock] = []
|
|
usage: MessagesUsage = MessagesUsage()
|
|
|
|
@property
|
|
def text(self) -> str:
|
|
return "".join(block.text or "" for block in self.content)
|
|
|
|
|
|
class CompletionChoice(BaseModel):
|
|
text: str | None = None
|
|
|
|
|
|
class CompletionsResult(BaseModel):
|
|
choices: list[CompletionChoice] = []
|
|
|
|
|
|
class EmbeddingItem(BaseModel):
|
|
embedding: list[float] = []
|
|
|
|
|
|
class EmbeddingsResult(BaseModel):
|
|
data: list[EmbeddingItem] = []
|
|
|
|
@property
|
|
def first_vector(self) -> tuple[float, ...]:
|
|
return tuple(self.data[0].embedding) if self.data else ()
|
|
|
|
|
|
class RerankItem(BaseModel):
|
|
index: int | None = None
|
|
relevance_score: float | None = None
|
|
|
|
|
|
class RerankResult(BaseModel):
|
|
results: list[RerankItem] = []
|
|
|
|
|
|
class ImageItem(BaseModel):
|
|
url: str | None = None
|
|
b64_json: str | None = None
|
|
|
|
|
|
class ImagesResult(BaseModel):
|
|
data: list[ImageItem] = []
|
|
|
|
|
|
class TranscriptionResult(BaseModel):
|
|
text: str = ""
|
|
|
|
|
|
class ModerationResultItem(BaseModel):
|
|
flagged: bool
|
|
categories: dict[str, bool] = {}
|
|
|
|
@property
|
|
def flagged_categories(self) -> tuple[str, ...]:
|
|
return tuple(name for name, hit in self.categories.items() if hit)
|
|
|
|
|
|
class ModerationResult(BaseModel):
|
|
results: list[ModerationResultItem] = []
|
|
|
|
@property
|
|
def first(self) -> ModerationResultItem | None:
|
|
return self.results[0] if self.results else None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class EndpointsClient:
|
|
proxy: ProxyClient
|
|
|
|
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
|
|
return self.proxy.create_model(model_name, litellm_params)
|
|
|
|
def delete_model(self, model_id: str) -> None:
|
|
self.proxy.delete_model(model_id)
|
|
|
|
def _send(
|
|
self, path: str, key: str, body: BaseModel, *, stream: bool = False
|
|
) -> StreamingResponse:
|
|
return self.proxy.transport.send(
|
|
path,
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=body,
|
|
stream=stream,
|
|
)
|
|
|
|
def responses(
|
|
self,
|
|
key: str,
|
|
model: str,
|
|
text: str,
|
|
*,
|
|
stream: bool = False,
|
|
guardrails: list[str] | None = None,
|
|
safety_identifier: str | None = None,
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/responses",
|
|
key,
|
|
ResponsesRequest(
|
|
model=model,
|
|
input=text,
|
|
instructions="You are a helpful assistant",
|
|
stream=stream,
|
|
guardrails=guardrails,
|
|
safety_identifier=safety_identifier,
|
|
),
|
|
stream=stream,
|
|
)
|
|
|
|
def responses_vision(
|
|
self, key: str, model: str, text: str, image_url: str
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/responses",
|
|
key,
|
|
ResponsesRequest(
|
|
model=model,
|
|
input=[
|
|
ResponsesInputMessage(
|
|
content=[
|
|
ResponsesInputTextPart(text=text),
|
|
ResponsesInputImagePart(image_url=image_url),
|
|
]
|
|
)
|
|
],
|
|
instructions="You are a helpful assistant",
|
|
),
|
|
)
|
|
|
|
def responses_with_tools(
|
|
self, key: str, model: str, text: str, tools: list[ResponsesFunctionTool]
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/responses",
|
|
key,
|
|
ResponsesRequest(
|
|
model=model,
|
|
input=text,
|
|
instructions="You are a helpful assistant",
|
|
tools=tools,
|
|
),
|
|
)
|
|
|
|
def messages(
|
|
self, key: str, model: str, text: str, *, max_tokens: int = 64
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/messages",
|
|
key,
|
|
MessagesRequest(
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
messages=[ChatMessage(role="user", content=text)],
|
|
),
|
|
)
|
|
|
|
def text_completions(
|
|
self, key: str, model: str, prompt: str, *, max_tokens: int = 32
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/completions",
|
|
key,
|
|
CompletionsRequest(model=model, prompt=prompt, max_tokens=max_tokens),
|
|
)
|
|
|
|
def embeddings(self, key: str, model: str, text: str) -> StreamingResponse:
|
|
return self._send("/embeddings", key, EmbeddingsRequest(model=model, input=text))
|
|
|
|
def rerank(
|
|
self, key: str, model: str, query: str, documents: list[str], top_n: int
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/rerank",
|
|
key,
|
|
RerankRequest(model=model, query=query, documents=documents, top_n=top_n),
|
|
)
|
|
|
|
def audio_speech(
|
|
self, key: str, model: str, text: str, *, voice: str = "alloy"
|
|
) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice)
|
|
)
|
|
|
|
def audio_speech_stream(
|
|
self, key: str, model: str, text: str, *, voice: str = "alloy"
|
|
) -> BinaryStream:
|
|
return self.proxy.transport.stream_binary(
|
|
"/v1/audio/speech",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=SpeechRequest(model=model, input=text, voice=voice),
|
|
)
|
|
|
|
def transcribe(
|
|
self, key: str, model: str, *, filename: str, content: bytes
|
|
) -> Result[TranscriptionResult]:
|
|
return self.proxy.transport.upload(
|
|
"/v1/audio/transcriptions",
|
|
headers=self.proxy.transport.bearer(key),
|
|
form=TranscriptionForm(model=model),
|
|
filename=filename,
|
|
content=content,
|
|
file_content_type="audio/wav",
|
|
response_type=TranscriptionResult,
|
|
)
|
|
|
|
def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]:
|
|
return self.proxy.transport.post(
|
|
"/v1/moderations",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=ModerationRequest(model=model, input=text),
|
|
response_type=ModerationResult,
|
|
)
|
|
|
|
def images(self, key: str, model: str, prompt: str) -> StreamingResponse:
|
|
return self._send(
|
|
"/v1/images/generations", key, ImageRequest(model=model, prompt=prompt)
|
|
)
|
|
|
|
def image_edit(
|
|
self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png"
|
|
) -> Result[ImagesResult]:
|
|
return self.proxy.transport.upload(
|
|
"/v1/images/edits",
|
|
headers=self.proxy.transport.bearer(key),
|
|
form=ImageEditForm(model=model, prompt=prompt),
|
|
filename=filename,
|
|
content=image,
|
|
file_content_type="image/png",
|
|
file_field="image",
|
|
response_type=ImagesResult,
|
|
timeout=SLOW_PROVIDER_TIMEOUT_SECONDS,
|
|
)
|
|
|
|
def generate_content(
|
|
self, key: str, model: str, text: str, *, stream: bool = False
|
|
) -> StreamingResponse:
|
|
operation = "streamGenerateContent" if stream else "generateContent"
|
|
return self._send(
|
|
f"/v1beta/models/{model}:{operation}",
|
|
key,
|
|
GenerateContentBody(
|
|
contents=(GenerateContentContent(parts=(GenerateContentPart(text=text),)),)
|
|
),
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient:
|
|
return EndpointsClient(proxy=proxy)
|