From c742a30e1b85faac6141a5852c11365017d367be Mon Sep 17 00:00:00 2001 From: root Date: Tue, 14 Apr 2026 11:49:39 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E7=BB=9F=E4=B8=80=20Packet=20?= =?UTF-8?q?=E9=80=9A=E4=BF=A1=E5=8D=8F=E8=AE=AE=20+=20=E5=90=8E=E7=AB=AF?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E7=AE=A1=E7=90=86=20+=20UI=20=E6=B8=85?= =?UTF-8?q?=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 合并 server.py 到 core/__init__.py,使用统一 Packet 协议 - 后端管理配置持久化 (~/.trulymem/config.json) - 前端移除 ConfigService,通过 BackendClient 与后端通信 - 删除 UI 中冗余的 AI 推理逻辑 (chat_service, tool_service, message_handler) - 删除 core/tools 重复文件 (tool_executor, tool_limiter) - 提示词管理器支持用户自定义 (~/.trulyemem/system_prompt.md) - 启动入口优化配置路径逻辑 - 更新测试覆盖 (42 tests) - 更新文档 --- core/__init__.py | 250 +----------- core/client.py | 100 ++--- core/graph_client.py | 10 +- core/prompts/prompt_manager.py | 17 + core/server.py | 588 ++++++++++++++--------------- core/tools/__init__.py | 4 +- core/tools/tool_executor.py | 307 --------------- core/tools/tool_limiter.py | 164 -------- docs/api.md | 152 ++++---- docs/一键启动指南.md | 42 ++- docs/架构.md | 206 ++++------ tests/test_core/test_packet.py | 488 ++++++++++++------------ tests/test_integration/__init__.py | 233 ++++-------- tests/test_ui/__init__.py | 275 +++++--------- trulymem_entry.py | 46 ++- ui/app.py | 75 ++-- ui/handlers/message_handler.py | 98 ----- ui/services/chat_service.py | 226 ----------- ui/services/tool_service.py | 88 ----- 19 files changed, 1020 insertions(+), 2349 deletions(-) delete mode 100644 core/tools/tool_executor.py delete mode 100644 core/tools/tool_limiter.py delete mode 100644 ui/handlers/message_handler.py delete mode 100644 ui/services/chat_service.py delete mode 100644 ui/services/tool_service.py diff --git a/core/__init__.py b/core/__init__.py index 9f6dee2..13ee61b 100644 --- a/core/__init__.py +++ b/core/__init__.py @@ -1,242 +1,12 @@ -import threading -import queue -import time -from typing import Any, Dict -from dataclasses import dataclass, field -from enum import Enum - - -class PacketType(Enum): - MESSAGE = "message" - CONFIG = "config" - TOOL = "tool" - STATUS = "status" - HISTORY = "history" - - -@dataclass -class Packet: - id: str - type: PacketType - body: Dict[str, Any] - created_at: float = field(default_factory=time.time) - - +from .server import BackendServer, Packet, PacketType, PacketResponse +from .client import BackendClient from .embedded_db import EmbeddedGraphDB - -class BackendServer: - def __init__(self, db_path: str = "graph_memory.db", use_embedded_db: bool = True): - self._db_path = db_path - self._use_embedded_db = use_embedded_db - self._graph = None - self._client = None - self._running = False - self._thread = None - self._input_queue: queue.Queue[Packet] = queue.Queue() - self._response_queues: Dict[str, queue.Queue] = {} - self._lock = threading.Lock() - self._config = {"api_key": "", "base_url": "https://api.deepseek.com"} - - def start(self, api_key: str = "", base_url: str = "https://api.deepseek.com") -> None: - if self._running: - return - self._init_graph() - self._config = {"api_key": api_key, "base_url": base_url} - if api_key: - from .graph_client import GraphMemoryClient - self._client = GraphMemoryClient(api_key=api_key, base_url=base_url, graph=self._graph) - self._running = True - self._thread = threading.Thread(target=self._run_loop, daemon=True) - self._thread.start() - - def _init_graph(self) -> None: - if self._use_embedded_db: - self._graph = EmbeddedGraphDB(db_path=self._db_path) - else: - from .graph_client import Neo4jGraph - self._graph = Neo4jGraph(uri="bolt://localhost:7687", user="neo4j", password="graphmemory123") - - def _run_loop(self) -> None: - while self._running: - try: - packet = self._input_queue.get(timeout=0.1) - except queue.Empty: - continue - self._process_packet(packet) - - def _process_packet(self, packet: Packet) -> None: - body = {"success": False, "error": "not implemented"} - try: - if packet.type == PacketType.MESSAGE: - body = self._handle_message(packet.body) - elif packet.type == PacketType.CONFIG: - body = self._handle_config(packet.body) - elif packet.type == PacketType.TOOL: - body = self._handle_tool(packet.body) - elif packet.type == PacketType.STATUS: - body = self._handle_status() - elif packet.type == PacketType.HISTORY: - body = self._handle_history(packet.body) - body["success"] = True - except Exception as e: - body["error"] = str(e) - finally: - self._send_response(packet.id, Packet(id=packet.id, type=packet.type, body=body)) - - def _handle_message(self, body: Dict) -> Dict: - from .tool_executor import execute_tool - from .tool_limiter import ToolLimiter - - limiter = ToolLimiter() - user_input = body.get("message", "") - if not self._client: - return {"error": "API Key not configured", "content": "请先配置 API Key"} - - messages = [{"role": "user", "content": user_input}] - response = self._client.send_message_with_history(messages) - message = response.choices[0].message - content = message.content or "(无回复)" - tool_calls = [] - - while message.tool_calls: - # 处理所有工具调用 - tool_results = [] - - for tool_call in message.tool_calls: - args = {} - try: - import json - args = json.loads(tool_call.function.arguments) - except: - pass - - allowed, reason = limiter.can_call(tool_call.function.name, args) - if not allowed: - tool_calls.append({"name": tool_call.function.name, "result": f"工具调用被拒绝: {reason}"}) - tool_results.append({"role": "tool", "tool_call_id": tool_call.id, "content": f"工具调用被拒绝: {reason}"}) - continue - - limiter.record_call(tool_call.function.name, args) - result = execute_tool(self._graph, tool_call.function.name, args) - tool_calls.append({"name": tool_call.function.name, "result": result}) - tool_results.append({"role": "tool", "tool_call_id": tool_call.id, "content": result}) - - # 添加 assistant 消息(包含完整的 tool_calls,必须有 type 字段) - assistant_msg = { - "role": "assistant", - "content": message.content or "", - "tool_calls": [{ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - } for tc in message.tool_calls] - } - messages.append(assistant_msg) - - # 添加工具结果 - messages.extend(tool_results) - - response = self._client.send_message_with_history(messages) - message = response.choices[0].message - content = message.content or content - - return {"content": content, "tool_calls": tool_calls} - - def _handle_config(self, body: Dict) -> Dict: - api_key = body.get("api_key", "") - base_url = body.get("base_url", "https://api.deepseek.com") - self._config = {"api_key": api_key, "base_url": base_url} - if api_key and self._graph: - from .graph_client import GraphMemoryClient - self._client = GraphMemoryClient(api_key=api_key, base_url=base_url, graph=self._graph) - return {"status": "config_updated"} - - def _handle_tool(self, body: Dict) -> Dict: - from .tool_executor import execute_tool - name = body.get("name", "") - args = body.get("arguments", {}) - result = execute_tool(self._graph, name, args) - return {"result": result} - - def _handle_status(self) -> Dict: - return {"running": self._running, "config": self._config, "client_ready": self._client is not None} - - def _handle_history(self, body: Dict) -> Dict: - action = body.get("action", "get") - if action == "save": - self._message_history = body.get("messages", []) - return {"status": "saved"} - return {"history": getattr(self, "_message_history", [])} - - def send(self, packet: Packet) -> Packet: - resp_q = queue.Queue() - with self._lock: - self._response_queues[packet.id] = resp_q - self._input_queue.put(packet) - try: - return resp_q.get(timeout=30.0) - except queue.Empty: - return Packet(id=packet.id, type=packet.type, body={"success": False, "error": "timeout"}) - finally: - with self._lock: - self._response_queues.pop(packet.id, None) - - def _send_response(self, request_id: str, packet: Packet) -> None: - with self._lock: - q = self._response_queues.pop(request_id, None) - if q: - q.put(packet) - - def shutdown(self) -> None: - self._running = False - if self._thread: - self._thread.join(timeout=2.0) - if self._graph: - self._graph.close() - self._graph = None - - -class BackendClient: - def __init__(self, server: BackendServer): - self._server = server - self._counter = 0 - self._lock = threading.Lock() - - def _next_id(self) -> str: - with self._lock: - self._counter += 1 - return f"{time.time()}_{self._counter}" - - def send_message(self, message: str) -> Dict: - packet = Packet(id=self._next_id(), type=PacketType.MESSAGE, body={"message": message}) - return self._server.send(packet).body - - def update_config(self, api_key: str, base_url: str = "https://api.deepseek.com") -> Dict: - packet = Packet(id=self._next_id(), type=PacketType.CONFIG, body={"api_key": api_key, "base_url": base_url}) - return self._server.send(packet).body - - def execute_tool(self, name: str, arguments: Dict) -> str: - packet = Packet(id=self._next_id(), type=PacketType.TOOL, body={"name": name, "arguments": arguments}) - return self._server.send(packet).body.get("result", "") - - def get_status(self) -> Dict: - packet = Packet(id=self._next_id(), type=PacketType.STATUS, body={}) - return self._server.send(packet).body - - def save_history(self, messages: list) -> None: - packet = Packet(id=self._next_id(), type=PacketType.HISTORY, body={"action": "save", "messages": messages}) - self._server.send(packet) - - def get_history(self) -> list: - packet = Packet(id=self._next_id(), type=PacketType.HISTORY, body={"action": "get"}) - return self._server.send(packet).body.get("history", []) - - def shutdown(self) -> None: - self._server.shutdown() - - -__all__ = ["BackendServer", "BackendClient", "EmbeddedGraphDB", "Packet", "PacketType"] \ No newline at end of file +__all__ = [ + "BackendServer", + "BackendClient", + "EmbeddedGraphDB", + "Packet", + "PacketType", + "PacketResponse" +] \ No newline at end of file diff --git a/core/client.py b/core/client.py index c94d002..592dc8a 100644 --- a/core/client.py +++ b/core/client.py @@ -1,61 +1,75 @@ import threading +import time from typing import Any, Dict -from .server import BackendServer, MessageType +from .server import BackendServer, Packet, PacketType class BackendClient: + def __init__(self, server: BackendServer): self._server = server - self._request_counter = 0 + self._counter = 0 self._lock = threading.Lock() - def process_message(self, user_input: str, timeout: float = 30.0) -> Dict[str, Any]: + def _next_id(self) -> str: with self._lock: - self._request_counter += 1 - return self._server.process_message(user_input, timeout) + self._counter += 1 + return f"{time.time()}_{self._counter}" - def execute_tool(self, tool_name: str, arguments: Dict[str, Any], timeout: float = 10.0) -> str: - with self._lock: - self._request_counter += 1 - return self._server.execute_tool(tool_name, arguments, timeout) + def send(self, message: str) -> Dict: + return self.process_message(message) - def update_config(self, api_key: str, base_url: str = "https://api.deepseek.com") -> Dict[str, Any]: - return self._send_request(MessageType.SET_CONFIG, {"api_key": api_key, "base_url": base_url}) + def process_message(self, user_input: str) -> Dict: + return self._server.process_message(user_input) - def get_config(self) -> Dict[str, str]: - result = self._send_request(MessageType.GET_CONFIG, {}) - return result.get("data", {"api_key": "", "base_url": "https://api.deepseek.com"}) - - def get_message_history(self) -> list: - result = self._send_request(MessageType.GET_HISTORY, {}) - return result.get("data", {}).get("history", []) - - def save_message_history(self, messages: list) -> None: - self._send_request(MessageType.SAVE_HISTORY, {"messages": messages}) - - def _send_request(self, message_type: MessageType, payload: Dict[str, Any]) -> Dict[str, Any]: - import queue - request_id = f"{id(self)}" - response_queue = queue.Queue() - - from .server import BackendRequest - request = BackendRequest( - request_id=request_id, - message_type=message_type, - payload=payload, - response_queue=response_queue + def update_config(self, api_key: str, base_url: str = "https://api.deepseek.com", model: str = "deepseek-chat") -> Dict: + packet = Packet( + id=self._next_id(), + type=PacketType.SET_CONFIG, + body={"api_key": api_key, "base_url": base_url, "model": model} ) - - self._server._request_queue.put(request) - - try: - response = response_queue.get(timeout=5.0) - if not response.success: - raise Exception(response.error) - return {"success": True, "data": response.data} - except queue.Empty: - return {"success": False, "error": "timeout"} + return self._server.send(packet).body + + def execute_tool(self, name: str, arguments: Dict) -> Dict: + packet = Packet( + id=self._next_id(), + type=PacketType.EXECUTE_TOOL, + body={"tool_name": name, "arguments": arguments} + ) + return self._server.send(packet).body + + def get_status(self) -> Dict: + packet = Packet( + id=self._next_id(), + type=PacketType.GET_STATUS, + body={} + ) + return self._server.send(packet).body + + def get_config(self) -> Dict: + packet = Packet( + id=self._next_id(), + type=PacketType.GET_CONFIG, + body={} + ) + return self._server.send(packet).body + + def save_history(self, messages: list) -> None: + packet = Packet( + id=self._next_id(), + type=PacketType.SAVE_HISTORY, + body={"messages": messages} + ) + self._server.send(packet) + + def get_history(self) -> list: + packet = Packet( + id=self._next_id(), + type=PacketType.GET_HISTORY, + body={} + ) + return self._server.send(packet).body.get("history", []) def shutdown(self) -> None: self._server.shutdown() \ No newline at end of file diff --git a/core/graph_client.py b/core/graph_client.py index 76c5c1b..706ff75 100644 --- a/core/graph_client.py +++ b/core/graph_client.py @@ -303,12 +303,12 @@ class Neo4jGraph: class GraphMemoryClient: """图记忆客户端""" - def __init__(self, api_key: str, base_url: str, graph): + def __init__(self, api_key: str, base_url: str, graph, model: str = "deepseek-chat"): self.client = OpenAI(api_key=api_key, base_url=base_url) self.graph = graph self.tools = TOOLS + self.model = model - # 使用新的提示词管理器 prompt_manager = PromptManager() self.system_prompt = prompt_manager.get_system_prompt() @@ -330,7 +330,7 @@ class GraphMemoryClient: messages.extend(tool_results) response = self.client.chat.completions.create( - model=MODEL_NAME, + model=self.model, messages=messages, tools=self.tools, tool_choice="auto" @@ -347,7 +347,7 @@ class GraphMemoryClient: messages.extend(messages_history) response = self.client.chat.completions.create( - model=MODEL_NAME, + model=self.model, messages=messages, tools=self.tools, tool_choice="auto" @@ -373,7 +373,7 @@ class GraphMemoryClient: messages.extend(tool_results) stream = self.client.chat.completions.create( - model=MODEL_NAME, + model=self.model, messages=messages, tools=self.tools, tool_choice="auto", diff --git a/core/prompts/prompt_manager.py b/core/prompts/prompt_manager.py index fb19377..fd30db2 100644 --- a/core/prompts/prompt_manager.py +++ b/core/prompts/prompt_manager.py @@ -10,6 +10,9 @@ class PromptManager: _instance = None _cached_prompt = None + # 用户自定义提示词路径 + USER_PROMPT_PATH = Path.home() / ".trulymem" / "system_prompt.md" + def __new__(cls): """单例模式,避免重复加载""" if cls._instance is None: @@ -26,11 +29,19 @@ class PromptManager: if PromptManager._cached_prompt is not None: return PromptManager._cached_prompt + # 1. 优先使用用户自定义提示词 + if self.USER_PROMPT_PATH.exists(): + with open(self.USER_PROMPT_PATH, "r", encoding="utf-8") as f: + PromptManager._cached_prompt = f.read() + return PromptManager._cached_prompt + + # 2. 使用打包的提示词 prompt_file = self.prompts_dir / "system_prompt.md" if prompt_file.exists(): with open(prompt_file, "r", encoding="utf-8") as f: PromptManager._cached_prompt = f.read() else: + # 3. 使用内置默认提示词 PromptManager._cached_prompt = self._build_default_prompt() return PromptManager._cached_prompt @@ -78,3 +89,9 @@ class PromptManager: - 如何使用工具 记住:灵活应对,保持自然对话体验。""" + + @staticmethod + def clear_cache(): + """清除缓存,强制重新加载提示词""" + PromptManager._cached_prompt = None + PromptManager._instance = None \ No newline at end of file diff --git a/core/server.py b/core/server.py index 5bd3d69..375746c 100644 --- a/core/server.py +++ b/core/server.py @@ -2,17 +2,16 @@ import threading import queue import time import json +import os +from pathlib import Path from typing import Any, Dict, Optional from dataclasses import dataclass, field from enum import Enum from .embedded_db import EmbeddedGraphDB -from .graph_client import GraphMemoryClient -from .tool_executor import execute_tool -from .tool_limiter import ToolLimiter -class MessageType(Enum): +class PacketType(Enum): PROCESS_MESSAGE = "process_message" EXECUTE_TOOL = "execute_tool" GET_STATUS = "get_status" @@ -24,46 +23,66 @@ class MessageType(Enum): @dataclass -class BackendRequest: - request_id: str - message_type: MessageType - payload: Dict[str, Any] - response_queue: queue.Queue = field(default=None) +class Packet: + id: str + type: PacketType + body: Dict[str, Any] + response_queue: Optional[queue.Queue] = field(default=None) + created_at: float = field(default_factory=time.time) @dataclass -class BackendResponse: - request_id: str +class PacketResponse: + id: str success: bool data: Any = None error: Optional[str] = None class BackendServer: - def __init__(self, db_path: str = "graph_memory.db", use_embedded_db: bool = True): + + DEFAULT_CONFIG_PATH = Path.home() / ".trulymem" / "config.json" + + def __init__(self, db_path: str = "graph_memory.db", use_embedded_db: bool = True, config_file: str = None): self._db_path = db_path self._use_embedded_db = use_embedded_db + self._config_file = Path(config_file) if config_file else self.DEFAULT_CONFIG_PATH self._graph = None self._client = None - self._tool_limiter = ToolLimiter() + self._tool_limiter = None - self._request_queue: queue.Queue[BackendRequest] = queue.Queue() + self._input_queue: queue.Queue[Packet] = queue.Queue() + self._response_queues: Dict[str, queue.Queue] = {} self._running = False self._thread: Optional[threading.Thread] = None self._lock = threading.Lock() + self._config = {"api_key": "", "base_url": "https://api.deepseek.com", "model": "deepseek-chat"} + self._message_history: list = [] - def start(self, api_key: str = "", base_url: str = "https://api.deepseek.com") -> None: + def start(self, api_key: str = "", base_url: str = "https://api.deepseek.com", model: str = "deepseek-chat") -> None: if self._running: return - self._init_graph() + self._load_config() if api_key: + self._config["api_key"] = api_key + if base_url: + self._config["base_url"] = base_url + if model: + self._config["model"] = model + + self._init_graph() + self._tool_limiter = self._create_tool_limiter() + + if self._config["api_key"]: + from .graph_client import GraphMemoryClient self._client = GraphMemoryClient( - api_key=api_key, - base_url=base_url, + api_key=self._config["api_key"], + base_url=self._config["base_url"], + model=self._config.get("model", "deepseek-chat"), graph=self._graph ) @@ -71,6 +90,24 @@ class BackendServer: self._thread = threading.Thread(target=self._run_loop, daemon=True) self._thread.start() + def _load_config(self) -> None: + if self._config_file.exists(): + try: + with open(self._config_file, 'r') as f: + saved = json.load(f) + self._config.update(saved) + except Exception: + pass + + def _save_config(self) -> None: + self._config_file.parent.mkdir(parents=True, exist_ok=True) + with open(self._config_file, 'w') as f: + json.dump(self._config, f, indent=2) + + def _create_tool_limiter(self): + from .tool_limiter import ToolLimiter + return ToolLimiter() + def _init_graph(self) -> None: if self._use_embedded_db: self._graph = EmbeddedGraphDB(db_path=self._db_path) @@ -85,310 +122,276 @@ class BackendServer: def _run_loop(self) -> None: while self._running: try: - request = self._request_queue.get(timeout=0.1) + packet = self._input_queue.get(timeout=0.1) except queue.Empty: continue - if request.message_type == MessageType.PROCESS_MESSAGE: - self._handle_process_message(request) - elif request.message_type == MessageType.EXECUTE_TOOL: - self._handle_execute_tool(request) - elif request.message_type == MessageType.GET_STATUS: - self._handle_get_status(request) - elif request.message_type == MessageType.GET_CONFIG: - self._handle_get_config(request) - elif request.message_type == MessageType.SET_CONFIG: - self._handle_set_config(request) - elif request.message_type == MessageType.GET_HISTORY: - self._handle_get_history(request) - elif request.message_type == MessageType.SAVE_HISTORY: - self._handle_save_history(request) - elif request.message_type == MessageType.SHUTDOWN: - self._running = False - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={"status": "shutdown"} - )) + self._process_packet(packet) - def _handle_process_message(self, request: BackendRequest) -> None: + def _process_packet(self, packet: Packet) -> None: + response_body = {"error": "not implemented"} + try: - user_input = request.payload.get("user_input", "") + if packet.type == PacketType.PROCESS_MESSAGE: + response_body = self._handle_process_message(packet.body) + elif packet.type == PacketType.EXECUTE_TOOL: + response_body = self._handle_execute_tool(packet.body) + elif packet.type == PacketType.GET_STATUS: + response_body = self._handle_get_status() + elif packet.type == PacketType.GET_CONFIG: + response_body = self._handle_get_config() + elif packet.type == PacketType.SET_CONFIG: + response_body = self._handle_set_config(packet.body) + elif packet.type == PacketType.GET_HISTORY: + response_body = self._handle_get_history() + elif packet.type == PacketType.SAVE_HISTORY: + response_body = self._handle_save_history(packet.body) + elif packet.type == PacketType.SHUTDOWN: + self._running = False + response_body = {"success": True, "status": "shutdown"} - if not self._client: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error="API Key 未配置" - )) - return + if "success" not in response_body: + response_body["success"] = True + except Exception as e: + response_body["success"] = False + response_body["error"] = str(e) + + self._send_response(packet.id, PacketResponse( + id=packet.id, + success=response_body.get("success", False), + data=response_body if response_body.get("success") else None, + error=response_body.get("error") + )) + + def _handle_process_message(self, body: Dict) -> Dict: + from .tool_executor import execute_tool + + user_input = body.get("user_input", "") + + if not self._client: + return {"success": False, "error": "API Key 未配置", "content": "请先配置 API Key"} + + self._tool_limiter.reset() + + messages_history = [{"role": "user", "content": user_input}] + + response = self._client.send_message_with_history(messages_history) + message = response.choices[0].message + + tool_calls = [] + accumulated_content = "" + rejected_tools = [] + + while message.tool_calls: + if message.content: + accumulated_content += message.content + "\n\n" - self._tool_limiter.reset() - - messages_history = [{"role": "user", "content": user_input}] - - response = self._client.send_message_with_history(messages_history) - message = response.choices[0].message - - tool_calls = [] - accumulated_content = "" - rejected_tools = [] - - while message.tool_calls: - if message.content: - accumulated_content += message.content + "\n\n" - - assistant_msg = { - "role": "assistant", - "content": message.content, - "tool_calls": [ - { - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - } for tc in message.tool_calls - ] - } - messages_history.append(assistant_msg) - - current_tool_results = [] - for tool_call in message.tool_calls: - args = json.loads(tool_call.function.arguments) - - allowed, reason = self._tool_limiter.can_call(tool_call.function.name, args) - - if not allowed: - rejected_tools.append((tool_call.function.name, reason)) - result = f"工具调用被拒绝: {reason}" - tool_result_msg = { - "role": "tool", - "tool_call_id": tool_call.id, - "content": result + assistant_msg = { + "role": "assistant", + "content": message.content, + "tool_calls": [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments } - current_tool_results.append(tool_result_msg) - continue - - self._tool_limiter.record_call(tool_call.function.name, args) - - result = execute_tool(self._graph, tool_call.function.name, args) - tool_calls.append({ - "name": tool_call.function.name, - "arguments": args, - "result": result - }) - + } for tc in message.tool_calls + ] + } + messages_history.append(assistant_msg) + + current_tool_results = [] + for tool_call in message.tool_calls: + args = json.loads(tool_call.function.arguments) + + allowed, reason = self._tool_limiter.can_call(tool_call.function.name, args) + + if not allowed: + rejected_tools.append((tool_call.function.name, reason)) + result = f"工具调用被拒绝: {reason}" tool_result_msg = { "role": "tool", "tool_call_id": tool_call.id, "content": result } current_tool_results.append(tool_result_msg) + continue - messages_history.extend(current_tool_results) + self._tool_limiter.record_call(tool_call.function.name, args) - response = self._client.send_message_with_history(messages_history) - message = response.choices[0].message - - final_content = message.content or "" - content = accumulated_content + final_content if accumulated_content else final_content - - if not content: - content = "(无回复)" - - if tool_calls: - tool_names = [tc["name"] for tc in tool_calls] - content = f"已执行工具: {', '.join(tool_names)}\n\n{content}" - - if rejected_tools: - rejected_info = "\n".join([f"{name}: {reason}" for name, reason in rejected_tools]) - content += f"\n\n部分工具调用被限制:\n{rejected_info}" - content += f"\n\n工具调用统计:\n{self._tool_limiter.get_summary()}" - - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={ - "content": content, - "tool_calls": tool_calls, - "rejected_tools": rejected_tools + result = execute_tool(self._graph, tool_call.function.name, args) + tool_calls.append({ + "name": tool_call.function.name, + "arguments": args, + "result": result + }) + + tool_result_msg = { + "role": "tool", + "tool_call_id": tool_call.id, + "content": result } - )) + current_tool_results.append(tool_result_msg) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) + messages_history.extend(current_tool_results) + + response = self._client.send_message_with_history(messages_history) + message = response.choices[0].message + + final_content = message.content or "" + content = accumulated_content + final_content if accumulated_content else final_content + + if not content: + content = "(无回复)" + + if tool_calls: + tool_names = [tc["name"] for tc in tool_calls] + content = f"已执行工具: {', '.join(tool_names)}\n\n{content}" + + if rejected_tools: + rejected_info = "\n".join([f"{name}: {reason}" for name, reason in rejected_tools]) + content += f"\n\n部分工具调用被限制:\n{rejected_info}" + content += f"\n\n工具调用统计:\n{self._tool_limiter.get_summary()}" + + return { + "success": True, + "content": content, + "tool_calls": tool_calls, + "rejected_tools": rejected_tools + } - def _handle_execute_tool(self, request: BackendRequest) -> None: - """处理直接工具调用请求(前端直接调用,不受次数限制)""" + def _handle_execute_tool(self, body: Dict) -> Dict: + from .tool_executor import execute_tool + try: - tool_name = request.payload.get("tool_name") - arguments = request.payload.get("arguments", {}) + tool_name = body.get("tool_name") + arguments = body.get("arguments", {}) - # 前端直接调用的工具不受次数限制,直接执行 result = execute_tool(self._graph, tool_name, arguments) - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={"result": result} - )) - + return {"success": True, "result": result} except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) + return {"success": False, "error": str(e)} - def _handle_get_status(self, request: BackendRequest) -> None: - try: - status = { - "graph_initialized": self._graph is not None, - "client_initialized": self._client is not None, - "running": self._running - } - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data=status - )) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) + def _handle_get_status(self) -> Dict: + return { + "running": self._running, + "config": self._config, + "graph_initialized": self._graph is not None, + "client_initialized": self._client is not None + } - def _handle_get_config(self, request: BackendRequest) -> None: - try: - config = self.get_config() - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data=config - )) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) + def _handle_get_config(self) -> Dict: + return self._config.copy() - def _handle_set_config(self, request: BackendRequest) -> None: - try: - api_key = request.payload.get("api_key", "") - base_url = request.payload.get("base_url", "https://api.deepseek.com") - self.update_config(api_key, base_url) - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={"status": "config_updated"} - )) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) - - def _handle_get_history(self, request: BackendRequest) -> None: - try: - history = self.get_message_history() - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={"history": history} - )) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) - - def _handle_save_history(self, request: BackendRequest) -> None: - try: - messages = request.payload.get("messages", []) - self.save_message_history(messages) - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=True, - data={"status": "history_saved"} - )) - except Exception as e: - self._send_response(request, BackendResponse( - request_id=request.request_id, - success=False, - error=str(e) - )) - - def _send_response(self, request: BackendRequest, response: BackendResponse) -> None: - if request.response_queue: - request.response_queue.put(response) - - def process_message(self, user_input: str, timeout: float = 30.0) -> Dict[str, Any]: - request_id = f"{time.time()}" - response_queue = queue.Queue() + def _handle_set_config(self, body: Dict) -> Dict: + api_key = body.get("api_key", "") + base_url = body.get("base_url", "https://api.deepseek.com") + model = body.get("model", "deepseek-chat") - request = BackendRequest( - request_id=request_id, - message_type=MessageType.PROCESS_MESSAGE, - payload={"user_input": user_input}, - response_queue=response_queue + self.update_config(api_key, base_url, model) + self._save_config() + return {"status": "config_updated"} + + def _handle_get_history(self) -> Dict: + return {"history": self._message_history} + + def _handle_save_history(self, body: Dict) -> Dict: + messages = body.get("messages", []) + self._message_history = messages + return {"status": "history_saved"} + + def _send_response(self, request_id: str, response: PacketResponse) -> None: + with self._lock: + q = self._response_queues.pop(request_id, None) + if q: + q.put(response) + + def send(self, packet: Packet) -> Packet: + resp_q = queue.Queue() + + with self._lock: + self._response_queues[packet.id] = resp_q + + self._input_queue.put(packet) + + try: + response = resp_q.get(timeout=30.0) + return Packet( + id=response.id, + type=packet.type, + body={ + "success": response.success, + "data": response.data, + "error": response.error + } + ) + except queue.Empty: + return Packet( + id=packet.id, + type=packet.type, + body={"success": False, "error": "timeout"} + ) + finally: + with self._lock: + self._response_queues.pop(packet.id, None) + + def process_message(self, user_input: str) -> Dict[str, Any]: + packet = Packet( + id=f"{time.time()}", + type=PacketType.PROCESS_MESSAGE, + body={"user_input": user_input} ) - self._request_queue.put(request) - - try: - response = response_queue.get(timeout=timeout) - if not response.success: - raise Exception(response.error) - return response.data - except queue.Empty: - raise TimeoutError("请求超时") + response = self.send(packet) + return response.body - def execute_tool(self, tool_name: str, arguments: Dict[str, Any], timeout: float = 10.0) -> str: - request_id = f"{time.time()}" - response_queue = queue.Queue() - - request = BackendRequest( - request_id=request_id, - message_type=MessageType.EXECUTE_TOOL, - payload={"tool_name": tool_name, "arguments": arguments}, - response_queue=response_queue + def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: + packet = Packet( + id=f"{time.time()}", + type=PacketType.EXECUTE_TOOL, + body={"tool_name": tool_name, "arguments": arguments} ) - self._request_queue.put(request) + response = self.send(packet) + return response.body + + def update_config(self, api_key: str, base_url: str = "https://api.deepseek.com", model: str = "deepseek-chat") -> None: + with self._lock: + self._config["api_key"] = api_key + self._config["base_url"] = base_url + self._config["model"] = model - try: - response = response_queue.get(timeout=timeout) - if not response.success: - raise Exception(response.error) - return response.data["result"] - except queue.Empty: - raise TimeoutError("工具执行超时") + if api_key and self._graph: + from .graph_client import GraphMemoryClient + self._client = GraphMemoryClient( + api_key=api_key, + base_url=base_url, + model=model, + graph=self._graph + ) + + def get_config(self) -> Dict[str, str]: + return self._config.copy() + + def save_message_history(self, messages: list) -> None: + self._message_history = messages + + def get_message_history(self) -> list: + return self._message_history.copy() def shutdown(self) -> None: if not self._running: return - request_id = f"{time.time()}" - response_queue = queue.Queue() - - request = BackendRequest( - request_id=request_id, - message_type=MessageType.SHUTDOWN, - payload={}, - response_queue=response_queue + packet = Packet( + id=f"{time.time()}", + type=PacketType.SHUTDOWN, + body={} ) - - self._request_queue.put(request) + self.send(packet) if self._thread: self._thread.join(timeout=2.0) @@ -396,22 +399,5 @@ class BackendServer: if self._graph: self._graph.close() self._graph = None - - def update_config(self, api_key: str, base_url: str = "https://api.deepseek.com") -> None: - with self._lock: - self._config = {"api_key": api_key, "base_url": base_url} - if api_key and self._graph: - self._client = GraphMemoryClient( - api_key=api_key, - base_url=base_url, - graph=self._graph - ) - - def get_config(self) -> Dict[str, str]: - return getattr(self, "_config", {"api_key": "", "base_url": "https://api.deepseek.com"}) - - def save_message_history(self, messages: list) -> None: - self._message_history = messages - - def get_message_history(self) -> list: - return getattr(self, "_message_history", []) \ No newline at end of file + + self._running = False \ No newline at end of file diff --git a/core/tools/__init__.py b/core/tools/__init__.py index 0a0ad7d..b562edc 100644 --- a/core/tools/__init__.py +++ b/core/tools/__init__.py @@ -2,7 +2,7 @@ 工具定义模块 """ from .memory_tools import TOOLS -from .tool_executor import execute_tool -from .tool_limiter import ToolLimiter, ToolLimits, ToolCallCount +from ..tool_executor import execute_tool +from ..tool_limiter import ToolLimiter, ToolLimits, ToolCallCount __all__ = ["TOOLS", "execute_tool", "ToolLimiter", "ToolLimits", "ToolCallCount"] diff --git a/core/tools/tool_executor.py b/core/tools/tool_executor.py deleted file mode 100644 index 28f70d6..0000000 --- a/core/tools/tool_executor.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -工具执行器 -""" -import json -from typing import Any, Dict - - -def execute_tool(graph: Any, tool_name: str, arguments: dict) -> str: - """执行工具调用""" - print(f"\n[工具调用] {tool_name}") - print(f"[参数] {json.dumps(arguments, ensure_ascii=False, indent=2)}") - - try: - # 基础记忆工具 - if tool_name == "memory_recall": - result = graph.recall( - query_intent=arguments.get("query_intent", ""), - seed_entities=arguments.get("seed_entities"), - depth=arguments.get("depth", 2), - time_range=arguments.get("time_range"), - session_filter=arguments.get("session_filter") - ) - return format_recall_result(result) - - elif tool_name == "memory_commit": - result = graph.commit( - triplets=arguments.get("triplets", []), - entity_types=arguments.get("entity_types"), - temporal_tag=arguments.get("temporal_tag") - ) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "memory_purge": - result = graph.purge( - criteria=arguments.get("criteria", {}), - mode=arguments.get("mode", "soft"), - new_relation=arguments.get("new_relation") - ) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "memory_introspect": - result = graph.introspect(session_id=arguments.get("session_id")) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "memory_archive": - result = graph.archive(days=arguments.get("days", 30)) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "memory_cleanup": - result = graph.cleanup(dry_run=arguments.get("dry_run", True)) - return json.dumps(result, ensure_ascii=False, default=str) - - # 人设图管理工具 - elif tool_name == "persona_update": - result = execute_persona_update(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "persona_clear": - result = execute_persona_clear(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - # 工作记忆链管理工具 - elif tool_name == "task_create": - result = execute_task_create(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "task_set_state": - result = execute_task_set_state(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "task_delete": - result = execute_task_delete(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - elif tool_name == "task_link_info": - result = execute_task_link_info(graph, arguments) - return json.dumps(result, ensure_ascii=False, default=str) - - return f"未知工具: {tool_name}" - - except Exception as e: - return f"工具执行错误: {str(e)}" - - -def format_recall_result(result: dict) -> str: - """格式化检索结果""" - lines = ["===== 记忆检索结果 ====="] - - if result.get("entities"): - lines.append(f"\n实体 ({len(result['entities'])} 个):") - for e in result["entities"]: - if e and isinstance(e, dict): - lines.append(f" - {e.get('name', 'N/A')} (类型: {e.get('type', 'unknown')}, 提及: {e.get('mention_count', 1)}次)") - - if result.get("relations"): - lines.append(f"\n关系 ({len(result['relations'])} 条):") - for r in result["relations"]: - if r and isinstance(r, dict): - lines.append(f" - {r.get('source', 'N/A')} --[{r.get('type', 'N/A')}]--> {r.get('target', 'N/A')}") - created = r.get("created_at", "N/A") - if created and created != "N/A": - created = created[:19] if "T" in str(created) else str(created) - session_id = r.get('session_id', 'N/A') - session_display = session_id[:20] if session_id and session_id != 'N/A' else 'N/A' - lines.append(f" 时间: {created}, 会话: {session_display}, 轮次: {r.get('turn_id', 0)}, 置信度: {r.get('confidence', 1.0)}") - - if not result.get("entities") and not result.get("relations"): - lines.append("\n(未找到相关记忆)") - - lines.append("=" * 30) - return "\n".join(lines) - - -# 人设图管理工具实现 -def execute_persona_update(graph: Any, arguments: dict) -> dict: - """更新人设""" - attributes = arguments.get("attributes", []) - mode = arguments.get("mode", "merge") - - if mode == "replace": - # 先清除旧人设 - graph.purge( - criteria={"subject_contains": "AI", "relation_type": "扮演角色"}, - mode="soft" - ) - graph.purge( - criteria={"subject_contains": "AI", "relation_type": "说话风格"}, - mode="soft" - ) - graph.purge( - criteria={"subject_contains": "AI", "relation_type": "性格特点"}, - mode="soft" - ) - - # 写入新人设 - triplets = [] - for attr in attributes: - triplets.append({ - "subject": "AI", - "relation": attr["attribute"], - "object": attr["value"], - "confidence": 1.0 - }) - - result = graph.commit(triplets=triplets) - return { - "status": "success", - "mode": mode, - "updated_attributes": len(attributes), - "details": result - } - - -def execute_persona_clear(graph: Any, arguments: dict) -> dict: - """清除人设""" - if not arguments.get("confirm", True): - return {"status": "cancelled", "message": "需要确认才能清除人设"} - - # 删除所有人设相关关系 - result1 = graph.purge( - criteria={"subject_contains": "AI", "relation_type": "扮演角色"}, - mode="soft" - ) - result2 = graph.purge( - criteria={"subject_contains": "AI", "relation_type": "说话风格"}, - mode="soft" - ) - result3 = graph.purge( - criteria={"subject_contains": "AI", "relation_type": "性格特点"}, - mode="soft" - ) - result4 = graph.purge( - criteria={"subject_contains": "AI", "relation_type": "语气特征"}, - mode="soft" - ) - - total_deleted = ( - result1.get("deleted_count", 0) + - result2.get("deleted_count", 0) + - result3.get("deleted_count", 0) + - result4.get("deleted_count", 0) - ) - - return { - "status": "success", - "deleted_count": total_deleted, - "message": "人设已清除,恢复默认身份" - } - - -# 工作记忆链管理工具实现 -def execute_task_create(graph: Any, arguments: dict) -> dict: - """创建任务节点""" - task_id = arguments.get("task_id") - description = arguments.get("description") - info_nodes = arguments.get("info_nodes", []) - - # 创建任务节点 - triplets = [ - {"subject": task_id, "relation": "is_type", "object": "TaskNode"}, - {"subject": task_id, "relation": "has_description", "object": description}, - {"subject": task_id, "relation": "HAS_STATE", "object": "State_进行中"} - ] - - result = graph.commit(triplets=triplets) - - # 关联信息节点 - if info_nodes: - link_triplets = [] - for node_name in info_nodes: - link_triplets.append({ - "subject": task_id, - "relation": "CONTAINS_INFO", - "object": node_name - }) - graph.commit(triplets=link_triplets) - - return { - "status": "success", - "task_id": task_id, - "description": description, - "info_nodes": info_nodes, - "details": result - } - - -def execute_task_set_state(graph: Any, arguments: dict) -> dict: - """设置任务状态""" - task_id = arguments.get("task_id") - state = arguments.get("state") - - # 删除旧状态 - graph.purge( - criteria={"subject_contains": task_id, "relation_type": "HAS_STATE"}, - mode="soft" - ) - - # 设置新状态 - state_node = f"State_{state}" - result = graph.commit( - triplets=[{"subject": task_id, "relation": "HAS_STATE", "object": state_node}] - ) - - return { - "status": "success", - "task_id": task_id, - "new_state": state, - "details": result - } - - -def execute_task_delete(graph: Any, arguments: dict) -> dict: - """删除任务节点""" - task_id = arguments.get("task_id") - delete_info_nodes = arguments.get("delete_info_nodes", True) - - # 查询关联的信息节点 - if delete_info_nodes: - recall_result = graph.recall( - query_intent=f"{task_id},CONTAINS_INFO", - depth=1 - ) - - # 删除信息节点 - for relation in recall_result.get("relations", []): - if relation.get("type") == "CONTAINS_INFO" and relation.get("source") == task_id: - info_node = relation.get("target") - graph.purge( - criteria={"subject_contains": info_node}, - mode="soft" - ) - - # 删除任务节点 - result = graph.purge( - criteria={"subject_contains": task_id}, - mode="soft" - ) - - return { - "status": "success", - "task_id": task_id, - "deleted_info_nodes": delete_info_nodes, - "details": result - } - - -def execute_task_link_info(graph: Any, arguments: dict) -> dict: - """关联信息节点""" - task_id = arguments.get("task_id") - info_node_names = arguments.get("info_node_names", []) - - triplets = [] - for node_name in info_node_names: - triplets.append({ - "subject": task_id, - "relation": "CONTAINS_INFO", - "object": node_name - }) - - result = graph.commit(triplets=triplets) - - return { - "status": "success", - "task_id": task_id, - "linked_nodes": info_node_names, - "details": result - } diff --git a/core/tools/tool_limiter.py b/core/tools/tool_limiter.py deleted file mode 100644 index 9f0f614..0000000 --- a/core/tools/tool_limiter.py +++ /dev/null @@ -1,164 +0,0 @@ -""" -工具调用限制器 - 限制每轮对话中各类工具的调用次数 -""" -from typing import Dict, List, Optional -from dataclasses import dataclass, field - - -@dataclass -class ToolLimits: - """工具调用限制配置""" - # 人设图限制 - persona_query_max: int = 1 # 每轮最多查询1次人设图 - persona_update_max: int = 1 # 每轮最多修改1次人设图 - - # 工作记忆链限制 - task_query_max: int = 4 # 每轮最多查询4次工作记忆链 - task_update_max: int = 2 # 每轮最多修改2次工作记忆链 - - # 一般记忆限制 - memory_query_max: int = 20 # 每轮最多查询20次一般记忆 - memory_update_max: int = 10 # 每轮最多修改10次一般记忆 - - -@dataclass -class ToolCallCount: - """工具调用计数""" - # 人设图 - persona_query: int = 0 - persona_update: int = 0 - - # 工作记忆链 - task_query: int = 0 - task_update: int = 0 - - # 一般记忆 - memory_query: int = 0 - memory_update: int = 0 - - -class ToolLimiter: - """工具调用限制器""" - - def __init__(self, limits: Optional[ToolLimits] = None): - self.limits = limits or ToolLimits() - self.counts = ToolCallCount() - - def _classify_tool(self, tool_name: str, arguments: dict) -> tuple: - """ - 分类工具调用 - 返回: (category, operation) - category: 'persona', 'task', 'memory' - operation: 'query', 'update' - """ - # 人设图工具 - if tool_name in ('persona_update', 'persona_clear'): - return ('persona', 'update') - - # 工作记忆链工具 - if tool_name in ('task_create', 'task_set_state', 'task_delete', 'task_link_info'): - # task_link_info 是关联操作,算作更新 - return ('task', 'update') - - # 一般记忆工具 - if tool_name == 'memory_recall': - # 判断是查询人设图、工作记忆链还是一般记忆 - query_intent = arguments.get('query_intent', '').lower() - - # 检查是否查询人设图 - if any(kw in query_intent for kw in ['人设', '角色', '性格', '语气', '说话风格', '扮演']): - return ('persona', 'query') - - # 检查是否查询工作记忆链 - if any(kw in query_intent for kw in ['tasknode', '工作记忆', '任务链', '任务', 'task']): - return ('task', 'query') - - # 一般记忆查询 - return ('memory', 'query') - - if tool_name == 'memory_commit': - return ('memory', 'update') - - if tool_name == 'memory_purge': - return ('memory', 'update') - - if tool_name == 'memory_introspect': - return ('memory', 'query') - - if tool_name in ('memory_archive', 'memory_cleanup'): - return ('memory', 'update') - - # 未知工具,归类为一般记忆更新 - return ('memory', 'update') - - def can_call(self, tool_name: str, arguments: dict) -> tuple: - """ - 检查是否允许调用工具 - 返回: (allowed, reason) - """ - category, operation = self._classify_tool(tool_name, arguments) - - # 获取当前计数和限制 - if category == 'persona': - if operation == 'query': - if self.counts.persona_query >= self.limits.persona_query_max: - return (False, f"人设图查询次数已达上限({self.limits.persona_query_max}次)") - else: # update - if self.counts.persona_update >= self.limits.persona_update_max: - return (False, f"人设图修改次数已达上限({self.limits.persona_update_max}次)") - - elif category == 'task': - if operation == 'query': - if self.counts.task_query >= self.limits.task_query_max: - return (False, f"工作记忆链查询次数已达上限({self.limits.task_query_max}次)") - else: # update - if self.counts.task_update >= self.limits.task_update_max: - return (False, f"工作记忆链修改次数已达上限({self.limits.task_update_max}次)") - - elif category == 'memory': - if operation == 'query': - if self.counts.memory_query >= self.limits.memory_query_max: - return (False, f"一般记忆查询次数已达上限({self.limits.memory_query_max}次)") - else: # update - if self.counts.memory_update >= self.limits.memory_update_max: - return (False, f"一般记忆修改次数已达上限({self.limits.memory_update_max}次)") - - return (True, "允许调用") - - def record_call(self, tool_name: str, arguments: dict) -> None: - """记录工具调用""" - category, operation = self._classify_tool(tool_name, arguments) - - if category == 'persona': - if operation == 'query': - self.counts.persona_query += 1 - else: - self.counts.persona_update += 1 - - elif category == 'task': - if operation == 'query': - self.counts.task_query += 1 - else: - self.counts.task_update += 1 - - elif category == 'memory': - if operation == 'query': - self.counts.memory_query += 1 - else: - self.counts.memory_update += 1 - - def get_summary(self) -> str: - """获取调用统计摘要""" - lines = [ - f"人设图: 查询{self.counts.persona_query}/{self.limits.persona_query_max}次, " - f"修改{self.counts.persona_update}/{self.limits.persona_update_max}次", - f"工作记忆链: 查询{self.counts.task_query}/{self.limits.task_query_max}次, " - f"修改{self.counts.task_update}/{self.limits.task_update_max}次", - f"一般记忆: 查询{self.counts.memory_query}/{self.limits.memory_query_max}次, " - f"修改{self.counts.memory_update}/{self.limits.memory_update_max}次" - ] - return "\n".join(lines) - - def reset(self) -> None: - """重置计数(新的一轮对话开始时调用)""" - self.counts = ToolCallCount() diff --git a/docs/api.md b/docs/api.md index d7cd31e..5967144 100644 --- a/docs/api.md +++ b/docs/api.md @@ -4,7 +4,7 @@ ## 概述 -TrulyMEM 后端采用**请求-响应队列模式**,通过 `queue.Queue` 实现线程安全通信。后端在独立线程中运行,处理来自客户端的请求。 +TrulyMEM 后端采用 **Packet 通信协议**,通过 `queue.Queue` 实现线程安全通信。后端在独立线程中运行,处理来自客户端的请求。 ### 核心组件 @@ -12,50 +12,51 @@ TrulyMEM 后端采用**请求-响应队列模式**,通过 `queue.Queue` 实现 |------|------| | `BackendServer` | 后端服务器,独立线程运行 | | `BackendClient` | 客户端封装,提供便捷方法 | -| `MessageType` | 请求类型枚举 | -| `BackendRequest` | 请求数据包 | -| `BackendResponse` | 响应数据包 | +| `PacketType` | 请求类型枚举 | +| `Packet` | 数据包(请求) | +| `PacketResponse` | 数据包响应 | --- -## 请求类型 (MessageType) +## 请求类型 (PacketType) ```python -class MessageType(Enum): +class PacketType(Enum): PROCESS_MESSAGE = "process_message" # 处理消息 EXECUTE_TOOL = "execute_tool" # 执行工具 GET_STATUS = "get_status" # 获取状态 GET_CONFIG = "get_config" # 获取配置 SET_CONFIG = "set_config" # 设置配置 - GET_HISTORY = "get_history" # 获取历史 - SAVE_HISTORY = "save_history" # 保存历史 - SHUTDOWN = "shutdown" # 关闭服务 + GET_HISTORY = "get_history" # 获取历史 + SAVE_HISTORY = "save_history" # 保存历史 + SHUTDOWN = "shutdown" # 关闭服务 ``` --- ## 数据包格式 -### BackendRequest +### Packet ```python @dataclass -class BackendRequest: - request_id: str # 请求唯一标识 - message_type: MessageType # 请求类型 - payload: Dict[str, Any] # 请求参数 - response_queue: queue.Queue # 响应队列(用于返回结果) +class Packet: + id: str # 唯一标识 + type: PacketType # 请求类型 + body: Dict[str, Any] # 请求参数 + response_queue: queue.Queue # 响应队列(可选) + created_at: float # 创建时间 ``` -### BackendResponse +### PacketResponse ```python @dataclass -class BackendResponse: - request_id: str # 对应的请求ID - success: bool # 是否成功 - data: Any = None # 返回数据 - error: Optional[str] = None # 错误信息 +class PacketResponse: + id: str # 对应的请求ID + success: bool # 是否成功 + data: Any = None # 返回数据 + error: Optional[str] = None # 错误信息 ``` --- @@ -68,14 +69,15 @@ class BackendResponse: **请求参数:** ```python -payload = { +body = { "user_input": str # 用户输入的消息 } ``` **响应数据:** ```python -data = { +{ + "success": True, "content": str, # AI 回复内容 "tool_calls": [ # 工具调用记录 { @@ -92,11 +94,16 @@ data = { **示例:** ```python -from core import BackendClient +from core import BackendServer, BackendClient + +server = BackendServer(db_path="graph_memory.db", use_embedded_db=True) +server.start(api_key="your-api-key") client = BackendClient(server) result = client.process_message("你好,请记住我的名字是小明") -# result = {"success": True, "data": {"content": "...", "tool_calls": [...]}} + +if result.get("success"): + print(result["content"]) ``` --- @@ -109,15 +116,16 @@ result = client.process_message("你好,请记住我的名字是小明") **请求参数:** ```python -payload = { +body = { "tool_name": str, # 工具名称 - "arguments": dict # 工具参数 + "arguments": dict # 工具参数 } ``` **响应数据:** ```python -data = { +{ + "success": True, "result": str # 工具执行结果 } ``` @@ -135,22 +143,23 @@ result = client.execute_tool("memory_recall", {"query_intent": "用户信息"}) **请求参数:** ```python -payload = {} # 无参数 +body = {} # 无参数 ``` **响应数据:** ```python -data = { - "graph_initialized": bool, # 图数据库是否初始化 - "client_initialized": bool, # API 客户端是否初始化 - "running": bool # 后端是否运行中 +{ + "running": bool, # 后端是否运行中 + "config": dict, # 当前配置 + "graph_initialized": bool, # 图数据库是否初始化 + "client_initialized": bool # API 客户端是否初始化 } ``` **示例:** ```python status = client.get_status() -# status = {"success": True, "data": {"running": True, ...}} +print(status["data"]["running"]) # True ``` --- @@ -161,12 +170,12 @@ status = client.get_status() **请求参数:** ```python -payload = {} # 无参数 +body = {} # 无参数 ``` **响应数据:** ```python -data = { +{ "api_key": str, # API Key "base_url": str # API Base URL } @@ -180,7 +189,7 @@ data = { **请求参数:** ```python -payload = { +body = { "api_key": str, # API Key "base_url": str # API Base URL (默认: https://api.deepseek.com) } @@ -188,7 +197,7 @@ payload = { **响应数据:** ```python -data = { +{ "status": "config_updated" } ``` @@ -209,12 +218,12 @@ result = client.update_config( **请求参数:** ```python -payload = {} # 无参数 +body = {} # 无参数 ``` **响应数据:** ```python -data = { +{ "history": list # 消息历史列表 } ``` @@ -227,14 +236,14 @@ data = { **请求参数:** ```python -payload = { +body = { "messages": list # 消息列表 } ``` **响应数据:** ```python -data = { +{ "status": "history_saved" } ``` @@ -247,12 +256,12 @@ data = { **请求参数:** ```python -payload = {} # 无参数 +body = {} # 无参数 ``` **响应数据:** ```python -data = { +{ "status": "shutdown" } ``` @@ -275,36 +284,36 @@ client = BackendClient(server) # 3. 发送消息 result = client.process_message("你好") -print(result["data"]["content"]) +print(result["content"]) # 4. 关闭 client.shutdown() ``` -### 直接使用请求队列 +### 使用 Packet 协议 ```python import queue -from core import BackendServer, MessageType, BackendRequest, BackendResponse +from core import BackendServer, Packet, PacketType server = BackendServer() server.start(api_key="your-key") -# 创建请求 +# 创建请求包 response_queue = queue.Queue() -request = BackendRequest( - request_id="req-001", - message_type=MessageType.PROCESS_MESSAGE, - payload={"user_input": "你好"}, +packet = Packet( + id="req-001", + type=PacketType.PROCESS_MESSAGE, + body={"user_input": "你好"}, response_queue=response_queue ) # 发送请求 -server._request_queue.put(request) +result = server.send(packet) +print(result.body) -# 等待响应 -response = response_queue.get(timeout=30.0) -print(response.data) +# 关闭 +server.shutdown() ``` --- @@ -357,7 +366,7 @@ client = BackendClient(server) async def handler(websocket): async for message in websocket: data = json.loads(message) - msg_type = data["type"] + msg_type = data.get("type") if msg_type == "message": result = client.process_message(data["content"]) @@ -373,41 +382,18 @@ async def handler(websocket): async def main(): server.start() async with websockets.serve(handler, "localhost", 8765): - await asyncio.Future() # run forever + await asyncio.Future() asyncio.run(main()) ``` -### 扩展为 gRPC - -```protobuf -// truly_mem.proto -syntax = "proto3"; - -service TrulyMEM { - rpc ProcessMessage(MessageRequest) returns (MessageResponse); - rpc UpdateConfig(ConfigRequest) returns (ConfigResponse); - rpc GetStatus(Empty) returns (StatusResponse); -} - -message MessageRequest { - string message = 1; -} - -message MessageResponse { - bool success = 1; - string content = 2; - string error = 3; -} -``` - --- ## 线程安全说明 - `BackendServer` 使用 `threading.Lock` 保护共享资源 - 所有请求通过 `queue.Queue` 传递,线程安全 -- 响应通过每个请求独立的 `response_queue` 返回 +- 响应通过每个请求独立的响应队列返回 - 默认超时时间:30 秒 --- @@ -463,4 +449,4 @@ message MessageResponse { |---------|------| | `API Key 未配置` | 未设置 API Key | | `timeout` | 请求超时 | -| `工具调用被拒绝: ...` | 工具调用频率超限 | +| `工具调用被拒绝: ...` | 工具调用频率超限 | \ No newline at end of file diff --git a/docs/一键启动指南.md b/docs/一键启动指南.md index ae34979..37fa885 100644 --- a/docs/一键启动指南.md +++ b/docs/一键启动指南.md @@ -5,20 +5,17 @@ ### 从源码运行 ```bash -# 克隆仓库 git clone cd TrulyMEM-TrueHumanMEM -# 安装依赖 pip install -r requirements.txt -# 运行 python trulymem_entry.py ``` ### 打包后运行 -打包后会生成可执行文件(Windows: TrulyMEM.exe, Linux/macOS: TrulyMEM): +打包后会生成可执行文件: ```bash # Linux/macOS @@ -38,9 +35,11 @@ TrulyMEM.exe 1. 运行应用 2. 按 **F2** 展开侧边栏 -3. 输入 **API Key** +3. 输入 **API Key**、**模型**、**Base URL** 4. 按 **Enter** 保存 +配置会自动保存到 `~/.trulymem/config.json`,下次启动自动加载。 + ## 快捷键 | 按键 | 功能 | @@ -53,9 +52,32 @@ TrulyMEM.exe ## 数据存储 -- **数据库**: `graph_memory.db`(应用目录) -- **配置**: `config.json`(应用目录) -- **格式**: SQLite +| 数据 | 位置 | +|------|------| +| 图数据库 | `graph_memory.db`(应用目录) | +| 配置文件 | `~/.trulymem/config.json` | +| 数据库格式 | SQLite | + +## 架构说明 + +### 通信协议 + +UI 与后端通过 **Packet 协议** 通信: + +``` +UI (Textual TUI) + ↓ BackendClient +Packet → queue.Queue → BackendServer (独立线程) + ↓ +处理请求 → 返回响应 +``` + +### 配置管理 + +- **存储位置**: `~/.trulymem/config.json` +- **自动加载**: 启动时从文件读取配置 +- **动态更新**: 运行时修改配置立即生效 +- **持久化**: 修改后自动保存到文件 ## 常见问题 @@ -85,7 +107,7 @@ pip install -r requirements.txt # 运行测试 pytest tests/ -# 打包(需 PyInstaller) +# 打包 bash build/build_windows.bat # Windows bash build/build_linux.sh # Linux -``` +``` \ No newline at end of file diff --git a/docs/架构.md b/docs/架构.md index a6f3a4a..0304c66 100644 --- a/docs/架构.md +++ b/docs/架构.md @@ -6,6 +6,7 @@ - 极简视觉,信息密度优先 - 工具痕迹默认隐藏,需要时可展开 - TUI 与后端分离,多线程通信 +- **一切皆图**,AI 推理全部在后端 ## 项目结构 @@ -14,53 +15,25 @@ TrulyMEM-TrueHumanMEM/ ├── trulymem_entry.py # 入口:先启动 core → 再启动 ui ├── core/ # 后端/业务逻辑 │ ├── __init__.py # 导出 BackendServer, BackendClient, EmbeddedGraphDB -│ ├── server.py # BackendServer (多线程队列通信) -│ ├── client.py # BackendClient +│ ├── server.py # BackendServer (Packet 通信协议) +│ ├── client.py # BackendClient (Packet 协议客户端) │ ├── embedded_db.py # SQLite 图数据库实现 -│ ├── graph_client.py -│ ├── tool_executor.py # 工具执行器 +│ ├── graph_client.py # OpenAI/DeepSeek API 客户端 +│ ├── tool_executor.py # 工具执行器 │ ├── tool_limiter.py # 工具调用限制器 -│ ├── memory_tools.py # 工具定义 -│ ├── prompts/ # 提示词管理 +│ ├── tools/ # 工具定义 │ │ ├── __init__.py -│ │ ├── prompt_manager.py -│ │ └── templates/ -│ │ └── system_prompt.md -│ └── tools/ # 工具模块 -│ ├── __init__.py -│ ├── memory_tools.py -│ ├── tool_executor.py -│ └── tool_limiter.py -├── ui/ # TUI 显示层 -│ ├── __init__.py # 导出 GraphMemoryApp, AppConfig -│ ├── app.py # GraphMemoryApp (纯显示) -│ ├── widgets/ # TUI 组件 -│ │ ├── left_panel.py -│ │ ├── right_panel.py -│ │ ├── message_history.py -│ │ ├── message_widget.py -│ │ ├── input_box.py -│ │ ├── config_section.py -│ │ ├── operation_log.py -│ │ ├── cypher_query_box.py -│ │ └── status_bar.py -│ ├── handlers/ # 事件处理 -│ │ ├── focus_handler.py -│ │ ├── key_handler.py -│ │ └── message_handler.py -│ ├── models/ # 数据模型 -│ │ ├── message.py -│ │ ├── config.py -│ │ └── log_entry.py -│ ├── services/ # 服务层 -│ │ ├── config_manager.py -│ │ ├── config_service.py -│ │ └── chat_service.py -│ └── styles/ # 样式文件 -│ ├── app.css -│ ├── components.css -│ └── messages.css -└── tests/ # 测试 (38 tests) +│ │ └── memory_tools.py +│ └── prompts/ # 提示词管理 +├── ui/ # TUI 显示层(仅显示,无 AI 逻辑) +│ ├── __init__.py # 导出 GraphMemoryApp, AppConfig +│ ├── app.py # GraphMemoryApp (通过 BackendClient 通信) +│ ├── widgets/ # TUI 组件 +│ ├── handlers/ # 事件处理 +│ ├── models/ # 数据模型 +│ ├── services/ # 服务层(仅配置管理) +│ └── styles/ # 样式文件 +└── tests/ # 测试 (42 tests) ``` ## 架构图 @@ -69,13 +42,14 @@ TrulyMEM-TrueHumanMEM/ trulymem_entry.py │ ├─ BackendServer.start() → 独立线程运行 - │ ├─ 处理 PROCESS_MESSAGE 请求 - │ ├─ 处理 EXECUTE_TOOL 请求 + │ ├─ 处理 PROCESS_MESSAGE 请求 → AI 推理 + 工具调用 + │ ├─ 处理 EXECUTE_TOOL 请求 → 外部工具调用(不限次数) + │ ├─ 处理 GET/SET_CONFIG 请求 │ └─ 管理 GraphMemoryClient, EmbeddedGraphDB │ └─ GraphMemoryApp(backend_server=server) │ - └─ BackendClient ← queue.Queue → BackendServer + └─ BackendClient ← Packet 通信 → BackendServer ``` ## 组件职责 @@ -84,22 +58,38 @@ trulymem_entry.py | 组件 | 职责 | |------|------| -| `server.py` | 多线程队列通信,处理消息和工具调用 | -| `client.py` | TUI 端的通信客户端 | +| `server.py` | Packet 协议处理,多线程队列通信,AI 推理,工具限制 | +| `client.py` | 客户端封装,UI 与后端通信桥梁 | | `embedded_db.py` | SQLite 图数据库 CRUD | | `graph_client.py` | OpenAI/DeepSeek API 客户端 | | `tool_executor.py` | 工具执行逻辑 | -| `tool_limiter.py` | 工具调用频率限制 | +| `tool_limiter.py` | 工具调用频率限制(仅限 AI 推理) | ### ui/ (显示层) | 组件 | 职责 | |------|------| -| `app.py` | Textual 应用主类 | -| `widgets/` | TUI 组件(面板、输入框等) | -| `handlers/` | 事件处理(键盘、焦点) | -| `models/` | 数据模型(消息、配置) | -| `services/` | 配置管理、服务层 | +| `app.py` | Textual 应用主类,仅通过 BackendClient 通信 | +| `services/` | 仅配置管理,无 AI 逻辑 | + +### 通信协议 + +UI 与后端通过 **Packet 通信协议** 交互: + +```python +from core import BackendServer, BackendClient, Packet, PacketType + +# 后端启动 +server = BackendServer(db_path="graph_memory.db", use_embedded_db=True) +server.start(api_key="your-key") + +# 客户端通信 +client = BackendClient(server) +result = client.process_message("你好") # AI 推理 +result = client.execute_tool("memory_introspect", {}) # 外部工具调用 +``` + +--- ## 数据流 @@ -108,21 +98,27 @@ trulymem_entry.py ↓ BackendClient.process_message(user_input) ↓ -queue.Queue → BackendServer (独立线程) +Packet (type=PROCESS_MESSAGE) → queue.Queue + ↓ +BackendServer (独立线程) ↓ GraphMemoryClient.send_message_with_history() ↓ OpenAI API / DeepSeek API ↓ -execute_tool() → EmbeddedGraphDB +execute_tool() + ToolLimiter (AI 推理时受限) + ↓ +EmbeddedGraphDB (图数据库) ↓ 循环调用 API 直到无 tool_calls ↓ -queue.Queue → 返回结果 +Packet 响应返回 ↓ MessageHistory 显示 ``` +--- + ## 启动流程 ```python @@ -146,82 +142,7 @@ def main(): backend_server.shutdown() ``` -## 布局结构 - -### 默认视图(右侧展开) - -``` -┌─────────────────────────────────────┬─────────────────────┐ -│ │ F2:隐藏侧边栏 │ -│ 🟠 14:30:25 │ ────────────────────│ -│ 用户: 量子力学是什么? │ API Key: *** │ -│ │ 模型: deepseek-chat │ -│ 🔵 14:30:26 │ Base URL: ... │ -│ 是的!根据记忆... │ ────────────────────│ -│ [工具:2次] (F3展开) │ [操作日志] │ -│ │ 14:30:26 recall │ -│ ┌─────────────────────────────┐ │ 实体: 量子力学 │ -│ │ 🟠 [输入框...] │ │ ───────────────────│ -│ └─────────────────────────────┘ │ >[查询...] │ -└─────────────────────────────────────┴─────────────────────┘ -``` - -### F2后(右侧折叠) - -``` -┌─────────────────────────────────────┐ -│ 🟠 14:30:25 │ -│ 用户: 量子力学是什么? │ -│ │ -│ 🔵 14:30:26 │ -│ 是的!根据记忆... │ -│ [工具:2次] (F3展开) │ -│ │ -│ ┌─────────────────────────────┐ │ -│ │ 🟠 [输入框...] │ │ -│ └─────────────────────────────┘ │ -│ F1:帮助 F2:展开 F5:清屏 F6:退出 │ -└─────────────────────────────────────┘ -``` - -## 快捷键 - -| 按键 | 功能 | -|------|------| -| F1 | 显示帮助 | -| F2 | 切换侧边栏 | -| F3 | 工具详情 | -| F5 | 清屏 | -| F6 | 退出 | - -## 技术栈 - -| 技术 | 用途 | -|------|------| -| Python 3.8+ | 编程语言 | -| Textual 0.47+ | TUI 框架 | -| SQLite | 图数据库(默认内嵌) | -| OpenAI SDK | API 调用(兼容 DeepSeek) | -| threading.Queue | 多线程通信 | -| PyInstaller | 打包 | - -## 数据库模式 - -### SQLite 内嵌(默认) - -```python -# core/embedded_db.py -class EmbeddedGraphDB: - def __init__(self, db_path="graph_memory.db"): - self.conn = sqlite3.connect(db_path, check_same_thread=False) -``` - -### 可选:Neo4j - -```bash -export USE_EMBEDDED_DB=false -docker run -d --name neo4j -p 7474:7474 -p 7687:7687 neo4j:latest -``` +--- ## 工具系统 @@ -241,4 +162,19 @@ docker run -d --name neo4j -p 7474:7474 -p 7687:7687 neo4j:latest - `task_create` - 创建任务 - `task_set_state` - 设置状态 - `task_delete` - 删除任务 -- `task_link_info` - 关联信息 \ No newline at end of file +- `task_link_info` - 关联信息 + +--- + +## 错误处理原则 + +所有 API **不抛出异常**,错误通过返回字典传递: + +```python +result = client.process_message("hello") + +if result.get("success"): + print(result["content"]) +else: + print(result["error"]) # 错误描述 +``` \ No newline at end of file diff --git a/tests/test_core/test_packet.py b/tests/test_core/test_packet.py index ad9d0c1..bccd9f2 100644 --- a/tests/test_core/test_packet.py +++ b/tests/test_core/test_packet.py @@ -9,304 +9,300 @@ os.environ["DEEPSEEK_API_KEY"] = "fake-test-key" class TestPacketTypeEnum: - def test_packet_type_message_exists(self): + """测试 PacketType 枚举""" + + def test_packet_type_process_message_exists(self): from core import PacketType - assert PacketType.MESSAGE is not None - assert PacketType.MESSAGE.value == "message" + assert PacketType.PROCESS_MESSAGE is not None + assert PacketType.PROCESS_MESSAGE.value == "process_message" - def test_packet_type_config_exists(self): + def test_packet_type_execute_tool_exists(self): from core import PacketType - assert PacketType.CONFIG is not None - assert PacketType.CONFIG.value == "config" + assert PacketType.EXECUTE_TOOL is not None + assert PacketType.EXECUTE_TOOL.value == "execute_tool" - def test_packet_type_tool_exists(self): + def test_packet_type_get_status_exists(self): from core import PacketType - assert PacketType.TOOL is not None - assert PacketType.TOOL.value == "tool" + assert PacketType.GET_STATUS is not None + assert PacketType.GET_STATUS.value == "get_status" - def test_packet_type_status_exists(self): + def test_packet_type_get_config_exists(self): from core import PacketType - assert PacketType.STATUS is not None - assert PacketType.STATUS.value == "status" + assert PacketType.GET_CONFIG is not None + assert PacketType.GET_CONFIG.value == "get_config" - def test_packet_type_history_exists(self): + def test_packet_type_set_config_exists(self): from core import PacketType - assert PacketType.HISTORY is not None - assert PacketType.HISTORY.value == "history" + assert PacketType.SET_CONFIG is not None + assert PacketType.SET_CONFIG.value == "set_config" + + def test_packet_type_get_history_exists(self): + from core import PacketType + assert PacketType.GET_HISTORY is not None + assert PacketType.GET_HISTORY.value == "get_history" + + def test_packet_type_save_history_exists(self): + from core import PacketType + assert PacketType.SAVE_HISTORY is not None + assert PacketType.SAVE_HISTORY.value == "save_history" + + def test_packet_type_shutdown_exists(self): + from core import PacketType + assert PacketType.SHUTDOWN is not None + assert PacketType.SHUTDOWN.value == "shutdown" def test_packet_type_all_values(self): from core import PacketType values = [pt.value for pt in PacketType] - assert "message" in values - assert "config" in values - assert "tool" in values - assert "status" in values - assert "history" in values - assert len(values) == 5 + assert "process_message" in values + assert "execute_tool" in values + assert "get_status" in values + assert "get_config" in values + assert "set_config" in values + assert "get_history" in values + assert "save_history" in values + assert "shutdown" in values + assert len(values) == 8 class TestPacketCreation: + """测试 Packet 创建""" + def test_packet_with_id_and_type(self): from core import Packet, PacketType - packet = Packet(id="test-1", type=PacketType.MESSAGE, body={"message": "hello"}) + packet = Packet(id="test-1", type=PacketType.PROCESS_MESSAGE, body={"user_input": "hello"}) assert packet.id == "test-1" - assert packet.type == PacketType.MESSAGE + assert packet.type == PacketType.PROCESS_MESSAGE def test_packet_body(self): from core import Packet, PacketType - body = {"message": "test", "extra": "data"} - packet = Packet(id="test-2", type=PacketType.CONFIG, body=body) + body = {"user_input": "test", "extra": "data"} + packet = Packet(id="test-2", type=PacketType.EXECUTE_TOOL, body=body) assert packet.body == body + def test_packet_with_empty_body(self): + from core import Packet, PacketType + packet = Packet(id="test-3", type=PacketType.GET_STATUS, body={}) + assert packet.body == {} + def test_packet_created_at_default(self): from core import Packet, PacketType before = time.time() - packet = Packet(id="test-3", type=PacketType.STATUS, body={}) + packet = Packet(id="test-4", type=PacketType.GET_CONFIG, body={}) after = time.time() assert before <= packet.created_at <= after - def test_packet_with_custom_created_at(self): - from core import Packet, PacketType - custom_time = 1234567890.0 - packet = Packet(id="test-4", type=PacketType.HISTORY, body={}, created_at=custom_time) - assert packet.created_at == custom_time + +class TestPacketResponse: + """测试 PacketResponse""" + + def test_packet_response_success(self): + from core import PacketResponse + response = PacketResponse(id="resp-1", success=True, data={"result": "ok"}) + assert response.id == "resp-1" + assert response.success is True + assert response.data == {"result": "ok"} + + def test_packet_response_error(self): + from core import PacketResponse + response = PacketResponse(id="resp-2", success=False, error="error msg") + assert response.id == "resp-2" + assert response.success is False + assert response.error == "error msg" -class TestBackendServerInit: - def test_create_server_defaults(self): +class TestBackendServerCreation: + """测试 BackendServer 创建""" + + def test_backend_server_init(self): + from core import BackendServer + server = BackendServer(db_path=":memory:", use_embedded_db=True) + assert server._db_path == ":memory:" + assert server._use_embedded_db is True + assert server._graph is None + assert server._client is None + assert server._tool_limiter is None + + def test_backend_server_default_params(self): from core import BackendServer server = BackendServer() assert server._db_path == "graph_memory.db" assert server._use_embedded_db is True - assert server._graph is None - assert server._client is None + + +class TestBackendServerLifecycle: + """测试 BackendServer 生命周期""" + + def test_backend_server_start_stop(self): + from core import BackendServer + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="") + + assert server._running is True + assert server._graph is not None + assert server._tool_limiter is not None + + server.shutdown() assert server._running is False - def test_create_server_custom_db(self): + def test_backend_server_start_with_api_key(self): from core import BackendServer - server = BackendServer(db_path="custom.db") - assert server._db_path == "custom.db" - - def test_create_server_no_embedded(self): - from core import BackendServer - server = BackendServer(use_embedded_db=False) - assert server._use_embedded_db is False + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="test-key", base_url="https://api.deepseek.com") + + assert server._client is not None + assert server._config["api_key"] == "test-key" + + server.shutdown() -class TestBackendServerStart: - def test_start_without_api_key(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - assert server._running is True - assert server._graph is not None - assert server._client is None - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - def test_start_with_api_key(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="test-key", base_url="https://api.test.com") - assert server._running is True - assert server._config["api_key"] == "test-key" - assert server._config["base_url"] == "https://api.test.com" - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - def test_start_twice_returns_early(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - running_before = server._running - server.start(api_key="") - running_after = server._running - assert running_before is True - assert running_after is True - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - -class TestBackendServerShutdown: - def test_shutdown_stops_server(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - assert server._running is True - server.shutdown() - assert server._running is False - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - def test_shutdown_closes_graph(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - server.shutdown() - assert server._graph is None - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - -class TestBackendServerPacketHandling: - def test_send_message_without_client(self): - from core import BackendServer, Packet, PacketType - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - packet = Packet(id="1", type=PacketType.MESSAGE, body={"message": "hello"}) - response = server.send(packet) - assert response.body["success"] is True - assert "API Key not configured" in response.body["error"] - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - def test_send_status_packet(self): - from core import BackendServer, Packet, PacketType - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - packet = Packet(id="2", type=PacketType.STATUS, body={}) - response = server.send(packet) - assert response.body["success"] is True - assert response.body["running"] is True - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - -class TestBackendClientInit: - def test_create_client_with_server(self): +class TestBackendClientCreation: + """测试 BackendClient 创建""" + + def test_backend_client_init(self): from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - client = BackendClient(server) - assert client._server is server - assert client._counter == 0 - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) + server = BackendServer(db_path=":memory:", use_embedded_db=True) + client = BackendClient(server) + + assert client._server is server + assert client._counter == 0 -class TestBackendClientMethods: +class TestBackendClientAPI: + """测试 BackendClient API""" + def test_get_status(self): from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - client = BackendClient(server) - status = client.get_status() - assert status["success"] is True - assert status["running"] is True - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="") + client = BackendClient(server) + + result = client.get_status() + assert result.get("success") is True + data = result.get("data", {}) + assert data.get("running") is True + + server.shutdown() def test_update_config(self): from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - client = BackendClient(server) - result = client.update_config(api_key="new-key", base_url="https://new-api.test.com") - assert result["success"] is True - assert result["status"] == "config_updated" - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="") + client = BackendClient(server) + + result = client.update_config(api_key="new-key", base_url="https://api.deepseek.com") + assert result.get("success") is True + + server.shutdown() - def test_save_and_get_history(self): + def test_get_config(self): from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - client = BackendClient(server) - - test_messages = [{"role": "user", "content": "hello"}] - client.save_history(test_messages) - - history = client.get_history() - assert history == test_messages - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="test-key") + client = BackendClient(server) + + result = client.get_config() + assert result.get("success") is True + assert result.get("data", {}).get("api_key") == "test-key" + + server.shutdown() - -class TestMultipleClients: - def test_multiple_clients_thread_safety(self): + def test_process_message_no_api_key(self): from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - client = BackendClient(server) - - status1 = client.get_status() - status2 = client.get_status() - - assert status1["success"] is True - assert status2["success"] is True - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="") + client = BackendClient(server) + + # 无 API key 应该返回错误(success 为 False) + result = client.process_message("hello") + # 由于 API 调用失败,success 应该是 False + assert result.get("success") is False + + server.shutdown() + + def test_execute_tool(self): + from core import BackendServer, BackendClient + server = BackendServer(db_path=":memory:", use_embedded_db=True) + server.start(api_key="") + client = BackendClient(server) + + result = client.execute_tool("memory_introspect", {}) + assert result.get("success") is True + + server.shutdown() -class TestServerStateAfterShutdown: - def test_server_state_not_running(self): - from core import BackendServer - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - assert server._running is True - server.shutdown() - assert server._running is False - finally: - if os.path.exists(db_path): - os.unlink(db_path) \ No newline at end of file +class TestToolLimiter: + """测试工具限制器""" + + def test_tool_limiter_init(self): + from core.tool_limiter import ToolLimiter + limiter = ToolLimiter() + assert limiter.counts.persona_query == 0 + assert limiter.counts.persona_update == 0 + + def test_tool_limiter_classify(self): + from core.tool_limiter import ToolLimiter + limiter = ToolLimiter() + + category, operation = limiter._classify_tool("memory_recall", {"query_intent": "test"}) + assert category == "memory" + assert operation == "query" + + category, operation = limiter._classify_tool("persona_update", {}) + assert category == "persona" + assert operation == "update" + + category, operation = limiter._classify_tool("task_create", {}) + assert category == "task" + assert operation == "update" + + def test_tool_limiter_can_call(self): + from core.tool_limiter import ToolLimiter + limiter = ToolLimiter() + + allowed, reason = limiter.can_call("persona_update", {}) + assert allowed is True + + limiter.record_call("persona_update", {}) + allowed, reason = limiter.can_call("persona_update", {}) + assert allowed is False + assert "已达上限" in reason + + def test_tool_limiter_reset(self): + from core.tool_limiter import ToolLimiter + limiter = ToolLimiter() + + limiter.record_call("persona_update", {}) + assert limiter.counts.persona_update == 1 + + limiter.reset() + assert limiter.counts.persona_update == 0 + + +class TestEmbeddedGraphDB: + """测试图数据库""" + + def test_embedded_db_init(self): + from core.embedded_db import EmbeddedGraphDB + db = EmbeddedGraphDB(db_path=":memory:") + assert db.conn is not None + db.close() + + def test_embedded_db_commit_and_recall(self): + from core.embedded_db import EmbeddedGraphDB + db = EmbeddedGraphDB(db_path=":memory:") + + # 写入记忆 (使用 triplets 参数) + result = db.commit( + triplets=[ + {"subject": "测试", "relation": "是", "object": "test"} + ], + session_id="test-session" + ) + + # 读取记忆 + results = db.recall("测试") + assert len(results.get("entities", [])) > 0 + + db.close() \ No newline at end of file diff --git a/tests/test_integration/__init__.py b/tests/test_integration/__init__.py index e817fd1..e70b2ef 100644 --- a/tests/test_integration/__init__.py +++ b/tests/test_integration/__init__.py @@ -6,8 +6,10 @@ os.environ["DEEPSEEK_API_KEY"] = "fake-test-key" class TestIntegrationPacketFlow: - def test_packet_round_trip_message(self): - from core import BackendServer, BackendClient, Packet, PacketType + """测试 Packet 通信流程""" + + def test_packet_round_trip_process_message(self): + from core import BackendServer, BackendClient with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: db_path = f.name try: @@ -15,9 +17,10 @@ class TestIntegrationPacketFlow: server.start(api_key="") client = BackendClient(server) - result = client.send_message("test message") - assert result["success"] is True - assert "API Key not configured" in result.get("error", "") + # 无 API key 时应该返回错误而非抛异常 + result = client.process_message("test message") + assert result.get("success") is False + assert "error" in result server.shutdown() finally: @@ -34,12 +37,12 @@ class TestIntegrationPacketFlow: client = BackendClient(server) result = client.update_config(api_key="test-api", base_url="https://test.com") - assert result["success"] is True - assert result["status"] == "config_updated" + assert result.get("success") is True status = client.get_status() - assert status["config"]["api_key"] == "test-api" - assert status["config"]["base_url"] == "https://test.com" + data = status.get("data", {}) + assert data.get("config", {}).get("api_key") == "test-api" + assert data.get("config", {}).get("base_url") == "https://test.com" server.shutdown() finally: @@ -55,17 +58,18 @@ class TestIntegrationPacketFlow: server.start(api_key="") client = BackendClient(server) - status = client.get_status() - assert status["success"] is True - assert status["running"] is True - assert status["client_ready"] is False + result = client.get_status() + assert result.get("success") is True + data = result.get("data", {}) + assert data.get("running") is True + assert data.get("graph_initialized") is True server.shutdown() finally: if os.path.exists(db_path): os.unlink(db_path) - def test_packet_round_trip_history(self): + def test_packet_round_trip_execute_tool(self): from core import BackendServer, BackendClient with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: db_path = f.name @@ -74,14 +78,8 @@ class TestIntegrationPacketFlow: server.start(api_key="") client = BackendClient(server) - messages = [ - {"role": "user", "content": "hello"}, - {"role": "assistant", "content": "hi there"} - ] - client.save_history(messages) - - retrieved = client.get_history() - assert retrieved == messages + result = client.execute_tool("memory_introspect", {}) + assert result.get("success") is True server.shutdown() finally: @@ -89,8 +87,10 @@ class TestIntegrationPacketFlow: os.unlink(db_path) -class TestIntegrationSequentialOperations: - def test_sequential_config_updates(self): +class TestIntegrationToolLimiter: + """测试工具限制器集成""" + + def test_external_tool_call_not_limited(self): from core import BackendServer, BackendClient with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: db_path = f.name @@ -99,36 +99,38 @@ class TestIntegrationSequentialOperations: server.start(api_key="") client = BackendClient(server) - client.update_config(api_key="key1", base_url="https://api1.com") - status1 = client.get_status() - assert status1["config"]["api_key"] == "key1" - - client.update_config(api_key="key2", base_url="https://api2.com") - status2 = client.get_status() - assert status2["config"]["api_key"] == "key2" + # 外部调用多次应该成功 + for i in range(5): + result = client.execute_tool("memory_introspect", {}) + assert result.get("success") is True server.shutdown() finally: if os.path.exists(db_path): os.unlink(db_path) - def test_save_history_override(self): + def test_internal_tool_call_limited(self): from core import BackendServer, BackendClient with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: db_path = f.name try: server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") + server.start(api_key="fake-key") # 假 key 会失败但不影响测试 client = BackendClient(server) - client.save_history([{"role": "user", "content": "first"}]) - history1 = client.get_history() - assert len(history1) == 1 + # 内部调用受限,tool_limiter 存在 + assert server._tool_limiter is not None - client.save_history([{"role": "user", "content": "second"}]) - history2 = client.get_history() - assert len(history2) == 1 - assert history2[0]["content"] == "second" + # 初始状态 + assert server._tool_limiter.counts.persona_update == 0 + + # 记录一次调用 + server._tool_limiter.record_call("persona_update", {}) + assert server._tool_limiter.counts.persona_update == 1 + + # 再次调用应该被拒绝 + allowed, reason = server._tool_limiter.can_call("persona_update", {}) + assert allowed is False server.shutdown() finally: @@ -136,117 +138,10 @@ class TestIntegrationSequentialOperations: os.unlink(db_path) -class TestIntegrationMultipleClients: - def test_two_clients_same_server(self): - from core import BackendServer, BackendClient - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - - client1 = BackendClient(server) - client2 = BackendClient(server) - - status1 = client1.get_status() - status2 = client2.get_status() - - assert status1["success"] is True - assert status2["success"] is True - - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - -class TestIntegrationDatabase: - def test_server_creates_database(self): - from core import BackendServer - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "test.db") - server = BackendServer(db_path=db_path, use_embedded_db=True) - server.start(api_key="") - - assert os.path.exists(db_path) - - server.shutdown() - assert os.path.exists(db_path) - - def test_server_persists_across_restart(self): - from core import BackendServer, BackendClient - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "persist.db") - - server1 = BackendServer(db_path=db_path, use_embedded_db=True) - server1.start(api_key="") - server1.shutdown() - - server2 = BackendServer(db_path=db_path, use_embedded_db=True) - server2.start(api_key="") - assert os.path.exists(db_path) - server2.shutdown() - - -class TestIntegrationEntryPoint: - def test_entry_import(self): - import trulymem_entry - assert trulymem_entry is not None - - def test_entry_has_main(self): - import trulymem_entry - assert hasattr(trulymem_entry, 'main') - assert callable(trulymem_entry.main) - - -class TestIntegrationAllPacketTypes: - def test_message_type_string(self): - from core import PacketType - assert PacketType.MESSAGE.value == "message" - - def test_config_type_string(self): - from core import PacketType - assert PacketType.CONFIG.value == "config" - - def test_tool_type_string(self): - from core import PacketType - assert PacketType.TOOL.value == "tool" - - def test_status_type_string(self): - from core import PacketType - assert PacketType.STATUS.value == "status" - - def test_history_type_string(self): - from core import PacketType - assert PacketType.HISTORY.value == "history" - - class TestIntegrationErrorHandling: - def test_timeout_on_slow_response(self): - from core import BackendServer, Packet, PacketType - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: - db_path = f.name - try: - server = BackendServer(db_path=db_path) - server.start(api_key="") - server._running = False - - packet = Packet(id="timeout-test", type=PacketType.MESSAGE, body={"message": "test"}) - - original_timeout = 30.0 - server.send = lambda p, timeout=original_timeout: ( - setattr(server, '_running', True), - server.send(p) - )[1] - - server.shutdown() - finally: - if os.path.exists(db_path): - os.unlink(db_path) - - -class TestIntegrationFullWorkflow: - def test_complete_workflow(self): + """测试错误处理""" + + def test_process_message_returns_error_not_raise(self): from core import BackendServer, BackendClient with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: db_path = f.name @@ -255,24 +150,30 @@ class TestIntegrationFullWorkflow: server.start(api_key="") client = BackendClient(server) - status_before = client.get_status() - assert status_before["success"] is True - - client.update_config(api_key="workflow-key", base_url="https://workflow.com") - status_after_config = client.get_status() - assert status_after_config["config"]["api_key"] == "workflow-key" - - history = [{"role": "user", "content": "test workflow"}] - client.save_history(history) - retrieved_history = client.get_history() - assert retrieved_history == history - - message_result = client.send_message("test") - assert message_result["success"] is True + # 应该返回错误,而不是抛出异常 + result = client.process_message("hello") + assert result.get("success") is False + assert "error" in result + + server.shutdown() + finally: + if os.path.exists(db_path): + os.unlink(db_path) + + def test_execute_tool_error_handling(self): + from core import BackendServer, BackendClient + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + db_path = f.name + try: + server = BackendServer(db_path=db_path, use_embedded_db=True) + server.start(api_key="") + client = BackendClient(server) + + # 不存在的工具应该返回错误 + result = client.execute_tool("nonexistent_tool", {}) + assert result.get("success") is False server.shutdown() - status_after_shutdown = client.get_status() - assert status_after_shutdown["running"] is False finally: if os.path.exists(db_path): os.unlink(db_path) \ No newline at end of file diff --git a/tests/test_ui/__init__.py b/tests/test_ui/__init__.py index 2097c4d..6497e77 100644 --- a/tests/test_ui/__init__.py +++ b/tests/test_ui/__init__.py @@ -6,7 +6,9 @@ from pathlib import Path os.environ["DEEPSEEK_API_KEY"] = "fake-test-key" -class TestUIImportFull: +class TestUIImport: + """测试 UI 模块导入""" + def test_import_graphmemoryapp(self): from ui import GraphMemoryApp assert GraphMemoryApp is not None @@ -15,52 +17,24 @@ class TestUIImportFull: from ui import AppConfig assert AppConfig is not None - def test_import_from_models(self): + def test_import_message(self): from ui.models.message import Message, ToolCall, ToolResult assert Message is not None assert ToolCall is not None assert ToolResult is not None - def test_import_from_models_config(self): + def test_import_config(self): from ui.models.config import AppConfig assert AppConfig is not None - def test_import_from_models_log_entry(self): + def test_import_log_entry(self): from ui.models.log_entry import LogEntry assert LogEntry is not None - def test_import_from_widgets(self): - from ui.widgets.left_panel import LeftPanel - from ui.widgets.right_panel import RightPanel - from ui.widgets.input_box import InputBox - from ui.widgets.message_history import MessageHistory - from ui.widgets.status_bar import StatusBar - assert LeftPanel is not None - assert RightPanel is not None - assert InputBox is not None - assert MessageHistory is not None - assert StatusBar is not None - def test_import_from_handlers(self): - from ui.handlers.focus_handler import FocusHandler - from ui.handlers.key_handler import KeyHandler - from ui.handlers.message_handler import MessageHandler - assert FocusHandler is not None - assert KeyHandler is not None - assert MessageHandler is not None - - def test_import_from_services(self): - from ui.services.config_manager import ConfigManager - from ui.services.config_service import ConfigService - from ui.services.chat_service import ChatService - from ui.services.tool_service import ToolService - assert ConfigManager is not None - assert ConfigService is not None - assert ChatService is not None - assert ToolService is not None - - -class TestAppConfigFull: +class TestAppConfig: + """测试配置模型""" + def test_config_default_values(self): from ui.models.config import AppConfig config = AppConfig() @@ -68,25 +42,15 @@ class TestAppConfigFull: assert config.model == "deepseek-chat" assert config.base_url == "https://api.deepseek.com" - def test_config_from_env_with_key(self): + def test_config_from_env(self): from ui.models.config import AppConfig config = AppConfig.from_env() assert "fake-test-key" in config.api_key - def test_config_from_env_custom(self): - os.environ["DEEPSEEK_API_KEY"] = "my-key" - os.environ["MODEL_NAME"] = "my-model" - os.environ["DEEPSEEK_BASE_URL"] = "https://my-api.com" - - from ui.models.config import AppConfig - config = AppConfig.from_env() - - assert config.api_key == "my-key" - assert config.model == "my-model" - assert config.base_url == "https://my-api.com" - -class TestMessageModelFull: +class TestMessageModel: + """测试消息模型""" + def test_message_creation_user(self): from ui.models.message import Message from datetime import datetime @@ -99,119 +63,36 @@ class TestMessageModelFull: from ui.models.message import Message msg = Message(role="assistant", content="assistant response") assert msg.role == "assistant" - assert msg.content == "assistant response" def test_message_with_tool_calls(self): from ui.models.message import Message, ToolCall tc = ToolCall(id="call-1", name="memory_recall", arguments={"query": "test"}) msg = Message(role="assistant", content="response", tool_calls=[tc]) - assert msg.role == "assistant" assert msg.tool_calls is not None assert len(msg.tool_calls) == 1 - assert msg.tool_calls[0].name == "memory_recall" - - def test_message_with_tool_results(self): - from ui.models.message import Message, ToolResult - tr = ToolResult(tool_call_id="call-1", name="memory_recall", arguments={}, result="result", success=True) - msg = Message(role="assistant", content="response", tool_results=[tr]) - assert msg.tool_results is not None - assert len(msg.tool_results) == 1 - assert msg.tool_results[0].success is True -class TestToolCallModelFull: - def test_toolcall_creation(self): - from ui.models.message import ToolCall - tc = ToolCall(id="call-1", name="memory_recall", arguments={"query": "test"}) - assert tc.id == "call-1" - assert tc.name == "memory_recall" - assert tc.arguments["query"] == "test" - - -class TestToolResultModelFull: - def test_toolresult_creation_success(self): - from ui.models.message import ToolResult - tr = ToolResult(tool_call_id="call-1", name="memory_recall", arguments={}, result="success result", success=True) - assert tr.tool_call_id == "call-1" - assert tr.name == "memory_recall" - assert tr.success is True - assert tr.result == "success result" - - def test_toolresult_creation_failure(self): - from ui.models.message import ToolResult - tr = ToolResult(tool_call_id="call-1", name="memory_recall", arguments={}, result="error", success=False) - assert tr.success is False - assert tr.result == "error" - - -class TestLogEntryModelFull: - def test_logentry_creation(self): - from ui.models.log_entry import LogEntry - from datetime import datetime - entry = LogEntry( - timestamp=datetime.now(), - tool_name="memory_recall", - arguments={"query": "test"}, - result="result", - duration=0.5 - ) - assert entry.tool_name == "memory_recall" - assert entry.duration == 0.5 - - def test_logentry_args_summary(self): - from ui.models.log_entry import LogEntry - from datetime import datetime - entry = LogEntry( - timestamp=datetime.now(), - tool_name="memory_recall", - arguments={"query": "test query with many characters"}, - result="result", - duration=0.5 - ) - summary = entry.args_summary - assert isinstance(summary, str) - assert "query" in summary - - def test_logentry_result_summary(self): - from ui.models.log_entry import LogEntry - from datetime import datetime - entry = LogEntry( - timestamp=datetime.now(), - tool_name="memory_recall", - arguments={}, - result="a" * 200, - duration=0.5 - ) - summary = entry.result_summary - assert len(summary) <= 103 - - -class TestAppCSSPathFull: +class TestAppCSSPath: + """测试 App CSS 配置""" + def test_app_has_css_path(self): from ui import GraphMemoryApp assert hasattr(GraphMemoryApp, 'CSS_PATH') assert len(GraphMemoryApp.CSS_PATH) > 0 - def test_css_path_are_paths(self): - from ui import GraphMemoryApp - for path in GraphMemoryApp.CSS_PATH: - assert isinstance(path, Path) - -class TestAppBindingsFull: +class TestAppBindings: + """测试 App 快捷键""" + def test_app_has_bindings(self): from ui import GraphMemoryApp assert hasattr(GraphMemoryApp, 'BINDINGS') assert len(GraphMemoryApp.BINDINGS) > 0 - def test_bindings_have_required_keys(self): - from ui import GraphMemoryApp - for binding in GraphMemoryApp.BINDINGS: - assert hasattr(binding, 'key') - assert hasattr(binding, 'action') - -class TestWidgetImportsFull: +class TestWidgetImports: + """测试组件导入""" + def test_import_left_panel(self): from ui.widgets.left_panel import LeftPanel assert LeftPanel is not None @@ -228,28 +109,14 @@ class TestWidgetImportsFull: from ui.widgets.message_history import MessageHistory assert MessageHistory is not None - def test_import_message_widget(self): - from ui.widgets.message_widget import MessageWidget - assert MessageWidget is not None - def test_import_status_bar(self): from ui.widgets.status_bar import StatusBar assert StatusBar is not None - def test_import_config_section(self): - from ui.widgets.config_section import ConfigSection - assert ConfigSection is not None - def test_import_operation_log(self): - from ui.widgets.operation_log import OperationLog - assert OperationLog is not None - - def test_import_cypher_query_box(self): - from ui.widgets.cypher_query_box import CypherQueryBox - assert CypherQueryBox is not None - - -class TestHandlerImportsFull: +class TestHandlerImports: + """测试处理器导入""" + def test_import_focus_handler(self): from ui.handlers.focus_handler import FocusHandler assert FocusHandler is not None @@ -258,30 +125,22 @@ class TestHandlerImportsFull: from ui.handlers.key_handler import KeyHandler assert KeyHandler is not None - def test_import_message_handler(self): - from ui.handlers.message_handler import MessageHandler - assert MessageHandler is not None - - -class TestServiceImportsFull: - def test_import_config_manager(self): - from ui.services.config_manager import ConfigManager - assert ConfigManager is not None +class TestServiceImports: + """测试服务导入""" + def test_import_config_service(self): from ui.services.config_service import ConfigService assert ConfigService is not None - def test_import_chat_service(self): - from ui.services.chat_service import ChatService - assert ChatService is not None - - def test_import_tool_service(self): - from ui.services.tool_service import ToolService - assert ToolService is not None + def test_import_config_manager(self): + from ui.services.config_manager import ConfigManager + assert ConfigManager is not None -class TestAppInitializationFull: +class TestAppInitialization: + """测试 App 初始化""" + def test_app_without_backend(self): from ui import GraphMemoryApp app = GraphMemoryApp() @@ -305,18 +164,56 @@ class TestAppInitializationFull: os.unlink(db_path) -class TestConfigManagerFull: - def test_config_manager_creation(self): - from ui.services.config_manager import ConfigManager - cm = ConfigManager() - assert cm is not None +class TestUIWithBackendClient: + """测试 UI 与后端通信""" + + def test_app_sends_message_via_backend_client(self): + from ui import GraphMemoryApp + from core import BackendServer, BackendClient + + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + db_path = f.name + try: + server = BackendServer(db_path=db_path) + server.start(api_key="") + + app = GraphMemoryApp(backend_server=server) + client = app._backend_client + + # 测试 get_status + status = client.get_status() + assert status.get("success") is True + + # 测试 update_config + result = client.update_config(api_key="sk-test", base_url="https://api.deepseek.com") + assert result.get("success") is True + + server.shutdown() + finally: + if os.path.exists(db_path): + os.unlink(db_path) - def test_config_manager_config_path(self): - from ui.services.config_manager import ConfigManager - cm = ConfigManager() - assert cm._config_path is not None - - def test_config_manager_exists_false(self): - from ui.services.config_manager import ConfigManager - cm = ConfigManager() - assert cm.exists() is False \ No newline at end of file + def test_ui_only_uses_backend_client(self): + from ui import GraphMemoryApp + from core import BackendServer + + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + db_path = f.name + try: + server = BackendServer(db_path=db_path) + server.start(api_key="") + + app = GraphMemoryApp(backend_server=server) + + # UI 不应该直接访问后端内部 + assert hasattr(app, '_backend_client') + assert app._backend_client is not None + + # 不应该有 _graph, _client 等直接访问 + assert not hasattr(app, '_graph') + assert not hasattr(app, '_client') + + server.shutdown() + finally: + if os.path.exists(db_path): + os.unlink(db_path) \ No newline at end of file diff --git a/trulymem_entry.py b/trulymem_entry.py index c121eba..d856a07 100644 --- a/trulymem_entry.py +++ b/trulymem_entry.py @@ -1,33 +1,43 @@ #!/usr/bin/env python3 import sys +import os from pathlib import Path -if getattr(sys, 'frozen', False): - application_path = Path(sys.executable).parent -else: - application_path = Path(__file__).parent +# 用户配置文件始终放在用户目录 +CONFIG_PATH = Path.home() / ".trulymem" / "config.json" +DB_PATH = Path.home() / ".trulymem" / "graph_memory.db" -sys.path.insert(0, str(application_path)) -import os -os.chdir(application_path) +# 源码运行时使用项目目录,打包后使用用户目录 +if getattr(sys, 'frozen', False): + # 打包版本:创建用户目录 + CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) +else: + # 源码版本:检查项目目录是否有配置(向后兼容) + project_dir = Path(__file__).parent + project_config = project_dir / "config.json" + project_db = project_dir / "graph_memory.db" + + if project_config.exists(): + CONFIG_PATH = project_config + if project_db.exists(): + DB_PATH = project_db + +sys.path.insert(0, str(Path(__file__).parent)) +os.chdir(Path(__file__).parent) from core import BackendServer from ui import GraphMemoryApp -from ui.services.config_service import ConfigService def main(): - # 配置文件保存在应用根目录(exe同级目录或源码根目录) - config_file = application_path / "config.json" + backend_server = BackendServer( + db_path=str(DB_PATH), + use_embedded_db=True, + config_file=str(CONFIG_PATH) + ) + backend_server.start() - # 加载配置 - config_service = ConfigService(config_file=config_file) - config = config_service.get_config() - - backend_server = BackendServer(db_path="graph_memory.db", use_embedded_db=True) - backend_server.start(api_key=config.api_key, base_url=config.base_url) - - app = GraphMemoryApp(backend_server=backend_server, config_service=config_service) + app = GraphMemoryApp(backend_server=backend_server, config_file=str(CONFIG_PATH)) try: app.run() diff --git a/ui/app.py b/ui/app.py index 2cd4a0f..2dbfc9f 100644 --- a/ui/app.py +++ b/ui/app.py @@ -5,7 +5,6 @@ from textual.binding import Binding from core import BackendServer, BackendClient from .models.message import Message -from .services.config_service import ConfigService class GraphMemoryApp(App): @@ -23,11 +22,10 @@ class GraphMemoryApp(App): Binding("f6", "quit", "退出"), ] - def __init__(self, backend_server: BackendServer = None, config_service: ConfigService = None, **kwargs): + def __init__(self, backend_server: BackendServer = None, config_file: str = None, **kwargs): super().__init__(**kwargs) self._backend_server = backend_server self._backend_client = BackendClient(backend_server) if backend_server else None - self._config_service = config_service self._api_configured = False def compose(self) -> ComposeResult: @@ -36,8 +34,16 @@ class GraphMemoryApp(App): from .widgets.status_bar import StatusBar from .models.config import AppConfig - # 获取初始配置 - initial_config = self._config_service.get_config() if self._config_service else AppConfig() + if self._backend_client: + result = self._backend_client.get_config() + config_data = result.get("data", {}) + initial_config = AppConfig( + api_key=config_data.get("api_key", ""), + base_url=config_data.get("base_url", "https://api.deepseek.com"), + model=config_data.get("model", "deepseek-chat") + ) + else: + initial_config = AppConfig() yield LeftPanel() yield RightPanel(config=initial_config) @@ -56,7 +62,8 @@ class GraphMemoryApp(App): return status = self._backend_client.get_status() - self._api_configured = status.get("config", {}).get("api_key", "") != "" + data = status.get("data", {}) + self._api_configured = data.get("config", {}).get("api_key", "") != "" status_bar.set_api_status(self._api_configured) from .widgets.message_history import MessageHistory @@ -118,24 +125,19 @@ class GraphMemoryApp(App): history = self.query_one(MessageHistory) status_bar = self.query_one(StatusBar) - try: - result = await asyncio.get_event_loop().run_in_executor( - None, - lambda: self._backend_client.send_message(user_input) - ) - - # 更新"处理中"消息为实际回复 - if result.get("success"): - content = result.get("content", "(无回复)") - history.update_latest_message(content) - else: - error = result.get("error", "未知错误") - history.update_latest_message(f"❌ 错误: {error}") - - except Exception as e: - history.update_latest_message(f"❌ 异常: {str(e)}") - finally: - status_bar.set_processing(False) + result = await asyncio.get_event_loop().run_in_executor( + None, + lambda: self._backend_client.process_message(user_input) + ) + + if result.get("success"): + content = result.get("content", "(无回复)") + history.update_latest_message(content) + else: + error = result.get("error", "未知错误") + history.update_latest_message(f"❌ 错误: {error}") + + status_bar.set_processing(False) def on_config_section_config_changed(self, event) -> None: """处理配置变更事件""" @@ -151,24 +153,41 @@ class GraphMemoryApp(App): async def _update_config_async(self, config) -> None: """异步更新配置""" from .widgets.status_bar import StatusBar + from .widgets.config_section import ConfigSection + status_bar = self.query_one(StatusBar) api_key = config.api_key base_url = config.base_url + model = getattr(config, 'model', 'deepseek-chat') try: result = await asyncio.get_event_loop().run_in_executor( None, - lambda: self._backend_client.update_config(api_key=api_key, base_url=base_url) + lambda: self._backend_client.update_config(api_key=api_key, base_url=base_url, model=model) ) if result.get("success"): self._api_configured = bool(api_key) status_bar.set_api_status(self._api_configured) - # 保存配置到文件(使用完整的config对象,保留model字段) - if self._config_service: - self._config_service.set_config(config) + # 从后端重新获取配置并刷新 UI + config_result = await asyncio.get_event_loop().run_in_executor( + None, + lambda: self._backend_client.get_config() + ) + config_data = config_result.get("data", {}) + + # 刷新输入框 + try: + config_section = self.query_one(ConfigSection) + config_section.set_config(AppConfig( + api_key=config_data.get("api_key", ""), + base_url=config_data.get("base_url", "https://api.deepseek.com"), + model=config_data.get("model", "deepseek-chat") + )) + except Exception: + pass self.notify("✅ 配置已保存并生效", title="配置成功", severity="information") else: diff --git a/ui/handlers/message_handler.py b/ui/handlers/message_handler.py deleted file mode 100644 index eeec026..0000000 --- a/ui/handlers/message_handler.py +++ /dev/null @@ -1,98 +0,0 @@ -"""消息处理器""" - -from datetime import datetime -from typing import TYPE_CHECKING -from ..models.message import Message, ToolCall, ToolResult - -if TYPE_CHECKING: - from ..services.chat_service import ChatService - from ..widgets.message_history import MessageHistory - from ..widgets.operation_log import OperationLog - - -class MessageHandler: - """消息处理器""" - - def __init__( - self, - chat_service: "ChatService", - message_history: "MessageHistory", - operation_log: "OperationLog" - ): - self._chat_service = chat_service - self._message_history = message_history - self._operation_log = operation_log - - async def handle_user_message(self, content: str) -> None: - """处理用户消息""" - # 创建用户消息 - user_message = Message( - role="user", - content=content, - timestamp=datetime.now() - ) - - # 添加到历史 - self._message_history.add_message(user_message) - - # 发送到聊天服务 - await self._process_response(content) - - async def _process_response(self, user_input: str) -> None: - """处理响应""" - streaming_message = None - - async for event in self._chat_service.send_message(user_input): - if event["type"] == "user_message": - # 用户消息已处理 - pass - - elif event["type"] == "content_delta": - # 流式内容更新 - if streaming_message is None: - # 创建流式消息 - streaming_message = Message( - role="assistant", - content="", - timestamp=datetime.now() - ) - self._message_history.add_message(streaming_message) - - # 更新消息内容 - self._message_history.update_latest_message(event["content"]) - - elif event["type"] == "assistant_message": - # 模型消息完成 - if streaming_message: - # 更新最终消息(包含工具调用信息) - streaming_message.content = event["content"] - streaming_message.tool_calls = event.get("tool_calls") - streaming_message.tool_results = event.get("tool_results") - else: - # 如果没有流式消息,直接添加 - message = Message( - role="assistant", - content=event["content"], - timestamp=datetime.now(), - tool_calls=event.get("tool_calls"), - tool_results=event.get("tool_results") - ) - self._message_history.add_message(message) - - elif event["type"] == "tool_call": - # 工具调用开始 - pass - - elif event["type"] == "tool_result": - # 工具执行结果 - log_entry = event["log_entry"] - self._operation_log.add_log(log_entry) - - elif event["type"] == "error": - # 错误处理 - error_message = Message( - role="assistant", - content=f"错误: {event['error']}", - timestamp=datetime.now() - ) - self._message_history.add_message(error_message) diff --git a/ui/services/chat_service.py b/ui/services/chat_service.py deleted file mode 100644 index fc96f43..0000000 --- a/ui/services/chat_service.py +++ /dev/null @@ -1,226 +0,0 @@ -"""聊天服务""" - -import asyncio -import json -from datetime import datetime -from typing import AsyncIterator, TYPE_CHECKING, List, Dict, Any -from ..core.imports import GraphMemoryClient -from ..models.message import ToolCall, ToolResult -from .tool_service import ToolService - -if TYPE_CHECKING: - from ..core.imports import Neo4jGraph - - -class ChatService: - """聊天业务服务""" - - def __init__( - self, - graph: "Neo4jGraph", - client: GraphMemoryClient, - tool_service: ToolService - ): - self._graph = graph - self._client = client - self._tool_service = tool_service - self._messages: List[Dict[str, Any]] = [] - - async def send_message(self, user_input: str) -> AsyncIterator[dict]: - """发送消息并流式返回事件""" - # 1. 发送用户消息事件 - yield { - "type": "user_message", - "content": user_input - } - - try: - # 2. 第一次API调用 - accumulated_content = "" - tool_calls_data = [] - - # 流式处理响应 - async for chunk in self._call_api_stream_async(user_input): - if chunk.get("content_delta"): - accumulated_content += chunk["content_delta"] - yield { - "type": "content_delta", - "content": accumulated_content - } - - if chunk.get("tool_calls"): - tool_calls_data = chunk["tool_calls"] - - # 3. 如果有工具调用,执行并继续调用API - tool_calls = None - tool_results = None - - if tool_calls_data: - tool_calls = [] - tool_results = [] - - # 执行所有工具 - for tool_call_data in tool_calls_data: - tool_call = ToolCall( - id=tool_call_data["id"], - name=tool_call_data["function"]["name"], - arguments=tool_call_data["function"]["arguments"] - ) - tool_calls.append(tool_call) - - yield { - "type": "tool_call", - "tool_call": tool_call - } - - result = await self._tool_service.execute(tool_call) - tool_results.append(result) - - log_entry = ToolService._create_log_entry(tool_call, result) - yield { - "type": "tool_result", - "tool_result": result, - "log_entry": log_entry - } - - # 构建工具结果消息 - tool_messages = [] - for tc, tr in zip(tool_calls, tool_results): - tool_messages.append({ - "role": "tool", - "tool_call_id": tc.id, - "content": tr.content - }) - - # 构建assistant消息(包含tool_calls) - assistant_message = { - "role": "assistant", - "content": accumulated_content, - "tool_calls": [ - { - "id": tc.id, - "type": "function", - "function": { - "name": tc.name, - "arguments": tc.arguments - } - } for tc in tool_calls - ] - } - - # 第二次API调用,传入工具结果 - final_content = "" - async for chunk in self._call_api_stream_with_tools( - user_input, - assistant_message, - tool_messages - ): - if chunk.get("content_delta"): - final_content += chunk["content_delta"] - yield { - "type": "content_delta", - "content": final_content - } - - accumulated_content = final_content - - # 4. 返回最终回复 - yield { - "type": "assistant_message", - "content": accumulated_content, - "tool_calls": tool_calls, - "tool_results": tool_results - } - - except Exception as e: - yield { - "type": "error", - "error": str(e) - } - - async def _call_api_stream_async(self, message: str) -> AsyncIterator[dict]: - """异步流式调用 API""" - loop = asyncio.get_event_loop() - - def process_stream(): - stream = self._client.send_message_stream(message) - tool_calls_accumulated = [] - - for chunk in stream: - delta = chunk.choices[0].delta - - if delta.content: - yield {"content_delta": delta.content} - - if delta.tool_calls: - for tc in delta.tool_calls: - if tc.index >= len(tool_calls_accumulated): - tool_calls_accumulated.append({ - "id": tc.id, - "type": "function", - "function": { - "name": "", - "arguments": "" - } - }) - - if tc.function: - if tc.function.name: - tool_calls_accumulated[tc.index]["function"]["name"] = tc.function.name - if tc.function.arguments: - tool_calls_accumulated[tc.index]["function"]["arguments"] += tc.function.arguments - - if tool_calls_accumulated: - yield {"tool_calls": tool_calls_accumulated} - - for result in await loop.run_in_executor(None, lambda: list(process_stream())): - yield result - - async def _call_api_stream_with_tools( - self, - user_input: str, - assistant_message: dict, - tool_messages: list - ) -> AsyncIterator[dict]: - """带工具结果的流式调用""" - loop = asyncio.get_event_loop() - - def process_stream(): - # 构建完整的消息列表 - messages = [ - {"role": "system", "content": self._client.system_prompt}, - {"role": "user", "content": user_input}, - assistant_message - ] - messages.extend(tool_messages) - - # 调用API - response = self._client.client.chat.completions.create( - model="deepseek-chat", - messages=messages, - tools=self._client.tools, - tool_choice="auto", - stream=True - ) - - for chunk in response: - delta = chunk.choices[0].delta - if delta.content: - yield {"content_delta": delta.content} - - # 处理可能的工具调用 - if delta.tool_calls: - # 如果还有工具调用,说明AI想继续调用工具 - # 但我们限制只调用一次,所以忽略 - pass - - for result in await loop.run_in_executor(None, lambda: list(process_stream())): - yield result - - def clear_history(self) -> None: - """清空消息历史""" - self._messages.clear() - - def get_history(self) -> list[dict]: - """获取消息历史""" - return self._messages.copy() diff --git a/ui/services/tool_service.py b/ui/services/tool_service.py deleted file mode 100644 index fb37659..0000000 --- a/ui/services/tool_service.py +++ /dev/null @@ -1,88 +0,0 @@ -"""工具服务""" - -import asyncio -import time -from datetime import datetime -from typing import Callable, TYPE_CHECKING -from ..core.imports import execute_tool -from ..models.log_entry import LogEntry -from ..models.message import ToolCall, ToolResult - -if TYPE_CHECKING: - from ..core.imports import Neo4jGraph - - -class ToolService: - """工具执行服务""" - - def __init__( - self, - graph: "Neo4jGraph", - log_callback: Callable[[LogEntry], None] | None = None - ): - self._graph = graph - self._log_callback = log_callback - - async def execute(self, tool_call: ToolCall) -> ToolResult: - """异步执行工具""" - start_time = time.time() - - try: - # 在线程池中执行同步工具 - loop = asyncio.get_event_loop() - result = await loop.run_in_executor( - None, - lambda: execute_tool(self._graph, tool_call.name, tool_call.arguments) - ) - - duration = time.time() - start_time - - # 创建日志条目 - log_entry = LogEntry( - timestamp=datetime.now(), - tool_name=tool_call.name, - arguments=tool_call.arguments, - result=result, - duration=duration - ) - - # 回调日志 - if self._log_callback: - self._log_callback(log_entry) - - # 返回结果 - return ToolResult( - tool_call_id=tool_call.id, - name=tool_call.name, - arguments=tool_call.arguments, - result=result, - success=not result.startswith("工具执行错误") - ) - - except Exception as e: - duration = time.time() - start_time - error_msg = f"工具执行异常: {str(e)}" - - # 创建错误日志 - log_entry = LogEntry( - timestamp=datetime.now(), - tool_name=tool_call.name, - arguments=tool_call.arguments, - result=error_msg, - duration=duration - ) - - if self._log_callback: - self._log_callback(log_entry) - - return ToolResult( - tool_call_id=tool_call.id, - name=tool_call.name, - arguments=tool_call.arguments, - result=error_msg, - success=False - ) - - def set_log_callback(self, callback: Callable[[LogEntry], None]) -> None: - """设置日志回调""" - self._log_callback = callback