seekrai 0.5.2__py3-none-any.whl → 0.5.24__py3-none-any.whl

seekrai/resources/alignment.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List
 
 from seekrai.abstract import api_requestor
 from seekrai.resources.resource_base import ResourceBase
@@ -7,14 +7,23 @@ from seekrai.types import (
     AlignmentEstimationRequest,
     AlignmentEstimationResponse,
     AlignmentList,
+    AlignmentOutput,
     AlignmentRequest,
     AlignmentResponse,
     AlignmentType,
     SeekrFlowRequest,
+    SystemPrompt,
+    SystemPromptCreateRequest,
+    SystemPromptUpdateRequest,
 )
+from seekrai.types.abstract import SeekrFlowClient
 
 
 class Alignment(ResourceBase):
+    def __init__(self, client: SeekrFlowClient) -> None:
+        super().__init__(client)
+        self.system_prompt = SystemPromptResource(client)
+
     def generate(
         self,
         instructions: str,
@@ -92,6 +101,86 @@ class Alignment(ResourceBase):
 
         return AlignmentResponse(**response.data)
 
+    def outputs(self, id: str) -> List[AlignmentOutput]:
+        """
+        Retrieves output files for an alignment job.
+
+        Args:
+            id (str): Alignment job ID whose outputs to fetch.
+
+        Returns:
+            list[AlignmentOutput]: Collection of alignment output metadata.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/alignment/{id}/outputs",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return [AlignmentOutput(**output) for output in response.data]  # type: ignore[arg-type]
+
+    def delete(self, id: str) -> None:
+        """
+        Deletes an alignment job.
+
+        Args:
+            id (str): Alignment job ID to delete.
+
+        Returns:
+            None
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"flow/alignment/{id}",
+            ),
+            stream=False,
+        )
+
+        # Endpoint returns 204 No Content
+        return None
+
+    def cancel(self, id: str) -> AlignmentResponse:
+        """
+        Method to cancel a running alignment job
+
+        Args:
+            id (str): Alignment job ID to cancel. A string that starts with `al-`.
+
+        Returns:
+            AlignmentResponse: Object containing information about cancelled alignment job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/{id}/cancel",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return AlignmentResponse(**response.data)
+
     def estimate(self, files: List[str]) -> AlignmentEstimationResponse:
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -115,6 +204,10 @@ class Alignment(ResourceBase):
 
 
 class AsyncAlignment(ResourceBase):
+    def __init__(self, client: SeekrFlowClient) -> None:
+        super().__init__(client)
+        self.system_prompt = AsyncSystemPromptResource(client)
+
     async def generate(
         self,
         instructions: str,
@@ -192,6 +285,86 @@ class AsyncAlignment(ResourceBase):
 
         return AlignmentResponse(**response.data)
 
+    async def outputs(self, id: str) -> List[AlignmentOutput]:
+        """
+        Retrieves output files for an alignment job asynchronously.
+
+        Args:
+            id (str): Alignment job ID whose outputs to fetch.
+
+        Returns:
+            list[AlignmentOutput]: Collection of alignment output metadata.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/alignment/{id}/outputs",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return [AlignmentOutput(**output) for output in response.data]  # type: ignore[arg-type]
+
+    async def delete(self, id: str) -> None:
+        """
+        Deletes an alignment job asynchronously.
+
+        Args:
+            id (str): Alignment job ID to delete.
+
+        Returns:
+            None
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"flow/alignment/{id}",
+            ),
+            stream=False,
+        )
+
+        # Endpoint returns 204 No Content
+        return None
+
+    async def cancel(self, id: str) -> AlignmentResponse:
+        """
+        Async method to cancel a running alignment job
+
+        Args:
+            id (str): Alignment job ID to cancel. A string that starts with `al-`.
+
+        Returns:
+            AlignmentResponse: Object containing information about cancelled alignment job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/{id}/cancel",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return AlignmentResponse(**response.data)
+
     async def estimate(self, files: List[str]) -> AlignmentEstimationResponse:
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -212,3 +385,289 @@ class AsyncAlignment(ResourceBase):
 
         assert isinstance(response, SeekrFlowResponse)
         return AlignmentEstimationResponse(**response.data)
+
+
+class SystemPromptResource(ResourceBase):
+    def create(self, source_id: str, instructions: str) -> SystemPrompt:
+        """
+        Creates a new AI-generated system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to create the system prompt for
+            instructions (str): Instructions for generating the system prompt
+
+        Returns:
+            SystemPrompt: The created system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptCreateRequest(
+            instructions=instructions
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/system_prompt/{source_id}",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    def get(self, source_id: str) -> SystemPrompt:
+        """
+        Retrieves the system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to retrieve the system prompt for
+
+        Returns:
+            SystemPrompt: The retrieved system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/alignment/system_prompt/{source_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    def regenerate(self, source_id: str, instructions: str) -> SystemPrompt:
+        """
+        Regenerates the AI-generated system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to regenerate the system prompt for
+            instructions (str): Instructions for regenerating the system prompt
+
+        Returns:
+            SystemPrompt: The regenerated system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptCreateRequest(
+            instructions=instructions
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/system_prompt/{source_id}/regenerate",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    def update(self, source_id: str, content: str) -> SystemPrompt:
+        """
+        Updates the system prompt for the given source_id with custom content.
+
+        Args:
+            source_id (str): The ID of the source to update the system prompt for
+            content (str): The custom content for the system prompt
+
+        Returns:
+            SystemPrompt: The updated system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptUpdateRequest(content=content).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="PUT",
+                url=f"flow/alignment/system_prompt/{source_id}",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    def delete(self, source_id: str) -> Dict[str, Any]:
+        """
+        Deletes the system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to delete the system prompt for
+
+        Returns:
+            dict: A dictionary with the deletion result
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"flow/alignment/system_prompt/{source_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return response.data
+
+
+class AsyncSystemPromptResource(ResourceBase):
+    async def create(self, source_id: str, instructions: str) -> SystemPrompt:
+        """
+        Asynchronously creates a new AI-generated system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to create the system prompt for
+            instructions (str): Instructions for generating the system prompt
+
+        Returns:
+            SystemPrompt: The created system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptCreateRequest(
+            instructions=instructions
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/system_prompt/{source_id}",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    async def get(self, source_id: str) -> SystemPrompt:
+        """
+        Asynchronously retrieves the system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to retrieve the system prompt for
+
+        Returns:
+            SystemPrompt: The retrieved system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/alignment/system_prompt/{source_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    async def regenerate(self, source_id: str, instructions: str) -> SystemPrompt:
+        """
+        Asynchronously regenerates the AI-generated system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to regenerate the system prompt for
+            instructions (str): Instructions for regenerating the system prompt
+
+        Returns:
+            SystemPrompt: The regenerated system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptCreateRequest(
+            instructions=instructions
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"flow/alignment/system_prompt/{source_id}/regenerate",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    async def update(self, source_id: str, content: str) -> SystemPrompt:
+        """
+        Asynchronously updates the system prompt for the given source_id with custom content.
+
+        Args:
+            source_id (str): The ID of the source to update the system prompt for
+            content (str): The custom content for the system prompt
+
+        Returns:
+            SystemPrompt: The updated system prompt
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = SystemPromptUpdateRequest(content=content).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="PUT",
+                url=f"flow/alignment/system_prompt/{source_id}",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return SystemPrompt(**response.data)
+
+    async def delete(self, source_id: str) -> Dict[str, Any]:
+        """
+        Asynchronously deletes the system prompt for the given source_id.
+
+        Args:
+            source_id (str): The ID of the source to delete the system prompt for
+
+        Returns:
+            dict: A dictionary with the deletion result
+        """
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="DELETE",
+                url=f"flow/alignment/system_prompt/{source_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return response.data

seekrai/resources/chat/completions.py

@@ -22,6 +22,7 @@ class ChatCompletions:
         *,
         messages: List[Dict[str, str]],
         model: str,
+        max_completion_tokens: int | None = None,
         max_tokens: int | None = 512,
         stop: List[str] | None = None,
         temperature: float = 0.7,
@@ -36,7 +37,7 @@ class ChatCompletions:
         safety_model: str | 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 = "auto",
+        tool_choice: str | Dict[str, str | Dict[str, str]] | None = None,
     ) -> ChatCompletionResponse | Iterator[ChatCompletionChunk]:
         """
         Method to generate completions based on a given prompt using a specified model.
@@ -45,6 +46,7 @@ class ChatCompletions:
             messages (List[Dict[str, str]]): A list of messages in the format
                 `[{"role": seekrai.types.chat_completions.MessageRole, "content": TEXT}, ...]`
             model (str): The name of the model to query.
+            max_completion_tokens (int, optional): The maximum number of tokens the output can contain.
             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.
@@ -99,6 +101,7 @@ class ChatCompletions:
             top_p=top_p,
             top_k=top_k,
             temperature=temperature,
+            max_completion_tokens=max_completion_tokens,
             max_tokens=max_tokens,
             stop=stop,
             repetition_penalty=repetition_penalty,
@@ -109,15 +112,16 @@ class ChatCompletions:
             n=n,
             safety_model=safety_model,
             response_format=response_format,
-            tools=tools or [],
-            tool_choice=tool_choice,
+            tools=tools or None,
         ).model_dump()
-
+        if tool_choice is not None:
+            parameter_payload["tool_choice"] = tool_choice
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
                 url="inference/chat/completions",
                 params=parameter_payload,
+                headers={"content-type": "application/json"},
             ),
             stream=stream,
         )
