speedy-utils 1.1.26__py3-none-any.whl → 1.1.28__py3-none-any.whl

llm_utils/lm/async_lm/async_lm_base.py

@@ -40,23 +40,23 @@ class AsyncLMBase:
     def __init__(
         self,
         *,
-        base_url: Optional[str] = None,
-        api_key: Optional[str] = None,
+        base_url: str | None = None,
+        api_key: str | None = None,
         cache: bool = True,
-        ports: Optional[List[int]] = None,
+        ports: list[int] | None = None,
     ) -> None:
         self.base_url = base_url
-        self.api_key = api_key or os.getenv("OPENAI_API_KEY", "abc")
+        self.api_key = api_key or os.getenv('OPENAI_API_KEY', 'abc')
         self._cache = cache
         self.ports = ports
 
     @property
-    def client(self) -> MAsyncOpenAI:
+    def client(self) -> 'MAsyncOpenAI': # type: ignore
         # 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*(/.*)?'
@@ -64,16 +64,16 @@ class AsyncLMBase:
             if match:
                 host_part = match.group(1)
                 path_part = match.group(2) or '/v1'
-                api_base = f"{host_part}:{port}{path_part}"
+                api_base = f'{host_part}:{port}{path_part}'
             else:
                 api_base = self.base_url
-            logger.debug(f"Using port: {port}")
+            logger.debug(f'Using port: {port}')
         else:
             api_base = self.base_url
-            
+
         if api_base is None:
-            raise ValueError("base_url must be provided")
-            
+            raise ValueError('base_url must be provided')
+
         client = MAsyncOpenAI(
             api_key=self.api_key,
             base_url=api_base,
@@ -89,8 +89,8 @@ class AsyncLMBase:
     async def __call__(  # type: ignore
         self,
         *,
-        prompt: Optional[str] = ...,
-        messages: Optional[RawMsgs] = ...,
+        prompt: str | None = ...,
+        messages: RawMsgs | None = ...,
         response_format: type[str] = str,
         return_openai_response: bool = ...,
         **kwargs: Any,
@@ -100,9 +100,9 @@ class AsyncLMBase:
     async def __call__(
         self,
         *,
-        prompt: Optional[str] = ...,
-        messages: Optional[RawMsgs] = ...,
-        response_format: Type[TModel],
+        prompt: str | None = ...,
+        messages: RawMsgs | None = ...,
+        response_format: type[TModel],
         return_openai_response: bool = ...,
         **kwargs: Any,
     ) -> TModel: ...
@@ -114,62 +114,62 @@ class AsyncLMBase:
     def _convert_messages(msgs: LegacyMsgs) -> Messages:
         converted: Messages = []
         for msg in msgs:
-            role = msg["role"]
-            content = msg["content"]
-            if role == "user":
+            role = msg['role']
+            content = msg['content']
+            if role == 'user':
                 converted.append(
-                    ChatCompletionUserMessageParam(role="user", content=content)
+                    ChatCompletionUserMessageParam(role='user', content=content)
                 )
-            elif role == "assistant":
+            elif role == 'assistant':
                 converted.append(
                     ChatCompletionAssistantMessageParam(
-                        role="assistant", content=content
+                        role='assistant', content=content
                     )
                 )
-            elif role == "system":
+            elif role == 'system':
                 converted.append(
-                    ChatCompletionSystemMessageParam(role="system", content=content)
+                    ChatCompletionSystemMessageParam(role='system', content=content)
                 )
-            elif role == "tool":
+            elif role == 'tool':
                 converted.append(
                     ChatCompletionToolMessageParam(
-                        role="tool",
+                        role='tool',
                         content=content,
-                        tool_call_id=msg.get("tool_call_id") or "",
+                        tool_call_id=msg.get('tool_call_id') or '',
                     )
                 )
             else:
-                converted.append({"role": role, "content": content})  # type: ignore[arg-type]
+                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: Any, response_format: type[str] | type[BaseModel]
+    ) -> 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 ""
+            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)
+        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 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")
+                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}"
+                    f'Failed to parse model output as JSON:\n{content}'
                 ) from exc
 
         if isinstance(raw_response, model_cls):
