返回 DeepSeek-Reasonix
search_index.go
根目录 / internal / session / search_index.go
1 package session
2
3 import (
4 "context"
5 "crypto/sha256"
6 "database/sql"
7 "errors"
8 "fmt"
9 "os"
10 "path/filepath"
11 "strings"
12 "unicode/utf8"
13
14 "reasonix/internal/projectiondb"
15 "reasonix/internal/provider"
16 "reasonix/internal/sessioncontent"
17 )
18
19 const searchIndexVersion = 4
20
21 var searchMigrations = []projectiondb.Migration{{Version: 1, Apply: func(ctx context.Context, tx *sql.Tx) error {
22 for _, statement := range []string{
23 `CREATE TABLE metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL)`,
24 `CREATE TABLE documents (rowid INTEGER PRIMARY KEY, message_id TEXT NOT NULL, version INTEGER NOT NULL, position INTEGER NOT NULL, event_sequence INTEGER NOT NULL, valid_to INTEGER NOT NULL DEFAULT 0, role TEXT NOT NULL, preview TEXT NOT NULL, text TEXT NOT NULL, current INTEGER NOT NULL, UNIQUE(message_id,version))`,
25 `CREATE INDEX documents_snapshot_position ON documents(position DESC,event_sequence,valid_to)`,
26 `CREATE INDEX documents_current_id ON documents(message_id) WHERE current=1`,
27 `CREATE VIRTUAL TABLE documents_fts USING fts5(text, content='documents', content_rowid='rowid', tokenize='trigram')`,
28 `CREATE TRIGGER documents_ai AFTER INSERT ON documents BEGIN INSERT INTO documents_fts(rowid,text) VALUES (new.rowid,new.text); END`,
29 `CREATE TRIGGER documents_ad AFTER DELETE ON documents BEGIN INSERT INTO documents_fts(documents_fts,rowid,text) VALUES('delete',old.rowid,old.text); END`,
30 `CREATE TRIGGER documents_au AFTER UPDATE OF text ON documents BEGIN INSERT INTO documents_fts(documents_fts,rowid,text) VALUES('delete',old.rowid,old.text); INSERT INTO documents_fts(rowid,text) VALUES(new.rowid,new.text); END`,
31 } {
32 if _, err := tx.ExecContext(ctx, statement); err != nil {
33 return err
34 }
35 }
36 return nil
37 }}}
38
39 type searchBuildState struct {
40 positions map[string]int64
41 versions map[string]int
42 nextPosition int64
43 }
44
45 type searchPreparation struct {
46 done chan struct{}
47 err error
48 }
49
50 func searchIndexPath(root, sessionID string) string {
51 return filepath.Join(root, ".query-cache", filepath.Base(sessionID), "search-v1.sqlite")
52 }
53
54 func (q *Query) SearchHistory(ctx context.Context, ref SessionRef, textQuery, cursor string, limit int) (SearchHistoryPage, error) {
55 if q == nil {
56 return SearchHistoryPage{}, errors.New("session: nil query")
57 }
58 if err := ref.validate(q.hostID); err != nil {
59 return SearchHistoryPage{}, err
60 }
61 textQuery = strings.TrimSpace(textQuery)
62 if textQuery == "" {
63 return SearchHistoryPage{}, errors.New("session: history search query is required")
64 }
65 if limit <= 0 {
66 limit = 50
67 }
68 limit = min(limit, 200)
69 filesystem, ok := q.persistence.(*FilesystemPersistence)
70 if !ok {
71 return SearchHistoryPage{}, errors.New("session: history search requires filesystem persistence")
72 }
73 path := searchIndexPath(filesystem.Root, ref.SessionID)
74 if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
75 preparation := q.prepareSearchIndex(filesystem, ref.SessionID, path)
76 select {
77 case <-preparation.done:
78 if preparation.err != nil {
79 return SearchHistoryPage{Status: "failed"}, preparation.err
80 }
81 default:
82 return SearchHistoryPage{Status: "preparing"}, nil
83 }
84 }
85 lock := q.projectionLock("search", ref.SessionID)
86 lock.Lock()
87 err := ensureSearchIndex(ctx, filesystem, ref.SessionID, path)
88 lock.Unlock()
89 if err != nil {
90 return SearchHistoryPage{}, err
91 }
92 handle, err := projectiondb.Open(ctx, projectiondb.OpenOptions{Path: path, Migrations: searchMigrations, RequireDisk: true, MaxOpenConns: 1})
93 if err != nil {
94 return SearchHistoryPage{}, err
95 }
96 defer handle.DB.Close()
97 metadata, err := readSearchMetadata(ctx, handle.DB)
98 if err != nil {
99 return SearchHistoryPage{}, err
100 }
101 snapshot := metadata.durableSequence
102 before := int64(^uint64(0) >> 1)
103 digest := fmt.Sprintf("%x", sha256.Sum256([]byte(textQuery)))
104 if cursor != "" {
105 parsed, err := decodeSearchHistoryCursor(cursor)
106 if err != nil {
107 return SearchHistoryPage{}, err
108 }
109 if parsed.SessionID != ref.SessionID || parsed.StorageRevision != StorageRevision || parsed.Projection != searchIndexVersion || parsed.QueryDigest != digest || parsed.SnapshotSequence > snapshot || parsed.BeforePosition <= 0 || parsed.Generation != metadata.generation {
110 return SearchHistoryPage{Status: "stale_cursor", CoverageSequence: metadata.durableSequence}, nil
111 }
112 snapshot, before = parsed.SnapshotSequence, parsed.BeforePosition
113 }
114 var rows *sql.Rows
115 if utf8.RuneCountInString(textQuery) >= 3 {
116 match := `"` + strings.ReplaceAll(textQuery, `"`, `""`) + `"`
117 rows, err = handle.DB.QueryContext(ctx, `SELECT d.message_id,d.position,d.role,d.preview,d.event_sequence FROM documents_fts JOIN documents d ON d.rowid=documents_fts.rowid WHERE documents_fts MATCH ? AND instr(d.text,?)>0 AND d.position<? AND d.event_sequence<=? AND (d.valid_to=0 OR d.valid_to>?) ORDER BY d.position DESC LIMIT ?`, match, textQuery, before, snapshot, snapshot, limit+1)
118 } else {
119 rows, err = handle.DB.QueryContext(ctx, `SELECT message_id,position,role,preview,event_sequence FROM documents WHERE instr(text,?)>0 AND position<? AND event_sequence<=? AND (valid_to=0 OR valid_to>?) ORDER BY position DESC LIMIT ?`, textQuery, before, snapshot, snapshot, limit+1)
120 }
121 if err != nil {
122 return SearchHistoryPage{}, err
123 }
124 defer rows.Close()
125 page := SearchHistoryPage{Hits: []SearchHistoryHit{}, SnapshotSequence: snapshot, CoverageSequence: metadata.durableSequence, Status: "ready"}
126 for rows.Next() {
127 var hit SearchHistoryHit
128 if err := rows.Scan(&hit.MessageID, &hit.Position, &hit.Role, &hit.Preview, &hit.EventSequence); err != nil {
129 return SearchHistoryPage{}, err
130 }
131 if len(page.Hits) == limit {
132 page.HasMore = true
133 break
134 }
135 page.Hits = append(page.Hits, hit)
136 }
137 if err := rows.Err(); err != nil {
138 return SearchHistoryPage{}, err
139 }
140 if page.HasMore && len(page.Hits) > 0 {
141 page.NextCursor, err = encodeSearchHistoryCursor(searchHistoryCursor{SessionID: ref.SessionID, StorageRevision: StorageRevision, SnapshotSequence: snapshot, BeforePosition: page.Hits[len(page.Hits)-1].Position, Projection: searchIndexVersion, QueryDigest: digest, Generation: metadata.generation})
142 if err != nil {
143 return SearchHistoryPage{}, err
144 }
145 }
146 return page, nil
147 }
148
149 func (q *Query) prepareSearchIndex(filesystem *FilesystemPersistence, sessionID, path string) *searchPreparation {
150 q.searchMu.Lock()
151 if current := q.searchBuilds[sessionID]; current != nil {
152 q.searchMu.Unlock()
153 return current
154 }
155 preparation := &searchPreparation{done: make(chan struct{})}
156 q.searchBuilds[sessionID] = preparation
157 q.searchMu.Unlock()
158 go func() {
159 if err := q.slots.acquire(q.rebuildCtx, rebuildPrioritySearch); err != nil {
160 preparation.err = err
161 close(preparation.done)
162 return
163 }
164 defer q.slots.release()
165 lock := q.projectionLock("search", sessionID)
166 lock.Lock()
167 preparation.err = ensureSearchIndex(q.rebuildCtx, filesystem, sessionID, path)
168 lock.Unlock()
169 close(preparation.done)
170 }()
171 return preparation
172 }
173
174 type searchMetadata struct {
175 sessionID string
176 logSize int64
177 storageRevision int
178 projection int
179 durableSequence uint64
180 generation string
181 }
182
183 func readSearchMetadata(ctx context.Context, db *sql.DB) (searchMetadata, error) {
184 values := map[string]string{}
185 rows, err := db.QueryContext(ctx, `SELECT key,value FROM metadata`)
186 if err != nil {
187 return searchMetadata{}, err
188 }
189 defer rows.Close()
190 for rows.Next() {
191 var key, value string
192 if err := rows.Scan(&key, &value); err != nil {
193 return searchMetadata{}, err
194 }
195 values[key] = value
196 }
197 metadata := searchMetadata{sessionID: values["session_id"], generation: values["generation"]}
198 if _, err := fmt.Sscan(values["log_size"], &metadata.logSize); err != nil {
199 return searchMetadata{}, err
200 }
201 if _, err := fmt.Sscan(values["storage_revision"], &metadata.storageRevision); err != nil {
202 return searchMetadata{}, err
203 }
204 if _, err := fmt.Sscan(values["projection_version"], &metadata.projection); err != nil {
205 return searchMetadata{}, err
206 }
207 if _, err := fmt.Sscan(values["durable_sequence"], &metadata.durableSequence); err != nil {
208 return searchMetadata{}, err
209 }
210 return metadata, rows.Err()
211 }
212
213 func ensureSearchIndex(ctx context.Context, persistence *FilesystemPersistence, sessionID, path string) error {
214 dir := filepath.Join(persistence.Root, sessionID)
215 revision, err := revisionOfLog(dir)
216 if err != nil {
217 return err
218 }
219 if handle, err := projectiondb.Open(ctx, projectiondb.OpenOptions{Path: path, Migrations: searchMigrations, RequireDisk: true, MaxOpenConns: 1}); err == nil {
220 metadata, metaErr := readSearchMetadata(ctx, handle.DB)
221 _ = handle.DB.Close()
222 if metaErr == nil && metadata.sessionID == sessionID && metadata.storageRevision == StorageRevision && metadata.projection == searchIndexVersion {
223 if metadata.logSize == revision.Size {
224 return nil
225 }
226 if metadata.logSize >= 0 && metadata.logSize < revision.Size {
227 if err := incrementSearchIndex(ctx, dir, path, revision, metadata); err == nil {
228 return nil
229 }
230 }
231 }
232 }
233 return rebuildSearchIndex(ctx, dir, path, sessionID, revision)
234 }
235
236 func rebuildSearchIndex(ctx context.Context, dir, path, sessionID string, revision logRevision) error {
237 return projectiondb.Rebuild(ctx, projectiondb.OpenOptions{Path: path, Migrations: searchMigrations, RequireDisk: true, MaxOpenConns: 1, QuickCheck: true}, func(ctx context.Context, db *sql.DB) error {
238 if err := configureHistoryRebuild(ctx, db); err != nil {
239 return err
240 }
241 metadata := searchMetadata{sessionID: sessionID, storageRevision: StorageRevision, projection: searchIndexVersion, generation: randomID()}
242 return populateSearchIndex(ctx, dir, db, 0, 1, revision, metadata, searchBuildState{positions: map[string]int64{}, versions: map[string]int{}})
243 })
244 }
245
246 func incrementSearchIndex(ctx context.Context, dir, path string, revision logRevision, metadata searchMetadata) error {
247 handle, err := projectiondb.Open(ctx, projectiondb.OpenOptions{Path: path, Migrations: searchMigrations, RequireDisk: true, MaxOpenConns: 1})
248 if err != nil {
249 return err
250 }
251 defer handle.DB.Close()
252 state := searchBuildState{positions: map[string]int64{}, versions: map[string]int{}}
253 rows, err := handle.DB.QueryContext(ctx, `SELECT message_id,position,version FROM documents WHERE current=1`)
254 if err != nil {
255 return err
256 }
257 for rows.Next() {
258 var id string
259 var position int64
260 var version int
261 if err := rows.Scan(&id, &position, &version); err != nil {
262 _ = rows.Close()
263 return err
264 }
265 state.positions[id], state.versions[id] = position, version
266 state.nextPosition = max(state.nextPosition, position)
267 }
268 if err := errors.Join(rows.Err(), rows.Close()); err != nil {
269 return err
270 }
271 versionRows, err := handle.DB.QueryContext(ctx, `SELECT message_id,MAX(version) FROM documents GROUP BY message_id`)
272 if err != nil {
273 return err
274 }
275 for versionRows.Next() {
276 var id string
277 var version int
278 if err := versionRows.Scan(&id, &version); err != nil {
279 _ = versionRows.Close()
280 return err
281 }
282 state.versions[id] = version
283 }
284 if err := errors.Join(versionRows.Err(), versionRows.Close()); err != nil {
285 return err
286 }
287 return populateSearchIndex(ctx, dir, handle.DB, metadata.logSize, metadata.durableSequence+1, revision, metadata, state)
288 }
289
290 func populateSearchIndex(ctx context.Context, dir string, db *sql.DB, startOffset int64, nextSequence uint64, revision logRevision, metadata searchMetadata, state searchBuildState) error {
291 manifest, err := readManifest(filepath.Join(dir, "manifest.json"))
292 if err != nil {
293 return err
294 }
295 log, err := os.Open(logPathForManifest(dir, manifest))
296 if err != nil {
297 return err
298 }
299 defer log.Close()
300 tx, err := db.BeginTx(ctx, nil)
301 if err != nil {
302 return err
303 }
304 defer func() { _ = tx.Rollback() }()
305 content := contentStoreForSessionDir(dir)
306 var buildErr error
307 progress, err := scanHistoryLog(ctx, log, startOffset, nextSequence, revision.Size, content, func(commit Commit) bool {
308 for _, event := range commit.Events {
309 if err := indexSearchEvent(ctx, tx, content, &state, event); err != nil {
310 buildErr = err
311 return false
312 }
313 }
314 return true
315 })
316 if err != nil || buildErr != nil {
317 return errors.Join(err, buildErr)
318 }
319 values := map[string]string{
320 "session_id": metadata.sessionID, "log_size": fmt.Sprint(progress.end),
321 "log_mtime_ns": fmt.Sprint(progress.modTimeNS), "storage_revision": fmt.Sprint(StorageRevision),
322 "projection_version": fmt.Sprint(searchIndexVersion), "durable_sequence": fmt.Sprint(progress.sequence),
323 "generation": metadata.generation,
324 }
325 for key, value := range values {
326 if _, err := tx.ExecContext(ctx, `INSERT INTO metadata(key,value) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, key, value); err != nil {
327 return err
328 }
329 }
330 return tx.Commit()
331 }
332
333 func indexSearchEvent(ctx context.Context, tx *sql.Tx, content *sessioncontent.Store, state *searchBuildState, event Event) error {
334 if event.Kind != "message/complete" && event.Kind != "message/upsert" && event.Kind != "message/retract" && event.Kind != "history/replace" && event.Kind != "legacy/import" {
335 return nil
336 }
337 payload := event.Payload
338 if event.PayloadRef != nil {
339 var err error
340 payload, err = resolveContentPayload(ctx, content, *event.PayloadRef)
341 if err != nil {
342 return err
343 }
344 }
345 switch event.Kind {
346 case "message/retract":
347 ids, err := retractedMessageIDs(event, payload)
348 if err != nil {
349 return err
350 }
351 for _, id := range ids {
352 if _, err := tx.ExecContext(ctx, `UPDATE documents SET current=0,valid_to=? WHERE message_id=? AND current=1`, event.Sequence, id); err != nil {
353 return err
354 }
355 delete(state.positions, id)
356 }
357 return nil
358 case "message/complete", "message/upsert":
359 var body struct {
360 Message *provider.Message `json:"message"`
361 }
362 if err := strictPayload(payload, &body); err != nil || body.Message == nil {
363 return damagedPayload(event, err)
364 }
365 return indexSearchMessage(ctx, tx, state, *body.Message, event.Sequence, event.Kind == "message/upsert")
366 case "history/replace", "legacy/import":
367 messages, err := replacementEventMessages(event, payload)
368 if err != nil {
369 return err
370 }
371 if _, err := tx.ExecContext(ctx, `UPDATE documents SET current=0,valid_to=? WHERE current=1`, event.Sequence); err != nil {
372 return err
373 }
374 state.positions = map[string]int64{}
375 state.nextPosition = 0
376 for _, message := range messages {
377 if err := indexSearchMessage(ctx, tx, state, message, event.Sequence, false); err != nil {
378 return err
379 }
380 }
381 }
382 return nil
383 }
384
385 func indexSearchMessage(ctx context.Context, tx *sql.Tx, state *searchBuildState, message provider.Message, sequence uint64, upsert bool) error {
386 id := strings.TrimSpace(message.ID)
387 if id == "" {
388 return errors.New("session: search document has no stable message id")
389 }
390 position, exists := state.positions[id]
391 if !exists {
392 state.nextPosition++
393 position = state.nextPosition
394 state.positions[id] = position
395 } else if !upsert {
396 return fmt.Errorf("session: duplicate search message id %q", id)
397 }
398 if exists {
399 if _, err := tx.ExecContext(ctx, `UPDATE documents SET current=0,valid_to=? WHERE message_id=? AND current=1`, sequence, id); err != nil {
400 return err
401 }
402 }
403 version := state.versions[id] + 1
404 state.versions[id] = version
405 _, err := tx.ExecContext(ctx, `INSERT INTO documents(message_id,version,position,event_sequence,valid_to,role,preview,text,current) VALUES(?,?,?,?,0,?,?,?,1)`, id, version, position, sequence, string(message.Role), messagePreview(message), messageSearchText(message))
406 return err
407 }
408
408 lines GO