speedy-utils 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

llm_utils/lm/text_lm.py

@@ -0,0 +1,130 @@
+from typing import Any, List, Optional, Union, Tuple
+from .base_lm import LM
+import random
+import logging
+import json
+from speedy_utils import identify_uuid
+
+logger = logging.getLogger(__name__)
+
+
+class TextLM(LM):
+    """
+    Language model that returns outputs as plain text (str).
+    """
+
+    def _generate_cache_key(
+        self,
+        messages: List[Any],
+        effective_kwargs: dict,
+    ) -> str:
+        """Generate a cache key based on input parameters."""
+        cache_key_list = self._generate_cache_key_base(messages, effective_kwargs)
+        return identify_uuid(str(cache_key_list))
+
+    def _check_cache(
+        self,
+        effective_cache: bool,
+        messages: List[Any],
+        effective_kwargs: dict,
+    ) -> Tuple[Optional[str], Optional[str]]:
+        """Check if result is in cache and return it if available."""
+        if not effective_cache:
+            return None, None
+
+        id_for_cache = self._generate_cache_key(messages, effective_kwargs)
+        cached_result = self.load_cache(id_for_cache)
+
+        if cached_result is not None:
+            if isinstance(cached_result, str):
+                return id_for_cache, cached_result
+            elif (
+                isinstance(cached_result, list)
+                and cached_result
+                and isinstance(cached_result[0], str)
+            ):
+                return id_for_cache, cached_result[0]
+            else:
+                logger.warning(
+                    f"Cached result for {id_for_cache} has unexpected type {type(cached_result)}. Ignoring cache."
+                )
+
+        return id_for_cache, None
+
+    def _parse_llm_output(self, llm_output: Any) -> str:
+        """Parse the LLM output into the correct format."""
+        if isinstance(llm_output, dict):
+            return json.dumps(llm_output)
+
+        if isinstance(llm_output, str):
+            return llm_output
+        else:
+            logger.warning(
+                f"LLM output type mismatch. Expected str, got {type(llm_output)}. Returning str."
+            )
+            return str(llm_output)
+
+    def _store_in_cache(
+        self, effective_cache: bool, id_for_cache: Optional[str], result: str
+    ):
+        """Store the result in cache if caching is enabled."""
+        self._store_in_cache_base(effective_cache, id_for_cache, result)
+
+    def forward_messages(
+        self,
+        messages: List[Any],
+        cache: Optional[bool] = None,
+        port: Optional[int] = None,
+        use_loadbalance: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> str:
+        # 1. Prepare inputs
+        effective_kwargs, effective_cache, current_port, dspy_main_input = (
+            self._prepare_call_inputs(
+                messages, max_tokens, port, use_loadbalance, cache, **kwargs
+            )
+        )
+
+        # 2. Check cache
+        cache_id, cached_result = self._check_cache(
+            effective_cache, messages, effective_kwargs
+        )
+        if cached_result:
+            return cached_result
+
+        # 3. Call LLM
+        llm_output = self._call_llm(
+            dspy_main_input, current_port, use_loadbalance, **effective_kwargs
+        )
+
+        # 4. Parse output
+        result: str = self._parse_llm_output(llm_output)
+
+        # 5. Store in cache
+        self._store_in_cache(effective_cache, cache_id, result)
+
+        return result
+
+    def _call_llm(
+        self,
+        dspy_main_input: List[Any],
+        current_port: Optional[int],
+        use_loadbalance: Optional[bool],
+        **kwargs,
+    ):
+        """Call the LLM and return the output as plain text (no retry)."""
+        # Use messages directly
+        messages = dspy_main_input
+
+        response = self.openai_client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            **kwargs,
+        )
+
+        # Update port usage stats if needed
+        if current_port and use_loadbalance is True:
+            self._update_port_use(current_port, -1)
+
+        return response.choices[0].message.content

llm_utils/lm/utils.py