@@ -139,6 +143,7 @@ class AsyncChatCompletions:
         *,
         messages: List[Dict[str, str]],
         model: str,
+        max_completion_tokens: int | None = None,
         max_tokens: int | None = 512,
         stop: List[str] | None = None,
         temperature: float = 0.7,
@@ -152,8 +157,8 @@ class AsyncChatCompletions:
         n: int = 1,
         safety_model: str | 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,
+        tools: Dict[str, str | Dict[str, Any]] | None = None,
+        tool_choice: str | Dict[str, str | Dict[str, str]] | None = None,
     ) -> AsyncGenerator[ChatCompletionChunk, None] | ChatCompletionResponse:
         """
         Async method to generate completions based on a given prompt using a specified model.
@@ -162,6 +167,7 @@ class AsyncChatCompletions:
             messages (List[Dict[str, str]]): A list of messages in the format
                 `[{"role": seekrai.types.chat_completions.MessageRole, "content": TEXT}, ...]`
             model (str): The name of the model to query.
+            max_completion_tokens (int, optional): The maximum number of tokens the output can contain.
             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.
@@ -217,6 +223,7 @@ class AsyncChatCompletions:
             top_p=top_p,
             top_k=top_k,
             temperature=temperature,
+            max_completion_tokens=max_completion_tokens,
             max_tokens=max_tokens,
             stop=stop,
             repetition_penalty=repetition_penalty,
@@ -227,15 +234,17 @@ class AsyncChatCompletions:
             n=n,
             safety_model=safety_model,
             response_format=response_format,
-            # tools=tools,
-            # tool_choice=tool_choice,
+            tools=tools or None,
         ).model_dump()
+        if tool_choice is not None:
+            parameter_payload["tool_choice"] = tool_choice
 
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
                 url="inference/chat/completions",
                 params=parameter_payload,
+                headers={"content-type": "application/json"},
             ),
             stream=stream,
         )

seekrai/resources/embeddings.py

@@ -29,7 +29,6 @@ class Embeddings(ResourceBase):
         Returns:
             EmbeddingResponse: Object containing embeddings
         """
-        raise NotImplementedError("Function not implemented yet")
 
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -45,6 +44,7 @@ class Embeddings(ResourceBase):
                 method="POST",
                 url="inference/embeddings",
                 params=parameter_payload,
+                headers={"content-type": "application/json"},
             ),
             stream=False,
         )
@@ -71,7 +71,6 @@ class AsyncEmbeddings(ResourceBase):
         Returns:
             EmbeddingResponse: Object containing embeddings
         """
-        raise NotImplementedError("Function not implemented yet")
 
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -87,6 +86,7 @@ class AsyncEmbeddings(ResourceBase):
                 method="POST",
                 url="inference/embeddings",
                 params=parameter_payload,
+                headers={"content-type": "application/json"},
             ),
             stream=False,
         )