Merge pull request #59 from JigsawStack/fix/v3checks

Fix/v3checks

Author: Khurdhula-Harshavardhan

.gitignore

@@ -30,3 +30,4 @@ pyproject.toml
 uv.lock
 
 .ruff_cache/
+local_tests/

jigsawstack/__init__.py

@@ -107,7 +107,7 @@ def __init__(
             api_url=api_url,
             disable_request_logging=disable_request_logging,
         )
-        
+
         self.embedding = Embedding(
             api_key=api_key,
             api_url=api_url,
@@ -123,8 +123,7 @@ def __init__(
             api_key=api_key,
             api_url=api_url,
             disable_request_logging=disable_request_logging,
-        )
-
+        ).classify
 
 
 class AsyncJigsawStack:
@@ -218,7 +217,6 @@ def __init__(
             disable_request_logging=disable_request_logging,
         )
 
-
         self.embedding = AsyncEmbedding(
             api_key=api_key,
             api_url=api_url,
@@ -235,8 +233,7 @@ def __init__(
             api_key=api_key,
             api_url=api_url,
             disable_request_logging=disable_request_logging,
-        )
-
+        ).classify
 
 
 # Create a global instance of the Web class

jigsawstack/_types.py

@@ -0,0 +1,13 @@
+from typing_extensions import NotRequired, TypedDict
+
+
+class UsageInfo(TypedDict):
+    input_tokens: int
+    output_tokens: int
+    inference_time_tokens: int
+    total_tokens: int
+
+
+class BaseResponse(TypedDict):
+    success: bool
+    _usage: NotRequired[UsageInfo]

jigsawstack/audio.py

@@ -7,6 +7,7 @@
 from typing_extensions import NotRequired, TypedDict, Literal
 from .custom_typing import SupportedAccents
 from .helpers import build_path
+from ._types import BaseResponse
 
 
 class SpeechToTextParams(TypedDict):
@@ -27,15 +28,28 @@ class ChunkParams(TypedDict):
 
 class BySpeakerParams(ChunkParams):
     speaker: str
+    timestamp: tuple[int, int]
+    text: str
 
 
-class SpeechToTextResponse(TypedDict):
-    success: bool
+class SpeechToTextResponse(BaseResponse):
     text: str
     chunks: List[ChunkParams]
     speakers: Optional[List[BySpeakerParams]]
 
 
+class SpeechToTextWebhookResponse(BaseResponse):
+    status: Literal["processing", "error"]
+    """
+    the status of the transcription process
+    """
+
+    id: str
+    """
+    the id of the transcription process
+    """
+
+
 class Audio(ClientConfig):
     config: RequestConfig
 
@@ -53,17 +67,19 @@ def __init__(
         )
 
     @overload
-    def speech_to_text(self, params: SpeechToTextParams) -> SpeechToTextResponse: ...
+    def speech_to_text(
+        self, params: SpeechToTextParams
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]: ...
     @overload
     def speech_to_text(
         self, blob: bytes, options: Optional[SpeechToTextParams] = None
-    ) -> SpeechToTextResponse: ...
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]: ...
 
     def speech_to_text(
         self,
         blob: Union[SpeechToTextParams, bytes],
         options: Optional[SpeechToTextParams] = None,
-    ) -> SpeechToTextResponse:
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]:
         if isinstance(
             blob, dict
         ):  # If params is provided as a dict, we assume it's the first argument
@@ -110,17 +126,17 @@ def __init__(
     @overload
     async def speech_to_text(
         self, params: SpeechToTextParams
-    ) -> SpeechToTextResponse: ...
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]: ...
     @overload
     async def speech_to_text(
         self, blob: bytes, options: Optional[SpeechToTextParams] = None
-    ) -> SpeechToTextResponse: ...
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]: ...
 
     async def speech_to_text(
         self,
         blob: Union[SpeechToTextParams, bytes],
         options: Optional[SpeechToTextParams] = None,
-    ) -> SpeechToTextResponse:
+    ) -> Union[SpeechToTextResponse, SpeechToTextWebhookResponse]:
         if isinstance(blob, dict):
             resp = await AsyncRequest(
                 config=self.config,

jigsawstack/classification.py

@@ -3,110 +3,63 @@
 from .request import Request, RequestConfig
 from .async_request import AsyncRequest, AsyncRequestConfig
 from ._config import ClientConfig
+from ._types import BaseResponse
 
 
-class DatasetItemText(TypedDict):
-    type: Literal["text"]
+class DatasetItem(TypedDict):
+    type: Literal["text", "image"]
     """
     Type of the dataset item: text
     """
-    
-    value: str
-    """
-    Value of the dataset item
-    """
-
 
-class DatasetItemImage(TypedDict):
-    type: Literal["image"]
-    """
-    Type of the dataset item: image
-    """
-    
     value: str
     """
     Value of the dataset item
     """
 
 
-class LabelItemText(TypedDict):
+class LabelItem(TypedDict):
     key: NotRequired[str]
     """
     Optional key for the label
     """
-    
-    type: Literal["text"]
+
+    type: Literal["text", "image"]
     """
     Type of the label: text
     """
-    
-    value: str
-    """
-    Value of the label
-    """
 
-
-class LabelItemImage(TypedDict):
-    key: NotRequired[str]
-    """
-    Optional key for the label
-    """
-    
-    type: Literal["image", "text"]
-    """
-    Type of the label: image or text
-    """
-    
     value: str
     """
     Value of the label
     """
 
 
-class ClassificationTextParams(TypedDict):
-    dataset: List[DatasetItemText]
+class ClassificationParams(TypedDict):
+    dataset: List[DatasetItem]
     """
     List of text dataset items to classify
     """
-    
-    labels: List[LabelItemText]
+
+    labels: List[LabelItem]
     """
     List of text labels for classification
     """
-    
-    multiple_labels: NotRequired[bool]
-    """
-    Whether to allow multiple labels per item
-    """
 
-
-class ClassificationImageParams(TypedDict):
-    dataset: List[DatasetItemImage]
-    """
-    List of image dataset items to classify
-    """
-    
-    labels: List[LabelItemImage]
-    """
-    List of labels for classification
-    """
-    
     multiple_labels: NotRequired[bool]
     """
     Whether to allow multiple labels per item
     """
 
 
-class ClassificationResponse(TypedDict):
+class ClassificationResponse(BaseResponse):
     predictions: List[Union[str, List[str]]]
     """
     Classification predictions - single labels or multiple labels per item
     """
 
 
-
 class Classification(ClientConfig):
-
     config: RequestConfig
 
     def __init__(
@@ -122,16 +75,7 @@ def __init__(
             disable_request_logging=disable_request_logging,
         )
 
-    def text(self, params: ClassificationTextParams) -> ClassificationResponse:
-        path = "/classification"
-        resp = Request(
-            config=self.config,
-            path=path,
-            params=cast(Dict[Any, Any], params),
-            verb="post",
-        ).perform_with_content()
-        return resp
-    def image(self, params: ClassificationImageParams) -> ClassificationResponse:
+    def classify(self, params: ClassificationParams) -> ClassificationResponse:
         path = "/classification"
         resp = Request(
             config=self.config,
@@ -142,7 +86,6 @@ def image(self, params: ClassificationImageParams) -> ClassificationResponse:
         return resp
 
 
-
 class AsyncClassification(ClientConfig):
     config: AsyncRequestConfig
 
@@ -159,7 +102,7 @@ def __init__(
             disable_request_logging=disable_request_logging,
         )
 
-    async def text(self, params: ClassificationTextParams) -> ClassificationResponse:
+    async def classify(self, params: ClassificationParams) -> ClassificationResponse:
         path = "/classification"
         resp = await AsyncRequest(
             config=self.config,
@@ -168,13 +111,3 @@ async def text(self, params: ClassificationTextParams) -> ClassificationResponse
             verb="post",
         ).perform_with_content()
         return resp
-
-    async def image(self, params: ClassificationImageParams) -> ClassificationResponse:
-        path = "/classification"
-        resp = await AsyncRequest(
-            config=self.config,
-            path=path,
-            params=cast(Dict[Any, Any], params),
-            verb="post",
-        ).perform_with_content()
-        return resp

jigsawstack/embedding.py

@@ -5,6 +5,7 @@
 from typing import List, Union
 from ._config import ClientConfig
 from .helpers import build_path
+from ._types import BaseResponse
 
 
 class EmbeddingParams(TypedDict):
@@ -13,17 +14,20 @@ class EmbeddingParams(TypedDict):
     type: Literal["text", "text-other", "image", "audio", "pdf"]
     url: NotRequired[str]
     file_store_key: NotRequired[str]
-    token_overflow_mode: NotRequired[Literal["truncate", "chunk", "error"]] = "chunk"
+    token_overflow_mode: NotRequired[Literal["truncate", "error"]]
 
 
-class EmbeddingResponse(TypedDict):
-    success: bool
+class Chunk(TypedDict):
+    text: str
+    timestamp: List[int]
+
+
+class EmbeddingResponse(BaseResponse):
     embeddings: List[List[float]]
-    chunks: List[str]
+    chunks: List[Chunk]
 
 
 class Embedding(ClientConfig):
-
     config: RequestConfig
 
     def __init__(
@@ -42,14 +46,16 @@ def __init__(
     @overload
     def execute(self, params: EmbeddingParams) -> EmbeddingResponse: ...
     @overload
-    def execute(self, blob: bytes, options: EmbeddingParams = None) -> EmbeddingResponse: ...
+    def execute(
+        self, blob: bytes, options: EmbeddingParams = None
+    ) -> EmbeddingResponse: ...
 
     def execute(
         self,
         blob: Union[EmbeddingParams, bytes],
         options: EmbeddingParams = None,
     ) -> EmbeddingResponse:
-        path="/embedding"
+        path = "/embedding"
         if isinstance(blob, dict):
             resp = Request(
                 config=self.config,
@@ -76,7 +82,6 @@ def execute(
 
 
 class AsyncEmbedding(ClientConfig):
-
     config: RequestConfig
 
     def __init__(
@@ -95,14 +100,16 @@ def __init__(
     @overload
     async def execute(self, params: EmbeddingParams) -> EmbeddingResponse: ...
     @overload
-    async def execute(self, blob: bytes, options: EmbeddingParams = None) -> EmbeddingResponse: ...
+    async def execute(
+        self, blob: bytes, options: EmbeddingParams = None
+    ) -> EmbeddingResponse: ...
 
     async def execute(
         self,
         blob: Union[EmbeddingParams, bytes],
         options: EmbeddingParams = None,
     ) -> EmbeddingResponse:
-        path="/embedding"
+        path = "/embedding"
         if isinstance(blob, dict):
             resp = await AsyncRequest(
                 config=self.config,