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

llm_utils/lm/async_lm.py

@@ -82,6 +82,7 @@ from functools import lru_cache
 from typing import (
     Any,
     Dict,
+    Generic,
     List,
     Literal,
     Optional,
@@ -95,6 +96,7 @@ from typing import (
 
 from httpx import URL
 from loguru import logger
+from numpy import isin
 from openai import AsyncOpenAI, AuthenticationError, BadRequestError, RateLimitError
 from openai.pagination import AsyncPage as AsyncSyncPage
 
@@ -462,6 +464,7 @@ class AsyncLM:
         max_tokens: Optional[int] = None,
         return_openai_response: bool = False,
         cache: Optional[bool] = True,
+        return_messages: bool = False,
         **kwargs,
     ):
         """Parse response using guided JSON generation."""
@@ -532,9 +535,19 @@ class AsyncLM:
             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
-        return self._parse_complete_output(completion, response_model)
+            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
 
     def _parse_complete_output(
         self, completion: Any, response_model: Type[BaseModel]
@@ -839,8 +852,11 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
 # Async LLMTask class
 # --------------------------------------------------------------------------- #
 
+InputModelType = TypeVar("InputModelType", bound=BaseModel)
+OutputModelType = TypeVar("OutputModelType", bound=BaseModel)
 
-class AsyncLLMTask(ABC):
+
+class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
     """
     Async callable wrapper around an AsyncLM endpoint.
 
@@ -885,30 +901,58 @@ class AsyncLLMTask(ABC):
     think: bool = False
     add_json_schema: bool = False
 
-    async def __call__(self, data: BaseModel | dict) -> BaseModel:
+    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]]]:
+        # Get the input and output model types from the generic parameters
+        type_args = getattr(self.__class__, "__orig_bases__", None)
         if (
-            not hasattr(self, "InputModel")
-            or not hasattr(self, "OutputModel")
-            or not hasattr(self, "lm")
+            type_args
+            and hasattr(type_args[0], "__args__")
+            and len(type_args[0].__args__) >= 2
         ):
-            raise NotImplementedError(
-                f"{self.__class__.__name__} must define lm, InputModel, and OutputModel as class attributes."
-            )
+            input_model = type_args[0].__args__[0]
+            output_model = type_args[0].__args__[1]
+        else:
+            # Fallback to the old way if type introspection fails
+            if (
+                not hasattr(self, "InputModel")
+                or not hasattr(self, "OutputModel")
+                or not hasattr(self, "lm")
+            ):
+                raise NotImplementedError(
+                    f"{self.__class__.__name__} must define lm, InputModel, and OutputModel as class attributes or use proper generic typing."
+                )
+            input_model = self.InputModel
+            output_model = self.OutputModel
 
-        item = data if isinstance(data, BaseModel) else self.InputModel(**data)
+        item = data if isinstance(data, BaseModel) else input_model(**data)
 
-        return await self.lm.parse(
+        result = await self.lm.parse(
             prompt=item.model_dump_json(),
             instruction=self.__doc__ or "",
-            response_model=self.OutputModel,
-            temperature=self.temperature,
+            response_model=output_model,
+            temperature=temperature or self.temperature,
             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)
+
     def generate_training_data(
         self, input_dict: Dict[str, Any], output: Dict[str, Any]
-    ):
+    ) -> Dict[str, Any]:
         """Return share gpt like format"""
         system_prompt = self.__doc__ or ""
         user_msg = self.InputModel(**input_dict).model_dump_json()  # type: ignore[attr-defined]
@@ -917,4 +961,5 @@ class AsyncLLMTask(ABC):
             system_msg=system_prompt, user_msg=user_msg, assistant_msg=assistant_msg
         )
         return {"messages": messages}
-    arun = __call__  # alias for compatibility with other LLMTask implementations
+
+    # arun = __call__  # alias for compatibility with other LLMTask implementations

llm_utils/lm/utils.py

@@ -7,11 +7,6 @@ import numpy as np
 from loguru import logger
 
 
-def _clear_port_use(ports):
-    for port in ports:
-        file_counter: str = f"/tmp/port_use_counter_{port}.npy"
-        if os.path.exists(file_counter):
-            os.remove(file_counter)
 
 
 def _atomic_save(array: np.ndarray, filename: str):

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

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

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

@@ -5,11 +5,11 @@ 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=_NmWEp_jCbD6soexXo489L40KS8xJPgtY5QxXLDYsis,34174
+llm_utils/lm/async_lm.py,sha256=xUOlAOivxrI-KRaUv1V0l5y7ajYTKTA6QZkqpK6Uue8,35847
 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
-llm_utils/lm/utils.py,sha256=GMvs64DRzVnXAki4SZ-A6mx2Fi9IgeF11BA-5FB-CYg,4777
+llm_utils/lm/utils.py,sha256=gUejbVZPYg97g4ftYEptYN52WhH3TAKOFW81sjLvi08,4585
 llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
 llm_utils/scripts/vllm_load_balancer.py,sha256=GjMdoZrdT9cSLos0qSdkLg2dwZgW1enAMsD3aTZAfNs,20845
 llm_utils/scripts/vllm_serve.py,sha256=4NaqpVs7LBvxtvTCMPsNCAOfqiWkKRttxWMmWY7SitA,14729
@@ -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.22.dist-info/METADATA,sha256=Ll1EUWmXjsgvn2nK2NZ-uSrf6SbkTY1mLaHRKWfgR2Q,7442
-speedy_utils-1.0.22.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.22.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
-speedy_utils-1.0.22.dist-info/RECORD,,
+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,,