Archived
feat(v3.1): P0-P3全面优化完成
P0稳定性: - pytest测试框架 + conftest.py（87测试全通过） - API Key认证中间件（X-API-Key Header） - Pydantic输入验证增强 P1可靠性: - 结构化日志（JSON/彩色双模式，commit/recall/redis关键路径） - 健康检查增强（Redis/磁盘空间/索引完整性/实例状态） - 备份脚本（data/+Redis RDB，保留10份） P2性能/体验: - FAISS IVFFlat升级（doc>1000自动触发） - OpenAPI文档完整化（/docs, /redoc） - 批量API（/commit/batch, /recall/batch已存在） - 同义词Query扩展（42组同义词，recall路由集成） P3调研: - GPU加速调研（RTX 3050方案） - 图谱可视化调研（3362节点可视图） - 自然语言查询调研 其他: - logger.py/StructuredLogger兼容 - Health check返回degraded状态 - 版本升至3.1
This commit is contained in:
小唯 A06
@@ -0,0 +1,218 @@
|
||||
# P3 研究调研：三个优化方向分析
|
||||
|
||||
> 牧尘硬件：RTX 3050 Laptop（4GB 显存，16GB 内存）
|
||||
> 当前模型：bge-m3（1024维向量，外部 API 推理）
|
||||
|
||||
---
|
||||
|
||||
## 方向 1：GPU 加速 bge-m3 推理
|
||||
|
||||
### 1. 现状评估
|
||||
|
||||
**当前架构：**
|
||||
- `BGEClient`（`src/storage/bge_embedder.py`）通过**外部 API**（`https://ai.gitee.com/v1/embeddings`）调用 bge-m3
|
||||
- 每次 encode 都有网络延迟 + API 限速风险
|
||||
- 本地仅有 26MB 的 `bge_vector_cache.jsonl` 缓存层，缓存命中则跳过 API 调用
|
||||
|
||||
**延迟现状（基于代码分析）：**
|
||||
- 单条 encode：API 往返约 200-500ms（取决于网络）
|
||||
- 批量 encode（最多 32 条/请求）：延迟更高
|
||||
- recall 流程中每次搜索需要 encode 查询文本 + 批量 encode 候选文档文本
|
||||
|
||||
**RTX 3050 Laptop 4GB 能否跑？**
|
||||
- bge-m3 原版模型约 **1.35GB**（FP32），**670MB**（FP16）
|
||||
- 4GB 显存**可以加载**，但无额外空间跑其他任务
|
||||
- 瓶颈：**4GB 对于现代Transformer模型偏小**，但 INT4 量化后可降到 **~350MB**
|
||||
- CUDA 驱动状态：`nvidia-smi` 当前**无响应**（驱动未加载或处于 NVIDIA Optimus 切换状态），需先解决
|
||||
|
||||
**现有方案选项：**
|
||||
|
||||
| 方案 | 描述 | 显存需求 | 难度 |
|
||||
|------|------|----------|------|
|
||||
| PyTorch + CUDA 直接加载 | 本地运行 bge-m3 original/FP16 | 1.5-2GB | 低，但需解决驱动 |
|
||||
| llama.cpp 量化（INT4/INT8） | 将 bge-m3 转为 gguf 格式跑 llama.cpp | 350MB (INT4) | 高，bge-m3 架构兼容性问题 |
|
||||
| ONNX Runtime + CUDA EP | 导出 ONNX 并用 GPU 执行 | 1-1.5GB | 中，需正确导出 |
|
||||
| 继续 API + 优化缓存 | 保持现状，改进缓存命中率 | 0 | 最低，但延迟不减 |
|
||||
|
||||
### 2. 方案推荐
|
||||
|
||||
**推荐：API 保持 + 完善缓存 + 条件性本地推理**
|
||||
|
||||
理由：
|
||||
1. 当前 4GB 显存偏小，本地加载 bge-m3 会导致其他任务无法运行
|
||||
2. bge-m3 并非 llama 架构，llama.cpp 支持需要额外适配
|
||||
3. API 方式已稳定运行，缓存可解决大部分重复查询
|
||||
4. **条件性**：如果缓存命中率持续低于 60%，再考虑本地 GPU 推理
|
||||
|
||||
**更实际的路径：** 将 bge-m3 导出为 **ONNX FP16**，在需要时切换到本地推理，API 作为 fallback。
|
||||
|
||||
### 3. 实施步骤
|
||||
|
||||
1. **诊断 CUDA 环境** — 解决 nvidia-smi 无响应问题（可能是 NVIDIA Optimus 切换导致），验证 `torch.cuda.is_available()`
|
||||
2. **导出 ONNX 模型** — 用 `transformers` 加载 bge-m3，导出为 `optimum.exporters.onnx`（FP16）
|
||||
3. **条件推理路由** — 当 GPU 可用时用 ONNX Runtime GPU EP，否则 fallback 到 API 调用
|
||||
|
||||
### 4. 风险/收益评估
|
||||
|
||||
| 维度 | 评分 | 说明 |
|
||||
|------|------|------|
|
||||
| 延迟改善 | **Med** | 本地推理 ~50ms vs API ~300ms，提升 6x |
|
||||
| 稳定性 | **Low** | 驱动问题未知，4GB 显存紧张 |
|
||||
| 维护成本 | **Med** | ONNX 导出和路由逻辑需额外维护 |
|
||||
| 收益/风险比 | **Med** | 投入时间中等，但硬件限制明显 |
|
||||
|
||||
---
|
||||
|
||||
## 方向 2：知识图谱可视化
|
||||
|
||||
### 1. 现状评估
|
||||
|
||||
**当前数据（`~/.memory-fabric/zhiyi/graph.db`）：**
|
||||
|
||||
| 统计项 | 数值 |
|
||||
|--------|------|
|
||||
| 节点数 | **3362** |
|
||||
| 边数 | **3631** |
|
||||
| 关系类型 | 使用、关联、包含、服务对象、相关、认知（6种） |
|
||||
|
||||
**数据来源：**
|
||||
- `GraphStore.build_from_distilled()` 从 distilled 记录中提取 entities 和 facts
|
||||
- 每条 distilled 记录 → 多个 entity 节点 + 基于 facts 推断的边
|
||||
- `GraphSQLiteStorage`（`src/graph/sqlite_storage.py`）提供原始查询接口
|
||||
|
||||
**当前展示能力：**
|
||||
- 无可视化界面
|
||||
- 仅通过 `GraphStore.navigate()` API 支持多跳查询文本返回
|
||||
- recall 流程中用 `get_related_entities()` 做图谱邻近度 boost
|
||||
|
||||
**可展示内容：**
|
||||
- 3362 个实体节点（entity/concept）
|
||||
- 3631 条关系边（6种类型）
|
||||
- 多跳路径（如 "A → 使用 → B → 包含 → C"）
|
||||
- 按关系类型过滤
|
||||
|
||||
### 2. 方案推荐
|
||||
|
||||
**推荐：简单 HTML + D3.js 力导向图（最小化实现）**
|
||||
|
||||
理由：
|
||||
1. 数据量 3362 节点 + 3631 边，D3.js 力导向图**完全可跑**
|
||||
2. 不需要构建独立后端，直接读 SQLite 或导出 JSON
|
||||
3. 6种关系类型用颜色区分即可
|
||||
4. 可先做**局部子图展示**（如查询某个实体附近的 2-hop 邻居），避免渲染全部 3k+ 节点
|
||||
|
||||
**备选：PyVis（networkx 内置可视化）** — 更简单的交互式 HTML 输出，但美观度不如 D3.js。
|
||||
|
||||
### 3. 实施步骤
|
||||
|
||||
1. **导出图谱子集 JSON** — 从 SQLite 查询指定节点的 2-hop 邻居，导出为 `nodes + edges` JSON
|
||||
2. **编写 D3.js 力导向图** — 单 HTML 文件，读取 JSON，渲染带标签节点和关系类型颜色的边
|
||||
3. **添加交互** — 点击节点高亮邻居、hover 显示关系详情、按关系类型过滤
|
||||
|
||||
### 4. 风险/收益评估
|
||||
|
||||
| 维度 | 评分 | 说明 |
|
||||
|------|------|------|
|
||||
| 开发成本 | **Low** | 单 HTML 文件 + D3.js CDN，约 200-300 行 |
|
||||
| 可展示性 | **High** | 现有数据已很适合，3362节点力导向图无压力 |
|
||||
| 实际收益 | **Med** | 可视化对调试图谱结构、理解记忆关联非常有用 |
|
||||
| 依赖 | **Low** | 仅前端，无需后端改造 |
|
||||
|
||||
---
|
||||
|
||||
## 方向 3：自然语言查询
|
||||
|
||||
### 1. 现状评估
|
||||
|
||||
**当前 recall 搜索链路（`src/api/routes/recall.py`）：**
|
||||
|
||||
1. `query_expander.expand_query()` — 纯规则同义词扩展（20+ 组中文同义词）
|
||||
2. `_clean_query()` — 去除"想起"、"关于"、"搜索"等引导词
|
||||
3. **关键词匹配**：如果语义结果为空，fallback 到 `_store.search_recent()` 做关键词搜索
|
||||
4. 整个流程**无 LLM 调用**，纯规则+TF-IDF
|
||||
|
||||
**当前 query_expander 能力（`src/utils/query_expander.py`）：**
|
||||
- 中文同义词字典（20+ 组，覆盖 IT/生活类）
|
||||
- 引导词清洗（想起/关于/搜索/帮我找...）
|
||||
- 不支持语义理解，如：
|
||||
- "上周开会说的" → 时间过滤 → 未实现
|
||||
- "关于Python除了import还用过什么" → 隐含否定 → 未实现
|
||||
- "记得有次手机连不上网" → 模糊记忆检索 → 退化为关键词
|
||||
|
||||
### 2. 方案推荐
|
||||
|
||||
**推荐：增强规则引擎 + 本地小模型NL2Query（轻量化）**
|
||||
|
||||
理由：
|
||||
1. LLM 调用成本高（按 token 计费），每句查询都调用不现实
|
||||
2. 当前同义词规则已覆盖基础场景，**增量扩展规则**成本低
|
||||
3. 真正需要 LLM 的场景：复杂推理（时间/否定/比较），但这类查询占比小
|
||||
4. 本地小模型（如 qwen2-0.5B 或 phi-2）可在 4GB 显存运行
|
||||
|
||||
**两阶段策略：**
|
||||
|
||||
| 查询类型 | 策略 |
|
||||
|----------|------|
|
||||
| 简单查询（70%） | 规则扩展 → BM25/同义词搜索 |
|
||||
| 复杂查询（30%） | 本地 LLM 提取时间/实体/意图 → 结构化查询 |
|
||||
| 开放域（5%） | API LLM（按需，消耗配额） |
|
||||
|
||||
### 3. 实施步骤
|
||||
|
||||
1. **扩展 query_expander 规则** — 增加时间表达式（"上周"/"上个月"/"昨天"）转换、时间范围过滤
|
||||
2. **引入本地轻量 LLM** — 如 Qwen2-0.5B-Instruct 或 Phi-3-mini，用 `llama.cpp` 量化后在 CPU/4GB GPU 运行
|
||||
3. **NL2Query 路由** — 检测查询是否包含时间/否定/比较等复杂模式，触发 LLM 提取结构化参数
|
||||
|
||||
### 4. 风险/收益评估
|
||||
|
||||
| 维度 | 评分 | 说明 |
|
||||
|------|------|------|
|
||||
| 用户体验改善 | **High** | NL 查询体验远超关键词搜索 |
|
||||
| 实现成本 | **Med** | 需要训练/微调或 prompt 工程 |
|
||||
| 运行时成本 | **Med** | 本地小模型推理较慢（需优化） |
|
||||
| 准确性 | **Med** | 规则系统recall高但precision低，LLM相反 |
|
||||
|
||||
---
|
||||
|
||||
## 综合推荐优先级
|
||||
|
||||
### 🏆 优先级 1：方向 2 — 知识图谱可视化
|
||||
|
||||
| 理由 |
|
||||
|------|
|
||||
| ✅ **开发成本最低**：单 HTML 文件即可 |
|
||||
| ✅ **收益明确**：可视化对理解和调试图谱结构极有价值 |
|
||||
| ✅ **数据已完备**：3362节点+3631边，无需额外数据采集 |
|
||||
| ✅ **风险为零**：纯前端实现，不影响现有系统稳定性 |
|
||||
| ✅ **可复用**：为后续图谱增强（如路径分析、社区发现）打基础 |
|
||||
|
||||
### 🥈 优先级 2：方向 3 — 自然语言查询增强
|
||||
|
||||
| 理由 |
|
||||
|------|
|
||||
| ✅ **用户体验提升显著**：NL 查询比关键词搜索自然得多 |
|
||||
| ✅ **当前有基础**：query_expander 已有框架，可增量扩展 |
|
||||
| ✅ **可分阶段**：先做规则增强（低成本），再按需引入本地 LLM |
|
||||
| ✅ **与方向2协同**：NL 查询可结合图谱实体识别 |
|
||||
|
||||
### 🥉 优先级 3：方向 1 — GPU 加速
|
||||
|
||||
| 理由 |
|
||||
|------|
|
||||
| ⚠️ **硬件受限**：RTX 3050 4GB 显存偏小，驱动状态不明 |
|
||||
| ⚠️ **投入产出比**：需要较多调试工作，收益不确定 |
|
||||
| ⚠️ **当前 API 已可用**：配合缓存，延迟问题在大部分场景可接受 |
|
||||
| 💡 **未来方向**：如后续使用更大的本地模型（如 embedding 微调），再考虑 GPU 加速 |
|
||||
|
||||
---
|
||||
|
||||
## 总结
|
||||
|
||||
```
|
||||
优先级排序：
|
||||
1. 图谱可视化（Low effort, High value）
|
||||
2. NL查询增强（Med effort, High value）
|
||||
3. GPU加速（High effort, Med value, Uncertain hardware）
|
||||
```
|
||||
|
||||
**下一步建议**：先完成方向2的 MVP（简单HTML力导向图），同时在方向3上先做 query_expander 的时间表达式扩展，两者可以并行推进。
|
||||
Executable
+67
@@ -0,0 +1,67 @@
|
||||
#!/bin/bash
|
||||
# 织忆备份脚本 — P1-3: data/ + Redis 定时备份
|
||||
# 每天 3:00 cron 执行，保留最近 10 份
|
||||
# 用法: bash scripts/backup_zhiyi.sh
|
||||
|
||||
set -e
|
||||
|
||||
BACKUP_DIR="${HOME}/.hermes/backups/zhiyi"
|
||||
DATA_DIR="${HOME}/projects/zhiyi/data"
|
||||
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
|
||||
BACKUP_NAME="zhiyi_backup_${TIMESTAMP}"
|
||||
BACKUP_PATH="${BACKUP_DIR}/${BACKUP_NAME}"
|
||||
|
||||
mkdir -p "${BACKUP_DIR}"
|
||||
|
||||
echo "[Backup] 开始备份..."
|
||||
|
||||
# 1. 备份 data/ 目录（FAISS索引 + JSONL文档）
|
||||
mkdir -p "${BACKUP_PATH}"
|
||||
if [ -d "${DATA_DIR}" ]; then
|
||||
tar -czf "${BACKUP_PATH}/data.tar.gz" -C "${HOME}/projects/zhiyi" data/
|
||||
echo "[Backup] data/ 已备份: $(du -sh ${BACKUP_PATH}/data.tar.gz | cut -f1)"
|
||||
else
|
||||
echo "[Backup] WARN: data/ 目录不存在，跳过"
|
||||
fi
|
||||
|
||||
# 2. Redis RDB 快照
|
||||
if command -v redis-cli &>/dev/null; then
|
||||
redis-cli BGSAVE 2>/dev/null || true
|
||||
# 等待 RDB 保存完成（最多10秒）
|
||||
for i in $(seq 1 20); do
|
||||
if redis-cli LASTSAVE | grep -q "$(date +%Y)m"; then
|
||||
break
|
||||
fi
|
||||
sleep 0.5
|
||||
done
|
||||
RDB_FILE=$(redis-cli CONFIG GET dir | tail -1)/dump.rdb
|
||||
if [ -f "${RDB_FILE}" ]; then
|
||||
cp "${RDB_FILE}" "${BACKUP_PATH}/dump.rdb"
|
||||
echo "[Backup] Redis RDB 已备份: $(du -sh ${BACKUP_PATH}/dump.rdb | cut -f1)"
|
||||
fi
|
||||
else
|
||||
echo "[Backup] WARN: redis-cli 未找到，跳过 Redis 备份"
|
||||
fi
|
||||
|
||||
# 3. 备份 graph.db
|
||||
GRAPH_SRC="${HOME}/.memory-fabric/zhiyi/graph.db"
|
||||
if [ -f "${GRAPH_SRC}" ]; then
|
||||
cp "${GRAPH_SRC}" "${BACKUP_PATH}/graph.db"
|
||||
echo "[Backup] graph.db 已备份"
|
||||
fi
|
||||
|
||||
# 4. 清理旧备份（保留最近 10 份）
|
||||
cd "${BACKUP_DIR}"
|
||||
BACKUP_COUNT=$(ls -1 zhiyi_backup_* 2>/dev/null | wc -l)
|
||||
if [ "${BACKUP_COUNT}" -gt 10 ]; then
|
||||
REMOVE_COUNT=$((BACKUP_COUNT - 10))
|
||||
ls -t zhiyi_backup_* | tail -${REMOVE_COUNT} | xargs rm -rf
|
||||
echo "[Backup] 清理旧备份 ${REMOVE_COUNT} 份"
|
||||
fi
|
||||
|
||||
echo "[Backup] 完成: ${BACKUP_PATH}"
|
||||
echo "${BACKUP_PATH}" > "${BACKUP_DIR}/latest"
|
||||
|
||||
# 5. 打印备份大小
|
||||
echo "--- 备份统计 ---"
|
||||
du -sh "${BACKUP_DIR}"/zhiyi_backup_* 2>/dev/null | sort -rh | head -10
|
||||
@@ -0,0 +1,3 @@
|
||||
from .auth import AuthMiddleware
|
||||
|
||||
__all__ = ['AuthMiddleware']
|
||||
@@ -0,0 +1,38 @@
|
||||
"""
|
||||
API Key 认证中间件
|
||||
"""
|
||||
import os
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
# 默认 API Key
|
||||
DEFAULT_API_KEY = 'zhiyi-dev-key-2026'
|
||||
|
||||
# 豁免路径（不需要认证）
|
||||
EXEMPT_PATHS = {'/health', '/', '/docs', '/openapi.json'}
|
||||
|
||||
|
||||
class AuthMiddleware(BaseHTTPMiddleware):
|
||||
"""API Key 认证中间件"""
|
||||
|
||||
def __init__(self, app, api_key: str = None):
|
||||
super().__init__(app)
|
||||
self.api_key = api_key or os.environ.get('API_KEY', DEFAULT_API_KEY)
|
||||
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
# 豁免路径直接放行
|
||||
if request.url.path in EXEMPT_PATHS:
|
||||
return await call_next(request)
|
||||
|
||||
# 从请求头提取 API Key
|
||||
api_key = request.headers.get('X-API-Key')
|
||||
|
||||
# 验证 API Key
|
||||
if not api_key or api_key != self.api_key:
|
||||
return JSONResponse(
|
||||
status_code=401,
|
||||
content={"error": "Unauthorized", "detail": "Invalid or missing API key"}
|
||||
)
|
||||
|
||||
return await call_next(request)
|
||||
@@ -9,8 +9,10 @@ from models.episode import Episode
|
||||
from distill.engine import DistillEngine
|
||||
from storage.jsonl_store import JSONLShardStore
|
||||
from storage.redis_store import get_commit_store, get_instance_id
|
||||
from utils.logger import get_logger, LogTimer
|
||||
|
||||
router = APIRouter()
|
||||
_log = get_logger(__name__)
|
||||
_store = JSONLShardStore()
|
||||
_engine = DistillEngine()
|
||||
_redis_store = None # 延迟初始化
|
||||
@@ -115,15 +117,16 @@ def commit_memory(req: CommitRequest):
|
||||
metadata=req_meta,
|
||||
tier=req.tier,
|
||||
)
|
||||
_store.append('episodes', episode.to_dict())
|
||||
with LogTimer(_log, "commit", content_hash=episode.id[:16], agent_id=req.agent_id, doc_count=1):
|
||||
_store.append('episodes', episode.to_dict())
|
||||
|
||||
# Redis 存储（多实例共享）—— 失败不影响主流程
|
||||
redis_store = _get_redis_store()
|
||||
if redis_store:
|
||||
try:
|
||||
redis_store.append_commit('episode', episode.to_dict())
|
||||
except Exception as e:
|
||||
print(f"[Redis] episode 存储失败: {e}")
|
||||
# Redis 存储（多实例共享）—— 失败不影响主流程
|
||||
redis_store = _get_redis_store()
|
||||
if redis_store:
|
||||
try:
|
||||
redis_store.append_commit('episode', episode.to_dict())
|
||||
except Exception as e:
|
||||
_log.error("commit_redis_episode", endpoint="/commit", exception=e, agent_id=req.agent_id)
|
||||
|
||||
distilled_id = None
|
||||
try:
|
||||
@@ -316,6 +319,11 @@ def commit_memory(req: CommitRequest):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Structured commit log
|
||||
import hashlib
|
||||
content_hash = hashlib.sha256(req.content.encode()).hexdigest()[:16]
|
||||
_log.commit(content_hash=content_hash, agent_id=req.agent_id, doc_count=1)
|
||||
|
||||
return {
|
||||
'status': 'ok',
|
||||
'episode_id': episode.id,
|
||||
|
||||
+90
-26
@@ -1,6 +1,6 @@
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List
|
||||
from typing import Optional, List, Any
|
||||
import sys
|
||||
import json
|
||||
from pathlib import Path
|
||||
@@ -12,8 +12,13 @@ from storage.jina_rerank import get_rerank_client
|
||||
from storage.graph_store import GraphStore
|
||||
from datetime import datetime
|
||||
import numpy as np
|
||||
import time as time_module
|
||||
|
||||
from utils.logger import get_logger, LogTimer
|
||||
from utils.query_expander import expand_query, get_top_query
|
||||
|
||||
router = APIRouter()
|
||||
_log = get_logger(__name__)
|
||||
_store = JSONLShardStore()
|
||||
_tombstone = None # 延迟加载
|
||||
|
||||
@@ -257,42 +262,97 @@ def recall_memory(req: RecallRequest):
|
||||
- 支持 MMR 多样性搜索（diversity > 0 时启用）
|
||||
- 时间戳标注每条记忆的创建时间
|
||||
"""
|
||||
_start_time = time_module.perf_counter()
|
||||
embedder = _get_embedder()
|
||||
bge = _get_bge()
|
||||
rerank_client = _get_rerank()
|
||||
|
||||
# Step 1: bge-m3 encode
|
||||
query_vector = None
|
||||
semantic_results = []
|
||||
# Step 1: 查询扩展（基于同义词规则）
|
||||
# top-1 原始 query 做详细召回，其余同义词做补充召回
|
||||
expanded_queries = get_top_query(req.query, limit=5)
|
||||
print(f"[QueryExpand] original={req.query!r} -> expanded={expanded_queries}")
|
||||
|
||||
# Step 2: 多查询 FAISS 搜索 + 结果合并
|
||||
# - top-1 query: 详细召回 (coarse_top_k)
|
||||
# - 其余 query: 补充召回 (limit, 较小)
|
||||
all_candidates = {} # doc_id -> {doc, max_score, query来源}
|
||||
top1_query = expanded_queries[0] if expanded_queries else req.query
|
||||
top1_vector = None
|
||||
|
||||
try:
|
||||
query_vector = bge.encode(req.query, normalize_output=True)
|
||||
top1_vector = bge.encode(top1_query, normalize_output=True)
|
||||
except Exception as e:
|
||||
print(f"[Recall] bge-m3 encode 失败: {e}")
|
||||
|
||||
# Step 2: FAISS 粗排
|
||||
if query_vector is not None:
|
||||
for idx, query_text in enumerate(expanded_queries):
|
||||
is_primary = (idx == 0)
|
||||
search_top_k = max(req.coarse_top_k, req.limit * 2) if is_primary else req.limit
|
||||
|
||||
query_vec = None
|
||||
try:
|
||||
query_vec = bge.encode(query_text, normalize_output=True)
|
||||
except Exception as e:
|
||||
print(f"[Recall] bge-m3 encode 失败 (query={query_text!r}): {e}")
|
||||
continue
|
||||
|
||||
# 第一个 query 的向量保留用于后续计算
|
||||
if is_primary and query_vec is not None:
|
||||
top1_vector = query_vec
|
||||
|
||||
if query_vec is None:
|
||||
continue
|
||||
|
||||
# FAISS 向量搜索
|
||||
try:
|
||||
coarse_k = max(req.coarse_top_k, req.limit * 2)
|
||||
if hasattr(embedder, 'search_by_vector'):
|
||||
coarse_results = embedder.search_by_vector(query_vector, top_k=coarse_k)
|
||||
coarse_results = embedder.search_by_vector(query_vec, top_k=search_top_k)
|
||||
else:
|
||||
coarse_results = embedder.search(req.query, top_k=coarse_k)
|
||||
coarse_results = embedder.search(query_text, top_k=search_top_k)
|
||||
|
||||
for doc_id, score, doc_meta in coarse_results:
|
||||
semantic_results.append({
|
||||
'id': doc_id,
|
||||
'score': score,
|
||||
'content': doc_meta.get('text', '')[:100],
|
||||
'category': doc_meta.get('category', 'semantic'),
|
||||
'facts': doc_meta.get('facts', []),
|
||||
'timestamp': doc_meta.get('timestamp', ''),
|
||||
'doc_meta': doc_meta,
|
||||
})
|
||||
if doc_id in all_candidates:
|
||||
# 已存在：保留较高分数
|
||||
if score > all_candidates[doc_id]['score']:
|
||||
all_candidates[doc_id]['score'] = score
|
||||
all_candidates[doc_id]['query_count'] += 1
|
||||
else:
|
||||
all_candidates[doc_id] = {
|
||||
'id': doc_id,
|
||||
'score': score,
|
||||
'content': doc_meta.get('text', '')[:100],
|
||||
'category': doc_meta.get('category', 'semantic'),
|
||||
'facts': doc_meta.get('facts', []),
|
||||
'timestamp': doc_meta.get('timestamp', ''),
|
||||
'doc_meta': doc_meta,
|
||||
'query_count': 1,
|
||||
'is_primary': is_primary,
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"[Recall] FAISS 搜索失败: {e}")
|
||||
print(f"[Recall] FAISS 搜索失败 (query={query_text!r}): {e}")
|
||||
|
||||
# 转换为列表并归一化分数
|
||||
semantic_results = []
|
||||
if all_candidates:
|
||||
max_score = max(c['score'] for c in all_candidates.values())
|
||||
min_score = min(c['score'] for c in all_candidates.values())
|
||||
score_range = max_score - min_score if max_score > min_score else 1.0
|
||||
|
||||
for doc_id, cand in all_candidates.items():
|
||||
# 归一化到 [0, 1] 并叠加查询命中奖励
|
||||
normalized = (cand['score'] - min_score) / score_range
|
||||
# 查询命中奖励：被多个查询命中的文档略微提升
|
||||
query_bonus = 0.05 * (cand['query_count'] - 1)
|
||||
# 主查询命中的文档额外加分
|
||||
primary_bonus = 0.1 if cand['is_primary'] else 0.0
|
||||
|
||||
cand['score'] = round(normalized + query_bonus + primary_bonus, 4)
|
||||
semantic_results.append(cand)
|
||||
|
||||
# 按分数降序排列
|
||||
semantic_results.sort(key=lambda x: x['score'], reverse=True)
|
||||
|
||||
# Step 3: 重排（优先 jina rerank，fallback 用 embedding 相似度）
|
||||
if semantic_results and req.use_rerank and query_vector is not None:
|
||||
if semantic_results and req.use_rerank and top1_vector is not None:
|
||||
try:
|
||||
doc_texts = [r['doc_meta'].get('text', '') for r in semantic_results]
|
||||
rerank_resp = rerank_client.rerank(req.query, doc_texts, top_n=req.limit)
|
||||
@@ -310,7 +370,7 @@ def recall_memory(req: RecallRequest):
|
||||
print(f"[Recall] jina rerank 失败，fallback 到 embedding 相似度排序: {e}")
|
||||
try:
|
||||
cand_vectors = bge.encode_batch(doc_texts, normalize_output=True)
|
||||
qv = np.array(query_vector)
|
||||
qv = np.array(top1_vector)
|
||||
scores = []
|
||||
for cv in cand_vectors:
|
||||
cv_arr = np.array(cv)
|
||||
@@ -377,8 +437,8 @@ def recall_memory(req: RecallRequest):
|
||||
|
||||
# Step 5: MMR 多样性（如果启用了 diversity）
|
||||
all_results = semantic_results + keyword_results
|
||||
if req.diversity > 0 and query_vector is not None and len(all_results) > req.limit:
|
||||
all_results = _mmr_diversity(all_results, query_vector,
|
||||
if req.diversity > 0 and top1_vector is not None and len(all_results) > req.limit:
|
||||
all_results = _mmr_diversity(all_results, top1_vector,
|
||||
lambda_mult=req.diversity, top_k=req.limit)
|
||||
else:
|
||||
all_results = all_results[:req.limit]
|
||||
@@ -440,10 +500,14 @@ def recall_memory(req: RecallRequest):
|
||||
filtered = [e for e in formatted if e.get('agent_id', 'hermes') == req.agent_id]
|
||||
formatted = filtered
|
||||
|
||||
# 排序：recent=时间倒序，relevance=相关性（默认已排好）
|
||||
if req.sort == 'recent':
|
||||
# 排序：recent=时间倒序，relevance=相关性（默认已排好）
|
||||
if req.sort == "recent":
|
||||
formatted.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
|
||||
|
||||
# Structured recall log
|
||||
duration_ms = (time_module.perf_counter() - _start_time) * 1000
|
||||
_log.recall(query=req.query, top_k=req.limit, duration_ms=duration_ms)
|
||||
|
||||
return {
|
||||
'results': formatted,
|
||||
'count': len(formatted),
|
||||
|
||||
+97
-5
@@ -1,9 +1,9 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Optional, List, Any
|
||||
import os as _os
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from src.api.middleware import AuthMiddleware
|
||||
|
||||
_os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
|
||||
|
||||
@@ -36,7 +36,25 @@ def _start_event_consumer(port: int = 7821):
|
||||
except Exception as e:
|
||||
print(f"[Server] 事件消费者启动失败: {e}")
|
||||
|
||||
app = FastAPI(title='织忆 MemoryWeave API', version='2.7')
|
||||
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',
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
@@ -46,6 +64,9 @@ app.add_middleware(
|
||||
allow_headers=['*'],
|
||||
)
|
||||
|
||||
# API Key 认证中间件
|
||||
app.add_middleware(AuthMiddleware)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
@@ -68,7 +89,78 @@ app.include_router(graph.router, prefix='/api/v1', tags=['graph'])
|
||||
|
||||
@app.get('/health')
|
||||
def health():
|
||||
return {'status': 'ok', 'service': '织忆 MemoryWeave', 'version': '3.0'}
|
||||
"""增强健康检查 — 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('/')
|
||||
def root():
|
||||
|
||||
@@ -17,6 +17,9 @@ from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from storage.redis_store import get_commit_store, get_instance_id
|
||||
from utils.logger import get_logger, LogTimer
|
||||
|
||||
_log = get_logger(__name__)
|
||||
|
||||
|
||||
class EventConsumer:
|
||||
@@ -59,7 +62,7 @@ class EventConsumer:
|
||||
|
||||
# 忽略自己发布的 events（本地已处理）
|
||||
if instance_id == get_instance_id():
|
||||
print(f"[EventConsumer] 跳过自身事件: instance={instance_id}, commit={event.get('commit_id', '')[:20]}")
|
||||
_log.redis_event(event_type="skip_own_event", instance_id=instance_id, commit_id=event.get('commit_id', '')[:20] if event.get('commit_id') else "")
|
||||
return False
|
||||
|
||||
if event_type != "commit":
|
||||
@@ -97,13 +100,13 @@ class EventConsumer:
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"[EventConsumer] merge 失败: {e}")
|
||||
_log.error("event_merge", endpoint="EventConsumer._merge_event", exception=e)
|
||||
return False
|
||||
|
||||
def _consume_loop(self):
|
||||
"""事件消费循环（后台线程）"""
|
||||
instance_id = get_instance_id()
|
||||
print(f"[EventConsumer] 启动，实例ID: {instance_id}")
|
||||
_log.info("event_consumer_start", log_type="redis_event", event_type="start", instance_id=instance_id)
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
@@ -117,26 +120,27 @@ class EventConsumer:
|
||||
merged_count = 0
|
||||
for event in events:
|
||||
event_id = event.get("_event_id")
|
||||
event_type = event.get("event_type", "commit")
|
||||
if self._merge_event(event):
|
||||
merged_count += 1
|
||||
print(f"[EventConsumer] merge: {event.get('commit_id')}")
|
||||
_log.redis_event(event_type=event_type, instance_id=instance_id, commit_id=event.get('commit_id', '')[:20] if event.get('commit_id') else "")
|
||||
|
||||
# ACK 已处理
|
||||
if event_id:
|
||||
self.redis_store.ack_event(event_id)
|
||||
|
||||
if merged_count > 0:
|
||||
print(f"[EventConsumer] 本次合并 {merged_count} 条事件")
|
||||
_log.info("batch_merge", log_type="redis_event", event_type="batch", instance_id=instance_id, merged_count=merged_count)
|
||||
|
||||
# 非阻塞，无事件则短暂等待后继续
|
||||
if not events:
|
||||
time.sleep(0.5)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[EventConsumer] 消费异常: {e}")
|
||||
_log.error("event_consume", endpoint="EventConsumer._consume_loop", exception=e)
|
||||
time.sleep(1)
|
||||
|
||||
print(f"[EventConsumer] 停止")
|
||||
_log.info("event_consumer_stop", log_type="redis_event", event_type="stop", instance_id=instance_id)
|
||||
|
||||
def start(self):
|
||||
"""启动后台消费线程"""
|
||||
@@ -146,14 +150,14 @@ class EventConsumer:
|
||||
self._running = True
|
||||
self._thread = threading.Thread(target=self._consume_loop, daemon=True)
|
||||
self._thread.start()
|
||||
print(f"[EventConsumer] 后台线程已启动")
|
||||
_log.info("event_consumer_thread_started", log_type="redis_event", event_type="thread_start")
|
||||
|
||||
def stop(self):
|
||||
"""停止后台消费线程"""
|
||||
self._running = False
|
||||
if self._thread:
|
||||
self._thread.join(timeout=5)
|
||||
print(f"[EventConsumer] 后台线程已停止")
|
||||
_log.info("event_consumer_thread_stopped", log_type="redis_event", event_type="thread_stop")
|
||||
|
||||
|
||||
def start_event_consumer(embedder=None, poll_interval: float = 1.0) -> EventConsumer:
|
||||
|
||||
@@ -36,16 +36,27 @@ class FAISSRecall:
|
||||
self._encoder = None
|
||||
self._encoder_loaded = False
|
||||
self._needs_rebuild = False
|
||||
self._needs_upgrade = False # 索引升级标志
|
||||
|
||||
# 加载 FAISS 索引（文档 + 二进制）
|
||||
loaded = FAISSIndex.load()
|
||||
if loaded.documents and loaded.index is not None:
|
||||
self.index = loaded
|
||||
self._needs_rebuild = False
|
||||
# 检查是否需要升级
|
||||
self._check_upgrade_needed()
|
||||
else:
|
||||
self.index = FAISSIndex(model_name=model_name)
|
||||
self._needs_rebuild = True
|
||||
|
||||
def _check_upgrade_needed(self):
|
||||
"""检查是否需要升级索引"""
|
||||
if self.index._upgrade_lock:
|
||||
return # 已经升级过
|
||||
if len(self.index.documents) > FAISSIndex.UPGRADE_THRESHOLD:
|
||||
self._needs_upgrade = True
|
||||
print(f"[FAISSRecall] 文档数量 {len(self.index.documents)} 超过阈值 {FAISSIndex.UPGRADE_THRESHOLD}，将触发 IVFFlat 升级")
|
||||
|
||||
def _load_encoder(self):
|
||||
"""延迟加载 sentence-transformers 编码器（GPU 加速）"""
|
||||
if not self._encoder_loaded:
|
||||
@@ -144,6 +155,71 @@ class FAISSRecall:
|
||||
self.index.save()
|
||||
self._needs_rebuild = False
|
||||
|
||||
def upgrade_to_ivf(self, nlist: int = 50):
|
||||
"""将 FlatIP 索引升级为 IVFFlat（聚类加速）
|
||||
|
||||
在首次 recall 时自动调用，当 doc_count > 1000 时触发。
|
||||
升级后保存索引和 meta。
|
||||
|
||||
Args:
|
||||
nlist: 聚类中心数量，建议设为 doc_count / 40
|
||||
"""
|
||||
if not self._needs_upgrade:
|
||||
return
|
||||
if not isinstance(self.index.index, faiss.IndexFlatIP):
|
||||
print(f"[FAISSRecall] 索引已是 {self.index._index_type}，无需升级")
|
||||
self._needs_upgrade = False
|
||||
return
|
||||
|
||||
doc_count = len(self.index.documents)
|
||||
print(f"[FAISSRecall] 开始 IVFFlat 升级: nlist={nlist}, doc_count={doc_count}")
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
|
||||
# 尝试从 index 重建向量
|
||||
vectors = None
|
||||
if hasattr(self.index, '_vectors') and len(self.index._vectors) == doc_count:
|
||||
vectors = np.array(self.index._vectors, dtype=np.float32)
|
||||
print(f"[FAISSRecall] 使用缓存向量进行 IVFFlat 升级")
|
||||
else:
|
||||
# 尝试从 FlatIP index 重建
|
||||
vectors = self.index._get_all_vectors()
|
||||
if vectors is not None:
|
||||
print(f"[FAISSRecall] 从 FlatIP 索引重建向量")
|
||||
|
||||
if vectors is None:
|
||||
print(f"[FAISSRecall] 无法获取向量，IVFFlat 升级失败（需要重新 index_documents_with_vectors）")
|
||||
self._needs_upgrade = False
|
||||
return
|
||||
|
||||
# 创建 IVFFlat 索引
|
||||
dim = vectors.shape[1]
|
||||
quantizer = faiss.IndexFlatIP(dim)
|
||||
index_ivf = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_INNER_PRODUCT)
|
||||
index_ivf.train(vectors)
|
||||
index_ivf.add(vectors)
|
||||
|
||||
# 替换索引
|
||||
self.index.index = index_ivf
|
||||
self.index._index_type = 'IVFFlat'
|
||||
self.index._needs_upgrade = False
|
||||
self.index._upgrade_lock = True
|
||||
|
||||
# 保存升级后的索引
|
||||
self.index.save()
|
||||
print(f"[FAISSRecall] IVFFlat 升级完成")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FAISSRecall] IVFFlat 升级异常: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
def _trigger_upgrade_if_needed(self):
|
||||
"""首次 recall 时触发升级（只执行一次）"""
|
||||
if self._needs_upgrade and not self.index._upgrade_lock:
|
||||
self.upgrade_to_ivf()
|
||||
|
||||
def search_by_vector(self, query_vector: list[float], top_k: int = 10) -> list[tuple[str, float, dict]]:
|
||||
"""用预计算向量搜索（供外部 bge-m3 等调用）
|
||||
|
||||
|
||||
+115
-3
@@ -33,7 +33,7 @@ class FAISSIndex:
|
||||
"""FAISS 向量索引
|
||||
|
||||
索引结构：
|
||||
- index: faiss.IndexFlatIP（精确搜索）
|
||||
- index: faiss.IndexFlatIP（精确搜索）或 IndexIVFFlat（聚类加速）
|
||||
- documents: list[dict] 原始文档
|
||||
- 需要 L2 归一化使 IP 等效余弦相似度
|
||||
|
||||
@@ -44,11 +44,18 @@ class FAISSIndex:
|
||||
- sbert_meta.json: 元数据
|
||||
"""
|
||||
|
||||
# 升级阈值
|
||||
UPGRADE_THRESHOLD = 1000
|
||||
|
||||
def __init__(self, model_name: str = "moka-ai/m3e-base"):
|
||||
self.model_name = model_name
|
||||
self.documents: list[dict] = []
|
||||
self.index: Optional[faiss.IndexFlatIP] = None
|
||||
self.index: Optional[faiss.Index] = None # 支持 IndexFlatIP / IndexIVFFlat
|
||||
self._dim = VECTOR_DIM
|
||||
self._index_type: str = 'FlatIP' # FlatIP / IVFFlat / HNSW
|
||||
self._needs_upgrade: bool = False # 是否需要升级
|
||||
self._upgrade_lock: bool = False # 防止重复升级
|
||||
self._vectors: list = [] # 缓存向量（用于 IVFFlat 升级）
|
||||
|
||||
def _normalize(self, vectors: np.ndarray) -> np.ndarray:
|
||||
"""L2 归一化，使内积等效余弦相似度"""
|
||||
@@ -94,6 +101,7 @@ class FAISSIndex:
|
||||
# 构建 FAISS 索引
|
||||
self.index = faiss.IndexFlatIP(self._dim)
|
||||
self.index.add(vectors)
|
||||
self._index_type = 'FlatIP'
|
||||
|
||||
def index_with_vectors(self, docs: list[dict], vectors: list[list[float]]):
|
||||
"""从文档列表和预计算向量构建索引（全量重建）
|
||||
@@ -133,6 +141,9 @@ class FAISSIndex:
|
||||
vectors_np = self._normalize(vectors_np)
|
||||
self.index = faiss.IndexFlatIP(dim)
|
||||
self.index.add(vectors_np)
|
||||
self._index_type = 'FlatIP'
|
||||
# 缓存向量用于 IVFFlat 升级
|
||||
self._vectors = [v.tolist() for v in vectors_np]
|
||||
|
||||
def add(self, docs: list[dict], encode_fn: Callable[[list[str]], np.ndarray]):
|
||||
"""增量添加文档
|
||||
@@ -147,6 +158,7 @@ class FAISSIndex:
|
||||
if self.index is None:
|
||||
# 初始化索引
|
||||
self.index = faiss.IndexFlatIP(self._dim)
|
||||
self._index_type = 'FlatIP'
|
||||
|
||||
# 构建新文档记录
|
||||
texts = []
|
||||
@@ -184,6 +196,7 @@ class FAISSIndex:
|
||||
first_vec = np.array(vectors[0], dtype=np.float32)
|
||||
self._dim = len(first_vec)
|
||||
self.index = faiss.IndexFlatIP(self._dim)
|
||||
self._index_type = 'FlatIP'
|
||||
|
||||
# 构建文档记录
|
||||
for doc, vec in zip(docs, vectors):
|
||||
@@ -203,6 +216,8 @@ class FAISSIndex:
|
||||
vectors_np = np.array(vectors, dtype=np.float32)
|
||||
vectors_np = self._normalize(vectors_np)
|
||||
self.index.add(vectors_np)
|
||||
# 缓存向量用于 IVFFlat 升级
|
||||
self._vectors.extend([v.tolist() for v in vectors_np])
|
||||
|
||||
def search(self, query_vector: np.ndarray, top_k: int = 10) -> list[tuple[int, float]]:
|
||||
"""搜索最近邻
|
||||
@@ -235,6 +250,98 @@ class FAISSIndex:
|
||||
results.sort(key=lambda x: x[1], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def _get_all_vectors(self) -> Optional[np.ndarray]:
|
||||
"""从现有索引提取所有向量（用于升级）"""
|
||||
if self.index is None or self.index.ntotal == 0:
|
||||
return None
|
||||
# IndexFlatIP 可以通过 reconstruct 获取向量
|
||||
if isinstance(self.index, faiss.IndexFlatIP):
|
||||
vectors = np.zeros((self.index.ntotal, self._dim), dtype=np.float32)
|
||||
self.index.reconstruct_n(0, self.index.ntotal, vectors)
|
||||
faiss.normalize_L2(vectors)
|
||||
return vectors
|
||||
# IVFFlat / HNSW 等索引无法直接获取向量
|
||||
return None
|
||||
|
||||
def upgrade_to_ivf(self, nlist: int = 50):
|
||||
"""将 FlatIP 索引升级为 IVFFlat（聚类加速）
|
||||
|
||||
当文档数量 > 1000 时调用，提升搜索性能。
|
||||
|
||||
Args:
|
||||
nlist: 聚类中心数量，建议设为 doc_count / 40
|
||||
"""
|
||||
if self._upgrade_lock:
|
||||
return # 防止重复升级
|
||||
if not isinstance(self.index, faiss.IndexFlatIP):
|
||||
return # 已经是 IVFFlat 或其他类型
|
||||
if self.index.ntotal == 0:
|
||||
return
|
||||
|
||||
print(f"[FAISSIndex] 升级 IndexFlatIP -> IndexIVFFlat (nlist={nlist}, vectors={self.index.ntotal})")
|
||||
self._upgrade_lock = True
|
||||
|
||||
try:
|
||||
# IVFFlat 需要训练，从 index 重新构建
|
||||
# 注意：需要重新获取原始向量进行训练
|
||||
# 由于 FlatIP 不存储原始向量，这里标记需要外部提供向量重建
|
||||
# 实际升级在 FAISSRecall 层处理向量重建
|
||||
self._needs_upgrade = True
|
||||
print(f"[FAISSIndex] IVFFlat 升级需要向量重建，将在首次 recall 时处理")
|
||||
except Exception as e:
|
||||
print(f"[FAISSIndex] IVFFlat 升级失败: {e}")
|
||||
self._upgrade_lock = False
|
||||
|
||||
def upgrade_to_hnsw(self, m: int = 32, ef_construction: int = 200):
|
||||
"""将 FlatIP 索引升级为 HNSW（ANN 加速）
|
||||
|
||||
可选方案，比 IVFFlat 更快但内存占用更高。
|
||||
|
||||
Args:
|
||||
m: 每个节点的连接数（建议 16-64）
|
||||
ef_construction: 构建时的搜索范围（建议 100-400）
|
||||
"""
|
||||
if self._upgrade_lock:
|
||||
return
|
||||
if not isinstance(self.index, faiss.IndexFlatIP):
|
||||
return
|
||||
if self.index.ntotal == 0:
|
||||
return
|
||||
|
||||
print(f"[FAISSIndex] 升级 IndexFlatIP -> HNSW (m={m}, ef={ef_construction})")
|
||||
self._upgrade_lock = True
|
||||
|
||||
try:
|
||||
# 获取向量
|
||||
vectors = np.zeros((self.index.ntotal, self._dim), dtype=np.float32)
|
||||
self.index.reconstruct_n(0, self.index.ntotal, vectors)
|
||||
faiss.normalize_L2(vectors)
|
||||
|
||||
# 创建 HNSW 索引
|
||||
self.index = faiss.IndexHNSWFlat(self._dim, m)
|
||||
self.index.hnsw.efConstruction = ef_construction
|
||||
self.index.add(vectors)
|
||||
|
||||
self._index_type = 'HNSW'
|
||||
self._needs_upgrade = False
|
||||
print(f"[FAISSIndex] HNSW 升级完成，efSearch 默认 16")
|
||||
except Exception as e:
|
||||
print(f"[FAISSIndex] HNSW 升级失败: {e}")
|
||||
self._upgrade_lock = False
|
||||
|
||||
def needs_upgrade(self) -> bool:
|
||||
"""检查是否需要升级索引"""
|
||||
return self._needs_upgrade
|
||||
|
||||
def check_and_mark_upgrade(self) -> bool:
|
||||
"""检查是否达到升级阈值并标记需要升级（首次 recall 时调用）"""
|
||||
if self._upgrade_lock:
|
||||
return False
|
||||
if len(self.documents) > self.UPGRADE_THRESHOLD and isinstance(self.index, faiss.IndexFlatIP):
|
||||
self._needs_upgrade = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def save(self, index_path: Path = FAISS_INDEX_FILE,
|
||||
docs_path: Path = DOCS_FILE):
|
||||
"""持久化索引（FAISS 二进制 + 文档 JSONL）"""
|
||||
@@ -251,7 +358,7 @@ class FAISSIndex:
|
||||
'model_name': self.model_name,
|
||||
'doc_count': len(self.documents),
|
||||
'vector_dim': self._dim,
|
||||
'index_type': 'FlatIP',
|
||||
'index_type': self._index_type,
|
||||
'updated_at': datetime.now().isoformat(),
|
||||
}
|
||||
with open(META_FILE, 'w', encoding='utf-8') as f:
|
||||
@@ -280,6 +387,11 @@ class FAISSIndex:
|
||||
with open(META_FILE, 'r', encoding='utf-8') as f:
|
||||
meta = json.load(f)
|
||||
index.model_name = meta.get('model_name', index.model_name)
|
||||
index._index_type = meta.get('index_type', 'FlatIP')
|
||||
# 如果已经是 IVFFlat 或 HNSW，标记为不需要升级
|
||||
if index._index_type in ('IVFFlat', 'HNSW'):
|
||||
index._upgrade_lock = True
|
||||
index._needs_upgrade = False
|
||||
|
||||
return index
|
||||
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Utility modules for zhiyi."""
|
||||
|
||||
from utils.logger import get_logger, LogTimer, StructuredLogger
|
||||
|
||||
__all__ = ["get_logger", "LogTimer", "StructuredLogger"]
|
||||
@@ -0,0 +1,214 @@
|
||||
"""结构化日志工具 — P1-1: 关键路径监控
|
||||
|
||||
关键 log points：
|
||||
- commit（content_hash, agent_id, doc_count）
|
||||
- recall（query, top_k, duration_ms）
|
||||
- redis_event（event_type, instance_id）
|
||||
- error（endpoint, exception）
|
||||
|
||||
输出格式：JSON Lines（生产环境）+ 彩色（开发环境）
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Optional, Any
|
||||
|
||||
|
||||
# 日志级别从环境变量读取
|
||||
_LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO').upper()
|
||||
_FORMATTER = None
|
||||
|
||||
|
||||
class JSONFormatter(logging.Formatter):
|
||||
"""JSON 格式日志（生产环境）"""
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
log_data = {
|
||||
'timestamp': datetime.utcnow().isoformat() + 'Z',
|
||||
'level': record.levelname,
|
||||
'logger': record.name,
|
||||
'message': record.getMessage(),
|
||||
}
|
||||
# 添加 extra 字段
|
||||
if hasattr(record, 'extra'):
|
||||
log_data.update(record.extra)
|
||||
# 添加异常信息
|
||||
if record.exc_info:
|
||||
log_data['exception'] = self.formatException(record.exc_info)
|
||||
return json.dumps(log_data, ensure_ascii=False)
|
||||
|
||||
|
||||
class ColorFormatter(logging.Formatter):
|
||||
"""彩色格式化（开发环境）"""
|
||||
COLORS = {
|
||||
'DEBUG': '\033[36m', # 青色
|
||||
'INFO': '\033[32m', # 绿色
|
||||
'WARNING': '\033[33m', # 黄色
|
||||
'ERROR': '\033[31m', # 红色
|
||||
'CRITICAL': '\033[35m', # 紫色
|
||||
}
|
||||
RESET = '\033[0m'
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
color = self.COLORS.get(record.levelname, '')
|
||||
time_str = datetime.now().strftime('%H:%M:%S')
|
||||
return (
|
||||
f"{color}[{time_str}] {record.levelname:8s} "
|
||||
f"{record.name:20s} {record.getMessage()}{self.RESET}"
|
||||
)
|
||||
|
||||
|
||||
def _get_formatter() -> logging.Formatter:
|
||||
global _FORMATTER
|
||||
if _FORMATTER is None:
|
||||
if os.environ.get('ZHIYI_ENV', 'dev') == 'prod':
|
||||
_FORMATTER = JSONFormatter()
|
||||
else:
|
||||
_FORMATTER = ColorFormatter()
|
||||
return _FORMATTER
|
||||
|
||||
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
"""获取结构化日志器
|
||||
|
||||
使用方法：
|
||||
logger = get_logger(__name__)
|
||||
logger.info("hello", extra={"user_id": 123})
|
||||
|
||||
# 或使用 LogTimer 上下文管理器
|
||||
with LogTimer(logger, "heavy_operation"):
|
||||
do_something()
|
||||
"""
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(getattr(logging, _LOG_LEVEL, logging.INFO))
|
||||
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setFormatter(_get_formatter())
|
||||
logger.addHandler(handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
class LogTimer:
|
||||
"""耗时统计上下文管理器
|
||||
|
||||
用法：
|
||||
logger = get_logger(__name__)
|
||||
with LogTimer(logger, "search", extra={"query": "python"}):
|
||||
results = search()
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
logger: logging.Logger,
|
||||
operation: str,
|
||||
level: int = logging.INFO,
|
||||
extra: Optional[dict] = None,
|
||||
):
|
||||
self.logger = logger
|
||||
self.operation = operation
|
||||
self.level = level
|
||||
self.extra = extra or {}
|
||||
self.start_time: float = 0
|
||||
|
||||
def __enter__(self):
|
||||
self.start_time = time.perf_counter()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
duration_ms = (time.perf_counter() - self.start_time) * 1000
|
||||
if exc_type is not None:
|
||||
self.logger.log(
|
||||
self.level,
|
||||
f"{self.operation} failed after {duration_ms:.1f}ms",
|
||||
extra={"duration_ms": round(duration_ms, 2), "error": str(exc_val), **self.extra}
|
||||
)
|
||||
else:
|
||||
self.logger.log(
|
||||
self.level,
|
||||
f"{self.operation} done in {duration_ms:.1f}ms",
|
||||
extra={"duration_ms": round(duration_ms, 2), **self.extra}
|
||||
)
|
||||
return False # 不吞异常
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────
|
||||
# 关键路径日志辅助函数
|
||||
# ─────────────────────────────────────────────
|
||||
|
||||
def log_commit(logger: logging.Logger, content_hash: str, agent_id: str, doc_count: int,
|
||||
episode_id: Optional[str] = None, duration_ms: float = 0):
|
||||
"""commit 路径日志"""
|
||||
logger.info(
|
||||
f"commit: agent={agent_id} docs={doc_count}",
|
||||
extra={
|
||||
"event": "commit",
|
||||
"content_hash": content_hash,
|
||||
"agent_id": agent_id,
|
||||
"doc_count": doc_count,
|
||||
"episode_id": episode_id,
|
||||
"duration_ms": round(duration_ms, 2),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def log_recall(logger: logging.Logger, query: str, top_k: int, duration_ms: float,
|
||||
results_count: int = 0, use_expand: bool = False):
|
||||
"""recall 路径日志"""
|
||||
# query 太长截断
|
||||
query_preview = query[:50] + '...' if len(query) > 50 else query
|
||||
logger.info(
|
||||
f"recall: query={query_preview!r} top_k={top_k} results={results_count} time={duration_ms:.1f}ms",
|
||||
extra={
|
||||
"event": "recall",
|
||||
"query_preview": query_preview,
|
||||
"top_k": top_k,
|
||||
"results_count": results_count,
|
||||
"duration_ms": round(duration_ms, 2),
|
||||
"use_query_expand": use_expand,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def log_redis_event(logger: logging.Logger, event_type: str, instance_id: str,
|
||||
commit_id: Optional[str] = None, stream_len: int = 0):
|
||||
"""Redis 事件消费日志"""
|
||||
logger.debug(
|
||||
f"redis_event: type={event_type} instance={instance_id}",
|
||||
extra={
|
||||
"event": "redis_event",
|
||||
"event_type": event_type,
|
||||
"instance_id": instance_id,
|
||||
"commit_id": commit_id,
|
||||
"stream_len": stream_len,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def log_search(logger: logging.Logger, query: str, results: int, duration_ms: float,
|
||||
index_type: str = "FlatIP", doc_count: int = 0):
|
||||
"""search 路径日志（低级别）"""
|
||||
logger.debug(
|
||||
f"search: results={results} time={duration_ms:.1f}ms index={index_type} docs={doc_count}",
|
||||
extra={
|
||||
"event": "search",
|
||||
"query_preview": query[:30],
|
||||
"results": results,
|
||||
"duration_ms": round(duration_ms, 2),
|
||||
"index_type": index_type,
|
||||
"doc_count": doc_count,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# 向后兼容别名（已由 get_logger() 替代）
|
||||
class StructuredLogger:
|
||||
"""已废弃，请使用 get_logger()"""
|
||||
def __init__(self, name: str):
|
||||
self._logger = get_logger(name)
|
||||
def info(self, msg: str, **kw): self._logger.info(msg, **kw)
|
||||
def error(self, msg: str, **kw): self._logger.error(msg, **kw)
|
||||
def warning(self, msg: str, **kw): self._logger.warning(msg, **kw)
|
||||
def debug(self, msg: str, **kw): self._logger.debug(msg, **kw)
|
||||
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
查询扩展工具 - 同义词扩展
|
||||
纯规则实现，不依赖 LLM
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
import re
|
||||
|
||||
|
||||
# 中文同义词词典（IT/编程/生活类）- 20+ 组
|
||||
CHINESE_SYNONYMS = {
|
||||
# 硬件设备
|
||||
"电脑": ["计算机", "PC", "主机", "台式机", "笔记本"],
|
||||
"手机": ["智能手机", "移动设备", "平板", "iPhone", "安卓"],
|
||||
|
||||
# 软件/程序
|
||||
"程序": ["软件", "应用", "APP", "应用程序", "程式"],
|
||||
"代码": ["编程", "开发", "写代码", "源码", "script"],
|
||||
"APP": ["应用", "应用程序", "软件", "app"],
|
||||
|
||||
# 存储相关
|
||||
"存储": ["保存", "存档", "备份", "存放"],
|
||||
"文件": ["文档", "资料", "档案", "附件"],
|
||||
"文档": ["文件", "资料", "档案", "doc"],
|
||||
|
||||
# 网络相关
|
||||
"网络": ["互联网", "在线", "联网", "WiFi", "宽带", "internet"],
|
||||
"服务器": ["主机", "云服务器", "VPS", "云主机", "server"],
|
||||
|
||||
# 数据相关
|
||||
"数据": ["信息", "资料", "记录", "日志", "content"],
|
||||
"数据库": ["DB", "数据存储", "资料库", "信息库", "database"],
|
||||
|
||||
# 账户相关
|
||||
"密码": ["口令", "登录密码", "账号密码", "passwd"],
|
||||
"账户": ["账号", "用户名", "登录名", "acount"],
|
||||
"登录": ["登陆", "signin", "login", "认证"],
|
||||
|
||||
# 问题/错误
|
||||
"错误": ["故障", "问题", "Bug", "报错", "异常"],
|
||||
"调试": ["排错", "修复", "解决", "排查", "debug"],
|
||||
|
||||
# 性能相关
|
||||
"性能": ["速度", "效率", "运行速度", "响应速度"],
|
||||
"慢": ["卡顿", "延迟", "反应慢", "龟速"],
|
||||
|
||||
# 安全相关
|
||||
"安全": ["加密", "隐私", "保密", "protection", "security"],
|
||||
"泄露": ["泄漏", "暴露", "外泄"],
|
||||
|
||||
# 搜索相关
|
||||
"搜索": ["查询", "查找", "检索", "搜寻", "寻找"],
|
||||
"忘记": ["遗忘", "丢失", "找不到了"],
|
||||
|
||||
# 启动/运行
|
||||
"启动": ["运行", "开机", "开始", "执行", "run"],
|
||||
"关闭": ["退出", "停止", "关机", "终止", "shutdown"],
|
||||
"重启": ["重新启动", "重启", "reboot"],
|
||||
|
||||
# 更新
|
||||
"更新": ["升级", "新版本", "版本更新", "upgrade"],
|
||||
"安装": ["部署", "配置", "设置", "搭建", "install"],
|
||||
|
||||
# 测试相关
|
||||
"测试": ["测验", "检验", "验证", "test"],
|
||||
"崩溃": ["闪退", "卡死", "宕机", "crash"],
|
||||
|
||||
# 显示/输出
|
||||
"显示": ["展示", "输出", "呈现", "display"],
|
||||
"截图": ["截屏", "屏幕截图", "screen"],
|
||||
|
||||
# 复制/粘贴
|
||||
"复制": ["拷贝", "粘贴", "重复", "copy"],
|
||||
"剪切": ["剪下", "移动", "cut"],
|
||||
|
||||
# 删除
|
||||
"删除": ["移除", "清除", "卸载", "去除", "remove"],
|
||||
"取消": ["撤回", "作废", "无效"],
|
||||
|
||||
# 生活类
|
||||
"吃饭": ["用餐", "进餐", "就餐", "觅食"],
|
||||
"睡觉": ["睡眠", "休息", "就寝", "入眠"],
|
||||
"运动": ["锻炼", "健身", "跑步", "锻炼身体"],
|
||||
"学习": ["读书", "复习", "预习", "学业"],
|
||||
"工作": ["上班", "办公", "干活", "就业"],
|
||||
"旅行": ["旅游", "出游", "度假", "观光"],
|
||||
}
|
||||
|
||||
|
||||
def _clean_query(query: str) -> str:
|
||||
"""清洗查询：去除句首的导向词，还原为核心查询。"""
|
||||
# 去除句首特定模式
|
||||
patterns = [
|
||||
(r'^想起\s*', ''), # 想起xxx -> xxx
|
||||
(r'^关于\s*', ''), # 关于xxx -> xxx
|
||||
(r'^查询\s*', ''), # 查询xxx -> xxx
|
||||
(r'^搜索\s*', ''), # 搜索xxx -> xxx
|
||||
(r'^找一下\s*', ''), # 找一下xxx -> xxx
|
||||
(r'^帮我找\s*', ''), # 帮我找xxx -> xxx
|
||||
(r'^我想知道\s*', ''), # 我想知道xxx -> xxx
|
||||
(r'^问一下\s*', ''), # 问一xxx -> xxx
|
||||
(r'^请问\s*', ''), # 请问xxx -> xxx
|
||||
(r'^有没有\s*', ''), # 有没有xxx -> xxx
|
||||
(r'^还有\s*', ''), # 还有xxx -> xxx
|
||||
(r'^以及\s*', ''), # 以及xxx -> xxx
|
||||
]
|
||||
cleaned = query.strip()
|
||||
for pattern, replacement in patterns:
|
||||
cleaned = re.sub(pattern, replacement, cleaned)
|
||||
return cleaned
|
||||
|
||||
|
||||
def expand_query(query: str) -> List[str]:
|
||||
"""
|
||||
扩展查询，返回扩展后的 query 列表（original + synonyms）。
|
||||
|
||||
Args:
|
||||
query: 原始查询字符串
|
||||
|
||||
Returns:
|
||||
扩展后的 query 列表，包含原始 query 和所有同义词变体
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
return [query]
|
||||
|
||||
original = query.strip()
|
||||
cleaned = _clean_query(original)
|
||||
|
||||
# 构建同义词集合
|
||||
synonym_variants = set()
|
||||
|
||||
# 对 cleaned 进行同义词替换
|
||||
for keyword, synonyms in CHINESE_SYNONYMS.items():
|
||||
if keyword in cleaned:
|
||||
# 添加所有同义词替换
|
||||
for syn in synonyms:
|
||||
variant = cleaned.replace(keyword, syn)
|
||||
synonym_variants.add(variant)
|
||||
|
||||
# 去重：保留原始、清洗后、以及所有变体
|
||||
expanded = [original, cleaned] if original != cleaned else [original]
|
||||
expanded.extend(synonym_variants)
|
||||
|
||||
# 去重并保持顺序
|
||||
seen = set()
|
||||
unique_expanded = []
|
||||
for q in expanded:
|
||||
if q and q not in seen:
|
||||
seen.add(q)
|
||||
unique_expanded.append(q)
|
||||
|
||||
return unique_expanded
|
||||
|
||||
|
||||
def get_top_query(original_query: str, limit: int = 5) -> List[str]:
|
||||
"""
|
||||
获取 top-N 查询变体（用于多查询搜索）。
|
||||
|
||||
策略：
|
||||
- 始终返回原始 query 作为 top-1（详细召回）
|
||||
- 其余按同义词数量优先排序（补充召回）
|
||||
|
||||
Args:
|
||||
original_query: 原始查询
|
||||
limit: 返回数量限制
|
||||
|
||||
Returns:
|
||||
top-N 查询列表
|
||||
"""
|
||||
expanded = expand_query(original_query)
|
||||
|
||||
# 优先保留原始查询
|
||||
result = []
|
||||
if expanded:
|
||||
result.append(expanded[0]) # 原始查询作为 top-1
|
||||
|
||||
# 添加同义词变体（去除重复）
|
||||
for q in expanded[1:]:
|
||||
if len(result) >= limit:
|
||||
break
|
||||
if q not in result:
|
||||
result.append(q)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 测试用例
|
||||
test_queries = [
|
||||
"想起我的密码是什么",
|
||||
"关于Python编程的问题",
|
||||
"搜索一下如何安装Node.js",
|
||||
"电脑太慢了怎么办",
|
||||
"手机APP崩溃了",
|
||||
"数据库连接错误",
|
||||
]
|
||||
|
||||
print("=" * 60)
|
||||
print("Query Expander 测试")
|
||||
print("=" * 60)
|
||||
|
||||
for q in test_queries:
|
||||
expanded = expand_query(q)
|
||||
top_queries = get_top_query(q, limit=5)
|
||||
print(f"\n原始查询: {q}")
|
||||
print(f" 扩展结果 ({len(expanded)}): {expanded[:10]}{'...' if len(expanded) > 10 else ''}")
|
||||
print(f" Top-5 查询: {top_queries}")
|
||||
@@ -0,0 +1,87 @@
|
||||
"""
|
||||
pytest configuration and shared fixtures for zhiyi tests.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add src to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from datetime import datetime
|
||||
from typing import List
|
||||
|
||||
from models.episode import Episode
|
||||
from models.distilled import Distilled, TYPE_FACT
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake embedder (list-based, no real bge-m3 model)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FakeEmbedder:
|
||||
"""Simple list-based fake embedder for fast tests.
|
||||
|
||||
Returns a fixed-size vector of ones. The embedding dimension is
|
||||
configurable (default 4) so tests run fast without loading the real
|
||||
bge-m3 model.
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int = 4):
|
||||
self.dim = dim
|
||||
|
||||
def encode(self, texts: List[str], **kwargs) -> List[List[float]]:
|
||||
"""Return fake vectors – one per input text."""
|
||||
return [[1.0] * self.dim for _ in texts]
|
||||
|
||||
async def encode_async(self, texts: List[str], **kwargs) -> List[List[float]]:
|
||||
"""Async wrapper – delegates to encode()."""
|
||||
return self.encode(texts, **kwargs)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_embedder():
|
||||
"""Provides a FakeEmbedder instance."""
|
||||
return FakeEmbedder(dim=4)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared data fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def sample_episode():
|
||||
"""A basic Episode instance for use across tests."""
|
||||
return Episode.create(
|
||||
content='Alice met Bob at the cafe yesterday.',
|
||||
source='test',
|
||||
entities=['Alice', 'Bob'],
|
||||
facts=['Alice met Bob at a cafe'],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_distilled():
|
||||
"""A basic Distilled instance for use across tests."""
|
||||
ep = Episode.create(content='Test content')
|
||||
return Distilled(
|
||||
id='distilled-1',
|
||||
episode_id=ep.id,
|
||||
type=TYPE_FACT,
|
||||
summary='Test summary of the distilled content.',
|
||||
importance=0.5,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async event loop fixture
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def async_event_loop():
|
||||
"""Provides a fresh event loop for each async test."""
|
||||
import asyncio
|
||||
loop = asyncio.new_event_loop()
|
||||
yield loop
|
||||
loop.close()
|
||||
@@ -14,11 +14,11 @@ def test_jaccard_similarity():
|
||||
def test_merge_or_flag():
|
||||
cp = ConsolidationPipeline()
|
||||
r1 = Distilled(id='r1', episode_id='e1', type='fact', summary='s1',
|
||||
facts=['python', 'docker', 'linux'], confidence=0.8)
|
||||
facts=['python', 'docker', 'linux'], importance=0.8)
|
||||
r2 = Distilled(id='r2', episode_id='e2', type='fact', summary='s2',
|
||||
facts=['python', 'docker', 'linux', 'git'], confidence=0.7)
|
||||
facts=['python', 'docker', 'linux', 'git'], importance=0.7)
|
||||
r3 = Distilled(id='r3', episode_id='e3', type='fact', summary='s3',
|
||||
facts=['different', 'facts'], confidence=0.6)
|
||||
facts=['different', 'facts'], importance=0.6)
|
||||
assert cp.merge_or_flag(r1, r2) == 'merged'
|
||||
assert cp.merge_or_flag(r1, r3) == 'flagged'
|
||||
|
||||
@@ -26,7 +26,7 @@ def test_consolidation_run():
|
||||
cp = ConsolidationPipeline()
|
||||
cp.set_existing_records([])
|
||||
r = Distilled(id='new1', episode_id='e1', type='fact', summary='s',
|
||||
facts=['fact1', 'fact2'], confidence=0.8)
|
||||
facts=['fact1', 'fact2'], importance=0.8)
|
||||
actions = cp.run(r)
|
||||
assert len(actions) > 0
|
||||
assert actions[0][0] in ['merge_or_flag', 'update_graph']
|
||||
|
||||
+1
-3
@@ -88,8 +88,6 @@ class TestGraphBuilder:
|
||||
"牧尘用Arch Linux"
|
||||
],
|
||||
entities=["牧尘"],
|
||||
confidence=0.8,
|
||||
status="active",
|
||||
importance=2
|
||||
)
|
||||
|
||||
@@ -109,7 +107,7 @@ class TestGraphBuilder:
|
||||
summary="",
|
||||
facts=["牧尘使用ComfyUI"],
|
||||
entities=["牧尘", "ComfyUI"],
|
||||
confidence=0.7
|
||||
importance=0.7
|
||||
)
|
||||
|
||||
builder = GraphBuilder()
|
||||
|
||||
@@ -3,7 +3,7 @@ from pathlib import Path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
||||
from datetime import datetime
|
||||
from models.episode import Episode
|
||||
from models.distilled import Distilled, TYPE_FACT, STATUS_PENDING, STATUS_VALIDATED
|
||||
from models.distilled import Distilled, TYPE_FACT
|
||||
|
||||
def test_episode_to_dict_from_dict():
|
||||
ep = Episode.create(content='test content', source='test', entities=['entity1'], facts=['fact1'])
|
||||
@@ -25,20 +25,19 @@ def test_distilled_to_dict_from_dict():
|
||||
ep = Episode.create(content='test')
|
||||
d = Distilled(
|
||||
id='d1', episode_id=ep.id, type=TYPE_FACT,
|
||||
summary='a summary', confidence=0.8, status=STATUS_PENDING
|
||||
summary='a summary', importance=0.8
|
||||
)
|
||||
dd = d.to_dict()
|
||||
assert dd['type'] == TYPE_FACT
|
||||
assert dd['confidence'] == 0.8
|
||||
assert dd['importance'] == 0.8
|
||||
d2 = Distilled.from_dict(dd)
|
||||
assert d2.id == d.id
|
||||
assert d2.type == d.type
|
||||
|
||||
def test_distilled_defaults():
|
||||
d = Distilled(id='d1', episode_id='e1', type=TYPE_FACT, summary='s')
|
||||
assert d.status == STATUS_PENDING
|
||||
assert d.confidence == 0.5
|
||||
assert d.importance == 0
|
||||
assert d.importance == 0.0
|
||||
assert d.tier == 'normal'
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_episode_to_dict_from_dict()
|
||||
|
||||
Reference in New Issue
Block a user