yaicli 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

yaicli/printer.py

@@ -1,20 +1,31 @@
-import itertools
 import time
 import traceback
 from typing import (
     Any,
     Dict,
     Iterator,
+    List,
     Optional,
     Tuple,
 )
 
-from rich import get_console
-from rich.console import Console
+from rich.console import Console, Group, RenderableType
 from rich.live import Live
-from rich.markdown import Markdown
 
-from yaicli.const import DEFAULT_CODE_THEME, EventTypeEnum
+from yaicli.console import get_console
+from yaicli.const import EventTypeEnum
+from yaicli.render import JustifyMarkdown as Markdown
+from yaicli.render import plain_formatter
+
+
+def cursor_animation() -> Iterator[str]:
+    """Generate a cursor animation for the console."""
+    cursors = ["_", " "]
+    while True:
+        # Use current time to determine cursor state (changes every 0.5 seconds)
+        current_time = time.time()
+        # Alternate between cursors based on time
+        yield cursors[int(current_time * 2) % 2]
 
 
 class Printer:
@@ -23,12 +34,40 @@ class Printer:
     _REASONING_PREFIX = "> "
     _CURSOR_ANIMATION_SLEEP = 0.005
 
-    def __init__(self, config: Dict[str, Any], console: Console, verbose: bool = False):
+    def __init__(
+        self,
+        config: Dict[str, Any],
+        console: Console,
+        verbose: bool = False,
+        markdown: bool = True,
+        reasoning_markdown: Optional[bool] = None,
+        content_markdown: Optional[bool] = None,
+    ):
+        """Initialize the Printer class.
+
+        Args:
+            config (Dict[str, Any]): The configuration dictionary.
+            console (Console): The console object.
+            verbose (bool): Whether to print verbose output.
+            markdown (bool): Whether to use Markdown formatting for all output (legacy).
+            reasoning_markdown (Optional[bool]): Whether to use Markdown for reasoning sections.
+            content_markdown (Optional[bool]): Whether to use Markdown for content sections.
+        """
         self.config = config
         self.console = console or get_console()
         self.verbose = verbose
-        self.code_theme = config.get("CODE_THEME", DEFAULT_CODE_THEME)
+        self.code_theme = config["CODE_THEME"]
         self.in_reasoning: bool = False
+        # Print reasoning content or not
+        self.show_reasoning = config["SHOW_REASONING"]
+
+        # Use explicit settings if provided, otherwise fall back to the global markdown setting
+        self.reasoning_markdown = reasoning_markdown if reasoning_markdown is not None else markdown
+        self.content_markdown = content_markdown if content_markdown is not None else markdown
+
+        # Set formatters for reasoning and content
+        self.reasoning_formatter = Markdown if self.reasoning_markdown else plain_formatter
+        self.content_formatter = Markdown if self.content_markdown else plain_formatter
 
     def _reset_state(self) -> None:
         """Resets the printer state for a new stream."""
@@ -118,7 +157,7 @@ class Printer:
 
         return content, reasoning
 
-    def _format_display_text(self, content: str, reasoning: str) -> str:
+    def _format_display_text(self, content: str, reasoning: str) -> RenderableType:
         """Format the text for display, combining content and reasoning if needed.
 
         Args:
@@ -126,26 +165,39 @@ class Printer:
             reasoning (str): The reasoning text.
 
         Returns:
-            str: The formatted text for display.
+            RenderableType: The formatted text ready for display as a Rich renderable.
         """
-        display_text = ""
-
-        # Add reasoning with proper formatting if it exists
-        if reasoning:
-            formatted_reasoning = reasoning.replace("\n", f"\n{self._REASONING_PREFIX}")
-            if not formatted_reasoning.startswith(self._REASONING_PREFIX):
-                formatted_reasoning = self._REASONING_PREFIX + formatted_reasoning
-            # Reasoning prefix is now added per line
-            display_text += f"\nThinking:\n{formatted_reasoning}"
-
-            # Only add newlines if there is content following the reasoning
-            if content:
-                display_text += "\n\n"
-
-        # Add content
-        display_text += content
-
-        return display_text
+        # Create list of display elements to avoid type issues with concatenation
+        display_elements: List[RenderableType] = []
+
+        reasoning = reasoning.strip()
+        # Format reasoning with proper formatting if it exists
+        if reasoning and self.show_reasoning:
+            raw_reasoning = reasoning.replace("\n", f"\n{self._REASONING_PREFIX}")
+            if not raw_reasoning.startswith(self._REASONING_PREFIX):
+                raw_reasoning = self._REASONING_PREFIX + raw_reasoning
+
+            # Format the reasoning section
+            reasoning_header = "\nThinking:\n"
+            formatted_reasoning = self.reasoning_formatter(reasoning_header + raw_reasoning, code_theme=self.code_theme)
+            display_elements.append(formatted_reasoning)
+
+        content = content.strip()
+        # Format content if it exists
+        if content:
+            formatted_content = self.content_formatter(content, code_theme=self.code_theme)
+
+            # Add spacing between reasoning and content if both exist
+            if reasoning and self.show_reasoning:
+                display_elements.append("")
+
+            display_elements.append(formatted_content)
+
+        # Return based on what we have
+        if not display_elements:
+            return ""
+        # Use Rich Group to combine multiple renderables
+        return Group(*display_elements)
 
     def _update_live_display(self, live: Live, content: str, reasoning: str, cursor: Iterator[str]) -> None:
         """Update live display content and execute cursor animation
@@ -157,52 +209,57 @@ class Printer:
             reasoning (str): The current reasoning text.
             cursor (Iterator[str]): The cursor animation iterator.
         """
-        # Format display text without cursor first
-        display_text = self._format_display_text(content, reasoning)
+
         cursor_char = next(cursor)
 
-        # Add cursor at the end of reasoning or content depending on current state
-        if self.in_reasoning:
-            # Add cursor at the end of reasoning content
+        # Handle cursor placement based on current state
+        if self.in_reasoning and self.show_reasoning:
+            # For reasoning, add cursor in plaintext to reasoning section
             if reasoning:
-                # Append cursor directly if reasoning ends without newline, otherwise append after prefix
                 if reasoning.endswith("\n"):
-                    display_text += f"\n{self._REASONING_PREFIX}{cursor_char}"
+                    cursor_line = f"\n{self._REASONING_PREFIX}{cursor_char}"
                 else:
-                    display_text += cursor_char
+                    cursor_line = cursor_char
+
+                # Re-format with cursor added
+                raw_reasoning = reasoning + cursor_line.replace(self._REASONING_PREFIX, "")
+                formatted_display = self._format_display_text(content, raw_reasoning)
             else:
-                # If reasoning just started and no content yet
-                # Updated to match formatting
-                display_text = f"\nThinking:\n{self._REASONING_PREFIX}{cursor_char}"
+                # If reasoning just started with no content yet
+                reasoning_header = f"\nThinking:\n{self._REASONING_PREFIX}{cursor_char}"
+                formatted_reasoning = self.reasoning_formatter(reasoning_header, code_theme=self.code_theme)
+                formatted_display = Group(formatted_reasoning)
         else:
-            # Add cursor at the end of normal content
-            display_text += cursor_char
+            # For content, add cursor to content section
+            formatted_content_with_cursor = content + cursor_char
+            formatted_display = self._format_display_text(formatted_content_with_cursor, reasoning)
 
-        markdown = Markdown(display_text, code_theme=self.code_theme)
-        live.update(markdown)
+        live.update(formatted_display)
         time.sleep(self._CURSOR_ANIMATION_SLEEP)
 
-    def display_stream(self, stream_iterator: Iterator[Dict[str, Any]]) -> Tuple[Optional[str], Optional[str]]:
+    def display_stream(
+        self, stream_iterator: Iterator[Dict[str, Any]], with_assistant_prefix: bool = True
+    ) -> Tuple[Optional[str], Optional[str]]:
         """Display streaming response content
         Handle stream events and update the live display accordingly.
         This method separates content and reasoning blocks for display and further processing.
 
         Args:
             stream_iterator (Iterator[Dict[str, Any]]): The stream iterator to process.
+            with_assistant_prefix (bool): Whether to display the "Assistant:" prefix.
         Returns:
             Tuple[Optional[str], Optional[str]]: The final content and reasoning texts if successful, None otherwise.
         """
