a820afada9
- Add GET_CONFIG, SET_CONFIG, GET_HISTORY, SAVE_HISTORY to BackendServer - Update BackendClient with get_config, get_message_history, save_message_history - All API key and message history now stored in core layer
426 lines
15 KiB
Python
426 lines
15 KiB
Python
import threading
|
|
import queue
|
|
import time
|
|
import json
|
|
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):
|
|
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"
|
|
|
|
|
|
@dataclass
|
|
class BackendRequest:
|
|
request_id: str
|
|
message_type: MessageType
|
|
payload: Dict[str, Any]
|
|
response_queue: queue.Queue = field(default=None)
|
|
|
|
|
|
@dataclass
|
|
class BackendResponse:
|
|
request_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):
|
|
self._db_path = db_path
|
|
self._use_embedded_db = use_embedded_db
|
|
|
|
self._graph = None
|
|
self._client = None
|
|
self._tool_limiter = ToolLimiter()
|
|
|
|
self._request_queue: queue.Queue[BackendRequest] = queue.Queue()
|
|
self._running = False
|
|
self._thread: Optional[threading.Thread] = None
|
|
|
|
self._lock = threading.Lock()
|
|
|
|
def start(self, api_key: str = "", base_url: str = "https://api.deepseek.com") -> None:
|
|
if self._running:
|
|
return
|
|
|
|
self._init_graph()
|
|
|
|
if api_key:
|
|
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:
|
|
request = self._request_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"}
|
|
))
|
|
|
|
def _handle_process_message(self, request: BackendRequest) -> None:
|
|
try:
|
|
user_input = request.payload.get("user_input", "")
|
|
|
|
if not self._client:
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
success=False,
|
|
error="API Key 未配置"
|
|
))
|
|
return
|
|
|
|
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
|
|
}
|
|
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
|
|
})
|
|
|
|
tool_result_msg = {
|
|
"role": "tool",
|
|
"tool_call_id": tool_call.id,
|
|
"content": result
|
|
}
|
|
current_tool_results.append(tool_result_msg)
|
|
|
|
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()}"
|
|
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
success=True,
|
|
data={
|
|
"content": content,
|
|
"tool_calls": tool_calls,
|
|
"rejected_tools": rejected_tools
|
|
}
|
|
))
|
|
|
|
except Exception as e:
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
success=False,
|
|
error=str(e)
|
|
))
|
|
|
|
def _handle_execute_tool(self, request: BackendRequest) -> None:
|
|
try:
|
|
tool_name = request.payload.get("tool_name")
|
|
arguments = request.payload.get("arguments", {})
|
|
|
|
allowed, reason = self._tool_limiter.can_call(tool_name, arguments)
|
|
if not allowed:
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
success=False,
|
|
error=f"工具调用被拒绝: {reason}"
|
|
))
|
|
return
|
|
|
|
self._tool_limiter.record_call(tool_name, arguments)
|
|
|
|
result = execute_tool(self._graph, tool_name, arguments)
|
|
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
success=True,
|
|
data={"result": result}
|
|
))
|
|
|
|
except Exception as e:
|
|
self._send_response(request, BackendResponse(
|
|
request_id=request.request_id,
|
|
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_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_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()
|
|
|
|
request = BackendRequest(
|
|
request_id=request_id,
|
|
message_type=MessageType.PROCESS_MESSAGE,
|
|
payload={"user_input": user_input},
|
|
response_queue=response_queue
|
|
)
|
|
|
|
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("请求超时")
|
|
|
|
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
|
|
)
|
|
|
|
self._request_queue.put(request)
|
|
|
|
try:
|
|
response = response_queue.get(timeout=timeout)
|
|
if not response.success:
|
|
raise Exception(response.error)
|
|
return response.data["result"]
|
|
except queue.Empty:
|
|
raise TimeoutError("工具执行超时")
|
|
|
|
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
|
|
)
|
|
|
|
self._request_queue.put(request)
|
|
|
|
if self._thread:
|
|
self._thread.join(timeout=2.0)
|
|
|
|
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", []) |