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

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "yaicli"
-version = "0.1.0"
+version = "0.2.0"
 description = "A simple CLI tool to interact with LLM"
 authors = [{ name = "belingud", email = "im.victor@qq.com" }]
 readme = "README.md"

yaicli/chat_manager.py

@@ -0,0 +1,263 @@
+import json
+import time
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, TypedDict
+
+from rich import get_console
+from rich.console import Console
+
+from yaicli.config import Config
+
+
+class ChatsMap(TypedDict):
+    title: Dict[str, dict]
+    index: Dict[int, dict]
+
+
+class ChatManager(ABC):
+    """Abstract base class that defines the chat manager interface"""
+
+    @abstractmethod
+    def make_chat_title(self, prompt: Optional[str] = None) -> str:
+        """Make a chat title from a given full prompt"""
+        pass
+
+    @abstractmethod
+    def save_chat(self, history: List[Dict[str, Any]], title: Optional[str] = None) -> str:
+        """Save a chat and return the chat title"""
+        pass
+
+    @abstractmethod
+    def list_chats(self) -> List[Dict[str, Any]]:
+        """List all saved chats and return the chat list"""
+        pass
+
+    @abstractmethod
+    def refresh_chats(self) -> None:
+        """Force refresh the chat list"""
+        pass
+
+    @abstractmethod
+    def load_chat_by_index(self, index: int) -> Dict[str, Any]:
+        """Load a chat by index and return the chat data"""
+        pass
+
+    @abstractmethod
+    def load_chat_by_title(self, title: str) -> Dict[str, Any]:
+        """Load a chat by title and return the chat data"""
+        pass
+
+    @abstractmethod
+    def delete_chat(self, index: int) -> bool:
+        """Delete a chat by index and return success status"""
+        pass
+
+    @abstractmethod
+    def validate_chat_index(self, index: int) -> bool:
+        """Validate a chat index and return success status"""
+        pass
+
+
+class FileChatManager(ChatManager):
+    """File system based chat manager implementation"""
+
+    def __init__(self, config: Config, console: Optional[Console] = None):
+        self.config = config
+        self.chat_dir = Path(self.config["CHAT_HISTORY_DIR"])
+        self.chat_dir.mkdir(parents=True, exist_ok=True)
+        self.max_saved_chats = self.config["MAX_SAVED_CHATS"]
+        self._chats_map: Optional[ChatsMap] = None  # Cache for chat map
+        self.console = console or get_console()
+
+    @property
+    def chats_map(self) -> ChatsMap:
+        """Get the map of chats, loading from disk only when needed"""
+        if self._chats_map is None:
+            self._load_chats()
+        return self._chats_map or {"index": {}, "title": {}}
+
+    def make_chat_title(self, prompt: Optional[str] = None) -> str:
+        """Make a chat title from a given full prompt"""
+        if prompt:
+            return prompt[:100]
+        else:
+            return f"Chat-{int(time.time())}"
+
+    def validate_chat_index(self, index: int) -> bool:
+        """Validate a chat index and return success status"""
+        return index > 0 and index in self.chats_map["index"]
+
+    def refresh_chats(self) -> None:
+        """Force refresh the chat list from disk"""
+        self._load_chats()
+
+    def _parse_filename(self, chat_file: Path, index: int) -> Dict[str, Any]:
+        """Parse a chat filename and extract metadata"""
+        # filename: "20250421-214005-title-meaning of life"
+        filename = chat_file.stem
+        parts = filename.split("-")
+        title_str_len = 6  # "title-" marker length
+
+        # Check if the filename has the expected format
+        if len(parts) >= 4 and "title" in parts:
+            str_title_index = filename.find("title")
+            if str_title_index == -1:
+                # If "title" is not found, use full filename as the title
+                # Just in case, fallback to use fullname, but this should never happen when `len(parts) >= 4 and "title" in parts`
+                str_title_index = 0
+                title_str_len = 0
+
+            # "20250421-214005-title-meaning of life" ==> "meaning of life"
+            title = filename[str_title_index + title_str_len :]
+            date_ = parts[0]
+            time_ = parts[1]
+            # Format date
+            date_str = f"{date_[:4]}-{date_[4:6]}-{date_[6:]} {time_[:2]}:{time_[2:4]}"
+
+            # Calculate timestamp from date parts
+            try:
+                date_time_str = f"{date_}{time_}"
+                timestamp = int(datetime.strptime(date_time_str, "%Y%m%d%H%M%S").timestamp())
+            except ValueError:
+                timestamp = 0
+        else:
+            # Fallback for files that don't match expected format
+            title = filename
+            date_str = ""
+            timestamp = 0
+
+        # The actual title is stored in the JSON file, so we'll use that when loading
+        # This is just for the initial listing before the file is opened
+        return {
+            "index": index,
+            "path": str(chat_file),
+            "title": title,
+            "date": date_str,
+            "timestamp": timestamp,
+        }
+
+    def _load_chats(self) -> None:
+        """Load chats from disk into memory"""
+        chat_files = sorted(list(self.chat_dir.glob("*.json")), reverse=True)
+        chats_map: ChatsMap = {"title": {}, "index": {}}
+
+        for i, chat_file in enumerate(chat_files[: self.max_saved_chats]):
+            try:
+                info = self._parse_filename(chat_file, i + 1)
+                chats_map["title"][info["title"]] = info
+                chats_map["index"][i + 1] = info
+            except Exception as e:
+                # Log the error but continue processing other files
+                self.console.print(f"Error parsing session file {chat_file}: {e}", style="dim")
+                continue
+
+        self._chats_map = chats_map
+
+    def list_chats(self) -> List[Dict[str, Any]]:
+        """List all saved chats and return the chat list"""
+        return list(self.chats_map["index"].values())
+
+    def save_chat(self, history: List[Dict[str, Any]], title: Optional[str] = None) -> str:
+        """Save chat history to the file system, overwriting existing chats with the same title.
+
+        If no title is provided, the chat will be saved with a default title.
+        The default title is "Chat-{current timestamp}".
+
+        Args:
+            history (List[Dict[str, Any]]): The chat history to save
+            title (Optional[str]): The title of the chat provided by the user
+
+        Returns:
+            str: The title of the saved chat
+        """
+        history = history or []
+
+        save_title = title or f"Chat-{int(time.time())}"
+        save_title = self.make_chat_title(save_title)
+
+        # Check for existing session with the same title and delete it
+        existing_chat = self.chats_map["title"].get(save_title)
+        if existing_chat:
+            try:
+                existing_path = Path(existing_chat["path"])
+                existing_path.unlink()
+            except OSError as e:
+                self.console.print(
+                    f"Warning: Could not delete existing chat file {existing_chat['path']}: {e}",
+                    style="dim",
+                )
+
+        timestamp = datetime.now().astimezone().strftime("%Y%m%d-%H%M%S")
+        filename = f"{timestamp}-title-{save_title}.json"
+        filepath = self.chat_dir / filename
+
+        try:
+            with open(filepath, "w", encoding="utf-8") as f:
+                json.dump({"history": history, "title": save_title}, f, ensure_ascii=False, indent=2)
+            # Force refresh the chat list after saving
+            self.refresh_chats()
+            return save_title
+        except Exception as e:
+            self.console.print(f"Error saving chat '{save_title}': {e}", style="dim")
+            return ""
+
+    def _load_chat_data(self, chat_info: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+        """Common method to load chat data from a chat info dict"""
+        if not chat_info:
+            return {}
+
+        try:
+            chat_file = Path(chat_info["path"])
+            with open(chat_file, "r", encoding="utf-8") as f:
+                chat_data = json.load(f)
+
+            return {
+                "title": chat_data.get("title", chat_info["title"]),
+                "timestamp": chat_info["timestamp"],
+                "history": chat_data.get("history", []),
+            }
+        except FileNotFoundError:
+            self.console.print(f"Chat file not found: {chat_info['path']}", style="dim")
+            return {}
+        except json.JSONDecodeError as e:
+            self.console.print(f"Invalid JSON in chat file {chat_info['path']}: {e}", style="dim")
+            return {}
+        except Exception as e:
+            self.console.print(f"Error loading chat from {chat_info['path']}: {e}", style="dim")
+            return {}
+
+    def load_chat_by_index(self, index: int) -> Dict[str, Any]:
+        """Load a chat by index and return the chat data"""
+        if not self.validate_chat_index(index):
+            return {}
+        chat_info = self.chats_map.get("index", {}).get(index)
+        return self._load_chat_data(chat_info)
+
+    def load_chat_by_title(self, title: str) -> Dict[str, Any]:
+        """Load a chat by title and return the chat data"""
+        chat_info = self.chats_map.get("title", {}).get(title)
+        return self._load_chat_data(chat_info)
+
+    def delete_chat(self, index: int) -> bool:
+        """Delete a chat by index and return success status"""
+        if not self.validate_chat_index(index):
+            return False
+
+        chat_info = self.chats_map["index"].get(index)
+        if not chat_info:
+            return False
+
+        try:
+            chat_file = Path(chat_info["path"])
+            chat_file.unlink()
+            # Force refresh the chat list
+            self.refresh_chats()
+            return True
+        except FileNotFoundError:
+            self.console.print(f"Chat file not found: {chat_info['path']}", style="dim")
+            return False
+        except Exception as e:
+            self.console.print(f"Error deleting chat {index}: {e}", style="dim")
+            return False

yaicli/cli.py

@@ -1,5 +1,6 @@
 import subprocess
 import sys
+import time
 import traceback
 from os.path import devnull
 from pathlib import Path
@@ -17,13 +18,18 @@ from rich.panel import Panel
 from rich.prompt import Prompt
 
 from yaicli.api import ApiClient
+from yaicli.chat_manager import ChatManager, FileChatManager
 from yaicli.config import CONFIG_PATH, Config
 from yaicli.const import (
     CHAT_MODE,
     CMD_CLEAR,
+    CMD_DELETE_CHAT,
     CMD_EXIT,
     CMD_HISTORY,
+    CMD_LIST_CHATS,
+    CMD_LOAD_CHAT,
     CMD_MODE,
+    CMD_SAVE_CHAT,
     DEFAULT_CODE_THEME,
     DEFAULT_INTERACTIVE_ROUND,
     DEFAULT_OS_NAME,
@@ -42,15 +48,32 @@ class CLI:
     HISTORY_FILE = Path("~/.yaicli_history").expanduser()
 
     def __init__(
-        self, verbose: bool = False, api_client: Optional[ApiClient] = None, printer: Optional[Printer] = None
+        self,
+        verbose: bool = False,
+        stdin: Optional[str] = None,
+        api_client: Optional[ApiClient] = None,
+        printer: Optional[Printer] = None,
+        chat_manager: Optional[ChatManager] = None,
     ):
         self.verbose = verbose
+        self.stdin = stdin
         self.console = get_console()
         self.bindings = KeyBindings()
         self.config: Config = Config(self.console)
         self.current_mode: str = TEMP_MODE
+
         self.history = []
         self.interactive_max_history = self.config.get("INTERACTIVE_MAX_HISTORY", DEFAULT_INTERACTIVE_ROUND)
+        self.chat_title = None
+        self.chat_start_time = None
+        self.is_temp_session = True
+
+        # Get and create chat history directory from configuration
+        self.chat_history_dir = Path(self.config["CHAT_HISTORY_DIR"])
+        self.chat_history_dir.mkdir(parents=True, exist_ok=True)
+
+        # Initialize session manager
+        self.chat_manager = chat_manager or FileChatManager(self.config)
 
         # Detect OS and Shell if set to auto
         if self.config.get("OS_NAME") == DEFAULT_OS_NAME:
@@ -63,7 +86,7 @@ class CLI:
             self.console.print(f"Config file path: {CONFIG_PATH}")
             for key, value in self.config.items():
                 display_value = "****" if key == "API_KEY" and value else value
-                self.console.print(f"  {key:<16}: {display_value}")
+                self.console.print(f"  {key:<17}: {display_value}")
             self.console.print(Markdown("---", code_theme=self.config.get("CODE_THEME", DEFAULT_CODE_THEME)))
 
         self.api_client = api_client or ApiClient(self.config, self.console, self.verbose)
@@ -77,19 +100,104 @@ class CLI:
             self.session = PromptSession(key_bindings=self.bindings)
         finally:
             if _origin_stderr:
+                sys.stderr.flush()
                 sys.stderr.close()
                 sys.stderr = _origin_stderr
 
     def get_prompt_tokens(self) -> List[Tuple[str, str]]:
         """Return prompt tokens for current mode"""
-        qmark = "💬" if self.current_mode == CHAT_MODE else "🚀" if self.current_mode == EXEC_MODE else "📝"
-        return [("class:qmark", f" {qmark} "), ("class:prompt", "> ")]
+        mode_icon = "💬" if self.current_mode == CHAT_MODE else "🚀" if self.current_mode == EXEC_MODE else "📝"
+        return [("class:qmark", f" {mode_icon} "), ("class:prompt", "> ")]
 
     def _check_history_len(self) -> None:
         """Check history length and remove oldest messages if necessary"""
         target_len = self.interactive_max_history * 2
         if len(self.history) > target_len:
             self.history = self.history[-target_len:]
+            if self.verbose:
+                self.console.print(f"History trimmed to {target_len} messages.", style="dim")
+
+    def _save_chat(self, title: Optional[str] = None) -> None:
+        """Save current chat history to a file using session manager."""
+        saved_title = self.chat_manager.save_chat(self.history, title)
+
+        if not saved_title:
+            self.console.print("Failed to save chat.", style="bold red")
+            return
+
+        # Session list will be refreshed automatically by the save method
+        self.console.print(f"Chat saved as: {saved_title}", style="bold green")
+
+        # If this was a temporary session, mark it as non-temporary now that it's saved
+        if self.is_temp_session:
+            self.is_temp_session = False
+            self.chat_title = saved_title
+            self.chat_start_time = int(time.time())
+            self.console.print(
+                "Session is now marked as persistent and will be auto-saved on exit.", style="bold green"
+            )
+
+    def _list_chats(self) -> None:
+        """List all saved chat sessions using session manager."""
+        chats = self.chat_manager.list_chats()
+
+        if not chats:
+            self.console.print("No saved chats found.", style="yellow")
+            return
+
+        self.console.print("Saved Chats:", style="bold underline")
+        for chat in chats:
+            index = chat["index"]
+            title = chat["title"]
+            date = chat.get("date", "")
+
+            if date:
+                self.console.print(f"[dim]{index}.[/dim] [bold blue]{title}[/bold blue] - {date}")
+            else:
+                self.console.print(f"[dim]{index}.[/dim] [bold blue]{title}[/bold blue]")
+
+    def _refresh_chats(self) -> None:
+        """Force refresh the chat list."""
+        self.chat_manager.refresh_chats()
+
+    def _load_chat_by_index(self, index: int) -> bool:
+        """Load a chat session by its index using session manager."""
+        if not self.chat_manager.validate_chat_index(index):
+            self.console.print("Invalid chat index.", style="bold red")
+            return False
+
+        chat_data = self.chat_manager.load_chat_by_index(index)
+
+        if not chat_data:
+            self.console.print("Invalid chat index or chat not found.", style="bold red")
+            return False
+
+        self.history = chat_data.get("history", [])
+        self.chat_title = chat_data.get("title")
+        self.chat_start_time = chat_data.get("timestamp", int(time.time()))
+        self.is_temp_session = False
+
+        self.console.print(f"Loaded chat: {self.chat_title}", style="bold green")
+        return True
+
+    def _delete_chat_by_index(self, index: int) -> bool:
+        """Delete a chat session by its index using session manager."""
+        if not self.chat_manager.validate_chat_index(index):
+            self.console.print("Invalid chat index.", style="bold red")
+            return False
+
+        chat_data = self.chat_manager.load_chat_by_index(index)
+
+        if not chat_data:
+            self.console.print("Invalid chat index or chat not found.", style="bold red")
+            return False
+
+        if self.chat_manager.delete_chat(index):
+            self.console.print(f"Deleted chat: {chat_data['title']}", style="bold green")
+            return True
+        else:
+            self.console.print(f"Failed to delete chat: {chat_data['title']}", style="bold red")
+            return False
 
     def _handle_special_commands(self, user_input: str) -> Optional[bool]:
         """Handle special command return: True-continue loop, False-exit loop, None-non-special command"""
@@ -105,17 +213,52 @@ class CLI:
                 self.console.print("History is empty.", style="yellow")
             else:
                 self.console.print("Chat History:", style="bold underline")
-                code_theme = self.config.get("CODE_THEME", "monokai")
                 for i in range(0, len(self.history), 2):
                     user_msg = self.history[i]
                     assistant_msg = self.history[i + 1] if (i + 1) < len(self.history) else None
                     self.console.print(f"[dim]{i // 2 + 1}[/dim] [bold blue]User:[/bold blue] {user_msg['content']}")
                     if assistant_msg:
-                        md = Markdown(assistant_msg["content"], code_theme=code_theme)
+                        md = Markdown(assistant_msg["content"], code_theme=self.config["CODE_THEME"])
                         padded_md = Padding(md, (0, 0, 0, 4))
                         self.console.print("    Assistant:", style="bold green")
                         self.console.print(padded_md)
             return True
+
+        # Handle /save command - optional title parameter
+        if command.startswith(CMD_SAVE_CHAT):
+            parts = command.split(maxsplit=1)
+            title = parts[1] if len(parts) > 1 else self.chat_title
+            self._save_chat(title)
+            return True
+
+        # Handle /load command - requires index parameter
+        if command.startswith(CMD_LOAD_CHAT):
+            parts = command.split(maxsplit=1)
+            if len(parts) == 2 and parts[1].isdigit():
+                # Try to parse as an index first
+                index = int(parts[1])
+                self._load_chat_by_index(index=index)
+            else:
+                self.console.print(f"Usage: {CMD_LOAD_CHAT} <index>", style="yellow")
+                self._list_chats()
+            return True
+
+        # Handle /delete command - requires index parameter
+        if command.startswith(CMD_DELETE_CHAT):
+            parts = command.split(maxsplit=1)
+            if len(parts) == 2 and parts[1].isdigit():
+                index = int(parts[1])
+                self._delete_chat_by_index(index=index)
+            else:
+                self.console.print(f"Usage: {CMD_DELETE_CHAT} <index>", style="yellow")
+                self._list_chats()
+            return True
+
+        # Handle /list command to list saved chats
+        if command == CMD_LIST_CHATS:
+            self._list_chats()
+            return True
+
         # Handle /mode command
         if command.startswith(CMD_MODE):
             parts = command.split(maxsplit=1)
@@ -176,9 +319,10 @@ class CLI:
     def get_system_prompt(self) -> str:
         """Return system prompt for current mode"""
         prompt_template = SHELL_PROMPT if self.current_mode == EXEC_MODE else DEFAULT_PROMPT
-        return prompt_template.format(
-            _os=self.config.get("OS_NAME", "Unknown OS"), _shell=self.config.get("SHELL_NAME", "Unknown Shell")
-        )
+        stdin = f"\n\nSTDIN\n{self.stdin}" if self.stdin else ""
+        if self.verbose and stdin:
+            self.console.print("Added STDIN to prompt", style="dim")
+        return prompt_template.format(_os=self.config["OS_NAME"], _shell=self.config["SHELL_NAME"]) + stdin
 
     def _build_messages(self, user_input: str) -> List[dict]:
         """Build the list of messages for the API call."""
@@ -251,9 +395,25 @@ class CLI:
             style="bold cyan",
         )
         self.console.print("Welcome to YAICLI!", style="bold")
+
+        # Display session type
+        if self.is_temp_session:
+            self.console.print("Current: [bold yellow]Temporary Session[/bold yellow] (use /save to make persistent)")
+        else:
+            self.console.print(
+                f"Current: [bold green]Persistent Session[/bold green]{f': {self.chat_title}' if self.chat_title else ''}"
+            )
+
         self.console.print("Press [bold yellow]TAB[/bold yellow] to switch mode")
         self.console.print(f"{CMD_CLEAR:<19}: Clear chat history")
         self.console.print(f"{CMD_HISTORY:<19}: Show chat history")
+        self.console.print(f"{CMD_LIST_CHATS:<19}: List saved chats")
+        save_cmd = f"{CMD_SAVE_CHAT} <title>"
+        self.console.print(f"{save_cmd:<19}: Save current chat")
+        load_cmd = f"{CMD_LOAD_CHAT} <index>"
+        self.console.print(f"{load_cmd:<19}: Load a saved chat")
+        delete_cmd = f"{CMD_DELETE_CHAT} <index>"
+        self.console.print(f"{delete_cmd:<19}: Delete a saved chat")
         cmd_exit = f"{CMD_EXIT}|Ctrl+D|Ctrl+C"
         self.console.print(f"{cmd_exit:<19}: Exit")
         cmd_mode = f"{CMD_MODE} {CHAT_MODE}|{EXEC_MODE}"
@@ -264,7 +424,7 @@ class CLI:
         self.prepare_chat_loop()
         self._print_welcome_message()
         while True:
-            self.console.print(Markdown("---", code_theme=self.config.get("CODE_THEME", "monokai")))
+            self.console.print(Markdown("---", code_theme=self.config["CODE_THEME"]))
             try:
                 user_input = self.session.prompt(self.get_prompt_tokens)
                 user_input = user_input.strip()
@@ -279,6 +439,11 @@ class CLI:
                     break
             except (KeyboardInterrupt, EOFError):
                 break
+
+        # Auto-save chat history when exiting if there are messages and not a temporary session
+        if not self.is_temp_session:
+            self._save_chat(self.chat_title)
+
         self.console.print("\nExiting YAICLI... Goodbye!", style="bold green")
 
     def _run_once(self, prompt_arg: str, is_shell_mode: bool) -> None:
@@ -310,6 +475,9 @@ class CLI:
         except Exception as e:
             self.console.print(f"[red]Error initializing prompt session history: {e}[/red]")
             self.session = PromptSession(key_bindings=self.bindings)
+        if self.chat_title:
+            chat_info = self.chat_manager.load_chat_by_title(self.chat_title)
+            self.history = chat_info.get("history", [])
 
     def _setup_key_bindings(self) -> None:
         """Setup keyboard shortcuts (e.g., TAB for mode switching)."""
@@ -325,6 +493,20 @@ class CLI:
                 self.console.print("[bold red]Error:[/bold red] API key not found. Cannot start chat mode.")
                 return
             self.current_mode = CHAT_MODE
+
+            # Only set chat start time and title if a prompt is provided
+            if prompt:
+                self.chat_start_time = int(time.time())
+                self.chat_title = self.chat_manager.make_chat_title(prompt)
+                self.console.print(f"Chat title: [bold]{self.chat_title}[/bold]")
+                self.is_temp_session = False
+            else:
+                # Mark as temporary session if no prompt
+                self.is_temp_session = True
+                self.console.print(
+                    "Starting a temporary chat session (will not be saved automatically)", style="yellow"
+                )
+
             self._run_repl()
         elif prompt:
             self._run_once(prompt, shell)