From ffdc0230654656dad8f04d04f64e925428cb8755 Mon Sep 17 00:00:00 2001 From: didier amyot Date: Fri, 5 Jun 2026 16:02:31 -0400 Subject: [PATCH] refactor: SearchEngine interface, remove nilable qdrantStore --- fts5.go | 23 ++++++++ main.go | 159 +++++++++++++++++++++++------------------------------- qdrant.go | 41 +++++--------- search.go | 9 ++++ types.go | 20 +++++++ 5 files changed, 134 insertions(+), 118 deletions(-) create mode 100644 fts5.go create mode 100644 search.go create mode 100644 types.go diff --git a/fts5.go b/fts5.go new file mode 100644 index 0000000..839eb6e --- /dev/null +++ b/fts5.go @@ -0,0 +1,23 @@ +package main + +import "context" + +type fts5Engine struct { + sqlite *sqliteStore +} + +func newFTS5Engine(sqlite *sqliteStore) *fts5Engine { + return &fts5Engine{sqlite: sqlite} +} + +func (e *fts5Engine) Index(ctx context.Context, id, _ int64, _ int, note string) error { + return e.sqlite.indexFTS(ctx, id, note) +} + +func (e *fts5Engine) Delete(ctx context.Context, id, _ int64) error { + return e.sqlite.deleteFTS(ctx, id) +} + +func (e *fts5Engine) Search(ctx context.Context, query string, limit uint64) ([]searchResult, error) { + return e.sqlite.searchFTS(ctx, query, limit) +} diff --git a/main.go b/main.go index 2031694..ad304ff 100644 --- a/main.go +++ b/main.go @@ -30,43 +30,10 @@ func main() { } defer logger.Sync() - qdrantURL := os.Getenv("QDRANT_URL") - openAIKey := requireEnv(logger, "OPENAI_API_KEY") token := requireEnv(logger, "MCP_TOKEN") - embeddingModel := os.Getenv("EMBEDDING_MODEL") - embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL") + openAIKey := requireEnv(logger, "OPENAI_API_KEY") port := envOrDefault("PORT", "8080") - qdrantPort := 6334 - if p, err := strconv.Atoi(envOrDefault("QDRANT_PORT", "6334")); err == nil { - qdrantPort = p - } - qdrantAPIKey := os.Getenv("QDRANT_API_KEY") - qdrantTLS := os.Getenv("QDRANT_TLS") == "true" - - 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) - 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) if err != nil { @@ -74,6 +41,8 @@ func main() { } defer sqliteStore.close() + searchEngine := createSearchEngine(logger, sqliteStore, openAIKey) + mcpServer := server.NewMCPServer( "flashcard-mcp", version, @@ -81,10 +50,10 @@ func main() { server.WithLogging(), ) - registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger) - registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr) - registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger) - registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger) + registerStoreFlashcard(mcpServer, sqliteStore, searchEngine, logger) + registerSearchFlashcard(mcpServer, searchEngine) + registerUpdateFlashcard(mcpServer, sqliteStore, searchEngine, logger) + registerDeleteFlashcard(mcpServer, sqliteStore, searchEngine, logger) streamableServer := server.NewStreamableHTTPServer(mcpServer) @@ -128,7 +97,44 @@ func main() { } } -func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { +func createSearchEngine(logger *zap.Logger, sqliteStore *sqliteStore, openAIKey string) SearchEngine { + qdrantURL := os.Getenv("QDRANT_URL") + if qdrantURL == "" { + logger.Warn("QDRANT_URL not set, using SQLite FTS5 for search") + return newFTS5Engine(sqliteStore) + } + + embeddingModel := os.Getenv("EMBEDDING_MODEL") + embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL") + + qdrantPort := 6334 + if p, err := strconv.Atoi(envOrDefault("QDRANT_PORT", "6334")); err == nil { + qdrantPort = p + } + qdrantAPIKey := os.Getenv("QDRANT_API_KEY") + qdrantTLS := os.Getenv("QDRANT_TLS") == "true" + + client, 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)) + } + + emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL) + store := newQdrantStore(client, emb, sqliteStore, logger) + + if err := store.ensureCollection(context.Background()); err != nil { + logger.Fatal("failed to ensure collection", zap.Error(err)) + } + + return store +} + +func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) { mcpServer.AddTool( mcp.NewTool("store_flashcard", mcp.WithDescription("Store a flashcard for later retrieval. Use this to persist useful information discovered during a session."), @@ -155,11 +161,11 @@ func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qd mcp.Description("Optional context — what triggered this"), ), ), - handleStoreFlashcard(sqlite, qdrant, logger), + handleStoreFlashcard(sqlite, engine, logger), ) } -func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore) { +func registerSearchFlashcard(mcpServer *server.MCPServer, engine SearchEngine) { mcpServer.AddTool( mcp.NewTool("search_flashcard", mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant entries for the given query."), @@ -171,11 +177,11 @@ func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, q mcp.Description("Maximum number of results to return (default: 5)"), ), ), - handleSearchFlashcard(sqlite, qdrant), + handleSearchFlashcard(engine), ) } -func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { +func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) { mcpServer.AddTool( mcp.NewTool("update_flashcard", mcp.WithDescription("Update an existing flashcard. Creates a new version while preserving history."), @@ -202,11 +208,11 @@ func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, q mcp.Description("Optional context — what triggered this"), ), ), - handleUpdateFlashcard(sqlite, qdrant, logger), + handleUpdateFlashcard(sqlite, engine, logger), ) } -func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc { +func handleUpdateFlashcard(sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { args, ok := request.Params.Arguments.(map[string]any) if !ok { @@ -258,16 +264,11 @@ func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap 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 := engine.Delete(ctx, latest.ID, publicID); err != nil { + logger.Warn("search engine 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)) - } + if err := engine.Index(ctx, id, publicID, version, note); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("search engine index failed: %v", err)), nil } entry, err := sqlite.getByID(ctx, id) @@ -284,20 +285,20 @@ func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap } } -func registerDeleteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { +func registerDeleteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) { mcpServer.AddTool( mcp.NewTool("delete_flashcard", - mcp.WithDescription("Delete a flashcard. Creates a tombstone version and removes it from the vector index."), + mcp.WithDescription("Delete a flashcard. Creates a tombstone version and removes it from the search index."), mcp.WithNumber("public_id", mcp.Required(), mcp.Description("Public identifier of the flashcard to delete"), ), ), - handleDeleteFlashcard(sqlite, qdrant, logger), + handleDeleteFlashcard(sqlite, engine, logger), ) } -func handleDeleteFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc { +func handleDeleteFlashcard(sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { args, ok := request.Params.Arguments.(map[string]any) if !ok { @@ -324,20 +325,15 @@ func handleDeleteFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap 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 { - logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err)) - } + if err := engine.Delete(ctx, latest.ID, publicID); err != nil { + logger.Warn("search engine delete failed", zap.Int64("id", latest.ID), zap.Error(err)) } return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil } } -func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc { +func handleStoreFlashcard(sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { args, ok := request.Params.Arguments.(map[string]any) if !ok { @@ -364,13 +360,8 @@ func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap. if publicID == 0 { 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 { - logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) - } + if err := engine.Index(ctx, id, publicID, 1, note); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("search engine index failed: %v", err)), nil } entry, err := sqlite.getByID(ctx, id) @@ -387,7 +378,7 @@ func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap. } } -func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.ToolHandlerFunc { +func handleSearchFlashcard(engine SearchEngine) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { args, ok := request.Params.Arguments.(map[string]any) if !ok { @@ -411,23 +402,9 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool } } - 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 - } + results, err := engine.Search(ctx, query, limit) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil } output, err := json.MarshalIndent(results, "", " ") diff --git a/qdrant.go b/qdrant.go index efa7403..70bfed3 100644 --- a/qdrant.go +++ b/qdrant.go @@ -6,36 +6,20 @@ import ( "strconv" "github.com/qdrant/go-client/qdrant" + "go.uber.org/zap" ) const collectionName = "flashcard" -type entry struct { - ID int64 `json:"_id"` - PublicID int64 `json:"public_id"` - Version int `json:"version"` - Note string `json:"note"` - Repos []string `json:"repos,omitempty"` - Host string `json:"host,omitempty"` - Tags []string `json:"tags,omitempty"` - User string `json:"user,omitempty"` - Context string `json:"context,omitempty"` - DeleteTimestamp string `json:"delete_timestamp,omitempty"` - Created string `json:"created"` -} - -type searchResult struct { - Entry entry `json:",inline"` - Score float32 `json:"score"` -} - type qdrantStore struct { client *qdrant.Client embedder *embedder + sqlite *sqliteStore + logger *zap.Logger } -func newQdrantStore(client *qdrant.Client, embedder *embedder) *qdrantStore { - return &qdrantStore{client: client, embedder: embedder} +func newQdrantStore(client *qdrant.Client, embedder *embedder, sqlite *sqliteStore, logger *zap.Logger) *qdrantStore { + return &qdrantStore{client: client, embedder: embedder, sqlite: sqlite, logger: logger} } func (s *qdrantStore) ensureCollection(ctx context.Context) error { @@ -61,7 +45,7 @@ func (s *qdrantStore) ensureCollection(ctx context.Context) error { return nil } -func (s *qdrantStore) index(ctx context.Context, id, publicID int64, version int, note string) error { +func (s *qdrantStore) Index(ctx context.Context, id, publicID int64, version int, note string) error { vec, err := s.embedder.embedText(ctx, note) if err != nil { return fmt.Errorf("embed text: %w", err) @@ -85,13 +69,12 @@ func (s *qdrantStore) index(ctx context.Context, id, publicID int64, version int }, }) if err != nil { - return fmt.Errorf("upsert point: %w", err) + s.logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) } - return nil } -func (s *qdrantStore) delete(ctx context.Context, publicID int64) error { +func (s *qdrantStore) Delete(ctx context.Context, id, publicID int64) error { wait := true _, err := s.client.Delete(ctx, &qdrant.DeletePoints{ CollectionName: collectionName, @@ -101,12 +84,12 @@ func (s *qdrantStore) delete(ctx context.Context, publicID int64) error { ), }) if err != nil { - return fmt.Errorf("delete point: %w", err) + s.logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err)) } return nil } -func (s *qdrantStore) search(ctx context.Context, query string, limit uint64) ([]searchResult, error) { +func (s *qdrantStore) Search(ctx context.Context, query string, limit uint64) ([]searchResult, error) { vec, err := s.embedder.embedText(ctx, query) if err != nil { return nil, fmt.Errorf("embed query: %w", err) @@ -140,6 +123,10 @@ func (s *qdrantStore) search(ctx context.Context, query string, limit uint64) ([ } } + if cached, err := s.sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil { + r.Entry = *cached + } + searchResults = append(searchResults, r) } diff --git a/search.go b/search.go new file mode 100644 index 0000000..92546db --- /dev/null +++ b/search.go @@ -0,0 +1,9 @@ +package main + +import "context" + +type SearchEngine interface { + Index(ctx context.Context, id, publicID int64, version int, note string) error + Delete(ctx context.Context, id, publicID int64) error + Search(ctx context.Context, query string, limit uint64) ([]searchResult, error) +} diff --git a/types.go b/types.go new file mode 100644 index 0000000..45d33a3 --- /dev/null +++ b/types.go @@ -0,0 +1,20 @@ +package main + +type entry struct { + ID int64 `json:"_id"` + PublicID int64 `json:"public_id"` + Version int `json:"version"` + Note string `json:"note"` + Repos []string `json:"repos,omitempty"` + Host string `json:"host,omitempty"` + Tags []string `json:"tags,omitempty"` + User string `json:"user,omitempty"` + Context string `json:"context,omitempty"` + DeleteTimestamp string `json:"delete_timestamp,omitempty"` + Created string `json:"created"` +} + +type searchResult struct { + Entry entry `json:",inline"` + Score float32 `json:"score"` +}