quantalogic 0.30.8__py3-none-any.whl → 0.31.0__py3-none-any.whl

quantalogic/interactive_text_editor.py

@@ -1,4 +1,4 @@
-"""Utility functions for the QuantaLogic AI Assistant."""
+from typing import List
 
 from prompt_toolkit import PromptSession
 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
@@ -9,70 +9,70 @@ from rich.panel import Panel
 
 
 class InputHistoryManager:
-    """Manages the history of input states for undo functionality."""
-
     def __init__(self):
-        """Initialize the InputHistoryManager with an empty history stack and current state."""
         self.history_stack = []
         self.current_state = []
 
-    def push_state(self, lines):
-        """Push the current state to the history stack and update the current state.
-
-        Args:
-            lines (list): The current lines of input to be saved in the history.
-        """
+    def push_state(self, lines: List[str]) -> None:
         self.history_stack.append(self.current_state.copy())
         self.current_state = lines.copy()
 
-    def undo(self, lines):
-        """Revert to the previous state from the history stack.
-
-        Args:
-            lines (list): The current lines of input to be updated to the previous state.
-
-        Returns:
-            bool: True if undo was successful, False otherwise.
-        """
+    def undo(self, lines: List[str]) -> bool:
         if self.history_stack:
             self.current_state = self.history_stack.pop()
             lines[:] = self.current_state
             return True
         return False
 
-
-def handle_edit_command(lines, args, console, session, history_manager):
-    """Handle the 'edit' command to modify a specific line.
-
-    Args:
-        lines (list): The current lines of input.
-        args (list): The arguments provided with the command.
-        console (Console): The console object for output.
-        session (PromptSession): The prompt session for user input.
-        history_manager (InputHistoryManager): The history manager for state tracking.
-    """
+class CommandRegistry:
+    def __init__(self):
+        self.commands = {}
+        
+    def register(self, name: str, help_text: str = "") -> callable:
+        def decorator(func: callable) -> callable:
+            self.commands[name] = {
+                "handler": func,
+                "help": func.__doc__ or help_text
+            }
+            return func
+        return decorator
+
+registry = CommandRegistry()
+
+@registry.register("/help", "Show available commands")
+def handle_help_command(lines: List[str], args: List[str], console: Console, 
+                      session: PromptSession, history_manager: InputHistoryManager) -> None:
+    """Display auto-generated help from registered commands."""
+    help_content = "\n".join([f"  {name}: {cmd['help']}" for name, cmd in registry.commands.items()])
+    console.print(Panel(f"Available commands:\n{help_content}", title="Help Menu", border_style="green"))
+
+@registry.register("/date", "Show current date/time")
+def handle_date_command(lines: List[str], args: List[str], console: Console,
+                      session: PromptSession, history_manager: InputHistoryManager) -> None:
+    """Display current date and time."""
+    from datetime import datetime
+    console.print(f"[bold #ffaa00]Current datetime: {datetime.now().isoformat()}[/bold #ffaa00]")
+
+@registry.register("/edit", "Edit specific line: /edit <line_number>")
+def handle_edit_command(lines: List[str], args: List[str], console: Console,
+                      session: PromptSession, history_manager: InputHistoryManager) -> None:
+    """Edit a specific line in the input buffer."""
     try:
         edit_line_num = int(args[0]) - 1
         if 0 <= edit_line_num < len(lines):
-            console.print(f"[bold]Editing Line {edit_line_num + 1}:[/bold] {lines[edit_line_num]}")
+            console.print(f"[bold #1d3557]Editing Line {edit_line_num + 1}:[/bold #1d3557] {lines[edit_line_num]}")  # Dark blue
             new_line = session.prompt("New content: ")
             history_manager.push_state(lines)
             lines[edit_line_num] = new_line
         else:
-            console.print("[red]Invalid line number.[/red]")
+            console.print("[bold #ff4444]Invalid line number.[/bold #ff4444]")
     except (ValueError, IndexError):
-        console.print("[red]Invalid edit command. Usage: edit <line_number>[/red]")
+            console.print("[bold #ff4444]Invalid edit command. Usage: /edit <line_number>[/bold #ff4444]")
 
-
-def handle_delete_command(lines, args, console, history_manager):
-    """Handle the 'delete' command to remove a specific line.
-
-    Args:
-        lines (list): The current lines of input.
-        args (list): The arguments provided with the command.
-        console (Console): The console object for output.
-        history_manager (InputHistoryManager): The history manager for state tracking.
-    """
+@registry.register("/delete", "Delete specific line: /delete <line_number>")
+def handle_delete_command(lines: List[str], args: List[str], console: Console,
+                        session: PromptSession, history_manager: InputHistoryManager) -> None:
+    """Delete a specific line from the input buffer."""
     try:
         delete_line_num = int(args[0]) - 1
         if 0 <= delete_line_num < len(lines):
