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

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "yaicli"
-version = "0.0.19"
+version = "0.2.0"
 description = "A simple CLI tool to interact with LLM"
 authors = [{ name = "belingud", email = "im.victor@qq.com" }]
 readme = "README.md"
@@ -44,13 +44,19 @@ Repository = "https://github.com/belingud/yaicli"
 Documentation = "https://github.com/belingud/yaicli"
 
 [project.scripts]
-ai = "yaicli:app"
+ai = "yaicli.entry:app"
+yai = "yaicli.entry:app"
 
 [tool.uv]
 resolution = "highest"
 
 [dependency-groups]
-dev = ["bump2version>=1.0.1", "pytest>=8.3.5", "ruff>=0.11.2"]
+dev = [
+    "bump2version>=1.0.1",
+    "pytest>=8.3.5",
+    "pytest-cov>=6.1.1",
+    "ruff>=0.11.2",
+]
 
 [tool.isort]
 profile = "black"
@@ -65,4 +71,4 @@ build-backend = "hatchling.build"
 
 [tool.hatch.build]
 exclude = ["test_*.py", "tests/*", ".gitignore"]
-include = ["yaicli.py", "pyproject.toml"]
+include = ["yaicli", "pyproject.toml"]

yaicli/api.py

@@ -0,0 +1,324 @@
+import json
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+import httpx
+import jmespath
+from rich.console import Console
+
+from yaicli.const import (
+    DEFAULT_BASE_URL,
+    DEFAULT_COMPLETION_PATH,
+    DEFAULT_MAX_TOKENS,
+    DEFAULT_MODEL,
+    DEFAULT_TEMPERATURE,
+    DEFAULT_TIMEOUT,
+    DEFAULT_TOP_P,
+    EventTypeEnum,
+)
+
+
+def parse_stream_line(line: Union[bytes, str], console: Console, verbose: bool) -> Optional[dict]:
+    """(Helper Function) Parse a single line from the SSE stream response."""
+    line_str: str
+    if isinstance(line, bytes):
+        try:
+            line_str = line.decode("utf-8")
+        except UnicodeDecodeError:
+            if verbose:
+                console.print(f"Warning: Could not decode stream line bytes: {line!r}", style="yellow")
+            return None
+    elif isinstance(line, str):
+        line_str = line
+    else:
+        # Handle unexpected line types
+        if verbose:
+            console.print(f"Warning: Received unexpected line type: {type(line)}", style="yellow")
+        return None
+
+    line_str = line_str.strip()
+    if not line_str or not line_str.startswith("data: "):
+        return None
+
+    data_part = line_str[6:]
+    if data_part.lower() == "[done]":
+        return {"done": True}  # Use a specific dictionary to signal DONE
+
+    try:
+        json_data = json.loads(data_part)
+        if not isinstance(json_data, dict) or "choices" not in json_data:
+            if verbose:
+                console.print(f"Warning: Invalid stream data format (missing 'choices'): {data_part}", style="yellow")
+            return None
+        return json_data
+    except json.JSONDecodeError:
+        console.print("Error decoding response JSON", style="red")
+        if verbose:
+            console.print(f"Invalid JSON data: {data_part}", style="red")
+        return None
+
+
+class ApiClient:
+    """Handles communication with the LLM API."""
+
+    def __init__(self, config: Dict[str, Any], console: Console, verbose: bool, client: Optional[httpx.Client] = None):
+        """Initialize the API client with configuration."""
+        self.config = config
+        self.console = console
+        self.verbose = verbose
+        self.base_url = str(config.get("BASE_URL", DEFAULT_BASE_URL))
+        self.completion_path = str(config.get("COMPLETION_PATH", DEFAULT_COMPLETION_PATH))
+        self.api_key = str(config.get("API_KEY", ""))
+        self.model = str(config.get("MODEL", DEFAULT_MODEL))
+        self.timeout = self.config.get("TIMEOUT", DEFAULT_TIMEOUT)
+        self.client = client or httpx.Client(timeout=self.config.get("TIMEOUT", DEFAULT_TIMEOUT))
+
+    def _prepare_request_body(self, messages: List[Dict[str, str]], stream: bool) -> Dict[str, Any]:
+        """Prepare the common request body for API calls."""
+        return {
+            "messages": messages,
+            "model": self.model,
+            "stream": stream,
+            "temperature": self.config.get("TEMPERATURE", DEFAULT_TEMPERATURE),
+            "top_p": self.config.get("TOP_P", DEFAULT_TOP_P),
+            "max_tokens": self.config.get("MAX_TOKENS", DEFAULT_MAX_TOKENS),
+        }
+
+    def _handle_api_error(self, e: httpx.HTTPError) -> None:
+        """Handle and print HTTP errors consistently."""
+        if isinstance(e, httpx.TimeoutException):
+            self.console.print(f"Error: API request timed out after {self.timeout} seconds. {e}", style="red")
+        elif isinstance(e, httpx.HTTPStatusError):
+            self.console.print(f"Error calling API: {e.response.status_code} {e.response.reason_phrase}", style="red")
+            if self.verbose:
+                self.console.print(f"Response Text: {e.response.text}")
+        elif isinstance(e, httpx.RequestError):
+            api_url = self.get_completion_url()
+            self.console.print(f"Error: Could not connect to API endpoint '{api_url}'. {e}", style="red")
+        else:
+            self.console.print(f"An unexpected HTTP error occurred: {e}", style="red")
+
+    def get_completion_url(self) -> str:
+        """Get the full completion URL."""
+        base_url = self.base_url.rstrip("/")
+        completion_path = self.completion_path.lstrip("/")
+        return f"{base_url}/{completion_path}"
+
+    def get_headers(self) -> Dict[str, str]:
+        """Get the request headers."""
+        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+    def _process_completion_response(self, response_json: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]:
+        """Process the JSON response from a non-streamed completion request."""
+        answer_path = self.config["ANSWER_PATH"]
+        message_path = answer_path.rsplit(".", 1)[0]
+
+        # Extract content and reasoning using JMESPath
+        content = jmespath.search(answer_path, response_json)
+        message = jmespath.search(message_path, response_json)
+        reasoning = self._get_reasoning_content(
+            message
+        )  # Reuse reasoning extraction if applicable to the whole message
+
+        # Process string content and extract reasoning from <think> tags if present
+        if isinstance(content, str):
+            content = content.lstrip()
+            if content.startswith("<think>"):
+                think_end = content.find("</think>")
+                if think_end != -1:
+                    # Extract reasoning from <think> tag only if not already found via message path
+                    if reasoning is None:
+                        reasoning = content[7:think_end].strip()  # Start after <think>
+                    # Remove the <think> block from the main content
+                    content = content[think_end + 8 :].strip()  # Start after </think>
+            # If it doesn't start with <think>, or if </think> wasn't found, return content as is
+            return content, reasoning
+        elif content:
+            self.console.print(
+                f"Warning: Unexpected content type from API: {type(content)}. Path: {answer_path}", style="yellow"
+            )
+            # Attempt to convert unexpected content to string, return existing reasoning
+            return str(content), reasoning
+        else:
+            self.console.print(f"Warning: Could not extract content using JMESPath '{answer_path}'.", style="yellow")
+            if self.verbose:
+                self.console.print(f"API Response: {response_json}")
+            return None, reasoning
+
+    def completion(self, messages: List[Dict[str, str]]) -> Tuple[Optional[str], Optional[str]]:
+        """Get a complete non-streamed response from the API."""
+        url = self.get_completion_url()
+        body = self._prepare_request_body(messages, stream=False)
+        headers = self.get_headers()
+
+        try:
+            response = self.client.post(url, json=body, headers=headers)
+            response.raise_for_status()
+            response_json = response.json()
+            # Delegate processing to the helper method
+            return self._process_completion_response(response_json)
+
+        except httpx.HTTPError as e:
+            self._handle_api_error(e)
+            return None, None
+
+    def _handle_http_error(self, e: httpx.HTTPStatusError) -> Dict[str, Any]:
+        """Handle HTTP errors during streaming and return an error event.
+
+        Args:
+            e: The HTTP status error that occurred
+
+        Returns:
+            An error event dictionary to be yielded to the client
+        """
+        error_body = e.response.read()
+        self._handle_api_error(e)
+
+        try:
+            error_json = json.loads(error_body)
+            error_message = error_json.get("error", {}).get("message")
+        except (json.JSONDecodeError, AttributeError):
+            error_message = None
+
+        if not error_message:
+            error_message = error_body.decode() if error_body else str(e)
+
+        return {"type": EventTypeEnum.ERROR, "message": error_message}
+
+    def _process_stream_chunk(
+        self, parsed_data: Dict[str, Any], in_reasoning: bool
+    ) -> Iterator[Tuple[Dict[str, Any], bool]]:
+        """Process a single chunk from the stream and yield events with updated reasoning state.
+
+        Args:
+            parsed_data: The parsed JSON data from a stream line
+            in_reasoning: Whether we're currently in a reasoning state
+
+        Yields:
+            A tuple containing:
+                - An event dictionary to yield to the client
+                - The updated reasoning state
+        """
+        # Handle stream errors
+        if "error" in parsed_data:
+            error_msg = parsed_data["error"].get("message", "Unknown error in stream data")
+            self.console.print(f"Error in stream data: {error_msg}", style="red")
+            yield {"type": EventTypeEnum.ERROR, "message": error_msg}, in_reasoning
+            return
+
+        # Get and validate the choice
+        choices = parsed_data.get("choices", [])
+        if not choices or not isinstance(choices, list):
+            if self.verbose:
+                self.console.print(f"Skipping stream chunk with no choices: {parsed_data}", style="dim")
+            return
+
+        choice = choices[0]
+        if not isinstance(choice, dict):
+            if self.verbose:
+                self.console.print(f"Skipping stream chunk with invalid choice structure: {choice}", style="dim")
+            return
+
+        # Get content from delta
+        delta = choice.get("delta", {})
+        if not isinstance(delta, dict):
+            if self.verbose:
+                self.console.print(f"Skipping stream chunk with invalid delta structure: {delta}", style="dim")
+            return
+
+        # Process content
+        reason = self._get_reasoning_content(delta)
+        content_chunk = delta.get("content", "")
+        finish_reason = choice.get("finish_reason")
+
+        # Yield events based on content type
+        if reason is not None:
+            in_reasoning = True
+            yield {"type": EventTypeEnum.REASONING, "chunk": reason}, in_reasoning
+        elif in_reasoning and content_chunk and isinstance(content_chunk, str):
+            # Signal the end of reasoning before yielding content
+            in_reasoning = False
+            yield {"type": EventTypeEnum.REASONING_END, "chunk": ""}, in_reasoning
+            yield {"type": EventTypeEnum.CONTENT, "chunk": content_chunk}, in_reasoning
+        elif content_chunk and isinstance(content_chunk, str):
+            yield {"type": EventTypeEnum.CONTENT, "chunk": content_chunk}, in_reasoning
+
+        if finish_reason:
+            yield {"type": EventTypeEnum.FINISH, "reason": finish_reason}, in_reasoning
+
+    def stream_completion(self, messages: List[Dict[str, str]]) -> Iterator[Dict[str, Any]]:
+        """Connect to the API and yield parsed stream events.
+
+        This method handles the streaming API connection and processes the response,
+        yielding events that can be consumed by the client. It handles various types
+        of content including regular content and reasoning content.
+
+        Args:
+            messages: The list of message dictionaries to send to the API
+
+        Yields:
+            Event dictionaries with the following structure:
+                - type: The event type (from EventTypeEnum)
+                - chunk/message/reason: The content of the event
+        """
+        url = self.get_completion_url()
+        body = self._prepare_request_body(messages, stream=True)
+        headers = self.get_headers()
+        in_reasoning = False
+
+        try:
+            with self.client.stream("POST", url, json=body, headers=headers) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as e:
+                    yield self._handle_http_error(e)
+                    return
+
+                # Process the stream line by line
+                for line in response.iter_lines():
+                    parsed_data = parse_stream_line(line, self.console, self.verbose)
+                    if parsed_data is None:
+                        continue
+                    if parsed_data.get("done"):
+                        break
+
+                    # Process chunks and yield events
+                    for event, updated_state in self._process_stream_chunk(parsed_data, in_reasoning):
+                        in_reasoning = updated_state
+                        # event: {type: str, Optional[chunk]: str, Optional[message]: str, Optional[reason]: str}
+                        yield event
+
+        except httpx.HTTPError as e:
+            self._handle_api_error(e)
+            yield {"type": EventTypeEnum.ERROR, "message": str(e)}
+        except Exception as e:
+            self.console.print(f"An unexpected error occurred during streaming: {e}", style="red")
+            if self.verbose:
+                import traceback
+
+                traceback.print_exc()
+            yield {"type": EventTypeEnum.ERROR, "message": f"Unexpected stream error: {e}"}
+
+    def _get_reasoning_content(self, delta: dict) -> Optional[str]:
+        """Extract reasoning content from delta if available based on specific keys.
+
+        This method checks for various keys that might contain reasoning content
+        in different API implementations.
+
+        Args:
+            delta: The delta dictionary from the API response
+
+        Returns:
+            The reasoning content string if found, None otherwise
+        """
+        if not delta:
+            return None
+        # reasoning_content: deepseek/infi-ai
+        # reasoning: openrouter
+        # <think> block implementation not in here
+        for key in ("reasoning_content", "reasoning", "metadata"):
+            # Check if the key exists and its value is a non-empty string
+            value = delta.get(key)
+            if isinstance(value, str) and value:
+                return value
+
+        return None  # Return None if no relevant key with a string value is found

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