seekrai 0.0.1__py3-none-any.whl

seekrai/types/__init__.py

@@ -0,0 +1,67 @@
+from seekrai.types.abstract import SeekrFlowClient
+from seekrai.types.chat_completions import (
+    ChatCompletionChunk,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+)
+from seekrai.types.common import SeekrFlowRequest
+from seekrai.types.completions import (
+    CompletionChunk,
+    CompletionRequest,
+    CompletionResponse,
+)
+from seekrai.types.embeddings import EmbeddingRequest, EmbeddingResponse
+from seekrai.types.files import (
+    FileDeleteResponse,
+    FileList,
+    FileObject,
+    FilePurpose,
+    FileRequest,
+    FileResponse,
+    FileType,
+)
+from seekrai.types.finetune import (
+    FinetuneDownloadResult,
+    FinetuneList,
+    FinetuneListEvents,
+    FinetuneRequest,
+    FinetuneResponse,
+InfrastructureConfig,
+TrainingConfig,
+)
+from seekrai.types.images import (
+    ImageRequest,
+    ImageResponse,
+)
+from seekrai.types.models import ModelObject
+
+
+__all__ = [
+    "SeekrFlowClient",
+    "SeekrFlowRequest",
+    "CompletionChunk",
+    "CompletionRequest",
+    "CompletionResponse",
+    "ChatCompletionChunk",
+    "ChatCompletionRequest",
+    "ChatCompletionResponse",
+    "EmbeddingRequest",
+    "EmbeddingResponse",
+    "FinetuneRequest",
+    "FinetuneResponse",
+    "FinetuneList",
+    "FinetuneListEvents",
+    "FinetuneDownloadResult",
+    "InfrastructureConfig",
+    "TrainingConfig",
+    "FileRequest",
+    "FileResponse",
+    "FileList",
+    "FileDeleteResponse",
+    "FileObject",
+    "FilePurpose",
+    "FileType",
+    "ImageRequest",
+    "ImageResponse",
+    "ModelObject",
+]

seekrai/types/abstract.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+import pydantic
+from pydantic import ConfigDict
+from typing_extensions import ClassVar
+
+from seekrai.constants import BASE_URL, MAX_RETRIES, TIMEOUT_SECS
+
+
+PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
+
+
+@dataclass
+class SeekrFlowClient:
+    api_key: str | None = None
+    base_url: str | None = BASE_URL
+    timeout: float | None = TIMEOUT_SECS
+    max_retries: int | None = MAX_RETRIES
+    supplied_headers: Dict[str, str] | None = None
+
+
+class BaseModel(pydantic.BaseModel):
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

