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

llm_utils/__init__.py

@@ -1,5 +1,10 @@
 from llm_utils.lm.openai_memoize import MOpenAI
+from llm_utils.lm import  LLMTask, AsyncLM, AsyncLLMTask
 from llm_utils.vector_cache import VectorCache
+from llm_utils.lm.lm_base import get_model_name
+from llm_utils.lm.base_prompt_builder import BasePromptBuilder
+
+LLM = LLMTask
 
 from .chat_format import (
     build_chatml_input,
@@ -12,7 +17,6 @@ from .chat_format import (
     transform_messages,
     transform_messages_to_chatml,
 )
-from .lm.async_lm import AsyncLLMTask, AsyncLM
 
 __all__ = [
     "transform_messages",
@@ -26,6 +30,10 @@ __all__ = [
     "display_chat_messages_as_html",
     "AsyncLM",
     "AsyncLLMTask",
+    "LLMTask",
     "MOpenAI",
+    "get_model_name",
     "VectorCache",
+    "BasePromptBuilder",
+    "LLM"
 ]

llm_utils/chat_format/display.py

@@ -1,19 +1,93 @@
 from __future__ import annotations
 
+import json
 from difflib import SequenceMatcher
 from typing import Any, Optional
 
 from IPython.display import HTML, display
 
 
+def _preprocess_as_json(content: str) -> str:
+    """
+    Preprocess content as JSON with proper formatting and syntax highlighting.
+    """
+    try:
+        # Try to parse and reformat JSON
+        parsed = json.loads(content)
+        return json.dumps(parsed, indent=2, ensure_ascii=False)
+    except (json.JSONDecodeError, TypeError):
+        # If not valid JSON, return as-is
+        return content
+
+
+def _preprocess_as_markdown(content: str) -> str:
+    """
+    Preprocess content as markdown with proper formatting.
+    """
+    # Basic markdown preprocessing - convert common patterns
+    lines = content.split('\n')
+    processed_lines = []
+    
+    for line in lines:
+        # Convert **bold** to span with bold styling
+        while '**' in line:
+            first_pos = line.find('**')
+            if first_pos != -1:
+                second_pos = line.find('**', first_pos + 2)
+                if second_pos != -1:
+                    before = line[:first_pos]
+                    bold_text = line[first_pos + 2:second_pos]
+                    after = line[second_pos + 2:]
+                    line = f'{before}<span style="font-weight: bold;">{bold_text}</span>{after}'
+                else:
+                    break
+            else:
+                break
+        
+        # Convert *italic* to span with italic styling
+        while '*' in line and line.count('*') >= 2:
+            first_pos = line.find('*')
+            if first_pos != -1:
+                second_pos = line.find('*', first_pos + 1)
+                if second_pos != -1:
+                    before = line[:first_pos]
+                    italic_text = line[first_pos + 1:second_pos]
+                    after = line[second_pos + 1:]
+                    line = f'{before}<span style="font-style: italic;">{italic_text}</span>{after}'
+                else:
+                    break
+            else:
+                break
+        
+        # Convert # headers to bold headers
+        if line.strip().startswith('#'):
+            level = len(line) - len(line.lstrip('#'))
+            header_text = line.lstrip('# ').strip()
+            line = f'<span style="font-weight: bold; font-size: 1.{min(4, level)}em;">{header_text}</span>'
+        
+        processed_lines.append(line)
+    
+    return '\n'.join(processed_lines)
+
+
 def show_chat(
     msgs: Any,
     return_html: bool = False,
     file: str = "/tmp/conversation.html",
     theme: str = "default",
+    as_markdown: bool = False,
+    as_json: bool = False,
 ) -> Optional[str]:
     """
     Display chat messages as HTML.
+    
+    Args:
+        msgs: Chat messages in various formats
+        return_html: If True, return HTML string instead of displaying
+        file: Path to save HTML file
+        theme: Color theme ('default', 'light', 'dark')
+        as_markdown: If True, preprocess content as markdown
+        as_json: If True, preprocess content as JSON
     """
     if isinstance(msgs, dict) and "messages" in msgs:
         msgs = msgs["messages"]
@@ -74,45 +148,66 @@ def show_chat(
                 name = tool_call["name"]
                 args = tool_call["arguments"]
                 content += f"Tool: {name}\nArguments: {args}"
-        content = content.replace("\n", "<br>")
-        content = content.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
-        content = content.replace("  ", "&nbsp;&nbsp;")
-        content = (
-            content.replace("<br>", "TEMP_BR")
-            .replace("<", "&lt;")
-            .replace(">", "&gt;")
-            .replace("TEMP_BR", "<br>")
-        )
+        
+        # Preprocess content based on format options
+        if as_json:
+            content = _preprocess_as_json(content)
+        elif as_markdown:
+            content = _preprocess_as_markdown(content)
+        
+        # Handle HTML escaping differently for markdown vs regular content
+        if as_markdown:
+            # For markdown, preserve HTML tags but escape other characters carefully
+            content = content.replace("\n", "<br>")
+            content = content.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
+            content = content.replace("  ", "&nbsp;&nbsp;")
+            # Don't escape < and > for markdown since we want to preserve our span tags
+        else:
+            # Regular escaping for non-markdown content
+            content = content.replace("\n", "<br>")
+            content = content.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
+            content = content.replace("  ", "&nbsp;&nbsp;")
+            content = (
+                content.replace("<br>", "TEMP_BR")
+                .replace("<", "&lt;")
+                .replace(">", "&gt;")
+                .replace("TEMP_BR", "<br>")
+            )
         if role in color_scheme:
             background_color = color_scheme[role]["background"]
             text_color = color_scheme[role]["text"]
         else:
             background_color = color_scheme["default"]["background"]
             text_color = color_scheme["default"]["text"]
+        
+        # Choose container based on whether we have markdown formatting
+        content_container = "div" if as_markdown else "pre"
+        container_style = 'style="white-space: pre-wrap;"' if as_markdown else ""
+        
         if role == "system":
             conversation_html += (
                 f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
-                f'<strong>System:</strong><br><pre id="system-{i}">{content}</pre></div>'
+                f'<strong>System:</strong><br><{content_container} id="system-{i}" {container_style}>{content}</{content_container}></div>'
             )
         elif role == "user":
             conversation_html += (
                 f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
-                f'<strong>User:</strong><br><pre id="user-{i}">{content}</pre></div>'
+                f'<strong>User:</strong><br><{content_container} id="user-{i}" {container_style}>{content}</{content_container}></div>'
             )
         elif role == "assistant":
             conversation_html += (
                 f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
-                f'<strong>Assistant:</strong><br><pre id="assistant-{i}">{content}</pre></div>'
+                f'<strong>Assistant:</strong><br><{content_container} id="assistant-{i}" {container_style}>{content}</{content_container}></div>'
             )
         elif role == "function":
             conversation_html += (
                 f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
-                f'<strong>Function:</strong><br><pre id="function-{i}">{content}</pre></div>'
+                f'<strong>Function:</strong><br><{content_container} id="function-{i}" {container_style}>{content}</{content_container}></div>'
             )
         else:
             conversation_html += (
                 f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
-                f'<strong>{role}:</strong><br><pre id="{role}-{i}">{content}</pre><br>'
+                f'<strong>{role}:</strong><br><{content_container} id="{role}-{i}" {container_style}>{content}</{content_container}><br>'
                 f"<button onclick=\"copyContent('{role}-{i}')\">Copy</button></div>"
             )
     html: str = f"""

llm_utils/lm/__init__.py

@@ -1,12 +1,13 @@
-# from .async_lm import AsyncLLMTask, AsyncLM
-# from .lm import LM, LLMTask
+from .async_lm.async_lm import AsyncLM
+from .async_lm.async_llm_task import AsyncLLMTask
+from .lm_base import LMBase, get_model_name
+from .llm_task import  LLMTask
+from .base_prompt_builder import BasePromptBuilder
 
-# OAI_LM = LM
-
-# __all__ = [
-#     "LM",
-#     "OAI_LM",
-#     "AsyncLM",
-#     "LLMTask",
-#     "AsyncLLMTask",
-# ]
+__all__ = [
+    "LMBase",
+    "LLMTask",
+    "AsyncLM",
+    "AsyncLLMTask",
+    "BasePromptBuilder",
+]

llm_utils/lm/async_lm/async_llm_task.py

@@ -1,3 +1,4 @@
+# type: ignore
 """
 Async LLM Task module for handling language model interactions with structured input/output.
 """
@@ -33,8 +34,6 @@ class LMConfiguration:
     model: Optional[str] = None
     temperature: Optional[float] = None
     max_tokens: Optional[int] = None
-    host: Optional[str] = None
-    port: Optional[Union[int, str]] = None
     base_url: Optional[str] = None
     api_key: Optional[str] = None
     cache: Optional[bool] = True
@@ -53,8 +52,6 @@ class LMConfiguration:
             "model": self.model,
             "temperature": self.temperature,
             "max_tokens": self.max_tokens,
-            "host": self.host,
-            "port": self.port,
             "base_url": self.base_url,
             "api_key": self.api_key,
             "cache": self.cache,
@@ -90,8 +87,6 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
     DEFAULT_CACHE_DIR: Optional[pathlib.Path] = None
     DEFAULT_TEMPERATURE: Optional[float] = None
     DEFAULT_MAX_TOKENS: Optional[int] = None
-    DEFAULT_HOST: Optional[str] = None
-    DEFAULT_PORT: Optional[Union[int, str]] = None
     DEFAULT_TOP_P: Optional[float] = None
     DEFAULT_PRESENCE_PENALTY: Optional[float] = None
     DEFAULT_TOP_K: Optional[int] = None
@@ -112,8 +107,6 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
         model: Optional[str] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
-        host: Optional[str] = None,
-        port: Optional[Union[int, str]] = None,
         base_url: Optional[str] = None,
         api_key: Optional[str] = None,
         cache: Optional[bool] = None,
@@ -139,8 +132,6 @@ class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
             max_tokens=max_tokens
             if max_tokens is not None
             else self.DEFAULT_MAX_TOKENS,
-            host=host if host is not None else self.DEFAULT_HOST,
-            port=port if port is not None else self.DEFAULT_PORT,
             base_url=base_url if base_url is not None else self.DEFAULT_BASE_URL,
             api_key=api_key if api_key is not None else self.DEFAULT_API_KEY,
             cache=cache if cache is not None else self.DEFAULT_CACHE,

llm_utils/lm/async_lm/async_lm.py

@@ -27,7 +27,6 @@ from ._utils import (
 
 
 def jloads_safe(content: str) -> Any:
-    # if contain ```json, remove it
     if "```json" in content:
         content = content.split("```json")[1].strip().split("```")[0].strip()
     try:
@@ -72,8 +71,6 @@ class AsyncLM(AsyncLMBase):
             print(f"Using model: {model}")
 
         super().__init__(
-            host=host,
-            port=port,
             ports=ports,
             base_url=base_url,
             cache=cache,
@@ -231,6 +228,7 @@ class AsyncLM(AsyncLMBase):
     def _extract_assistant_message(self, choice):  # -> dict[str, str] | dict[str, Any]:
         # TODO this current assume choice is a dict with "reasoning_content" and "content"
         has_reasoning = False
+        reasoning_content = ""
         if "reasoning_content" in choice and isinstance(
             choice["reasoning_content"], str
         ):
@@ -249,7 +247,7 @@ class AsyncLM(AsyncLMBase):
 
         return assistant_msg
 
-    async def __call__(
+    async def call_with_messages(
         self,
         prompt: Optional[str] = None,
         messages: Optional[RawMsgs] = None,
@@ -295,6 +293,17 @@ class AsyncLM(AsyncLMBase):
             msg_dump = dict(assistant_msg)
         return msg_dump, full_messages
 
+    
+    def call_sync(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[RawMsgs] = None,
+        max_tokens: Optional[int] = None,
+    ):
+        """Synchronous wrapper around the async __call__ method."""
+        import asyncio
+        return asyncio.run(self.__call__(prompt=prompt, messages=messages, max_tokens=max_tokens))
+    
     async def parse(
         self,
         instruction,

llm_utils/lm/async_lm/async_lm_base.py

@@ -40,32 +40,40 @@ class AsyncLMBase:
     def __init__(
         self,
         *,
-        host: str = "localhost",
-        port: Optional[Union[int, str]] = None,
         base_url: Optional[str] = None,
         api_key: Optional[str] = None,
         cache: bool = True,
         ports: Optional[List[int]] = None,
     ) -> None:
-        self._port = port
-        self._host = host
-        self.base_url = base_url or (f"http://{host}:{port}/v1" if port else None)
+        self.base_url = base_url
         self.api_key = api_key or os.getenv("OPENAI_API_KEY", "abc")
         self._cache = cache
         self.ports = ports
-        self._init_port = port  # <-- store the port provided at init
 
     @property
     def client(self) -> MAsyncOpenAI:
         # if have multiple ports
-        if self.ports:
+        if self.ports and self.base_url:
             import random
-
+            import re
+            
             port = random.choice(self.ports)
-            api_base = f"http://{self._host}:{port}/v1"
+            # 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 or f"http://{self._host}:{self._port}/v1"
+            api_base = self.base_url
+            
+        if api_base is None:
+            raise ValueError("base_url must be provided")
+            
         client = MAsyncOpenAI(
             api_key=self.api_key,
             base_url=api_base,
@@ -182,11 +190,13 @@ class AsyncLMBase:
     # ------------------------------------------------------------------ #
 
     @staticmethod
-    async def list_models(port=None, host="localhost") -> List[str]:
+    async def list_models(base_url: Optional[str] = None) -> List[str]:
         try:
-            client = AsyncLMBase(port=port, host=host).client  # type: ignore[arg-type]
-            base_url: URL = client.base_url
-            logger.debug(f"Base URL: {base_url}")
+            if base_url is None:
+                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}")
             models: AsyncSyncPage[Model] = await client.models.list()  # type: ignore[assignment]
             return [model.id for model in models.data]
         except Exception as exc: