feat: added fts5 backend

This commit is contained in:
didier committed 2026-06-05 15:51:41 -04:00
1 parent d5cae6efa9
commit 98aa7dd68c
6 files changed
+118 -14

No files matched your search

+1
View File
@@ -10,3 +10,4 @@ result
# IDE # IDE
.idea/ .idea/
.vscode/ .vscode/
*.db
+1
View File
@@ -14,6 +14,7 @@ require (
github.com/dustin/go-humanize v1.0.1 // indirect github.com/dustin/go-humanize v1.0.1 // indirect
github.com/google/jsonschema-go v0.4.2 // indirect github.com/google/jsonschema-go v0.4.2 // indirect
github.com/google/uuid v1.6.0 // indirect github.com/google/uuid v1.6.0 // indirect
github.com/joho/godotenv v1.5.1 // indirect
github.com/mattn/go-isatty v0.0.21 // indirect github.com/mattn/go-isatty v0.0.21 // indirect
github.com/mfridman/interpolate v0.0.2 // indirect github.com/mfridman/interpolate v0.0.2 // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect
+2
View File
@@ -24,6 +24,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
+47 -14
View File
@@ -12,6 +12,7 @@ import (
"syscall" "syscall"
"time" "time"
"github.com/joho/godotenv"
"github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server" "github.com/mark3labs/mcp-go/server"
"github.com/qdrant/go-client/qdrant" "github.com/qdrant/go-client/qdrant"
@@ -21,13 +22,15 @@ import (
var version = "0.1.0" var version = "0.1.0"
func main() { func main() {
godotenv.Load()
logger, err := zap.NewProduction() logger, err := zap.NewProduction()
if err != nil { if err != nil {
panic(fmt.Sprintf("failed to init logger: %v", err)) panic(fmt.Sprintf("failed to init logger: %v", err))
} }
defer logger.Sync() defer logger.Sync()
qdrantURL := requireEnv(logger, "QDRANT_URL") qdrantURL := os.Getenv("QDRANT_URL")
openAIKey := requireEnv(logger, "OPENAI_API_KEY") openAIKey := requireEnv(logger, "OPENAI_API_KEY")
token := requireEnv(logger, "MCP_TOKEN") token := requireEnv(logger, "MCP_TOKEN")
embeddingModel := os.Getenv("EMBEDDING_MODEL") embeddingModel := os.Getenv("EMBEDDING_MODEL")
@@ -41,6 +44,8 @@ func main() {
qdrantAPIKey := os.Getenv("QDRANT_API_KEY") qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
qdrantTLS := os.Getenv("QDRANT_TLS") == "true" qdrantTLS := os.Getenv("QDRANT_TLS") == "true"
var qdrantStr *qdrantStore
if qdrantURL != "" {
qdrantClient, err := qdrant.NewClient(&qdrant.Config{ qdrantClient, err := qdrant.NewClient(&qdrant.Config{
Host: qdrantURL, Host: qdrantURL,
Port: qdrantPort, Port: qdrantPort,
@@ -53,7 +58,14 @@ func main() {
defer qdrantClient.Close() defer qdrantClient.Close()
emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL) emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
qdrantStore := newQdrantStore(qdrantClient, emb) qdrantStr = newQdrantStore(qdrantClient, emb)
if err := qdrantStr.ensureCollection(context.Background()); err != nil {
logger.Fatal("failed to ensure collection", zap.Error(err))
}
} else {
logger.Warn("QDRANT_URL not set, using SQLite FTS5 for search")
}
sqlitePath := envOrDefault("SQLITE_PATH", "flashcard.db") sqlitePath := envOrDefault("SQLITE_PATH", "flashcard.db")
sqliteStore, err := newSqliteStore(sqlitePath) sqliteStore, err := newSqliteStore(sqlitePath)
@@ -62,10 +74,6 @@ func main() {
} }
defer sqliteStore.close() defer sqliteStore.close()
if err := qdrantStore.ensureCollection(context.Background()); err != nil {
logger.Fatal("failed to ensure collection", zap.Error(err))
}
mcpServer := server.NewMCPServer( mcpServer := server.NewMCPServer(
"flashcard-mcp", "flashcard-mcp",
version, version,
@@ -73,10 +81,10 @@ func main() {
server.WithLogging(), server.WithLogging(),
) )
registerStoreFlashcard(mcpServer, sqliteStore, qdrantStore, logger) registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
registerSearchFlashcard(mcpServer, sqliteStore, qdrantStore, logger) registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr)
registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStore, logger) registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStore, logger) registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
streamableServer := server.NewStreamableHTTPServer(mcpServer) streamableServer := server.NewStreamableHTTPServer(mcpServer)
@@ -111,7 +119,6 @@ func main() {
if err := streamableServer.Shutdown(shutdownCtx); err != nil { if err := streamableServer.Shutdown(shutdownCtx); err != nil {
logger.Error("shutdown error", zap.Error(err)) logger.Error("shutdown error", zap.Error(err))
} }
qdrantClient.Close()
sqliteStore.close() sqliteStore.close()
os.Exit(0) os.Exit(0)
}() }()
@@ -152,7 +159,7 @@ func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qd
) )
} }
func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore) {
mcpServer.AddTool( mcpServer.AddTool(
mcp.NewTool("search_flashcard", mcp.NewTool("search_flashcard",
mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant entries for the given query."), mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant entries for the given query."),
@@ -251,9 +258,17 @@ func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap
return mcp.NewToolResultError(fmt.Sprintf("failed to persist update: %v", err)), nil return mcp.NewToolResultError(fmt.Sprintf("failed to persist update: %v", err)), nil
} }
if err := sqlite.deleteFTS(ctx, latest.ID); err != nil {
logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err))
}
if err := sqlite.indexFTS(ctx, id, note); err != nil {
return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil
}
if qdrant != nil {
if err := qdrant.index(ctx, id, publicID, version, note); err != nil { if err := qdrant.index(ctx, id, publicID, version, note); err != nil {
logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
} }
}
entry, err := sqlite.getByID(ctx, id) entry, err := sqlite.getByID(ctx, id)
if err != nil { if err != nil {
@@ -309,9 +324,14 @@ func handleDeleteFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap
return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", err)), nil return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", err)), nil
} }
if err := sqlite.deleteFTS(ctx, latest.ID); err != nil {
logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err))
}
if qdrant != nil {
if err := qdrant.delete(ctx, publicID); err != nil { if err := qdrant.delete(ctx, publicID); err != nil {
logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err)) logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err))
} }
}
return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil
} }
@@ -344,9 +364,14 @@ func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.
if publicID == 0 { if publicID == 0 {
publicID = id publicID = id
} }
if err := sqlite.indexFTS(ctx, id, note); err != nil {
return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil
}
if qdrant != nil {
if err := qdrant.index(ctx, id, publicID, 1, note); err != nil { if err := qdrant.index(ctx, id, publicID, 1, note); err != nil {
logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
} }
}
entry, err := sqlite.getByID(ctx, id) entry, err := sqlite.getByID(ctx, id)
if err != nil { if err != nil {
@@ -386,16 +411,24 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool
} }
} }
results, err := qdrant.search(ctx, query, limit) var results []searchResult
var err error
if qdrant != nil {
results, err = qdrant.search(ctx, query, limit)
if err != nil { if err != nil {
return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
} }
for i, r := range results { for i, r := range results {
if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil { if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
results[i].Entry = *cached results[i].Entry = *cached
} }
} }
} else {
results, err = sqlite.searchFTS(ctx, query, limit)
if err != nil {
return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
}
}
output, err := json.MarshalIndent(results, "", " ") output, err := json.MarshalIndent(results, "", " ")
if err != nil { if err != nil {
+5
View File
@@ -0,0 +1,5 @@
-- +goose Up
CREATE VIRTUAL TABLE IF NOT EXISTS flashcard_fts USING fts5(note);
-- +goose Down
DROP TABLE IF EXISTS flashcard_fts;
+62
View File
@@ -184,6 +184,68 @@ func (s *sqliteStore) tombstone(ctx context.Context, publicID int64, version int
return id, nil return id, nil
} }
func (s *sqliteStore) indexFTS(ctx context.Context, id int64, note string) error {
_, err := s.db.ExecContext(ctx,
`INSERT INTO flashcard_fts(rowid, note) VALUES (?, ?)`,
id, note,
)
if err != nil {
return fmt.Errorf("index fts: %w", err)
}
return nil
}
func (s *sqliteStore) deleteFTS(ctx context.Context, id int64) error {
_, err := s.db.ExecContext(ctx,
`DELETE FROM flashcard_fts WHERE rowid = ?`,
id,
)
if err != nil {
return fmt.Errorf("delete fts: %w", err)
}
return nil
}
func (s *sqliteStore) searchFTS(ctx context.Context, query string, limit uint64) ([]searchResult, error) {
if limit == 0 {
limit = 5
}
rows, err := s.db.QueryContext(ctx,
`SELECT rowid, rank FROM flashcard_fts WHERE note MATCH ? ORDER BY rank LIMIT ?`,
query, limit,
)
if err != nil {
return nil, fmt.Errorf("fts query: %w", err)
}
defer rows.Close()
var results []searchResult
for rows.Next() {
var id int64
var rank float64
if err := rows.Scan(&id, &rank); err != nil {
return nil, fmt.Errorf("scan: %w", err)
}
entry, err := s.getByID(ctx, id)
if err != nil {
return nil, fmt.Errorf("get entry: %w", err)
}
if entry == nil {
continue
}
results = append(results, searchResult{
Entry: *entry,
Score: float32(-rank),
})
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("rows: %w", err)
}
return results, nil
}
func (s *sqliteStore) close() error { func (s *sqliteStore) close() error {
return s.db.Close() return s.db.Close()
} }