jarvis-ai-assistant 0.7.8__py3-none-any.whl → 1.0.2__py3-none-any.whl

jarvis/jarvis_platform/base.py

@@ -1,40 +1,52 @@
+import os
+
 # -*- coding: utf-8 -*-
 import re
-import os
+from abc import ABC
+from abc import abstractmethod
 from datetime import datetime
-from abc import ABC, abstractmethod
 from types import TracebackType
-from typing import Dict, Generator, List, Optional, Tuple, Type
-
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Type
+
+from rich import box
+from rich.live import Live
+from rich.panel import Panel
+from rich.status import Status
+from rich.text import Text
 from typing_extensions import Self
 
-from rich import box  # type: ignore
-from rich.live import Live  # type: ignore
-from rich.panel import Panel  # type: ignore
-from rich.status import Status  # type: ignore
-from rich.text import Text  # type: ignore
-
-from jarvis.jarvis_utils.config import (
-    get_pretty_output,
-    is_print_prompt,
-    is_immediate_abort,
-    is_save_session_history,
-    get_data_dir,
-    get_max_input_token_count,
-    get_conversation_turn_threshold,
-)
-from jarvis.jarvis_utils.globals import set_in_chat, get_interrupt, console
 import jarvis.jarvis_utils.globals as G
-from jarvis.jarvis_utils.output import OutputType, PrettyOutput  # 保留用于语法高亮
-from jarvis.jarvis_utils.tag import ct, ot
-from jarvis.jarvis_utils.utils import while_success, while_true
+from jarvis.jarvis_utils.config import get_cheap_max_input_token_count
+from jarvis.jarvis_utils.config import get_conversation_turn_threshold
+from jarvis.jarvis_utils.config import get_data_dir
+from jarvis.jarvis_utils.config import get_max_input_token_count
+from jarvis.jarvis_utils.config import get_pretty_output
+from jarvis.jarvis_utils.config import get_smart_max_input_token_count
+from jarvis.jarvis_utils.config import is_immediate_abort
+from jarvis.jarvis_utils.config import is_print_prompt
+from jarvis.jarvis_utils.config import is_save_session_history
 from jarvis.jarvis_utils.embedding import get_context_token_count
+from jarvis.jarvis_utils.globals import console
+from jarvis.jarvis_utils.globals import get_interrupt
+from jarvis.jarvis_utils.globals import set_in_chat
+from jarvis.jarvis_utils.output import OutputType  # 保留用于语法高亮
+from jarvis.jarvis_utils.output import PrettyOutput
+from jarvis.jarvis_utils.tag import ct
+from jarvis.jarvis_utils.tag import ot
+from jarvis.jarvis_utils.utils import while_success
+from jarvis.jarvis_utils.utils import while_true
 
 
 class BasePlatform(ABC):
     """大语言模型基类"""
 
-    def __init__(self):
+    def __init__(self, llm_config: Optional[Dict[str, Any]] = None):
         """初始化模型"""
         self.suppress_output = True  # 添加输出控制标志
         self.web = False  # 添加web属性，默认false
@@ -42,6 +54,7 @@ class BasePlatform(ABC):
         self.model_group: Optional[str] = None
         self._session_history_file: Optional[str] = None
         self._conversation_turn = 0  # 对话轮次计数器
+        self.platform_type: str = "normal"  # 平台类型：normal/cheap/smart
 
     def __enter__(self) -> Self:
         """进入上下文管理器"""
@@ -84,21 +97,21 @@ class BasePlatform(ABC):
 
     def _format_progress_bar(self, percent: float, width: int = 20) -> str:
         """格式化进度条字符串
-        
+
         参数:
             percent: 百分比 (0-100)
             width: 进度条宽度（字符数）
-            
+
         返回:
             str: 格式化的进度条字符串
         """
         # 限制百分比范围
         percent = max(0, min(100, percent))
-        
+
         # 计算填充的字符数
         filled = int(width * percent / 100)
         empty = width - filled
-        
+
         # 根据百分比选择颜色
         if percent >= 90:
             color = "red"
