seekrai 0.0.1__py3-none-any.whl

seekrai/cli/api/chat.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import cmd
+import json
+from typing import List, Tuple
+
+import click
+
+from seekrai import seekrflow
+from seekrai.types.chat_completions import (
+    ChatCompletionChoicesChunk,
+    ChatCompletionChunk,
+    ChatCompletionResponse,
+)
+
+
+class ChatShell(cmd.Cmd):
+    intro = "Type /exit to exit, /help, or /? to list commands.\n"
+    prompt = ">>> "
+
+    def __init__(
+        self,
+        client: SeekrFlow,
+        model: str,
+        max_tokens: int | None = None,
+        stop: List[str] | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        top_k: int | None = None,
+        repetition_penalty: float | None = None,
+        safety_model: str | None = None,
+        system_message: str | None = None,
+    ) -> None:
+        super().__init__()
+        self.client = client
+        self.model = model
+        self.max_tokens = max_tokens
+        self.stop = stop
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+        self.repetition_penalty = repetition_penalty
+        self.safety_model = safety_model
+        self.system_message = system_message
+
+        self.messages = (
+            [{"role": "system", "content": self.system_message}]
+            if self.system_message
+            else []
+        )
+
+    def precmd(self, line: str) -> str:
+        if line.startswith("/"):
+            return line[1:]
+        else:
+            return "say " + line
+
+    def do_say(self, arg: str) -> None:
+        self.messages.append({"role": "user", "content": arg})
+
+        output = ""
+
+        for chunk in self.client.chat.completions.create(
+            messages=self.messages,
+            model=self.model,
+            max_tokens=self.max_tokens,
+            stop=self.stop,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            repetition_penalty=self.repetition_penalty,
+            safety_model=self.safety_model,
+            stream=True,
+        ):
+            # assertions for type checking
+            assert isinstance(chunk, ChatCompletionChunk)
+            assert chunk.choices
+            assert chunk.choices[0].delta
+            assert chunk.choices[0].delta.content
+
+            token = chunk.choices[0].delta.content
+
+            click.echo(token, nl=False)
+
+            output += token
+
+        click.echo("\n")
+
+        self.messages.append({"role": "assistant", "content": output})
+
+    def do_reset(self, arg: str) -> None:
+        self.messages = (
+            [{"role": "system", "content": self.system_message}]
+            if self.system_message
+            else []
+        )
+
+    def do_exit(self, arg: str) -> bool:
+        return True
+
+
+@click.command(name="chat.interactive")
+@click.pass_context
+@click.option("--model", type=str, required=True, help="Model name")
+@click.option("--max-tokens", type=int, help="Max tokens to generate")
+@click.option(
+    "--stop", type=str, multiple=True, help="List of strings to stop generation"
+)
+@click.option("--temperature", type=float, help="Sampling temperature")
+@click.option("--top-p", type=int, help="Top p sampling")
+@click.option("--top-k", type=float, help="Top k sampling")
+@click.option("--safety-model", type=str, help="Moderation model")
+@click.option("--system-message", type=str, help="System message to use for the chat")
+def interactive(
+    ctx: click.Context,
+    model: str,
+    max_tokens: int | None = None,
+    stop: List[str] | None = None,
+    temperature: float | None = None,
+    top_p: float | None = None,
+    top_k: int | None = None,
+    repetition_penalty: float | None = None,
+    safety_model: str | None = None,
+    system_message: str | None = None,
+) -> None:
+    """Interactive chat shell"""
+    client: SeekrFlow = ctx.obj
+
+    ChatShell(
+        client=client,
+        model=model,
+        max_tokens=max_tokens,
+        stop=stop,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        safety_model=safety_model,
+        system_message=system_message,
+    ).cmdloop()
+
+
+@click.command(name="chat.completions")
+@click.pass_context
+@click.option(
+    "--message",
+    type=(str, str),
+    multiple=True,
+    required=True,
+    help="Message to generate chat completions from",
+)
+@click.option("--model", type=str, required=True, help="Model name")
+@click.option("--max-tokens", type=int, help="Max tokens to generate")
+@click.option(
+    "--stop", type=str, multiple=True, help="List of strings to stop generation"
+)
+@click.option("--temperature", type=float, help="Sampling temperature")
+@click.option("--top-p", type=int, help="Top p sampling")
+@click.option("--top-k", type=float, help="Top k sampling")
+@click.option("--repetition-penalty", type=float, help="Repetition penalty")
+@click.option("--no-stream", is_flag=True, help="Disable streaming")
+@click.option("--logprobs", type=int, help="Return logprobs. Only works with --raw.")
+@click.option("--echo", is_flag=True, help="Echo prompt. Only works with --raw.")
+@click.option("--n", type=int, help="Number of output generations")
+@click.option("--safety-model", type=str, help="Moderation model")
+@click.option("--raw", is_flag=True, help="Output raw JSON")
+def chat(
+    ctx: click.Context,
+    message: List[Tuple[str, str]],
+    model: str,
+    max_tokens: int | None = None,
+    stop: List[str] | None = None,
+    temperature: float | None = None,
+    top_p: float | None = None,
+    top_k: int | None = None,
+    repetition_penalty: float | None = None,
+    no_stream: bool = False,
+    logprobs: int | None = None,
+    echo: bool | None = None,
+    n: int | None = None,
+    safety_model: str | None = None,
+    raw: bool = False,
+) -> None:
+    """Generate chat completions from messages"""
+    client: SeekrFlow = ctx.obj
+
+    messages = [{"role": msg[0], "content": msg[1]} for msg in message]
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stop=stop,
+        repetition_penalty=repetition_penalty,
+        stream=not no_stream,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model,
+    )
+
+    if not no_stream:
+        for chunk in response:
+            # assertions for type checking
+            assert isinstance(chunk, ChatCompletionChunk)
+            assert chunk.choices
+
+            if raw:
+                click.echo(f"{json.dumps(chunk.model_dump())}")
+                continue
+
+            should_print_header = len(chunk.choices) > 1
+            for stream_choice in sorted(chunk.choices, key=lambda c: c.index):  # type: ignore
+                assert isinstance(stream_choice, ChatCompletionChoicesChunk)
+                assert stream_choice.delta
+
+                if should_print_header:
+                    click.echo(f"\n===== Completion {stream_choice.index} =====\n")
+                click.echo(f"{stream_choice.delta.content}", nl=False)
+
+                if should_print_header:
+                    click.echo("\n")
+
+        # new line after stream ends
+        click.echo("\n")
+    else:
+        # assertions for type checking
+        assert isinstance(response, ChatCompletionResponse)
+        assert isinstance(response.choices, list)
+
+        if raw:
+            click.echo(f"{json.dumps(response.model_dump(), indent=4)}")
+            return
+
+        should_print_header = len(response.choices) > 1
+        for i, choice in enumerate(response.choices):
+            if should_print_header:
+                click.echo(f"===== Completion {i} =====")
+            click.echo(choice.message.content)  # type: ignore
+
+            if should_print_header:
+                click.echo("\n")

seekrai/cli/api/completions.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import json
+from typing import List
+
+import click
+
+from seekrai import seekrflow
+from seekrai.types import CompletionChunk
+from seekrai.types.completions import CompletionChoicesChunk, CompletionResponse
+
+
+@click.command()
+@click.pass_context
+@click.argument("prompt", type=str, required=True)
+@click.option("--model", type=str, required=True, help="Model name")
+@click.option("--no-stream", is_flag=True, help="Disable streaming")
+@click.option("--max-tokens", type=int, help="Max tokens to generate")
+@click.option(
+    "--stop", type=str, multiple=True, help="List of strings to stop generation"
+)
+@click.option("--temperature", type=float, help="Sampling temperature")
+@click.option("--top-p", type=int, help="Top p sampling")
+@click.option("--top-k", type=float, help="Top k sampling")
+@click.option("--logprobs", type=int, help="Return logprobs. Only works with --raw.")
+@click.option("--echo", is_flag=True, help="Echo prompt. Only works with --raw.")
+@click.option("--n", type=int, help="Number of output generations")
+@click.option("--safety-model", type=str, help="Moderation model")
+@click.option("--raw", is_flag=True, help="Return raw JSON response")
+def completions(
+    ctx: click.Context,
+    prompt: str,
+    model: str,
+    max_tokens: int | None = 512,
+    stop: List[str] | None = None,
+    temperature: float | None = None,
+    top_p: float | None = None,
+    top_k: int | None = None,
+    repetition_penalty: float | None = None,
+    no_stream: bool = False,
+    logprobs: int | None = None,
+    echo: bool | None = None,
+    n: int | None = None,
+    safety_model: str | None = None,
+    raw: bool = False,
+) -> None:
+    """Generate text completions"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.completions.create(
+        model=model,
+        prompt=prompt,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stop=stop,
+        repetition_penalty=repetition_penalty,
+        stream=not no_stream,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model,
+    )
+
+    if not no_stream:
+        for chunk in response:
+            # assertions for type checking
+            assert isinstance(chunk, CompletionChunk)
+            assert chunk.choices
+
+            if raw:
+                click.echo(f"{json.dumps(chunk.model_dump())}")
+                continue
+
+            should_print_header = len(chunk.choices) > 1
+            for stream_choice in sorted(chunk.choices, key=lambda c: c.index):  # type: ignore
+                # assertions for type checking
+                assert isinstance(stream_choice, CompletionChoicesChunk)
+                assert stream_choice.delta
+
+                if should_print_header:
+                    click.echo(f"\n===== Completion {stream_choice.index} =====\n")
+                click.echo(f"{stream_choice.delta.content}", nl=False)
+
+                if should_print_header:
+                    click.echo("\n")
+
+        # new line after stream ends
+        click.echo("\n")
+    else:
+        # assertions for type checking
+        assert isinstance(response, CompletionResponse)
+        assert isinstance(response.choices, list)
+
+        if raw:
+            click.echo(f"{json.dumps(response.model_dump(), indent=4)}")
+            return
+
+        should_print_header = len(response.choices) > 1
+        for i, choice in enumerate(response.choices):
+            if should_print_header:
+                click.echo(f"===== Completion {i} =====")
+            click.echo(choice.text)
+
+            if should_print_header or not choice.text.endswith("\n"):
+                click.echo("\n")

seekrai/cli/api/files.py

@@ -0,0 +1,125 @@
+import json
+import pathlib
+from textwrap import wrap
+
+import click
+from tabulate import tabulate
+
+from seekrai import seekrflow
+from seekrai.types import FilePurpose
+from seekrai.utils import check_file, convert_bytes, convert_unix_timestamp
+
+
+@click.group()
+@click.pass_context
+def files(ctx: click.Context) -> None:
+    """File API commands"""
+    pass
+
+
+@files.command()
+@click.pass_context
+@click.argument(
+    "file",
+    type=click.Path(
+        exists=True, file_okay=True, resolve_path=True, readable=True, dir_okay=False
+    ),
+    required=True,
+)
+@click.option(
+    "--purpose",
+    type=str,
+    default=FilePurpose.FineTune.value,
+    help="Purpose of file upload. Acceptable values in enum `seekrai.types.FilePurpose`. Defaults to `fine-tunes`.",
+)
+def upload(ctx: click.Context, file: pathlib.Path, purpose: str) -> None:
+    """Upload file"""
+
+    client: SeekrFlow = ctx.obj
+
+    response = client.files.upload(file=file, purpose=purpose)
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@files.command()
+@click.pass_context
+def list(ctx: click.Context) -> None:
+    """List files"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.files.list()
+
+    display_list = []
+    for i in response.data or []:
+        display_list.append(
+            {
+                "File name": "\n".join(wrap(i.filename or "", width=30)),
+                "File ID": i.id,
+                "Size": convert_bytes(
+                    float(str(i.bytes))
+                ),  # convert to string for mypy typing
+                "Created At": convert_unix_timestamp(i.created_at or 0),
+                "Line Count": i.line_count,
+            }
+        )
+    table = tabulate(display_list, headers="keys", tablefmt="grid", showindex=True)
+
+    click.echo(table)
+
+
+@files.command()
+@click.pass_context
+@click.argument("id", type=str, required=True)
+def retrieve(ctx: click.Context, id: str) -> None:
+    """Upload file"""
+
+    client: SeekrFlow = ctx.obj
+
+    response = client.files.retrieve(id=id)
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@files.command()
+@click.pass_context
+@click.argument("id", type=str, required=True)
+@click.option("--output", type=str, default=None, help="Output filename")
+def retrieve_content(ctx: click.Context, id: str, output: str) -> None:
+    """Retrieve file content and output to file"""
+
+    client: SeekrFlow = ctx.obj
+
+    response = client.files.retrieve_content(id=id, output=output)
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@files.command()
+@click.pass_context
+@click.argument("id", type=str, required=True)
+def delete(ctx: click.Context, id: str) -> None:
+    """Delete remote file"""
+
+    client: SeekrFlow = ctx.obj
+
+    response = client.files.delete(id=id)
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@files.command()
+@click.pass_context
+@click.argument(
+    "file",
+    type=click.Path(
+        exists=True, file_okay=True, resolve_path=True, readable=True, dir_okay=False
+    ),
+    required=True,
+)
+def check(ctx: click.Context, file: pathlib.Path) -> None:
+    """Check file for issues"""
+
+    report = check_file(file)
+
+    click.echo(json.dumps(report, indent=4))

