speedy-utils 1.0.24__tar.gz → 1.1.2__tar.gz

{speedy_utils-1.0.24 → speedy_utils-1.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.0.24
+Version: 1.1.2
 Summary: Fast and easy-to-use package for data science
 Author: AnhVTH
 Author-email: anhvth.226@gmail.com

{speedy_utils-1.0.24 → speedy_utils-1.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "speedy-utils"
-version = "1.0.24"
+version = "1.1.2"
 description = "Fast and easy-to-use package for data science"
 authors = ["AnhVTH <anhvth.226@gmail.com>"]
 readme = "README.md"

{speedy_utils-1.0.24 → speedy_utils-1.1.2}/src/llm_utils/lm/async_lm.py

@@ -78,7 +78,7 @@ import hashlib
 import json
 import os
 from abc import ABC
-from functools import lru_cache
+from functools import cache, lru_cache
 from typing import (
     Any,
     Dict,
@@ -93,7 +93,7 @@ from typing import (
     cast,
     overload,
 )
-
+from typing_extensions import TypedDict
 from httpx import URL
 from loguru import logger
 from numpy import isin
@@ -146,6 +146,15 @@ def _yellow(t):
     return _color(33, t)
 
 
+
+TParsed = TypeVar('TParsed', bound=BaseModel)
+
+class ParsedOutput(TypedDict, Generic[TParsed]):
+    messages: List
+    completion: Any
+    parsed: TParsed
+
+
 class AsyncLM:
     """Unified **async** language‑model wrapper with optional JSON parsing."""
 
@@ -454,7 +463,7 @@ class AsyncLM:
     # ------------------------------------------------------------------ #
     async def parse(
         self,
-        response_model: Type[BaseModel],
+        response_model: Type[TParsed],
         instruction: Optional[str] = None,
         prompt: Optional[str] = None,
         messages: Optional[RawMsgs] = None,
@@ -462,11 +471,9 @@ class AsyncLM:
         add_json_schema_to_instruction: bool = False,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
-        return_openai_response: bool = False,
         cache: Optional[bool] = True,
-        return_messages: bool = False,
         **kwargs,
-    ):
+    ) -> ParsedOutput[TParsed]:
         """Parse response using guided JSON generation."""
         if messages is None:
             assert instruction is not None, "Instruction must be provided."
@@ -509,6 +516,7 @@ class AsyncLM:
 
         use_cache = self.do_cache if cache is None else cache
         cache_key = None
+        completion = None
         if use_cache:
             cache_data = {
                 "messages": messages,
@@ -517,37 +525,29 @@ class AsyncLM:
                 "response_format": response_model.__name__,
             }
             cache_key = self._cache_key(cache_data, {}, response_model)
-            cached_response = self._load_cache(cache_key)
-            self.last_log = [prompt, messages, cached_response]
-            if cached_response is not None:
-                if return_openai_response:
-                    return cached_response
-                return self._parse_complete_output(cached_response, response_model)
-
-        completion = await self.client.chat.completions.create(
-            model=self.model,  # type: ignore
-            messages=messages,  # type: ignore
-            extra_body={"guided_json": json_schema},
-            **model_kwargs,
+            completion = self._load_cache(cache_key)  # dict
+        if not completion:
+            completion = await self.client.chat.completions.create(
+                model=self.model,  # type: ignore
+                messages=messages,  # type: ignore
+                extra_body={"guided_json": json_schema},
+                **model_kwargs,
+            )
+            completion = completion.model_dump()
+            if cache_key:
+                self._dump_cache(cache_key, completion)
+        assert isinstance(completion, dict), (
+            "Completion must be a dictionary with OpenAI response format."
         )
-
-        if cache_key:
-            self._dump_cache(cache_key, completion)
-
         self.last_log = [prompt, messages, completion]
 
-        output = self._parse_complete_output(completion, response_model)
-        if return_openai_response:
-            return {"completion": completion, "parsed": output}
-        if return_messages:
-            # content = completion.model_dump()
-            full_messages = messages + [completion.model_dump()]
-            return {
-                "messages": full_messages,
-                "completion": completion,
-                "parsed": output,
-            }
-        return output
+        output = cast(TParsed, self._parse_complete_output(completion, response_model))
+        full_messages = messages + [completion]
+        return ParsedOutput(
+            messages=full_messages,
+            completion=completion,
+            parsed=output,
+        )
 
     def _parse_complete_output(
         self, completion: Any, response_model: Type[BaseModel]
@@ -894,20 +894,20 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
     """
 
     lm: "AsyncLM"
-    InputModel: Type[BaseModel]
-    OutputModel: Type[BaseModel]
+    InputModel: InputModelType
+    OutputModel: OutputModelType
 
     temperature: float = 0.6
     think: bool = False
     add_json_schema: bool = False
+    cache: bool = False
 
     async def __call__(
         self,
         data: BaseModel | dict,
         temperature: float = 0.1,
         cache: bool = False,
-        collect_messages: bool = False,
-    ) -> OutputModelType | tuple[OutputModelType, List[Dict[str, Any]]]:
+    ) -> tuple[OutputModelType, List[Dict[str, Any]]]:
         # Get the input and output model types from the generic parameters
         type_args = getattr(self.__class__, "__orig_bases__", None)
         if (
@@ -930,7 +930,17 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
             input_model = self.InputModel
             output_model = self.OutputModel
 
-        item = data if isinstance(data, BaseModel) else input_model(**data)
+        # Ensure input_model is a class before calling
+        if isinstance(data, BaseModel):
+            item = data
+        elif isinstance(input_model, type) and issubclass(input_model, BaseModel):
+            item = input_model(**data)
+        else:
+            raise TypeError("InputModel must be a subclass of BaseModel")
+
+        assert isinstance(output_model, type) and issubclass(output_model, BaseModel), (
+            "OutputModel must be a subclass of BaseModel"
+        )
 
         result = await self.lm.parse(
             prompt=item.model_dump_json(),
@@ -939,16 +949,13 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
             temperature=temperature or self.temperature,
             think=self.think,
             add_json_schema_to_instruction=self.add_json_schema,
-            cache=cache,
-            return_messages=True,
+            cache=self.cache or cache,
         )
 
-        if collect_messages:
-            return (
-                cast(OutputModelType, result),
-                result["messages"],
-            )
-        return cast(OutputModelType, result)
+        return (
+            cast(OutputModelType, result["parsed"]),  # type: ignore
+            cast(List[dict], result["messages"]),  # type: ignore
+        )
 
     def generate_training_data(
         self, input_dict: Dict[str, Any], output: Dict[str, Any]
@@ -962,4 +969,4 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
         )
         return {"messages": messages}
 
-    # arun = __call__  # alias for compatibility with other LLMTask implementations
+    arun = __call__  # alias for compatibility with other LLMTask implementations