seekrai 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

seekrai/client.py

@@ -34,8 +34,8 @@ class SeekrFlow:
         """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`
+        - `api_key` from `SEEKR_API_KEY`
+        - `base_url` from `SEEKR_BASE_URL`
         """
 
         # get api key
@@ -50,7 +50,7 @@ class SeekrFlow:
 
         # get base url
         if not base_url:
-            base_url = os.environ.get("SEEKRFLOW_BASE_URL")
+            base_url = os.environ.get("SEEKR_BASE_URL")
 
         if not base_url:
             base_url = BASE_URL
@@ -64,7 +64,7 @@ class SeekrFlow:
         # SeekrFlowClient object
         self.client = SeekrFlowClient(
             api_key=api_key,
-            base_url=enforce_trailing_slash(base_url),
+            base_url=enforce_trailing_slash(base_url),  # type: ignore
             timeout=timeout,
             max_retries=max_retries,
             supplied_headers=supplied_headers,
@@ -77,7 +77,6 @@ class SeekrFlow:
         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:
@@ -104,23 +103,23 @@ class AsyncSeekrFlow:
         """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`
+        - `api_key` from `SEEKR_API_KEY`
+        - `base_url` from `SEEKR_BASE_URL`
         """
 
         # get api key
         if not api_key:
-            api_key = os.environ.get("SEEKRFLOW_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 "
-                "SEEKRFLOW_API_KEY environment variable"
+                "SEEKR_API_KEY environment variable"
             )
 
         # get base url
         if not base_url:
-            base_url = os.environ.get("SEEKRFLOW_BASE_URL")
+            base_url = os.environ.get("SEEKR_BASE_URL")
 
         if not base_url:
             base_url = BASE_URL
@@ -134,7 +133,7 @@ class AsyncSeekrFlow:
         # SeekrFlowClient object
         self.client = SeekrFlowClient(
             api_key=api_key,
-            base_url=enforce_trailing_slash(base_url),
+            base_url=enforce_trailing_slash(base_url),  # type: ignore
             timeout=timeout,
             max_retries=max_retries,
             supplied_headers=supplied_headers,

seekrai/constants.py

@@ -1,32 +1,52 @@
+import os
+from typing import Any
+
+
+# Function to read from environment or default to hardcoded value
+def env_or_default(var_name: str, default: Any) -> Any:
+    value = os.getenv(f"SEEKR_{var_name}")
+    if value is not None:
+        return value.split(",") if isinstance(default, list) else value
+    return default
+
+
 # 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
+TIMEOUT_SECS = int(env_or_default("TIMEOUT_SECS", 600))
+MAX_SESSION_LIFETIME_SECS = int(env_or_default("MAX_SESSION_LIFETIME_SECS", 180))
+MAX_CONNECTION_RETRIES = int(env_or_default("MAX_CONNECTION_RETRIES", 2))
+MAX_RETRIES = int(env_or_default("MAX_RETRIES", 5))
+INITIAL_RETRY_DELAY = float(env_or_default("INITIAL_RETRY_DELAY", 0.5))
+MAX_RETRY_DELAY = float(env_or_default("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"
+BASE_URL = env_or_default("BASE_URL", "https://build.seekr.com/v1")
 
 # Download defaults
-DOWNLOAD_BLOCK_SIZE = 10 * 1024 * 1024  # 10 MB
-DISABLE_TQDM = False
+DOWNLOAD_BLOCK_SIZE = int(
+    env_or_default("DOWNLOAD_BLOCK_SIZE", 10 * 1024 * 1024)
+)  # 10 MB
+DISABLE_TQDM = bool(
+    int(env_or_default("DISABLE_TQDM", 0))
+)  # Assumes DISABLE_TQDM set as 0 (False) or 1 (True)
 
 # Messages
-MISSING_API_KEY_MESSAGE = """SEEKR_API_KEY not found.
+MISSING_API_KEY_MESSAGE = env_or_default(
+    "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"""
+Find your SEEKR_API_KEY at https://seekr.com/xxx""",
+)
 
 # Minimum number of samples required for fine-tuning file
-MIN_SAMPLES = 100
+MIN_SAMPLES = int(env_or_default("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
+NUM_BYTES_IN_GB = int(env_or_default("NUM_BYTES_IN_GB", 2**30))
 
 # maximum number of GB sized files we support finetuning for
-MAX_FILE_SIZE_GB = 4.9
+MAX_FILE_SIZE_GB = float(env_or_default("MAX_FILE_SIZE_GB", 4.9))
 
 # expected columns for Parquet files
-PARQUET_EXPECTED_COLUMNS = ["input_ids", "attention_mask", "labels"]
+PARQUET_EXPECTED_COLUMNS = env_or_default(
+    "PARQUET_EXPECTED_COLUMNS", ["input_ids", "attention_mask", "labels"]
+)

seekrai/error.py

@@ -79,14 +79,7 @@ class JSONError(SeekrFlowException):
 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.
-                """
+        self.message = f"""No running instances for {model}."""
 
 
 class RateLimitError(SeekrFlowException):

seekrai/filemanager.py

@@ -5,19 +5,17 @@ import shutil
 import stat
 import tempfile
 import uuid
-from functools import partial
 from pathlib import Path
 from typing import Tuple
 
+import httpx
 import requests
-from filelock import FileLock
 from requests.structures import CaseInsensitiveDict
 from tqdm import tqdm
-from tqdm.utils import CallbackIOWrapper
 
 import seekrai.utils
 from seekrai.abstract import api_requestor
-from seekrai.constants import DISABLE_TQDM, DOWNLOAD_BLOCK_SIZE, MAX_RETRIES
+from seekrai.constants import DISABLE_TQDM
 from seekrai.error import (
     APIError,
     AuthenticationError,
@@ -148,7 +146,6 @@ class DownloadManager:
                 url=url,
                 headers={"Range": "bytes=0-1"},
             ),
-            remaining_retries=MAX_RETRIES,
             stream=False,
         )
 
@@ -189,47 +186,43 @@ class DownloadManager:
             url, output, remote_name, fetch_metadata
         )
 
-        temp_file_manager = partial(
-            tempfile.NamedTemporaryFile, mode="wb", dir=file_path.parent, delete=False
-        )
-
         # Prevent parallel downloads of the same file with a lock.
         lock_path = Path(file_path.as_posix() + ".lock")
 
-        with FileLock(lock_path.as_posix()):
-            with temp_file_manager() as temp_file:
-                response = requestor.request_raw(
-                    options=SeekrFlowRequest(
-                        method="GET",
-                        url=url,
-                    ),
-                    remaining_retries=MAX_RETRIES,
-                    stream=True,
-                )
+        with tempfile.NamedTemporaryFile() as temp_file:
+            response = requestor.request_raw(
+                options=SeekrFlowRequest(
+                    method="GET",
+                    url=url,
+                ),
+                stream=True,
+            )
 
-                try:
-                    response.raise_for_status()
-                except Exception as e:
-                    os.remove(lock_path)
-                    raise APIError(
-                        "Error downloading file", http_status=response.status_code
-                    ) from e
-
-                if not fetch_metadata:
-                    file_size = int(response.headers.get("content-length", 0))
-
-                assert file_size != 0, "Unable to retrieve remote file."
-
-                with tqdm(
-                    total=file_size,
-                    unit="B",
-                    unit_scale=True,
-                    desc=f"Downloading file {file_path.name}",
-                    disable=bool(DISABLE_TQDM),
-                ) as pbar:
-                    for chunk in response.iter_content(DOWNLOAD_BLOCK_SIZE):
-                        pbar.update(len(chunk))
-                        temp_file.write(chunk)
+            try:
+                response.raise_for_status()
+            except Exception as e:
+                os.remove(lock_path)
+                raise APIError(
+                    "Error downloading file", http_status=response.status_code
+                ) from e
+
+            if not fetch_metadata:
+                file_size = int(response.headers.get("content-length", 0))
+
+            assert file_size != 0, "Unable to retrieve remote file."
+
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                desc=f"Downloading file {file_path.name}",
+                disable=bool(DISABLE_TQDM),
+            ) as pbar:
+                num_bytes_downloaded = response.num_bytes_downloaded
+                for chunk in response.iter_bytes():
+                    temp_file.write(chunk)
+                    pbar.update(response.num_bytes_downloaded - num_bytes_downloaded)
+                    num_bytes_downloaded = response.num_bytes_downloaded
 
             # Raise exception if remote file size does not match downloaded file size
             if os.stat(temp_file.name).st_size != file_size:
@@ -252,7 +245,7 @@ class UploadManager:
 
     @classmethod
     def _redirect_error_handler(
-        cls, requestor: api_requestor.APIRequestor, response: requests.Response
+        cls, requestor: api_requestor.APIRequestor, response: httpx.Response
     ) -> None:
         if response.status_code == 401:
             raise AuthenticationError(
@@ -296,7 +289,6 @@ class UploadManager:
                 override_headers=True,
                 headers=headers,
             ),
-            remaining_retries=MAX_RETRIES,
         )
 
         self._redirect_error_handler(requestor, response)
@@ -327,24 +319,16 @@ class UploadManager:
         purpose: FilePurpose,
         redirect: bool = False,
     ) -> FileResponse:
-        file_id = None
-
         requestor = api_requestor.APIRequestor(
             client=self._client,
         )
 
-        redirect_url = None
         if redirect:
-            if file.suffix == ".jsonl":
-                filetype = FileType.jsonl
-            elif file.suffix == ".parquet":
-                filetype = FileType.parquet
-            else:
+            if file.suffix not in [".jsonl", ".parquet", ".pt"]:
                 raise FileTypeError(
                     f"Unknown extension of file {file}. "
-                    "Only files with extensions .jsonl and .parquet are supported."
+                    "Only files with extensions .jsonl, .parquet, and .pt are supported."
                 )
-            # redirect_url, file_id = self.get_upload_url(url, file, purpose, filetype)
 
         file_size = os.stat(file.as_posix()).st_size
 
@@ -354,40 +338,17 @@ class UploadManager:
             unit_scale=True,
             desc=f"Uploading file {file.name}",
             disable=bool(DISABLE_TQDM),
-        ) as pbar:
+        ):
             with file.open("rb") as f:
-                # wrapped_file = CallbackIOWrapper(pbar.update, f, "read")
-
-                # if redirect:
-                #     callback_response = requestor.request_raw(
-                #         options=SeekrFlowRequest(
-                #             method="PUT",
-                #             url=redirect_url,
-                #             params=wrapped_file,
-                #             override_headers=True,
-                #         ),
-                #         absolute=True,
-                #         remaining_retries=MAX_RETRIES,
-                #     )
-                # else:
                 response, _, _ = requestor.request(
                     options=SeekrFlowRequest(
                         method="PUT",
                         url=url,
                         files={"files": f, "filename": file.name},
+                        params={"purpose": purpose.value},
                     ),
                 )
 
-        # if redirect:
-        #     assert isinstance(callback_response, requests.Response)
-        #
-        #     if not callback_response.status_code == 200:
-        #         raise APIError(
-        #             f"Error code: {callback_response.status_code} - Failed to process uploaded file"
-        #         )
-        #
-        #     response = self.callback(f"{url}/{file_id}/preprocess")
-
         assert isinstance(response, SeekrFlowResponse)
 
         return FileResponse(**response.data)

seekrai/resources/chat/completions.py

@@ -22,7 +22,7 @@ class ChatCompletions:
         *,
         messages: List[Dict[str, str]],
         model: str,
-        max_tokens: int | None = None,
+        max_tokens: int | None = 512,
         stop: List[str] | None = None,
         temperature: float = 0.7,
         top_p: float = 1,
@@ -113,7 +113,7 @@ class ChatCompletions:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url="chat/completions",
+                url="inference/chat/completions",
                 params=parameter_payload,
             ),
             stream=stream,
@@ -136,19 +136,19 @@ class AsyncChatCompletions:
         *,
         messages: List[Dict[str, str]],
         model: str,
-        max_tokens: int | None = None,
+        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,
+        temperature: float = 0.7,
+        top_p: float = 1,
+        top_k: int = 5,
+        repetition_penalty: float = 1,
         stream: bool = False,
-        logprobs: int | None = None,
-        echo: bool | None = None,
-        n: int | None = None,
+        logprobs: int = 0,
+        echo: bool = False,
+        n: int = 1,
         safety_model: str | None = None,
-        response_format: Dict[str, Any] | None = None,
-        # tools: Dict[str, str | Dict[str, str | Dict[str, Any]]] | None = None,
+        response_format: Dict[str, str | Dict[str, Any]] | None = None,
+        # tools: Dict[str, str | Dict[str, Any]] | None = None,
         # tool_choice: str | Dict[str, str | Dict[str, str]] | None = None,
     ) -> AsyncGenerator[ChatCompletionChunk, None] | ChatCompletionResponse:
         """
@@ -227,7 +227,7 @@ class AsyncChatCompletions:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url="chat/completions",
+                url="inference/chat/completions",
                 params=parameter_payload,
             ),
             stream=stream,

seekrai/resources/completions.py

@@ -71,7 +71,7 @@ class Completions:
             CompletionResponse | Iterator[CompletionChunk]: Object containing the completions
             or an iterator over completion chunks.
         """
-
+        raise NotImplementedError("function not yet implemented")
         requestor = api_requestor.APIRequestor(
             client=self._client,
         )
@@ -95,7 +95,7 @@ class Completions:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url="completions",
+                url="inference/completions",
                 params=parameter_payload,
             ),
             stream=stream,
@@ -167,7 +167,7 @@ class AsyncCompletions:
             AsyncGenerator[CompletionChunk, None] | CompletionResponse: Object containing the completions
             or an iterator over completion chunks.
         """
-
+        raise NotImplementedError("function not yet implemented")
         requestor = api_requestor.APIRequestor(
             client=self._client,
         )
@@ -191,7 +191,7 @@ class AsyncCompletions:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url="completions",
+                url="inference/completions",
                 params=parameter_payload,
             ),
             stream=stream,

seekrai/resources/embeddings.py

@@ -32,6 +32,7 @@ class Embeddings:
         Returns:
             EmbeddingResponse: Object containing embeddings
         """
+        raise NotImplementedError("Function not implemented yet")
 
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -45,7 +46,7 @@ class Embeddings:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url="embeddings",
+                url="inference/embeddings",
                 params=parameter_payload,
             ),
             stream=False,
@@ -76,6 +77,7 @@ class AsyncEmbeddings:
         Returns:
             EmbeddingResponse: Object containing embeddings
         """
+        raise NotImplementedError("Function not implemented yet")
 
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -89,7 +91,7 @@ class AsyncEmbeddings:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url="embeddings",
+                url="inference/embeddings",
                 params=parameter_payload,
             ),
             stream=False,

seekrai/resources/files.py

@@ -34,7 +34,7 @@ class Files:
 
         assert isinstance(purpose, FilePurpose)
 
-        return upload_manager.upload("files", file, purpose=purpose, redirect=True)
+        return upload_manager.upload("flow/files", file, purpose=purpose, redirect=True)
 
     def list(self) -> FileList:
         requestor = api_requestor.APIRequestor(
@@ -44,13 +44,21 @@ class Files:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url="files",
+                url="flow/files",
             ),
             stream=False,
         )
 
         assert isinstance(response, SeekrFlowResponse)
-        files = [FileResponse(id=file['id'], filename= file['filename'], created_at = file['created_at'], object="file") for file in response.data['data']]
+        files = [
+            FileResponse(
+                id=file["id"],
+                filename=file["filename"],
+                created_at=file["created_at"],
+                object="file",
+            )
+            for file in response.data["data"]
+        ]
         return FileList(object="list", data=files)
 
     def retrieve(self, id: str) -> FileResponse:
@@ -61,7 +69,7 @@ class Files:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"files/{id}",
+                url=f"flow/files/{id}",
             ),
             stream=False,
         )
@@ -79,7 +87,7 @@ class Files:
             output = Path(output)
 
         downloaded_filename, file_size = download_manager.download(
-            f"files/{id}/content", output, normalize_key(f"{id}.jsonl")
+            f"flow/files/{id}/content", output, normalize_key(f"{id}.jsonl")
         )
 
         return FileObject(
@@ -97,7 +105,7 @@ class Files:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="DELETE",
-                url=f"files/{id}",
+                url=f"flow/files/{id}",
             ),
             stream=False,
         )
@@ -124,7 +132,7 @@ class AsyncFiles:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url="files",
+                url="flow/files",
             ),
             stream=False,
         )
@@ -141,7 +149,7 @@ class AsyncFiles:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"files/{id}",
+                url=f"flow/files/{id}",
             ),
             stream=False,
         )
@@ -163,7 +171,7 @@ class AsyncFiles:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="DELETE",
-                url=f"files/{id}",
+                url=f"flow/files/{id}",
             ),
             stream=False,
         )