speedy-utils 1.0.23__py3-none-any.whl → 1.1.0__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,
@@ -92,9 +93,10 @@ from typing import (
     cast,
     overload,
 )
-
+from typing_extensions import TypedDict
 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
 
@@ -144,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."""
 
@@ -460,10 +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,
         **kwargs,
-    ):
+    ) -> ParsedOutput:  # -> dict[str, Any]:
         """Parse response using guided JSON generation."""
         if messages is None:
             assert instruction is not None, "Instruction must be provided."
@@ -514,27 +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]
-        if return_openai_response:
-            return completion
-        return self._parse_complete_output(completion, response_model)
+
+        output = 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]
@@ -839,8 +846,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.
 
@@ -878,37 +888,71 @@ class AsyncLLMTask(ABC):
     """
 
     lm: "AsyncLM"
-    InputModel: Type[BaseModel]
-    OutputModel: Type[BaseModel]
+    InputModel: InputModelType
+    OutputModel: OutputModelType
 
     temperature: float = 0.6
     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,
+    ) -> 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
+
+        # 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")
 
-        item = data if isinstance(data, BaseModel) else self.InputModel(**data)
+        assert isinstance(output_model, type) and issubclass(output_model, BaseModel), (
+            "OutputModel must be a subclass of BaseModel"
+        )
 
-        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 (
+            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]
-    ):
+    ) -> 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.23.dist-info → speedy_utils-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.0.23
+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.23.dist-info → speedy_utils-1.1.0.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=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
-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.23.dist-info/METADATA,sha256=E2NtrXhJt45XHFz5cv9BuxmdlPHqSZwHdq2vHJG7xqk,7442
-speedy_utils-1.0.23.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.23.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
-speedy_utils-1.0.23.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,,