speedy-utils 1.1.7__tar.gz → 1.1.9__tar.gz

{speedy_utils-1.1.7 → speedy_utils-1.1.9}/PKG-INFO

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

{speedy_utils-1.1.7 → speedy_utils-1.1.9}/pyproject.toml

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

{speedy_utils-1.1.7 → speedy_utils-1.1.9}/src/llm_utils/lm/async_lm/async_llm_task.py

@@ -7,6 +7,7 @@ import pathlib
 from abc import ABC
 from dataclasses import dataclass
 from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union, cast
+from venv import logger
 
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
@@ -507,3 +508,13 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
 
     # Compatibility alias for other LLMTask implementations
     arun = __call__
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if hasattr(self._lm, "_last_client"):
+            last_client = self._lm._last_client   # type: ignore
+            await last_client._client.aclose()
+        else:
+            logger.warning("No last client to close")

{speedy_utils-1.1.7 → speedy_utils-1.1.9}/src/llm_utils/lm/async_lm/async_lm_base.py

@@ -70,6 +70,7 @@ class AsyncLMBase:
             api_key=self.api_key,
             base_url=api_base,
         )
+        self._last_client = client
         return client
 
     # ------------------------------------------------------------------ #