@@ -106,18 +119,20 @@ class BasePlatform(ABC):
             color = "yellow"
         else:
             color = "green"
-        
+
         # 构建进度条：使用 █ 表示已填充，░ 表示未填充
         bar = "█" * filled + "░" * empty
-        
+
         return f"[{color}]{bar}[/{color}]"
 
-    def _get_token_usage_info(self, current_response: str = "") -> Tuple[float, str, str]:
+    def _get_token_usage_info(
+        self, current_response: str = ""
+    ) -> Tuple[float, str, str]:
         """获取 token 使用信息
-        
+
         参数:
             current_response: 当前响应内容（用于计算流式输出时的 token）
-            
+
         返回:
             Tuple[float, str, str]: (usage_percent, percent_color, progress_bar)
         """
@@ -125,8 +140,8 @@ class BasePlatform(ABC):
             history_tokens = self.get_used_token_count()
             current_response_tokens = get_context_token_count(current_response)
             total_tokens = history_tokens + current_response_tokens
-            max_tokens = get_max_input_token_count(self.model_group)
-            
+            max_tokens = self._get_platform_max_input_token_count()
+
             if max_tokens > 0:
                 usage_percent = (total_tokens / max_tokens) * 100
                 if usage_percent >= 90:
@@ -142,10 +157,14 @@ class BasePlatform(ABC):
             return 0.0, "green", ""
 
     def _update_panel_subtitle_with_token(
-        self, panel: Panel, response: str, is_completed: bool = False, duration: float = 0.0
+        self,
+        panel: Panel,
+        response: str,
+        is_completed: bool = False,
+        duration: float = 0.0,
     ) -> None:
         """更新面板的 subtitle，包含 token 使用信息
-        
+
         参数:
             panel: 要更新的面板
             response: 当前响应内容
@@ -153,13 +172,17 @@ class BasePlatform(ABC):
             duration: 耗时（秒）
         """
         from datetime import datetime
-        
+
         current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         try:
-            usage_percent, percent_color, progress_bar = self._get_token_usage_info(response)
-            max_tokens = get_max_input_token_count(self.model_group)
-            total_tokens = self.get_used_token_count() + get_context_token_count(response)
-            
+            usage_percent, percent_color, progress_bar = self._get_token_usage_info(
+                response
+            )
+            max_tokens = self._get_platform_max_input_token_count()
+            total_tokens = self.get_used_token_count() + get_context_token_count(
+                response
+            )
+
             threshold = get_conversation_turn_threshold()
             if is_completed:
                 if max_tokens > 0 and progress_bar:
@@ -188,21 +211,21 @@ class BasePlatform(ABC):
 
     def _chat_with_pretty_output(self, message: str, start_time: float) -> str:
         """使用 pretty output 模式进行聊天
-        
+
         参数:
             message: 用户消息
             start_time: 开始时间
-            
+
         返回:
             str: 模型响应
         """
         import time
-        
+
         chat_iterator = self.chat(message)
         first_chunk = None
 
         with Status(
-            f"🤔 {(G.current_agent_name + ' · ') if G.current_agent_name else ''}{self.name()} 正在思考中...",
+            f"🤔 {(G.get_current_agent_name() + ' · ') if G.get_current_agent_name() else ''}{self.name()} 正在思考中...",
             spinner="dots",
             console=console,
         ):
