yaicli 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

yaicli/printer.py

@@ -1,176 +1,131 @@
 import time
-import traceback
-from typing import (
-    Any,
-    Dict,
-    Iterator,
-    List,
-    Optional,
-    Tuple,
-)
-
-from rich.console import Console, Group, RenderableType
-from rich.live import Live
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Iterator, List, Tuple, Union
 
-from yaicli.console import get_console
-from yaicli.const import EventTypeEnum
-from yaicli.render import JustifyMarkdown as Markdown
-from yaicli.render import plain_formatter
+from rich.console import Group, RenderableType
+from rich.live import Live
 
+from .client import RefreshLive
+from .config import Config, get_config
+from .console import YaiConsole, get_console
+from .schemas import ChatMessage
+from .render import Markdown, 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]
+if TYPE_CHECKING:
+    from .schemas import LLMResponse
 
 
+@dataclass
 class Printer:
-    """Handles printing responses to the console, including stream processing."""
-
-    _REASONING_PREFIX = "> "
-    _CURSOR_ANIMATION_SLEEP = 0.005
-
-    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["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
+    console: YaiConsole = field(default_factory=get_console)
+    config: Config = field(default_factory=get_config)
+    content_markdown: bool = True
+
+    _REASONING_PREFIX: str = "> "
+    _UPDATE_INTERVAL: float = 0.01
+
+    def __post_init__(self):
+        self.code_theme: str = self.config["CODE_THEME"]
+        self.show_reasoning: bool = self.config["SHOW_REASONING"]
+        # Set formatter for reasoning and content
+        self.reasoning_formatter = Markdown
         self.content_formatter = Markdown if self.content_markdown else plain_formatter
+        # Track if we're currently processing reasoning content
+        self.in_reasoning: bool = False
 
     def _reset_state(self) -> None:
-        """Resets the printer state for a new stream."""
+        """Reset printer state for a new stream."""
         self.in_reasoning = False
 
-    def _process_reasoning_chunk(self, chunk: str, content: str, reasoning: str) -> Tuple[str, str]:
-        """Adds a reasoning chunk to the reasoning text.
-        This method handles the processing of reasoning chunks, and update the reasoning state
-        when <think> tag is closed.
+    def _check_and_update_think_tags(self, content: str, reasoning: str) -> Tuple[str, str]:
+        """Check for <think> tags in the accumulated content and reasoning.
 
-        Args:
-            chunk (str): The reasoning chunk to process.
-            content (str): The current content text.
-            reasoning (str): The current reasoning text.
-
-        Returns:
-            Tuple[str, str]: The updated content text and reasoning text.
-        """
-        if not self.in_reasoning:
-            self.in_reasoning = True
-            reasoning = ""
-
-        tmp = chunk.replace("\n", f"\n{self._REASONING_PREFIX}")
-        tmp_reasoning = reasoning + tmp
-
-        reasoning += chunk
-        if "</think>" in tmp_reasoning:
-            self.in_reasoning = False
-            reasoning, content = reasoning.split("</think>", maxsplit=1)
-        return content, reasoning
-
-    def _process_content_chunk(self, chunk: str, content: str, reasoning: str) -> Tuple[str, str]:
-        """Adds a content chunk to the content text.
-        This method handles the processing of content chunks, and update the reasoning state
-        when <think> tag is opened.
+        This function checks the entire accumulated text for <think> tags
+        and updates state accordingly.
 
         Args:
-            chunk (str): The content chunk to process.
-            content (str): The current content text.
-            reasoning (str): The current reasoning text.
+            content: Current accumulated content text
+            reasoning: Current accumulated reasoning text
 
         Returns:
-            Tuple[str, str]: The updated content text and reasoning text.
+            Updated content and reasoning after tag processing
         """
-        if content == "":
-            chunk = chunk.lstrip()  # Remove leading whitespace from first chunk
+        # First, check if we have a <think> opener in content
+        if "<think>" in content and not self.in_reasoning:
+            parts = content.split("<think>", 1)
+            new_content = parts[0]
+            new_reasoning = parts[1]
+            self.in_reasoning = True
 
-        if self.in_reasoning:
+            # Check if the new reasoning has a </think> closer
+            if "</think>" in new_reasoning:
+                closer_parts = new_reasoning.split("</think>", 1)
+                reasoning += closer_parts[0]
+                new_content += closer_parts[1]
+                self.in_reasoning = False
+                return new_content, reasoning
+            else:
+                # No closer yet
+                reasoning += new_reasoning
+                return new_content, reasoning
+
+        # Check if we have a </think> closer in reasoning
+        if "</think>" in reasoning and self.in_reasoning:
+            parts = reasoning.split("</think>", 1)
+            new_reasoning = parts[0]
+            content += parts[1]
             self.in_reasoning = False
-        content += chunk
-
-        if content.startswith("<think>"):
-            # Remove <think> tag and leading whitespace
-            self.in_reasoning = True
-            reasoning = content[7:].lstrip()
-            content = ""  # Content starts after the initial <think> tag
+            return content, new_reasoning
 
         return content, reasoning
 
-    def _handle_event(self, event: Dict[str, Any], content: str, reasoning: str) -> Tuple[str, str]:
-        """Process a single stream event and return the updated content and reasoning.
+    def _process_chunk(self, chunk_content: str, chunk_reasoning: str, content: str, reasoning: str) -> Tuple[str, str]:
+        """Process a single chunk and update content and reasoning.
 
         Args:
-            event (Dict[str, Any]): The stream event to process.
-            content (str): The current content text (non-reasoning).
-            reasoning (str): The current reasoning text.
+            chunk_content: Content from the current chunk
+            chunk_reasoning: Reasoning from the current chunk
+            content: Current accumulated content
+            reasoning: Current accumulated reasoning
+
         Returns:
-            Tuple[str, str]: The updated content text and reasoning text.
+            Updated content and reasoning
         """
-        event_type = event.get("type")
-        chunk = event.get("chunk")
-
-        if event_type == EventTypeEnum.ERROR and self.verbose:
-            self.console.print(f"Stream error: {event.get('message')}", style="dim")
-            return content, reasoning
-
-        # Handle explicit reasoning end event
-        if event_type == EventTypeEnum.REASONING_END:
+        # Process reasoning field first (if present)
+        if chunk_reasoning:
             if self.in_reasoning:
-                self.in_reasoning = False
-            return content, reasoning
+                # Already in reasoning mode, append to reasoning
+                reasoning += chunk_reasoning
+            else:
+                # Force reasoning mode for explicit reasoning field
+                self.in_reasoning = True
+                reasoning += chunk_reasoning
 
-        if event_type in (EventTypeEnum.REASONING, EventTypeEnum.CONTENT) and chunk:
-            if event_type == EventTypeEnum.REASONING or self.in_reasoning:
-                return self._process_reasoning_chunk(str(chunk), content, reasoning)
-            return self._process_content_chunk(str(chunk), content, reasoning)
+        # Then process content field (if present)
+        if chunk_content:
+            if self.in_reasoning:
+                # In reasoning mode, append to reasoning
+                reasoning += chunk_content
+            else:
+                # Normal content mode
+                content += chunk_content
 
-        return content, reasoning
+        # Check for any <think> tags in the updated content/reasoning
+        return self._check_and_update_think_tags(content, reasoning)
 
     def _format_display_text(self, content: str, reasoning: str) -> RenderableType:
         """Format the text for display, combining content and reasoning if needed.
 
         Args:
-            content (str): The content text.
-            reasoning (str): The reasoning text.
+            content: The content text.
+            reasoning: The reasoning text.
 
         Returns:
-            RenderableType: The formatted text ready for display as a Rich renderable.
+            The formatted text ready for display as a Rich renderable.
         """
         # 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}")
@@ -182,7 +137,6 @@ class Printer:
             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)
@@ -199,102 +153,68 @@ class Printer:
         # 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
-        Sleep for a short duration to control the cursor animation speed.
-
-        Args:
-            live (Live): The live display object.
-            content (str): The current content text.
-            reasoning (str): The current reasoning text.
-            cursor (Iterator[str]): The cursor animation iterator.
-        """
+    def display_normal(
+        self, content_iterator: Iterator[Union["LLMResponse", RefreshLive]], messages: list["ChatMessage"]
+    ) -> tuple[str, str]:
+        """Process and display non-stream LLMContent, including reasoning and content parts."""
+        self._reset_state()
+        full_content = full_reasoning = ""
 
-        cursor_char = next(cursor)
+        for chunk in content_iterator:
+            if not isinstance(chunk, LLMResponse):
+                continue
 
-        # 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:
-                if reasoning.endswith("\n"):
-                    cursor_line = f"\n{self._REASONING_PREFIX}{cursor_char}"
-                else:
-                    cursor_line = cursor_char
+            # Process chunk and update content/reasoning
+            full_content, full_reasoning = self._process_chunk(
+                chunk.content or "", chunk.reasoning or "", full_content, full_reasoning
+            )
 
-                # 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 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:
-            # 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)
-
-        live.update(formatted_display)
-        time.sleep(self._CURSOR_ANIMATION_SLEEP)
+            # Display reasoning
+            if self.show_reasoning and full_reasoning:
+                reasoning = full_reasoning.replace("\n", f"\n{self._REASONING_PREFIX}")
+                self.console.print("Thinking:")
+                self.console.print(self.reasoning_formatter(reasoning))
 
-    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.
+            # Display content
+            if full_content:
+                self.console.print()
+                self.console.print(self.content_formatter(full_content))
 
-        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.
-        """
-        if with_assistant_prefix:
-            self.console.print("Assistant:", style="bold green")
-        self._reset_state()  # Reset state for the new stream
-        content = ""
-        reasoning = ""
-        cursor = cursor_animation()
-
-        with Live(console=self.console) as live:
-            try:
-                for event in stream_iterator:
-                    content, reasoning = self._handle_event(event, content, reasoning)
-
-                    if event.get("type") in (
-                        EventTypeEnum.CONTENT,
-                        EventTypeEnum.REASONING,
-                        EventTypeEnum.REASONING_END,
-                    ):
-                        self._update_live_display(live, content, reasoning, cursor)
-
-                # Remove cursor and finalize display
-                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")
-                if self.verbose:
-                    traceback.print_exc()
-                return None, None
+            messages.append(ChatMessage(role="assistant", content=full_content))
 
-    def display_normal(
-        self, content: Optional[str], reasoning: Optional[str] = None, with_assistant_prefix: bool = True
-    ) -> None:
-        """Display a complete, non-streamed response.
+        return full_content, full_reasoning
 
-        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.
-        """
-        if with_assistant_prefix:
-            self.console.print("Assistant:", style="bold green")
-        if content or reasoning:
-            # Use the existing _format_display_text method
-            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")
+    def display_stream(
+        self, stream_iterator: Iterator[Union["LLMResponse", RefreshLive]], messages: list["ChatMessage"]
+    ) -> tuple[str, str]:
+        """Process and display LLMContent stream, including reasoning and content parts."""
+        self._reset_state()
+        full_content = full_reasoning = ""
+        live = Live(console=self.console)
+        live.start()
+
+        for chunk in stream_iterator:
+            if isinstance(chunk, RefreshLive):
+                # Refresh live display when in next completion
+                live.stop()
+                messages.append(ChatMessage(role="assistant", content=full_content))
+                live = Live(console=self.console)
+                live.start()
+                # Initialize full_content and full_reasoning for the next completion
+                full_content = full_reasoning = ""
+                self._reset_state()
+                continue
+
+            # Process chunk and update content/reasoning
+            full_content, full_reasoning = self._process_chunk(
+                chunk.content or "", chunk.reasoning or "", full_content, full_reasoning
+            )
+
+            # Update display
+            formatted_display = self._format_display_text(full_content, full_reasoning)
+            live.update(formatted_display)
+            time.sleep(self._UPDATE_INTERVAL)
+
+        live.stop()
+        messages.append(ChatMessage(role="assistant", content=full_content))
+        return full_content, full_reasoning

yaicli/render.py

@@ -2,7 +2,7 @@ from typing import Any
 
 from rich.markdown import Markdown
 
-from yaicli.config import cfg
+from .config import cfg
 
 
 class JustifyMarkdown(Markdown):

yaicli/role.py

@@ -0,0 +1,231 @@
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Dict, TypeVar
+
+import typer
+from rich.table import Table
+
+from .config import cfg
+from .console import YaiConsole, get_console
+from .const import DEFAULT_ROLES, ROLES_DIR, DefaultRoleNames
+from .utils import detect_os, detect_shell, option_callback
+
+T = TypeVar("T")
+
+
+@dataclass
+class Role:
+    name: str
+    prompt: str
+    variables: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.name or not isinstance(self.name, str):
+            raise ValueError("Role must have a non-empty name")
+
+        if not self.prompt or not isinstance(self.prompt, str):
+            raise ValueError("Role must have a non-empty description")
+
+        if not self.variables:
+            self.variables = {"_os": detect_os(cfg), "_shell": detect_shell(cfg)}
+        self.prompt = self.prompt.format(**self.variables)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class RoleManager:
+    roles: Dict[str, Role] = field(default_factory=dict)
+    roles_dir: Path = ROLES_DIR
+    console: YaiConsole = get_console()
+
+    def __post_init__(self) -> None:
+        self._ensure_roles_dir()
+        self._load_default_roles()
+        self._load_user_roles()
+
+    def _ensure_roles_dir(self) -> None:
+        """Ensure the roles directory exists, and create default roles if they don't exist"""
+        self.roles_dir.mkdir(parents=True, exist_ok=True)
+        for role in DEFAULT_ROLES.values():
+            if not (self.roles_dir / f"{role['name']}.json").exists():
+                with open(self.roles_dir / f"{role['name']}.json", "w") as f:
+                    json.dump(role, f, indent=2)
+
+    def _load_default_roles(self) -> None:
+        """Load default roles"""
+        for name, role_dict in DEFAULT_ROLES.items():
+            self.roles[name] = Role(**role_dict)
+
+    def _load_user_roles(self) -> None:
+        """Load user-defined roles, user can overwrite default roles"""
+        if not self.roles_dir.exists():
+            return
+
+        for filename in self.roles_dir.glob("*.json"):
+            try:
+                with open(filename, "r") as f:
+                    role_dict = json.load(f)
+                    role = Role(**role_dict)
+                    self.roles[role.name] = role
+            except (json.JSONDecodeError, KeyError, TypeError) as e:
+                self.console.print(f"Error loading role from {filename}: {e}", style="red")
+
+    def get_role(self, name: str) -> Role:
+        """Get a role by name"""
+        if name not in self.roles:
+            raise ValueError(f"Role '{name}' does not exist.")
+        return self.roles[name]
+
+    def create_role(self, name: str, description: str) -> Role:
+        """Create and save a new role"""
+        role = Role(name=name, prompt=description)
+        self.roles[name] = role
+
+        # Save to file
+        role_path = self.roles_dir / f"{name}.json"
+        with open(role_path, "w") as f:
+            json.dump(role.to_dict(), f, indent=2)
+
+        return role
+
+    def delete_role(self, name: str) -> bool:
+        """Delete a role"""
+
+        # Delete role file
+        role_path = self.roles_dir / f"{name}.json"
+        if role_path.exists():
+            role_path.unlink()
+
+        # Delete role from memory
+        if name in self.roles:
+            del self.roles[name]
+            return True
+
+        return False
+
+    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": self.roles_dir / f"{role_id}.json",
+                }
+            )
+        return roles_list
+
+    def print_roles(self) -> None:
+        """Print all role information"""
+        table = Table("Name", "Description", "Temperature", "Top-P", title="Available Roles")
+
+        for role in self.list_roles():
+            table.add_row(
+                role.name,
+                role.prompt,
+                str(role.temperature),
+                str(role.top_p),
+            )
+
+        self.console.print(table)
+
+    @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")
+
+    @classmethod
+    @option_callback
+    def create_role_option(cls, value: str) -> None:
+        """Create role option callback"""
+        if not value:
+            return
+
+        role_manager = RoleManager()
+
+        # Check if role name already exists
+        if value in role_manager.roles:
+            cls.console.print(f"Role '{value}' already exists.", style="red")
+            return
+
+        # Get role description
+        description = typer.prompt("Enter role description")
+
+        # Create role
+        role = role_manager.create_role(value, description)
+        cls.console.print(f"Created role: {role.name}", style="green")
+
+    @classmethod
+    @option_callback
+    def delete_role_option(cls, value: str) -> None:
+        """Delete role option callback"""
+        if not value:
+            return
+
+        role_manager = RoleManager()
+
+        # Check if role exists
+        if value not in role_manager.roles:
+            cls.console.print(f"Role '{value}' does not exist.", style="yellow")
+            return
+
+        # Delete role
+        if role_manager.delete_role(value):
+            cls.console.print(f"Deleted role: {value}", style="green")
+        else:
+            cls.console.print(f"Failed to delete role: {value}", style="red")
+
+    @classmethod
+    @option_callback
+    def show_role_option(cls, value: str) -> None:
+        """Show role option callback"""
+        if not value:
+            return
+
+        role_manager = RoleManager()
+
+        # Check if role exists
+        role = role_manager.get_role(value)
+        if not role:
+            cls.console.print(f"Role '{value}' does not exist.", style="red")
+            return
+
+        # Show role information
+        cls.console.print(f"[bold]Name:[/bold] {role.name}")
+        cls.console.print(f"[bold]Description:[/bold] {role.prompt}")
+
+    @classmethod
+    def check_id_ok(cls, value: str) -> str:
+        """Check if role ID is valid option callback"""
+        if not value:
+            # Empty value is valid
+            return value
+        if value in DEFAULT_ROLES:
+            # Built-in role is valid
+            return value
+
+        role_manager = RoleManager()
+
+        if value not in role_manager.roles:
+            cls.console.print(f"Role '{value}' does not exist. Using default role.", style="red")
+            return DefaultRoleNames.DEFAULT
+
+        return value
+
+
+role_mgr = RoleManager()