star

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/__init__.py

@@ -0,0 +1,3 @@
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+__all__ = ["chat", "completions", "embeddings"]

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/chat.py

@@ -0,0 +1,142 @@
+"""
+This module defines the schemas for the MLflow AI Gateway's chat endpoint.
+
+The schemas must be compatible with OpenAI's Chat Completion API.
+https://platform.openai.com/docs/api-reference/chat
+
+NB: These Pydantic models just alias the models defined in mlflow.types.chat to avoid code
+    duplication, but with the addition of RequestModel and ResponseModel base classes.
+"""
+
+from typing import Literal, Optional
+
+from pydantic import Field
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+
+# Import marked with noqa is for backward compatibility
+from mlflow.types.chat import (
+    ChatChoice,
+    ChatChoiceDelta,
+    ChatChunkChoice,
+    ChatCompletionChunk,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatMessage,
+    ChatUsage,  # noqa F401
+    Function,  # noqa F401
+    FunctionToolDefinition,
+    ToolCall,  # noqa F401
+)
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+# NB: `import x as y` does not work and will cause a Pydantic error.
+StreamDelta = ChatChoiceDelta
+StreamChoice = ChatChunkChoice
+RequestMessage = ChatMessage
+
+
+class UnityCatalogFunctionToolDefinition(RequestModel):
+    name: str
+
+
+class ChatToolWithUC(RequestModel):
+    """
+    A tool definition for the chat endpoint with Unity Catalog integration.
+    The Gateway request accepts a special tool type 'uc_function' for Unity Catalog integration.
+    https://mlflow.org/docs/latest/llms/deployments/uc_integration.html
+    """
+
+    type: Literal["function", "uc_function"]
+    function: Optional[FunctionToolDefinition] = None
+    uc_function: Optional[UnityCatalogFunctionToolDefinition] = None
+
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "messages": [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Hello!"},
+    ],
+    "temperature": 0.0,
+    "max_tokens": 64,
+    "stop": ["END"],
+    "n": 1,
+    "stream": False,
+}
+
+
+class RequestPayload(ChatCompletionRequest, RequestModel):
+    messages: list[RequestMessage] = Field(..., min_items=1)
+    tools: Optional[list[ChatToolWithUC]] = None
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "3cdb958c-e4cc-4834-b52b-1d1a7f324714",
+        "object": "chat.completion",
+        "created": 1700173217,
+        "model": "llama-2-70b-chat-hf",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hello! I am an AI assistant"},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
+    }
+}
+
+
+class ResponseMessage(ChatMessage, ResponseModel):
+    # Override the `tool_call_id` field to be excluded from the response.
+    # This is a band-aid solution to avoid exposing the tool_call_id in the response,
+    # while we use the same ChatMessage model for both request and response.
+    tool_call_id: Optional[str] = Field(None, exclude=True)
+
+
+class Choice(ChatChoice, ResponseModel):
+    # Override the `message` field to use the ResponseMessage model.
+    message: ResponseMessage
+
+
+class ResponsePayload(ChatCompletionResponse, ResponseModel):
+    # Override the `choices` field to use the Choice model
+    choices: list[Choice]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+
+
+_STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "3cdb958c-e4cc-4834-b52b-1d1a7f324714",
+        "object": "chat.completion",
+        "created": 1700173217,
+        "model": "llama-2-70b-chat-hf",
+        "choices": [
+            {
+                "index": 6,
+                "finish_reason": "stop",
+                "delta": {"role": "assistant", "content": "you?"},
+            }
+        ],
+    }
+}
+
+
+class StreamResponsePayload(ChatCompletionChunk, ResponseModel):
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/completions.py

@@ -0,0 +1,109 @@
+from typing import Optional
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+from mlflow.types.chat import BaseRequestPayload
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "prompt": "hello",
+        "temperature": 0.0,
+        "max_tokens": 64,
+        "stop": ["END"],
+        "n": 1,
+    }
+}
+
+
+class RequestPayload(BaseRequestPayload, RequestModel):
+    prompt: str
+    model: Optional[str] = None
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+class Choice(ResponseModel):
+    index: int
+    text: str
+    finish_reason: Optional[str] = None
+
+
+class CompletionsUsage(ResponseModel):
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "cmpl-123",
+        "object": "text_completion",
+        "created": 1589478378,
+        "model": "gpt-4",
+        "choices": [
+            {"text": "Hello! I am an AI Assistant!", "index": 0, "finish_reason": "length"}
+        ],
+        "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12},
+    }
+}
+
+
+class ResponsePayload(ResponseModel):
+    id: Optional[str] = None
+    object: str = "text_completion"
+    created: int
+    model: str
+    choices: list[Choice]
+    usage: CompletionsUsage
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+
+
+class StreamDelta(ResponseModel):
+    role: Optional[str] = None
+    content: Optional[str] = None
+
+
+class StreamChoice(ResponseModel):
+    index: int
+    finish_reason: Optional[str] = None
+    text: Optional[str] = None
+
+
+_STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "cmpl-123",
+        "object": "text_completion",
+        "created": 1589478378,
+        "model": "gpt-4",
+        "choices": [
+            {
+                "index": 6,
+                "finish_reason": "stop",
+                "delta": {"role": "assistant", "content": "you?"},
+            }
+        ],
+    }
+}
+
+
+class StreamResponsePayload(ResponseModel):
+    id: Optional[str] = None
+    object: str = "text_completion_chunk"
+    created: int
+    model: str
+    choices: list[StreamChoice]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/embeddings.py

@@ -0,0 +1,97 @@
+from typing import Optional, Union
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "input": ["hello", "world"],
+    }
+}
+
+
+class RequestPayload(RequestModel):
+    input: Union[str, list[str], list[int], list[list[int]]]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+class EmbeddingObject(ResponseModel):
+    object: str = "embedding"
+    embedding: Union[list[float], str]
+    index: int
+
+
+class EmbeddingsUsage(ResponseModel):
+    prompt_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "object": "list",
+    "data": [
+        {
+            "object": "embedding",
+            "index": 0,
+            "embedding": [
+                0.017291732,
+                -0.017291732,
+                0.014577783,
+                -0.02902633,
+                -0.037271563,
+                0.019333655,
+                -0.023055641,
+                -0.007359971,
+                -0.015818445,
+                -0.030654699,
+                0.008348623,
+                0.018312693,
+                -0.017149571,
+                -0.0044424757,
+                -0.011165961,
+                0.01018377,
+            ],
+        },
+        {
+            "object": "embedding",
+            "index": 1,
+            "embedding": [
+                0.0060126893,
+                -0.008691099,
+                -0.0040095365,
+                0.019889368,
+                0.036211833,
+                -0.0013270887,
+                0.013401738,
+                -0.0036735237,
+                -0.0049594184,
+                0.035229642,
+                -0.03435084,
+                0.019798903,
+                -0.0006110424,
+                0.0073793563,
+                0.005657291,
+                0.022487005,
+            ],
+        },
+    ],
+    "model": "text-embedding-ada-002-v2",
+    "usage": {"prompt_tokens": 400, "total_tokens": 400},
+}
+
+
+class ResponsePayload(ResponseModel):
+    object: str = "list"
+    data: list[EmbeddingObject]
+    model: str
+    usage: EmbeddingsUsage
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA