zqf/memory-system
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

初始化：记忆系统源代码上传（已脱敏）

Browse Source 
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html）
- config.py 默认密码已替换为占位符 CHANGE_ME_*
- init_db.sql 移除生产数据库用户 GRANT 段
- README.html 数据库用户名已脱敏
- 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
This commit is contained in:
zqf
2026-08-03 04:53:48 +08:00
commit b1f93588f1
30 changed files with 5591 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+19
View File
@@ -0,0 +1,19 @@
# Environment & secrets
.env
.env.*
*.bak
*.bak.*
# Python
__pycache__/
*.py[cod]
venv/
.venv/
# IDE
.idea/
.vscode/
# Runtime
*.log
*.pid
README.html
+353
View File
@@ -0,0 +1,353 @@
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>多智能体记忆系统 — 内部运维文档</title>
<style>
:root { --bg: #0d1117; --card: #161b22; --border: #30363d; --text: #c9d1d9; --heading: #58a6ff; --accent: #3fb950; --warn: #d29922; --code-bg: #0d1117; }
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: var(--bg); color: var(--text); line-height: 1.7; padding: 2rem; max-width: 960px; margin: 0 auto; }
h1 { color: var(--heading); margin-bottom: 0.5rem; font-size: 1.8rem; }
h2 { color: var(--heading); margin: 2.5rem 0 0.5rem; border-bottom: 1px solid var(--border); padding-bottom: 0.3rem; font-size: 1.3rem; }
h3 { color: var(--accent); margin: 1.5rem 0 0.3rem; font-size: 1.1rem; }
.meta { color: #8b949e; margin-bottom: 2rem; font-size: 0.9rem; }
table { width: 100%; border-collapse: collapse; margin: 1rem 0; }
th, td { text-align: left; padding: 0.5rem 0.75rem; border: 1px solid var(--border); }
th { background: var(--card); color: var(--heading); font-weight: 600; }
td { background: var(--card); }
code { font-family: 'Fira Code', monospace; background: var(--code-bg); border: 1px solid var(--border); padding: 0.15rem 0.4rem; border-radius: 4px; font-size: 0.85rem; color: var(--accent); }
pre { background: var(--code-bg); border: 1px solid var(--border); border-radius: 6px; padding: 1rem; overflow-x: auto; margin: 0.75rem 0; font-size: 0.85rem; line-height: 1.5; }
pre code { border: none; padding: 0; background: none; color: var(--text); }
.method { display: inline-block; padding: 0.15rem 0.5rem; border-radius: 4px; font-weight: bold; font-size: 0.8rem; margin-right: 0.5rem; }
.get { background: #1f6feb33; color: #58a6ff; }
.post { background: #3fb95033; color: #3fb950; }
.put { background: #d2992233; color: #d29922; }
.delete { background: #f8514933; color: #f85149; }
.endpoint { font-family: monospace; color: var(--text); }
.warn-box { background: #d2992222; border: 1px solid var(--warn); border-radius: 6px; padding: 1rem; margin: 1rem 0; }
.ok-box { background: #3fb95022; border: 1px solid var(--accent); border-radius: 6px; padding: 1rem; margin: 1rem 0; }
.danger-box { background: #f8514922; border: 1px solid #f85149; border-radius: 6px; padding: 1rem; margin: 1rem 0; }
ul, ol { margin: 0.5rem 0 0.5rem 1.5rem; }
li { margin: 0.25rem 0; }
a { color: var(--heading); }
.tag { display: inline-block; padding: 0.1rem 0.4rem; border-radius: 3px; font-size: 0.75rem; margin-left: 0.3rem; }
.tag-new { background: #3fb95033; color: #3fb950; }
.tag-auth { background: #d2992233; color: #d29922; }
</style>
</head>
<body>
<h1>🧠 多智能体记忆系统 — 内部运维文档</h1>
<p class="meta">
服务器: 43.156.57.141 (新加坡腾讯云) &nbsp;|&nbsp;
域名: <a href="https://memory.lsz.name">memory.lsz.name</a> &nbsp;|&nbsp;
版本: 3.0.2 &nbsp;|&nbsp;
最后更新: 2026-06-22
</p>
<div class="ok-box">
<strong>✅ 系统状态:</strong> 所有组件运行正常<br>
MySQL 8.0.46 · Redis 6.0.16 · BGE-small-zh-v1.5 (512维) · Gunicorn + Flask · Nginx + SSL (Let's Encrypt)<br>
2 个团队 · 4 个 Agent · 32+ 个人记忆 · 4 条团队记忆 · v3.0.2 性能优化（LLM thinking mode 禁用）
</div>
<h2>🏗️ 架构概览</h2>
<table>
<tr><th>组件</th><th>技术</th><th>端口</th><th>说明</th></tr>
<tr><td>API 服务</td><td>Flask + Gunicorn</td><td>8081 (local)</td><td>HTTP JSON API</td></tr>
<tr><td>Nginx</td><td>nginx 1.18</td><td>443 (SSL)</td><td>反向代理 + SSL 终止</td></tr>
<tr><td>持久存储</td><td>MySQL 8.0.46</td><td>3306</td><td>记忆数据 + 向量 BLOB</td></tr>
<tr><td>缓存层</td><td>Redis 6.0.16</td><td>6379</td><td>工作记忆 (TTL 24h)</td></tr>
<tr><td>Embedding</td><td>BGE-small-zh-v1.5</td><td>8080 (local)</td><td>512维语义向量</td></tr>
</table>
<h2>🔐 API 认证 <span class="tag tag-auth">HMAC-SHA256</span></h2>
<p>除 <code>GET /health</code> 和 <code>/admin/*</code> 外，所有 API 端点均需以下 HTTP 头：</p>
<table>
<tr><th>Header</th><th>说明</th></tr>
<tr><td><code>X-API-Key</code></td><td>API Key（以 <code>msk_</code> 前缀开头）</td></tr>
<tr><td><code>X-Timestamp</code></td><td>Unix 时间戳（秒，5分钟内有效）</td></tr>
<tr><td><code>X-Signature</code></td><td><code>HMAC-SHA256(secret, timestamp + request_body)</code></td></tr>
</table>
<h3>认证流程</h3>
<ol>
<li>获取 API Key 和 Secret（通过 <code>/admin/api-keys</code> 端点，仅限 localhost）</li>
<li>构造请求体（JSON），取当前 Unix 时间戳</li>
<li>计算签名：<code>signature = HMAC-SHA256(secret, str(timestamp) + json_body)</code></li>
<li>在请求头中附加三个认证字段</li>
</ol>
<h3>Python 认证示例</h3>
<pre><code>import time, hmac, hashlib, json, requests
API_KEY = "msk_xxx..."
SECRET = "xxx..."
BASE = "https://memory.lsz.name"
def api(method, path, body=None):
ts = str(time.time())
body_bytes = json.dumps(body).encode() if body else b""
sig = hmac.new(
SECRET.encode(), ts.encode() + body_bytes, hashlib.sha256
).hexdigest()
headers = {
"Content-Type": "application/json",
"X-API-Key": API_KEY,
"X-Timestamp": ts,
"X-Signature": sig,
}
resp = requests.request(method, BASE + path, data=body_bytes, headers=headers)
return resp.json()
# 使用示例
result = api("POST", "/memories/personal/search", {
"agent_id": "yihui",
"query": "策略回测",
"limit": 5
})</code></pre>
<div class="warn-box">
<strong>⚠️ 注意:</strong> Secret 仅在创建时显示一次，不会再次展示。请妥善保管。<br>
Admin 端点仅限 <code>127.0.0.1</code> 访问，外部请求会被拒绝（403）。
</div>
<h2>👥 多团队隔离</h2>
<p>所有记忆数据通过 <code>team_id</code> 字段实现硬隔离。查询时强制 <code>WHERE team_id = ?</code>。每个团队的 API Key 绑定固定的 <code>team_id</code>，跨团队数据完全不可见。</p>
<h3>当前团队</h3>
<table>
<tr><th>Team ID</th><th>名称</th><th>Agent 数</th><th>API Key</th></tr>
<tr><td><code>hermes1</code></td><td>策略研究部</td><td>2 (弈回, 云策)</td><td><code>msk_a192...85</code></td></tr>
<tr><td><code>wangcheng</code></td><td>王成</td><td>2 (Hermes Agent, OpenClaw Main)</td><td><code>msk_6804...18</code></td></tr>
</table>
<h2>📋 记忆类型</h2>
<table>
<tr><th>类型</th><th>存储</th><th>生命周期</th><th>说明</th></tr>
<tr><td>个人长期记忆</td><td>MySQL</td><td>持久化</td><td>每个 Agent 独立的持久记忆，支持语义检索</td></tr>
<tr><td>个人工作记忆</td><td>Redis List</td><td>TTL 24h</td><td>短期上下文缓存</td></tr>
<tr><td>团队共享记忆</td><td>MySQL</td><td>持久化</td><td>同团队内所有 Agent 可访问的知识库</td></tr>
</table>
<h2>📡 API 端点</h2>
<h3>公开端点（无需认证）</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/health</td><td>健康检查</td></tr>
</table>
<h3>团队管理</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/teams</td><td>创建团队</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/teams/{team_id}</td><td>获取团队信息</td></tr>
<tr><td><span class="method delete">DELETE</span></td><td class="endpoint">/teams/{team_id}</td><td>删除团队（级联删除所有关联数据）</td></tr>
</table>
<h3>Agent 管理</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/agents</td><td>创建 Agent</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/agents/{agent_id}</td><td>获取 Agent 信息</td></tr>
</table>
<h3>个人记忆</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/memories/personal</td><td>添加个人记忆（自动生成 embedding）</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/memories/personal/search</td><td>语义检索个人记忆（团队隔离）</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/memories/personal/recent/{agent_id}</td><td>获取最近个人记忆</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/memories/personal/{memory_id}</td><td>获取单条个人记忆</td></tr>
<tr><td><span class="method put">PUT</span></td><td class="endpoint">/memories/personal/{memory_id}</td><td>更新个人记忆</td></tr>
<tr><td><span class="method delete">DELETE</span></td><td class="endpoint">/memories/personal/{memory_id}</td><td>删除个人记忆</td></tr>
</table>
<h3>工作记忆</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/memories/working</td><td>添加工作记忆（Redis，TTL 24h）</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/memories/working/{agent_id}</td><td>获取工作记忆</td></tr>
<tr><td><span class="method delete">DELETE</span></td><td class="endpoint">/memories/working/{agent_id}</td><td>清空工作记忆</td></tr>
</table>
<h3>团队记忆</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/memories/team</td><td>添加团队共享记忆</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/memories/team/search</td><td>语义检索团队记忆</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/memories/team/recent/{team_id}</td><td>获取最近团队记忆</td></tr>
<tr><td><span class="method put">PUT</span></td><td class="endpoint">/memories/team/{memory_id}</td><td>更新团队记忆</td></tr>
<tr><td><span class="method delete">DELETE</span></td><td class="endpoint">/memories/team/{memory_id}</td><td>删除团队记忆</td></tr>
</table>
<h3>生命周期管理</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/compress</td><td>压缩工作记忆到长期记忆</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/cleanup</td><td>清理过期/低重要性记忆</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/extract-facts</td><td>原子事实提取（LLM, 5-15s）</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/aggregate-scenes</td><td>场景聚合（LLM, 5-15s）</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/generate-persona</td><td>画像生成（LLM, 5-15s）</td></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/lifecycle/auto-config</td><td>Pipeline 自动化配置</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/lifecycle/scenarios/{agent_id}</td><td>获取场景块</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/lifecycle/persona/{agent_id}</td><td>获取用户画像</td></tr>
</table>
<h3>统计</h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/stats?team_id=xxx</td><td>获取系统统计信息</td></tr>
</table>
<h3>Admin 端点 <span class="tag tag-new">仅 localhost</span></h3>
<table>
<tr><th>方法</th><th>路径</th><th>说明</th></tr>
<tr><td><span class="method post">POST</span></td><td class="endpoint">/admin/api-keys</td><td>创建 API Key（需 team_id, name）</td></tr>
<tr><td><span class="method get">GET</span></td><td class="endpoint">/admin/api-keys</td><td>列出所有 API Key</td></tr>
<tr><td><span class="method delete">DELETE</span></td><td class="endpoint">/admin/api-keys/{key_id}</td><td>禁用 API Key</td></tr>
</table>
<h2>⚡ 性能指标（2026-06-08 验证）</h2>
<table>
<tr><th>指标</th><th>实测值</th><th>目标</th><th>状态</th></tr>
<tr><td>写入延迟 P50</td><td>14.4ms</td><td>—</td><td>✅</td></tr>
<tr><td>写入延迟 P99</td><td>16.3ms</td><td>&lt; 100ms</td><td>✅ PASS</td></tr>
<tr><td>检索延迟 P50</td><td>17.9ms</td><td>—</td><td>✅</td></tr>
<tr><td>检索延迟 P99</td><td>27.5ms</td><td>&lt; 50ms</td><td>✅ PASS</td></tr>
<tr><td>进程 RSS</td><td>51 MB</td><td>&lt; 1.5 GiB</td><td>✅ PASS</td></tr>
<tr><td>团队隔离</td><td>100%</td><td>100%</td><td>✅ PASS</td></tr>
</table>
<h2>🗄️ 数据库表结构</h2>
<table>
<tr><th>表名</th><th>说明</th></tr>
<tr><td><code>teams</code></td><td>团队信息（id, name, description, config JSON）</td></tr>
<tr><td><code>agents</code></td><td>Agent 信息（id, team_id FK, name, role）</td></tr>
<tr><td><code>personal_memories</code></td><td>个人长期记忆（agent_id, team_id, content, embedding BLOB, importance, metadata JSON）</td></tr>
<tr><td><code>team_memories</code></td><td>团队共享记忆（team_id, content, embedding BLOB, importance, category, metadata JSON）</td></tr>
<tr><td><code>api_keys</code></td><td>API 认证密钥（team_id, api_key, secret, is_active）</td></tr>
<tr><td><code>pipeline_config</code></td><td>Pipeline 自动化配置（compress_every_n, cleanup_idle_days, warmup 等）</td></tr>
<tr><td><code>memory_scenarios</code></td><td>场景块（team_id, agent_id, name, summary, memory_ids JSON）</td></tr>
<tr><td><code>user_personas</code></td><td>用户画像（team_id, agent_id, preferences/habits/expertise JSON, summary）</td></tr>
</table>
<h2>🔧 运维操作</h2>
<h3>服务管理</h3>
<pre><code># 查看服务状态
sudo systemctl status memory-system
# 重启服务
sudo systemctl restart memory-system
# 查看日志
sudo journalctl -u memory-system -f --no-pager -n 100
# Nginx 配置测试
sudo nginx -t
# 重载 Nginx
sudo systemctl reload nginx</code></pre>
<h3>项目路径</h3>
<pre><code>/opt/memory-system/ # 项目根目录
/opt/memory-system/venv/ # Python 虚拟环境
/opt/memory-system/config.py # 配置文件
/opt/memory-system/app.py # Flask 入口 + JSON 序列化
/opt/memory-system/api/routes.py # API 路由
/opt/memory-system/auth/api_auth.py # HMAC 认证中间件
/opt/memory-system/storage/ # MySQL + Redis + 向量检索
/opt/memory-system/lifecycle/ # 记忆生命周期管理
/opt/memory-system/tests/ # 测试
/opt/memory-system/docs/ # 公网文档</code></pre>
<h3>关键配置</h3>
<pre><code># config.py
MYSQL_HOST = "127.0.0.1"
MYSQL_PORT = 3306
MYSQL_USER = "your_db_user"
MYSQL_DATABASE = "memory_system"
MYSQL_UNIX_SOCKET = "/var/run/mysqld/mysqld.sock"
REDIS_HOST = "127.0.0.1"
REDIS_PORT = 6379
EMBEDDING_SERVICE_URL = "http://127.0.0.1:8080"
EMBEDDING_DIM = 512
API_PORT = 8081
WORKING_MEMORY_TTL = 86400 # 24小时</code></pre>
<h3>创建新 API Key（SSH 到服务器）</h3>
<pre><code># 通过 localhost 访问 admin 端点
curl -X POST http://127.0.0.1:8081/admin/api-keys \
-H 'Content-Type: application/json' \
-d '{"team_id": "your_team_id", "name": "Your Key Name"}'
# 响应中会包含 api_key 和 secret（仅显示一次）</code></pre>
<h3>SSL 证书</h3>
<pre><code># 自动续期已配置，手动测试:
sudo certbot renew --dry-run
# 证书路径:
/etc/letsencrypt/live/memory.lsz.name/fullchain.pem
/etc/letsencrypt/live/memory.lsz.name/privkey.pem</code></pre>
<h2>🔒 安全说明</h2>
<ul>
<li>Redis 已禁用 FLUSHALL/FLUSHDB/DEBUG/SHUTDOWN 命令</li>
<li>Redis CONFIG/KEYS 命令已重命名（CONFIG_b9f7e2d4, KEYS_a3c8f1e6）</li>
<li>MySQL 使用 unix socket 连接</li>
<li>Nginx 强制 HTTPS (Let's Encrypt)</li>
<li>API 仅监听 127.0.0.1，通过 Nginx 暴露公网</li>
<li>Admin 端点仅限 localhost 访问</li>
<li>API 使用 HMAC-SHA256 签名认证，时间戳窗口 5 分钟</li>
</ul>
<h2>🧪 测试验证</h2>
<pre><code># 运行测试
cd /opt/memory-system
./venv/bin/python tests/test_memory.py</code>
<p>测试覆盖: 创建团队/Agent → 写入记忆 → 语义检索 → 跨团队隔离 → 工作记忆 → 性能基准</p>
<div class="ok-box">
<strong>✅ 2026-06-08 验证结果：</strong><br>
14/14 项测试全部通过，包括：<br>
• 健康检查 · 统计端点 · 团队信息查询<br>
• 个人记忆语义检索 · 团队记忆语义检索<br>
• 写入+搜索+更新+删除全流程<br>
• 跨团队隔离验证（wangcheng 看不到 hermes1 数据）<br>
• 工作记忆 CRUD · 认证拒绝测试
</div>
<h2>📋 变更日志</h2>
<h3>v3.0.2 (2026-06-22) — 性能优化</h3>
<div class="ok-box">
<strong>核心修复</strong>：禁用 DeepSeek thinking mode<br>
根因：DeepSeek v4-flash 默认开启思考模式，每次 LLM 调用输出 3000-7500 字符中文推理链，token 预算被吃光<br>
修复：llm_client.py 新增 thinking 参数透传；llm_parse.py 所有 JSON 生成任务添加 thinking={"type": "disabled"}<br>
<strong>实测提速</strong>：aggregate-scenes 7.9x, generate-persona 4.9x, extract-facts 2.0s
</div>
<h3>v3.0.1-bugfix (2026-06-22) — Bug 修复</h3>
<div class="ok-box">
• BGE 长文本双保险截断（服务端 512 tokens + 客户端 2000 字符）<br>
• 空 query fallback 到 get_recent<br>
• embedding 失败返回友好错误，不暴露内部 URL<br>
• GET /memories/personal/{id} 405 → 补上 GET 路由<br>
• aggregate-scenes JSON 解析失败 → thinking mode 禁用
</div>
<h3>v3.0.0 (2026-06-08) — 生命周期管理</h3>
<div class="ok-box">
• Pipeline 自动化引擎（自动压缩、清理、暖机）<br>
• 原子事实提取（/lifecycle/extract-facts）<br>
• 场景聚合（/lifecycle/aggregate-scenes）<br>
• 用户画像生成（/lifecycle/generate-persona）
</div>
</body>
</html>
api/__init__.py
+1
View File
@@ -0,0 +1 @@
# API module
api/routes.py
+489
View File
@@ -0,0 +1,489 @@
"""Flask API Routes for Memory System"""
import logging
from flask import Blueprint, request, current_app, g
logger = logging.getLogger(__name__)
api = Blueprint("api", __name__)
def _get_ms():
"""Get MemorySystem instance from app context."""
return current_app.memory_system
def _ok(data=None, status=200):
return current_app.json_response({"ok": True, "data": data}, status=status)
def _err(msg, status=400, code=None):
body = {"ok": False, "error": msg}
if code:
body["code"] = code
return current_app.json_response(body, status=status)
# ── Health ───────────────────────────────────────────────────────
@api.route("/health", methods=["GET"])
def health():
return _ok({"status": "ok"})
# ── Team management ─────────────────────────────────────────────
@api.route("/teams", methods=["POST"])
def create_team():
d = request.json or {}
team_id = d.get("team_id")
name = d.get("name")
if not team_id or not name:
return _err("team_id and name required", code="VALIDATION_ERROR")
try:
team = _get_ms().create_team(team_id, name, d.get("description", ""), d.get("config"))
return _ok(team)
except Exception as e:
return _err(str(e))
@api.route("/teams/<team_id>", methods=["GET"])
def get_team(team_id):
team = _get_ms().get_team(team_id)
if not team:
return _err("Team not found", 404, code="TEAM_NOT_FOUND")
return _ok(team)
@api.route("/teams/<team_id>", methods=["DELETE"])
def delete_team(team_id):
_get_ms().delete_team(team_id)
return _ok()
# ── Agent management ────────────────────────────────────────────
@api.route("/agents", methods=["POST"])
def create_agent():
d = request.json or {}
agent_id = d.get("agent_id")
name = d.get("name")
if not agent_id or not name:
return _err("agent_id and name required", code="VALIDATION_ERROR")
try:
agent = _get_ms().create_agent(agent_id, g.team_id, name, d.get("role", ""))
return _ok(agent)
except Exception as e:
return _err(str(e))
@api.route("/agents/<team_id>", methods=["GET"])
def list_agents(team_id):
if team_id != g.team_id:
return _err("Access denied: cannot view other team's agents", 403, code="FORBIDDEN_CROSS_TEAM")
agents = _get_ms().get_agents_by_team(team_id)
return _ok(agents)
@api.route("/agents/<team_id>/<agent_id>", methods=["GET"])
def get_agent(team_id, agent_id):
if team_id != g.team_id:
return _err("Access denied: cannot view other team's agents", 403, code="FORBIDDEN_CROSS_TEAM")
agent = _get_ms().get_agent(agent_id)
if not agent:
return _err("Agent not found", 404, code="AGENT_NOT_FOUND")
return _ok(agent)
@api.route("/agents/<team_id>/<agent_id>", methods=["DELETE"])
def delete_agent(team_id, agent_id):
if team_id != g.team_id:
return _err("Access denied: cannot delete other team's agents", 403, code="FORBIDDEN_CROSS_TEAM")
_get_ms().delete_agent(agent_id)
return _ok()
# ── Personal memories ───────────────────────────────────────────
@api.route("/memories/personal", methods=["POST"])
def add_personal_memory():
d = request.json or {}
agent_id = d.get("agent_id")
content = d.get("content")
if not agent_id or not content:
return _err("agent_id and content required", code="VALIDATION_ERROR")
try:
mem = _get_ms().add_personal_memory(
agent_id, content,
importance=d.get("importance", 0.5),
metadata=d.get("metadata"),
enable_dedup=d.get("enable_dedup", True),
dedup_threshold=d.get("dedup_threshold", 0.85),
)
status = 200 if mem.get("dedup_skipped") else 201
return _ok(mem, status)
except Exception as e:
return _err(str(e))
@api.route("/memories/personal/search", methods=["POST"])
def search_personal_memories():
d = request.json or {}
agent_id = d.get("agent_id")
query = d.get("query")
if not agent_id:
return _err("agent_id required")
# Empty query: fallback to recent memories
if not query:
limit = d.get("limit", 10)
results = _get_ms().get_recent_personal_memories(agent_id, limit)
return _ok(results)
try:
results = _get_ms().search_personal_memories(
agent_id, query,
limit=d.get("limit", 10),
min_score=d.get("min_score", 0.3),
max_chars_per_memory=d.get("max_chars_per_memory", 0),
max_total_chars=d.get("max_total_chars", 0),
)
return _ok(results)
except Exception as e:
return _err(str(e))
@api.route("/memories/personal/recent/<agent_id>", methods=["GET"])
def get_recent_personal_memories(agent_id):
limit = request.args.get("limit", 20, type=int)
results = _get_ms().get_recent_personal_memories(agent_id, limit)
return _ok(results)
@api.route("/memories/personal/<memory_id>", methods=["GET"])
def get_personal_memory(memory_id):
mem = _get_ms().get_personal_memory(memory_id)
if not mem:
return _err("Memory not found", 404, code="MEMORY_NOT_FOUND")
return _ok(mem)
@api.route("/memories/personal/<memory_id>", methods=["PUT"])
def update_personal_memory(memory_id):
d = request.json or {}
try:
ok = _get_ms().update_personal_memory(
memory_id,
content=d.get("content"),
importance=d.get("importance"),
metadata=d.get("metadata"),
)
if ok:
mem = _get_ms().get_personal_memory(memory_id)
return _ok(mem)
return _err("Memory not found", 404, code="MEMORY_NOT_FOUND")
except Exception as e:
return _err(str(e))
@api.route("/memories/personal/<memory_id>", methods=["DELETE"])
def delete_personal_memory(memory_id):
_get_ms().delete_personal_memory(memory_id)
return _ok()
# ── Working memory ──────────────────────────────────────────────
@api.route("/memories/working", methods=["POST"])
def add_working_memory():
d = request.json or {}
agent_id = d.get("agent_id")
content = d.get("content")
if not agent_id or not content:
return _err("agent_id and content required", code="VALIDATION_ERROR")
item = _get_ms().add_working_memory(agent_id, content, ttl=d.get("ttl"))
compress_result = _get_ms().check_auto_compress(agent_id)
if compress_result:
item["auto_compressed"] = compress_result
return _ok(item, 201)
@api.route("/memories/working/<agent_id>", methods=["GET"])
def get_working_memories(agent_id):
limit = request.args.get("limit", 20, type=int)
results = _get_ms().get_working_memories(agent_id, limit)
return _ok(results)
@api.route("/memories/working/<agent_id>", methods=["DELETE"])
def clear_working_memory(agent_id):
_get_ms().clear_working_memory(agent_id)
return _ok()
# ── Team memories ───────────────────────────────────────────────
@api.route("/memories/team", methods=["POST"])
def add_team_memory():
d = request.json or {}
content = d.get("content")
if not content:
return _err("content required", code="VALIDATION_ERROR")
try:
mem = _get_ms().add_team_memory(
g.team_id, content,
importance=d.get("importance", 0.5),
category=d.get("category", "general"),
metadata=d.get("metadata"),
enable_dedup=d.get("enable_dedup", True),
dedup_threshold=d.get("dedup_threshold", 0.85),
)
status = 200 if mem.get("dedup_skipped") else 201
return _ok(mem, status)
except Exception as e:
return _err(str(e))
@api.route("/memories/team/search", methods=["POST"])
def search_team_memories():
d = request.json or {}
query = d.get("query")
# Empty query: fallback to recent memories
if not query:
limit = d.get("limit", 10)
results = _get_ms().get_recent_team_memories(g.team_id, limit)
return _ok(results)
try:
results = _get_ms().search_team_memories(
g.team_id, query,
limit=d.get("limit", 10),
min_score=d.get("min_score", 0.3),
max_chars_per_memory=d.get("max_chars_per_memory", 0),
max_total_chars=d.get("max_total_chars", 0),
)
return _ok(results)
except Exception as e:
return _err(str(e))
@api.route("/memories/team/recent", methods=["GET"])
def get_recent_team_memories():
limit = request.args.get("limit", 20, type=int)
results = _get_ms().get_recent_team_memories(g.team_id, limit)
return _ok(results)
@api.route("/memories/team/<memory_id>", methods=["GET"])
def get_team_memory(memory_id):
mem = _get_ms().get_team_memory(memory_id)
if not mem:
return _err("Memory not found", 404, code="MEMORY_NOT_FOUND")
return _ok(mem)
@api.route("/memories/team/<memory_id>", methods=["PUT"])
def update_team_memory(memory_id):
d = request.json or {}
try:
ok = _get_ms().update_team_memory(
memory_id,
content=d.get("content"),
importance=d.get("importance"),
)
if ok:
mem = _get_ms().get_team_memory(memory_id)
return _ok(mem)
return _err("Memory not found", 404, code="MEMORY_NOT_FOUND")
except Exception as e:
return _err(str(e))
@api.route("/memories/team/<memory_id>", methods=["DELETE"])
def delete_team_memory(memory_id):
_get_ms().delete_team_memory(memory_id)
return _ok()
# ── Lifecycle ───────────────────────────────────────────────────
@api.route("/lifecycle/compress", methods=["POST"])
def compress_working_memories():
d = request.json or {}
agent_id = d.get("agent_id")
if not agent_id:
return _err("agent_id required", code="VALIDATION_ERROR")
try:
result = _get_ms().compress_working_memories(
agent_id, max_items=d.get("target_count", 5),
)
return _ok(result)
except Exception as e:
return _err(str(e))
@api.route("/lifecycle/cleanup", methods=["POST"])
def cleanup_memories():
d = request.json or {}
result = _get_ms().cleanup_memories(
team_id=d.get("team_id"),
max_age_days=d.get("max_age_days", 90),
min_importance=d.get("min_importance", 0.2),
)
return _ok(result)
# ── Stats ───────────────────────────────────────────────────────
@api.route("/stats", methods=["GET"])
def get_stats():
team_id = request.args.get("team_id")
if team_id and team_id != g.team_id:
return _err("Access denied: cannot view other team's stats", 403, code="FORBIDDEN_CROSS_TEAM")
stats = _get_ms().get_stats(g.team_id)
return _ok(stats)
# -- Pipeline Automation --
@api.route("/lifecycle/auto-config", methods=["GET"])
def get_pipeline_config():
config = _get_ms().get_pipeline_config(g.team_id)
if config:
return _ok(config.to_dict())
return _ok({"team_id": g.team_id, "compress_every_n": 0, "cleanup_idle_days": 0, "enabled": False})
@api.route("/lifecycle/auto-config", methods=["POST"])
def set_pipeline_config():
d = request.json or {}
try:
config = _get_ms().set_pipeline_config(
team_id=g.team_id,
compress_every_n=d.get("compress_every_n"),
compress_target_count=d.get("compress_target_count"),
cleanup_idle_days=d.get("cleanup_idle_days"),
cleanup_min_importance=d.get("cleanup_min_importance"),
enabled=d.get("enabled"),
warmup_max_memories=d.get("warmup_max_memories", 0),
warmup_compress_every_n=d.get("warmup_compress_every_n", 1),
)
return _ok(config.to_dict())
except Exception as e:
return _err(str(e))
@api.route("/lifecycle/auto-run", methods=["POST"])
def run_pipeline_cleanup():
d = request.json or {}
team_id = d.get("team_id", g.team_id)
result = _get_ms().run_auto_cleanup(team_id)
if result:
return _ok(result)
return _ok({"message": "No cleanup rules configured"})
# -- Scenario Aggregation --
@api.route("/lifecycle/aggregate-scenes", methods=["POST"])
def aggregate_scenes():
"""Aggregate personal memories into scenario blocks using LLM.
Optional body: {"max_memories": 30}
"""
d = request.json or {}
agent_id = d.get("agent_id", "")
if not agent_id:
return _err("agent_id required", code="VALIDATION_ERROR")
try:
result = _get_ms().aggregate_scenarios(
agent_id=agent_id,
team_id=g.team_id,
max_memories=d.get("max_memories", 50),
)
except Exception as e:
logger.exception("aggregate-scenes failed")
return _err(str(e), 500, code="INTERNAL_ERROR")
if "error" in result:
return _err(result["error"])
return _ok(result)
@api.route("/lifecycle/scenarios/<agent_id>", methods=["GET"])
def get_scenes(agent_id):
"""Get stored scenarios for an agent."""
scenarios = _get_ms().get_scenarios(agent_id)
return _ok(scenarios)
@api.route("/lifecycle/scenarios/<scenario_id>", methods=["DELETE"])
def delete_scene(scenario_id):
"""Delete a scenario."""
ok = _get_ms().delete_scenario(scenario_id, g.team_id)
if ok:
return _ok({"deleted": scenario_id})
return _err("Scenario not found or access denied", 404, code="NOT_FOUND")
# -- Persona Generation --
@api.route("/lifecycle/generate-persona", methods=["POST"])
def generate_persona():
"""Generate user persona from memories and scenarios using LLM.
Body: {"agent_id": "...", "max_items": 30}
"""
d = request.json or {}
agent_id = d.get("agent_id", "")
if not agent_id:
return _err("agent_id required", code="VALIDATION_ERROR")
try:
result = _get_ms().generate_persona(
agent_id=agent_id,
team_id=g.team_id,
max_items=d.get("max_items", 30),
)
except Exception as e:
logger.exception("generate-persona failed")
return _err(str(e), 500, code="INTERNAL_ERROR")
if "error" in result:
return _err(result["error"])
return _ok(result)
@api.route("/lifecycle/persona/<agent_id>", methods=["GET"])
def get_persona(agent_id):
"""Get the latest persona for an agent."""
persona = _get_ms().get_persona(agent_id)
if persona:
return _ok(persona)
return _ok({"message": "No persona found", "agent_id": agent_id})
@api.route("/lifecycle/persona/<persona_id>", methods=["DELETE"])
def delete_persona(persona_id):
"""Delete a persona."""
ok = _get_ms().delete_persona(persona_id, g.team_id)
if ok:
return _ok({"deleted": persona_id})
return _err("Persona not found or access denied", 404, code="NOT_FOUND")
# -- Atomic Fact Extraction --
@api.route("/lifecycle/extract-facts", methods=["POST"])
def extract_facts():
"""Extract structured atomic facts from working memory using LLM.
Body: {"agent_id": "...", "max_memories": 20, "delete_after": false}
Returns: {"facts": [{"id": "...", "content": "...", "importance": 0.8, "category": "..."}]}
"""
d = request.json or {}
agent_id = d.get("agent_id", "")
if not agent_id:
return _err("agent_id required", code="VALIDATION_ERROR")
try:
result = _get_ms().extract_facts(
agent_id=agent_id,
team_id=g.team_id,
max_memories=d.get("max_memories", 20),
delete_after=d.get("delete_after", False),
)
except Exception as e:
logger.exception("extract-facts failed")
return _err(str(e), 500, code="INTERNAL_ERROR")
if "error" in result:
return _err(result["error"])
return _ok(result)
app.py
+120
View File
@@ -0,0 +1,120 @@
"""Flask Application Entry Point"""
import logging
import sys
import os
# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import json
from datetime import datetime, date
from flask import Flask, Response
from config import Config
from memory_system import MemorySystem
from api.routes import api
from auth.api_auth import require_auth, api_auth_bp
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)
def _sanitize(obj):
"""Recursively sanitize data for JSON serialization."""
if isinstance(obj, dict):
return {k: _sanitize(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_sanitize(v) for v in obj]
if isinstance(obj, bytes):
return None
if isinstance(obj, (datetime, date)):
return obj.isoformat()
return obj
def json_response(data, status=200):
"""Create a JSON response with proper serialization."""
return Response(
json.dumps(_sanitize(data), ensure_ascii=False),
status=status,
mimetype="application/json",
)
LANDING_HTML = """<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Memory System - 多智能体记忆系统</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Helvetica, Arial, sans-serif;
background: #0d1117; color: #c9d1d9; display: flex; justify-content: center;
align-items: center; min-height: 100vh; }
.container { text-align: center; max-width: 600px; padding: 40px 24px; }
h1 { font-size: 32px; color: #f0f6fc; margin-bottom: 12px; }
p { font-size: 16px; color: #8b949e; line-height: 1.8; margin-bottom: 24px; }
.links { display: flex; gap: 16px; justify-content: center; flex-wrap: wrap; }
a { display: inline-block; padding: 10px 24px; border-radius: 8px;
text-decoration: none; font-size: 15px; font-weight: 500; transition: all 0.2s; }
a.primary { background: #238636; color: #fff; }
a.primary:hover { background: #2ea043; }
a.secondary { border: 1px solid #30363d; color: #c9d1d9; }
a.secondary:hover { border-color: #58a6ff; color: #58a6ff; }
.footer { margin-top: 48px; font-size: 13px; color: #484f58; }
</style>
</head>
<body>
<div class="container">
<h1>🧠 Memory System</h1>
<p>多智能体记忆系统 —— 为 AI Agent 提供双层次记忆、语义检索和智能生命周期管理的基础设施。</p>
<div class="links">
<a class="primary" href="/docs/">📖 API 文档</a>
<a class="secondary" href="/health">💚 健康检查</a>
</div>
<p class="footer">Memory System v3.0 | Powered by Flask + MySQL + Redis + BGE</p>
</div>
</body>
</html>"""
def create_app():
app = Flask(__name__)
app.config["JSON_SORT_KEYS"] = False
app.json_response = json_response
app.sanitize = _sanitize
# Initialize memory system
cfg = Config()
app.memory_system = MemorySystem(cfg)
logger.info("MemorySystem initialized")
# Register auth before_request handler
app.before_request(require_auth)
logger.info("API Key auth middleware registered")
# Register API blueprint
app.register_blueprint(api)
# Register admin blueprint (localhost-only)
app.register_blueprint(api_auth_bp)
logger.info("Admin endpoints registered at /admin/*")
return app
app = create_app()
@app.route("/")
def index():
from flask import Response
return Response(LANDING_HTML, mimetype="text/html")
if __name__ == "__main__":
app.run(host="127.0.0.1", port=8081, debug=False)
auth/__init__.py
+2
View File
@@ -0,0 +1,2 @@
"""Authentication module for Memory System API"""
from auth.api_auth import require_auth, api_auth_bp
auth/api_auth.py
+214
View File
@@ -0,0 +1,214 @@
"""API Key + HMAC Signature Authentication for Memory System
Auth flow:
1. Client sends X-API-Key, X-Timestamp, X-Signature headers
2. X-Signature = HMAC-SHA256(api_secret, timestamp + request_body)
3. Server validates key exists, is active, timestamp within 5min, signature matches
4. Authenticated team_id is injected into flask.g for downstream use
"""
import hashlib
import hmac
import json
import logging
import os
import secrets
import time
from datetime import datetime
from functools import wraps
import pymysql
from flask import Blueprint, current_app, g, request
logger = logging.getLogger(__name__)
# Timestamp tolerance in seconds (5 minutes)
TIMESTAMP_TOLERANCE = 300
api_auth_bp = Blueprint("admin", __name__, url_prefix="/admin")
# Helpers
def _get_db():
ms = current_app.memory_system
return ms.mysql._get_conn()
def _verify_signature(api_secret, timestamp, body):
message = timestamp.encode() + body
return hmac.new(api_secret.encode(), message, hashlib.sha256).hexdigest()
def _lookup_api_key(api_key):
conn = _get_db()
try:
with conn.cursor(pymysql.cursors.DictCursor) as cur:
cur.execute(
"SELECT id, team_id, api_key, secret, name, is_active "
"FROM api_keys WHERE api_key = %s",
(api_key,),
)
return cur.fetchone()
finally:
conn.close()
def _touch_api_key(key_id):
conn = _get_db()
try:
with conn.cursor() as cur:
cur.execute(
"UPDATE api_keys SET last_used = NOW() WHERE id = %s",
(key_id,),
)
conn.commit()
finally:
conn.close()
# Auth middleware
def require_auth():
path = request.path
# Skip health check, root landing page, and admin endpoints
if path in ("/", "/health") or path.startswith("/admin"):
return None
api_key = request.headers.get("X-API-Key")
timestamp = request.headers.get("X-Timestamp")
signature = request.headers.get("X-Signature")
if not api_key or not timestamp or not signature:
return current_app.json_response(
{"ok": False, "error": "Missing authentication headers. Required: X-API-Key, X-Timestamp, X-Signature",
"code": "AUTH_MISSING_HEADERS"}, 401)
try:
ts = float(timestamp)
except (ValueError, TypeError):
return current_app.json_response(
{"ok": False, "error": "Invalid timestamp format",
"code": "AUTH_INVALID_TIMESTAMP"}, 401)
now = time.time()
if abs(now - ts) > TIMESTAMP_TOLERANCE:
return current_app.json_response(
{"ok": False, "error": "Timestamp expired (5 min tolerance)",
"code": "AUTH_TIMESTAMP_EXPIRED"}, 401)
key_record = _lookup_api_key(api_key)
if not key_record:
return current_app.json_response(
{"ok": False, "error": "Invalid API key",
"code": "AUTH_INVALID_KEY"}, 401)
if not key_record["is_active"]:
return current_app.json_response(
{"ok": False, "error": "API key is disabled",
"code": "AUTH_KEY_DISABLED"}, 401)
body = request.get_data()
expected = _verify_signature(key_record["secret"], timestamp, body)
if not hmac.compare_digest(signature, expected):
return current_app.json_response(
{"ok": False, "error": "Invalid signature",
"code": "AUTH_INVALID_SIGNATURE"}, 401)
g.team_id = key_record["team_id"]
g.api_key_id = key_record["id"]
_touch_api_key(key_record["id"])
return None
# Admin endpoints (localhost only)
@api_auth_bp.before_request
def _restrict_admin_to_localhost():
remote = request.remote_addr or ""
if remote not in ("127.0.0.1", "::1", "localhost"):
return current_app.json_response(
{"ok": False, "error": "Admin endpoints are localhost-only",
"code": "FORBIDDEN_ADMIN"}, 403)
@api_auth_bp.route("/api-keys", methods=["POST"])
def create_api_key():
d = request.json or {}
team_id = d.get("team_id")
name = d.get("name", "")
if not team_id:
return current_app.json_response(
{"ok": False, "error": "team_id required",
"code": "VALIDATION_ERROR"}, 400)
if not current_app.memory_system.get_team(team_id):
return current_app.json_response(
{"ok": False, "error": f"Team {team_id} not found",
"code": "TEAM_NOT_FOUND"}, 404)
key_id = secrets.token_hex(16)
api_key = "msk_" + secrets.token_hex(24)
secret = secrets.token_hex(32)
conn = _get_db()
try:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO api_keys (id, team_id, api_key, secret, name, is_active) "
"VALUES (%s, %s, %s, %s, %s, TRUE)",
(key_id, team_id, api_key, secret, name),
)
conn.commit()
finally:
conn.close()
return current_app.json_response({
"ok": True,
"data": {
"id": key_id, "team_id": team_id, "api_key": api_key,
"secret": secret, "name": name,
"message": "Store the secret securely. It will NOT be shown again.",
},
}, 201)
@api_auth_bp.route("/api-keys", methods=["GET"])
def list_api_keys():
conn = _get_db()
try:
with conn.cursor(pymysql.cursors.DictCursor) as cur:
cur.execute(
"SELECT id, team_id, api_key, name, is_active, created_at, last_used "
"FROM api_keys ORDER BY created_at DESC")
rows = cur.fetchall()
finally:
conn.close()
for row in rows:
for k, v in row.items():
if isinstance(v, datetime):
row[k] = v.isoformat()
return current_app.json_response({"ok": True, "data": rows})
@api_auth_bp.route("/api-keys/<key_id>", methods=["DELETE"])
def disable_api_key(key_id):
conn = _get_db()
try:
with conn.cursor() as cur:
cur.execute(
"UPDATE api_keys SET is_active = FALSE WHERE id = %s",
(key_id,))
conn.commit()
affected = cur.rowcount
finally:
conn.close()
if affected == 0:
return current_app.json_response(
{"ok": False, "error": "API key not found",
"code": "NOT_FOUND"}, 404)
return current_app.json_response({"ok": True, "message": "API key disabled"})
config.py
+51
View File
@@ -0,0 +1,51 @@
"""Memory System Configuration"""
import os
import json
class Config:
# MySQL
MYSQL_HOST = os.getenv("MYSQL_HOST", "127.0.0.1")
MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
MYSQL_USER = os.getenv("MYSQL_USER", "your_db_user")
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "CHANGE_ME_MYSQL_PASSWORD")
MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "memory_system")
MYSQL_CHARSET = "utf8mb4"
MYSQL_UNIX_SOCKET = os.getenv("MYSQL_UNIX_SOCKET", "/var/run/mysqld/mysqld.sock")
# Redis
REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1")
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "CHANGE_ME_REDIS_PASSWORD")
REDIS_DB = int(os.getenv("REDIS_DB", "0"))
# BGE Embedding Service
EMBEDDING_SERVICE_URL = os.getenv("EMBEDDING_SERVICE_URL", "http://127.0.0.1:8080")
EMBEDDING_DIM = 512
# Working memory TTL (seconds)
WORKING_MEMORY_TTL = int(os.getenv("WORKING_MEMORY_TTL", "259200")) # 72h
# Cleanup defaults
CLEANUP_MAX_AGE_DAYS = 90
CLEANUP_MIN_IMPORTANCE = 0.2
# API
API_HOST = os.getenv("API_HOST", "127.0.0.1")
API_PORT = int(os.getenv("API_PORT", "8081"))
# LLM providers (JSON list, priority=1 is primary, higher = fallback)
LLM_PROVIDERS_JSON = os.getenv("LLM_PROVIDERS", "[]")
@property
def llm_providers(self) -> list:
try:
val = self.LLM_PROVIDERS_JSON
providers = json.loads(val) if val else []
for p in providers:
key = p.get("api_key", "")
if isinstance(key, str) and key.startswith("${") and key.endswith("}"):
p["api_key"] = os.getenv(key[2:-1], "") or ""
return providers
except (json.JSONDecodeError, TypeError):
return []
doc.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 24 KiB

docs/doc.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 24 KiB

docs/index.html
+1809
View File
File diff suppressed because it is too large Load Diff
embedding_service.py
+67
View File
@@ -0,0 +1,67 @@
"""BGE Embedding Service Client"""
import requests
import logging
from typing import List
logger = logging.getLogger(__name__)
# BGE-small-zh-v1.5: 512 tokens, roughly ~1500 Chinese chars
# Conservative limit: 2000 chars (most text under this is safe)
MAX_TEXT_CHARS = 2000
class EmbeddingService:
"""Client for BGE-small-zh-v1.5 embedding service."""
def __init__(self, base_url: str = "http://127.0.0.1:8080", timeout: int = 30):
self.base_url = base_url.rstrip("/")
self.timeout = timeout
def embed(self, text: str) -> List[float]:
"""Get embedding for a single text.
Text is truncated to MAX_TEXT_CHARS if too long.
"""
original_len = len(text)
if len(text) > MAX_TEXT_CHARS:
logger.warning(f"Text too long ({original_len} chars), truncating to {MAX_TEXT_CHARS}")
text = text[:MAX_TEXT_CHARS]
try:
resp = requests.post(
f"{self.base_url}/embed",
json={"text": text},
timeout=self.timeout,
)
resp.raise_for_status()
data = resp.json()
return data["embedding"]
except requests.exceptions.ConnectionError:
msg = "Embedding service unavailable (connection refused)"
logger.error(msg)
raise RuntimeError(msg)
except requests.exceptions.Timeout:
msg = f"Embedding request timeout ({self.timeout}s). Text may be too long ({original_len} chars, max {MAX_TEXT_CHARS})"
logger.error(msg)
raise RuntimeError(msg)
except requests.exceptions.HTTPError as e:
status = e.response.status_code if e.response is not None else "unknown"
msg = f"Embedding service error (HTTP {status}): {e}"
logger.error(msg)
raise RuntimeError(msg)
except Exception as e:
msg = f"Embedding failed: {e}"
logger.error(msg)
raise RuntimeError(msg)
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""Get embeddings for multiple texts."""
return [self.embed(t) for t in texts]
def health_check(self) -> bool:
"""Check if embedding service is alive."""
try:
resp = requests.get(f"{self.base_url}/health", timeout=5)
return resp.status_code == 200
except Exception:
return False
gunicorn.conf.py
+15
View File
@@ -0,0 +1,15 @@
"""Gunicorn configuration for Memory System"""
import multiprocessing
import os
bind = os.getenv("API_BIND", "127.0.0.1:8081")
workers = min(multiprocessing.cpu_count(), 4)
worker_class = "sync"
timeout = 120
keepalive = 5
max_requests = 2000
max_requests_jitter = 200
accesslog = "-"
errorlog = "-"
loglevel = "info"
preload_app = True
init_db.sql
+140
View File
@@ -0,0 +1,140 @@
-- Memory System Database Initialization
CREATE DATABASE IF NOT EXISTS memory_system DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
USE memory_system;
-- Teams table
CREATE TABLE IF NOT EXISTS teams (
id VARCHAR(64) PRIMARY KEY,
name VARCHAR(255) NOT NULL,
description TEXT,
config JSON,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- Agents table
CREATE TABLE IF NOT EXISTS agents (
id VARCHAR(64) PRIMARY KEY,
team_id VARCHAR(64) NOT NULL,
name VARCHAR(255) NOT NULL,
role VARCHAR(255),
config JSON,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
INDEX idx_team (team_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- Personal long-term memories
CREATE TABLE IF NOT EXISTS personal_memories (
id VARCHAR(64) PRIMARY KEY,
agent_id VARCHAR(64) NOT NULL,
team_id VARCHAR(64) NOT NULL,
content TEXT NOT NULL,
summary TEXT,
embedding BLOB,
importance FLOAT DEFAULT 0.5,
memory_type VARCHAR(32) DEFAULT 'long_term',
metadata JSON,
access_count INT DEFAULT 0,
last_accessed TIMESTAMP NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
INDEX idx_agent (agent_id),
INDEX idx_team (team_id),
INDEX idx_importance (importance),
INDEX idx_created (created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- Team shared memories
CREATE TABLE IF NOT EXISTS team_memories (
id VARCHAR(64) PRIMARY KEY,
team_id VARCHAR(64) NOT NULL,
content TEXT NOT NULL,
summary TEXT,
embedding BLOB,
importance FLOAT DEFAULT 0.5,
category VARCHAR(64) DEFAULT 'general',
metadata JSON,
access_count INT DEFAULT 0,
last_accessed TIMESTAMP NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
INDEX idx_team (team_id),
INDEX idx_category (category),
INDEX idx_importance (importance),
INDEX idx_created (created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
FLUSH PRIVILEGES;
-- ============================================================
-- v3.0 新增表 - API Keys 认证
-- ============================================================
CREATE TABLE IF NOT EXISTS api_keys (
id VARCHAR(64) PRIMARY KEY,
team_id VARCHAR(64) NOT NULL,
api_key VARCHAR(128) NOT NULL UNIQUE,
secret VARCHAR(128) NOT NULL,
name VARCHAR(255) DEFAULT '',
is_active TINYINT DEFAULT 1,
last_used TIMESTAMP NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
INDEX idx_team (team_id),
INDEX idx_api_key (api_key)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- ============================================================
-- v3.0 新增表 - Pipeline 自动化配置
-- ============================================================
CREATE TABLE IF NOT EXISTS pipeline_config (
id INT AUTO_INCREMENT PRIMARY KEY,
team_id VARCHAR(64) NOT NULL UNIQUE,
compress_every_n INT DEFAULT 0,
compress_target_count INT DEFAULT 5,
cleanup_idle_days INT DEFAULT 0,
cleanup_min_importance FLOAT DEFAULT 0.2,
enabled TINYINT DEFAULT 1,
warmup_max_memories INT DEFAULT 0,
warmup_compress_every_n INT DEFAULT 1,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
INDEX idx_team (team_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- ============================================================
-- v3.0 新增表 - 场景块（LLM 聚合）
-- ============================================================
CREATE TABLE IF NOT EXISTS memory_scenarios (
id INT AUTO_INCREMENT PRIMARY KEY,
team_id VARCHAR(64) NOT NULL,
agent_id VARCHAR(64) NOT NULL,
name VARCHAR(255),
summary TEXT,
memory_ids JSON,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE,
INDEX idx_agent (team_id, agent_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-- ============================================================
-- v3.0 新增表 - 用户画像（LLM 生成）
-- ============================================================
CREATE TABLE IF NOT EXISTS user_personas (
id INT AUTO_INCREMENT PRIMARY KEY,
team_id VARCHAR(64) NOT NULL,
agent_id VARCHAR(64) NOT NULL,
preferences JSON,
habits JSON,
expertise JSON,
communication_style TEXT,
summary TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE CASCADE,
FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE,
INDEX idx_agent (team_id, agent_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
lifecycle/__init__.py
+1
View File
@@ -0,0 +1 @@
# Lifecycle management
lifecycle/aggregator.py
+92
View File
@@ -0,0 +1,92 @@
"""Scenario Aggregator - Group related memories into scenario blocks."""
import json
import logging
from typing import Optional, List
from lifecycle.llm_parse import parse_llm_json_array
logger = logging.getLogger(__name__)
SCENE_SYSTEM_PROMPT = """你是JSON生成器。直接输出JSON数组，禁止分析、思考、解释、讨论、markdown。
任务：将记忆按主题聚合成场景块。
规则：
1. 将相关记忆归入同一场景
2. 每个场景：name（名称）、summary（100字以内摘要）、memory_ids（ID列表）
3. 最多10个场景
4. 第一个字符必须是[，最后一个字符必须是]
输出格式：
[{"name":"场景名","summary":"摘要","memory_ids":["id1","id2"]}]
只输出JSON，不要有任何其他文字。"""
class ScenarioAggregator:
"""Aggregate personal memories into scenario blocks."""
def __init__(self, mysql_store, llm_client):
self.mysql = mysql_store
self.llm = llm_client
def aggregate(self, agent_id: str, team_id: str,
max_memories: int = 50) -> dict:
"""Aggregate memories into scenarios using LLM."""
if not self.llm.available:
return {"error": "LLM not configured. Set LLM_API_URL, LLM_API_KEY, LLM_MODEL environment variables."}
memories = self.mysql.get_personal_memories_by_agent(agent_id, limit=max_memories)
if not memories:
return {"scenarios": [], "count": 0}
memory_lines = []
for m in memories:
mid = m["id"]
content = m.get("content", "")[:100]
memory_lines.append(f"[{mid}] {content}")
user_prompt = "请将以下记忆聚合为JSON场景数组：\n\n" + "\n".join(memory_lines)
try:
scenarios = parse_llm_json_array(
self.llm, SCENE_SYSTEM_PROMPT, user_prompt,
temperature=0.3, context="aggregate-scenes",
max_tokens=2000,
)
except ValueError as e:
return {"error": str(e)}
stored = []
for scene in scenarios:
if not isinstance(scene, dict):
logger.warning("Skipping non-dict item in scenarios: %r", scene)
continue
name = scene.get("name", "unnamed")
summary = scene.get("summary", "")
memory_ids = scene.get("memory_ids", [])
valid_ids = [mid for mid in memory_ids if isinstance(mid, str) and len(mid) > 0]
scene_id = self.mysql.add_scenario(
team_id=team_id,
agent_id=agent_id,
name=name,
summary=summary,
memory_ids=valid_ids,
)
stored.append({
"id": scene_id,
"name": name,
"summary": summary,
"memory_ids": valid_ids,
})
return {"scenarios": stored, "count": len(stored)}
def get_scenarios(self, agent_id: str) -> List[dict]:
return self.mysql.get_scenarios_by_agent(agent_id)
def delete_scenario(self, scenario_id: str, team_id: str) -> bool:
scenario = self.mysql.get_scenario(scenario_id)
if not scenario or scenario["team_id"] != team_id:
return False
return self.mysql.delete_scenario(scenario_id)
lifecycle/cleaner.py
+43
View File
@@ -0,0 +1,43 @@
"""Memory Cleaner - Periodic cleanup of old/low-importance memories"""
import logging
logger = logging.getLogger(__name__)
class MemoryCleaner:
"""Clean up old, low-importance memories."""
def __init__(self, mysql_store):
self.mysql = mysql_store
def cleanup(self, team_id: str = None, max_age_days: int = 90,
min_importance: float = 0.2) -> dict:
"""
Remove old memories with low importance scores.
Args:
team_id: If set, only clean this team's memories
max_age_days: Remove memories older than this
min_importance: Remove memories with importance below this
Returns:
Dict with counts of deleted memories
"""
personal_deleted = self.mysql.cleanup_personal_memories(
team_id=team_id,
max_age_days=max_age_days,
min_importance=min_importance,
)
team_deleted = self.mysql.cleanup_team_memories(
team_id=team_id,
max_age_days=max_age_days,
min_importance=min_importance,
)
result = {
"personal_deleted": personal_deleted,
"team_deleted": team_deleted,
"total_deleted": personal_deleted + team_deleted,
}
logger.info(f"Cleanup result: {result}")
return result
lifecycle/compressor.py
+89
View File
@@ -0,0 +1,89 @@
"""Working Memory Compressor"""
import logging
from typing import Optional, Callable
logger = logging.getLogger(__name__)
class MemoryCompressor:
"""Compress working memories when they exceed threshold."""
def __init__(self, mysql_store, redis_cache, embedding_service):
self.mysql = mysql_store
self.redis = redis_cache
self.embedder = embedding_service
def compress(self, agent_id: str, max_items: int = 20,
summary_callback: Optional[Callable] = None) -> str:
"""
Compress working memories into a single long-term memory.
If working memory exceeds max_items, take the oldest items,
summarize them (via callback or simple concatenation), and
store as a personal long-term memory.
Returns: Summary text of compressed memories.
"""
# Get agent info to find team_id
agent = self.mysql.get_agent(agent_id)
if not agent:
raise ValueError(f"Agent {agent_id} not found")
items = self.redis.get_working_memories(agent_id, limit=100)
if len(items) <= max_items:
return {
"compressed": 0,
"remaining": len(items),
"total_before": len(items),
"target_count": max_items,
"summary": None,
}
# Take items beyond max_items (oldest)
to_compress = items[max_items:]
remaining = items[:max_items]
# Build summary
contents = [item.get("content", "") for item in to_compress]
if summary_callback:
summary = summary_callback(contents)
else:
summary = " | ".join(contents[:50]) # Simple concat, limit 50 items
# Store as personal long-term memory
embedding = self.embedder.embed(summary)
from storage.vector_search import embedding_to_bytes
emb_bytes = embedding_to_bytes(embedding)
self.mysql.add_personal_memory(
agent_id=agent_id,
team_id=agent["team_id"],
content=summary,
embedding=emb_bytes,
importance=0.6, # Slightly elevated importance for compressed memories
metadata={
"source": "compression",
"item_count": len(to_compress),
"source_items": [
{"content": item.get("content", "")[:200], "timestamp": item.get("timestamp")}
for item in to_compress
],
},
)
# Replace working memory with remaining items
self.redis.clear_working_memory(agent_id)
for item in reversed(remaining): # lpush reverses order
self.redis.add_working_memory(
agent_id, item["content"],
metadata=item.get("metadata", {}),
)
logger.info(f"Compressed {len(to_compress)} working memories for agent {agent_id}")
return {
"compressed": len(to_compress),
"remaining": len(remaining),
"total_before": len(items),
"target_count": max_items,
"summary": summary,
}
lifecycle/fact_extractor.py
+101
View File
@@ -0,0 +1,101 @@
"""Fact Extractor - Extract structured atomic facts from working memory using LLM."""
import json
import logging
from typing import Optional, List
from lifecycle.llm_parse import parse_llm_json_array
logger = logging.getLogger(__name__)
FACT_SYSTEM_PROMPT = """你是JSON生成器。输入工作记忆，直接提取结构化事实。
严格规则：
1. 禁止输出分析、推理、思考、解释
2. 第一个字符必须是 [，最后一个字符必须是 ]
3. 每条事实：content（50-200字）、importance（0.0-1.0）、category（分类）
4. 最多20条，无值得提取的事实返回 []
输出格式（只输出JSON，不输出其他任何文字）：
[{"content":"事实","importance":0.8,"category":"分类"}]"""
class FactExtractor:
"""Extract atomic facts from working memory using LLM."""
def __init__(self, mysql_store, redis_cache, llm_client, embedding_service):
self.mysql = mysql_store
self.redis = redis_cache
self.llm = llm_client
self.embedder = embedding_service
def extract(self, agent_id: str, team_id: str,
max_memories: int = 20,
delete_after: bool = False) -> dict:
"""Extract structured facts from working memory."""
if not self.llm.available:
return {"error": "LLM not configured. Set LLM_API_URL, LLM_API_KEY, LLM_MODEL environment variables."}
items = self.redis.get_working_memories(agent_id, limit=max_memories)
if not items:
return {"facts": [], "count": 0}
memory_lines = []
for i, item in enumerate(items):
if not isinstance(item, dict):
logger.warning('Skipping non-dict working memory item: %r', item)
continue
content = item.get("content", "")[:300]
memory_lines.append(f"[{i+1}] {content}")
user_prompt = "以下是需要提取事实的工作记忆：\n\n" + "\n".join(memory_lines)
try:
facts = parse_llm_json_array(
self.llm, FACT_SYSTEM_PROMPT, user_prompt,
temperature=0.2, context="extract-facts",
)
except ValueError as e:
return {"error": str(e)}
stored = []
for fact in facts:
if not isinstance(fact, dict):
logger.warning('Skipping non-dict fact from LLM: %r', fact)
continue
content = fact.get("content", "").strip()
if not content or len(content) < 10:
continue
importance = fact.get("importance", 0.6)
if not isinstance(importance, (int, float)) or importance < 0 or importance > 1:
importance = 0.6
category = fact.get("category", "")
embedding = self.embedder.embed(content)
from storage.vector_search import embedding_to_bytes
emb_bytes = embedding_to_bytes(embedding)
memory_id = self.mysql.add_personal_memory(
agent_id=agent_id,
team_id=team_id,
content=content,
embedding=emb_bytes,
importance=float(importance),
metadata={
"source": "fact_extraction",
"category": category,
"extracted_from": "working_memory",
},
)
stored.append({
"id": memory_id,
"content": content,
"importance": importance,
"category": category,
})
if delete_after and stored:
self.redis.clear_working_memory(agent_id)
logger.info(f"Deleted working memories for {agent_id} after extracting {len(stored)} facts")
return {"facts": stored, "count": len(stored)}
lifecycle/llm_parse.py
+344
View File
@@ -0,0 +1,344 @@
"""Shared LLM response parsing utilities for lifecycle modules."""
import json
import logging
import re
import time
from typing import Optional
logger = logging.getLogger(__name__)
MAX_PARSE_RETRIES = 2 # 3 total attempts (1 initial + 2 retry)
RETRY_SYSTEM_PROMPT = (
'只输出原始JSON。不要有任何其他文字、分析、思考、解释、markdown。\n'
'输出格式示例：\n'
'[{"name": "场景名", "summary": "摘要", "memory_ids": ["id1", "id2"]}]\n'
'重要：你回复的第一个字符必须是 [，最后一个字符必须是 ]。'
)
RETRY_OBJECT_PROMPT = (
'只输出原始JSON对象。不要有任何其他文字、分析、思考、解释、markdown。\n'
'输出格式：\n'
'{"data": [{"name": "场景名", "summary": "摘要", "memory_ids": ["id1", "id2"]}]}\n'
'重要：你回复的第一个字符必须是 {，最后一个字符必须是 }。'
)
def _strip_markdown_code_block(text: str) -> str:
text = text.strip()
if not text.startswith("```"):
return text
text = text.split("\n", 1)[1] if "\n" in text else text[3:]
if text.endswith("```"):
text = text[:-3].strip()
if text.lower().startswith("json"):
text = text[4:].strip()
elif text.lower().startswith("json\n"):
text = text[5:].strip()
return text.strip()
def _strip_reasoning_prefix(text: str) -> str:
"""Strip reasoning/analysis text that some models output before JSON.
Models like deepseek-v4-flash often output chain-of-thought analysis
before the actual JSON. This function tries to find where the JSON
actually starts by looking for lines that start with [ or { after
stripping analysis text.
"""
text = text.strip()
# If text starts with [ or { it might already be JSON
if text and text[0] in ('[', '{'):
return text
# Try to find the last occurrence of a JSON-like pattern
lines = text.split('\n')
for i in range(len(lines) - 1, -1, -1):
stripped = lines[i].strip()
if stripped.startswith('[') or stripped.startswith('{'):
candidate = '\n'.join(lines[i:])
try:
json.loads(candidate)
return candidate
except (json.JSONDecodeError, ValueError):
continue
# Try finding first [ or { that starts a valid JSON
for start_char in ['[', '{']:
idx = text.find(start_char)
while idx >= 0:
candidate = text[idx:]
end_char = ']' if start_char == '[' else '}'
if candidate.rstrip().endswith(end_char):
try:
json.loads(candidate)
return candidate
except (json.JSONDecodeError, ValueError):
pass
idx = text.find(start_char, idx + 1)
return text
def _find_json_string(text: str) -> Optional[str]:
"""Find the first valid JSON array or object in text.
Iterates through ALL bracket-delimited segments, not just from the
first bracket. This handles cases where the LLM outputs analysis
text containing bracketed IDs before the actual JSON.
"""
text = text.strip()
if not text:
return None
# Fast path: entire text is valid JSON
if (text.startswith("[") or text.startswith("{")) and text[-1] in ("}", "]"):
try:
json.loads(text)
return text
except (json.JSONDecodeError, ValueError):
pass
# Strip markdown code blocks
cleaned = _strip_markdown_code_block(text)
if cleaned != text:
return _find_json_string(cleaned)
# Strip reasoning prefix
cleaned = _strip_reasoning_prefix(text)
if cleaned != text:
result = _find_json_string(cleaned)
if result:
return result
# Scan all bracket-delimited segments (for "[" / "]" and "{" / "}")
for start_char, end_char in [("[", "]"), ("{", "}")]:
search_start = 0
while True:
start_idx = text.find(start_char, search_start)
if start_idx == -1:
break
depth = 0
for i in range(start_idx, len(text)):
if text[i] == start_char:
depth += 1
elif text[i] == end_char:
depth -= 1
if depth == 0:
candidate = text[start_idx:i + 1]
try:
json.loads(candidate)
return candidate
except (json.JSONDecodeError, ValueError):
# Move past this bracket pair and continue scanning
search_start = i + 1
break
else:
# No matching end bracket found
search_start = start_idx + 1
# Regex fallback: find ALL bracket-delimited segments
for pattern in [r"\[.*\]", r"\{.*\}"]:
for match in re.finditer(pattern, text, re.DOTALL):
candidate = match.group(0)
try:
json.loads(candidate)
return candidate
except (json.JSONDecodeError, ValueError):
continue
return None
def _try_fix_json(text: str) -> Optional[str]:
"""Try to fix common JSON issues from LLM output."""
# Remove trailing commas before } or ]
text = re.sub(r",\s*}", "}", text)
text = re.sub(r",\s*\]", "]", text)
# Replace single quotes with double quotes only if no double quotes present
if '"' not in text:
text = text.replace("'", '"')
# Remove BOM and zero-width characters
text = text.replace("\ufeff", "").replace("\u200b", "")
# Try parsing as-is first
try:
json.loads(text)
return text
except json.JSONDecodeError:
pass
# Try fixing unescaped newlines inside JSON string values
fixed = re.sub(r'(?<=: ")((?:[^"]|\\")*?)(\n)((?:[^"]|\\")*?(?="))', r'\1\\n\3', text)
try:
json.loads(fixed)
return fixed
except json.JSONDecodeError:
pass
# Try wrapping in array if it looks like a bare object
stripped = text.strip()
if stripped.startswith("{") and not stripped.startswith("["):
wrapped = "[" + stripped + "]"
try:
json.loads(wrapped)
return wrapped
except json.JSONDecodeError:
pass
# Try to extract JSON from truncated output: find last complete object in array
if stripped.startswith("["):
depth = 0
objects = []
obj_start = -1
for i, ch in enumerate(stripped):
if ch == "{":
if depth == 0:
obj_start = i
depth += 1
elif ch == "}":
depth -= 1
if depth == 0 and obj_start >= 0:
objects.append(stripped[obj_start:i + 1])
obj_start = -1
if objects:
candidate = "[" + ",".join(objects) + "]"
try:
json.loads(candidate)
return candidate
except json.JSONDecodeError:
pass
return None
def parse_llm_json_array(llm_client, system_prompt, user_prompt,
temperature=0.3, context="", max_tokens=4000):
last_raw = ""
total = 1 + MAX_PARSE_RETRIES
t_start = time.time()
for attempt in range(total):
is_retry = attempt > 0
if is_retry:
if attempt >= total - 1:
# Final attempt: json_mode as last resort with wrapper object
prompt = RETRY_OBJECT_PROMPT
use_json_mode = True
prefill_val = None
else:
# Intermediate retry: use prefill "["
prompt = RETRY_SYSTEM_PROMPT
use_json_mode = False
prefill_val = "["
else:
# First attempt: prefill "[" to force JSON output immediately
prompt = system_prompt
use_json_mode = False
prefill_val = "["
response = llm_client.chat(
prompt, user_prompt,
temperature=temperature, json_mode=use_json_mode,
max_tokens=max_tokens,
thinking={"type": "disabled"},
prefill=prefill_val,
)
if not response:
logger.warning("[%s] Attempt %d: LLM returned empty", context, attempt + 1)
continue
last_raw = response
json_str = _find_json_string(response)
if json_str:
try:
result = json.loads(json_str)
except json.JSONDecodeError as e:
logger.warning("[%s] Attempt %d: JSON parse error: %s", context, attempt + 1, e)
continue
if isinstance(result, list):
elapsed = time.time() - t_start
logger.info("[%s] Parsed JSON array in %.1fs (%d attempt(s), %d items)", context, elapsed, attempt + 1, len(result))
return result
if isinstance(result, dict):
if "data" in result and isinstance(result["data"], list):
elapsed = time.time() - t_start
logger.info("[%s] Parsed JSON array in %.1fs (%d attempt(s), %d items)", context, elapsed, attempt + 1, len(result["data"]))
return result["data"]
logger.warning("[%s] Attempt %d: LLM returned dict, wrapping in list", context, attempt + 1)
return [result]
fixed = _try_fix_json(response)
if fixed:
try:
result = json.loads(fixed)
except json.JSONDecodeError:
continue
if isinstance(result, list):
return result
if isinstance(result, dict):
if "data" in result and isinstance(result["data"], list):
return result["data"]
return [result]
logger.warning(
"[%s] Attempt %d/%d: No valid JSON array found (len=%d): %r",
context, attempt + 1, total, len(response), response[:500]
)
logger.error("[%s] All %d attempts failed. Last raw: %r", context, total, last_raw[:500])
raise ValueError("LLM failed to return valid JSON after retries")
def parse_llm_json_object(llm_client, system_prompt, user_prompt,
temperature=0.3, context="", max_tokens=4000):
last_raw = ""
total = 1 + MAX_PARSE_RETRIES
t_start = time.time()
for attempt in range(total):
is_retry = attempt > 0
if is_retry:
if attempt >= total - 1:
# Final attempt: json_mode as last resort
prompt = '只输出 JSON 对象。不要输出任何其他文字。\n输出格式：{"preferences":[],"habits":[],"expertise":[],"communication_style":"","summary":""}'
use_json_mode = True
prefill_val = None
else:
prompt = '只输出 JSON 对象。不要输出任何其他文字。\n禁止输出：分析、思考、解释、讨论、描述、markdown。'
use_json_mode = False
prefill_val = "{"
else:
# First attempt: prefill "{" to force JSON output
prompt = system_prompt
use_json_mode = False
prefill_val = "{"
response = llm_client.chat(
prompt, user_prompt,
temperature=temperature, json_mode=use_json_mode,
max_tokens=max_tokens,
thinking={"type": "disabled"},
prefill=prefill_val,
)
if not response:
logger.warning("[%s] Attempt %d: LLM returned empty", context, attempt + 1)
continue
last_raw = response
json_str = _find_json_string(response)
if json_str:
try:
result = json.loads(json_str)
except json.JSONDecodeError as e:
logger.warning("[%s] Attempt %d: JSON parse error: %s", context, attempt + 1, e)
continue
if isinstance(result, dict):
elapsed = time.time() - t_start
logger.info("[%s] Parsed JSON object in %.1fs (%d attempt(s))", context, elapsed, attempt + 1)
return result
fixed = _try_fix_json(response)
if fixed:
try:
result = json.loads(fixed)
except json.JSONDecodeError:
continue
if isinstance(result, dict):
return result
logger.warning(
"[%s] Attempt %d/%d: No valid JSON object found (len=%d): %r",
context, attempt + 1, total, len(response), response[:500]
)
logger.error("[%s] All %d attempts failed. Last raw: %r", context, total, last_raw[:500])
raise ValueError("LLM failed to return valid JSON after retries")
lifecycle/persona_generator.py
+80
View File
@@ -0,0 +1,80 @@
"""Persona Generator - Generate user persona from scenarios and memories."""
import json
import logging
from typing import Optional
from lifecycle.llm_parse import parse_llm_json_object
logger = logging.getLogger(__name__)
PERSONA_SYSTEM_PROMPT = """你是JSON生成器。直接输出JSON对象，禁止分析、思考、解释、markdown。
根据记忆和场景生成用户画像。
输出格式（第一个字符必须是{，最后一个字符必须是}）：
{"preferences":["偏好1"],"habits":["习惯1"],"expertise":["擅长1"],"communication_style":"风格","summary":"100字以内摘要"}
只输出JSON，不要有任何其他文字。"""
class PersonaGenerator:
"""Generate user persona from memories and scenarios."""
def __init__(self, mysql_store, llm_client):
self.mysql = mysql_store
self.llm = llm_client
def generate(self, agent_id: str, team_id: str,
max_items: int = 30) -> dict:
if not self.llm.available:
return {"error": "LLM not configured. Set LLM_API_URL, LLM_API_KEY, LLM_MODEL environment variables."}
memories = self.mysql.get_personal_memories_by_agent(agent_id, limit=max_items)
scenarios = self.mysql.get_scenarios_by_agent(agent_id)
if not memories and not scenarios:
return {"error": "No memories or scenarios found for this agent"}
context_parts = []
if memories:
mem_lines = [f"- {m.get('content', '')[:100]}" for m in memories]
context_parts.append("历史记忆：\n" + "\n".join(mem_lines))
if scenarios:
scene_lines = [f"- [{s.get('name', '')}] {s.get('summary', '')}" for s in scenarios]
context_parts.append("已识别场景：\n" + "\n".join(scene_lines))
user_prompt = "请根据以下内容生成用户画像JSON：\n\n" + "\n\n".join(context_parts)
try:
persona = parse_llm_json_object(
self.llm, PERSONA_SYSTEM_PROMPT, user_prompt,
temperature=0.3, context="generate-persona",
max_tokens=2000,
)
except ValueError as e:
return {"error": str(e)}
if not isinstance(persona, dict):
logger.warning("LLM returned non-dict persona: %r", persona)
return {"error": "LLM returned invalid persona format"}
pid = self.mysql.add_persona(
team_id=team_id,
agent_id=agent_id,
preferences=persona.get("preferences", []),
habits=persona.get("habits", []),
expertise=persona.get("expertise", []),
communication_style=persona.get("communication_style", ""),
summary=persona.get("summary", ""),
)
persona["id"] = pid
return {"persona": persona}
def get_persona(self, agent_id: str):
return self.mysql.get_persona_by_agent(agent_id)
def delete_persona(self, persona_id: str, team_id: str) -> bool:
persona = self.mysql.get_persona(persona_id)
if not persona or persona["team_id"] != team_id:
return False
return self.mysql.delete_persona(persona_id)
lifecycle/pipeline.py
+136
View File
@@ -0,0 +1,136 @@
"""Pipeline Automation - Auto-trigger compression and cleanup based on configurable rules."""
import logging
from typing import Optional
logger = logging.getLogger(__name__)
class PipelineConfig:
"""Pipeline configuration for a team."""
def __init__(self, team_id: str, compress_every_n: int = 0,
compress_target_count: int = 5,
cleanup_idle_days: int = 0,
cleanup_min_importance: float = 0.2,
enabled: bool = True, **kwargs):
self.team_id = team_id
self.compress_every_n = compress_every_n # 0 = disabled
self.compress_target_count = compress_target_count
self.cleanup_idle_days = cleanup_idle_days # 0 = disabled
self.cleanup_min_importance = cleanup_min_importance
self.enabled = enabled
self.warmup_max_memories = kwargs.get("warmup_max_memories", 0) # 0 = disabled
self.warmup_compress_every_n = kwargs.get("warmup_compress_every_n", 1)
def to_dict(self):
return {
"team_id": self.team_id,
"compress_every_n": self.compress_every_n,
"compress_target_count": self.compress_target_count,
"cleanup_idle_days": self.cleanup_idle_days,
"cleanup_min_importance": self.cleanup_min_importance,
"enabled": self.enabled,
"warmup_max_memories": self.warmup_max_memories,
"warmup_compress_every_n": self.warmup_compress_every_n,
}
class Pipeline:
"""Auto-pipeline: trigger compression and cleanup based on rules."""
def __init__(self, mysql_store, redis_cache, compressor, cleaner):
self.mysql = mysql_store
self.redis = redis_cache
self.compressor = compressor
self.cleaner = cleaner
def get_config(self, team_id: str) -> Optional[PipelineConfig]:
row = self.mysql.get_pipeline_config(team_id)
if not row:
return None
return PipelineConfig(
team_id=row["team_id"],
compress_every_n=row.get("compress_every_n", 0),
compress_target_count=row.get("compress_target_count", 5),
cleanup_idle_days=row.get("cleanup_idle_days", 0),
cleanup_min_importance=row.get("cleanup_min_importance", 0.2),
enabled=bool(row.get("enabled", 1)),
warmup_max_memories=row.get("warmup_max_memories", 0),
warmup_compress_every_n=row.get("warmup_compress_every_n", 1),
)
def set_config(self, team_id: str, **kwargs) -> PipelineConfig:
existing = self.mysql.get_pipeline_config(team_id)
if existing:
self.mysql.update_pipeline_config(team_id, **kwargs)
else:
self.mysql.create_pipeline_config(team_id, **kwargs)
return self.get_config(team_id)
def check_after_wm_write(self, agent_id: str) -> Optional[dict]:
agent = self.mysql.get_agent(agent_id)
if not agent:
return None
team_id = agent["team_id"]
config = self.get_config(team_id)
if not config or not config.enabled:
return None
if config.compress_every_n <= 0:
return None
wm_count = self.redis.get_working_memory_count(agent_id)
# Warm-up: use lower threshold for new agents
effective_threshold = config.compress_every_n
if config.warmup_max_memories > 0:
personal_count = len(self.mysql.get_personal_memories_by_agent(agent_id, limit=config.warmup_max_memories + 1))
if personal_count < config.warmup_max_memories:
effective_threshold = config.warmup_compress_every_n
if wm_count < effective_threshold:
return None
logger.info(f"Auto-compress triggered for {agent_id}: WM {wm_count} >= {config.compress_every_n}")
try:
result = self.compressor.compress(agent_id, max_items=config.compress_target_count)
result["auto_triggered"] = True
return result
except Exception as e:
logger.error(f"Auto-compress failed for {agent_id}: {e}")
return None
def run_cleanup(self, team_id: str = None) -> Optional[dict]:
configs = []
if team_id:
config = self.get_config(team_id)
if config and config.enabled and config.cleanup_idle_days > 0:
configs.append(config)
else:
all_configs = self.mysql.get_all_pipeline_configs()
for row in all_configs:
config = PipelineConfig(
team_id=row["team_id"],
compress_every_n=row.get("compress_every_n", 0),
compress_target_count=row.get("compress_target_count", 5),
cleanup_idle_days=row.get("cleanup_idle_days", 0),
cleanup_min_importance=row.get("cleanup_min_importance", 0.2),
enabled=bool(row.get("enabled", 1)),
warmup_max_memories=row.get("warmup_max_memories", 0),
warmup_compress_every_n=row.get("warmup_compress_every_n", 1),
)
if config.enabled and config.cleanup_idle_days > 0:
configs.append(config)
if not configs:
return None
results = []
for config in configs:
try:
result = self.cleaner.cleanup(
team_id=config.team_id,
max_age_days=config.cleanup_idle_days,
min_importance=config.cleanup_min_importance,
)
result["team_id"] = config.team_id
result["auto_triggered"] = True
results.append(result)
logger.info(f"Auto-cleanup for {config.team_id}: {result}")
except Exception as e:
logger.error(f"Auto-cleanup failed for {config.team_id}: {e}")
results.append({"team_id": config.team_id, "error": str(e)})
return {"results": results}
llm_client.py
+156
View File
@@ -0,0 +1,156 @@
"""LLM Client with multi-provider failover and JSON mode support."""
import json
import logging
import time
import requests
from typing import Optional, List, Dict, Any
logger = logging.getLogger(__name__)
COOLDOWN_SECONDS = 300 # 5 min before retrying primary after failure
class LLMClient:
"""LLM client supporting provider list with automatic failover."""
def __init__(self, providers_config: List[Dict[str, Any]] = None):
self.providers = []
if providers_config:
self.providers = sorted(providers_config, key=lambda p: p.get("priority", 999))
# Normalize API URLs
for p in self.providers:
url = p.get("api_url", "").rstrip("/")
if url and not url.endswith("/chat/completions"):
url = url + "/chat/completions"
p["api_url"] = url
self._fallback_mode = False
self._last_primary_failure = 0.0
@property
def available(self) -> bool:
return len(self.providers) > 0
def chat(self, system_prompt: str, user_prompt: str,
temperature: float = 0.3, max_tokens: int = 800,
json_mode: bool = False, thinking: Optional[dict] = None,
prefill: str = None) -> Optional[str]:
if not self.available:
logger.warning("No LLM providers configured")
return None
for i, provider in enumerate(self.providers):
is_primary = (i == 0)
if is_primary and self._fallback_mode:
if time.time() - self._last_primary_failure < COOLDOWN_SECONDS:
continue
logger.info("Primary cooldown expired, retrying...")
try:
result = self._call(provider, system_prompt, user_prompt,
temperature, max_tokens, json_mode, thinking, prefill)
if is_primary and self._fallback_mode:
logger.info("Primary provider recovered, switched back")
self._fallback_mode = False
return result
except requests.exceptions.RequestException as e:
logger.warning("Provider %s (%s) failed: %s",
provider.get('model', '?'),
provider.get('api_url', '?'), e)
if is_primary:
self._fallback_mode = True
self._last_primary_failure = time.time()
logger.info("Switched to fallback provider (cooldown: %ds)", COOLDOWN_SECONDS)
continue
except ValueError as e:
# json_mode returned non-JSON content (reasoning text from proxy)
logger.warning("Provider %s returned non-JSON for json_mode, trying next: %s",
provider.get('model', '?'), str(e)[:100])
if is_primary:
self._fallback_mode = True
self._last_primary_failure = time.time()
continue
logger.error("All LLM providers failed")
return None
def _call(self, provider: dict, system_prompt: str, user_prompt: str,
temperature: float, max_tokens: int, json_mode: bool,
thinking: Optional[dict] = None, prefill: str = None) -> str:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {provider['api_key']}",
}
payload = {
"model": provider["model"],
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
"temperature": temperature,
"max_tokens": max_tokens,
}
if json_mode:
payload["response_format"] = {"type": "json_object"}
if thinking is not None:
payload["thinking"] = thinking
if prefill:
payload["messages"].append({"role": "assistant", "content": prefill})
resp = requests.post(
provider["api_url"], headers=headers, json=payload, timeout=60
)
resp.raise_for_status()
data = resp.json()
msg = data["choices"][0]["message"]
reasoning = msg.get("reasoning_content", "") or ""
content = msg.get("content") or ""
# If content is empty but reasoning has content, use reasoning as fallback
if not content and reasoning:
content = reasoning
if reasoning and content:
logger.debug("LLM: reasoning=%d chars, content=%d chars", len(reasoning), len(content))
# Detect when json_mode was requested but response is not JSON.
# This catches two cases:
# 1. Proxy merged reasoning into content field (reasoning is empty, content starts with analysis text)
# 2. Model spent all tokens on reasoning (content was empty, reasoning was used as fallback)
if json_mode and content:
stripped = content.lstrip()
if stripped and stripped[0] not in ('[', '{', '"'):
raise ValueError(
"json_mode response is not JSON (starts with: %.50r)" % stripped
)
if json_mode and not content:
logger.warning("JSON mode returned empty content for %s", provider.get("model", "?"))
raise ValueError("JSON mode returned empty content")
return content
def get_status(self) -> dict:
status = []
for p in self.providers:
status.append({
"model": p["model"],
"priority": p.get("priority"),
"api_url": p["api_url"],
})
return {
"fallback_mode": self._fallback_mode,
"last_primary_failure_ago":
int(time.time() - self._last_primary_failure) if self._fallback_mode else 0,
"providers": status,
}
@property
def current_provider(self) -> str:
if not self.providers:
return "none"
idx = 1 if (self._fallback_mode and len(self.providers) > 1) else 0
return self.providers[idx].get("model", "unknown")
memory_system.py
+423
View File
@@ -0,0 +1,423 @@
"""Core MemorySystem - Orchestrates all storage, search, and lifecycle operations"""
import logging
from typing import Optional, List, Dict, Any
from config import Config
from embedding_service import EmbeddingService
from storage.mysql_store import MySQLStore
from storage.redis_cache import RedisCache
from storage.vector_search import embedding_to_bytes, search_similar, find_duplicates
from lifecycle.compressor import MemoryCompressor
from llm_client import LLMClient
from lifecycle.pipeline import Pipeline
from lifecycle.aggregator import ScenarioAggregator
from lifecycle.persona_generator import PersonaGenerator
from lifecycle.fact_extractor import FactExtractor
from lifecycle.cleaner import MemoryCleaner
logger = logging.getLogger(__name__)
class MemorySystem:
"""Multi-agent memory system with team isolation."""
def __init__(self, config: Config = None):
cfg = config or Config()
self.mysql = MySQLStore(
host=cfg.MYSQL_HOST,
port=cfg.MYSQL_PORT,
user=cfg.MYSQL_USER,
password=cfg.MYSQL_PASSWORD,
database=cfg.MYSQL_DATABASE,
unix_socket=cfg.MYSQL_UNIX_SOCKET,
)
self.redis = RedisCache(
host=cfg.REDIS_HOST,
port=cfg.REDIS_PORT,
password=cfg.REDIS_PASSWORD,
db=cfg.REDIS_DB,
)
self.embedder = EmbeddingService(base_url=cfg.EMBEDDING_SERVICE_URL)
self.compressor = MemoryCompressor(self.mysql, self.redis, self.embedder)
self.cleaner = MemoryCleaner(self.mysql)
self.default_ttl = cfg.WORKING_MEMORY_TTL
self.llm = LLMClient(providers_config=cfg.llm_providers)
self.aggregator = ScenarioAggregator(self.mysql, self.llm)
self.persona_gen = PersonaGenerator(self.mysql, self.llm)
self.fact_extractor = FactExtractor(self.mysql, self.redis, self.llm, self.embedder)
self.pipeline = Pipeline(self.mysql, self.redis, self.compressor, self.cleaner)
# ── Team management ──────────────────────────────────────────
def create_team(self, team_id: str, name: str, description: str = "",
config: dict = None) -> dict:
return self.mysql.create_team(team_id, name, description, config)
def delete_team(self, team_id: str) -> bool:
return self.mysql.delete_team(team_id)
def get_team(self, team_id: str) -> Optional[dict]:
return self.mysql.get_team(team_id)
# ── Agent management ─────────────────────────────────────────
def create_agent(self, agent_id: str, team_id: str, name: str,
role: str = "") -> dict:
# Verify team exists
if not self.mysql.get_team(team_id):
raise ValueError(f"Team {team_id} not found")
return self.mysql.create_agent(agent_id, team_id, name, role)
def delete_agent(self, agent_id: str) -> bool:
return self.mysql.delete_agent(agent_id)
def get_agent(self, agent_id: str) -> Optional[dict]:
return self.mysql.get_agent(agent_id)
def get_personal_memory(self, memory_id: str) -> Optional[dict]:
return self.mysql.get_personal_memory(memory_id)
def get_team_memory(self, memory_id: str) -> Optional[dict]:
return self.mysql.get_team_memory(memory_id)
def get_agents_by_team(self, team_id: str) -> list:
return self.mysql.get_agents_by_team(team_id)
# ── Personal memories ────────────────────────────────────────
def add_personal_memory(self, agent_id: str, content: str,
importance: float = 0.5,
metadata: dict = None,
enable_dedup: bool = True,
dedup_threshold: float = 0.85) -> dict:
agent = self.mysql.get_agent(agent_id)
if not agent:
raise ValueError(f"Agent {agent_id} not found")
# Generate embedding
embedding = self.embedder.embed(content)
emb_bytes = embedding_to_bytes(embedding)
# Dedup check: compare against existing memories
if enable_dedup:
stored = self.mysql.get_personal_memories_with_embeddings(agent_id)
if stored:
dupes = find_duplicates(embedding, stored, threshold=dedup_threshold)
if dupes:
best_id, best_score, _ = dupes[0]
logger.info(f"Dedup: memory {best_id} already covers this (score={best_score:.4f}), skipping")
existing = self.mysql.get_personal_memory(best_id)
if existing:
existing["dedup_skipped"] = True
existing["dedup_score"] = round(best_score, 4)
return existing
return self.mysql.add_personal_memory(
agent_id=agent_id,
team_id=agent["team_id"],
content=content,
embedding=emb_bytes,
importance=importance,
metadata=metadata,
)
def search_personal_memories(self, agent_id: str, query: str,
limit: int = 10,
min_score: float = 0.0,
max_chars_per_memory: int = 0,
max_total_chars: int = 0) -> list:
# Generate query embedding
query_embedding = self.embedder.embed(query)
# Get all memories with embeddings for this agent
stored = self.mysql.get_personal_memories_with_embeddings(agent_id)
if not stored:
return []
# Vector search
results = search_similar(query_embedding, stored, top_k=limit, min_score=min_score)
# Format output with char limits
output = []
total_chars = 0
for memory_id, score, item in results:
self.mysql.touch_personal_memory(memory_id)
content = item["content"]
# Truncate single memory if needed
if max_chars_per_memory > 0 and len(content) > max_chars_per_memory:
content = content[:max_chars_per_memory] + "..."
# Check total budget
if max_total_chars > 0:
if total_chars + len(content) > max_total_chars:
remaining = max_total_chars - total_chars
if remaining > 50: # only include if meaningful
content = content[:remaining] + "..."
else:
break
total_chars += len(content)
output.append({
"id": memory_id,
"content": content,
"score": round(score, 4),
"importance": item.get("importance", 0.5),
"metadata": item.get("metadata"),
})
return output
def get_recent_personal_memories(self, agent_id: str,
limit: int = 20) -> list:
rows = self.mysql.get_personal_memories_by_agent(agent_id, limit)
output = [
{
"id": r["id"],
"content": r["content"],
"importance": r.get("importance", 0.5),
"metadata": r.get("metadata"),
"created_at": str(r.get("created_at", "")),
}
for r in rows
]
for item in output:
self.mysql.touch_personal_memory(item["id"])
return output
def update_personal_memory(self, memory_id: str, content: str = None,
importance: float = None,
metadata: dict = None) -> bool:
embedding = None
if content is not None:
emb = self.embedder.embed(content)
embedding = embedding_to_bytes(emb)
result = self.mysql.update_personal_memory(
memory_id, content=content, importance=importance,
metadata=metadata, embedding=embedding,
)
if result:
self.mysql.touch_personal_memory(memory_id)
return result
def delete_personal_memory(self, memory_id: str) -> bool:
return self.mysql.delete_personal_memory(memory_id)
# ── Working memory ───────────────────────────────────────────
def add_working_memory(self, agent_id: str, content: str,
ttl: int = None) -> dict:
ttl = ttl or self.default_ttl
return self.redis.add_working_memory(agent_id, content, ttl=ttl)
def get_working_memories(self, agent_id: str,
limit: int = 20) -> list:
return self.redis.get_working_memories(agent_id, limit)
def clear_working_memory(self, agent_id: str) -> bool:
return self.redis.clear_working_memory(agent_id)
# ── Team memories ────────────────────────────────────────────
def add_team_memory(self, team_id: str, content: str,
importance: float = 0.5, category: str = "general",
metadata: dict = None,
enable_dedup: bool = True,
dedup_threshold: float = 0.85) -> dict:
if not self.mysql.get_team(team_id):
raise ValueError(f"Team {team_id} not found")
embedding = self.embedder.embed(content)
emb_bytes = embedding_to_bytes(embedding)
# Dedup check
if enable_dedup:
stored = self.mysql.get_team_memories_with_embeddings(team_id)
if stored:
dupes = find_duplicates(embedding, stored, threshold=dedup_threshold)
if dupes:
best_id, best_score, _ = dupes[0]
logger.info(f"Dedup: team memory {best_id} already covers this (score={best_score:.4f}), skipping")
existing = self.mysql.get_team_memory(best_id)
if existing:
existing["dedup_skipped"] = True
existing["dedup_score"] = round(best_score, 4)
return existing
return self.mysql.add_team_memory(
team_id=team_id,
content=content,
embedding=emb_bytes,
importance=importance,
category=category,
metadata=metadata,
)
def search_team_memories(self, team_id: str, query: str,
limit: int = 10,
min_score: float = 0.0,
max_chars_per_memory: int = 0,
max_total_chars: int = 0) -> list:
query_embedding = self.embedder.embed(query)
stored = self.mysql.get_team_memories_with_embeddings(team_id)
if not stored:
return []
results = search_similar(query_embedding, stored, top_k=limit, min_score=min_score)
output = []
total_chars = 0
for memory_id, score, item in results:
self.mysql.touch_team_memory(memory_id)
content = item["content"]
if max_chars_per_memory > 0 and len(content) > max_chars_per_memory:
content = content[:max_chars_per_memory] + "..."
if max_total_chars > 0:
if total_chars + len(content) > max_total_chars:
remaining = max_total_chars - total_chars
if remaining > 50:
content = content[:remaining] + "..."
else:
break
total_chars += len(content)
output.append({
"id": memory_id,
"content": content,
"score": round(score, 4),
"importance": item.get("importance", 0.5),
"metadata": item.get("metadata"),
})
return output
def get_recent_team_memories(self, team_id: str,
limit: int = 20) -> list:
rows = self.mysql.get_team_memories_by_team(team_id, limit)
output = [
{
"id": r["id"],
"content": r["content"],
"importance": r.get("importance", 0.5),
"category": r.get("category", "general"),
"created_at": str(r.get("created_at", "")),
}
for r in rows
]
for item in output:
self.mysql.touch_team_memory(item["id"])
return output
def update_team_memory(self, memory_id: str, content: str = None,
importance: float = None) -> bool:
embedding = None
if content is not None:
emb = self.embedder.embed(content)
embedding = embedding_to_bytes(emb)
result = self.mysql.update_team_memory(
memory_id, content=content, importance=importance,
embedding=embedding,
)
if result:
self.mysql.touch_team_memory(memory_id)
return result
def delete_team_memory(self, memory_id: str) -> bool:
return self.mysql.delete_team_memory(memory_id)
# ── Lifecycle ────────────────────────────────────────────────
def compress_working_memories(self, agent_id: str, max_items: int = 20,
summary_callback=None) -> str:
return self.compressor.compress(agent_id, max_items, summary_callback)
# -- Pipeline automation --
# -- Scenario aggregation --
# -- Persona generation --
# -- Fact extraction --
def extract_facts(self, agent_id: str, team_id: str,
max_memories: int = 20, delete_after: bool = False):
return self.fact_extractor.extract(agent_id, team_id, max_memories, delete_after)
def generate_persona(self, agent_id: str, team_id: str, max_items: int = 30):
return self.persona_gen.generate(agent_id, team_id, max_items)
def get_persona(self, agent_id: str):
return self.persona_gen.get_persona(agent_id)
def delete_persona(self, persona_id: str, team_id: str):
return self.persona_gen.delete_persona(persona_id, team_id)
def aggregate_scenarios(self, agent_id: str, team_id: str, max_memories: int = 30):
return self.aggregator.aggregate(agent_id, team_id, max_memories)
def get_scenarios(self, agent_id: str):
return self.aggregator.get_scenarios(agent_id)
def delete_scenario(self, scenario_id: str, team_id: str):
return self.aggregator.delete_scenario(scenario_id, team_id)
def get_pipeline_config(self, team_id: str):
return self.pipeline.get_config(team_id)
def set_pipeline_config(self, team_id: str, **kwargs):
return self.pipeline.set_config(team_id, **kwargs)
def check_auto_compress(self, agent_id: str):
return self.pipeline.check_after_wm_write(agent_id)
def run_auto_cleanup(self, team_id: str = None):
return self.pipeline.run_cleanup(team_id)
def cleanup_memories(self, team_id: str = None,
max_age_days: int = 90,
min_importance: float = 0.2) -> dict:
return self.cleaner.cleanup(team_id, max_age_days, min_importance)
def rebuild_vector_index(self, team_id: str = None) -> dict:
"""Re-embed all memories (useful after model update)."""
# Get all personal memories
conn = self.mysql._get_conn()
with conn.cursor() as cur:
if team_id:
cur.execute("SELECT id, content FROM personal_memories WHERE team_id = %s", (team_id,))
else:
cur.execute("SELECT id, content FROM personal_memories")
personals = cur.fetchall()
if team_id:
cur.execute("SELECT id, content FROM team_memories WHERE team_id = %s", (team_id,))
else:
cur.execute("SELECT id, content FROM team_memories")
teams = cur.fetchall()
rebuilt = 0
for row in personals:
try:
emb = self.embedder.embed(row["content"])
self.mysql.update_personal_memory(row["id"], embedding=embedding_to_bytes(emb))
rebuilt += 1
except Exception as e:
logger.warning(f"Failed to rebuild personal memory {row['id']}: {e}")
for row in teams:
try:
emb = self.embedder.embed(row["content"])
self.mysql.update_team_memory(row["id"], embedding=embedding_to_bytes(emb))
rebuilt += 1
except Exception as e:
logger.warning(f"Failed to rebuild team memory {row['id']}: {e}")
return {"rebuilt": rebuilt, "total": len(personals) + len(teams)}
# ── Stats ────────────────────────────────────────────────────
def get_stats(self, team_id: str = None) -> dict:
stats = self.mysql.get_stats(team_id)
stats["redis_ok"] = self.redis.health_check()
stats["embedding_ok"] = self.embedder.health_check()
return stats
requirements.txt
+7
View File
@@ -0,0 +1,7 @@
flask==3.0.0
gunicorn==21.2.0
pymysql==1.1.0
redis==5.0.1
numpy==1.24.4
requests==2.31.0
cryptography==41.0.7
storage/__init__.py
+1
View File
@@ -0,0 +1 @@
# Storage layer
storage/mysql_store.py
+367
View File
@@ -0,0 +1,367 @@
"""MySQL Storage Layer"""
import uuid
import json
import logging
from datetime import datetime
from typing import Optional, List, Dict, Any
import pymysql
from pymysql.cursors import DictCursor
logger = logging.getLogger(__name__)
class MySQLStore:
"""MySQL-backed persistent storage for memories."""
def __init__(self, host, port, user, password, database, charset="utf8mb4", unix_socket=None):
self.conn_kwargs = dict(
host=host,
port=port,
user=user,
password=password,
database=database,
charset=charset,
cursorclass=DictCursor,
autocommit=True,
)
if unix_socket:
self.conn_kwargs["unix_socket"] = unix_socket
self._conn = None
def _get_conn(self):
if self._conn is None or not self._conn.open:
self._conn = pymysql.connect(**self.conn_kwargs)
try:
self._conn.ping(reconnect=True)
except Exception:
self._conn = pymysql.connect(**self.conn_kwargs)
return self._conn
def _query(self, sql, args=None, fetch=False):
conn = self._get_conn()
with conn.cursor() as cur:
cur.execute(sql, args)
if fetch:
return cur.fetchall()
return None
# ── Team management ──────────────────────────────────────────
def create_team(self, team_id: str, name: str, description: str = "", config: dict = None):
self._query(
"INSERT INTO teams (id, name, description, config) VALUES (%s, %s, %s, %s)",
(team_id, name, description, json.dumps(config) if config else None),
)
return self.get_team(team_id)
def get_team(self, team_id: str) -> Optional[dict]:
rows = self._query("SELECT * FROM teams WHERE id = %s", (team_id,), fetch=True)
return rows[0] if rows else None
def delete_team(self, team_id: str) -> bool:
self._query("DELETE FROM teams WHERE id = %s", (team_id,))
return True
# ── Agent management ─────────────────────────────────────────
def create_agent(self, agent_id: str, team_id: str, name: str, role: str = ""):
self._query(
"INSERT INTO agents (id, team_id, name, role) VALUES (%s, %s, %s, %s)",
(agent_id, team_id, name, role),
)
return self.get_agent(agent_id)
def get_agent(self, agent_id: str) -> Optional[dict]:
rows = self._query("SELECT * FROM agents WHERE id = %s", (agent_id,), fetch=True)
return rows[0] if rows else None
def delete_agent(self, agent_id: str) -> bool:
self._query("DELETE FROM agents WHERE id = %s", (agent_id,))
return True
def get_agents_by_team(self, team_id: str) -> list:
return self._query(
"SELECT * FROM agents WHERE team_id = %s ORDER BY created_at DESC",
(team_id,), fetch=True,
)
# ── Personal memories ────────────────────────────────────────
def add_personal_memory(self, agent_id: str, team_id: str, content: str,
embedding: bytes, importance: float = 0.5,
metadata: dict = None) -> dict:
mid = str(uuid.uuid4())[:16]
self._query(
"""INSERT INTO personal_memories
(id, agent_id, team_id, content, embedding, importance, metadata)
VALUES (%s, %s, %s, %s, %s, %s, %s)""",
(mid, agent_id, team_id, content, embedding, importance,
json.dumps(metadata) if metadata else None),
)
return self.get_personal_memory(mid)
def get_personal_memory(self, memory_id: str) -> Optional[dict]:
rows = self._query("SELECT * FROM personal_memories WHERE id = %s", (memory_id,), fetch=True)
return rows[0] if rows else None
def get_personal_memories_by_agent(self, agent_id: str, limit: int = 20) -> List[dict]:
return self._query(
"SELECT * FROM personal_memories WHERE agent_id = %s ORDER BY created_at DESC LIMIT %s",
(agent_id, limit), fetch=True,
)
def get_personal_memories_with_embeddings(self, agent_id: str) -> List[dict]:
"""Get all personal memories with embeddings for vector search."""
return self._query(
"SELECT id, content, embedding, importance, metadata FROM personal_memories WHERE agent_id = %s",
(agent_id,), fetch=True,
)
def update_personal_memory(self, memory_id: str, content: str = None,
importance: float = None, metadata: dict = None,
embedding: bytes = None) -> bool:
sets, args = [], []
if content is not None:
sets.append("content = %s")
args.append(content)
if importance is not None:
sets.append("importance = %s")
args.append(importance)
if metadata is not None:
sets.append("metadata = %s")
args.append(json.dumps(metadata))
if embedding is not None:
sets.append("embedding = %s")
args.append(embedding)
if not sets:
return False
args.append(memory_id)
self._query(f"UPDATE personal_memories SET {', '.join(sets)} WHERE id = %s", args)
return True
def touch_personal_memory(self, memory_id: str):
self._query(
"UPDATE personal_memories SET access_count = access_count + 1, last_accessed = NOW() WHERE id = %s",
(memory_id,),
)
def delete_personal_memory(self, memory_id: str) -> bool:
self._query("DELETE FROM personal_memories WHERE id = %s", (memory_id,))
return True
# ── Team memories ────────────────────────────────────────────
def add_team_memory(self, team_id: str, content: str, embedding: bytes,
importance: float = 0.5, category: str = "general",
metadata: dict = None) -> dict:
mid = str(uuid.uuid4())[:16]
self._query(
"""INSERT INTO team_memories
(id, team_id, content, embedding, importance, category, metadata)
VALUES (%s, %s, %s, %s, %s, %s, %s)""",
(mid, team_id, content, embedding, importance, category,
json.dumps(metadata) if metadata else None),
)
return self.get_team_memory(mid)
def get_team_memory(self, memory_id: str) -> Optional[dict]:
rows = self._query("SELECT * FROM team_memories WHERE id = %s", (memory_id,), fetch=True)
return rows[0] if rows else None
def get_team_memories_by_team(self, team_id: str, limit: int = 20) -> List[dict]:
return self._query(
"SELECT * FROM team_memories WHERE team_id = %s ORDER BY created_at DESC LIMIT %s",
(team_id, limit), fetch=True,
)
def get_team_memories_with_embeddings(self, team_id: str) -> List[dict]:
"""Get all team memories with embeddings for vector search."""
return self._query(
"SELECT id, content, embedding, importance, metadata FROM team_memories WHERE team_id = %s",
(team_id,), fetch=True,
)
def update_team_memory(self, memory_id: str, content: str = None,
importance: float = None, embedding: bytes = None) -> bool:
sets, args = [], []
if content is not None:
sets.append("content = %s")
args.append(content)
if importance is not None:
sets.append("importance = %s")
args.append(importance)
if embedding is not None:
sets.append("embedding = %s")
args.append(embedding)
if not sets:
return False
args.append(memory_id)
self._query(f"UPDATE team_memories SET {', '.join(sets)} WHERE id = %s", args)
return True
def touch_team_memory(self, memory_id: str):
self._query(
"UPDATE team_memories SET access_count = access_count + 1, last_accessed = NOW() WHERE id = %s",
(memory_id,),
)
def delete_team_memory(self, memory_id: str) -> bool:
self._query("DELETE FROM team_memories WHERE id = %s", (memory_id,))
return True
# ── Cleanup ──────────────────────────────────────────────────
def cleanup_personal_memories(self, team_id: str = None, max_age_days: int = 90,
min_importance: float = 0.2) -> int:
sql = ("DELETE FROM personal_memories WHERE importance < %s "
"AND COALESCE(last_accessed, created_at) < DATE_SUB(NOW(), INTERVAL %s DAY)")
args = [min_importance, max_age_days]
if team_id:
sql += " AND team_id = %s"
args.append(team_id)
conn = self._get_conn()
with conn.cursor() as cur:
cur.execute(sql, args)
return cur.rowcount
def cleanup_team_memories(self, team_id: str = None, max_age_days: int = 90,
min_importance: float = 0.2) -> int:
sql = ("DELETE FROM team_memories WHERE importance < %s "
"AND COALESCE(last_accessed, created_at) < DATE_SUB(NOW(), INTERVAL %s DAY)")
args = [min_importance, max_age_days]
if team_id:
sql += " AND team_id = %s"
args.append(team_id)
conn = self._get_conn()
with conn.cursor() as cur:
cur.execute(sql, args)
return cur.rowcount
# ── Stats ────────────────────────────────────────────────────
def get_stats(self, team_id: str = None) -> dict:
stats = {}
if team_id:
row = self._query(
"SELECT COUNT(*) as cnt FROM personal_memories WHERE team_id = %s",
(team_id,), fetch=True,
)
stats["personal_memories"] = row[0]["cnt"] if row else 0
row = self._query(
"SELECT COUNT(*) as cnt FROM team_memories WHERE team_id = %s",
(team_id,), fetch=True,
)
stats["team_memories"] = row[0]["cnt"] if row else 0
row = self._query(
"SELECT COUNT(*) as cnt FROM agents WHERE team_id = %s",
(team_id,), fetch=True,
)
stats["agents"] = row[0]["cnt"] if row else 0
else:
row = self._query("SELECT COUNT(*) as cnt FROM personal_memories", fetch=True)
stats["personal_memories"] = row[0]["cnt"] if row else 0
row = self._query("SELECT COUNT(*) as cnt FROM team_memories", fetch=True)
stats["team_memories"] = row[0]["cnt"] if row else 0
row = self._query("SELECT COUNT(*) as cnt FROM agents", fetch=True)
stats["agents"] = row[0]["cnt"] if row else 0
row = self._query("SELECT COUNT(*) as cnt FROM teams", fetch=True)
stats["teams"] = row[0]["cnt"] if row else 0
return stats
# -- Pipeline config --
def get_pipeline_config(self, team_id: str):
rows = self._query(
"SELECT * FROM pipeline_config WHERE team_id = %s",
(team_id,), fetch=True,
)
return rows[0] if rows else None
def create_pipeline_config(self, team_id: str, compress_every_n: int = 0,
compress_target_count: int = 5,
cleanup_idle_days: int = 0,
cleanup_min_importance: float = 0.2,
enabled: bool = True,
warmup_max_memories: int = 0,
warmup_compress_every_n: int = 1):
self._query(
"INSERT INTO pipeline_config (team_id, compress_every_n, compress_target_count, cleanup_idle_days, cleanup_min_importance, enabled, warmup_max_memories, warmup_compress_every_n) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)",
(team_id, compress_every_n, compress_target_count, cleanup_idle_days, cleanup_min_importance, enabled, warmup_max_memories, warmup_compress_every_n),
)
def update_pipeline_config(self, team_id: str, **kwargs):
sets, args = [], []
for field in ["compress_every_n", "compress_target_count", "cleanup_idle_days", "cleanup_min_importance", "enabled", "warmup_max_memories", "warmup_compress_every_n"]:
if field in kwargs and kwargs[field] is not None:
sets.append(f"{field} = %s")
args.append(kwargs[field])
if not sets:
return
args.append(team_id)
self._query(f"UPDATE pipeline_config SET {', '.join(sets)} WHERE team_id = %s", tuple(args))
def get_all_pipeline_configs(self):
return self._query("SELECT * FROM pipeline_config WHERE enabled = 1", fetch=True) or []
# -- Scenario management --
def add_scenario(self, team_id: str, agent_id: str, name: str,
summary: str, memory_ids: list) -> str:
import uuid
sid = str(uuid.uuid4())[:16]
self._query(
"INSERT INTO memory_scenarios (id, team_id, agent_id, name, summary, memory_ids) VALUES (%s, %s, %s, %s, %s, %s)",
(sid, team_id, agent_id, name, summary, json.dumps(memory_ids)),
)
return sid
def get_scenario(self, scenario_id: str):
rows = self._query("SELECT * FROM memory_scenarios WHERE id = %s", (scenario_id,), fetch=True)
return rows[0] if rows else None
def get_scenarios_by_agent(self, agent_id: str):
rows = self._query("SELECT * FROM memory_scenarios WHERE agent_id = %s ORDER BY created_at DESC", (agent_id,), fetch=True)
for row in rows:
if isinstance(row.get("memory_ids"), str):
row["memory_ids"] = json.loads(row["memory_ids"])
return rows or []
def delete_scenario(self, scenario_id: str) -> bool:
self._query("DELETE FROM memory_scenarios WHERE id = %s", (scenario_id,))
return True
# -- Persona management --
def add_persona(self, team_id, agent_id, preferences, habits, expertise, communication_style, summary):
import uuid
pid = str(uuid.uuid4())[:16]
self._query(
'INSERT INTO user_personas (id, team_id, agent_id, preferences, habits, expertise, communication_style, summary) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)',
(pid, team_id, agent_id, json.dumps(preferences), json.dumps(habits), json.dumps(expertise), communication_style, summary),
)
return pid
def get_persona(self, persona_id):
rows = self._query('SELECT * FROM user_personas WHERE id = %s', (persona_id,), fetch=True)
if rows:
row = rows[0]
for field in ['preferences', 'habits', 'expertise']:
if isinstance(row.get(field), str):
row[field] = json.loads(row[field])
return row
return None
def get_persona_by_agent(self, agent_id):
rows = self._query('SELECT * FROM user_personas WHERE agent_id = %s ORDER BY created_at DESC LIMIT 1', (agent_id,), fetch=True)
if rows:
row = rows[0]
for field in ['preferences', 'habits', 'expertise']:
if isinstance(row.get(field), str):
row[field] = json.loads(row[field])
return row
return None
def delete_persona(self, persona_id):
self._query('DELETE FROM user_personas WHERE id = %s', (persona_id,))
return True
storage/redis_cache.py
+68
View File
@@ -0,0 +1,68 @@
"""Redis Cache Layer for Working Memory"""
import json
import time
import logging
from typing import List, Optional, Dict
import redis
logger = logging.getLogger(__name__)
class RedisCache:
"""Redis-backed working memory cache with TTL."""
def __init__(self, host: str = "127.0.0.1", port: int = 6379,
password: str = "", db: int = 0):
self.client = redis.Redis(
host=host, port=port, password=password, db=db,
decode_responses=True, socket_timeout=5,
)
def _key(self, agent_id: str) -> str:
return f"wm:{agent_id}"
def add_working_memory(self, agent_id: str, content: str,
metadata: dict = None, ttl: int = 259200) -> dict:
"""Add an item to agent's working memory (Redis List)."""
item = {
"content": content,
"metadata": metadata or {},
"timestamp": time.time(),
}
key = self._key(agent_id)
self.client.lpush(key, json.dumps(item))
self.client.expire(key, ttl)
return item
def get_working_memories(self, agent_id: str, limit: int = 20) -> List[dict]:
"""Get recent working memory items for an agent."""
key = self._key(agent_id)
items = self.client.lrange(key, 0, limit - 1)
result = []
for item_str in items:
try:
parsed = json.loads(item_str)
if isinstance(parsed, dict):
result.append(parsed)
except json.JSONDecodeError:
continue
return result
def clear_working_memory(self, agent_id: str) -> bool:
"""Clear all working memory for an agent."""
key = self._key(agent_id)
self.client.delete(key)
return True
def get_working_memory_count(self, agent_id: str) -> int:
"""Get count of working memory items."""
key = self._key(agent_id)
return self.client.llen(key)
def health_check(self) -> bool:
"""Check Redis connectivity."""
try:
return self.client.ping()
except Exception:
return False
storage/vector_search.py
+113
View File
@@ -0,0 +1,113 @@
"""Vector Search Layer - Pure NumPy Implementation"""
import struct
import logging
from typing import List, Tuple, Optional
import numpy as np
logger = logging.getLogger(__name__)
def embedding_to_bytes(embedding: List[float]) -> bytes:
"""Convert float list to compact bytes for MySQL BLOB storage."""
return struct.pack(f"{len(embedding)}f", *embedding)
def bytes_to_embedding(data: bytes) -> np.ndarray:
"""Convert bytes back to numpy array."""
n = len(data) // 4
return np.array(struct.unpack(f"{n}f", data), dtype=np.float32)
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Compute cosine similarity between two vectors."""
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)
if norm_a == 0 or norm_b == 0:
return 0.0
return float(np.dot(a, b) / (norm_a * norm_b))
def search_similar(
query_embedding: List[float],
stored_items: List[dict],
top_k: int = 10,
min_score: float = 0.0,
) -> List[Tuple[str, float, dict]]:
"""
Search for similar items using cosine similarity.
Args:
query_embedding: Query vector
stored_items: List of dicts with 'id', 'embedding' (bytes), 'content'
top_k: Max results
min_score: Minimum similarity score
Returns:
List of (id, score, item_dict) tuples sorted by score descending
"""
if not stored_items:
return []
query_vec = np.array(query_embedding, dtype=np.float32)
query_norm = np.linalg.norm(query_vec)
if query_norm == 0:
return []
scores = []
for item in stored_items:
emb_bytes = item.get("embedding")
if emb_bytes is None:
continue
try:
stored_vec = bytes_to_embedding(emb_bytes)
score = float(np.dot(query_vec, stored_vec) / (query_norm * np.linalg.norm(stored_vec)))
if score >= min_score:
scores.append((item["id"], score, item))
except Exception as e:
logger.warning(f"Failed to compute similarity for item {item.get('id')}: {e}")
continue
scores.sort(key=lambda x: x[1], reverse=True)
return scores[:top_k]
def find_duplicates(
new_embedding: List[float],
stored_items: List[dict],
threshold: float = 0.85,
) -> List[Tuple[str, float, dict]]:
"""
Find items whose similarity to new_embedding exceeds threshold.
Used for deduplication before storing new memories.
Args:
new_embedding: The embedding of the new memory to check
stored_items: Existing memories with 'id', 'embedding' (bytes), 'content'
threshold: Similarity threshold above which items are considered duplicates
Returns:
List of (id, score, item_dict) for items exceeding threshold, sorted by score desc
"""
if not stored_items:
return []
query_vec = np.array(new_embedding, dtype=np.float32)
query_norm = np.linalg.norm(query_vec)
if query_norm == 0:
return []
duplicates = []
for item in stored_items:
emb_bytes = item.get("embedding")
if emb_bytes is None:
continue
try:
stored_vec = bytes_to_embedding(emb_bytes)
score = float(np.dot(query_vec, stored_vec) / (query_norm * np.linalg.norm(stored_vec)))
if score >= threshold:
duplicates.append((item["id"], score, item))
except Exception as e:
continue
duplicates.sort(key=lambda x: x[1], reverse=True)
return duplicates
tests/test_memory.py
+290
View File
@@ -0,0 +1,290 @@
#!/usr/bin/env python3
"""
Memory System Validation Tests
Tests:
1. Create two teams (team_a, team_b)
2. Create agents in each team
3. Write personal memories with team-specific content
4. Semantic search with team isolation verification
5. Performance benchmarks (write latency, search latency, memory usage)
"""
import sys
import os
import time
import json
import statistics
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import Config
from memory_system import MemorySystem
# ANSI colors
GREEN = "\033[92m"
RED = "\033[91m"
YELLOW = "\033[93m"
CYAN = "\033[96m"
RESET = "\033[0m"
BOLD = "\033[1m"
def header(text):
print(f"\n{BOLD}{CYAN}{'='*60}{RESET}")
print(f"{BOLD}{CYAN}{text}{RESET}")
print(f"{BOLD}{CYAN}{'='*60}{RESET}")
def ok(msg):
print(f" {GREEN}✓{RESET} {msg}")
def fail(msg):
print(f" {RED}✗{RESET} {msg}")
def warn(msg):
print(f" {YELLOW}⚠{RESET} {msg}")
def measure(label, func, *args, **kwargs):
"""Run func and return (result, elapsed_ms)."""
start = time.perf_counter()
result = func(*args, **kwargs)
elapsed = (time.perf_counter() - start) * 1000
return result, elapsed
def main():
cfg = Config()
ms = MemorySystem(cfg)
# ── Health checks ────────────────────────────────────────────
header("Health Checks")
if ms.embedder.health_check():
ok("BGE embedding service is alive")
else:
fail("BGE embedding service not responding")
sys.exit(1)
if ms.redis.health_check():
ok("Redis connection OK")
else:
fail("Redis connection failed")
sys.exit(1)
ok("MySQL connection OK (implicit via queries)")
# ── Step 1: Create teams ─────────────────────────────────────
header("Step 1: Create Teams")
# Clean up first (idempotent)
try:
ms.delete_team("team_a")
ms.delete_team("team_b")
except:
pass
team_a, elapsed = measure("create_team", ms.create_team, "team_a", "Team Alpha", "First test team")
ok(f"Created team_a: {team_a['name']} ({elapsed:.1f}ms)")
team_b, elapsed = measure("create_team", ms.create_team, "team_b", "Team Beta", "Second test team")
ok(f"Created team_b: {team_b['name']} ({elapsed:.1f}ms)")
# ── Step 2: Create agents ────────────────────────────────────
header("Step 2: Create Agents")
agent_a, elapsed = measure("create_agent", ms.create_agent, "agent_alpha_1", "team_a", "Alpha Agent", "developer")
ok(f"Created agent_alpha_1 in team_a ({elapsed:.1f}ms)")
agent_b, elapsed = measure("create_agent", ms.create_agent, "agent_beta_1", "team_b", "Beta Agent", "developer")
ok(f"Created agent_beta_1 in team_b ({elapsed:.1f}ms)")
# ── Step 3: Write personal memories ──────────────────────────
header("Step 3: Write Personal Memories")
write_latencies = []
mem_a, elapsed = measure("add_personal_memory", ms.add_personal_memory,
"agent_alpha_1", "用户的编程偏好是 Python，喜欢用 Flask 和 FastAPI 框架",
importance=0.8, metadata={"topic": "preference"})
write_latencies.append(elapsed)
ok(f"team_a memory: '{mem_a['content'][:40]}...' ({elapsed:.1f}ms)")
mem_b, elapsed = measure("add_personal_memory", ms.add_personal_memory,
"agent_beta_1", "用户的编程偏好是 Java，喜欢用 Spring Boot 框架",
importance=0.8, metadata={"topic": "preference"})
write_latencies.append(elapsed)
ok(f"team_b memory: '{mem_b['content'][:40]}...' ({elapsed:.1f}ms)")
# Add more memories for richer testing
for i, content in enumerate([
"团队每周一开例会",
"项目截止日期是下个月底",
"用户不喜欢加班",
]):
_, elapsed = measure("add_personal_memory", ms.add_personal_memory,
"agent_alpha_1", content, importance=0.5 + i * 0.1)
write_latencies.append(elapsed)
for i, content in enumerate([
"团队使用 Jenkins 做 CI/CD",
"数据库用的是 PostgreSQL",
"代码审查需要至少两人通过",
]):
_, elapsed = measure("add_personal_memory", ms.add_personal_memory,
"agent_beta_1", content, importance=0.5 + i * 0.1)
write_latencies.append(elapsed)
ok(f"Wrote {len(write_latencies)} memories total")
# ── Step 4: Write team memories ──────────────────────────────
header("Step 4: Write Team Shared Memories")
tm_a, elapsed = measure("add_team_memory", ms.add_team_memory,
"team_a", "团队技术栈: Python, Flask, MySQL, Redis",
importance=0.9, category="tech_stack")
write_latencies.append(elapsed)
ok(f"team_a shared: '{tm_a['content'][:40]}...' ({elapsed:.1f}ms)")
tm_b, elapsed = measure("add_team_memory", ms.add_team_memory,
"team_b", "团队技术栈: Java, Spring Boot, PostgreSQL, Kafka",
importance=0.9, category="tech_stack")
write_latencies.append(elapsed)
ok(f"team_b shared: '{tm_b['content'][:40]}...' ({elapsed:.1f}ms)")
# ── Step 5: Semantic search with isolation ───────────────────
header("Step 5: Semantic Search + Team Isolation")
search_latencies = []
# Search team_a for "编程偏好"
results_a, elapsed = measure("search_personal_memories", ms.search_personal_memories,
"agent_alpha_1", "编程偏好", limit=5)
search_latencies.append(elapsed)
ok(f"team_a search '编程偏好' → {len(results_a)} results ({elapsed:.1f}ms)")
for r in results_a:
print(f" score={r['score']:.4f} | {r['content'][:60]}")
# Search team_b for "编程偏好"
results_b, elapsed = measure("search_personal_memories", ms.search_personal_memories,
"agent_beta_1", "编程偏好", limit=5)
search_latencies.append(elapsed)
ok(f"team_b search '编程偏好' → {len(results_b)} results ({elapsed:.1f}ms)")
for r in results_b:
print(f" score={r['score']:.4f} | {r['content'][:60]}")
# Verify isolation
header("Step 6: Cross-Team Isolation Verification")
a_contents = [r["content"] for r in results_a]
b_contents = [r["content"] for r in results_b]
python_in_a = any("Python" in c for c in a_contents)
java_in_a = any("Java" in c for c in a_contents)
python_in_b = any("Python" in c for c in b_contents)
java_in_b = any("Java" in c for c in b_contents)
if python_in_a and not java_in_a:
ok("team_a returns Python (not Java) ✓")
elif not a_contents:
warn("team_a search returned no results")
else:
fail(f"team_a isolation issue: found Java={java_in_a}, Python={python_in_a}")
if java_in_b and not python_in_b:
ok("team_b returns Java (not Python) ✓")
elif not b_contents:
warn("team_b search returned no results")
else:
fail(f"team_b isolation issue: found Python={python_in_b}, Java={java_in_b}")
# Also verify team shared memories
team_search_a, elapsed = measure("search_team_memories", ms.search_team_memories,
"team_a", "技术栈", limit=3)
search_latencies.append(elapsed)
ok(f"team_a shared search '技术栈' → {len(team_search_a)} results ({elapsed:.1f}ms)")
for r in team_search_a:
print(f" score={r['score']:.4f} | {r['content'][:60]}")
team_search_b, elapsed = measure("search_team_memories", ms.search_team_memories,
"team_b", "技术栈", limit=3)
search_latencies.append(elapsed)
ok(f"team_b shared search '技术栈' → {len(team_search_b)} results ({elapsed:.1f}ms)")
for r in team_search_b:
print(f" score={r['score']:.4f} | {r['content'][:60]}")
# ── Step 7: Working memory ──────────────────────────────────
header("Step 7: Working Memory (Redis)")
wm, elapsed = measure("add_working_memory", ms.add_working_memory,
"agent_alpha_1", "刚刚和用户讨论了部署方案")
ok(f"Added working memory ({elapsed:.1f}ms)")
wm_items, elapsed = measure("get_working_memories", ms.get_working_memories,
"agent_alpha_1", limit=10)
ok(f"Retrieved {len(wm_items)} working memories ({elapsed:.1f}ms)")
# ── Step 8: Recent memories ─────────────────────────────────
header("Step 8: Recent Memories")
recent_a, elapsed = measure("get_recent_personal_memories", ms.get_recent_personal_memories,
"agent_alpha_1", limit=5)
ok(f"Recent personal memories for agent_alpha_1: {len(recent_a)} ({elapsed:.1f}ms)")
recent_team, elapsed = measure("get_recent_team_memories", ms.get_recent_team_memories,
"team_a", limit=5)
ok(f"Recent team memories for team_a: {len(recent_team)} ({elapsed:.1f}ms)")
# ── Performance Report ──────────────────────────────────────
header("Performance Report")
write_p50 = statistics.median(write_latencies)
write_p99 = sorted(write_latencies)[int(len(write_latencies) * 0.99)] if len(write_latencies) > 1 else write_latencies[0]
write_max = max(write_latencies)
search_p50 = statistics.median(search_latencies)
search_p99 = sorted(search_latencies)[int(len(search_latencies) * 0.99)] if len(search_latencies) > 1 else search_latencies[0]
search_max = max(search_latencies)
print(f" Write latency: P50={write_p50:.1f}ms P99={write_p99:.1f}ms Max={write_max:.1f}ms")
print(f" Search latency: P50={search_p50:.1f}ms P99={search_p99:.1f}ms Max={search_max:.1f}ms")
if write_p99 < 100:
ok("Write P99 < 100ms ✓")
else:
fail(f"Write P99 = {write_p99:.1f}ms (target < 100ms)")
if search_p99 < 50:
ok("Search P99 < 50ms ✓")
else:
warn(f"Search P99 = {search_p99:.1f}ms (target < 50ms, but embedding call dominates)")
# Memory usage
try:
import resource
rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 # Linux: KB → MB
print(f" Process RSS: {rss_mb:.0f} MB")
if rss_mb < 1536:
ok("RSS < 1.5 GiB ✓")
else:
fail(f"RSS = {rss_mb:.0f} MB (target < 1536 MB)")
except:
warn("Could not measure RSS")
# Stats
header("System Stats")
stats, elapsed = measure("get_stats", ms.get_stats)
ok(f"Stats: {json.dumps(stats, indent=2)}")
# ── Summary ─────────────────────────────────────────────────
header("TEST SUMMARY")
print(f" Teams created: 2")
print(f" Agents created: 2")
print(f" Memories written: {len(write_latencies)}")
print(f" Searches executed: {len(search_latencies)}")
print(f" Team isolation: {'PASS' if (python_in_a and java_in_b and not java_in_a and not python_in_b) else 'CHECK MANUALLY'}")
print(f" Write P99: {write_p99:.1f}ms {'PASS' if write_p99 < 100 else 'WARN'}")
print(f" Search P99: {search_p99:.1f}ms {'PASS' if search_p99 < 50 else 'WARN'}")
print()
print(f"{GREEN}{BOLD}All validation tests completed!{RESET}")
if __name__ == "__main__":
main()