@@ -0,0 +1,130 @@
+import fcntl
+import os
+import tempfile
+import time
+from typing import List, Dict
+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):
+    tmp_dir = os.path.dirname(filename) or "."
+    with tempfile.NamedTemporaryFile(dir=tmp_dir, delete=False) as tmp:
+        np.save(tmp, array)
+        temp_name = tmp.name
+    os.replace(temp_name, filename)
+
+
+def _update_port_use(port: int, increment: int) -> None:
+    file_counter: str = f"/tmp/port_use_counter_{port}.npy"
+    file_counter_lock: str = f"/tmp/port_use_counter_{port}.lock"
+    with open(file_counter_lock, "w") as lock_file:
+        fcntl.flock(lock_file, fcntl.LOCK_EX)
+        try:
+            if os.path.exists(file_counter):
+                try:
+                    counter = np.load(file_counter)
+                except Exception as e:
+                    logger.warning(f"Corrupted usage file {file_counter}: {e}")
+                    counter = np.array([0])
+            else:
+                counter: np.ndarray = np.array([0], dtype=np.int64)
+            counter[0] += increment
+            _atomic_save(counter, file_counter)
+        finally:
+            fcntl.flock(lock_file, fcntl.LOCK_UN)
+
+
+def _pick_least_used_port(ports: List[int]) -> int:
+    global_lock_file = "/tmp/ports.lock"
+    with open(global_lock_file, "w") as lock_file:
+        fcntl.flock(lock_file, fcntl.LOCK_EX)
+        try:
+            port_use: Dict[int, int] = {}
+            for port in ports:
+                file_counter = f"/tmp/port_use_counter_{port}.npy"
+                if os.path.exists(file_counter):
+                    try:
+                        counter = np.load(file_counter)
+                    except Exception as e:
+                        logger.warning(f"Corrupted usage file {file_counter}: {e}")
+                        counter = np.array([0])
+                else:
+                    counter = np.array([0])
+                port_use[port] = counter[0]
+            if not port_use:
+                if ports:
+                    raise ValueError("Port usage data is empty, cannot pick a port.")
+                else:
+                    raise ValueError("No ports provided to pick from.")
+            lsp = min(port_use, key=lambda k: port_use[k])
+            _update_port_use(lsp, 1)
+        finally:
+            fcntl.flock(lock_file, fcntl.LOCK_UN)
+    return lsp
+
+
+def retry_on_exception(max_retries=10, exceptions=(Exception,), sleep_time=3):
+    def decorator(func):
+        from functools import wraps
+
+        def wrapper(self, *args, **kwargs):
+            retry_count = kwargs.get("retry_count", 0)
+            last_exception = None
+            while retry_count <= max_retries:
+                try:
+                    return func(self, *args, **kwargs)
+                except exceptions as e:
+                    import litellm
+
+                    if isinstance(
+                        e, (litellm.exceptions.APIError, litellm.exceptions.Timeout)
+                    ):
+                        base_url_info = kwargs.get(
+                            "base_url", getattr(self, "base_url", None)
+                        )
+                        logger.warning(
+                            f"[{base_url_info=}] {type(e).__name__}: {str(e)[:100]}, will sleep for {sleep_time}s and retry"
+                        )
+                        time.sleep(sleep_time)
+                        retry_count += 1
+                        kwargs["retry_count"] = retry_count
+                        last_exception = e
+                        continue
+                    elif hasattr(
+                        litellm.exceptions, "ContextWindowExceededError"
+                    ) and isinstance(e, litellm.exceptions.ContextWindowExceededError):
+                        logger.error(f"Context window exceeded: {e}")
+                        raise
+                    else:
+                        logger.error(f"Generic error during LLM call: {e}")
+                        import traceback
+
+                        traceback.print_exc()
+                        raise
+            logger.error(f"Retry limit exceeded, error: {last_exception}")
+            if last_exception:
+                raise last_exception
+            raise ValueError("Retry limit exceeded with no specific error.")
+
+        return wraps(func)(wrapper)
+
+    return decorator
+
+
+def forward_only(func):
+    from functools import wraps
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        kwargs["retry_count"] = 0
+        return func(self, *args, **kwargs)
+
+    return wrapper

