seekrai 0.0.1__py3-none-any.whl

seekrai/cli/api/images.py

@@ -0,0 +1,82 @@
+import base64
+import pathlib
+
+import click
+from PIL import Image
+
+from seekrai import seekrflow
+from seekrai.types import ImageResponse
+from seekrai.types.images import ImageChoicesData
+
+
+@click.group()
+@click.pass_context
+def images(ctx: click.Context) -> None:
+    """Images generations API commands"""
+    pass
+
+
+@images.command()
+@click.pass_context
+@click.argument("prompt", type=str, required=True)
+@click.option("--model", type=str, required=True, help="Model name")
+@click.option("--steps", type=int, default=20, help="Number of steps to run generation")
+@click.option("--seed", type=int, default=None, help="Random seed")
+@click.option("--n", type=int, default=1, help="Number of images to generate")
+@click.option("--height", type=int, default=1024, help="Image height")
+@click.option("--width", type=int, default=1024, help="Image width")
+@click.option("--negative-prompt", type=str, default=None, help="Negative prompt")
+@click.option(
+    "--output",
+    type=click.Path(exists=True, file_okay=False, resolve_path=True),
+    required=False,
+    default=pathlib.Path("."),
+    help="Output directory",
+)
+@click.option("--prefix", type=str, required=False, default="image-")
+@click.option("--no-show", is_flag=True, help="Do not open images in viewer")
+def generate(
+    ctx: click.Context,
+    prompt: str,
+    model: str,
+    steps: int,
+    seed: int,
+    n: int,
+    height: int,
+    width: int,
+    negative_prompt: str,
+    output: pathlib.Path,
+    prefix: str,
+    no_show: bool,
+) -> None:
+    """Generate image"""
+
+    client: SeekrFlow = ctx.obj
+
+    response = client.images.generate(
+        prompt=prompt,
+        model=model,
+        steps=steps,
+        seed=seed,
+        n=n,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt,
+    )
+
+    assert isinstance(response, ImageResponse)
+    assert isinstance(response.data, list)
+
+    for i, choice in enumerate(response.data):
+        assert isinstance(choice, ImageChoicesData)
+
+        with open(f"{output}/{prefix}{choice.index}.png", "wb") as f:
+            f.write(base64.b64decode(choice.b64_json))
+
+        click.echo(
+            f"Image [{i + 1}/{len(response.data)}] saved to {output}/{prefix}{choice.index}.png"
+        )
+
+        if not no_show:
+            image = Image.open(f"{output}/{prefix}{choice.index}.png")
+            image.show()

seekrai/cli/api/models.py