@@ -182,7 +182,7 @@ class AsyncLMBase:
             return model_cls.model_validate(data)
         except Exception as exc:
             raise ValueError(
-                f"Model did not return valid JSON:\n---\n{raw_response}"
+                f'Model did not return valid JSON:\n---\n{raw_response}'
             ) from exc
 
     # ------------------------------------------------------------------ #
@@ -190,17 +190,17 @@ class AsyncLMBase:
     # ------------------------------------------------------------------ #
 
     @staticmethod
-    async def list_models(base_url: Optional[str] = None) -> List[str]:
+    async def list_models(base_url: str | None = None) -> list[str]:
         try:
             if base_url is None:
-                raise ValueError("base_url must be provided")
+                raise ValueError('base_url must be provided')
             client = AsyncLMBase(base_url=base_url).client
             base_url_obj: URL = client.base_url
-            logger.debug(f"Base URL: {base_url_obj}")
+            logger.debug(f'Base URL: {base_url_obj}')
             models: AsyncSyncPage[Model] = await 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}")
+            logger.error(f'Failed to list models: {exc}')
             return []
 
     def build_system_prompt(
@@ -212,15 +212,15 @@ class AsyncLMBase:
         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>"
+            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:
+            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>",
-                    "",
+                    r'<output_json_schema>.*?</output_json_schema>',
+                    '',
                     system_content,
                     flags=re.DOTALL,
                 )
@@ -228,21 +228,20 @@ class AsyncLMBase:
             system_content += schema_block
 
         if think is True:
-            if "/think" in system_content:
+            if '/think' in system_content:
                 pass
-            elif "/no_think" in system_content:
-                system_content = system_content.replace("/no_think", "/think")
+            elif '/no_think' in system_content:
+                system_content = system_content.replace('/no_think', '/think')
             else:
-                system_content += "\n\n/think"
+                system_content += '\n\n/think'
         elif think is False:
-            if "/no_think" in system_content:
+            if '/no_think' in system_content:
                 pass
-            elif "/think" in system_content:
-                system_content = system_content.replace("/think", "/no_think")
+            elif '/think' in system_content:
+                system_content = system_content.replace('/think', '/no_think')
             else:
-                system_content += "\n\n/no_think"
+                system_content += '\n\n/no_think'
         return system_content
 
     async def inspect_history(self):
         """Inspect the history of the LLM calls."""
-        pass

llm_utils/lm/async_lm/lm_specific.py

@@ -2,6 +2,7 @@ from typing import List
 
 from .async_lm import AsyncLM
 
+
 KNOWN_CONFIG = {
     # Qwen3 family (see model card "Best Practices" section)
     "qwen3-think": {
@@ -83,7 +84,7 @@ KNOWN_CONFIG = {
     # Add more as needed...
 }
 
-KNOWN_KEYS: List[str] = list(KNOWN_CONFIG.keys())
+KNOWN_KEYS: list[str] = list(KNOWN_CONFIG.keys())
 
 
 class AsyncLMQwenThink(AsyncLM):
@@ -107,7 +108,7 @@ class AsyncLMQwenThink(AsyncLM):
             top_k=top_k,
             presence_penalty=presence_penalty,
             **other_kwargs,
-            think=True
+            think=True,
         )
 
 
@@ -132,5 +133,5 @@ class AsyncLMQwenNoThink(AsyncLM):
             top_k=top_k,
             presence_penalty=presence_penalty,
             **other_kwargs,
-            think=False
+            think=False,
         )

llm_utils/lm/base_prompt_builder.py

@@ -4,22 +4,23 @@
 Simplified LLM Task module for handling language model interactions with structured input/output.
 """
 
-from typing import Any, Dict, List, Optional, Type, Union, cast
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 from openai import OpenAI
 from openai.types.chat import ChatCompletionMessageParam
-from pydantic import BaseModel
-from pydantic import create_model
-from typing import Callable, Tuple
-from abc import ABC, abstractmethod
+from pydantic import BaseModel, create_model
+
 
 # Type aliases for better readability
-Messages = List[ChatCompletionMessageParam]
+Messages = list[ChatCompletionMessageParam]
 
 import json
-from typing import Type, TypeVar
+from typing import TypeVar
+
 
-B = TypeVar("B", bound="BasePromptBuilder")
+B = TypeVar('B', bound='BasePromptBuilder')
 
 
 class BasePromptBuilder(BaseModel, ABC):
@@ -40,7 +41,7 @@ class BasePromptBuilder(BaseModel, ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def get_io_keys(self) -> Tuple[List[str], List[Union[str, Tuple[str, str]]]]:
+    def get_io_keys(self) -> tuple[list[str], list[str | tuple[str, str]]]:
         """
         Return (input_keys, output_keys).
         Each key must match a field of the subclass.
