diff --git a/.gitignore b/.gitignore index c839cd7..83b316e 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ result # IDE .idea/ .vscode/ +*.db diff --git a/go.mod b/go.mod index f22ef0d..386df46 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,7 @@ require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/jsonschema-go v0.4.2 // 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/mfridman/interpolate v0.0.2 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect diff --git a/go.sum b/go.sum index de8e8b8..9f709c1 100644 --- a/go.sum +++ b/go.sum @@ -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/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/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/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= diff --git a/main.go b/main.go index 0d8bbfd..2031694 100644 --- a/main.go +++ b/main.go @@ -12,6 +12,7 @@ import ( "syscall" "time" + "github.com/joho/godotenv" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "github.com/qdrant/go-client/qdrant" @@ -21,13 +22,15 @@ import ( var version = "0.1.0" func main() { + godotenv.Load() + logger, err := zap.NewProduction() if err != nil { panic(fmt.Sprintf("failed to init logger: %v", err)) } defer logger.Sync() - qdrantURL := requireEnv(logger, "QDRANT_URL") + qdrantURL := os.Getenv("QDRANT_URL") openAIKey := requireEnv(logger, "OPENAI_API_KEY") token := requireEnv(logger, "MCP_TOKEN") embeddingModel := os.Getenv("EMBEDDING_MODEL") @@ -41,19 +44,28 @@ func main() { qdrantAPIKey := os.Getenv("QDRANT_API_KEY") qdrantTLS := os.Getenv("QDRANT_TLS") == "true" - qdrantClient, err := qdrant.NewClient(&qdrant.Config{ - Host: qdrantURL, - Port: qdrantPort, - APIKey: qdrantAPIKey, - UseTLS: qdrantTLS, - }) - if err != nil { - logger.Fatal("failed to create qdrant client", zap.Error(err)) - } - defer qdrantClient.Close() + var qdrantStr *qdrantStore + if qdrantURL != "" { + qdrantClient, err := qdrant.NewClient(&qdrant.Config{ + Host: qdrantURL, + Port: qdrantPort, + APIKey: qdrantAPIKey, + UseTLS: qdrantTLS, + }) + if err != nil { + logger.Fatal("failed to create qdrant client", zap.Error(err)) + } + defer qdrantClient.Close() - emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL) - qdrantStore := newQdrantStore(qdrantClient, emb) + emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL) + 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") sqliteStore, err := newSqliteStore(sqlitePath) @@ -62,10 +74,6 @@ func main() { } defer sqliteStore.close() - if err := qdrantStore.ensureCollection(context.Background()); err != nil { - logger.Fatal("failed to ensure collection", zap.Error(err)) - } - mcpServer := server.NewMCPServer( "flashcard-mcp", version, @@ -73,10 +81,10 @@ func main() { server.WithLogging(), ) - registerStoreFlashcard(mcpServer, sqliteStore, qdrantStore, logger) - registerSearchFlashcard(mcpServer, sqliteStore, qdrantStore, logger) - registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStore, logger) - registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStore, logger) + registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger) + registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr) + registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger) + registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger) streamableServer := server.NewStreamableHTTPServer(mcpServer) @@ -111,7 +119,6 @@ func main() { if err := streamableServer.Shutdown(shutdownCtx); err != nil { logger.Error("shutdown error", zap.Error(err)) } - qdrantClient.Close() sqliteStore.close() 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( mcp.NewTool("search_flashcard", mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant entries for the given query."), @@ -251,8 +258,16 @@ func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap return mcp.NewToolResultError(fmt.Sprintf("failed to persist update: %v", 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)) + 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 { + logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) + } } entry, err := sqlite.getByID(ctx, id) @@ -309,8 +324,13 @@ func handleDeleteFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", 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)) + 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 { + 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 @@ -344,8 +364,13 @@ func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap. if publicID == 0 { publicID = id } - 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)) + 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 { + logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) + } } entry, err := sqlite.getByID(ctx, id) @@ -386,14 +411,22 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool } } - results, err := qdrant.search(ctx, query, limit) - if err != nil { - return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil - } - - for i, r := range results { - if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil { - results[i].Entry = *cached + var results []searchResult + var err error + if qdrant != nil { + results, err = qdrant.search(ctx, query, limit) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil + } + for i, r := range results { + if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil { + 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 } } diff --git a/migrations/000002_add_fts5.sql b/migrations/000002_add_fts5.sql new file mode 100644 index 0000000..b878b15 --- /dev/null +++ b/migrations/000002_add_fts5.sql @@ -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; diff --git a/sqlite.go b/sqlite.go index 0b061b7..cce8ee9 100644 --- a/sqlite.go +++ b/sqlite.go @@ -184,6 +184,68 @@ func (s *sqliteStore) tombstone(ctx context.Context, publicID int64, version int 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 { return s.db.Close() }