seekrai 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

seekrai/client.py

@@ -19,8 +19,11 @@ class SeekrFlow:
     models: resources.Models
     fine_tuning: resources.FineTuning
     alignment: resources.Alignment
+    ingestion: resources.Ingestion
     projects: resources.Projects
     deployments: resources.Deployments
+    vector_database: resources.VectorDatabase
+    agents: resources.Agents
 
     # client options
     client: SeekrFlowClient
@@ -81,8 +84,11 @@ class SeekrFlow:
         self.models = resources.Models(self.client)
         self.fine_tuning = resources.FineTuning(self.client)
         self.alignment = resources.Alignment(self.client)
+        self.ingestion = resources.Ingestion(self.client)
         self.projects = resources.Projects(self.client)
         self.deployments = resources.Deployments(self.client)
+        self.vector_database = resources.VectorDatabase(self.client)
+        self.agents = resources.Agents(self.client)
 
 
 class AsyncSeekrFlow:
@@ -94,8 +100,11 @@ class AsyncSeekrFlow:
     models: resources.AsyncModels
     fine_tuning: resources.AsyncFineTuning
     alignment: resources.AsyncAlignment
+    ingestion: resources.AsyncIngestion
     projects: resources.AsyncProjects
     deployments: resources.AsyncDeployments
+    vector_database: resources.AsyncVectorDatabase
+    agents: resources.AsyncAgents
 
     # client options
     client: SeekrFlowClient
@@ -156,8 +165,11 @@ class AsyncSeekrFlow:
         self.models = resources.AsyncModels(self.client)
         self.fine_tuning = resources.AsyncFineTuning(self.client)
         self.alignment = resources.AsyncAlignment(self.client)
+        self.ingestion = resources.AsyncIngestion(self.client)
         self.projects = resources.AsyncProjects(self.client)
         self.deployments = resources.AsyncDeployments(self.client)
+        self.vector_database = resources.AsyncVectorDatabase(self.client)
+        self.agents = resources.AsyncAgents(self.client)
 
 
 Client = SeekrFlow

seekrai/filemanager.py

@@ -91,9 +91,9 @@ def _prepare_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."
+    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}"
@@ -341,3 +341,181 @@ class UploadManager:
         assert isinstance(response, SeekrFlowResponse)
 
         return FileResponse(**response.data)
+
+    def bulk_upload(
+        self,
+        url: str,
+        files: list[Path],
+        *,
+        purpose: FilePurpose,
+        redirect: bool = False,
+    ) -> list[FileResponse]:
+        """
+        Upload multiple files in a bulk request.
+
+        Args:
+            url: API endpoint to upload to
+            files: List of file paths to upload
+            purpose: The purpose of the files
+            redirect: Whether to redirect after upload
+
+        Returns:
+            List of FileResponse objects for each uploaded file
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        # Prepare files for multipart form upload
+        # Format needs to be a list of tuples: [('files', (filename, file_stream)), ...]
+        files_tuples = []
+        file_streams = []
+        progress_bars = []
+
+        try:
+            for file_path in files:
+                file_size = os.stat(file_path.as_posix()).st_size
+
+                # Create progress bar
+                progress_bar = tqdm(
+                    total=file_size,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Uploading file {file_path.name}",
+                    disable=bool(DISABLE_TQDM),
+                )
+                progress_bars.append(progress_bar)
+
+                # Open file and track for cleanup
+                file_stream = file_path.open("rb")
+                file_streams.append(file_stream)
+
+                # Create wrapper for progress tracking
+                reader_wrapper = CallbackIOWrapper(
+                    progress_bar.update, file_stream, "read"
+                )
+
+                # Add to files list as a tuple
+                files_tuples.append(("files", (file_path.name, reader_wrapper)))
+
+            # Make the request
+            response, _, _ = requestor.request(
+                options=SeekrFlowRequest(
+                    method="PUT",
+                    url=url,
+                    files=files_tuples,  # Pass as a list of tuples (field_name, (filename, file_stream))
+                    params={"purpose": purpose.value},
+                    allow_redirects=redirect,
+                ),
+                stream=False,
+            )
+
+            assert isinstance(response, SeekrFlowResponse)
+
+            # Parse the response
+            if isinstance(response.data, list):
+                file_responses = [
+                    FileResponse(**file_data) for file_data in response.data
+                ]
+            else:
+                # Handle case where response might be a single object
+                file_responses = [FileResponse(**response.data)]
+
+            return file_responses
+
+        finally:
+            # Clean up resources
+            for file_stream in file_streams:
+                file_stream.close()
+
+            for progress_bar in progress_bars:
+                progress_bar.close()
+
+    async def bulk_upload_async(
+        self,
+        url: str,
+        files: list[Path],
+        *,
+        purpose: FilePurpose,
+        redirect: bool = False,
+    ) -> list[FileResponse]:
+        """
+        Upload multiple files in a bulk request.
+
+        Args:
+            url: API endpoint to upload to
+            files: List of file paths to upload
+            purpose: The purpose of the files
+            redirect: Whether to redirect after upload
+
+        Returns:
+            List of FileResponse objects for each uploaded file
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        # Prepare files for multipart form upload
+        # Format needs to be a list of tuples: [('files', (filename, file_stream)), ...]
+        files_tuples = []
+        file_streams = []
+        progress_bars = []
+
+        try:
+            for file_path in files:
+                file_size = os.stat(file_path.as_posix()).st_size
+
+                # Create progress bar
+                progress_bar = tqdm(
+                    total=file_size,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Uploading file {file_path.name}",
+                    disable=bool(DISABLE_TQDM),
+                )
+                progress_bars.append(progress_bar)
+
+                # Open file and track for cleanup
+                file_stream = file_path.open("rb")
+                file_streams.append(file_stream)
+
+                # Create wrapper for progress tracking
+                reader_wrapper = CallbackIOWrapper(
+                    progress_bar.update, file_stream, "read"
+                )
+
+                # Add to files list as a tuple
+                files_tuples.append(("files", (file_path.name, reader_wrapper)))
+
+            # Make the request
+            response, _, _ = await requestor.arequest(
+                options=SeekrFlowRequest(
+                    method="PUT",
+                    url=url,
+                    files=files_tuples,  # Pass as a list of tuples (field_name, (filename, file_stream))
+                    params={"purpose": purpose.value},
+                    allow_redirects=redirect,
+                ),
+                stream=False,
+            )
+
+            assert isinstance(response, SeekrFlowResponse)
+
+            # Parse the response
+            if isinstance(response.data, list):
+                file_responses = [
+                    FileResponse(**file_data) for file_data in response.data
+                ]
+            else:
+                # Handle case where response might be a single object
+                file_responses = [FileResponse(**response.data)]
+
+            return file_responses
+
+        finally:
+            # Clean up resources
+            for file_stream in file_streams:
+                file_stream.close()
+
+            for progress_bar in progress_bars:
+                progress_bar.close()

seekrai/resources/__init__.py

@@ -1,3 +1,4 @@
+from seekrai.resources.agents import AgentInference, Agents, AsyncAgents
 from seekrai.resources.alignment import Alignment, AsyncAlignment
 from seekrai.resources.chat import AsyncChat, Chat
 from seekrai.resources.completions import AsyncCompletions, Completions
@@ -6,8 +7,10 @@ from seekrai.resources.embeddings import AsyncEmbeddings, Embeddings
 from seekrai.resources.files import AsyncFiles, Files
 from seekrai.resources.finetune import AsyncFineTuning, FineTuning
 from seekrai.resources.images import AsyncImages, Images
+from seekrai.resources.ingestion import AsyncIngestion, Ingestion
 from seekrai.resources.models import AsyncModels, Models
 from seekrai.resources.projects import AsyncProjects, Projects
