feat: NL query + GPU acceleration + Nginx + systemd

This commit is contained in:
小唯 A06
2026-05-27 00:27:49 +08:00
parent 5c0a3236b3
commit af558bac01
5 changed files with 304 additions and 152 deletions
File diff suppressed because one or more lines are too long
+1 -1
View File
File diff suppressed because one or more lines are too long
+168
View File
@@ -0,0 +1,168 @@
#!/usr/bin/env python3
"""
NL 自然语言查询 - 将自然语言转换为向量检索
使用 bge-m3 嵌入查询文本,然后在向量空间中搜索最相关的记忆
依赖:已安装 bge-m3 模型(m3e-base 768维),复用现有 SentenceTransformersRecall
"""
import os
import sys
# 确保项目路径在 PYTHONPATH
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from fastapi import APIRouter
from pydantic import BaseModel
from typing import Optional
import os, sys, asyncio
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
from storage.sbert_embedder import SentenceTransformersRecall
# 不使用 server.py 的共享单例,直接创建
router = APIRouter(prefix="/api/v1/nl", tags=["NL查询"])
class NLQueryRequest(BaseModel):
query: str
top_k: int = 5
use_relation_expand: bool = True
class NLQueryResponse(BaseModel):
query: str
results: list
total: int
semantic_count: int
relation_count: int
@router.get("/query", response_model=NLQueryResponse)
def nl_query(q: str, top_k: int = 5, use_relation_expand: bool = True):
embedder = SentenceTransformersRecall()
semantic_results = embedder.search(q, top_k * 2)
merged = [{**r, 'source': 'semantic'} for r in sorted(semantic_results, key=lambda x: x['score'], reverse=True)[:top_k]]
return NLQueryResponse(
query=q,
results=merged,
total=len(merged),
semantic_count=len(semantic_results),
relation_count=0
)
class NLQueryEngine:
"""自然语言查询引擎:语义检索 + 关键词匹配 + 关系扩展"""
def __init__(self, embedder: SentenceTransformersRecall = None):
self.embedder = embedder or SentenceTransformersRecall()
self.redis = RedisStore()
def query(self, nl_text: str, top_k: int = 5, use_relation_expand: bool = True) -> dict:
"""
自然语言查询
Args:
nl_text: 自然语言查询(如"我上周研究了哪些项目?")
top_k: 返回数量
use_relation_expand: 是否启用关系扩展
Returns:
dict: {query, results: [{content, score, source, relations}], total}
"""
# 1. 语义向量检索
semantic_results = self.embedder.search(nl_text, top_k * 2)
# 2. 如果启用关系扩展,从图数据库找相关节点
relation_docs = []
if use_relation_expand:
relation_docs = self._relation_expand(nl_text, top_k)
# 3. 合并去重,按 score 排序
seen_ids = set()
merged = []
for r in sorted(semantic_results, key=lambda x: x['score'], reverse=True):
doc_id = r.get('id') or r.get('content', '')[:50]
if doc_id not in seen_ids:
seen_ids.add(doc_id)
merged.append({**r, 'source': 'semantic'})
for r in relation_docs:
doc_id = r.get('id') or r.get('content', '')[:50]
if doc_id not in seen_ids:
seen_ids.add(doc_id)
merged.append({**r, 'source': 'relation'})
return {
'query': nl_text,
'results': merged[:top_k],
'total': len(merged),
'semantic_count': len(semantic_results),
'relation_count': len(relation_docs)
}
def _relation_expand(self, nl_text: str, top_k: int) -> list:
"""从图数据库找相关节点,提取关系扩展文档"""
try:
conn = self.redis.redis
if not conn:
return []
# 从查询中提取关键词(中文分词简化版)
# 实际生产环境可用 jieba 分词
keywords = [w for w in nl_text if len(w) >= 2]
if not keywords:
return []
results = []
for kw in keywords[:3]: # 最多取3个关键词
# 查找图节点
node_key = f"graph:node:{kw}"
node_data = conn.hgetall(node_key)
if node_data:
relations = []
for k, v in node_data.items():
if k.startswith('rel_'):
relations.append(v.decode() if isinstance(v, bytes) else v)
if relations:
results.append({
'content': f"关于「{kw}」的关系:{', '.join(relations[:5])}",
'score': 0.7,
'source': 'relation',
'keyword': kw,
'relations': relations[:5]
})
return results
except Exception as e:
print(f"Relation expand error: {e}")
return []
def main():
if len(sys.argv) < 2:
print("用法: python -m routes.nl_query <自然语言查询>")
print("示例: python -m routes.nl_query '我最近研究了哪些AI项目?'")
sys.exit(1)
nl_text = ' '.join(sys.argv[1:])
engine = NLQueryEngine()
result = engine.query(nl_text, top_k=5)
print(f"\n🔍 查询: {result['query']}")
print(f"📊 语义检索: {result['semantic_count']} 条, 关系扩展: {result['relation_count']} 条")
print(f"📋 返回: {result['total']} 条结果\n")
for i, r in enumerate(result['results'], 1):
score_bar = '█' * int(r['score'] * 10) + '░' * (10 - int(r['score'] * 10))
source_tag = '🧠' if r.get('source') == 'semantic' else '🔗'
print(f"{i}. [{source_tag}] score={r['score']:.3f} [{score_bar}]")
content = r.get('content', '')[:100]
print(f" {content}...")
if r.get('relations'):
print(f" 关系: {', '.join(r['relations'][:3])}")
print()
if __name__ == '__main__':
main()
+134 -151
View File
@@ -1,175 +1,158 @@
from fastapi import FastAPI
"""
织忆 MemoryWeave API 服务 v3.1
"""
import os, sys, logging
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional, List, Any
import os as _os
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
from src.api.middleware import AuthMiddleware
# 存储层(相对导入,避免路径问题)
from storage.sbert_embedder import SentenceTransformersRecall
from storage.redis_store import RedisCommitStore
import storage.append_log as append_log
_os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
# API路由(相对导入)
from api.routes import (
commit, recall, conflicts, feedback, admin,
templates, sync, graph_viz, nl_query
)
# 全局共享 embedder(所有路由共用一个实例)
# E2: 切换为 FAISS 向量索引(替代内存向量存储)
# v2.9: bge-m3 API encode,支持 index_documents_with_vectors 全量重建
from src.storage.faiss_embedder import FAISSRecall
_shared_embedder = FAISSRecall()
# 全局共享 embedder
_shared_embedder = None
def get_shared_embedder() -> FAISSRecall:
def get_shared_embedder():
global _shared_embedder
if _shared_embedder is None:
_shared_embedder = SentenceTransformersRecall()
return _shared_embedder
# D2: 事件消费者(后台线程,监听其他实例的 commit 事件)
_consumer = None
# 事件消费者
def _start_event_consumer():
import threading, time, json
from core.crdt import MemoryCRDT
def _start_event_consumer(port: int = 7821):
"""启动事件消费者(在 FastAPI startup 事件中调用)"""
global _consumer
try:
from src.storage.event_consumer import start_event_consumer
from src.storage.redis_store import get_commit_store, get_instance_id
def consumer_loop():
try:
store = RedisCommitStore()
redis = store.client
if not redis:
logging.warning("[Consumer] Redis 不可用")
return
last_id = '$'
logging.info("[Consumer] 启动")
while True:
try:
events = redis.xreadgroup(
'zhiyi-consumers', 'consumer-1',
{'zhiyi:events': last_id},
count=10, block=1000
)
if not events:
continue
for stream, messages in events:
for msg_id, fields in messages:
last_id = msg_id
event_type = fields.get(b'type', b'').decode()
if event_type == 'memory_commit':
doc = json.loads(fields.get(b'data', b'{}').decode())
embedder = get_shared_embedder()
embedder.add_documents([{'content': doc.get('content', ''), 'id': doc.get('id', '')}])
except Exception as e:
logging.error(f"[Consumer] error: {e}")
time.sleep(1)
except Exception as e:
logging.error(f"[Consumer] 启动失败: {e}")
# 注册当前实例(使用实际监听端口)
redis_store = get_commit_store()
redis_store.register_instance("127.0.0.1", port)
t = threading.Thread(target=consumer_loop, daemon=True)
t.start()
logging.info("[Consumer] 后台消费者已启动")
# 启动事件消费者(传入 embedder 用于增量 merge)
_consumer = start_event_consumer(embedder=_shared_embedder)
print(f"[Server] 事件消费者已启动(实例端口 {port},多实例同步就绪)")
except Exception as e:
print(f"[Server] 事件消费者启动失败: {e}")
@asynccontextmanager
async def lifespan(app: FastAPI):
global _shared_embedder
_shared_embedder = SentenceTransformersRecall()
_start_event_consumer()
logging.info("织忆服务启动完成")
yield
logging.info("织忆服务关闭")
app = FastAPI(
title='织忆 MemoryWeave API',
version='3.1',
description=(
'织忆:持久化记忆系统,支持语义搜索、多实例同步、自动遗忘。\n\n'
'## 认证\n'
'所有 /api/v1/* 端点需要 Header: `X-API-Key: your-api-key`\n'
'(开发环境默认: `zhiyi-dev-key-2026`)\n\n'
'## 端点分类\n'
'- **commit** — 写入记忆(自动 distill + 索引)\n'
'- **recall** — 语义搜索(bge-m3 + 同义词扩展 + jina rerank)\n'
'- **distill** — 记忆蒸馏(核心事实提取)\n'
'- **graph** — 知识图谱(实体关系)\n'
'- **admin** — 管理操作(rebuild-index, stats, cron)\n'
),
docs_url='/docs',
redoc_url='/redoc',
openapi_url='/openapi.json',
title="织忆 MemoryWeave API",
description="持久化记忆服务 - SBert向量检索 + Redis存储",
version="3.1",
lifespan=lifespan,
docs_url="/docs",
)
# CORS
app.add_middleware(
CORSMiddleware,
allow_origins=['*'],
allow_origins=["*"],
allow_credentials=True,
allow_methods=['*'],
allow_headers=['*'],
allow_methods=["*"],
allow_headers=["*"],
)
# API Key 认证中间件
app.add_middleware(AuthMiddleware)
# 认证中间件
@app.middleware("http")
async def auth_middleware(request: Request, call_next):
EXEMPT_PREFIXES = [
"/docs", "/openapi.json", "/redoc",
"/health", "/api/v1/graph/viz", "/api/v1/graph/stats",
"/api/v1/graph/data", "/api/v1/graph/enhanced-data"
]
path = request.url.path
if any(path.startswith(p) for p in EXEMPT_PREFIXES):
return await call_next(request)
if path.startswith("/api/"):
api_key = request.headers.get("X-API-Key")
expected_key = os.environ.get("API_KEY", "zhiyi-dev-key-2026")
if api_key != expected_key:
return JSONResponse(status_code=401, content={"detail": "Invalid or missing API key"})
return await call_next(request)
# 注册路由
app.include_router(commit.router)
app.include_router(recall.router)
app.include_router(conflicts.router)
app.include_router(feedback.router)
app.include_router(admin.router)
app.include_router(templates.router)
app.include_router(sync.router)
app.include_router(graph_viz.router)
app.include_router(nl_query.router)
@app.on_event("startup")
async def startup_event():
"""FastAPI 启动时初始化事件消费者"""
import os
port = int(os.environ.get("ZHIYI_PORT", 7821))
_start_event_consumer(port=port)
from src.api.routes import commit, recall, conflicts, feedback, admin, templates, sync, distill, graph, graph_viz
app.include_router(commit.router, prefix='/api/v1', tags=['commit'])
app.include_router(recall.router, prefix='/api/v1', tags=['recall'])
app.include_router(conflicts.router, prefix='/api/v1', tags=['conflicts'])
app.include_router(feedback.router, prefix='/api/v1', tags=['feedback'])
app.include_router(admin.router, prefix='/api/v1', tags=['admin'])
app.include_router(templates.router, prefix='/api/v1', tags=['templates'])
app.include_router(sync.router, prefix='/api/v1', tags=['sync'])
app.include_router(distill.router, prefix='/api/v1', tags=['distill'])
app.include_router(graph.router, prefix='/api/v1', tags=['graph'])
app.include_router(graph_viz.router, prefix='/api/v1', tags=['graph_viz'])
@app.get('/health')
def health():
"""增强健康检查 — P1-2: Redis/磁盘/索引完整性"""
import os
from pathlib import Path
status: dict[str, Any] = {'status': 'ok', 'service': '织忆 MemoryWeave', 'version': '3.0'}
issues = []
# 1. Redis 连通性
try:
from src.storage.redis_store import get_redis_client
r = get_redis_client()
r.ping()
status['redis'] = 'ok'
except Exception as e:
status['redis'] = f'error: {e}'
issues.append('redis')
# 2. 磁盘空间(data/ 目录)
try:
data_path = Path.home() / 'projects' / 'zhiyi' / 'data'
if data_path.exists():
stat = data_path.stat()
import shutil
total, used, free = shutil.disk_usage(data_path)
status['disk'] = {
'total_gb': round(total / (1024**3), 1),
'used_gb': round(used / (1024**3), 1),
'free_gb': round(free / (1024**3), 1),
'data_mb': round(sum(f.stat().st_size for f in data_path.rglob('*') if f.is_file()) / (1024**2), 1),
}
if free < 1024**3: # < 1GB
issues.append('disk_low')
else:
status['disk'] = 'not_found'
except Exception as e:
status['disk'] = f'error: {e}'
# 3. 向量索引完整性
try:
embedder = _shared_embedder
stats = embedder.get_stats()
status['index'] = {
'doc_count': stats.get('doc_count', 0),
'vector_dim': stats.get('vector_dim', 0),
'index_type': stats.get('index_type', 'unknown'),
'needs_rebuild': stats.get('needs_rebuild', False),
}
if stats.get('needs_rebuild'):
issues.append('index_rebuild')
except Exception as e:
status['index'] = f'error: {e}'
issues.append('index')
# 4. 实例状态
try:
from src.storage.redis_store import get_commit_store, get_instance_id
store = get_commit_store()
instances = store.discover_instances()
status['instances'] = {
'current': get_instance_id(),
'alive': list(instances.keys()),
'count': len(instances),
}
except Exception as e:
status['instances'] = f'error: {e}'
# 5. 发现问题标记
if issues:
status['status'] = 'degraded'
status['issues'] = issues
return status
@app.get('/')
@app.get("/")
def root():
return {'service': '织忆 MemoryWeave', 'version': '3.0', 'docs': '/docs'}
return {"service": "织忆 MemoryWeave", "version": "3.1", "docs": "/docs"}
def main():
@app.get("/health")
def health():
import os, subprocess
stats = {"status": "healthy", "service": "zhiyi", "version": "3.1"}
try:
store = RedisCommitStore()
if store.redis and store.redis.ping():
stats["redis"] = {"connected": True}
else:
stats["redis"] = {"connected": False}
except:
stats["redis"] = {"connected": False}
try:
embedder = get_shared_embedder()
stats["index"] = {
"doc_count": len(embedder.docs),
"vector_dim": embedder.dim,
"index_type": "FlatIP"
}
except Exception as e:
stats["index"] = {"error": str(e)}
return stats
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=7821)
if __name__ == '__main__':
main()
port = int(os.environ.get("ZHIYI_PORT", 7821))
uvicorn.run("api.server:app", host="0.0.0.0", port=port, reload=False)