@@ -218,7 +241,7 @@ class BasePlatform(ABC):
         text_content = Text(overflow="fold")
         panel = Panel(
             text_content,
-            title=f"[bold cyan]{(G.current_agent_name + ' · ') if G.current_agent_name else ''}{self.name()}[/bold cyan]",
+            title=f"[bold cyan]{(G.get_current_agent_name() + ' · ') if G.get_current_agent_name() else ''}{self.name()}[/bold cyan]",
             subtitle="[yellow]正在回答... (按 Ctrl+C 中断)[/yellow]",
             border_style="bright_blue",
             box=box.ROUNDED,
@@ -227,14 +250,17 @@ class BasePlatform(ABC):
 
         response = ""
         last_subtitle_update_time = time.time()
-        subtitle_update_interval = 3  # subtitle 更新间隔（秒），减少更新频率避免重复渲染标题
+        subtitle_update_interval = (
+            3  # subtitle 更新间隔（秒），减少更新频率避免重复渲染标题
+        )
         update_count = 0  # 更新计数器，用于控制 subtitle 更新频率
-        with Live(panel, refresh_per_second=4, transient=False) as live:
+        with Live(panel, refresh_per_second=4, transient=True) as live:
+
             def _update_panel_content(content: str, update_subtitle: bool = False):
                 nonlocal response, last_subtitle_update_time, update_count
                 text_content.append(content, style="bright_white")
                 update_count += 1
-                
+
                 # Scrolling Logic - 只在内容超过一定行数时才应用滚动
                 max_text_height = console.height - 5
                 if max_text_height <= 0:
@@ -250,27 +276,32 @@ class BasePlatform(ABC):
                     text_content.plain = "\n".join(
                         [line.plain for line in lines[-max_text_height:]]
                     )
-                
+
                 # 只在需要时更新 subtitle（减少更新频率，避免重复渲染标题）
                 # 策略：每 10 次内容更新或每 3 秒更新一次 subtitle
                 current_time = time.time()
                 should_update_subtitle = (
-                    update_subtitle 
+                    update_subtitle
                     or update_count % 10 == 0  # 每 10 次更新一次
-                    or (current_time - last_subtitle_update_time) >= subtitle_update_interval
+                    or (current_time - last_subtitle_update_time)
+                    >= subtitle_update_interval
                 )
-                
+
                 if should_update_subtitle:
-                    self._update_panel_subtitle_with_token(panel, response, is_completed=False)
+                    self._update_panel_subtitle_with_token(
+                        panel, response, is_completed=False
+                    )
                     last_subtitle_update_time = current_time
-                
+
                 # 更新 panel（只更新内容，subtitle 更新频率已降低）
                 live.update(panel)
 
             # Process first chunk
             response += first_chunk
             if first_chunk:
-                _update_panel_content(first_chunk, update_subtitle=True)  # 第一次更新时更新 subtitle
+                _update_panel_content(
+                    first_chunk, update_subtitle=True
+                )  # 第一次更新时更新 subtitle
 
             # 缓存机制：降低更新频率，减少界面闪烁
             buffer = ""
@@ -308,33 +339,22 @@ class BasePlatform(ABC):
 
             _flush_buffer()
             # 在结束前，将面板内容替换为完整响应，确保最后一次渲染的 panel 显示全部内容
-            if response:
-                text_content.plain = response
-            # 最后更新 subtitle 和 panel
-            end_time = time.time()
-            duration = end_time - start_time
-            self._update_panel_subtitle_with_token(panel, response, is_completed=True, duration=duration)
-            # 最后更新 panel，Live 上下文退出时会自动打印（transient=False）
-            live.update(panel)
-            # 注意：不要在这里调用 console.print()，因为 Live 退出时会自动打印 panel
-        # Live 退出后仅添加空行分隔，不再重复打印 panel，避免内容重复
-        console.print()
         return response
 
     def _chat_with_simple_output(self, message: str, start_time: float) -> str:
         """使用简单输出模式进行聊天
-        
+
         参数:
             message: 用户消息
             start_time: 开始时间
-            
+
         返回:
             str: 模型响应
         """
         import time
-        
+
         console.print(
-            f"🤖 模型输出 - {(G.current_agent_name + ' · ') if G.current_agent_name else ''}{self.name()}  (按 Ctrl+C 中断)",
+            f"🤖 模型输出 - {(G.get_current_agent_name() + ' · ') if G.get_current_agent_name() else ''}{self.name()}  (按 Ctrl+C 中断)",
             soft_wrap=False,
         )
         response = ""
@@ -352,10 +372,10 @@ class BasePlatform(ABC):
 
     def _chat_with_suppressed_output(self, message: str) -> str:
         """使用静默模式进行聊天
-        
+
         参数:
             message: 用户消息
-            
+
         返回:
             str: 模型响应
         """
@@ -369,10 +389,10 @@ class BasePlatform(ABC):
 
     def _process_response(self, response: str) -> str:
         """处理响应，移除 think 标签
-        
+
         参数:
             response: 原始响应
-            
+
         返回:
             str: 处理后的响应
         """
@@ -391,7 +411,7 @@ class BasePlatform(ABC):
 
         # 当输入为空白字符串时，打印警告并直接返回空字符串
         if message.strip() == "":
-            print("⚠️ 输入为空白字符串，已忽略本次请求")
+            PrettyOutput.auto_print("⚠️ 输入为空白字符串，已忽略本次请求")
             return ""
 
         # 检查并截断消息以避免超出剩余token限制
@@ -403,12 +423,18 @@ class BasePlatform(ABC):
                 response = self._chat_with_pretty_output(message, start_time)
             else:
                 response = self._chat_with_simple_output(message, start_time)
+
+            # 计算响应时间并打印总结
+            end_time = time.time()
+            duration = end_time - start_time
+            PrettyOutput.auto_print(f"✅ {self.name()}模型响应完成: {duration:.2f}秒")
         else:
             response = self._chat_with_suppressed_output(message)
 
         # 处理响应并保存会话历史
         response = self._process_response(response)
         self._append_session_history(message, response)
+
         # 增加对话轮次计数
         self._conversation_turn += 1
         return response
@@ -419,17 +445,25 @@ class BasePlatform(ABC):
             set_in_chat(True)
             if not self.suppress_output and is_print_prompt():
                 PrettyOutput.print(f"{message}", OutputType.USER)  # 保留用于语法高亮
-            
+
+            # 记录用户输入（模型输入）
+            from jarvis.jarvis_utils.dialogue_recorder import record_user_message
+
+            record_user_message(message)
+
             result: str = ""
-            result = while_true(
-                lambda: while_success(lambda: self._chat(message))
-            )
-            
+            result = while_true(lambda: while_success(lambda: self._chat(message)))
+
             # Check if result is empty or False (retry exhausted)
             # Convert False to empty string for type safety
             if result is False or result == "":
                 raise ValueError("返回结果为空")
-            
+
+            # 记录模型输出
+            from jarvis.jarvis_utils.dialogue_recorder import record_assistant_message
+
+            record_assistant_message(result)
+
             from jarvis.jarvis_utils.globals import set_last_message
 
             set_last_message(result)
@@ -437,6 +471,14 @@ class BasePlatform(ABC):
         finally:
             set_in_chat(False)
 
+    def get_conversation_turn(self) -> int:
+        """获取当前对话轮次数
+
+        返回:
+            int: 当前对话轮次数
+        """
+        return self._conversation_turn
+
     @abstractmethod
     def name(self) -> str:
         """模型名称"""
@@ -523,6 +565,27 @@ class BasePlatform(ABC):
         """设置网页标志"""
         self.web = web
 
+    def set_platform_type(self, platform_type: str):
+        """设置平台类型
+
+        参数:
+            platform_type: 平台类型，可选值为 'normal'、'cheap' 或 'smart'
+        """
+        self.platform_type = platform_type
+
+    def _get_platform_max_input_token_count(self) -> int:
+        """根据平台类型获取对应的最大输入token数量
+
+        返回:
+            int: 模型能处理的最大输入token数量
+        """
+        if self.platform_type == "cheap":
+            return get_cheap_max_input_token_count(self.model_group)
+        elif self.platform_type == "smart":
+            return get_smart_max_input_token_count(self.model_group)
+        else:
+            return get_max_input_token_count(self.model_group)
+
     def _append_session_history(self, user_input: str, model_output: str) -> None:
         """
         Append the user input and model output to a session history file if enabled.
@@ -554,11 +617,14 @@ class BasePlatform(ABC):
                 ts = datetime.now().strftime("%Y%m%d_%H%M%S")
 
                 self._session_history_file = os.path.join(
-                    session_dir, f"session_history_{safe_platform}_{safe_model}_{ts}.log"
+                    session_dir,
+                    f"session_history_{safe_platform}_{safe_model}_{ts}.log",
                 )
 
             # Append record
-            with open(self._session_history_file, "a", encoding="utf-8", errors="ignore") as f:
+            with open(
+                self._session_history_file, "a", encoding="utf-8", errors="ignore"
+            ) as f:
                 ts_line = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 f.write(f"===== {ts_line} =====\n")
                 f.write("USER:\n")
@@ -571,10 +637,10 @@ class BasePlatform(ABC):
 
     def get_conversation_history(self) -> List[Dict[str, str]]:
         """获取当前对话历史
-        
+
         返回:
             List[Dict[str, str]]: 对话历史列表，每个元素包含 role 和 content
-            
+
         注意:
             默认实现检查是否有 messages 属性，子类可以重写此方法以提供自定义实现
         """
@@ -584,88 +650,90 @@ class BasePlatform(ABC):
 
     def get_used_token_count(self) -> int:
         """计算当前对话历史使用的token数量
-        
+
         返回:
             int: 当前对话历史使用的token数量
         """
         history = self.get_conversation_history()
         if not history:
             return 0
-        
+
         total_tokens = 0
         for message in history:
             content = message.get("content", "")
             if content:
                 total_tokens += get_context_token_count(content)
-        
+
         return total_tokens
 
     def get_remaining_token_count(self) -> int:
         """获取剩余可用的token数量
-        
+
         返回:
             int: 剩余可用的token数量（输入窗口限制 - 当前使用的token数量）
         """
-        max_tokens = get_max_input_token_count(self.model_group)
+        max_tokens = self._get_platform_max_input_token_count()
         used_tokens = self.get_used_token_count()
         remaining = max_tokens - used_tokens
         return max(0, remaining)  # 确保返回值不为负数
 
     def _truncate_message_if_needed(self, message: str) -> str:
         """如果消息超出剩余token限制，则截断消息
-        
+
         参数:
             message: 原始消息
-            
+
         返回:
             str: 截断后的消息（如果不需要截断则返回原消息）
         """
         try:
             # 获取剩余token数量
             remaining_tokens = self.get_remaining_token_count()
-            
+
             # 如果剩余token为0或负数，返回空消息
             if remaining_tokens <= 0:
-                print("⚠️ 警告：剩余token为0，无法发送消息")
+                PrettyOutput.auto_print("⚠️ 警告：剩余token为0，无法发送消息")
                 return ""
-            
+
             # 计算消息的token数量
             message_tokens = get_context_token_count(message)
-            
+
             # 如果消息token数小于等于剩余token数，不需要截断
             if message_tokens <= remaining_tokens:
                 return message
-            
+
             # 需要截断：保留剩余token的80%用于消息，20%作为安全余量
             target_tokens = int(remaining_tokens * 0.8)
             if target_tokens <= 0:
-                print("⚠️ 警告：剩余token不足，无法发送消息")
+                PrettyOutput.auto_print("⚠️ 警告：剩余token不足，无法发送消息")
                 return ""
-            
+
             # 估算字符数（1 token ≈ 4字符）
             target_chars = target_tokens * 4
-            
+
             # 如果消息长度小于目标字符数，不需要截断（token估算可能有误差）
             if len(message) <= target_chars:
                 return message
-            
+
             # 截断消息：保留前面的内容，添加截断提示
             truncated_message = message[:target_chars]
             # 尝试在最后一个完整句子处截断
-            last_period = truncated_message.rfind('.')
-            last_newline = truncated_message.rfind('\n')
+            last_period = truncated_message.rfind(".")
+            last_newline = truncated_message.rfind("\n")
             last_break = max(last_period, last_newline)
-            
+
             if last_break > target_chars * 0.5:  # 如果找到的断点不太靠前
-                truncated_message = truncated_message[:last_break + 1]
-            
+                truncated_message = truncated_message[: last_break + 1]
+
             truncated_message += "\n\n... (消息过长，已截断以避免超出上下文限制)"
-            print(f"⚠️ 警告：消息过长（{message_tokens} tokens），已截断至约 {target_tokens} tokens")
-            
+            PrettyOutput.auto_print(
+                f"⚠️ 警告：消息过长（{message_tokens} tokens），已截断至约 {target_tokens} tokens"
+            )
+
             return truncated_message
         except Exception as e:
             # 如果截断过程中出错，返回原消息（避免阻塞对话）
-            print(f"⚠️ 警告：检查消息长度时出错: {e}，使用原消息")
+            PrettyOutput.auto_print(f"⚠️ 警告：检查消息长度时出错: {e}，使用原消息")
             return message
 
     @abstractmethod

jarvis/jarvis_platform/human.py

@@ -1,17 +1,23 @@
-# -*- coding: utf-8 -*-
-
-# 人类交互平台实现模块
-
-# 提供与真实人类交互的模拟接口
-
 import json
 import random
 import string
-from typing import Generator, List, Tuple
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Tuple
 
 from jarvis.jarvis_platform.base import BasePlatform
 from jarvis.jarvis_utils.clipboard import copy_to_clipboard
 from jarvis.jarvis_utils.input import get_multiline_input
+from jarvis.jarvis_utils.output import PrettyOutput
+
+# -*- coding: utf-8 -*-
+
+# 人类交互平台实现模块
+
+# 提供与真实人类交互的模拟接口
 
 
 class HumanPlatform(BasePlatform):
@@ -21,8 +27,13 @@ class HumanPlatform(BasePlatform):
         """获取支持的模型列表"""
         return [("human", "Human Interaction")]
 
-    def __init__(self):
-        """初始化人类交互平台"""
+    def __init__(self, llm_config: Optional[Dict[str, Any]] = None):
+        """
+        初始化人类交互平台
+
+        参数:
+            llm_config: LLM配置字典（人类平台不使用此配置，但为保持接口一致而接受）
+        """
         super().__init__()
         self.conversation_id = ""  # 会话ID，用于标识当前对话
         self.model_name = "human"  # 默认模型名称
@@ -38,7 +49,7 @@ class HumanPlatform(BasePlatform):
         if model_name == "human":
             self.model_name = model_name
         else:
-            print(f"❌ 错误：不支持的模型: {model_name}")
+            PrettyOutput.auto_print(f"❌ 错误：不支持的模型: {model_name}")
 
     def chat(self, message: str) -> Generator[str, None, None]:
         """发送消息并获取人类响应"""
@@ -65,7 +76,7 @@ class HumanPlatform(BasePlatform):
 
     def upload_files(self, file_list: List[str]) -> bool:
         """文件上传功能，人类平台不需要实际处理"""
-        print("⚠️ 人类交互平台不支持文件上传")
+        PrettyOutput.auto_print("⚠️ 人类交互平台不支持文件上传")
         return False
 
     def delete_chat(self) -> bool:
@@ -87,10 +98,10 @@ class HumanPlatform(BasePlatform):
             with open(file_path, "w", encoding="utf-8") as f:
                 json.dump(state, f, ensure_ascii=False, indent=4)
             self._saved = True
-            print(f"✅ 会话已成功保存到 {file_path}")
+            PrettyOutput.auto_print(f"✅ 会话已成功保存到 {file_path}")
             return True
         except Exception as e:
-            print(f"❌ 保存会话失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 保存会话失败: {str(e)}")
             return False
 
     def restore(self, file_path: str) -> bool:
@@ -105,13 +116,13 @@ class HumanPlatform(BasePlatform):
             self.first_message = state.get("first_message", True)
             self._saved = True
 
-            print(f"✅ 从 {file_path} 成功恢复会话")
+            PrettyOutput.auto_print(f"✅ 从 {file_path} 成功恢复会话")
             return True
         except FileNotFoundError:
-            print(f"❌ 会话文件未找到: {file_path}")
+            PrettyOutput.auto_print(f"❌ 会话文件未找到: {file_path}")
             return False
         except Exception as e:
-            print(f"❌ 恢复会话失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 恢复会话失败: {str(e)}")
             return False
 
     def name(self) -> str: