439 lines
14 KiB
Python
439 lines
14 KiB
Python
"""
|
||
内嵌图数据库 - 基于SQLite实现
|
||
无需Docker，开箱即用
|
||
"""
|
||
|
||
import sqlite3
|
||
import json
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Any
|
||
|
||
|
||
class EmbeddedGraphDB:
|
||
"""内嵌图数据库 - SQLite实现"""
|
||
|
||
def __init__(self, db_path: str = "graph_memory.db"):
|
||
"""
|
||
初始化数据库
|
||
|
||
Args:
|
||
db_path: 数据库文件路径
|
||
"""
|
||
self.db_path = Path(db_path)
|
||
self.conn = None
|
||
self._init_db()
|
||
|
||
def _init_db(self):
|
||
"""初始化数据库表"""
|
||
self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
|
||
self.conn.row_factory = sqlite3.Row
|
||
|
||
cursor = self.conn.cursor()
|
||
|
||
# 创建实体表
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS entities (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT UNIQUE NOT NULL,
|
||
type TEXT,
|
||
mention_count INTEGER DEFAULT 1,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
|
||
# 创建关系表
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS relations (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
source_id INTEGER NOT NULL,
|
||
target_id INTEGER NOT NULL,
|
||
relation_type TEXT NOT NULL,
|
||
confidence REAL DEFAULT 1.0,
|
||
status TEXT DEFAULT 'active',
|
||
session_id TEXT,
|
||
turn_id INTEGER,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
date_bucket TEXT,
|
||
superseded_by INTEGER,
|
||
FOREIGN KEY (source_id) REFERENCES entities(id),
|
||
FOREIGN KEY (target_id) REFERENCES entities(id)
|
||
)
|
||
""")
|
||
|
||
# 创建索引
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_entity_name ON entities(name)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_entity_type ON entities(type)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_relation_source ON relations(source_id)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_relation_target ON relations(target_id)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_relation_type ON relations(relation_type)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_relation_status ON relations(status)")
|
||
|
||
self.conn.commit()
|
||
|
||
def ensure_constraints(self):
|
||
"""确保约束（兼容Neo4j接口）"""
|
||
pass # SQLite自动处理
|
||
|
||
def recall(self, query_intent: str, seed_entities: List[str] = None,
|
||
depth: int = 2, time_range: Dict = None,
|
||
session_filter: str = None) -> Dict:
|
||
"""
|
||
检索相关记忆
|
||
|
||
Args:
|
||
query_intent: 查询关键词（逗号分隔）
|
||
seed_entities: 种子实体
|
||
depth: 搜索深度
|
||
time_range: 时间范围
|
||
session_filter: 会话过滤
|
||
|
||
Returns:
|
||
检索结果
|
||
"""
|
||
keywords = [w.strip().lower() for w in query_intent.replace(',', ' ').split() if w.strip()]
|
||
|
||
cursor = self.conn.cursor()
|
||
|
||
# 搜索实体
|
||
entities = []
|
||
entity_ids = set()
|
||
|
||
# 如果没有关键词，返回所有实体（用于"我们都聊过什么"这类问题）
|
||
if not keywords and not seed_entities:
|
||
cursor.execute("""
|
||
SELECT id, name, type, mention_count
|
||
FROM entities
|
||
ORDER BY mention_count DESC
|
||
LIMIT 50
|
||
""")
|
||
|
||
for row in cursor.fetchall():
|
||
entity_ids.add(row['id'])
|
||
entities.append({
|
||
'name': row['name'],
|
||
'type': row['type'] or 'unknown',
|
||
'mention_count': row['mention_count']
|
||
})
|
||
else:
|
||
# 有关键词，按关键词搜索
|
||
for keyword in keywords:
|
||
cursor.execute("""
|
||
SELECT id, name, type, mention_count
|
||
FROM entities
|
||
WHERE LOWER(name) LIKE ?
|
||
""", (f"%{keyword}%",))
|
||
|
||
for row in cursor.fetchall():
|
||
if row['id'] not in entity_ids:
|
||
entity_ids.add(row['id'])
|
||
entities.append({
|
||
'name': row['name'],
|
||
'type': row['type'] or 'unknown',
|
||
'mention_count': row['mention_count']
|
||
})
|
||
|
||
# 搜索关系
|
||
relations = []
|
||
|
||
if entity_ids:
|
||
placeholders = ','.join('?' * len(entity_ids))
|
||
|
||
query = f"""
|
||
SELECT r.id, e1.name as source, e2.name as target,
|
||
r.relation_type as type, r.confidence, r.session_id,
|
||
r.turn_id, r.created_at, r.status
|
||
FROM relations r
|
||
JOIN entities e1 ON r.source_id = e1.id
|
||
JOIN entities e2 ON r.target_id = e2.id
|
||
WHERE (r.source_id IN ({placeholders}) OR r.target_id IN ({placeholders}))
|
||
AND r.status = 'active'
|
||
"""
|
||
|
||
params = list(entity_ids) + list(entity_ids)
|
||
|
||
if session_filter:
|
||
query += " AND r.session_id = ?"
|
||
params.append(session_filter)
|
||
|
||
cursor.execute(query, params)
|
||
|
||
for row in cursor.fetchall():
|
||
relations.append({
|
||
'source': row['source'],
|
||
'target': row['target'],
|
||
'type': row['type'],
|
||
'confidence': row['confidence'],
|
||
'session_id': row['session_id'],
|
||
'turn_id': row['turn_id'],
|
||
'created_at': row['created_at'],
|
||
'status': row['status']
|
||
})
|
||
|
||
return {
|
||
"entities": entities,
|
||
"relations": relations,
|
||
"message": f"找到 {len(entities)} 个实体, {len(relations)} 条关系"
|
||
}
|
||
|
||
def commit(self, triplets: List[Dict], entity_types: Dict = None,
|
||
temporal_tag: str = None, session_id: str = None,
|
||
turn_id: int = None) -> Dict:
|
||
"""
|
||
写入记忆
|
||
|
||
Args:
|
||
triplets: 三元组列表
|
||
entity_types: 实体类型
|
||
temporal_tag: 时间标签
|
||
session_id: 会话ID
|
||
turn_id: 轮次ID
|
||
|
||
Returns:
|
||
写入结果
|
||
"""
|
||
cursor = self.conn.cursor()
|
||
|
||
created_entities = 0
|
||
created_relations = 0
|
||
|
||
for triplet in triplets:
|
||
subject = triplet.get('subject')
|
||
relation = triplet.get('relation')
|
||
obj = triplet.get('object')
|
||
confidence = triplet.get('confidence', 1.0)
|
||
|
||
if not all([subject, relation, obj]):
|
||
continue
|
||
|
||
# 创建或更新实体
|
||
for entity_name in [subject, obj]:
|
||
entity_type = entity_types.get(entity_name) if entity_types else None
|
||
|
||
cursor.execute("""
|
||
INSERT INTO entities (name, type)
|
||
VALUES (?, ?)
|
||
ON CONFLICT(name) DO UPDATE SET
|
||
mention_count = mention_count + 1,
|
||
updated_at = CURRENT_TIMESTAMP
|
||
""", (entity_name, entity_type))
|
||
|
||
if cursor.rowcount > 0:
|
||
created_entities += 1
|
||
|
||
# 获取实体ID
|
||
cursor.execute("SELECT id FROM entities WHERE name = ?", (subject,))
|
||
source_id = cursor.fetchone()['id']
|
||
|
||
cursor.execute("SELECT id FROM entities WHERE name = ?", (obj,))
|
||
target_id = cursor.fetchone()['id']
|
||
|
||
# 创建关系
|
||
date_bucket = datetime.now().strftime('%Y-%m-%d')
|
||
|
||
cursor.execute("""
|
||
INSERT INTO relations (
|
||
source_id, target_id, relation_type, confidence,
|
||
session_id, turn_id, date_bucket
|
||
)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
""", (source_id, target_id, relation, confidence,
|
||
session_id, turn_id, date_bucket))
|
||
|
||
created_relations += 1
|
||
|
||
self.conn.commit()
|
||
|
||
return {
|
||
"created_entities": created_entities,
|
||
"created_relations": created_relations,
|
||
"message": f"创建了 {created_entities} 个实体, {created_relations} 条关系"
|
||
}
|
||
|
||
def purge(self, criteria: Dict, mode: str = "soft",
|
||
new_relation: Dict = None) -> Dict:
|
||
"""
|
||
删除或修正记忆
|
||
|
||
Args:
|
||
criteria: 删除条件
|
||
mode: 删除模式 (soft/hard)
|
||
new_relation: 替代关系
|
||
|
||
Returns:
|
||
删除结果
|
||
"""
|
||
cursor = self.conn.cursor()
|
||
|
||
# 构建查询条件
|
||
conditions = []
|
||
params = []
|
||
|
||
if criteria.get('source'):
|
||
cursor.execute("SELECT id FROM entities WHERE name = ?", (criteria['source'],))
|
||
row = cursor.fetchone()
|
||
if row:
|
||
conditions.append("source_id = ?")
|
||
params.append(row['id'])
|
||
|
||
if criteria.get('target'):
|
||
cursor.execute("SELECT id FROM entities WHERE name = ?", (criteria['target'],))
|
||
row = cursor.fetchone()
|
||
if row:
|
||
conditions.append("target_id = ?")
|
||
params.append(row['id'])
|
||
|
||
if criteria.get('relation'):
|
||
conditions.append("relation_type = ?")
|
||
params.append(criteria['relation'])
|
||
|
||
if not conditions:
|
||
return {"deleted": 0, "message": "无删除条件"}
|
||
|
||
where_clause = " AND ".join(conditions)
|
||
|
||
if mode == "soft":
|
||
cursor.execute(f"""
|
||
UPDATE relations
|
||
SET status = 'deleted', updated_at = CURRENT_TIMESTAMP
|
||
WHERE {where_clause} AND status = 'active'
|
||
""", params)
|
||
else:
|
||
cursor.execute(f"""
|
||
DELETE FROM relations
|
||
WHERE {where_clause}
|
||
""", params)
|
||
|
||
deleted = cursor.rowcount
|
||
self.conn.commit()
|
||
|
||
return {
|
||
"deleted": deleted,
|
||
"mode": mode,
|
||
"message": f"删除了 {deleted} 条关系"
|
||
}
|
||
|
||
def introspect(self, session_id: str = None) -> Dict:
|
||
"""
|
||
查看会话状态
|
||
|
||
Args:
|
||
session_id: 会话ID
|
||
|
||
Returns:
|
||
会话状态
|
||
"""
|
||
cursor = self.conn.cursor()
|
||
|
||
# 统计实体
|
||
cursor.execute("SELECT COUNT(*) as count FROM entities")
|
||
entity_count = cursor.fetchone()['count']
|
||
|
||
# 统计关系
|
||
cursor.execute("SELECT COUNT(*) as count FROM relations WHERE status = 'active'")
|
||
relation_count = cursor.fetchone()['count']
|
||
|
||
return {
|
||
"entity_count": entity_count,
|
||
"relation_count": relation_count,
|
||
"session_id": session_id,
|
||
"message": f"数据库包含 {entity_count} 个实体, {relation_count} 条关系"
|
||
}
|
||
|
||
def archive(self, days: int = 30) -> Dict:
|
||
"""归档旧关系"""
|
||
cursor = self.conn.cursor()
|
||
|
||
cursor.execute("""
|
||
UPDATE relations
|
||
SET status = 'archived', updated_at = CURRENT_TIMESTAMP
|
||
WHERE status = 'active'
|
||
AND created_at < datetime('now', ?)
|
||
""", (f'-{days} days',))
|
||
|
||
archived = cursor.rowcount
|
||
self.conn.commit()
|
||
|
||
return {
|
||
"archived": archived,
|
||
"message": f"归档了 {archived} 条关系"
|
||
}
|
||
|
||
def cleanup(self, dry_run: bool = True) -> Dict:
|
||
"""清理已删除数据"""
|
||
cursor = self.conn.cursor()
|
||
|
||
if dry_run:
|
||
cursor.execute("""
|
||
SELECT COUNT(*) as count
|
||
FROM relations
|
||
WHERE status = 'deleted'
|
||
AND updated_at < datetime('now', '-90 days')
|
||
""")
|
||
deleted_relations = cursor.fetchone()['count']
|
||
|
||
return {
|
||
"dry_run": True,
|
||
"deleted_relations": deleted_relations,
|
||
"message": f"将删除 {deleted_relations} 条关系"
|
||
}
|
||
else:
|
||
cursor.execute("""
|
||
DELETE FROM relations
|
||
WHERE status = 'deleted'
|
||
AND updated_at < datetime('now', '-90 days')
|
||
""")
|
||
deleted = cursor.rowcount
|
||
self.conn.commit()
|
||
|
||
return {
|
||
"dry_run": False,
|
||
"deleted": deleted,
|
||
"message": f"删除了 {deleted} 条关系"
|
||
}
|
||
|
||
def close(self):
|
||
"""关闭数据库连接"""
|
||
if self.conn:
|
||
self.conn.close()
|
||
self.conn = None
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.close()
|
||
|
||
|
||
# 兼容性别名
|
||
Neo4jGraph = EmbeddedGraphDB
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 测试
|
||
print("Testing Embedded Graph Database...")
|
||
|
||
with EmbeddedGraphDB("test.db") as db:
|
||
# 写入测试
|
||
result = db.commit(
|
||
triplets=[
|
||
{"subject": "用户", "relation": "喜欢", "object": "Python"},
|
||
{"subject": "用户", "relation": "学习", "object": "AI"}
|
||
],
|
||
session_id="test-session",
|
||
turn_id=1
|
||
)
|
||
print(f"Commit: {result}")
|
||
|
||
# 检索测试
|
||
result = db.recall("Python,AI")
|
||
print(f"Recall: {result}")
|
||
|
||
# 状态测试
|
||
result = db.introspect()
|
||
print(f"Introspect: {result}")
|
||
|
||
print("\nTest completed!")
|