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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.1.0
+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.1.0 → speedy_utils-1.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "speedy-utils"
-version = "1.1.0"
+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.1.0 → 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,
@@ -146,10 +146,13 @@ def _yellow(t):
     return _color(33, t)
 
 
-class ParsedOutput(TypedDict):
+
+TParsed = TypeVar('TParsed', bound=BaseModel)
+
+class ParsedOutput(TypedDict, Generic[TParsed]):
     messages: List
     completion: Any
-    parsed: BaseModel
+    parsed: TParsed
 
 
 class AsyncLM:
@@ -460,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,
@@ -470,7 +473,7 @@ class AsyncLM:
         max_tokens: Optional[int] = None,
         cache: Optional[bool] = True,
         **kwargs,
-    ) -> ParsedOutput:  # -> dict[str, Any]:
+    ) -> ParsedOutput[TParsed]:
         """Parse response using guided JSON generation."""
         if messages is None:
             assert instruction is not None, "Instruction must be provided."
@@ -513,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,
@@ -522,7 +526,7 @@ class AsyncLM:
             }
             cache_key = self._cache_key(cache_data, {}, response_model)
             completion = self._load_cache(cache_key)  # dict
-        else:
+        if not completion:
             completion = await self.client.chat.completions.create(
                 model=self.model,  # type: ignore
                 messages=messages,  # type: ignore
@@ -532,10 +536,12 @@ class AsyncLM:
             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."
+        )
         self.last_log = [prompt, messages, completion]
 
-        output = self._parse_complete_output(completion, response_model)
+        output = cast(TParsed, self._parse_complete_output(completion, response_model))
         full_messages = messages + [completion]
         return ParsedOutput(
             messages=full_messages,
@@ -894,6 +900,7 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
     temperature: float = 0.6
     think: bool = False
     add_json_schema: bool = False
+    cache: bool = False
 
     async def __call__(
         self,
@@ -942,7 +949,7 @@ 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,
+            cache=self.cache or cache,
         )
 
         return (