speedy-utils 1.1.17__py3-none-any.whl → 1.1.18__py3-none-any.whl

llm_utils/lm/lm.py

@@ -0,0 +1,207 @@
+# # from ._utils import *
+# from typing import (
+#     Any,
+#     List,
+#     Literal,
+#     Optional,
+#     Type,
+#     Union,
+#     cast,
+# )
+
+# from loguru import logger
+# from openai import AuthenticationError, BadRequestError, OpenAI, RateLimitError
+# from pydantic import BaseModel
+# from speedy_utils import jloads
+
+# # from llm_utils.lm.async_lm.async_llm_task import OutputModelType
+# from llm_utils.lm.lm_base import LMBase
+
+# from .async_lm._utils import (
+#     LegacyMsgs,
+#     Messages,
+#     OutputModelType,
+#     ParsedOutput,
+#     RawMsgs,
+# )
+
+
+# class LM(LMBase):
+#     """Unified **sync** language‑model wrapper with optional JSON parsing."""
+
+#     def __init__(
+#         self,
+#         *,
+#         model: Optional[str] = None,
+#         response_model: Optional[type[BaseModel]] = None,
+#         temperature: float = 0.0,
+#         max_tokens: int = 2_000,
+#         base_url: Optional[str] = None,
+#         api_key: Optional[str] = None,
+#         cache: bool = True,
+#         ports: Optional[List[int]] = None,
+#         top_p: float = 1.0,
+#         presence_penalty: float = 0.0,
+#         top_k: int = 1,
+#         repetition_penalty: float = 1.0,
+#         frequency_penalty: Optional[float] = None,
+#     ) -> None:
+
+#         if model is None:
+#             if base_url is None:
+#                 raise ValueError("Either model or base_url must be provided")
+#             models = OpenAI(base_url=base_url, api_key=api_key or 'abc').models.list().data
+#             assert len(models) == 1, f"Found {len(models)} models, please specify one."
+#             model = models[0].id
+#             print(f"Using model: {model}")
+
+#         super().__init__(
+#             ports=ports,
+#             base_url=base_url,
+#             cache=cache,
+#             api_key=api_key,
+#         )
+
+#         # Model behavior options
+#         self.response_model = response_model
+
+#         # Store all model-related parameters in model_kwargs
+#         self.model_kwargs = dict(
+#             model=model,
+#             temperature=temperature,
+#             max_tokens=max_tokens,
+#             top_p=top_p,
+#             presence_penalty=presence_penalty,
+#         )
+#         self.extra_body = dict(
+#             top_k=top_k,
+#             repetition_penalty=repetition_penalty,
+#             frequency_penalty=frequency_penalty,
+#         )
+
+#     def _unified_client_call(
+#         self,
+#         messages: RawMsgs,
+#         extra_body: Optional[dict] = None,
+#         max_tokens: Optional[int] = None,
+#     ) -> dict:
+#         """Unified method for all client interactions (caching handled by MOpenAI)."""
+#         converted_messages: Messages = (
+#             self._convert_messages(cast(LegacyMsgs, messages))
+#             if messages and isinstance(messages[0], dict)
+#             else cast(Messages, messages)
+#         )
+#         if max_tokens is not None:
+#             self.model_kwargs["max_tokens"] = max_tokens
+
+#         try:
+#             # Get completion from API (caching handled by MOpenAI)
+#             call_kwargs = {
+#                 "messages": converted_messages,
+#                 **self.model_kwargs,
+#             }
+#             if extra_body:
+#                 call_kwargs["extra_body"] = extra_body
+
+#             completion = self.client.chat.completions.create(**call_kwargs)
+
+#             if hasattr(completion, "model_dump"):
+#                 completion = completion.model_dump()
+
+#         except (AuthenticationError, RateLimitError, BadRequestError) as exc:
+#             error_msg = f"OpenAI API error ({type(exc).__name__}): {exc}"
+#             logger.error(error_msg)
+#             raise
+
+#         return completion
+
+#     def __call__(
+#         self,
+#         prompt: Optional[str] = None,
+#         messages: Optional[RawMsgs] = None,
+#         max_tokens: Optional[int] = None,
+#     ):  # -> tuple[Any | dict[Any, Any], list[ChatCompletionMessagePar...:# -> tuple[Any | dict[Any, Any], list[ChatCompletionMessagePar...:
+#         """Unified sync call for language model, returns (assistant_message.model_dump(), messages)."""
+#         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
+
+#         openai_msgs: Messages = (
+#             self._convert_messages(cast(LegacyMsgs, messages))
+#             if isinstance(messages[0], dict)
+#             else cast(Messages, messages)
+#         )
+
+#         assert self.model_kwargs["model"] is not None, (
+#             "Model must be set before making a call."
+#         )
+
+#         # Use unified client call
+#         raw_response = self._unified_client_call(
+#             list(openai_msgs), max_tokens=max_tokens
+#         )
+
+#         if hasattr(raw_response, "model_dump"):
+#             raw_response = raw_response.model_dump()  # type: ignore
+
+#         # Extract the assistant's message
+#         assistant_msg = raw_response["choices"][0]["message"]
+#         # Build the full messages list (input + assistant reply)
+#         full_messages = list(messages) + [
+#             {"role": assistant_msg["role"], "content": assistant_msg["content"]}
+#         ]
+#         # Return the OpenAI message as model_dump (if available) and the messages list
+#         if hasattr(assistant_msg, "model_dump"):
+#             msg_dump = assistant_msg.model_dump()
+#         else:
+#             msg_dump = dict(assistant_msg)
+#         return msg_dump, full_messages
+
+#     def parse(
+#         self,
+#         messages: Messages,
+#         response_model: Optional[type[BaseModel]] = None,
+#     ) -> ParsedOutput[BaseModel]:
+#         """Parse response using OpenAI's native parse API."""
+#         # Use provided response_model or fall back to instance default
+#         model_to_use = response_model or self.response_model
+#         assert model_to_use is not None, "response_model must be provided or set at init."
+
+#         # Use OpenAI's native parse API directly
+#         response = self.client.chat.completions.parse(
+#             model=self.model_kwargs["model"],
+#             messages=messages,
+#             response_format=model_to_use,
+#             **{k: v for k, v in self.model_kwargs.items() if k != "model"}
+#         )
+        
+#         parsed = response.choices[0].message.parsed
+#         completion = response.model_dump() if hasattr(response, "model_dump") else {}
+#         full_messages = list(messages) + [
+#             {"role": "assistant", "content": parsed}
+#         ]
+
+#         return ParsedOutput(
+#             messages=full_messages,
+#             parsed=cast(BaseModel, parsed),
+#             completion=completion,
+#             model_kwargs=self.model_kwargs,
+#         )
+
+
+
+#     def __enter__(self):
+#         return self
+
+#     def __exit__(self, exc_type, exc_val, exc_tb):
+#         if hasattr(self, "_last_client"):
+#             last_client = self._last_client  # type: ignore
+#             if hasattr(last_client, "close"):
+#                 last_client.close()
+#         else:
+#             logger.warning("No last client to close")
+LM = None

