vscode_api_server.py•12.8 kB
#!/usr/bin/env python3
"""
VS Code扩展专用API服务器
运行在localhost:7788端口，提供/mem/search接口
基于SimpleMemoryOps实现，专为VS Code扩展设计
"""
import json
import time
from datetime import datetime
from typing import Dict, Any, List, Optional
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
import threading
import sys
from pathlib import Path
# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
try:
from simple_memory_ops_sdk import SimpleMemoryOps
except ImportError:
print("Error: Cannot import SimpleMemoryOps", file=sys.stderr)
sys.exit(1)
class VSCodeAPIHandler(BaseHTTPRequestHandler):
"""VS Code扩展API请求处理器"""
def __init__(self, *args, memory_ops=None, **kwargs):
self.memory_ops = memory_ops
super().__init__(*args, **kwargs)
def do_GET(self):
"""处理GET请求"""
try:
parsed_path = urlparse(self.path)
path = parsed_path.path
if path == "/":
self._send_redirect("/docs")
elif path == "/docs":
self._send_api_docs()
elif path == "/health":
self._handle_health_check()
elif path == "/mem/search":
# 支持GET方式的简单搜索
query_params = parse_qs(parsed_path.query)
query = query_params.get('q', [''])[0]
if query:
self._handle_search_memories({"query": query, "top_k": 5})
else:
self._send_error(400, "Bad Request", "缺少查询参数 'q'")
else:
self._send_error(404, "Not Found", "端点不存在")
except Exception as e:
self._send_error(500, "Internal Server Error", str(e))
def do_POST(self):
"""处理POST请求"""
try:
parsed_path = urlparse(self.path)
path = parsed_path.path
# 读取请求体
content_length = int(self.headers.get('Content-Length', 0))
post_data = self.rfile.read(content_length)
try:
request_data = json.loads(post_data.decode('utf-8')) if post_data else {}
except json.JSONDecodeError:
self._send_error(400, "Bad Request", "无效的JSON格式")
return
if path == "/mem/search":
self._handle_search_memories(request_data)
elif path == "/mem/add":
self._handle_add_memory(request_data)
else:
self._send_error(404, "Not Found", "端点不存在")
except Exception as e:
self._send_error(500, "Internal Server Error", str(e))
def do_OPTIONS(self):
"""处理OPTIONS请求（CORS预检）"""
self._send_cors_headers()
self.send_response(200)
self.end_headers()
def _send_cors_headers(self):
"""发送CORS头"""
self.send_header('Access-Control-Allow-Origin', '*')
self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
self.send_header('Access-Control-Allow-Headers', 'Content-Type, Accept')
def _send_json_response(self, data: Dict[str, Any], status_code: int = 200):
"""发送JSON响应"""
response_json = json.dumps(data, ensure_ascii=False, indent=2)
self.send_response(status_code)
self.send_header('Content-Type', 'application/json; charset=utf-8')
self.send_header('Content-Length', str(len(response_json.encode('utf-8'))))
self._send_cors_headers()
self.end_headers()
self.wfile.write(response_json.encode('utf-8'))
def _send_error(self, status_code: int, error_type: str, message: str):
"""发送错误响应"""
error_response = {
"success": False,
"error_code": error_type.upper().replace(" ", "_"),
"error_message": message,
"timestamp": datetime.now().isoformat()
}
self._send_json_response(error_response, status_code)
def _send_redirect(self, location: str):
"""发送重定向响应"""
self.send_response(307)
self.send_header('Location', location)
self.end_headers()
def _send_api_docs(self):
"""发送API文档页面"""
docs_html = """
<!DOCTYPE html>
<html>
<head>
<title>VS Code MemOS API文档</title>
<meta charset="utf-8">
<style>
body { font-family: Arial, sans-serif; margin: 40px; }
.endpoint { margin: 20px 0; padding: 15px; border: 1px solid #ddd; border-radius: 5px; }
.method { font-weight: bold; color: #007bff; }
.path { font-family: monospace; background: #f8f9fa; padding: 2px 5px; }
pre { background: #f8f9fa; padding: 10px; border-radius: 3px; overflow-x: auto; }
</style>
</head>
<body>
<h1>VS Code MemOS API文档</h1>
<p>专为VS Code扩展设计的MemOS API接口</p>
<h2>搜索接口</h2>
<div class="endpoint">
<div><span class="method">GET</span> <span class="path">/mem/search?q=查询内容</span></div>
<p>简单搜索 - 通过URL参数进行记忆搜索</p>
</div>
<div class="endpoint">
<div><span class="method">POST</span> <span class="path">/mem/search</span></div>
<p>高级搜索 - 通过JSON请求体进行记忆搜索</p>
<pre>
{
"query": "搜索查询",
"memory_type": "code_snippet_mem",
"limit": 5
}
</pre>
</div>
<div class="endpoint">
<div><span class="method">POST</span> <span class="path">/mem/add</span></div>
<p>添加记忆 - 添加新的记忆到系统中</p>
<pre>
{
"text": "记忆内容",
"memory_type": "code_snippet_mem",
"tags": ["标签1", "标签2"],
"metadata": {}
}
</pre>
</div>
<h2>使用示例</h2>
<pre>
# 简单搜索
curl "http://localhost:7788/mem/search?q=Python"
# 高级搜索
curl -X POST http://localhost:7788/mem/search \\
-H "Content-Type: application/json" \\
-d '{"query": "Python", "memory_type": "code_snippet_mem", "limit": 5}'
# 添加记忆
curl -X POST http://localhost:7788/mem/add \\
-H "Content-Type: application/json" \\
-d '{"text": "def hello(): print(\\"Hello World\\")", "memory_type": "code_snippet_mem", "tags": ["Python", "function"]}'
</pre>
</body>
</html>
"""
self.send_response(200)
self.send_header('Content-Type', 'text/html; charset=utf-8')
self.send_header('Content-Length', str(len(docs_html.encode('utf-8'))))
self.end_headers()
self.wfile.write(docs_html.encode('utf-8'))
def _handle_health_check(self):
"""处理健康检查"""
try:
# 测试Memory Ops连接
stats = self.memory_ops.stats()
health_data = {
"api_status": "healthy",
"memory_ops_status": "ok",
"memory_types": list(self.memory_ops.list_memory_types().keys()),
"total_memories": sum(info.get("count", 0) for info in stats.values() if isinstance(info, dict))
}
response = {
"success": True,
"message": "VS Code API服务健康状态正常",
"data": health_data,
"timestamp": datetime.now().isoformat()
}
self._send_json_response(response)
except Exception as e:
self._send_error(503, "Service Unavailable", f"服务不可用: {str(e)}")
def _handle_search_memories(self, request_data: Dict[str, Any]):
"""处理搜索记忆"""
try:
# 验证必需字段
if "query" not in request_data:
self._send_error(400, "Bad Request", "缺少必需字段: query")
return
query = request_data["query"]
memory_type = request_data.get("memory_type")
limit = request_data.get("limit", 5)
# 记录操作开始时间
start_time = time.time()
# 执行搜索
results = self.memory_ops.query(query, memory_type=memory_type, limit=limit)
# 格式化搜索结果
formatted_results = []
for result in results:
formatted_result = {
"id": result.get("id", ""),
"content": result.get("memory", ""),
"memory_type": result.get("memory_type", ""),
"tags": result.get("metadata", {}).get("tags", []),
"metadata": result.get("metadata", {}),
"created_at": result.get("metadata", {}).get("timestamp", "")
}
formatted_results.append(formatted_result)
# 记录操作时间
duration = time.time() - start_time
response = {
"success": True,
"message": f"搜索完成，找到 {len(formatted_results)} 条相关记忆",
"data": {
"query": query,
"memory_type": memory_type,
"results": formatted_results,
"total_count": len(formatted_results),
"processing_time": round(duration, 3)
},
"timestamp": datetime.now().isoformat()
}
self._send_json_response(response)
except Exception as e:
self._send_error(500, "Internal Server Error", f"搜索记忆失败: {str(e)}")
def _handle_add_memory(self, request_data: Dict[str, Any]):
"""处理添加记忆"""
try:
# 验证必需字段
if "text" not in request_data:
self._send_error(400, "Bad Request", "缺少必需字段: text")
return
text = request_data["text"]
memory_type = request_data.get("memory_type", "code_snippet_mem")
tags = request_data.get("tags", [])
metadata = request_data.get("metadata", {})
# 记录操作开始时间
start_time = time.time()
# 添加记忆
success = self.memory_ops.add(
text=text,
memory_type=memory_type,
tags=tags,
metadata=metadata
)
if not success:
self._send_error(500, "Internal Server Error", "记忆创建失败")
return
# 记录操作时间
duration = time.time() - start_time
response = {
"success": True,
"message": "记忆创建成功",
"data": {
"content_preview": text[:50] + "..." if len(text) > 50 else text,
"memory_type": memory_type,
"tags": tags,
"processing_time": round(duration, 3)
},
"timestamp": datetime.now().isoformat()
}
self._send_json_response(response)
except Exception as e:
self._send_error(500, "Internal Server Error", f"创建记忆失败: {str(e)}")
def create_api_handler(memory_ops):
"""创建API处理器的工厂函数"""
def handler(*args, **kwargs):
return VSCodeAPIHandler(*args, memory_ops=memory_ops, **kwargs)
return handler
def start_vscode_api_server(host: str = "localhost", port: int = 7788):
"""启动VS Code API服务器"""
print("🚀 启动VS Code MemOS API服务器...")
print(f"📖 API文档: http://{host}:{port}/docs")
print(f"🔍 健康检查: http://{host}:{port}/health")
print(f"🔎 搜索接口: http://{host}:{port}/mem/search")
# 初始化Memory Ops
memory_ops = SimpleMemoryOps(verbose=True)
# 创建HTTP服务器
handler_class = create_api_handler(memory_ops)
server = HTTPServer((host, port), handler_class)
print(f"✅ 服务器已启动在 http://{host}:{port}")
print("按 Ctrl+C 停止服务器")
try:
server.serve_forever()
except KeyboardInterrupt:
print("\n⏹️ 服务器已停止")
memory_ops.close()
server.shutdown()
if __name__ == "__main__":
start_vscode_api_server()