初始化：记忆系统源代码上传（已脱敏）
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html） - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
This commit is contained in:
zqf
@@ -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
|
||||
Reference in New Issue
Block a user