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

seekrai/client.py

@@ -22,6 +22,7 @@ class SeekrFlow:
     ingestion: resources.Ingestion
     projects: resources.Projects
     deployments: resources.Deployments
+    vector_database: resources.VectorDatabase
     agents: resources.Agents
 
     # client options
@@ -86,6 +87,7 @@ class SeekrFlow:
         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)
 
 
@@ -101,7 +103,8 @@ class AsyncSeekrFlow:
     ingestion: resources.AsyncIngestion
     projects: resources.AsyncProjects
     deployments: resources.AsyncDeployments
-    agents: resources.Agents
+    vector_database: resources.AsyncVectorDatabase
+    agents: resources.AsyncAgents
 
     # client options
     client: SeekrFlowClient
@@ -165,7 +168,8 @@ class AsyncSeekrFlow:
         self.ingestion = resources.AsyncIngestion(self.client)
         self.projects = resources.AsyncProjects(self.client)
         self.deployments = resources.AsyncDeployments(self.client)
-        self.agents = resources.Agents(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}"
@@ -430,3 +430,92 @@ class UploadManager:
 
             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,4 +1,4 @@
-from seekrai.resources.agents import Agents
+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
@@ -10,6 +10,7 @@ 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__ = [
@@ -35,5 +36,9 @@ __all__ = [
     "Projects",
     "AsyncDeployments",
     "Deployments",
+    "AsyncAgents",
     "Agents",
+    "VectorDatabase",
+    "AsyncVectorDatabase",
+    "AgentInference",
 ]

seekrai/resources/agents/__init__.py

@@ -1,8 +1,13 @@
-from typing import Any
+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
 
-from .agent_inference import AgentInference
 
-
-class Agents:
-    def __init__(self, client: Any) -> None:
-        self.agent = AgentInference(client)
+__all__ = [
+    "Agents",
+    "AgentInference",
+    "AsyncAgentInference",
+    "AsyncAgents",
+    "AgentThreads",
+    "AsyncAgentThreads",
+]

seekrai/resources/agents/agent_inference.py

@@ -1,70 +1,277 @@
-from typing import Any, Dict, Iterator, Union
+from typing import Any, AsyncGenerator, Iterator, Union
 
 from seekrai.abstract import api_requestor
 from seekrai.seekrflow_response import SeekrFlowResponse
-from seekrai.types import SeekrFlowRequest
+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: Union[int, str],
-        query: str,
+        agent_id: str,
+        thread_id: str,
         *,
         stream: bool = False,
-        thread_id: Union[str, None] = None,
-        headers: Union[Dict[str, str], None] = None,
         **model_settings: Any,
-    ) -> Union[Dict[str, Any], Iterator[Dict[str, Any]]]:
+    ) -> Union[RunResponse, Iterator[Any]]:
         """
         Run an inference call on a deployed agent.
 
         Args:
-            agent_id (Union[int, str]): The unique identifier of the deployed agent.
-            query (str): The user query or prompt.
+            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.
-            thread_id (str, optional): An optional thread identifier.
-            headers (dict, optional): Optional HTTP headers to include in the request.
-                                    If provided, these will be merged with default headers.
             **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: Dict[str, Any] = {"prompt": query}
-        if thread_id is not None:
-            payload["thread_id"] = thread_id
+        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)
 
-        default_headers = {}
-        if self._client.api_key:
-            default_headers["x-api-key"] = self._client.api_key
-        if self._client.supplied_headers:
-            default_headers.update(self._client.supplied_headers)
+    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.
 
-        request_headers = default_headers.copy()
-        if headers:
-            request_headers.update(headers)
+        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).
 
-        requestor = api_requestor.APIRequestor(client=self._client)
-        endpoint = f"agents/{agent_id}/run"
+        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, _, _ = requestor.request(
+        response, _, _ = await self._requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
                 url=endpoint,
                 params=payload,
-                headers=request_headers,
             ),
             stream=stream,
         )
 
         if stream:
             assert not isinstance(response, SeekrFlowResponse)
-            return (chunk.data for chunk in response)
+
+            async def output() -> AsyncGenerator[Any, None]:
+                async for chunk in response:
+                    yield chunk.data
+
+            return output()
         else:
             assert isinstance(response, SeekrFlowResponse)
-            return response.data
+            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