Source code for imagine.types.chat_completions

from __future__ import annotations

from typing import Any

from pydantic import BaseModel

from imagine.types.common import (
    DeltaMessage,
    FinishReason,
    LLMSamplingParams,
    UsageInfo,
)


class FunctionTool(BaseModel):
    description: str | None = None
    name: str
    parameters: dict[str, Any] | None = None


class Tools(BaseModel):
    type: str
    function: FunctionTool


class Function(BaseModel):
    name: str | None = None
    arguments: str | None = None


class ChatMessageToolCall(BaseModel):
    id: str | None = None
    type: str | None = None
    function: Function | None = None




class ChatMessage(BaseModel):
    #: The role of the message `user`, `assistant`, `system`, `tool`
    role: str

    #: The content of the message
    content: str

    #: The tool calls generated by the model, such as function calls
    tool_calls: list[ChatMessageToolCall] | None = None

    #: The ID of the tool call.
    tool_call_id: str | None = None

    #: An optional name for the participant.
    name: str | None = None





class ChatCompletionResponseStreamChoice(BaseModel):
    #: The index of the choice in the list of choices.
    index: int

    #: A chat completion delta generated by streamed model responses.
    delta: DeltaMessage

    #: The reason the model stopped generating tokens.
    #: This will be `stop` if the model hit a natural stop point or a provided stop sequence,
    #: `length` if the maximum number of tokens specified in the request was reached
    #: `error` in case of error
    finish_reason: FinishReason | None





class ChatCompletionStreamResponse(BaseModel):
    #: A unique identifier for the chat completion
    id: str

    #: The model used for the chat completion.
    model: str

    #: A list of chat completion choices
    choices: list[ChatCompletionResponseStreamChoice]

    #: The Unix timestamp of when the chat completion was created.
    created: float | None = None

    #: The object type, which is always `chat.completion.chunk`.
    object: str | None = None

    #: Usage statistics for the completion request.
    usage: UsageInfo | None = None

    @property
    def first_content(self) -> str | None:
        """
        Gets the first content from the response

        :return: message content
        """
        return self.choices[0].delta.content





class ChatCompletionResponseChoice(BaseModel):
    #: The index of the choice in the list of choices.
    index: int

    #: A chat completion message generated by the model.
    message: ChatMessage

    #: The reason the model stopped generating tokens.
    #: This will be `stop` if the model hit a natural stop point or a provided stop sequence,
    #: `length` if the maximum number of tokens specified in the request was reached
    #: `error` in case of error
    finish_reason: FinishReason | None





class ChatCompletionResponse(BaseModel):
    #: A unique identifier for the chat completion.
    id: str

    #: The object type, which is always `chat.completion`.
    object: str

    #: The Unix timestamp of when the chat completion was created.
    created: float

    #: The model used for the chat completion.
    model: str

    #: A list of chat completion choices
    choices: list[ChatCompletionResponseChoice]

    #: Usage statistics for the completion request.
    usage: UsageInfo

    @property
    def first_content(self) -> str | None:
        """
        Gets the first content from the response

        :return: message content
        """
        return self.choices[0].message.content





class ChatCompletionRequest(LLMSamplingParams):
    #: A list of [messages]
    messages: list[ChatMessage]

    #: The model to be used for the chat completion.
    model: str

    #: If set, partial message deltas will be sent
    stream: bool

    # List of tools that can be used for chat completions
    tools: list[Tools] | None = None