seekrai 0.0.1__py3-none-any.whl

seekrai/resources/completions.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import AsyncGenerator, Iterator, List
+
+from seekrai.abstract import api_requestor
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    CompletionChunk,
+    CompletionRequest,
+    CompletionResponse,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+
+
+class Completions:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def create(
+        self,
+        *,
+        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,
+        stream: bool = False,
+        logprobs: int | None = None,
+        echo: bool | None = None,
+        n: int | None = None,
+        safety_model: str | None = None,
+    ) -> CompletionResponse | Iterator[CompletionChunk]:
+        """
+        Method to generate completions based on a given prompt using a specified model.
+
+        Args:
+            prompt (str): A string providing context for the model to complete.
+            model (str): The name of the model to query.
+            max_tokens (int, optional): The maximum number of tokens to generate.
+                Defaults to 512.
+            stop (List[str], optional): List of strings at which to stop generation.
+                Defaults to None.
+            temperature (float, optional): A decimal number that determines the degree of randomness in the response.
+                Defaults to None.
+            top_p (float, optional): The top_p (nucleus) parameter is used to dynamically adjust the number
+                    of choices for each predicted token based on the cumulative probabilities.
+                Defaults to None.
+            top_k (int, optional): The top_k parameter is used to limit the number of choices for the
+                    next predicted word or token.
+                Defaults to None.
+            repetition_penalty (float, optional): A number that controls the diversity of generated text
+                    by reducing the likelihood of repeated sequences. Higher values decrease repetition.
+                Defaults to None.
+            stream (bool, optional): Flag indicating whether to stream the generated completions.
+                Defaults to False.
+            logprobs (int, optional): Number of top-k logprobs to return
+                Defaults to None.
+            echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
+                Defaults to None.
+            n (int, optional): Number of completions to generate. Setting to None will return a single generation.
+                Defaults to None.
+            safety_model (str, optional): A moderation model to validate tokens. Choice between available moderation
+                    models found [here](https://docs.seekrflow.ai/docs/inference-models#moderation-models).
+                Defaults to None.
+
+        Returns:
+            CompletionResponse | Iterator[CompletionChunk]: Object containing the completions
+            or an iterator over completion chunks.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = CompletionRequest(
+            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=stream,
+            logprobs=logprobs,
+            echo=echo,
+            n=n,
+            safety_model=safety_model,
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="completions",
+                params=parameter_payload,
+            ),
+            stream=stream,
+        )
+
+        if stream:
+            # must be an iterator
+            assert not isinstance(response, SeekrFlowResponse)
+            return (CompletionChunk(**line.data) for line in response)
+        assert isinstance(response, SeekrFlowResponse)
+        return CompletionResponse(**response.data)
+
+
+class AsyncCompletions:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def create(
+        self,
+        *,
+        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,
+        stream: bool = False,
+        logprobs: int | None = None,
+        echo: bool | None = None,
+        n: int | None = None,
+        safety_model: str | None = None,
+    ) -> AsyncGenerator[CompletionChunk, None] | CompletionResponse:
+        """
+        Async method to generate completions based on a given prompt using a specified model.
+
+        Args:
+            prompt (str): A string providing context for the model to complete.
+            model (str): The name of the model to query.
+            max_tokens (int, optional): The maximum number of tokens to generate.
+                Defaults to 512.
+            stop (List[str], optional): List of strings at which to stop generation.
+                Defaults to None.
+            temperature (float, optional): A decimal number that determines the degree of randomness in the response.
+                Defaults to None.
+            top_p (float, optional): The top_p (nucleus) parameter is used to dynamically adjust the number
+                    of choices for each predicted token based on the cumulative probabilities.
+                Defaults to None.
+            top_k (int, optional): The top_k parameter is used to limit the number of choices for the
+                    next predicted word or token.
+                Defaults to None.
+            repetition_penalty (float, optional): A number that controls the diversity of generated text
+                    by reducing the likelihood of repeated sequences. Higher values decrease repetition.
+                Defaults to None.
+            stream (bool, optional): Flag indicating whether to stream the generated completions.
+                Defaults to False.
+            logprobs (int, optional): Number of top-k logprobs to return
+                Defaults to None.
+            echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
+                Defaults to None.
+            n (int, optional): Number of completions to generate. Setting to None will return a single generation.
+                Defaults to None.
+            safety_model (str, optional): A moderation model to validate tokens. Choice between available moderation
+                    models found [here](https://docs.seekrflow.ai/docs/inference-models#moderation-models).
+                Defaults to None.
+
+        Returns:
+            AsyncGenerator[CompletionChunk, None] | CompletionResponse: Object containing the completions
+            or an iterator over completion chunks.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = CompletionRequest(
+            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=stream,
+            logprobs=logprobs,
+            echo=echo,
+            n=n,
+            safety_model=safety_model,
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="completions",
+                params=parameter_payload,
+            ),
+            stream=stream,
+        )
+
+        if stream:
+            # must be an iterator
+            assert not isinstance(response, SeekrFlowResponse)
+            return (CompletionChunk(**line.data) async for line in response)
+        assert isinstance(response, SeekrFlowResponse)
+        return CompletionResponse(**response.data)

