seekrai 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

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/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/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/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/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/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

{seekrai-0.3.2.dist-info → seekrai-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: seekrai
-Version: 0.3.2
+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.2.dist-info → seekrai-0.3.3.dist-info}/RECORD

@@ -1,14 +1,14 @@
 seekrai/__init__.py,sha256=HC6iy-IdwqecabH-6a80Lsy9qO2PBToAI0WqEErV41c,935
 seekrai/abstract/__init__.py,sha256=wNiOTW9TJpUgfCJCG-wAbhhWWH2PtoVpAuL3nxvQGps,56
-seekrai/abstract/api_requestor.py,sha256=tI43JdbUfNQHWjvRrenUM7C7Ea-zD-7yd2JNS-Lki04,18297
+seekrai/abstract/api_requestor.py,sha256=NT41M0_Utfw61bMqibP1PEdIIQmmKWfzj-Y-xk4DNY0,18395
 seekrai/client.py,sha256=Yhejl-2a-Uoc8nWi-XxETZT4a4Ou_t_TRLcp0e9APIY,5376
 seekrai/constants.py,sha256=hoR2iF5te5Ydjt_lxIOSGID4vESIakG4F-3xAWdwxaU,1854
 seekrai/error.py,sha256=rAYL8qEd8INwYMMKvhS-HKeC3QkWL4Wq-zfazFU-zBg,4861
 seekrai/filemanager.py,sha256=8RuSzJvELD-fCI2Wd_t0jSKeVrmFwF7E5AzXIgDFxNA,9572
 seekrai/resources/__init__.py,sha256=-OpvrZp0c_ro_e1G14gjOILtcSfpSOeRZeK8kt41WLo,1014
-seekrai/resources/alignment.py,sha256=6qQm9w0Em0q3zVeOzs8cX3wulr-B-wh4Pcr3pboIeTE,4468
+seekrai/resources/alignment.py,sha256=U2wMIg45Bo_EAuThn_YotPCHIS_3smHyAf3EhdyPovM,5826
 seekrai/resources/chat/__init__.py,sha256=KmtPupgECtEN80NyvcnSmieTAFXhwmVxhMHP0qhspA4,618
-seekrai/resources/chat/completions.py,sha256=pGBVrIvUfY-wr0ooiywzuhIjZ4GwiojH3ofmRE6bYlw,11643
+seekrai/resources/chat/completions.py,sha256=v_gaeryfOCPL_vR0ysGEvSejai2dNdLcZh2gC6567rk,11646
 seekrai/resources/completions.py,sha256=w3La3zPMlN00y-b-tJwLgvZVH-xK_dKC6ktI5Ggn1us,8564
 seekrai/resources/deployments.py,sha256=HmP7MuxAlLUoQl6z705_d1Y53MDvGgSQXhlgN3DKX2A,6078
 seekrai/resources/embeddings.py,sha256=3lohUrkdFqzSg8FgS7p4r87jwjE0NXU1PilWv278quk,2705
@@ -18,9 +18,9 @@ seekrai/resources/images.py,sha256=E48lAe7YsZ2WXBHR_qz4SF7P4Y-U7t61m_bWNS91pM0,4
 seekrai/resources/models.py,sha256=Pdd0S0gZdratWcHJPKNb7LkEdUGjr3xNR06W6GDiyxk,5000
 seekrai/resources/projects.py,sha256=AWJUeUDSzkbxBksHjJ4a3c83UR62TlMGHutm2NdV6Xk,3790
 seekrai/seekrflow_response.py,sha256=5RFEQzamDy7sTSDkxSsZQThZ3biNmeCPeHWdrFId5Go,1320
-seekrai/types/__init__.py,sha256=E41__vAq-tnm5aVYW0Ef1Ky-GMPnMpUKlyLr3pYhLSM,2358
+seekrai/types/__init__.py,sha256=XJ8lP4XDGL5mr1SbKJFFeponmlCRNp1HrvxscoxfRlU,2492
 seekrai/types/abstract.py,sha256=TqWFQV_6bPblywfCH-r8FCkXWvPkc9KlJ4QVgyrnaMc,642
-seekrai/types/alignment.py,sha256=-lxF0Cyj3w4brMft8WAoc2EDMAlbX7zdRYD04ot98-Y,1129
+seekrai/types/alignment.py,sha256=P4HZbmFjgtZkWsj6-ZjZ1v-FZkHusEaIEiRWREla-wA,1632
 seekrai/types/chat_completions.py,sha256=xRTHBbDJDbz0HgW042WX3csQDolhjEuO81w0rzFSeBU,3691
 seekrai/types/common.py,sha256=OH3l3u-0_5oz1KYrcHMybFESzivDySocYlJAsLSLOWU,1940
 seekrai/types/completions.py,sha256=lm9AFdZR3Xg5AHPkV-qETHikkwMJmkHrLGr5GG-YR-M,2171
@@ -28,7 +28,7 @@ seekrai/types/deployments.py,sha256=n7_t7DEeBSC8cDJSjIfvWtcgUql1DaEn89zuGYN_RaI,
 seekrai/types/embeddings.py,sha256=OANoLNOs0aceS8NppVvvcNYQbF7-pAOAmcr30pw64OU,749
 seekrai/types/error.py,sha256=uTKISs9aRC4_6zwirtNkanxepN8KY-SqCq0kNbfZylQ,370
 seekrai/types/files.py,sha256=XmtiM6d9i3tnYS-Kii3QpxZJRqemJi2rvLJ32GsECXQ,2602
-seekrai/types/finetune.py,sha256=Izh7NfaW7pGzSN5GpWYSaZpkHV2TilE1CF62DrB8vE0,6089
+seekrai/types/finetune.py,sha256=VNuAZx-Lq4Q55hY47obWnFLDr9if93v5Lwr_qcyCtfU,6124
 seekrai/types/images.py,sha256=Fusj8OhVYFsT8kz636lRGGivLbPXo_ZNgakKwmzJi3U,914
 seekrai/types/models.py,sha256=1ZfW9WwayApkISRizDntjkWhYNv-wkbrRVIfHn2QuC4,1242
 seekrai/types/projects.py,sha256=JFgpZdovia8Orcnhp6QkIEAXzyPCfKT_bUiwjxUaHHQ,670
@@ -38,8 +38,8 @@ seekrai/utils/api_helpers.py,sha256=0Y8BblNIr9h_R12zdmhkxgTlxgoRkbq84QNi4nNWGu8,
 seekrai/utils/files.py,sha256=B61Pwra49MVVWjPtdkx4hBtAuUe9UI63hdNus87Uq0o,7164
 seekrai/utils/tools.py,sha256=jgJTL-dOIouDbEJLdQpQfpXhqaz_poQYS52adyUtBjo,1781
 seekrai/version.py,sha256=q6iGQVFor8zXiPP5F-3vy9TndOxKv5JXbaNJ2kdOQws,125
-seekrai-0.3.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-seekrai-0.3.2.dist-info/METADATA,sha256=M1kwCRaiObFvV3IpKSYII0kx3_FqGyPRZ8tU1ukz3e8,4748
-seekrai-0.3.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-seekrai-0.3.2.dist-info/entry_points.txt,sha256=N49yOEGi1sK7Xr13F_rkkcOxQ88suyiMoOmRhUHTZ_U,48
-seekrai-0.3.2.dist-info/RECORD,,
+seekrai-0.3.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+seekrai-0.3.3.dist-info/METADATA,sha256=avqvw86pT87l1EEYJHqURviZ_IMmWWYVViRcBQxojOE,4748
+seekrai-0.3.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+seekrai-0.3.3.dist-info/entry_points.txt,sha256=N49yOEGi1sK7Xr13F_rkkcOxQ88suyiMoOmRhUHTZ_U,48
+seekrai-0.3.3.dist-info/RECORD,,