speedy-utils 1.0.24__py3-none-any.whl → 1.1.0__py3-none-any.whl

llm_utils/lm/async_lm.py

@@ -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,12 @@ def _yellow(t):
     return _color(33, t)
 
 
+class ParsedOutput(TypedDict):
+    messages: List
+    completion: Any
+    parsed: BaseModel
+
+
 class AsyncLM:
     """Unified **async** language‑model wrapper with optional JSON parsing."""
 
@@ -462,11 +468,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:  # -> dict[str, Any]:
         """Parse response using guided JSON generation."""
         if messages is None:
             assert instruction is not None, "Instruction must be provided."
@@ -517,37 +521,27 @@ 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,
-        )
-
-        if cache_key:
-            self._dump_cache(cache_key, completion)
+            completion = self._load_cache(cache_key)  # dict
+        else:
+            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)
 
         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
+        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,8 +888,8 @@ 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
@@ -906,8 +900,7 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
         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 +923,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(),
@@ -940,15 +943,12 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
             think=self.think,
             add_json_schema_to_instruction=self.add_json_schema,
             cache=cache,
-            return_messages=True,
         )
 
-        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 +962,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

{speedy_utils-1.0.24.dist-info → speedy_utils-1.1.0.dist-info}/METADATA

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

{speedy_utils-1.0.24.dist-info → speedy_utils-1.1.0.dist-info}/RECORD

@@ -5,7 +5,7 @@ llm_utils/chat_format/transform.py,sha256=8TZhvUS5DrjUeMNtDIuWY54B_QZ7jjpXEL9c8F
 llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
 llm_utils/group_messages.py,sha256=8CU9nKOja3xeuhdrX5CvYVveSqSKb2zQ0eeNzA88aTQ,3621
 llm_utils/lm/__init__.py,sha256=rX36_MsnekM5GHwWS56XELbm4W5x2TDwnPERDTfo0eU,194
-llm_utils/lm/async_lm.py,sha256=xUOlAOivxrI-KRaUv1V0l5y7ajYTKTA6QZkqpK6Uue8,35847
+llm_utils/lm/async_lm.py,sha256=kiWEecrkCTTQFlQj5JiHNziFeLOF1-7G_2xC2Dra1bw,35806
 llm_utils/lm/chat_html.py,sha256=FkGo0Dv_nAHYBMZzXfMu_bGQKaCx302goh3XaT-_ETc,8674
 llm_utils/lm/lm_json.py,sha256=fMt42phzFV2f6ulrtWcDXsWHi8WcG7gGkCzpIq8VSSM,1975
 llm_utils/lm/sync_lm.py,sha256=ANw_m5KiWcRwwoeQ5no6dzPFLc6j9o2oEcJtkMKqrn8,34640
@@ -31,7 +31,7 @@ speedy_utils/multi_worker/thread.py,sha256=u_hTwXh7_FciMa5EukdEA1fDCY_vUC4moDceB
 speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 speedy_utils/scripts/mpython.py,sha256=73PHm1jqbCt2APN4xuNjD0VDKwzOj4EZsViEMQiZU2g,3853
 speedy_utils/scripts/openapi_client_codegen.py,sha256=f2125S_q0PILgH5dyzoKRz7pIvNEjCkzpi4Q4pPFRZE,9683
-speedy_utils-1.0.24.dist-info/METADATA,sha256=p8QjQuz3u1B50Sj5qoPR0p-FsrYGvhd19dwfYM_MlwA,7442
-speedy_utils-1.0.24.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.24.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
-speedy_utils-1.0.24.dist-info/RECORD,,
+speedy_utils-1.1.0.dist-info/METADATA,sha256=h1Alzm4q92GSiw5GNZWn6d8sHaSJS4X8RTMXStjkqHY,7441
+speedy_utils-1.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+speedy_utils-1.1.0.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
+speedy_utils-1.1.0.dist-info/RECORD,,