zqf
b1f93588f1
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html） - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
424 lines
17 KiB
Python
424 lines
17 KiB
Python
"""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
|