seekrai 0.0.1__py3-none-any.whl

seekrai/filemanager.py

@@ -0,0 +1,393 @@
+from __future__ import annotations
+
+import os
+import shutil
+import stat
+import tempfile
+import uuid
+from functools import partial
+from pathlib import Path
+from typing import Tuple
+
+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.error import (
+    APIError,
+    AuthenticationError,
+    DownloadError,
+    FileTypeError,
+)
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    FilePurpose,
+    FileResponse,
+    FileType,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+
+
+def chmod_and_replace(src: Path, dst: Path) -> None:
+    """Set correct permission before moving a blob from tmp directory to cache dir.
+
+    Do not take into account the `umask` from the process as there is no convenient way
+    to get it that is thread-safe.
+    """
+
+    # Get umask by creating a temporary file in the cache folder.
+    tmp_file = dst.parent / f"tmp_{uuid.uuid4()}"
+
+    try:
+        tmp_file.touch()
+
+        cache_dir_mode = Path(tmp_file).stat().st_mode
+
+        os.chmod(src.as_posix(), stat.S_IMODE(cache_dir_mode))
+
+    finally:
+        tmp_file.unlink()
+
+    shutil.move(src.as_posix(), dst.as_posix())
+
+
+def _get_file_size(
+    headers: CaseInsensitiveDict[str],
+) -> int:
+    """
+    Extracts file size from header
+    """
+    total_size_in_bytes = 0
+
+    parts = headers.get("Content-Range", "").split(" ")
+
+    if len(parts) == 2:
+        range_parts = parts[1].split("/")
+
+        if len(range_parts) == 2:
+            total_size_in_bytes = int(range_parts[1])
+
+    assert total_size_in_bytes != 0, "Unable to retrieve remote file."
+
+    return total_size_in_bytes
+
+
+def _prepare_output(
+    headers: CaseInsensitiveDict[str],
+    step: int = -1,
+    output: Path | None = None,
+    remote_name: str | None = None,
+) -> Path:
+    """
+    Generates output file name from remote name and headers
+    """
+    if output:
+        return output
+
+    content_type = str(headers.get("content-type"))
+
+    assert remote_name, (
+        "No model name found in fine_tune object. "
+        "Please specify an `output` file name."
+    )
+
+    if step > 0:
+        remote_name += f"-checkpoint-{step}"
+
+    if "x-tar" in content_type.lower():
+        remote_name += ".tar.gz"
+
+    elif "zstd" in content_type.lower() or step != -1:
+        remote_name += ".tar.zst"
+
+    else:
+        raise FileTypeError(
+            f"Unknown file type {content_type} found. Aborting download."
+        )
+
+    return Path(remote_name)
+
+
+class DownloadManager:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def get_file_metadata(
+        self,
+        url: str,
+        output: Path | None = None,
+        remote_name: str | None = None,
+        fetch_metadata: bool = False,
+    ) -> Tuple[Path, int]:
+        """
+        gets remote file head and parses out file name and file size
+        """
+
+        if not fetch_metadata:
+            if isinstance(output, Path):
+                file_path = output
+            else:
+                assert isinstance(remote_name, str)
+                file_path = Path(remote_name)
+
+            return file_path, 0
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response = requestor.request_raw(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=url,
+                headers={"Range": "bytes=0-1"},
+            ),
+            remaining_retries=MAX_RETRIES,
+            stream=False,
+        )
+
+        try:
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            raise APIError(
+                "Error fetching file metadata", http_status=response.status_code
+            ) from e
+
+        headers = response.headers
+
+        assert isinstance(headers, CaseInsensitiveDict)
+
+        file_path = _prepare_output(
+            headers=headers,
+            output=output,
+            remote_name=remote_name,
+        )
+
+        file_size = _get_file_size(headers)
+
+        return file_path, file_size
+
+    def download(
+        self,
+        url: str,
+        output: Path | None = None,
+        remote_name: str | None = None,
+        fetch_metadata: bool = False,
+    ) -> Tuple[str, int]:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        # pre-fetch remote file name and file size
+        file_path, file_size = self.get_file_metadata(
+            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,
+                )
+
+                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)
+
+            # Raise exception if remote file size does not match downloaded file size
+            if os.stat(temp_file.name).st_size != file_size:
+                DownloadError(
+                    f"Downloaded file size `{pbar.n}` bytes does not match "
+                    f"remote file size `{file_size}` bytes."
+                )
+
+            # Moves temp file to output file path
+            chmod_and_replace(Path(temp_file.name), file_path)
+
+        os.remove(lock_path)
+
+        return str(file_path.resolve()), file_size
+
+
+class UploadManager:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    @classmethod
+    def _redirect_error_handler(
+        cls, requestor: api_requestor.APIRequestor, response: requests.Response
+    ) -> None:
+        if response.status_code == 401:
+            raise AuthenticationError(
+                "This job would exceed your free trial credits. "
+                "Please upgrade to a paid account through "
+                "Settings -> Billing on api.seekrai.ai to continue.",
+            )
+        elif response.status_code != 302:
+            raise APIError(
+                f"Unexpected error raised by endpoint: {response.content.decode()}, headers: {response.headers}",
+                http_status=response.status_code,
+            )
+
+    def get_upload_url(
+        self,
+        url: str,
+        file: Path,
+        purpose: FilePurpose,
+        filetype: FileType,
+    ) -> Tuple[str, str]:
+        data = {
+            "purpose": purpose.value,
+            "file_name": file.name,
+            "file_type": filetype.value,
+        }
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        method = "POST"
+
+        headers = seekrai.utils.get_headers(method, requestor.api_key)
+
+        response = requestor.request_raw(
+            options=SeekrFlowRequest(
+                method=method,
+                url=url,
+                params=data,
+                allow_redirects=False,
+                override_headers=True,
+                headers=headers,
+            ),
+            remaining_retries=MAX_RETRIES,
+        )
+
+        self._redirect_error_handler(requestor, response)
+
+        redirect_url = response.headers["Location"]
+        file_id = response.headers["X-SeekrFlow-File-Id"]
+
+        return redirect_url, file_id
+
+    def callback(self, url: str) -> SeekrFlowResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=url,
+            ),
+        )
+
+        return response
+
+    def upload(
+        self,
+        url: str,
+        file: Path,
+        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:
+                raise FileTypeError(
+                    f"Unknown extension of file {file}. "
+                    "Only files with extensions .jsonl and .parquet are supported."
+                )
+            # redirect_url, file_id = self.get_upload_url(url, file, purpose, filetype)
+
+        file_size = os.stat(file.as_posix()).st_size
+
+        with tqdm(
+            total=file_size,
+            unit="B",
+            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},
+                    ),
+                )
+
+        # 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/legacy/base.py

