129 lines
3.4 KiB
Python
129 lines
3.4 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
from mlflow.models import ModelSignature
|
|
from mlflow.types.schema import (
|
|
Array,
|
|
ColSpec,
|
|
DataType,
|
|
Object,
|
|
Property,
|
|
Schema,
|
|
)
|
|
from mlflow.utils.annotations import deprecated
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatMessage")
|
|
@dataclass
|
|
class Message:
|
|
role: str = "user" # "system", "user", or "assistant"
|
|
content: str = "What is mlflow?"
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionRequest")
|
|
@dataclass
|
|
class ChatCompletionRequest:
|
|
messages: list[Message] = field(default_factory=lambda: [Message()])
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionRequest")
|
|
@dataclass
|
|
class SplitChatMessagesRequest:
|
|
query: str = "What is mlflow?"
|
|
history: Optional[list[Message]] = field(default_factory=list)
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionRequest")
|
|
@dataclass
|
|
class MultiturnChatRequest:
|
|
query: str = "What is mlflow?"
|
|
history: Optional[list[Message]] = field(default_factory=list)
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatChoice")
|
|
@dataclass
|
|
class ChainCompletionChoice:
|
|
index: int = 0
|
|
message: Message = field(
|
|
default_factory=lambda: Message(
|
|
role="assistant",
|
|
content="MLflow is an open source platform for the machine learning lifecycle.",
|
|
)
|
|
)
|
|
finish_reason: str = "stop"
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionChunk")
|
|
@dataclass
|
|
class ChainCompletionChunk:
|
|
index: int = 0
|
|
delta: Message = field(
|
|
default_factory=lambda: Message(
|
|
role="assistant",
|
|
content="MLflow is an open source platform for the machine learning lifecycle.",
|
|
)
|
|
)
|
|
finish_reason: str = "stop"
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionResponse")
|
|
@dataclass
|
|
class ChatCompletionResponse:
|
|
choices: list[ChainCompletionChoice] = field(default_factory=lambda: [ChainCompletionChoice()])
|
|
object: str = "chat.completion"
|
|
# TODO: support ChainCompletionChunk in the future
|
|
|
|
|
|
@deprecated("mlflow.types.llm.ChatCompletionResponse")
|
|
@dataclass
|
|
class StringResponse:
|
|
content: str = "MLflow is an open source platform for the machine learning lifecycle."
|
|
|
|
|
|
CHAT_COMPLETION_REQUEST_SCHEMA = Schema(
|
|
[
|
|
ColSpec(
|
|
name="messages",
|
|
type=Array(
|
|
Object(
|
|
[
|
|
Property("role", DataType.string),
|
|
Property("content", DataType.string),
|
|
]
|
|
)
|
|
),
|
|
),
|
|
]
|
|
)
|
|
|
|
CHAT_COMPLETION_RESPONSE_SCHEMA = Schema(
|
|
[
|
|
ColSpec(
|
|
name="choices",
|
|
type=Array(
|
|
Object(
|
|
[
|
|
Property("index", DataType.long),
|
|
Property(
|
|
"message",
|
|
Object(
|
|
[
|
|
Property("role", DataType.string),
|
|
Property("content", DataType.string),
|
|
]
|
|
),
|
|
),
|
|
Property("finish_reason", DataType.string),
|
|
]
|
|
)
|
|
),
|
|
),
|
|
]
|
|
)
|
|
|
|
SIGNATURE_FOR_LLM_INFERENCE_TASK = {
|
|
"llm/v1/chat": ModelSignature(
|
|
inputs=CHAT_COMPLETION_REQUEST_SCHEMA, outputs=CHAT_COMPLETION_RESPONSE_SCHEMA
|
|
),
|
|
}
|