{speedy_utils-1.0.0.dist-info → speedy_utils-1.0.1.dist-info}/METADATA

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

{speedy_utils-1.0.0.dist-info → speedy_utils-1.0.1.dist-info}/RECORD

@@ -1,7 +1,12 @@
-llm_utils/__init__.py,sha256=4IO_cP7FlFvV57W7_AkJGdXWcZfWALkr6UwR9s9QEF4,701
+llm_utils/__init__.py,sha256=ujANp90z6cBoRAufbDbRs85KsDaghKc4-wzEugl6meg,717
 llm_utils/chat_format.py,sha256=ZY2HYv3FPL2xiMxbbO-huIwT5LZrcJm_if_us-2eSZ4,15094
 llm_utils/group_messages.py,sha256=GKMQkenQf-6DD_1EJa11UBj7-VfkGT7xVhR_B_zMzqY,3868
-llm_utils/lm.py,sha256=3d3b9UMtIeKcG6vpXHwuDZ3QP46X1aSGeMvayy3tmHs,29044
+llm_utils/lm/__init__.py,sha256=UZ5Ij-d5Xf64luuf1aAiDE1CFRoqEgoCU8xWuuLMBis,226
+llm_utils/lm/base_lm.py,sha256=daT9mueYSZtHE0QMmwRQT_kvjC82KJmZ-u6_1gOIsrw,12560
+llm_utils/lm/chat_session.py,sha256=rkjj1caGu03o7I3d8q_HJzHA9CVzV-Q-mz6iWgqHno0,3301
+llm_utils/lm/pydantic_lm.py,sha256=f-eFFCoCwbPUTZAoS3s217upzUSUHxu-kKlZz6Jl5Ck,6628
+llm_utils/lm/text_lm.py,sha256=xt41ZKDFjIfAe0e2rdvTWePNq5tdvdPzW64UJ8thuXw,4069
+llm_utils/lm/utils.py,sha256=-fDNueiXKQI6RDoNHJYNyORomf2XlCf2doJZ3GEV2Io,4762
 llm_utils/lm_classification.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 llm_utils/load_chat_dataset.py,sha256=hsPPlOmZEDqsg7GQD7SgxTEGJky6JDm6jnG3JHbpjb4,1895
 llm_utils/scripts/vllm_load_balancer.py,sha256=uSjGd_jOmI9W9eVOhiOXbeUnZkQq9xG4bCVzhmpupcA,16096
@@ -21,7 +26,7 @@ speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 speedy_utils/multi_worker/process.py,sha256=XwQlffxzRFnCVeKjDNBZDwFfUQHiJiuFA12MRGJVru8,6708
 speedy_utils/multi_worker/thread.py,sha256=9pXjvgjD0s0Hp0cZ6I3M0ndp1OlYZ1yvqbs_bcun_Kw,12775
 speedy_utils/scripts/mpython.py,sha256=ZzkBWI5Xw3vPoMx8xQt2x4mOFRjtwWqfvAJ5_ngyWgw,3816
-speedy_utils-1.0.0.dist-info/METADATA,sha256=nIcZcNBPA9_eUs7on52UMVtcV43lmVFRJLHLptk5AA0,7165
-speedy_utils-1.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.0.dist-info/entry_points.txt,sha256=fsv8_lMg62BeswoUHrqfj2u6q2l4YcDCw7AgQFg6GRw,61
-speedy_utils-1.0.0.dist-info/RECORD,,
+speedy_utils-1.0.1.dist-info/METADATA,sha256=U-8IfrZTleZAsydTg-22oTyjpzsW0pKTPV3P7Jp2WVw,7165
+speedy_utils-1.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+speedy_utils-1.0.1.dist-info/entry_points.txt,sha256=fsv8_lMg62BeswoUHrqfj2u6q2l4YcDCw7AgQFg6GRw,61
+speedy_utils-1.0.1.dist-info/RECORD,,