seekrai 0.3.1__tar.gz → 0.3.3__tar.gz

{seekrai-0.3.1 → seekrai-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: seekrai
-Version: 0.3.1
+Version: 0.3.3
 Summary: Python client for SeekrAI
 Home-page: https://gitlab.cb.ntent.com/ml/seekr-py
 License: Apache-2.0

{seekrai-0.3.1 → seekrai-0.3.3}/pyproject.toml

@@ -14,7 +14,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "seekrai"
-version = "0.3.1"
+version = "0.3.3"
 authors = [
     "SeekrFlow <support@seekr.com>"
 ]
@@ -73,6 +73,7 @@ tox = "^4.14.1"
 [tool.poetry.group.dev.dependencies]
 ipdb = "^0.13.13"
 pre-commit = "^3.7.1"
+jupyter = "^1.0.0"
 
 [tool.poetry.urls]
 "Homepage" = "https://www.seekr.com/"

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/abstract/api_requestor.py

@@ -524,6 +524,8 @@ class APIRequestor:
                 if rbody.strip().endswith("[DONE]"):
                     # TODO
                     rbody = rbody.replace("data: [DONE]", "")
+                if rbody.startswith("data: "):
+                    rbody = rbody[len("data: ") :]
                 data = json.loads(rbody)
         except (JSONDecodeError, UnicodeDecodeError) as e:
             raise error.APIError(

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/resources/alignment.py

@@ -3,6 +3,8 @@ from typing import List
 from seekrai.abstract import api_requestor
 from seekrai.seekrflow_response import SeekrFlowResponse
 from seekrai.types import (
+    AlignmentEstimationRequest,
+    AlignmentEstimationResponse,
     AlignmentList,
     AlignmentRequest,
     AlignmentResponse,
@@ -92,6 +94,27 @@ class Alignment:
 
         return AlignmentResponse(**response.data)
 
+    def estimate(self, files: List[str]) -> AlignmentEstimationResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = AlignmentEstimationRequest(
+            files=files,
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="flow/alignment/estimate",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return AlignmentEstimationResponse(**response.data)
+
 
 class AsyncAlignment:
     def __init__(self, client: SeekrFlowClient) -> None:
@@ -173,3 +196,24 @@ class AsyncAlignment:
         assert isinstance(response, SeekrFlowResponse)
 
         return AlignmentResponse(**response.data)
+
+    async def estimate(self, files: List[str]) -> AlignmentEstimationResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = AlignmentEstimationRequest(
+            files=files,
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="flow/alignment/estimate",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        return AlignmentEstimationResponse(**response.data)

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/resources/chat/completions.py

@@ -30,13 +30,13 @@ class ChatCompletions:
         repetition_penalty: float = 1,
         stream: bool = False,
         logprobs: bool | None = False,
-        top_logprobs: int | None = 0,
+        top_logprobs: int | None = None,
         echo: bool = False,
         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 = [],
+        tool_choice: str | Dict[str, str | Dict[str, str]] | None = "auto",
     ) -> ChatCompletionResponse | Iterator[ChatCompletionChunk]:
         """
         Method to generate completions based on a given prompt using a specified model.

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/resources/finetune.py

@@ -129,7 +129,7 @@ class FineTuning:
 
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
-                method="POST",
+                method="PUT",
                 url=f"flow/fine-tunes/{id}/cancel",
             ),
             stream=False,

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/types/__init__.py

@@ -1,5 +1,7 @@
 from seekrai.types.abstract import SeekrFlowClient
 from seekrai.types.alignment import (
+    AlignmentEstimationRequest,
+    AlignmentEstimationResponse,
     AlignmentJobStatus,
     AlignmentList,
     AlignmentRequest,
@@ -87,6 +89,8 @@ __all__ = [
     "ModelResponse",
     "ModelList",
     "AlignmentRequest",
+    "AlignmentEstimationRequest",
+    "AlignmentEstimationResponse",
     "AlignmentResponse",
     "AlignmentJobStatus",
     "AlignmentList",

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/types/alignment.py

@@ -18,6 +18,24 @@ class AlignmentRequest(BaseModel):
     )
 
 
+class AlignmentEstimationRequest(BaseModel):
+    files: List[str] = Field(
+        default=...,
+        description="List of file ids to use to generate an alignment estiamte",
+    )
+
+
+class AlignmentEstimationResponse(BaseModel):
+    input_tokens: int = Field(
+        default=...,
+        description="Estimated number of input tokens for the given files",
+    )
+    output_tokens: int = Field(
+        default=...,
+        description="Estimated number of output tokens for the given files",
+    )
+
+
 class AlignmentJobStatus(str, Enum):
     STATUS_PENDING = "pending"
     STATUS_QUEUED = "queued"

{seekrai-0.3.1 → seekrai-0.3.3}/src/seekrai/types/finetune.py

@@ -172,6 +172,7 @@ class FinetuneResponse(BaseModel):
     inference_available: bool = False
     project_id: Optional[int] = None  # TODO - fix this
     completed_at: datetime | None = None
+    description: str | None = None
 
     # dataset token count
     # TODO