264 lines
6.9 KiB
Python
Executable File
264 lines
6.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
FTS5 记忆索引系统
|
|
用于全文搜索 memory/*.md 文件
|
|
"""
|
|
|
|
import sqlite3
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
|
|
# 配置
|
|
MEMORY_DIR = Path.home() / ".openclaw" / "workspace" / "memory"
|
|
DB_PATH = MEMORY_DIR / ".fts5" / "memory.db"
|
|
|
|
|
|
def init_db():
|
|
"""初始化 FTS5 数据库"""
|
|
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# 创建文件表
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
id INTEGER PRIMARY KEY,
|
|
path TEXT UNIQUE NOT NULL,
|
|
modified_at TIMESTAMP,
|
|
indexed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# 创建 FTS5 虚拟表用于全文搜索
|
|
cursor.execute("""
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS memory_index USING fts5(
|
|
content,
|
|
file_id,
|
|
chunk_index,
|
|
tokenize='porter unicode61'
|
|
)
|
|
""")
|
|
|
|
# 创建 chunks 表存储分块信息
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS chunks (
|
|
id INTEGER PRIMARY KEY,
|
|
file_id INTEGER,
|
|
chunk_index INTEGER,
|
|
content TEXT,
|
|
start_line INTEGER,
|
|
end_line INTEGER,
|
|
FOREIGN KEY (file_id) REFERENCES files(id)
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
print(f"✓ 数据库初始化完成: {DB_PATH}")
|
|
|
|
|
|
def chunk_content(content, chunk_size=500, overlap=100):
|
|
"""
|
|
将内容分块，带重叠以保持上下文
|
|
|
|
Args:
|
|
content: 文件内容
|
|
chunk_size: 每块字符数
|
|
overlap: 重叠字符数
|
|
|
|
Returns:
|
|
list of (chunk_text, start_line, end_line)
|
|
"""
|
|
lines = content.split('\n')
|
|
chunks = []
|
|
|
|
current_chunk = []
|
|
current_size = 0
|
|
start_line = 0
|
|
|
|
for i, line in enumerate(lines):
|
|
line_with_newline = line + '\n'
|
|
|
|
if current_size + len(line_with_newline) > chunk_size and current_chunk:
|
|
# 保存当前块
|
|
chunk_text = ''.join(current_chunk)
|
|
chunks.append((chunk_text, start_line + 1, i))
|
|
|
|
# 重叠部分
|
|
overlap_text = chunk_text[-overlap:] if len(chunk_text) > overlap else chunk_text
|
|
current_chunk = [overlap_text]
|
|
current_size = len(overlap_text)
|
|
start_line = i - len(overlap_text.split('\n')) + 1
|
|
|
|
current_chunk.append(line_with_newline)
|
|
current_size += len(line_with_newline)
|
|
|
|
# 最后一块
|
|
if current_chunk:
|
|
chunk_text = ''.join(current_chunk)
|
|
chunks.append((chunk_text, start_line + 1, len(lines)))
|
|
|
|
return chunks
|
|
|
|
|
|
def index_file(file_path, conn=None):
|
|
"""索引单个文件"""
|
|
close_conn = False
|
|
if conn is None:
|
|
conn = sqlite3.connect(DB_PATH)
|
|
close_conn = True
|
|
|
|
cursor = conn.cursor()
|
|
|
|
# 读取文件
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
except Exception as e:
|
|
print(f"✗ 读取失败 {file_path}: {e}")
|
|
return
|
|
|
|
# 获取文件信息
|
|
stat = os.stat(file_path)
|
|
modified_at = datetime.fromtimestamp(stat.st_mtime)
|
|
rel_path = str(Path(file_path).relative_to(MEMORY_DIR.parent))
|
|
|
|
# 插入或更新文件记录
|
|
cursor.execute("""
|
|
INSERT OR REPLACE INTO files (path, modified_at, indexed_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
""", (rel_path, modified_at))
|
|
|
|
file_id = cursor.lastrowid
|
|
|
|
# 删除旧索引
|
|
cursor.execute("DELETE FROM memory_index WHERE file_id = ?", (file_id,))
|
|
cursor.execute("DELETE FROM chunks WHERE file_id = ?", (file_id,))
|
|
|
|
# 分块并索引
|
|
chunks = chunk_content(content)
|
|
|
|
for chunk_idx, (chunk_text, start_line, end_line) in enumerate(chunks):
|
|
# 插入 chunks 表
|
|
cursor.execute("""
|
|
INSERT INTO chunks (file_id, chunk_index, content, start_line, end_line)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""", (file_id, chunk_idx, chunk_text, start_line, end_line))
|
|
|
|
# 插入 FTS5 索引
|
|
cursor.execute("""
|
|
INSERT INTO memory_index (content, file_id, chunk_index)
|
|
VALUES (?, ?, ?)
|
|
""", (chunk_text, file_id, chunk_idx))
|
|
|
|
conn.commit()
|
|
|
|
if close_conn:
|
|
conn.close()
|
|
|
|
print(f"✓ 已索引: {rel_path} ({len(chunks)} 块)")
|
|
|
|
|
|
def index_all():
|
|
"""索引所有 memory 文件"""
|
|
init_db()
|
|
|
|
conn = sqlite3.connect(DB_PATH)
|
|
|
|
# 查找所有 .md 文件
|
|
md_files = list(MEMORY_DIR.glob("*.md"))
|
|
|
|
print(f"\n开始索引 {len(md_files)} 个文件...")
|
|
|
|
for file_path in md_files:
|
|
if file_path.name.startswith('.'):
|
|
continue
|
|
index_file(file_path, conn)
|
|
|
|
conn.close()
|
|
print(f"\n✓ 索引完成")
|
|
|
|
|
|
def search(query, limit=10):
|
|
"""
|
|
全文搜索
|
|
|
|
Args:
|
|
query: 搜索关键词
|
|
limit: 返回结果数量
|
|
|
|
Returns:
|
|
list of dict with path, snippet, rank
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# 使用 FTS5 的 snippet 功能生成摘要
|
|
cursor.execute("""
|
|
SELECT
|
|
f.path,
|
|
snippet(memory_index, 0, '[', ']', '...', 32) as snippet,
|
|
c.start_line,
|
|
c.end_line,
|
|
rank
|
|
FROM memory_index
|
|
JOIN files f ON memory_index.file_id = f.id
|
|
JOIN chunks c ON memory_index.file_id = c.file_id
|
|
AND memory_index.chunk_index = c.chunk_index
|
|
WHERE memory_index MATCH ?
|
|
ORDER BY rank
|
|
LIMIT ?
|
|
""", (query, limit))
|
|
|
|
results = []
|
|
for row in cursor.fetchall():
|
|
results.append({
|
|
'path': row[0],
|
|
'snippet': row[1],
|
|
'start_line': row[2],
|
|
'end_line': row[3],
|
|
'rank': row[4]
|
|
})
|
|
|
|
conn.close()
|
|
return results
|
|
|
|
|
|
def print_search_results(query, results):
|
|
"""打印搜索结果"""
|
|
print(f"\n🔍 搜索: '{query}'")
|
|
print(f"找到 {len(results)} 条结果:\n")
|
|
|
|
for i, r in enumerate(results, 1):
|
|
print(f"{i}. {r['path']} (行 {r['start_line']}-{r['end_line']})")
|
|
print(f" {r['snippet']}")
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) < 2:
|
|
print("用法:")
|
|
print(f" {sys.argv[0]} index # 索引所有文件")
|
|
print(f" {sys.argv[0]} search <query> # 搜索")
|
|
sys.exit(1)
|
|
|
|
command = sys.argv[1]
|
|
|
|
if command == "index":
|
|
index_all()
|
|
elif command == "search":
|
|
if len(sys.argv) < 3:
|
|
print("错误: 需要提供搜索关键词")
|
|
sys.exit(1)
|
|
query = sys.argv[2]
|
|
results = search(query)
|
|
print_search_results(query, results)
|
|
else:
|
|
print(f"未知命令: {command}")
|
|
sys.exit(1)
|