+from seekrai.resources.vectordb import AsyncVectorDatabase, VectorDatabase
 
 
 __all__ = [
@@ -25,10 +28,17 @@ __all__ = [
     "Files",
     "AsyncImages",
     "Images",
+    "Ingestion",
+    "AsyncIngestion",
     "AsyncModels",
     "Models",
     "AsyncProjects",
     "Projects",
     "AsyncDeployments",
     "Deployments",
+    "AsyncAgents",
+    "Agents",
+    "VectorDatabase",
+    "AsyncVectorDatabase",
+    "AgentInference",
 ]

seekrai/resources/agents/__init__.py

@@ -0,0 +1,13 @@
+from seekrai.resources.agents.agent_inference import AgentInference, AsyncAgentInference
+from seekrai.resources.agents.agents import Agents, AsyncAgents
+from seekrai.resources.agents.threads import AgentThreads, AsyncAgentThreads
+
+
+__all__ = [
+    "Agents",
+    "AgentInference",
+    "AsyncAgentInference",
+    "AsyncAgents",
+    "AgentThreads",
+    "AsyncAgentThreads",
+]

seekrai/resources/agents/agent_inference.py

@@ -0,0 +1,277 @@
+from typing import Any, AsyncGenerator, Iterator, Union
+
+from seekrai.abstract import api_requestor
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import Run, RunRequest, RunResponse, SeekrFlowRequest
+
+
+class AgentInference:
+    def __init__(self, client: Any) -> None:
+        self._client = client
+        self._requestor = api_requestor.APIRequestor(client=self._client)
+
+    def run(
+        self,
+        agent_id: str,
+        thread_id: str,
+        *,
+        stream: bool = False,
+        **model_settings: Any,
+    ) -> Union[RunResponse, Iterator[Any]]:
+        """
+        Run an inference call on a deployed agent.
+
+        Args:
+            agent_id (str): The unique identifier of the deployed agent.
+            thread_id (str): A thread identifier.
+            stream (bool, optional): Whether to stream the response. Defaults to False.
+            **model_settings: Additional parameters (such as temperature, max_tokens, etc).
+
+        Returns:
+            A dictionary with the response (if non-streaming) or an iterator over response chunks.
+        """
+        payload = RunRequest(agent_id=agent_id).model_dump()
+        payload.update(model_settings)
+        endpoint = f"threads/{thread_id}/runs"
+        if stream:
+            endpoint += "/stream"
+
+        response, _, _ = self._requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=endpoint,
+                params=payload,
+            ),
+            stream=stream,
+        )
+
+        if stream:
+            assert not isinstance(response, SeekrFlowResponse)
+            return (chunk.data for chunk in response)
+        else:
+            assert isinstance(response, SeekrFlowResponse)
+            return RunResponse(**response.data)
+
+    def cancel(self, agent_id: str, run_id: str, thread_id: str) -> dict[str, Any]:
+        """Cancels a Run that is in progress.
+
+        Args:
+            agent_id: Identifier for the agent performing the run.
+            run_id: Identifier for the run to be cancelled.
+            thread_id: Identifier for the thread used by the run.
+
+        Returns:
+            {'status': 'canceled', 'run_id': run_id} on success.
+        """
+        response, _, _ = self._requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"threads/{thread_id}/runs/{run_id}/cancel",
+                params=RunRequest(agent_id=agent_id).model_dump(),
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return response.data
+
+    def attach(self, run_id: str, thread_id: str) -> Iterator[Any]:
+        """Returns a stream of output from a Run.
+
+        Args:
+            run_id: Identifier for the Run.
+            thread_id: Identifier for the Thread used by the Run.
+
+        Returns:
+            An Iterator of streamed output from the Run.
+        """
+        response, _, _ = self._requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs/{run_id}/attach",
+            ),
+            stream=True,
+        )
+
+        assert not isinstance(response, SeekrFlowResponse)
+        return (chunk.data for chunk in response)
+
+    def retrieve(self, run_id: str, thread_id: str) -> Run:
+        """Retrieves a Run.
+
+        Args:
+            run_id: Identifier for the Run.
+            thread_id: Identifier for the Thread used by the Run.
+
+        Returns:
+            The Run whose id matches run_id.
+        """
+        response, _, _ = self._requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs/{run_id}",
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return Run(**response.data)
+
+    def list(self, thread_id: str) -> list[Run]:
+        """Retrieves a list of Runs relevant to a referenced Thread.
+
+        Args:
+            thread_id: Identifier for a Thread.
+
+        Returns:
+            A list of Runs that have leveraged the referenced Thread.
+        """
+        response, _, _ = self._requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs",
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return [Run(**run) for run in response.data]  # type: ignore
+
+
+class AsyncAgentInference:
+    def __init__(self, client: Any) -> None:
+        self._client = client
+        self._requestor = api_requestor.APIRequestor(client=self._client)
+
+    async def run(
+        self,
+        agent_id: str,
+        thread_id: str,
+        *,
+        stream: bool = False,
+        **model_settings: Any,
+    ) -> Union[RunResponse, AsyncGenerator[Any, None]]:
+        """
+        Run an inference call on a deployed agent.
+
+        Args:
+            agent_id (str): The unique identifier of the deployed agent.
+            thread_id (str): A thread identifier.
+            stream (bool, optional): Whether to stream the response. Defaults to False.
+            **model_settings: Additional parameters (such as temperature, max_tokens, etc).
+
+        Returns:
+            A dictionary with the response (if non-streaming) or an iterator over response chunks.
+        """
+        payload = RunRequest(agent_id=agent_id).model_dump()
+        payload.update(model_settings)
+        endpoint = f"threads/{thread_id}/runs"
+        if stream:
+            endpoint += "/stream"
+
+        response, _, _ = await self._requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=endpoint,
+                params=payload,
+            ),
+            stream=stream,
+        )
+
+        if stream:
+            assert not isinstance(response, SeekrFlowResponse)
+
+            async def output() -> AsyncGenerator[Any, None]:
+                async for chunk in response:
+                    yield chunk.data
+
+            return output()
+        else:
+            assert isinstance(response, SeekrFlowResponse)
+            return RunResponse(**response.data)
+
+    async def cancel(
+        self, agent_id: str, run_id: str, thread_id: str
+    ) -> dict[str, Any]:
+        """Cancels a Run that is in progress.
+
+        Args:
+            agent_id: Identifier for the agent performing the run.
+            run_id: Identifier for the run to be cancelled.
+            thread_id: Identifier for the thread used by the run.
+
+        Returns:
+            {'status': 'canceled', 'run_id': run_id} on success.
+        """
+        response, _, _ = await self._requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"threads/{thread_id}/runs/{run_id}/cancel",
+                params=RunRequest(agent_id=agent_id).model_dump(),
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return response.data
+
+    async def attach(self, run_id: str, thread_id: str) -> AsyncGenerator[Any, None]:
+        """Returns a stream of output from a Run.
+
+        Args:
+            run_id: Identifier for the Run.
+            thread_id: Identifier for the Thread used by the Run.
+
+        Returns:
+            An Iterator of streamed output from the Run.
+        """
+        response, _, _ = await self._requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs/{run_id}/attach",
+            ),
+            stream=True,
+        )
+
+        assert not isinstance(response, SeekrFlowResponse)
+
+        async def output() -> AsyncGenerator[Any, None]:
+            async for chunk in response:
+                yield chunk.data
+
+        return output()
+
+    async def retrieve(self, run_id: str, thread_id: str) -> Run:
+        """Retrieves a Run.
+
+        Args:
+            run_id: Identifier for the Run.
+            thread_id: Identifier for the Thread used by the Run.
+
+        Returns:
+            The Run whose id matches run_id.
+        """
+        response, _, _ = await self._requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs/{run_id}",
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return Run(**response.data)
+
+    async def list(self, thread_id: str) -> list[Run]:
+        """Retrieves a list of Runs relevant to a referenced Thread.
+
+        Args:
+            thread_id: Identifier for a Thread.
+
+        Returns:
+            A list of Runs that have leveraged the referenced Thread.
+        """
+        response, _, _ = await self._requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"threads/{thread_id}/runs",
+            )
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return [Run(**run) for run in response.data]  # type: ignore