llm_utils/lm/lm_base.py

@@ -0,0 +1,285 @@
+# from ._utils import *
+import json
+import os
+from typing import (
+    Any,
+    List,
+    Optional,
+    Type,
+    Union,
+    cast,
+    overload,
+)
+
+from httpx import URL
+from loguru import logger
+from openai import OpenAI
+from openai.pagination import SyncPage
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.model import Model
+from pydantic import BaseModel
+
+from llm_utils.lm.openai_memoize import MOpenAI
+
+from .async_lm._utils import (
+    LegacyMsgs,
+    Messages,
+    RawMsgs,
+    TModel,
+)
+
+
+class LMBase:
+    """Unified **sync** language‑model wrapper with optional JSON parsing."""
+
+    def __init__(
+        self,
+        *,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        cache: bool = True,
+        ports: Optional[List[int]] = None,
+    ) -> None:
+        self.base_url = base_url
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY", "abc")
+        self._cache = cache
+        self.ports = ports
+
+    @property
+    def client(self) -> MOpenAI:
+        # if have multiple ports
+        if self.ports and self.base_url:
+            import random
+            import re
+            
+            port = random.choice(self.ports)
+            # Replace port in base_url if it exists
+            base_url_pattern = r'(https?://[^:/]+):?\d*(/.*)?'
+            match = re.match(base_url_pattern, self.base_url)
+            if match:
+                host_part = match.group(1)
+                path_part = match.group(2) or '/v1'
+                api_base = f"{host_part}:{port}{path_part}"
+            else:
+                api_base = self.base_url
+            logger.debug(f"Using port: {port}")
+        else:
+            api_base = self.base_url
+            
+        if api_base is None:
+            raise ValueError("base_url must be provided")
+            
+        client = MOpenAI(
+            api_key=self.api_key,
+            base_url=api_base,
+            cache=self._cache,
+        )
+        self._last_client = client
+        return client
+
+    # ------------------------------------------------------------------ #
+    # Public API – typed overloads
+    # ------------------------------------------------------------------ #
+    @overload
+    def __call__(  # type: ignore
+        self,
+        *,
+        prompt: Optional[str] = ...,
+        messages: Optional[RawMsgs] = ...,
+        response_format: type[str] = str,
+        return_openai_response: bool = ...,
+        **kwargs: Any,
+    ) -> str: ...
+
+    @overload
+    def __call__(
+        self,
+        *,
+        prompt: Optional[str] = ...,
+        messages: Optional[RawMsgs] = ...,
+        response_format: Type[TModel],
+        return_openai_response: bool = ...,
+        **kwargs: Any,
+    ) -> TModel: ...
+
+    # ------------------------------------------------------------------ #
+    # Utilities below are unchanged (sync I/O is acceptable)
+    # ------------------------------------------------------------------ #
+    @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 "",
+                    )
+                )
+            else:
+                converted.append({"role": role, "content": content})  # type: ignore[arg-type]
+        return converted
+
+    @staticmethod
+    def _parse_output(
+        raw_response: Any, response_format: Union[type[str], Type[BaseModel]]
+    ) -> Union[str, BaseModel]:
+        if hasattr(raw_response, "model_dump"):
+            raw_response = raw_response.model_dump()
+
+        if response_format is str:
+            if isinstance(raw_response, dict) and "choices" in raw_response:
+                message = raw_response["choices"][0]["message"]
+                return message.get("content", "") or ""
+            return cast(str, raw_response)
+
+        model_cls = cast(Type[BaseModel], response_format)
+
+        if isinstance(raw_response, dict) and "choices" in raw_response:
+            message = raw_response["choices"][0]["message"]
+            if "parsed" in message:
+                return model_cls.model_validate(message["parsed"])
+            content = message.get("content")
+            if content is None:
+                raise ValueError("Model returned empty content")
+            try:
+                data = json.loads(content)
+                return model_cls.model_validate(data)
+            except Exception as exc:
+                raise ValueError(
+                    f"Failed to parse model output as JSON:\n{content}"
+                ) from exc
+
+        if isinstance(raw_response, model_cls):
+            return raw_response
+        if isinstance(raw_response, dict):
+            return model_cls.model_validate(raw_response)
+
+        try:
+            data = json.loads(raw_response)
+            return model_cls.model_validate(data)
+        except Exception as exc:
+            raise ValueError(
+                f"Model did not return valid JSON:\n---\n{raw_response}"
+            ) from exc
+
+    # ------------------------------------------------------------------ #
+    # Misc helpers
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def list_models(base_url: Optional[str] = None) -> List[str]:
+        try:
+            if base_url is None:
+                raise ValueError("base_url must be provided")
+            client = LMBase(base_url=base_url).client
+            base_url_obj: URL = client.base_url
+            logger.debug(f"Base URL: {base_url_obj}")
+            models: SyncPage[Model] = client.models.list()  # type: ignore[assignment]
+            return [model.id for model in models.data]
+        except Exception as exc:
+            logger.error(f"Failed to list models: {exc}")
+            return []
+
+    def build_system_prompt(
+        self,
+        response_model,
+        add_json_schema_to_instruction,
+        json_schema,
+        system_content,
+        think,
+    ):
+        if add_json_schema_to_instruction and response_model:
+            schema_block = f"\n\n<output_json_schema>\n{json.dumps(json_schema, indent=2)}\n</output_json_schema>"
+            # if schema_block not in system_content:
+            if "<output_json_schema>" in system_content:
+                # remove exsting schema block
+                import re  # replace
+
+                system_content = re.sub(
+                    r"<output_json_schema>.*?</output_json_schema>",
+                    "",
+                    system_content,
+                    flags=re.DOTALL,
+                )
+                system_content = system_content.strip()
+            system_content += schema_block
+
+        if think is True:
+            if "/think" in system_content:
+                pass
+            elif "/no_think" in system_content:
+                system_content = system_content.replace("/no_think", "/think")
+            else:
+                system_content += "\n\n/think"
+        elif think is False:
+            if "/no_think" in system_content:
+                pass
+            elif "/think" in system_content:
+                system_content = system_content.replace("/think", "/no_think")
+            else:
+                system_content += "\n\n/no_think"
+        return system_content
+
+    def inspect_history(self):
+        """Inspect the history of the LLM calls."""
+        pass
+    
+
+def get_model_name(client: OpenAI|str|int) -> str:
+    """
+    Get the first available model name from the client.
+    
+    Args:
+        client: OpenAI client, base_url string, or port number
+        
+    Returns:
+        Name of the first available model
+        
+    Raises:
+        ValueError: If no models are available or client is invalid
+    """
+    try:
+        if isinstance(client, OpenAI):
+            openai_client = client
+        elif isinstance(client, str):
+            # String base_url
+            openai_client = OpenAI(base_url=client, api_key='abc')
+        elif isinstance(client, int):
+            # Port number
+            base_url = f"http://localhost:{client}/v1"
+            openai_client = OpenAI(base_url=base_url, api_key='abc')
+        else:
+            raise ValueError(f"Unsupported client type: {type(client)}")
+            
+        models = openai_client.models.list()
+        if not models.data:
+            raise ValueError("No models available")
+            
+        return models.data[0].id
+        
+    except Exception as exc:
+        logger.error(f"Failed to get model name: {exc}")
+        raise ValueError(f"Could not retrieve model name: {exc}") from exc