speedy-utils 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

llm_utils/lm/base_lm.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import os
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+    cast,
+)
+
+from httpx import URL
+from loguru import logger
+from openai import OpenAI, AuthenticationError, RateLimitError
+from openai.pagination import SyncPage
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.parsed_chat_completion import ParsedChatCompletion
+from openai.types.model import Model
+from pydantic import BaseModel
+import warnings
+
+# --------------------------------------------------------------------------- #
+# type helpers
+# --------------------------------------------------------------------------- #
+TModel = TypeVar("TModel", bound=BaseModel)
+Messages = List[ChatCompletionMessageParam]  # final, already-typed messages
+LegacyMsgs = List[Dict[str, str]]  # old “…role/content…” dicts
+RawMsgs = Union[Messages, LegacyMsgs]  # what __call__ accepts
+
+
+class LM:
+    """
+    Unified language-model wrapper.
+
+    • `response_format=str`               → returns `str`
+    • `response_format=YourPydanticModel` → returns that model instance
+    """
+
+    # --------------------------------------------------------------------- #
+    # ctor / plumbing
+    # --------------------------------------------------------------------- #
+    def __init__(
+        self,
+        model: str | None = None,
+        *,
+        temperature: float = 0.0,
+        max_tokens: int = 2_000,
+        host: str = "localhost",
+        port: Optional[int] = None,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        cache: bool = True,
+        **openai_kwargs: Any,
+    ) -> None:
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.base_url = base_url or (f"http://{host}:{port}/v1" if port else None)
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY", "abc")
+        self.openai_kwargs = openai_kwargs
+        self.do_cache = cache
+
+        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    def set_model(self, model: str) -> None:
+        """Set the model name after initialization."""
+        self.model = model
+
+    # --------------------------------------------------------------------- #
+    # public API – typed overloads
+    # --------------------------------------------------------------------- #
+    @overload
+    def __call__(
+        self,
+        *,
+        prompt: str | None = ...,
+        messages: RawMsgs | None = ...,
+        response_format: type[str] = str,
+        **kwargs: Any,
+    ) -> str: ...
+
+    @overload
+    def __call__(
+        self,
+        *,
+        prompt: str | None = ...,
+        messages: RawMsgs | None = ...,
+        response_format: Type[TModel],
+        **kwargs: Any,
+    ) -> TModel: ...
+
+    # single implementation
+    def __call__(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[RawMsgs] = None,
+        response_format: Union[type[str], Type[BaseModel]] = str,
+        cache: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        # argument validation ------------------------------------------------
+        if (prompt is None) == (messages is None):
+            raise ValueError("Provide *either* `prompt` or `messages` (but not both).")
+
+        if prompt is not None:
+            messages = [{"role": "user", "content": prompt}]
+
+        assert messages is not None  # for type-checker
+        assert self.model is not None, "Model must be set before calling."
+        openai_msgs: Messages = (
+            self._convert_messages(cast(LegacyMsgs, messages))
+            if isinstance(messages[0], dict)  # legacy style
+            else cast(Messages, messages)  # already typed
+        )
+
+        kw = dict(
+            self.openai_kwargs,
+            temperature=self.temperature,
+            max_tokens=max_tokens or self.max_tokens,
+            **kwargs,
+        )
+        use_cache = self.do_cache if cache is None else cache
+
+        raw = self._call_raw(
+            openai_msgs,
+            response_format=response_format,
+            use_cache=use_cache,
+            **kw,
+        )
+        return self._parse_output(raw, response_format)
+
+    # --------------------------------------------------------------------- #
+    # low-level OpenAI call
+    # --------------------------------------------------------------------- #
+    def _call_raw(
+        self,
+        messages: Sequence[ChatCompletionMessageParam],
+        response_format: Union[type[str], Type[BaseModel]],
+        use_cache: bool,
+        **kw: Any,
+    ):
+        assert self.model is not None, "Model must be set before making a call."
+        model: str = self.model
+        cache_key = (
+            self._cache_key(messages, kw, response_format) if use_cache else None
+        )
+        if cache_key and (hit := self._load_cache(cache_key)) is not None:
+            return hit
+
+        try:
+            # structured mode
+            if response_format is not str and issubclass(response_format, BaseModel):
+                rsp: ParsedChatCompletion[BaseModel] = (
+                    self.client.beta.chat.completions.parse(
+                        model=model,
+                        messages=list(messages),
+                        response_format=response_format,  # type: ignore[arg-type]
+                        **kw,
+                    )
+                )
+                result: Any = rsp.choices[0].message.parsed  # already a model
+            # plain-text mode
+            else:
+                rsp = self.client.chat.completions.create(
+                    model=model,
+                    messages=list(messages),
+                    **kw,
+                )
+                result = rsp.choices[0].message.content  # str
+        except (AuthenticationError, RateLimitError) as exc:  # pragma: no cover
+            logger.error(exc)
+            raise
+
+        if cache_key:
+            self._dump_cache(cache_key, result)
+
+        return result
+
+    # --------------------------------------------------------------------- #
+    # legacy → typed messages
+    # --------------------------------------------------------------------- #
+    @staticmethod
+    def _convert_messages(msgs: LegacyMsgs) -> Messages:
+        converted: Messages = []
+        for msg in msgs:
+            role = msg["role"]
+            content = msg["content"]
+            if role == "user":
+                converted.append(
+                    ChatCompletionUserMessageParam(role="user", content=content)
+                )
+            elif role == "assistant":
+                converted.append(
+                    ChatCompletionAssistantMessageParam(
+                        role="assistant", content=content
+                    )
+                )
+            elif role == "system":
+                converted.append(
+                    ChatCompletionSystemMessageParam(role="system", content=content)
+                )
+            elif role == "tool":
+                converted.append(
+                    ChatCompletionToolMessageParam(
+                        role="tool",
+                        content=content,
+                        tool_call_id=msg.get("tool_call_id") or "",  # str, never None
+                    )
+                )
+            else:
+                # fall back to raw dict for unknown roles
+                converted.append({"role": role, "content": content})  # type: ignore[arg-type]
+        return converted
+
+    # --------------------------------------------------------------------- #
+    # final parse (needed for plain-text or cache hits only)
+    # --------------------------------------------------------------------- #
+    @staticmethod
+    def _parse_output(
+        raw: Any,
+        response_format: Union[type[str], Type[BaseModel]],
+    ) -> str | BaseModel:
+        if response_format is str:
+            return cast(str, raw)
+
+        # For the type-checker: we *know* it's a BaseModel subclass here.
+        model_cls = cast(Type[BaseModel], response_format)
+
+        if isinstance(raw, model_cls):
+            return raw
+        if isinstance(raw, dict):
+            return model_cls.model_validate(raw)
+        try:
+            data = json.loads(raw)
+        except Exception as exc:  # noqa: BLE001
+            raise ValueError(f"Model did not return JSON:\n---\n{raw}") from exc
+        return model_cls.model_validate(data)
+
+    # --------------------------------------------------------------------- #
+    # tiny disk cache
+    # --------------------------------------------------------------------- #
+    @staticmethod
+    def _cache_key(
+        messages: Any, kw: Any, response_format: Union[type[str], Type[BaseModel]]
+    ) -> str:
+        tag = response_format.__name__ if response_format is not str else "text"
+        blob = json.dumps([messages, kw, tag], sort_keys=True).encode()
+        return base64.urlsafe_b64encode(hashlib.sha256(blob).digest()).decode()[:22]
+
+    @staticmethod
+    def _cache_path(key: str) -> str:
+        return os.path.expanduser(f"~/.cache/lm/{key}.json")
+
+    def _dump_cache(self, key: str, val: Any) -> None:
+        try:
+            path = self._cache_path(key)
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            with open(path, "w") as fh:
+                if isinstance(val, BaseModel):
+                    json.dump(val.model_dump(mode="json"), fh)
+                else:
+                    json.dump(val, fh)
+        except Exception as exc:  # pragma: no cover
+            logger.debug(f"cache write skipped: {exc}")
+
+    def _load_cache(self, key: str) -> Any | None:
+        path = self._cache_path(key)
+        if not os.path.exists(path):
+            return None
+        try:
+            with open(path) as fh:
+                return json.load(fh)
+        except Exception:  # pragma: no cover
+            return None
+
+    @staticmethod
+    def list_models(port=None) -> List[str]:
+        """
+        List available models.
+        """
+        try:
+            client: OpenAI = LM(port=port).client
+            base_url: URL = client.base_url
+            logger.debug(f"Base URL: {base_url}")
+            models: SyncPage[Model] = client.models.list()
+            return [model.id for model in models.data]
+        except Exception as exc:
+            logger.error(f"Failed to list models: {exc}")
+            return []

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