seekrai/types/chat_completions.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Dict, List
+
+from pydantic import Field
+
+from seekrai.types.abstract import BaseModel
+from seekrai.types.common import (
+    DeltaContent,
+    FinishReason,
+    LogprobsPart,
+    ObjectType,
+    PromptPart,
+    UsageData,
+)
+
+
+class MessageRole(str, Enum):
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    USER = "user"
+
+
+class ResponseFormatType(str, Enum):
+    JSON_OBJECT = "json_object"
+
+
+class FunctionCall(BaseModel):
+    name: str | None = None
+    arguments: str | None = None
+
+
+class ToolCalls(BaseModel):
+    id: str | None = None
+    type: str | None = None
+    function: FunctionCall | None = None
+
+
+class ChatCompletionMessage(BaseModel):
+    role: MessageRole
+    content: str | None = None
+    # tool_calls: List[ToolCalls] | None = None
+
+
+class ResponseFormat(BaseModel):
+    type: ResponseFormatType
+    schema_: Dict[str, Any] | None = Field(None, alias="schema")
+
+
+class FunctionTool(BaseModel):
+    description: str | None = None
+    name: str
+    parameters: Dict[str, Any] | None = None
+
+
+class FunctionToolChoice(BaseModel):
+    name: str
+
+
+class Tools(BaseModel):
+    type: str
+    function: FunctionTool
+
+
+class ToolChoice(BaseModel):
+    type: str
+    function: FunctionToolChoice
+
+
+class ToolChoiceEnum(str, Enum):
+    Auto = "auto"
+
+
+class ChatCompletionRequest(BaseModel):
+    # list of messages
+    messages: List[ChatCompletionMessage]
+    # model name
+    model: str
+    # stopping criteria: max tokens to generate
+    max_tokens: int | None = None
+    # stopping criteria: list of strings to stop generation
+    stop: List[str] | None = None
+    # sampling hyperparameters
+    temperature: float = 0.7
+    top_p: float = 1
+    top_k: int = -1
+    repetition_penalty: float = 1
+    # stream SSE token chunks
+    stream: bool = False
+    # return logprobs
+    logprobs: int = 0
+    # echo prompt.
+    # can be used with logprobs to return prompt logprobs
+    echo: bool = False
+    # number of output generations
+    n: int = 1
+    # moderation model
+    safety_model: str | None = None
+    # constraints
+    response_format: ResponseFormat | None = None
+    # tools: List[Tools] | None = None
+    # tool_choice: ToolChoice | ToolChoiceEnum | None = None
+
+
+class ChatCompletionChoicesData(BaseModel):
+    index: int | None = None
+    logprobs: LogprobsPart | None = None
+    finish_reason: FinishReason | None = None
+    message: ChatCompletionMessage | None = None
+
+
+class ChatCompletionResponse(BaseModel):
+    # request id
+    id: str | None = None
+    # object type
+    object: ObjectType | None = None
+    # created timestamp
+    created: int | None = None
+    # model name
+    model: str | None = None
+    # choices list
+    choices: List[ChatCompletionChoicesData] | None = None
+    # prompt list
+    prompt: List[PromptPart] | List[None] | None = None
+    # token usage data
+    usage: UsageData | None = None
+
+
+class ChatCompletionChoicesChunk(BaseModel):
+    index: int | None = None
+    logprobs: float | None = None
+    finish_reason: FinishReason | None = None
+    delta: DeltaContent | None = None
+
+
+class ChatCompletionChunk(BaseModel):
+    # request id
+    id: str | None = None
+    # object type
+    object: ObjectType | None = None
+    # created timestamp
+    created: int | None = None
+    # model name
+    model: str | None = None
+    # delta content
+    choices: List[ChatCompletionChoicesChunk] | None = None
+    # finish reason
+    finish_reason: FinishReason | None = None
+    # token usage data
+    usage: UsageData | None = None

seekrai/types/common.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Dict, List
+
+from pydantic import ConfigDict
+from tqdm.utils import CallbackIOWrapper
+
+from seekrai.types.abstract import BaseModel
+
+
+# Generation finish reason
+class FinishReason(str, Enum):
+    Length = "length"
+    StopSequence = "stop"
+    EOS = "eos"
+    ToolCalls = "tool_calls"
+
+
+class UsageData(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class ObjectType(str, Enum):
+    Completion = "text.completion"
+    CompletionChunk = "completion.chunk"
+    ChatCompletion = "chat.completion"
+    ChatCompletionChunk = "chat.completion.chunk"
+    Embedding = "embedding"
+    FinetuneEvent = "fine-tune-event"
+    File = "file"
+    Model = "model"
+
+
+class LogprobsPart(BaseModel):
+    # token list
+    tokens: List[str | None] | None = None
+    # token logprob list
+    token_logprobs: List[float | None] | None = None
+
+
+class PromptPart(BaseModel):
+    # prompt string
+    text: str | None = None
+    # list of prompt logprobs
+    logprobs: LogprobsPart | None = None
+
+
+class DeltaContent(BaseModel):
+    content: str | None = None
+
+
+class SeekrFlowRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    method: str
+    url: str
+    headers: Dict[str, str] | None = None
+    params: Dict[str, Any] | CallbackIOWrapper | None = None
+    files: Dict[str, Any] | None = None
+    allow_redirects: bool = True
+    override_headers: bool = False

seekrai/types/completions.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import List
+
+from seekrai.types.abstract import BaseModel
+from seekrai.types.common import (
+    DeltaContent,
+    FinishReason,
+    LogprobsPart,
+    ObjectType,
+    PromptPart,
+    UsageData,
+)
+
+
+class CompletionRequest(BaseModel):
+    # prompt to complete
+    prompt: str
+    # query model
+    model: str
+    # stopping criteria: max tokens to generate
+    max_tokens: int | None = None
+    # stopping criteria: list of strings to stop generation
+    stop: List[str] | None = None
+    # sampling hyperparameters
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    repetition_penalty: float | None = None
+    # stream SSE token chunks
+    stream: bool = False
+    # return logprobs
+    logprobs: int | None = None
+    # echo prompt.
+    # can be used with logprobs to return prompt logprobs
+    echo: bool | None = None
+    # number of output generations
+    n: int | None = None
+    # moderation model
+    safety_model: str | None = None
+
+
+class CompletionChoicesData(BaseModel):
+    index: int
+    logprobs: LogprobsPart | None = None
+    finish_reason: FinishReason
+    text: str
+
+
+class CompletionChoicesChunk(BaseModel):
+    index: int | None = None
+    logprobs: float | None = None
+    finish_reason: FinishReason | None = None
+    delta: DeltaContent | None = None
+
+
+class CompletionResponse(BaseModel):
+    # request id
+    id: str | None = None
+    # object type
+    object: ObjectType | None = None
+    # created timestamp
+    created: int | None = None
+    # model name
+    model: str | None = None
+    # choices list
+    choices: List[CompletionChoicesData] | None = None
+    # prompt list
+    prompt: List[PromptPart] | None = None
+    # token usage data
+    usage: UsageData | None = None
+
+
+class CompletionChunk(BaseModel):
+    # request id
+    id: str | None = None
+    # object type
+    object: ObjectType | None = None
+    # created timestamp
+    created: int | None = None
+    # model name
+    model: str | None = None
+    # choices list
+    choices: List[CompletionChoicesChunk] | None = None
+    # token usage data
+    usage: UsageData | None = None

seekrai/types/embeddings.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import List, Literal
+
+from seekrai.types.abstract import BaseModel
+from seekrai.types.common import (
+    ObjectType,
+)
+
+
+class EmbeddingRequest(BaseModel):
+    # input or list of inputs
+    input: str | List[str]
+    # model to query
+    model: str
+
+
+class EmbeddingChoicesData(BaseModel):
+    # response index
+    index: int
+    # object type
+    object: ObjectType
+    # embedding response
+    embedding: List[float] | None = None
+
+
+class EmbeddingResponse(BaseModel):
+    # job id
+    id: str | None = None
+    # query model
+    model: str | None = None
+    # object type
+    object: Literal["list"] | None = None
+    # list of embedding choices
+    data: List[EmbeddingChoicesData] | None = None

seekrai/types/error.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from pydantic import Field
+
+from seekrai.types.abstract import BaseModel
+
+
+class SeekrFlowErrorResponse(BaseModel):
+    # error message
+    message: str | None = None
+    # error type
+    type_: str | None = Field(None, alias="type")
+    # param causing error
+    param: str | None = None
+    # error code
+    code: str | None = None

seekrai/types/files.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import List, Literal
+
+from pydantic import Field
+
+from seekrai.types.abstract import BaseModel
+from seekrai.types.common import (
+    ObjectType,
+)
+from datetime import datetime
+
+class FilePurpose(str, Enum):
+    FineTune = "fine-tune"
+
+
+class FileType(str, Enum):
+    jsonl = "jsonl"
+    parquet = "parquet"
+
+
+class FileRequest(BaseModel):
+    """
+    Files request type
+    """
+
+    # training file ID
+    training_file: str
+    # base model string
+    model: str
+    # number of epochs to train for
+    n_epochs: int
+    # training learning rate
+    learning_rate: float
+    # number of checkpoints to save
+    n_checkpoints: int | None = None
+    # training batch size
+    batch_size: int | None = None
+    # up to 40 character suffix for output model name
+    suffix: str | None = None
+    # weights & biases api key
+    wandb_api_key: str | None = None
+
+
+class FileResponse(BaseModel):
+    """
+    Files API response type
+    """
+
+    id: str
+    object: Literal[ObjectType.File]
+    # created timestamp
+    created_at: datetime | None = None
+    type: FileType | None = None
+    purpose: FilePurpose | None = None
+    filename: str | None = None
+    # file byte size
+    bytes: int | None = None
+    # JSONL/Parquet line count
+    line_count: int | None = None
+
+
+class FileList(BaseModel):
+    # object type
+    object: Literal["list"] | None = None
+    # list of fine-tune job objects
+    data: List[FileResponse] | None = None
+
+
+class FileDeleteResponse(BaseModel):
+    # file id
+    id: str
+    # object type
+    object: Literal[ObjectType.File]
+    # is deleted
+    deleted: bool
+
+
+class FileObject(BaseModel):
+    # object type
+    object: Literal["local"] | None = None
+    # fine-tune job id
+    id: str | None = None
+    # local path filename
+    filename: str | None = None
+    # size in bytes
+    size: int | None = None