seekrai/cli/api/finetune.py

@@ -0,0 +1,175 @@
+import json
+from textwrap import wrap
+
+import click
+from tabulate import tabulate
+
+from seekrai import seekrflow
+from seekrai.utils import finetune_price_to_dollars, parse_timestamp
+
+
+@click.group(name="fine-tuning")
+@click.pass_context
+def fine_tuning(ctx: click.Context) -> None:
+    """Fine-tunes API commands"""
+    pass
+
+
+@fine_tuning.command()
+@click.pass_context
+@click.option(
+    "--training-file", type=str, required=True, help="Training file ID from Files API"
+)
+@click.option("--model", type=str, required=True, help="Base model name")
+@click.option("--n-epochs", type=int, default=1, help="Number of epochs to train for")
+@click.option(
+    "--n-checkpoints", type=int, default=1, help="Number of checkpoints to save"
+)
+@click.option("--batch-size", type=int, default=32, help="Train batch size")
+@click.option("--learning-rate", type=float, default=3e-5, help="Learning rate")
+@click.option(
+    "--suffix", type=str, default=None, help="Suffix for the fine-tuned model name"
+)
+@click.option("--wandb-api-key", type=str, default=None, help="Wandb API key")
+def create(
+    ctx: click.Context,
+    training_file: str,
+    model: str,
+    n_epochs: int,
+    n_checkpoints: int,
+    batch_size: int,
+    learning_rate: float,
+    suffix: str,
+    wandb_api_key: str,
+) -> None:
+    """Start fine-tuning"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.create(
+        training_file=training_file,
+        model=model,
+        n_epochs=n_epochs,
+        n_checkpoints=n_checkpoints,
+        batch_size=batch_size,
+        learning_rate=learning_rate,
+        suffix=suffix,
+        wandb_api_key=wandb_api_key,
+    )
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@fine_tuning.command()
+@click.pass_context
+def list(ctx: click.Context) -> None:
+    """List fine-tuning jobs"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.list()
+
+    response.data = response.data or []
+
+    response.data.sort(key=lambda x: parse_timestamp(x.created_at or ""))
+
+    display_list = []
+    for i in response.data:
+        display_list.append(
+            {
+                "Fine-tune ID": i.id,
+                "Model Output Name": "\n".join(wrap(i.output_name or "", width=30)),
+                "Status": i.status,
+                "Created At": i.created_at,
+                "Price": f"""${finetune_price_to_dollars(
+                    float(str(i.total_price))
+                )}""",  # convert to string for mypy typing
+            }
+        )
+    table = tabulate(display_list, headers="keys", tablefmt="grid", showindex=True)
+
+    click.echo(table)
+
+
+@fine_tuning.command()
+@click.pass_context
+@click.argument("fine_tune_id", type=str, required=True)
+def retrieve(ctx: click.Context, fine_tune_id: str) -> None:
+    """Retrieve fine-tuning job details"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.retrieve(fine_tune_id)
+
+    # remove events from response for cleaner output
+    response.events = None
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@fine_tuning.command()
+@click.pass_context
+@click.argument("fine_tune_id", type=str, required=True)
+def cancel(ctx: click.Context, fine_tune_id: str) -> None:
+    """Cancel fine-tuning job"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.cancel(fine_tune_id)
+
+    click.echo(json.dumps(response.model_dump(), indent=4))
+
+
+@fine_tuning.command()
+@click.pass_context
+@click.argument("fine_tune_id", type=str, required=True)
+def list_events(ctx: click.Context, fine_tune_id: str) -> None:
+    """List fine-tuning events"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.list_events(fine_tune_id)
+
+    response.data = response.data or []
+
+    display_list = []
+    for i in response.data:
+        display_list.append(
+            {
+                "Message": "\n".join(wrap(i.message or "", width=50)),
+                "Type": i.type,
+                "Created At": parse_timestamp(i.created_at or ""),
+                "Hash": i.hash,
+            }
+        )
+    table = tabulate(display_list, headers="keys", tablefmt="grid", showindex=True)
+
+    click.echo(table)
+
+
+@fine_tuning.command()
+@click.pass_context
+@click.argument("fine_tune_id", type=str, required=True)
+@click.option(
+    "--output_dir",
+    type=click.Path(exists=True, file_okay=False, resolve_path=True),
+    required=False,
+    default=None,
+    help="Output directory",
+)
+@click.option(
+    "--checkpoint-step",
+    type=int,
+    required=False,
+    default=-1,
+    help="Download fine-tuning checkpoint. Defaults to latest.",
+)
+def download(
+    ctx: click.Context,
+    fine_tune_id: str,
+    output_dir: str,
+    checkpoint_step: int,
+) -> None:
+    """Download fine-tuning checkpoint"""
+    client: SeekrFlow = ctx.obj
+
+    response = client.fine_tuning.download(
+        fine_tune_id, output=output_dir, checkpoint_step=checkpoint_step
+    )
+
+    click.echo(json.dumps(response.model_dump(), indent=4))