Coverage for merco/memory/session_search.py: 82%

1"""Session FTS5 全文搜索 — 借鉴 hermes 双 tokenizer 设计 

2 

3当前: unicode61 tokenizer (英文/Latin 友好) 

4后期: 加 trigram tokenizer (CJK/子串友好) 

5""" 

6 

7import logging 

8import re 

9 

10logger = logging.getLogger("merco.session_search") 

11 

12 

13class SessionSearch: 

14 """会话全文搜索 — FTS5 索引 + 片段高亮""" 

15 

16 def __init__(self, store): 

17 """store: SessionStore 实例，复用其 DB 路径""" 

18 self._store = store 

19 self._ensure_index() 

20 

21 def _conn(self): 

22 import sqlite3 

23 conn = sqlite3.connect(self._store.db_path) 

24 conn.row_factory = sqlite3.Row 

25 return conn 

26 

27 def _ensure_index(self): 

28 with self._conn() as conn: 

29 conn.executescript(""" 

30 CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts 

31 USING fts5(content, tool_name, session_id, 

32 content_rowid='id', 

33 tokenize='unicode61'); 

34 

35 CREATE TRIGGER IF NOT EXISTS msg_fts_insert 

36 AFTER INSERT ON messages BEGIN 

37 INSERT INTO messages_fts(rowid, content, tool_name, session_id) 

38 VALUES (new.id, new.content, '', new.session_id); 

39 END; 

40 

41 CREATE TRIGGER IF NOT EXISTS msg_fts_delete 

42 AFTER DELETE ON messages BEGIN 

43 INSERT INTO messages_fts(messages_fts, rank) 

44 VALUES ('delete', old.id, 1); 

45 END; 

46 """) 

47 

48 def rebuild(self): 

49 """全量重建 FTS 索引""" 

50 with self._conn() as conn: 

51 conn.execute("DELETE FROM messages_fts") 

52 conn.execute( 

53 "INSERT INTO messages_fts(rowid, content, session_id) " 

54 "SELECT id, content, session_id FROM messages" 

55 ) 

56 logger.info("FTS index rebuilt") 

57 

58 def search(self, query: str, session_id: str | None = None, 

59 limit: int = 20) -> list[dict]: 

60 """搜索消息，返回排序结果 + 片段""" 

61 import sqlite3 

62 

63 sanitized = self._sanitize(query) 

64 sql = """ 

65 SELECT 

66 m.id, m.session_id, m.role, m.timestamp, 

67 snippet(messages_fts, 1, '>>>', '<<<', '...', 40) as snippet, 

68 s.title as session_title 

69 FROM messages_fts f 

70 JOIN messages m ON m.id = f.rowid 

71 JOIN sessions s ON s.id = m.session_id 

72 WHERE messages_fts MATCH ? 

73 """ 

74 params = [sanitized] 

75 

76 if session_id: 

77 sql += " AND m.session_id = ?" 

78 params.append(session_id) 

79 

80 sql += " ORDER BY rank LIMIT ?" 

81 params.append(limit) 

82 

83 with self._conn() as conn: 

84 try: 

85 rows = conn.execute(sql, params).fetchall() 

86 except sqlite3.OperationalError: 

87 # FTS5 syntax error despite sanitization — return empty 

88 logger.debug("FTS5 query error for: %s", sanitized) 

89 return [] 

90 

91 return [ 

92 { 

93 "id": r["id"], 

94 "session_id": r["session_id"], 

95 "session_title": r["session_title"], 

96 "role": r["role"], 

97 "snippet": r["snippet"], 

98 "timestamp": r["timestamp"], 

99 } 

100 for r in rows 

101 ] 

102 

103 @staticmethod 

104 def _sanitize(query: str) -> str: 

105 """Port of Hermes's _sanitize_fts5_query — preserve meaning, escape danger. 

106 

107 Strategy (6 steps): 

108 1. Protect balanced double-quoted phrases with placeholders 

109 2. Strip remaining FTS5-special characters (+, {}, (), ", ^) 

110 3. Normalise wildcard * (collapse repeats, remove leading *) 

111 4. Remove dangling boolean operators (AND/OR/NOT at start/end) 

112 5. Wrap hyphenated/dotted/underscored terms as FTS5 phrase literals 

113 6. Restore preserved quoted phrases 

114 """ 

115 if not query or not query.strip(): 

116 return "*" 

117 

118 # Step 1: Extract balanced double-quoted phrases, protect them 

119 _quoted_parts: list[str] = [] 

120 

121 def _preserve_quoted(m: re.Match) -> str: 

122 _quoted_parts.append(m.group(0)) 

123 return f"\x00Q{len(_quoted_parts) - 1}\x00" 

124 

125 sanitized = re.sub(r'"[^"]*"', _preserve_quoted, query) 

126 

127 # Step 2: Strip remaining FTS5-special characters 

128 sanitized = re.sub(r'[+{}()\"^]', " ", sanitized) 

129 

130 # Step 3: Collapse repeated *, remove leading * (prefix needs char before *) 

131 sanitized = re.sub(r"\*+", "*", sanitized) 

132 sanitized = re.sub(r"(^|\s)\*", r"\1", sanitized) 

133 

134 # Step 4: Remove dangling boolean operators at start/end 

135 sanitized = re.sub(r"(?i)^(AND|OR|NOT)\b\s*", "", sanitized.strip()) 

136 sanitized = re.sub(r"(?i)\s+(AND|OR|NOT)\s*$", "", sanitized.strip()) 

137 

138 # Step 5: Wrap hyphenated/dotted/underscored terms in quotes 

139 # Single pass avoids double-quoting e.g. my-app.config 

140 sanitized = re.sub(r"\b(\w+(?:[._-]\w+)+)\b", r'"\1"', sanitized) 

141 

142 # Step 6: Restore preserved quoted phrases 

143 for i, quoted in enumerate(_quoted_parts): 

144 sanitized = sanitized.replace(f"\x00Q{i}\x00", quoted) 

145 

146 return sanitized.strip() or "*"