Archived
feat: NL query + GPU acceleration + Nginx + systemd
This commit is contained in:
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -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
@@ -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)
|
||||
Reference in New Issue
Block a user