@@ -54,8 +55,10 @@ class BasePromptBuilder(BaseModel, ABC):
     # ------------------------------------------------------------------ #
     # Auto-build models from keys
     # ------------------------------------------------------------------ #
-    def _build_model_from_keys(self, keys: Union[List[str], List[Union[str, Tuple[str, str]]]], name: str) -> Type[BaseModel]:
-        fields: Dict[str, tuple[Any, Any]] = {}
+    def _build_model_from_keys(
+        self, keys: list[str] | list[str | tuple[str, str]], name: str
+    ) -> type[BaseModel]:
+        fields: dict[str, tuple[Any, Any]] = {}
         for key in keys:
             if isinstance(key, tuple):
                 # Handle tuple: (original_field_name, renamed_field_name)
@@ -63,7 +66,11 @@ class BasePromptBuilder(BaseModel, ABC):
                 if original_key not in self.model_fields:
                     raise ValueError(f"Key '{original_key}' not found in model fields")
                 field_info = self.model_fields[original_key]
-                field_type = field_info.annotation if field_info.annotation is not None else (Any,)
+                field_type = (
+                    field_info.annotation
+                    if field_info.annotation is not None
+                    else (Any,)
+                )
                 default = field_info.default if field_info.default is not None else ...
                 fields[renamed_key] = (field_type, default)
             else:
@@ -71,28 +78,37 @@ class BasePromptBuilder(BaseModel, ABC):
                 if key not in self.model_fields:
                     raise ValueError(f"Key '{key}' not found in model fields")
                 field_info = self.model_fields[key]
-                field_type = field_info.annotation if field_info.annotation is not None else (Any,)
+                field_type = (
+                    field_info.annotation
+                    if field_info.annotation is not None
+                    else (Any,)
+                )
                 default = field_info.default if field_info.default is not None else ...
                 fields[key] = (field_type, default)
         return create_model(name, **fields)  # type: ignore
 
-    def get_input_model(self) -> Type[BaseModel]:
+    def get_input_model(self) -> type[BaseModel]:
         input_keys, _ = self.get_io_keys()
-        return self._build_model_from_keys(input_keys, "InputModel")
+        return self._build_model_from_keys(input_keys, 'InputModel')
 
-    def get_output_model(self) -> Type[BaseModel]:
+    def get_output_model(self) -> type[BaseModel]:
         _, output_keys = self.get_io_keys()
-        return self._build_model_from_keys(output_keys, "OutputModel")
+        return self._build_model_from_keys(output_keys, 'OutputModel')
 
     # ------------------------------------------------------------------ #
     # Dump methods (JSON)
     # ------------------------------------------------------------------ #
-    def _dump_json_unique(self, schema_model: Type[BaseModel], keys: Union[List[str], List[Union[str, Tuple[str, str]]]], **kwargs) -> str:
+    def _dump_json_unique(
+        self,
+        schema_model: type[BaseModel],
+        keys: list[str] | list[str | tuple[str, str]],
+        **kwargs,
+    ) -> str:
         allowed = list(schema_model.model_fields.keys())
         seen = set()
         unique_keys = [k for k in allowed if not (k in seen or seen.add(k))]
         data = self.model_dump()
-        
+
         # Handle key mapping for renamed fields
         filtered = {}
         for key in keys:
@@ -103,7 +119,7 @@ class BasePromptBuilder(BaseModel, ABC):
             else:
                 if key in data and key in unique_keys:
                     filtered[key] = data[key]
-                    
+
         return schema_model(**filtered).model_dump_json(**kwargs)
 
     def model_dump_json_input(self, **kwargs) -> str:
@@ -117,68 +133,70 @@ class BasePromptBuilder(BaseModel, ABC):
     # ------------------------------------------------------------------ #
     # Markdown helpers
     # ------------------------------------------------------------------ #
-    def _to_markdown(self, obj: Any, level: int = 1, title: Optional[str] = None) -> str:
+    def _to_markdown(self, obj: Any, level: int = 1, title: str | None = None) -> str:
         """
         Recursively convert dict/list/primitive into clean, generic Markdown.
         """
-        md: List[str] = []
-        
+        md: list[str] = []
+
         # Format title if provided
         if title is not None:
             formatted_title = title.replace('_', ' ').title()
             if level <= 2:
-                md.append(f"{'#' * level} {formatted_title}")
+                md.append(f'{"#" * level} {formatted_title}')
             else:
-                md.append(f"**{formatted_title}:**")
-        
+                md.append(f'**{formatted_title}:**')
+
         if isinstance(obj, dict):
             if not obj:  # Empty dict
-                md.append("None")
+                md.append('None')
             else:
                 for k, v in obj.items():
                     if isinstance(v, (str, int, float, bool)) and len(str(v)) < 100:
                         # Short values inline
                         key_name = k.replace('_', ' ').title()
                         if level <= 2:
-                            md.append(f"**{key_name}:** {v}")
+                            md.append(f'**{key_name}:** {v}')
                         else:
-                            md.append(f"- **{key_name}:** {v}")
+                            md.append(f'- **{key_name}:** {v}')
                     else:
                         # Complex values get recursive handling
                         md.append(self._to_markdown(v, level=level + 1, title=k))
         elif isinstance(obj, list):
             if not obj:  # Empty list
-                md.append("None")
+                md.append('None')
             elif all(isinstance(i, dict) for i in obj):
                 # List of objects
                 for i, item in enumerate(obj, 1):
                     if level <= 2:
-                        md.append(f"### {title or 'Item'} {i}")
+                        md.append(f'### {title or "Item"} {i}')
                     else:
-                        md.append(f"**{title or 'Item'} {i}:**")
+                        md.append(f'**{title or "Item"} {i}:**')
                     # Process dict items inline for cleaner output
                     for k, v in item.items():
                         key_name = k.replace('_', ' ').title()
-                        md.append(f"- **{key_name}:** {v}")
+                        md.append(f'- **{key_name}:** {v}')
                     if i < len(obj):  # Add spacing between items
-                        md.append("")
+                        md.append('')
             else:
                 # Simple list
                 for item in obj:
-                    md.append(f"- {item}")
+                    md.append(f'- {item}')
         else:
             # Primitive value
-            value_str = str(obj) if obj is not None else "None"
+            value_str = str(obj) if obj is not None else 'None'
             if title is None:
                 md.append(value_str)
             else:
                 md.append(value_str)
 
-        return "\n".join(md)
+        return '\n'.join(md)
 
-    def _dump_markdown_unique(self, keys: Union[List[str], List[Union[str, Tuple[str, str]]]]) -> str:
+    def _dump_markdown_unique(
+        self, keys: list[str] | list[str | tuple[str, str]]
+    ) -> str:
         data = self.model_dump()
-        filtered: Dict[str, Any] = {}
+        filtered: dict[str, Any] = {}
         for key in keys:
             if isinstance(key, tuple):
                 original_key, renamed_key = key
@@ -187,7 +205,7 @@ class BasePromptBuilder(BaseModel, ABC):
             else:
                 if key in data:
                     filtered[key] = data[key]
-        
+
         # Generate markdown without top-level headers to avoid duplication
         parts = []
         for key, value in filtered.items():
@@ -195,10 +213,10 @@ class BasePromptBuilder(BaseModel, ABC):
                 continue
             formatted_key = key.replace('_', ' ').title()
             if isinstance(value, (str, int, float, bool)) and len(str(value)) < 200:
-                parts.append(f"**{formatted_key}:** {value}")
+                parts.append(f'**{formatted_key}:** {value}')
             else:
                 parts.append(self._to_markdown(value, level=2, title=key))
-        
+
         return '\n'.join(parts)
 
     def model_dump_markdown_input(self) -> str:
@@ -212,42 +230,46 @@ class BasePromptBuilder(BaseModel, ABC):
     # ------------------------------------------------------------------ #
     # Training & preview (JSON or Markdown)
     # ------------------------------------------------------------------ #
-    def build_training_data(self, format: str = "json", indent=None) -> dict[str, Any]:
+    def build_training_data(self, format: str = 'json', indent=None) -> dict[str, Any]:
         """
         Build training data in either JSON (dict for OpenAI-style messages)
         or Markdown (clean format without role prefixes).
         """
-        if format == "json":
+        if format == 'json':
             return {
-                "messages": [
-                    {"role": "system", "content": self.get_instruction()},
-                    {"role": "user", "content": self.model_dump_json_input(indent=indent)},
-                    {"role": "assistant", "content": self.model_dump_json_output(indent=indent)},
+                'messages': [
+                    {'role': 'system', 'content': self.get_instruction()},
+                    {
+                        'role': 'user',
+                        'content': self.model_dump_json_input(indent=indent),
+                    },
+                    {
+                        'role': 'assistant',
+                        'content': self.model_dump_json_output(indent=indent),
+                    },
                 ]
             }
-        elif format == "markdown":
+        if format == 'markdown':
             system_content = self.get_instruction()
-            
+
             return {
                 'messages': [
-                    {"role": "system", "content": system_content},
-                    {"role": "user", "content": self.model_dump_markdown_input()},
-                    {"role": "assistant", "content": self.model_dump_markdown_output()},
+                    {'role': 'system', 'content': system_content},
+                    {'role': 'user', 'content': self.model_dump_markdown_input()},
+                    {'role': 'assistant', 'content': self.model_dump_markdown_output()},
                 ]
             }
         raise ValueError("format must be either 'json' or 'markdown'")
 
     def __str__(self) -> str:
         # Return clean format without explicit role prefixes
-        training_data = self.build_training_data(format="markdown")
+        training_data = self.build_training_data(format='markdown')
         messages = training_data['messages']  # type: ignore[index]
-        
+
         parts = []
         for msg in messages:
             content = msg['content']
-            if msg['role'] == 'system':
-                parts.append(content)
-            elif msg['role'] == 'user':
+            if msg['role'] == 'system' or msg['role'] == 'user':
                 parts.append(content)
             elif msg['role'] == 'assistant':
                 # Get output keys to determine the main output field name
@@ -256,33 +278,34 @@ class BasePromptBuilder(BaseModel, ABC):
                 if isinstance(main_output, tuple):
                     main_output = main_output[1]  # Use renamed key
                 title = main_output.replace('_', ' ').title()
-                parts.append(f"## {title}\n{content}")
-        
+                parts.append(f'## {title}\n{content}')
+
         return '\n\n'.join(parts)
 
     @classmethod
-    def from_messages(cls: Type[B], messages: list[dict]) -> B:
+    def from_messages(cls: type[B], messages: list[dict]) -> B:
         """
         Reconstruct a prompt builder instance from OpenAI-style messages.
         """
-        user_msg = next((m for m in messages if m.get("role") == "user"), None)
-        assistant_msg = next((m for m in messages if m.get("role") == "assistant"), None)
+        user_msg = next((m for m in messages if m.get('role') == 'user'), None)
+        assistant_msg = next(
+            (m for m in messages if m.get('role') == 'assistant'), None
+        )
 
         if user_msg is None:
-            raise ValueError("No user message found")
+            raise ValueError('No user message found')
         if assistant_msg is None:
-            raise ValueError("No assistant message found")
+            raise ValueError('No assistant message found')
 
         try:
-            user_data = json.loads(user_msg["content"])  # type: ignore[index]
+            user_data = json.loads(user_msg['content'])  # type: ignore[index]
         except Exception as e:
-            raise ValueError(f"Invalid user JSON content: {e}")
+            raise ValueError(f'Invalid user JSON content: {e}') from e
 
         try:
-            assistant_data = json.loads(assistant_msg["content"])  # type: ignore[index]
+            assistant_data = json.loads(assistant_msg['content'])  # type: ignore[index]
         except Exception as e:
-            raise ValueError(f"Invalid assistant JSON content: {e}")
+            raise ValueError(f'Invalid assistant JSON content: {e}') from e
 
         combined_data = {**user_data, **assistant_data}
         return cast(B, cls(**combined_data))
-