@@ -0,0 +1,42 @@
+from textwrap import wrap
+
+import click
+from tabulate import tabulate
+
+from seekrai import seekrflow
+from seekrai.types.models import ModelObject
+
+
+@click.group()
+@click.pass_context
+def models(ctx: click.Context) -> None:
+    """Models API commands"""
+    pass
+
+
+@models.command()
+@click.pass_context
+def list(ctx: click.Context) -> None:
+    """List models"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.models.list()
+
+    display_list = []
+
+    model: ModelObject
+    for model in response:
+        display_list.append(
+            {
+                "ID": "\n".join(wrap(model.id or "", width=30)),
+                "Name": "\n".join(wrap(model.display_name or "", width=30)),
+                "Organization": model.organization,
+                "Type": model.type,
+                "Context Length": model.context_length,
+                "License": "\n".join(wrap(model.license or "", width=30)),
+                "Input per 1M token": model.pricing.input,
+                "Output per 1M token": model.pricing.output,
+            }
+        )
+
+    click.echo(tabulate(display_list, headers="keys", tablefmt="grid"))

seekrai/cli/cli.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import click
+
+import seekrai
+from seekrai.cli.api.chat import chat, interactive
+from seekrai.cli.api.completions import completions
+from seekrai.cli.api.files import files
+from seekrai.cli.api.finetune import fine_tuning
+from seekrai.cli.api.images import images
+from seekrai.cli.api.models import models
+from seekrai.constants import MAX_RETRIES, TIMEOUT_SECS
+
+
+def print_version(ctx: click.Context, params: Any, value: Any) -> None:
+    if not value or ctx.resilient_parsing:
+        return
+    click.echo(f"Version {seekrai.version}")
+    ctx.exit()
+
+
+@click.group()
+@click.pass_context
+@click.option(
+    "--api-key",
+    type=str,
+    help="API Key. Defaults to environment variable `SEEKRFLOW_API_KEY`",
+    default=os.getenv("SEEKRFLOW_API_KEY"),
+)
+@click.option(
+    "--base-url", type=str, help="API Base URL. Defaults to SeekrFlow AI endpoint."
+)
+@click.option(
+    "--timeout", type=int, help=f"Request timeout. Defaults to {TIMEOUT_SECS} seconds"
+)
+@click.option(
+    "--max-retries",
+    type=int,
+    help=f"Maximum number of HTTP retries. Defaults to {MAX_RETRIES}.",
+)
+@click.option(
+    "--version",
+    is_flag=True,
+    callback=print_version,
+    expose_value=False,
+    is_eager=True,
+    help="Print version",
+)
+@click.option("--debug", help="Debug mode", is_flag=True)
+def main(
+    ctx: click.Context,
+    api_key: str | None,
+    base_url: str | None,
+    timeout: int | None,
+    max_retries: int | None,
+    debug: bool | None,
+) -> None:
+    """This is a sample CLI tool."""
+    seekrai.log = "debug" if debug else None
+    ctx.obj = seekrai.SeekrFlow(
+        api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries
+    )
+
+
+main.add_command(chat)
+main.add_command(interactive)
+main.add_command(completions)
+main.add_command(images)
+main.add_command(files)
+main.add_command(fine_tuning)
+main.add_command(models)
+
+if __name__ == "__main__":
+    main()

seekrai/client.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import os
+from typing import Dict
+
+from seekrai import resources
+from seekrai.constants import BASE_URL, MAX_RETRIES, TIMEOUT_SECS
+from seekrai.error import AuthenticationError
+from seekrai.types import SeekrFlowClient
+from seekrai.utils import enforce_trailing_slash
+
+
+class SeekrFlow:
+    completions: resources.Completions
+    chat: resources.Chat
+    embeddings: resources.Embeddings
+    files: resources.Files
+    images: resources.Images
+    models: resources.Models
+    fine_tuning: resources.FineTuning
+
+    # client options
+    client: SeekrFlowClient
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        timeout: float | None = None,
+        max_retries: int | None = None,
+        supplied_headers: Dict[str, str] | None = None,
+    ) -> None:
+        """Construct a new synchronous seekrai client instance.
+
+        This automatically infers the following arguments from their corresponding environment variables if they are not provided:
+        - `api_key` from `SEEKRFLOW_API_KEY`
+        - `base_url` from `SEEKRFLOW_BASE_URL`
+        """
+
+        # get api key
+        if not api_key:
+            api_key = os.environ.get("SEEKR_API_KEY")
+
+        if not api_key:
+            raise AuthenticationError(
+                "The api_key client option must be set either by passing api_key to the client or by setting the "
+                "SEEKR_API_KEY environment variable"
+            )
+
+        # get base url
+        if not base_url:
+            base_url = os.environ.get("SEEKRFLOW_BASE_URL")
+
+        if not base_url:
+            base_url = BASE_URL
+
+        if timeout is None:
+            timeout = TIMEOUT_SECS
+
+        if max_retries is None:
+            max_retries = MAX_RETRIES
+
+        # SeekrFlowClient object
+        self.client = SeekrFlowClient(
+            api_key=api_key,
+            base_url=enforce_trailing_slash(base_url),
+            timeout=timeout,
+            max_retries=max_retries,
+            supplied_headers=supplied_headers,
+        )
+
+        self.completions = resources.Completions(self.client)
+        self.chat = resources.Chat(self.client)
+        self.embeddings = resources.Embeddings(self.client)
+        self.files = resources.Files(self.client)
+        self.images = resources.Images(self.client)
+        self.models = resources.Models(self.client)
+        self.fine_tuning = resources.FineTuning(self.client)
+        # self.instruction_fine_tuning = resources.InstructionFineTuning(self.client) # todo
+
+
+class AsyncSeekrFlow:
+    completions: resources.AsyncCompletions
+    chat: resources.AsyncChat
+    embeddings: resources.AsyncEmbeddings
+    files: resources.AsyncFiles
+    images: resources.AsyncImages
+    models: resources.AsyncModels
+    fine_tuning: resources.AsyncFineTuning
+
+    # client options
+    client: SeekrFlowClient
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        timeout: float | None = None,
+        max_retries: int | None = None,
+        supplied_headers: Dict[str, str] | None = None,
+    ) -> None:
+        """Construct a new async seekrai client instance.
+
+        This automatically infers the following arguments from their corresponding environment variables if they are not provided:
+        - `api_key` from `SEEKRFLOW_API_KEY`
+        - `base_url` from `SEEKRFLOW_BASE_URL`
+        """
+
+        # get api key
+        if not api_key:
+            api_key = os.environ.get("SEEKRFLOW_API_KEY")
+
+        if not api_key:
+            raise AuthenticationError(
+                "The api_key client option must be set either by passing api_key to the client or by setting the "
+                "SEEKRFLOW_API_KEY environment variable"
+            )
+
+        # get base url
+        if not base_url:
+            base_url = os.environ.get("SEEKRFLOW_BASE_URL")
+
+        if not base_url:
+            base_url = BASE_URL
+
+        if timeout is None:
+            timeout = TIMEOUT_SECS
+
+        if max_retries is None:
+            max_retries = MAX_RETRIES
+
+        # SeekrFlowClient object
+        self.client = SeekrFlowClient(
+            api_key=api_key,
+            base_url=enforce_trailing_slash(base_url),
+            timeout=timeout,
+            max_retries=max_retries,
+            supplied_headers=supplied_headers,
+        )
+
+        self.completions = resources.AsyncCompletions(self.client)
+        self.chat = resources.AsyncChat(self.client)
+        self.embeddings = resources.AsyncEmbeddings(self.client)
+        self.files = resources.AsyncFiles(self.client)
+        self.images = resources.AsyncImages(self.client)
+        self.models = resources.AsyncModels(self.client)
+        self.fine_tuning = resources.AsyncFineTuning(self.client)
+
+
+Client = SeekrFlow
+
+AsyncClient = AsyncSeekrFlow

seekrai/constants.py

@@ -0,0 +1,32 @@
+# Session constants
+TIMEOUT_SECS = 600
+MAX_SESSION_LIFETIME_SECS = 180
+MAX_CONNECTION_RETRIES = 2
+MAX_RETRIES = 5
+INITIAL_RETRY_DELAY = 0.5
+MAX_RETRY_DELAY = 8.0
+
+# API defaults
+# BASE_URL = "http://seekrflow-llm-training-svc.llm.k8s.prd.int.seekr.com/v1"
+BASE_URL = "http://localhost:8000/v1"
+
+# Download defaults
+DOWNLOAD_BLOCK_SIZE = 10 * 1024 * 1024  # 10 MB
+DISABLE_TQDM = False
+
+# Messages
+MISSING_API_KEY_MESSAGE = """SEEKR_API_KEY not found.
+Please set it as an environment variable or set it as seekrai.api_key
+Find your SEEKR_API_KEY at https://seekr.com/xxx"""
+
+# Minimum number of samples required for fine-tuning file
+MIN_SAMPLES = 100
+
+# the number of bytes in a gigabyte, used to convert bytes to GB for readable comparison
+NUM_BYTES_IN_GB = 2**30
+
+# maximum number of GB sized files we support finetuning for
+MAX_FILE_SIZE_GB = 4.9
+
+# expected columns for Parquet files
+PARQUET_EXPECTED_COLUMNS = ["input_ids", "attention_mask", "labels"]

seekrai/error.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Dict
+
+from requests import RequestException
+
+from seekrai.types.error import SeekrFlowErrorResponse
+
+
+class SeekrFlowException(Exception):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        headers: str | Dict[Any, Any] | None = None,
+        request_id: str | None = None,
+        http_status: int | None = None,
+    ) -> None:
+        _message = (
+            json.dumps(message.model_dump())
+            if isinstance(message, SeekrFlowErrorResponse)
+            else message
+        )
+        self._message = f"Error code: {http_status} - {_message}"
+
+        super(SeekrFlowException, self).__init__(self._message)
+
+        self.http_status = http_status
+        self.headers = headers or {}
+        self.request_id = request_id
+
+    def __repr__(self) -> str:
+        repr_message = json.dumps(
+            {
+                "response": self._message,
+                "status": self.http_status,
+                "request_id": self.request_id,
+                "headers": self.headers,
+            }
+        )
+        return "%s(%r)" % (self.__class__.__name__, repr_message)
+
+
+class AuthenticationError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class ResponseError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class JSONError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class InstanceError(SeekrFlowException):
+    def __init__(self, model: str | None = "model", **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.message = f"""No running instances for {model}.
+                You can start an instance with one of the following methods:
+                  1. navigating to the SeekrFlow Playground at api.seekrflow.ai
+                  2. starting one in python using seekrflow.Models.start(model_name)
+                  3. `$ seekrflow models start <MODEL_NAME>` at the command line.
+                See `seekrflow.Models.list()` in python or `$ seekrflow models list` in command line
+                to get an updated list of valid model names.
+                """
+
+
+class RateLimitError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class FileTypeError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class AttributeError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class Timeout(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class APIConnectionError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class InvalidRequestError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class APIError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class ServiceUnavailableError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)
+
+
+class DownloadError(SeekrFlowException):
+    def __init__(
+        self,
+        message: (
+            SeekrFlowErrorResponse | Exception | str | RequestException | None
+        ) = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(message=message, **kwargs)