-        self.console.print("Assistant:", style="bold green")
+        if with_assistant_prefix:
+            self.console.print("Assistant:", style="bold green")
         self._reset_state()  # Reset state for the new stream
         content = ""
         reasoning = ""
-        # Removed in_reasoning local variable initialization
-        cursor = itertools.cycle(["_", " "])
+        cursor = cursor_animation()
 
         with Live(console=self.console) as live:
             try:
                 for event in stream_iterator:
-                    # Pass only content and reasoning, rely on self.in_reasoning
                     content, reasoning = self._handle_event(event, content, reasoning)
 
                     if event.get("type") in (
@@ -210,35 +267,34 @@ class Printer:
                         EventTypeEnum.REASONING,
                         EventTypeEnum.REASONING_END,
                     ):
-                        # Pass only necessary variables, rely on self.in_reasoning
                         self._update_live_display(live, content, reasoning, cursor)
 
                 # Remove cursor and finalize display
-                display_text = self._format_display_text(content, reasoning)
-                live.update(Markdown(display_text, code_theme=self.code_theme))
+                live.update(self._format_display_text(content, reasoning))
                 return content, reasoning
 
             except Exception as e:
                 self.console.print(f"An error occurred during stream display: {e}", style="red")
-                display_text = self._format_display_text(content, reasoning) + " [Display Error]"
-                live.update(Markdown(markup=display_text, code_theme=self.code_theme))
                 if self.verbose:
                     traceback.print_exc()
                 return None, None
 
-    def display_normal(self, content: Optional[str], reasoning: Optional[str] = None) -> None:
+    def display_normal(
+        self, content: Optional[str], reasoning: Optional[str] = None, with_assistant_prefix: bool = True
+    ) -> None:
         """Display a complete, non-streamed response.
 
         Args:
             content (Optional[str]): The main content to display.
             reasoning (Optional[str]): The reasoning content to display.
+            with_assistant_prefix (bool): Whether to display the "Assistant:" prefix.
         """
-        self.console.print("Assistant:", style="bold green")
-
+        if with_assistant_prefix:
+            self.console.print("Assistant:", style="bold green")
         if content or reasoning:
             # Use the existing _format_display_text method
-            display_text = self._format_display_text(content or "", reasoning or "")
-            self.console.print(Markdown(display_text, code_theme=self.code_theme))
+            formatted_display = self._format_display_text(content or "", reasoning or "")
+            self.console.print(formatted_display)
             self.console.print()  # Add a newline for spacing
         else:
             self.console.print("Assistant did not provide any content.", style="yellow")

yaicli/render.py

@@ -0,0 +1,19 @@
+from typing import Any
+
+from rich.markdown import Markdown
+
+from yaicli.config import cfg
+
+
+class JustifyMarkdown(Markdown):
+    """Custom Markdown class that defaults to the configured justify value."""
+
+    def __init__(self, *args, **kwargs):
+        if "justify" not in kwargs:
+            kwargs["justify"] = cfg["JUSTIFY"]
+        super().__init__(*args, **kwargs)
+
+
+def plain_formatter(text: str, **kwargs: Any) -> str:
+    """Format the text for display, without Markdown formatting."""
+    return text

yaicli/roles.py