@@ -82,71 +82,89 @@ def handle_delete_command(lines, args, console, history_manager):
         else:
             console.print("[red]Invalid line number.[/red]")
     except (ValueError, IndexError):
-        console.print("[red]Invalid delete command. Usage: delete <line_number>[/red]")
-
-
-def handle_replace_command(lines, args, console, history_manager):
-    """Handle the 'replace' command to search and replace text in all lines.
+            console.print("[bold #ff4444]Invalid delete command. Usage: /delete <line_number>[/bold #ff4444]")
 
-    Args:
-        lines (list): The current lines of input.
-        args (list): The arguments provided with the command.
-        console (Console): The console object for output.
-        history_manager (InputHistoryManager): The history manager for state tracking.
-    """
+@registry.register("/replace", "Search and replace: /replace <search> <replace>")
+def handle_replace_command(lines: List[str], args: List[str], console: Console,
+    session: PromptSession, history_manager: InputHistoryManager) -> None:
     try:
-        search_str, replace_str = args
+        search_str = args[0]
+        replace_str = args[1]
         history_manager.push_state(lines)
         for i in range(len(lines)):
             lines[i] = lines[i].replace(search_str, replace_str)
-        console.print("[bold]Search and replace completed.[/bold]")
-    except ValueError:
-        console.print("[red]Invalid replace command. Usage: replace <search_str> <replace_str>[/red]")
-
-
-commands = {"edit": handle_edit_command, "delete": handle_delete_command, "replace": handle_replace_command}
-
-
-def handle_command(line, lines, console, session, history_manager):
-    """Handle a command entered by the user.
-
-    Args:
-        line (str): The command line entered by the user.
-        lines (list): The current lines of input.
-        console (Console): The console object for output.
-        session (PromptSession): The prompt session for user input.
-        history_manager (InputHistoryManager): The history manager for state tracking.
+            console.print("[bold #00cc66]Search and replace completed.[/bold #00cc66]")
+    except (ValueError, IndexError):
+            console.print("[bold #ff4444]Invalid replace command. Usage: /replace <search_str> <replace_str>[/bold #ff4444]")
 
-    Returns:
-        bool: True if the command was handled, False otherwise.
-    """
-    parts = line.split()
-    if not parts:
-        return False
-    command = parts[0]
-    args = parts[1:]
-    if command in commands:
-        commands[command](lines, args, console, session, history_manager)
-        return True
-    return False
+@registry.register("/model", "Show current AI model") 
+def handle_model_command(lines: List[str], args: List[str], console: Console,
+    session: PromptSession, history_manager: InputHistoryManager) -> None:
+    from quantalogic.agent_factory import AgentRegistry
+    try:
+        current_agent = AgentRegistry.get_agent("main_agent")
+        if current_agent:
+            console.print(f"[bold #ffaa00]Current AI model: {current_agent.model_name}[/bold #ffaa00]")
+        else:
+            console.print("[bold #ffaa00]No active agent found.[/bold #ffaa00]")
+    except ValueError as e:
+            console.print(f"[bold #ff4444]Error: {str(e)}[/bold #ff4444]")
+
+@registry.register("/setmodel", "Set AI model name: /setmodel <name>")
+def handle_set_model_command(lines: List[str], args: List[str], console: Console,
+    session: PromptSession, history_manager: InputHistoryManager) -> None:
+    from quantalogic.agent_factory import AgentRegistry
+    try:
+        if len(args) < 1:
+            console.print("[bold #ff4444]Error: Model name required. Usage: /setmodel <name>[/bold #ff4444]")
+            return
+            
+        model_name = args[0]
+        current_agent = AgentRegistry.get_agent("main_agent")
+        if current_agent:
+            current_agent.model_name = model_name
+            console.print(f"[bold #00cc66]Model name updated to: {model_name}[/bold #00cc66]")
+        else:
+            console.print("[yellow]No active agent found.[/yellow]")
+    except ValueError as e:
+        console.print(f"[red]Error: {str(e)}[/red]")
+
+@registry.register("/models", "List all available AI models")
+def handle_models_command(lines: List[str], args: List[str], console: Console,
+    session: PromptSession, history_manager: InputHistoryManager) -> None:
+    """Display all available AI models supported by the system."""
+    from quantalogic.utils.get_all_models import get_all_models
+    try:
+        models = get_all_models()
+        if models:
+            # Group models by provider
+            provider_groups = {}
+            for model in models:
+                provider = model.split('/')[0] if '/' in model else 'default'
+                if provider not in provider_groups:
+                    provider_groups[provider] = []
+                provider_groups[provider].append(model)
+            
+            # Create formatted output
+            output = "[bold #00cc66]Available AI Models:[/bold #00cc66]\n"
+            for provider, model_list in provider_groups.items():
+                output += f"\n[bold #ffaa00]{provider.upper()}[/bold #ffaa00]\n"
+                for model in sorted(model_list):
+                    output += f"  • {model}\n"
+            
+            console.print(Panel(output, border_style="green"))
+        else:
+            console.print("[yellow]No models available.[/yellow]")
+    except Exception as e:
+        console.print(f"[red]Error retrieving models: {str(e)}[/red]")
 
 
 def get_multiline_input(console: Console) -> str:
-    """Get multiline input from the user with enhanced UX.
-
-    Args:
-        console (Console): The console object for output.
-
-    Returns:
-        str: The multiline input provided by the user.
-    """
+    """Get multiline input with slash command support."""
     console.print(
         Panel(
             "Enter your task. Press [bold]Enter[/bold] twice to submit.\n"
-            "Available commands:\n"
-            "  edit <line_number> - Edit a specific line\n"
-            "  delete <line_number> - Delete a specific line\n"
-            "  replace <search_str> <replace_str> - Replace text in all lines",
+            "Type /help for available commands",
             title="Multi-line Input",
             border_style="blue",
         )
@@ -162,18 +180,174 @@ def get_multiline_input(console: Console) -> str:
     @bindings.add("c-z")
     def _(event):
         if history_manager.undo(lines):
-            console.print("[bold]Undo successful.[/bold]")
+            console.print("[bold #00cc66]Undo successful.[/bold #00cc66]")
+
+
+    from prompt_toolkit.completion import Completer, Completion
+    from prompt_toolkit.styles import Style
+    
+    class CommandCompleter(Completer):
+        def get_completions(self, document, complete_event):
+            text = document.text_before_cursor
+            line_parts = text.split()
+            
+            if len(line_parts) == 0 or (len(line_parts) == 1 and text.endswith(' ')):
+                for cmd, details in registry.commands.items():
+                    yield Completion(
+                        cmd,
+                        start_position=-len(text),
+                        display=f"[{cmd}] {details['help']}",
+                        style="fg:ansicyan bold",
+                    )
+            elif line_parts[0] in registry.commands:
+                cmd = registry.commands[line_parts[0]]
+                arg_index = len(line_parts) - 1
+                if arg_index == 1:
+                    doc = cmd['handler'].__doc__ or ""
+                    args_hint = next((line.split(':')[1].strip() for line in doc.split('\n') 
+                                    if 'Args:' in line), "")
+                    if args_hint:
+                        yield Completion(
+                            f"<{args_hint}>",
+                            start_position=-len(text.split()[-1]),
+                            style="fg:ansimagenta italic",
+                            display=f"Expected argument: {args_hint}",
+                        )
+            else:
+                if text.startswith('/'):
+                    partial = text[1:].lstrip('/')
+                    exact_matches = []
+                    prefix_matches = []
+                    
+                    for cmd in registry.commands:
+                        cmd_without_slash = cmd[1:]
+                        if cmd_without_slash.lower() == partial.lower():
+                            exact_matches.append(cmd)
+                        elif cmd_without_slash.lower().startswith(partial.lower()):
+                            prefix_matches.append(cmd)
+                    
+                    # Prioritize exact matches first
+                    for match in exact_matches:
+                        remaining = match[len('/' + partial):]
+                        yield Completion(
+                            remaining,
+                            start_position=0,  # Corrected from -len(partial)
+                            display=f"{match} - {registry.commands[match]['help']}",
+                            style="fg:ansiyellow bold",
+                        )
+                    
+                    # Then prefix matches
+                    for match in prefix_matches:
+                        remaining = match[len('/' + partial):]
+                        yield Completion(
+                            remaining,
+                            start_position=0,  # Corrected from -len(partial)
+                            display=f"{match} - {registry.commands[match]['help']}",
+                            style="fg:ansiyellow bold",
+                        )
+            def get_completions(self, document, complete_event):
+                text = document.text_before_cursor
+                line_parts = text.split()
+                
+                if len(line_parts) == 0 or (len(line_parts) == 1 and text.endswith(' ')):
+                    for cmd, details in registry.commands.items():
+                        yield Completion(
+                            cmd,
+                            start_position=-len(text),
+                            display=f"[{cmd}] {details['help']}",
+                            style="fg:ansicyan bold",
+                        )
+                elif line_parts[0] in registry.commands:
+                    cmd = registry.commands[line_parts[0]]
+                    arg_index = len(line_parts) - 1
+                    if arg_index == 1:
+                        doc = cmd['handler'].__doc__ or ""
+                        args_hint = next((line.split(':')[1].strip() for line in doc.split('\n') 
+                                        if 'Args:' in line), "")
+                        if args_hint:
+                            yield Completion(
+                                f"<{args_hint}>",
+                                start_position=-len(text.split()[-1]),
+                                style="fg:ansimagenta italic",
+                                display=f"Expected argument: {args_hint}",
+                            )
+                else:
+                    if text.startswith('/'):
+                        partial = text[1:].lstrip('/')
+                        exact_matches = []
+                        prefix_matches = []
+                        
+                        for cmd in registry.commands:
+                            cmd_without_slash = cmd[1:]
+                            if cmd_without_slash.lower() == partial.lower():
+                                exact_matches.append(cmd)
+                            elif cmd_without_slash.lower().startswith(partial.lower()):
+                                prefix_matches.append(cmd)
+                        
+                        for match in exact_matches:
+                            remaining = match[len('/' + partial):]
+                            yield Completion(
+                                remaining,
+                                start_position=-len(partial),
+                                display=f"{match} - {registry.commands[match]['help']}",
+                                style="fg:ansiyellow bold",
+                            )
+                        
+                        for match in prefix_matches:
+                            remaining = match[len('/' + partial):]
+                            yield Completion(
+                                remaining,
+                                start_position=-len(partial),
+                                display=f"{match} - {registry.commands[match]['help']}",
+                                style="fg:ansiyellow bold",
+                            )
+
+    command_completer = CommandCompleter()
+    
+    def get_command_help(cmd_name: str) -> str:
+        """Get formatted help text for a command"""
+        if cmd := registry.commands.get(cmd_name):
+            return f"[bold]{cmd_name}[/bold]: {cmd['help']}\n{cmd['handler'].__doc__ or ''}"
+        return ""
 
-    session = PromptSession(history=InMemoryHistory(), auto_suggest=AutoSuggestFromHistory(), key_bindings=bindings)
+    custom_style = Style.from_dict({
+        'completion-menu.completion': 'bg:#005577 #ffffff',
+        'completion-menu.completion.current': 'bg:#007799 #ffffff bold',
+        'scrollbar.background': 'bg:#6699aa',
+        'scrollbar.button': 'bg:#444444',
+        'documentation': 'bg:#003366 #ffffff',
+    })
+
+    session = PromptSession(
+        history=InMemoryHistory(),
+        auto_suggest=AutoSuggestFromHistory(),
+        key_bindings=bindings,
+        completer=command_completer,
+        complete_while_typing=True,
+        style=custom_style,
+        bottom_toolbar=lambda: get_command_help(session.default_buffer.document.text.split()[0][1:] 
+                                              if session.default_buffer.document.text.startswith('/') 
+                                              else ""),
+    )
 
     try:
         while True:
             prompt_text = f"{line_number:>3}: "
             line = session.prompt(prompt_text, rprompt="Press Enter twice to submit")
 
-            # Handle commands with single return
             if line.strip().startswith('/'):
-                return line.strip()
+                cmd_parts = line.strip().split()
+                cmd_name = cmd_parts[0].lower()
+                args = cmd_parts[1:]
+                
+                if cmd_handler := registry.commands.get(cmd_name):
+                    try:
+                        cmd_handler["handler"](lines, args, console, session, history_manager)
+                    except Exception as e:
+                        console.print(f"[red]Error executing {cmd_name}: {str(e)}[/red]")
+                else:
+                    console.print(f"[red]Unknown command: {cmd_name}[/red]")
+                continue
 
             if line.strip() == "":
                 blank_lines += 1
@@ -181,7 +355,7 @@ def get_multiline_input(console: Console) -> str:
                     break
             else:
                 blank_lines = 0
-                if not handle_command(line, lines, console, session, history_manager):
+                if not any(line.strip().startswith(cmd) for cmd in registry.commands):
                     history_manager.push_state(lines)
                     lines.append(line)
                     line_number += 1
@@ -191,4 +365,4 @@ def get_multiline_input(console: Console) -> str:
         console.print("\n[bold]Input cancelled by user.[/bold]")
         return ""
 
-    return "\n".join(lines)
+    return "\n".join(lines)

quantalogic/llm.py

@@ -0,0 +1,135 @@
+"""LLM wrapper module for handling LiteLLM operations."""
+
+__all__ = [
+    "generate_completion",
+    "generate_image",
+    "count_tokens",
+    "get_model_max_input_tokens",
+    "get_model_max_output_tokens",
+]
+
+import os
+from typing import Any, Dict, List
+
+from litellm import (
+    completion,
+    image_generation,
+    token_counter,
+)
+from loguru import logger
+
+from quantalogic.get_model_info import (
+    get_max_input_tokens,
+    get_max_output_tokens,
+    model_info,
+)
+
+
+def get_model_info(model_name: str) -> dict | None:
+    """Get model information for a given model name."""
+    return model_info.get(model_name, None)
+
+
+def generate_completion(**kwargs: Dict[str, Any]) -> Any:
+    """Wraps litellm completion with proper type hints."""
+    model = kwargs.get("model", "")
+    if model.startswith("dashscope/"):
+        # Remove prefix and configure for OpenAI-compatible endpoint
+        kwargs["model"] = model.replace("dashscope/", "")
+        kwargs["custom_llm_provider"] = "openai"  # Explicitly specify OpenAI provider
+        kwargs["base_url"] = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
+        api_key = os.getenv("DASHSCOPE_API_KEY")
+        if not api_key:
+            raise ValueError("DASHSCOPE_API_KEY is not set in the environment variables.")
+        kwargs["api_key"] = api_key
+    return completion(**kwargs)
+
+
+def generate_image(**kwargs: Dict[str, Any]) -> Any:
+    """Wraps litellm image_generation with proper type hints."""
+    return image_generation(**kwargs)
+
+
+def count_tokens(model: str, messages: List[Dict[str, Any]]) -> int:
+    """Wraps litellm token_counter with proper type hints."""
+    return token_counter(model=model, messages=messages)
+
+
+def _get_model_info_impl(model_name: str) -> dict:
+    """Get information about the model with prefix fallback logic."""
+    original_model = model_name
+    tried_models = [model_name]
+
+    while True:
+        try:
+            logger.debug(f"Attempting to retrieve model info for: {model_name}")
+            # Try direct lookup from model_info dictionary first
+            if model_name in model_info:
+                logger.debug(f"Found model info for {model_name} in model_info")
+                return model_info[model_name]
+
+            # Try get_model_info as fallback
+            info = get_model_info(model_name)
+            if info:
+                logger.debug(f"Found model info for {model_name} via get_model_info")
+                return info
+        except Exception as e:
+            logger.debug(f"Failed to get model info for {model_name}: {str(e)}")
+            pass
+
+        # Try removing one prefix level
+        parts = model_name.split("/")
+        if len(parts) <= 1:
+            break
+        model_name = "/".join(parts[1:])
+        tried_models.append(model_name)
+
+    error_msg = f"Could not find model info for {original_model} after trying: {' → '.join(tried_models)}"
+    logger.error(error_msg)
+    raise ValueError(error_msg)
+
+
+def get_model_max_input_tokens(model_name: str) -> int | None:
+    """Get the maximum number of input tokens for the model."""
+    try:
+        # First try direct lookup
+        max_tokens = get_max_input_tokens(model_name)
+        if max_tokens is not None:
+            return max_tokens
+
+        # If not found, try getting from model info
+        model_info = _get_model_info_impl(model_name)
+        max_input = model_info.get("max_input_tokens")
+        if max_input is not None:
+            return max_input
+
+        # If still not found, log warning and return default
+        logger.warning(f"No max input tokens found for {model_name}. Using default.")
+        return 8192  # A reasonable default for many models
+
+    except Exception as e:
+        logger.error(f"Error getting max input tokens for {model_name}: {e}")
+        return None
+
+
+def get_model_max_output_tokens(model_name: str) -> int | None:
+    """Get the maximum number of output tokens for the model."""
+    try:
+        # First try direct lookup
+        max_tokens = get_max_output_tokens(model_name)
+        if max_tokens is not None:
+            return max_tokens
+
+        # If not found, try getting from model info
+        model_info = _get_model_info_impl(model_name)
+        max_output = model_info.get("max_output_tokens")
+        if max_output is not None:
+            return max_output
+
+        # If still not found, log warning and return default
+        logger.warning(f"No max output tokens found for {model_name}. Using default.")
+        return 4096  # A reasonable default for many models
+
+    except Exception as e:
+        logger.error(f"Error getting max output tokens for {model_name}: {e}")
+        return None