seekrai/resources/embeddings.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import List
+
+from seekrai.abstract import api_requestor
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    EmbeddingRequest,
+    EmbeddingResponse,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+
+
+class Embeddings:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def create(
+        self,
+        *,
+        input: str | List[str],
+        model: str,
+    ) -> EmbeddingResponse:
+        """
+        Method to generate completions based on a given prompt using a specified model.
+
+        Args:
+            input (str | List[str]): A string or list of strings to embed
+            model (str): The name of the model to query.
+
+        Returns:
+            EmbeddingResponse: Object containing embeddings
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = EmbeddingRequest(
+            input=input,
+            model=model,
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="embeddings",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return EmbeddingResponse(**response.data)
+
+
+class AsyncEmbeddings:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def create(
+        self,
+        *,
+        input: str | List[str],
+        model: str,
+    ) -> EmbeddingResponse:
+        """
+        Async method to generate completions based on a given prompt using a specified model.
+
+        Args:
+            input (str | List[str]): A string or list of strings to embed
+            model (str): The name of the model to query.
+
+        Returns:
+            EmbeddingResponse: Object containing embeddings
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = EmbeddingRequest(
+            input=input,
+            model=model,
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="embeddings",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return EmbeddingResponse(**response.data)

seekrai/resources/files.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from seekrai.abstract import api_requestor
+from seekrai.filemanager import DownloadManager, UploadManager
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    FileDeleteResponse,
+    FileList,
+    FileObject,
+    FilePurpose,
+    FileResponse,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+from seekrai.utils import normalize_key
+
+
+class Files:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def upload(
+        self, file: Path | str, *, purpose: FilePurpose | str = FilePurpose.FineTune
+    ) -> FileResponse:
+        upload_manager = UploadManager(self._client)
+
+        if isinstance(file, str):
+            file = Path(file)
+
+        if isinstance(purpose, str):
+            purpose = FilePurpose(purpose)
+
+        assert isinstance(purpose, FilePurpose)
+
+        return upload_manager.upload("files", file, purpose=purpose, redirect=True)
+
+    def list(self) -> FileList:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="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']]
+        return FileList(object="list", data=files)
+
+    def retrieve(self, id: str) -> FileResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"files/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FileResponse(**response.data)
+
+    def retrieve_content(
+        self, id: str, *, output: Path | str | None = None
+    ) -> FileObject:
+        download_manager = DownloadManager(self._client)
+
+        if isinstance(output, str):
+            output = Path(output)
+
+        downloaded_filename, file_size = download_manager.download(
+            f"files/{id}/content", output, normalize_key(f"{id}.jsonl")
+        )
+
+        return FileObject(
+            object="local",
+            id=id,
+            filename=downloaded_filename,
+            size=file_size,
+        )
+
+    def delete(self, id: str) -> FileDeleteResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"files/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FileDeleteResponse(**response.data)
+
+
+class AsyncFiles:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def upload(
+        self, file: Path | str, *, purpose: FilePurpose | str = FilePurpose.FineTune
+    ) -> None:
+        raise NotImplementedError()
+
+    async def list(self) -> FileList:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="files",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FileList(**response.data)
+
+    async def retrieve(self, id: str) -> FileResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"files/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FileResponse(**response.data)
+
+    async def retrieve_content(
+        self, id: str, *, output: Path | str | None = None
+    ) -> FileObject:
+        raise NotImplementedError()
+
+    async def delete(self, id: str) -> FileDeleteResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"files/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FileDeleteResponse(**response.data)