@@ -0,0 +1,248 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import typer
+from rich.console import Console
+from rich.prompt import Prompt
+from rich.table import Table
+
+from yaicli.config import cfg
+from yaicli.console import get_console
+from yaicli.const import DEFAULT_ROLES, ROLES_DIR, DefaultRoleNames
+from yaicli.exceptions import RoleAlreadyExistsError, RoleCreationError
+from yaicli.utils import detect_os, detect_shell, option_callback
+
+
+class Role:
+    def __init__(
+        self, name: str, prompt: str, variables: Optional[Dict[str, Any]] = None, filepath: Optional[str] = None
+    ):
+        self.name = name
+        self.prompt = prompt
+        if not variables:
+            variables = {"_os": detect_os(cfg), "_shell": detect_shell(cfg)}
+        self.variables = variables
+        self.filepath = filepath
+
+        self.prompt = self.prompt.format(**self.variables)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert Role to dictionary for serialization"""
+        return {
+            "name": self.name,
+            "prompt": self.prompt,
+        }
+
+    @classmethod
+    def from_dict(cls, role_id: str, data: Dict[str, Any], filepath: Optional[str] = None) -> "Role":
+        """Create Role object from dictionary"""
+        return cls(
+            name=data.get("name", role_id),
+            prompt=data.get("prompt", ""),
+            variables=data.get("variables", {}),
+            filepath=filepath,
+        )
+
+    def __str__(self):
+        return f"Role(name={self.name}, prompt={self.prompt[:30]}...)"
+
+
+class RoleManager:
+    roles_dir: Path = ROLES_DIR
+    console: Console = get_console()
+
+    def __init__(self):
+        self.roles: Dict[str, Role] = self._load_roles()
+
+    def _load_roles(self) -> Dict[str, Role]:
+        """Load all role configurations"""
+        roles = {}
+        self.roles_dir.mkdir(parents=True, exist_ok=True)
+
+        # Check if any role files exist
+        role_files: list[Path] = list(self.roles_dir.glob("*.json"))
+
+        if not role_files:
+            # Fast path: no existing roles, just create defaults
+            for role_id, role_config in DEFAULT_ROLES.items():
+                role_file = self.roles_dir / f"{role_id}.json"
+                filepath = str(role_file)
+                roles[role_id] = Role.from_dict(role_id, role_config, filepath)
+                with role_file.open("w", encoding="utf-8") as f:
+                    json.dump(role_config, f, indent=2)
+            return roles
+
+        # Load existing role files
+        for role_file in role_files:
+            role_id = role_file.stem
+            filepath = str(role_file)
+            try:
+                with role_file.open("r", encoding="utf-8") as f:
+                    role_data = json.load(f)
+                    roles[role_id] = Role.from_dict(role_id, role_data, filepath)
+            except Exception as e:
+                self.console.print(f"Error loading role {role_id}: {e}", style="red")
+
+        # Ensure default roles exist
+        for role_id, role_config in DEFAULT_ROLES.items():
+            if role_id not in roles:
+                role_file = self.roles_dir / f"{role_id}.json"
+                filepath = str(role_file)
+                roles[role_id] = Role.from_dict(role_id, role_config, filepath)
+                with role_file.open("w", encoding="utf-8") as f:
+                    json.dump(role_config, f, indent=2)
+
+        return roles
+
+    @classmethod
+    @option_callback
+    def print_list_option(cls, _: Any):
+        """Print the list of roles.
+        This method is a cli option callback.
+        """
+        table = Table(show_header=True, show_footer=False)
+        table.add_column("Name", style="dim")
+        table.add_column("Filepath", style="dim")
+        for file in sorted(cls.roles_dir.glob("*.json"), key=lambda f: f.stat().st_mtime):
+            table.add_row(file.stem, str(file))
+        cls.console.print(table)
+        cls.console.print("Use `ai --show-role <name>` to view a role.", style="dim")
+
+    def list_roles(self) -> list:
+        """List all available roles info"""
+        roles_list = []
+        for role_id, role in sorted(self.roles.items()):
+            roles_list.append(
+                {
+                    "id": role_id,
+                    "name": role.name,
+                    "prompt": role.prompt,
+                    "is_default": role_id in DEFAULT_ROLES,
+                    "filepath": role.filepath,
+                }
+            )
+        return roles_list
+
+    @classmethod
+    @option_callback
+    def show_role_option(cls, name: str):
+        """Show a role's prompt.
+        This method is a cli option callback.
+        """
+        self = cls()
+        role = self.get_role(name)
+        if not role:
+            self.console.print(f"Role '{name}' does not exist", style="red")
+            return
+        self.console.print(role.prompt)
+
+    def get_role(self, role_id: str) -> Optional[Role]:
+        """Get role by ID"""
+        return self.roles.get(role_id)
+
+    @classmethod
+    def check_id_ok(cls, role_id: str):
+        """Check if role exists by ID.
+        This method is a cli option callback.
+        If role does not exist, exit with error.
+        """
+        if not role_id:
+            return role_id
+        self = cls()
+        if not self.role_exists(role_id):
+            self.console.print(f"Role '{role_id}' does not exist", style="red")
+            raise typer.Abort()
+        return role_id
+
+    def role_exists(self, role_id: str) -> bool:
+        """Check if role exists"""
+        return role_id in self.roles
+
+    def save_role(self, role_id: str, role: Role) -> None:
+        """Save role configuration"""
+        try:
+            self.roles[role_id] = role
+            role_file = self.roles_dir / f"{role_id}.json"
+            role.filepath = str(role_file)
+            with role_file.open("w", encoding="utf-8") as f:
+                json.dump(role.to_dict(), f, indent=2)
+        except Exception as e:
+            raise RoleCreationError(f"RoleCreationError {e}") from e
+
+    @classmethod
+    @option_callback
+    def create_role_option(cls, name: str):
+        """Create a new role and save it to file.
+        This method is a cli option callback.
+        """
+        self = cls()
+        if name in self.roles:
+            self.console.print(f"Role '{name}' already exists", style="yellow")
+            return
+        description = Prompt.ask("Enter role description")
+
+        # Format the prompt as "You are {role_id}, {description}"
+        prompt = f"You are {name}, {description}"
+
+        role = Role(name=name, prompt=prompt)
+        self.create_role(name, role)
+        self.console.print(f"Role '{name}' created successfully", style="green")
+
+    def create_role(self, role_id: str, role: Union[Role, Dict[str, Any]]) -> None:
+        """Create a new role and save it to file"""
+        if role_id in self.roles:
+            raise RoleAlreadyExistsError(f"Role '{role_id}' already exists")
+        if isinstance(role, dict):
+            if "name" not in role or "prompt" not in role:
+                raise RoleCreationError("Role must have 'name' and 'prompt' keys")
+            # Convert dict to Role object
+            role = Role.from_dict(role_id, role)
+        self.save_role(role_id, role)
+
+    @classmethod
+    @option_callback
+    def delete_role_option(cls, name: str):
+        """Delete a role and its file.
+        This method is a cli option callback.
+        """
+        self = cls()
+        if self.delete_role(name):
+            self.console.print(f"Role '{name}' deleted successfully", style="green")
+
+    def delete_role(self, role_id: str) -> bool:
+        """Delete a role and its file"""
+        if role_id not in self.roles:
+            self.console.print(f"Role '{role_id}' does not exist", style="red")
+            return False
+
+        # Don't allow deleting default roles
+        if role_id in DEFAULT_ROLES:
+            self.console.print(f"Cannot delete default role: '{role_id}'", style="red")
+            return False
+
+        try:
+            role = self.roles[role_id]
+            if role.filepath:
+                Path(role.filepath).unlink(missing_ok=True)
+            del self.roles[role_id]
+            return True
+        except Exception as e:
+            self.console.print(f"Error deleting role: {e}", style="red")
+            return False
+
+    def get_system_prompt(self, role_id: str) -> str:
+        """Get prompt from file by role ID"""
+        role = self.get_role(role_id)
+        if not role:
+            # Fall back to default role if specified role doesn't exist
+            self.console.print(f"Role {role_id} not found, using default role", style="yellow")
+            role = self.get_role(DefaultRoleNames.DEFAULT)
+            if not role:
+                # Last resort fallback
+                default_config = DEFAULT_ROLES[DefaultRoleNames.DEFAULT]
+                role = Role.from_dict(DefaultRoleNames.DEFAULT, default_config)
+
+        # Create a copy of the role with system variables
+        system_role = Role(name=role.name, prompt=role.prompt)
+        return system_role.prompt

yaicli/utils.py

@@ -1,13 +1,32 @@
 import platform
 from os import getenv
 from os.path import basename, pathsep
-from typing import Any, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 
+import typer
 from distro import name as distro_name
 
 from yaicli.const import DEFAULT_OS_NAME, DEFAULT_SHELL_NAME
 
-T = TypeVar("T", int, float, str)
+T = TypeVar("T", int, float, str, bool)
+
+
+def option_callback(func: Callable) -> Callable:  # pragma: no cover
+    """
+    A decorator for Typer option callbacks that ensures the application exits
+    after the callback function is executed.
+
+    Args:
+        func (Callable): The callback classmethod to wrap.
+    """
+
+    def wrapper(cls, value: T) -> T:
+        if not value:
+            return value
+        func(cls, value)
+        raise typer.Exit()
+
+    return wrapper
 
 
 def detect_os(config: dict[str, Any]) -> str: