wehub-resource-sync
341 lines
10 KiB
Python
341 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import warnings
|
|
from typing import Annotated, Any, Literal
|
|
from uuid import uuid4
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_serializer
|
|
|
|
|
|
class TextContentPart(BaseModel):
|
|
type: Literal["text"]
|
|
text: str
|
|
|
|
|
|
class ImageUrl(BaseModel):
|
|
"""
|
|
Represents an image URL.
|
|
|
|
Attributes:
|
|
url: Either a URL of an image or base64 encoded data.
|
|
https://platform.openai.com/docs/guides/vision?lang=curl#uploading-base64-encoded-images
|
|
detail: The level of resolution for the image when the model receives it.
|
|
For example, when set to "low", the model will see a image resized to
|
|
512x512 pixels, which consumes fewer tokens. In OpenAI, this is optional
|
|
and defaults to "auto".
|
|
https://platform.openai.com/docs/guides/vision?lang=curl#low-or-high-fidelity-image-understanding
|
|
"""
|
|
|
|
url: str
|
|
detail: Literal["auto", "low", "high"] | None = None
|
|
|
|
|
|
class ImageContentPart(BaseModel):
|
|
type: Literal["image_url"]
|
|
image_url: ImageUrl
|
|
|
|
|
|
class InputAudio(BaseModel):
|
|
data: str # base64 encoded data
|
|
format: Literal["wav", "mp3"]
|
|
|
|
|
|
class AudioContentPart(BaseModel):
|
|
type: Literal["input_audio"]
|
|
input_audio: InputAudio
|
|
|
|
|
|
ContentPartsList = list[
|
|
Annotated[TextContentPart | ImageContentPart | AudioContentPart, Field(discriminator="type")]
|
|
]
|
|
|
|
|
|
ContentType = Annotated[str | ContentPartsList, Field(union_mode="left_to_right")]
|
|
|
|
|
|
class Function(BaseModel):
|
|
name: str | None = None
|
|
arguments: str | None = None
|
|
|
|
def to_tool_call(self, id=None) -> ToolCall:
|
|
if id is None:
|
|
id = str(uuid4())
|
|
return ToolCall(id=id, type="function", function=self)
|
|
|
|
|
|
class ToolCall(BaseModel):
|
|
id: str
|
|
type: str = Field(default="function")
|
|
function: Function
|
|
# Gemini thinking-mode models return a thought_signature with each
|
|
# function call that must be echoed back in subsequent turns.
|
|
# https://ai.google.dev/gemini-api/docs/thought-signatures
|
|
thought_signature: str | None = Field(default=None)
|
|
|
|
@model_serializer(mode="wrap")
|
|
def _serialize(self, handler):
|
|
data = handler(self)
|
|
if data.get("thought_signature") is None:
|
|
data.pop("thought_signature", None)
|
|
return data
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
"""
|
|
A chat request. ``content`` can be a string, or an array of content parts.
|
|
|
|
A content part is one of the following:
|
|
|
|
- :py:class:`TextContentPart <mlflow.types.chat.TextContentPart>`
|
|
- :py:class:`ImageContentPart <mlflow.types.chat.ImageContentPart>`
|
|
- :py:class:`AudioContentPart <mlflow.types.chat.AudioContentPart>`
|
|
"""
|
|
|
|
role: str
|
|
content: ContentType | None = None
|
|
# NB: In the actual OpenAI chat completion API spec, these fields only
|
|
# present in either the request or response message (tool_call_id is only in
|
|
# the request, while the other two are only in the response).
|
|
# Strictly speaking, we should separate the request and response message types
|
|
# to match OpenAI's API spec. However, we don't want to do that because we the
|
|
# request and response message types are not distinguished in many parts of the
|
|
# codebase, and also we don't want to ask users to use two different classes.
|
|
# Therefore, we include all fields in this class, while marking them as optional.
|
|
# TODO: Define a sub classes for different type of messages (request/response, and
|
|
# system/user/assistant/tool, etc), and create a factory function to allow users
|
|
# to create them without worrying about the details.
|
|
tool_calls: list[ToolCall] | None = None
|
|
refusal: str | None = None
|
|
tool_call_id: str | None = None
|
|
|
|
|
|
AllowedType = Literal["string", "number", "integer", "object", "array", "boolean", "null"]
|
|
|
|
|
|
class ParamType(BaseModel):
|
|
type: AllowedType | list[AllowedType] | None = None
|
|
|
|
|
|
class ParamProperty(ParamType):
|
|
"""
|
|
OpenAI uses JSON Schema (https://json-schema.org/) for function parameters.
|
|
See OpenAI function calling reference:
|
|
https://platform.openai.com/docs/guides/function-calling?&api-mode=responses#defining-functions
|
|
|
|
JSON Schema enum supports any JSON type (str, int, float, bool, null, arrays, objects),
|
|
but we restrict to basic scalar types for practical use cases and API safety.
|
|
"""
|
|
|
|
description: str | None = None
|
|
enum: list[str | int | float | bool] | None = None
|
|
# Recursive type so nested arrays (e.g. list[list[str]]) preserve their inner
|
|
# `items` schema through Pydantic round-trips. If this were `ParamType`, the
|
|
# inner `items` field would be silently stripped and downstream providers
|
|
# would reject the schema with "array schema missing items".
|
|
items: ParamProperty | None = None
|
|
|
|
|
|
class FunctionParams(BaseModel):
|
|
properties: dict[str, ParamProperty]
|
|
type: Literal["object"] = "object"
|
|
required: list[str] | None = None
|
|
additionalProperties: bool | None = None
|
|
|
|
|
|
class FunctionToolDefinition(BaseModel):
|
|
name: str
|
|
description: str | None = None
|
|
parameters: FunctionParams | None = None
|
|
strict: bool | None = None
|
|
|
|
|
|
class ChatTool(BaseModel):
|
|
"""
|
|
A tool definition passed to the chat completion API.
|
|
|
|
Ref: https://platform.openai.com/docs/guides/function-calling
|
|
"""
|
|
|
|
type: Literal["function"]
|
|
function: FunctionToolDefinition | None = None
|
|
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
message='Field name "schema" in "JsonSchemaSpec" shadows an attribute in parent '
|
|
'"BaseModel"',
|
|
category=UserWarning,
|
|
)
|
|
|
|
class JsonSchemaSpec(BaseModel):
|
|
"""
|
|
OpenAI-compatible JSON Schema envelope for structured outputs.
|
|
|
|
Attributes:
|
|
name: The schema name.
|
|
schema: A JSON Schema definition.
|
|
strict: Whether model output should strictly follow the schema.
|
|
"""
|
|
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
name: str
|
|
schema: dict[str, Any] = Field(...)
|
|
strict: bool = True
|
|
|
|
|
|
class ResponseFormat(BaseModel):
|
|
"""
|
|
Response format configuration for structured outputs. Compatible with
|
|
OpenAI's Chat Completion API.
|
|
|
|
Supported formats:
|
|
- {"type": "text"}
|
|
- {"type": "json_object"}
|
|
- {"type": "json_schema", "json_schema": {"name": ..., "schema": {...}, "strict": true}}
|
|
"""
|
|
|
|
type: Literal["text", "json_object", "json_schema"]
|
|
json_schema: JsonSchemaSpec | None = None
|
|
|
|
|
|
class ToolChoiceFunction(BaseModel):
|
|
"""Specifies a tool the model should use."""
|
|
|
|
name: str
|
|
|
|
|
|
class ToolChoice(BaseModel):
|
|
"""
|
|
Specifies a particular tool to use.
|
|
|
|
OpenAI format: {"type": "function", "function": {"name": "my_function"}}
|
|
"""
|
|
|
|
type: Literal["function"]
|
|
function: ToolChoiceFunction
|
|
|
|
|
|
class BaseRequestPayload(BaseModel):
|
|
"""Common parameters used for chat completions and completion endpoints."""
|
|
|
|
n: int = Field(1, ge=1)
|
|
stop: list[str] | None = Field(None, min_length=1)
|
|
max_tokens: int | None = Field(None, ge=1)
|
|
max_completion_tokens: int | None = Field(None, ge=1)
|
|
stream: bool | None = None
|
|
stream_options: dict[str, Any] | None = None
|
|
model: str | None = None
|
|
response_format: ResponseFormat | None = None
|
|
temperature: float | None = Field(None, ge=0, le=2)
|
|
top_p: float | None = Field(None, ge=0, le=1)
|
|
presence_penalty: float | None = Field(None, ge=-2, le=2)
|
|
frequency_penalty: float | None = Field(None, ge=-2, le=2)
|
|
top_k: int | None = Field(None, ge=1)
|
|
|
|
|
|
# NB: For interface constructs that rely on other BaseModel implementations, in
|
|
# pydantic 1 the **order** in which classes are defined in this module is absolutely
|
|
# critical to prevent ForwardRef errors. Pydantic 2 does not have this limitation.
|
|
# To maintain compatibility with Pydantic 1, ensure that all classes that are defined in
|
|
# this file have dependencies defined higher than the line of usage.
|
|
|
|
|
|
class ChatChoice(BaseModel):
|
|
index: int
|
|
message: ChatMessage
|
|
finish_reason: str | None = None
|
|
|
|
|
|
class PromptTokensDetails(BaseModel):
|
|
model_config = {"extra": "allow"}
|
|
|
|
cached_tokens: int | None = None
|
|
|
|
|
|
class ChatUsage(BaseModel):
|
|
model_config = {"extra": "allow"}
|
|
|
|
prompt_tokens: int | None = None
|
|
completion_tokens: int | None = None
|
|
total_tokens: int | None = None
|
|
prompt_tokens_details: PromptTokensDetails | None = None
|
|
|
|
@model_serializer(mode="wrap")
|
|
def _serialize(self, handler):
|
|
data = handler(self)
|
|
if data.get("prompt_tokens_details") is None:
|
|
data.pop("prompt_tokens_details", None)
|
|
return data
|
|
|
|
|
|
class ToolCallDelta(BaseModel):
|
|
index: int
|
|
id: str | None = None
|
|
type: str | None = None
|
|
function: Function
|
|
# Gemini thinking-mode models return a thought_signature with each
|
|
# function call that must be echoed back in subsequent turns.
|
|
# https://ai.google.dev/gemini-api/docs/thought-signatures
|
|
thought_signature: str | None = Field(default=None)
|
|
|
|
@model_serializer(mode="wrap")
|
|
def _serialize(self, handler):
|
|
data = handler(self)
|
|
if data.get("thought_signature") is None:
|
|
data.pop("thought_signature", None)
|
|
return data
|
|
|
|
|
|
class ChatChoiceDelta(BaseModel):
|
|
role: str | None = None
|
|
content: str | None = None
|
|
tool_calls: list[ToolCallDelta] | None = None
|
|
|
|
|
|
class ChatChunkChoice(BaseModel):
|
|
index: int
|
|
finish_reason: str | None = None
|
|
delta: ChatChoiceDelta
|
|
|
|
|
|
class ChatCompletionChunk(BaseModel):
|
|
"""A chunk of a chat completion stream response."""
|
|
|
|
id: str | None = None
|
|
object: str = "chat.completion.chunk"
|
|
created: int
|
|
model: str
|
|
choices: list[ChatChunkChoice]
|
|
usage: ChatUsage | None = None
|
|
|
|
|
|
class ChatCompletionRequest(BaseRequestPayload):
|
|
"""
|
|
A request to the chat completion API.
|
|
|
|
Must be compatible with OpenAI's Chat Completion API.
|
|
https://platform.openai.com/docs/api-reference/chat
|
|
"""
|
|
|
|
messages: list[ChatMessage] = Field(..., min_length=1)
|
|
tools: list[ChatTool] | None = Field(None, min_length=1)
|
|
tool_choice: Literal["none", "auto", "required"] | ToolChoice | None = None
|
|
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
"""
|
|
A response from the chat completion API.
|
|
|
|
Must be compatible with OpenAI's Chat Completion API.
|
|
https://platform.openai.com/docs/api-reference/chat
|
|
"""
|
|
|
|
id: str | None = None
|
|
object: str = "chat.completion"
|
|
created: int
|
|
model: str
|
|
choices: list[ChatChoice]
|
|
usage: ChatUsage
|