Source code for imagine.types.completions

from __future__ import annotations

from pydantic import BaseModel

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


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

    #: A 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 CompletionStreamResponse(BaseModel):
    #: A unique identifier for the completion
    id: str

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

    #: A list of completion choices
    choices: list[CompletionResponseStreamChoice]

    #: The Unix timestamp of when the 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 CompletionResponseChoice(BaseModel):
    #: The index of the choice in the list of choices.
    index: int

    #: Text content generated by the model
    text: str

    #: 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 CompletionResponse(BaseModel):
    #: A unique identifier for the completion.
    id: str

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

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

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

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

    #: Usage Statistics
    usage: UsageInfo | None = None

    #: generation time.
    generation_time: float | None = None

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

        :return: text
        """
        return self.choices[0].text





class CompletionRequest(LLMSamplingParams):
    #: Prompt for completion request
    prompt: str

    #: Model to be used for Completion Request
    model: str

    #: Should it be a Streaming Request
    stream: bool