@@ -0,0 +1,27 @@
+import functools
+import warnings
+
+
+API_KEY_WARNING = (
+    "The use of seekrai.api_key is deprecated and will be removed in the next major release. "
+    "Please set the SEEKRFLOW_API_KEY environment variable instead."
+)
+
+
+def deprecated(func):  # type: ignore
+    """
+    This is a decorator which can be used to mark functions
+    as deprecated. It will result in a warning being emitted
+    when the function is used.
+    """
+
+    @functools.wraps(func)
+    def new_func(*args, **kwargs):  # type: ignore
+        warnings.warn(
+            f"Call to deprecated function {func.__name__}.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        return func(*args, **kwargs)
+
+    return new_func

seekrai/legacy/complete.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import warnings
+from typing import Any, AsyncGenerator, Dict, Iterator
+
+import seekrai
+from seekrai.legacy.base import API_KEY_WARNING, deprecated
+from seekrai.types import CompletionChunk, CompletionResponse
+
+
+class Complete:
+    @classmethod
+    @deprecated  # type: ignore
+    def create(
+        cls,
+        prompt: str,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Legacy completion function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.completions.create(
+            prompt=prompt, stream=False, **kwargs
+        ).model_dump()  # type: ignore
+
+    @classmethod
+    @deprecated  # type: ignore
+    def create_streaming(
+        cls,
+        prompt: str,
+        **kwargs,
+    ) -> Iterator[Dict[str, Any]]:
+        """Legacy streaming completion function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return (
+            token.model_dump()  # type: ignore
+            for token in client.completions.create(prompt=prompt, stream=True, **kwargs)
+        )
+
+
+class Completion:
+    @classmethod
+    @deprecated  # type: ignore
+    def create(
+        cls,
+        prompt: str,
+        **kwargs,
+    ) -> CompletionResponse | Iterator[CompletionChunk]:
+        """Completion function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.completions.create(prompt=prompt, **kwargs)
+
+
+class AsyncComplete:
+    @classmethod
+    @deprecated  # type: ignore
+    async def create(
+        cls,
+        prompt: str,
+        **kwargs,
+    ) -> CompletionResponse | AsyncGenerator[CompletionChunk, None]:
+        """Async completion function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.AsyncSeekrFlow(api_key=api_key)
+
+        return await client.completions.create(prompt=prompt, **kwargs)

seekrai/legacy/embeddings.py

@@ -0,0 +1,25 @@
+import warnings
+from typing import Any, Dict
+
+import seekrai
+from seekrai.legacy.base import API_KEY_WARNING, deprecated
+
+
+class Embeddings:
+    @classmethod
+    @deprecated  # type: ignore
+    def create(
+        cls,
+        input: str,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Legacy embeddings function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.embeddings.create(input=input, **kwargs).model_dump()

seekrai/legacy/files.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import json
+import warnings
+from typing import Any, Dict, List
+
+import seekrai
+from seekrai.legacy.base import API_KEY_WARNING, deprecated
+from seekrai.utils.files import check_file as check_json
+
+
+class Files:
+    @classmethod
+    @deprecated  # type: ignore
+    def list(
+        cls,
+    ) -> Dict[str, Any]:
+        """Legacy file list function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.files.list().model_dump()
+
+    @classmethod
+    def check(self, file: str) -> Dict[str, object]:
+        return check_json(file)
+
+    @classmethod
+    @deprecated  # type: ignore
+    def upload(
+        cls,
+        file: str,
+        check: bool = True,
+    ) -> Dict[str, Any]:
+        """Legacy file upload function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        if check:
+            report_dict = check_json(file)
+            if not report_dict["is_check_passed"]:
+                raise seekrai.error.FileTypeError(
+                    f"Invalid file supplied. Failed to upload.\nReport:\n {report_dict}"
+                )
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        response = client.files.upload(file=file).model_dump()
+
+        response["report_dict"] = report_dict
+
+        return response
+
+    @classmethod
+    @deprecated  # type: ignore
+    def delete(
+        cls,
+        file_id: str,
+    ) -> Dict[str, Any]:
+        """Legacy file delete function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.files.delete(id=file_id).model_dump()
+
+    @classmethod
+    @deprecated  # type: ignore
+    def retrieve(
+        cls,
+        file_id: str,
+    ) -> Dict[str, Any]:
+        """Legacy file retrieve function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.files.retrieve(id=file_id).model_dump()
+
+    @classmethod
+    @deprecated  # type: ignore
+    def retrieve_content(
+        cls,
+        file_id: str,
+        output: str | None = None,
+    ) -> Dict[str, Any]:
+        """Legacy file retrieve content function."""
+
+        api_key = None
+        if seekrai.api_key:
+            warnings.warn(API_KEY_WARNING)
+            api_key = seekrai.api_key
+
+        client = seekrai.SeekrFlow(api_key=api_key)
+
+        return client.files.retrieve_content(id=file_id, output=output).model_dump()
+
+    @classmethod
+    @deprecated  # type: ignore
+    def save_jsonl(
+        self, data: Dict[str, str], output_path: str, append: bool = False
+    ) -> None:
+        """
+        Write list of objects to a JSON lines file.
+        """
+        mode = "a+" if append else "w"
+        with open(output_path, mode, encoding="utf-8") as f:
+            for line in data:
+                json_record = json.dumps(line, ensure_ascii=False)
+                f.write(json_record + "\n")
+        print("Wrote {} records to {}".format(len(data), output_path))
+
+    @classmethod
+    @deprecated  # type: ignore
+    def load_jsonl(self, input_path: str) -> List[Dict[str, str]]:
+        """
+        Read list of objects from a JSON lines file.
+        """
+        data = []
+        with open(input_path, "r", encoding="utf-8") as f:
+            for line in f:
+                data.append(json.loads(line.rstrip("\n|\r")))
+        print("Loaded {} records from {}".format(len(data), input_path))
+        return data