refactor: SearchEngine interface, remove nilable qdrantStore

This commit is contained in:
didier committed 2026-06-05 16:02:31 -04:00
1 parent 98aa7dd68c
commit ffdc023065
5 files changed
+134 -118

No files matched your search

+23
View File
@@ -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)
}
+68 -91
View File
@@ -30,43 +30,10 @@ func main() {
} }
defer logger.Sync() defer logger.Sync()
qdrantURL := os.Getenv("QDRANT_URL")
openAIKey := requireEnv(logger, "OPENAI_API_KEY")
token := requireEnv(logger, "MCP_TOKEN") token := requireEnv(logger, "MCP_TOKEN")
embeddingModel := os.Getenv("EMBEDDING_MODEL") openAIKey := requireEnv(logger, "OPENAI_API_KEY")
embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL")
port := envOrDefault("PORT", "8080") 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") sqlitePath := envOrDefault("SQLITE_PATH", "flashcard.db")
sqliteStore, err := newSqliteStore(sqlitePath) sqliteStore, err := newSqliteStore(sqlitePath)
if err != nil { if err != nil {
@@ -74,6 +41,8 @@ func main() {
} }
defer sqliteStore.close() defer sqliteStore.close()
searchEngine := createSearchEngine(logger, sqliteStore, openAIKey)
mcpServer := server.NewMCPServer( mcpServer := server.NewMCPServer(
"flashcard-mcp", "flashcard-mcp",
version, version,
@@ -81,10 +50,10 @@ func main() {
server.WithLogging(), server.WithLogging(),
) )
registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger) registerStoreFlashcard(mcpServer, sqliteStore, searchEngine, logger)
registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr) registerSearchFlashcard(mcpServer, searchEngine)
registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger) registerUpdateFlashcard(mcpServer, sqliteStore, searchEngine, logger)
registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger) registerDeleteFlashcard(mcpServer, sqliteStore, searchEngine, logger)
streamableServer := server.NewStreamableHTTPServer(mcpServer) 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( mcpServer.AddTool(
mcp.NewTool("store_flashcard", mcp.NewTool("store_flashcard",
mcp.WithDescription("Store a flashcard for later retrieval. Use this to persist useful information discovered during a session."), 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"), 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( 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."),
@@ -171,11 +177,11 @@ func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, q
mcp.Description("Maximum number of results to return (default: 5)"), 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( mcpServer.AddTool(
mcp.NewTool("update_flashcard", mcp.NewTool("update_flashcard",
mcp.WithDescription("Update an existing flashcard. Creates a new version while preserving history."), 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"), 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) { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
args, ok := request.Params.Arguments.(map[string]any) args, ok := request.Params.Arguments.(map[string]any)
if !ok { 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 return mcp.NewToolResultError(fmt.Sprintf("failed to persist update: %v", err)), nil
} }
if err := sqlite.deleteFTS(ctx, latest.ID); err != nil { if err := engine.Delete(ctx, latest.ID, publicID); err != nil {
logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err)) logger.Warn("search engine delete failed", zap.Int64("id", latest.ID), zap.Error(err))
} }
if err := sqlite.indexFTS(ctx, id, note); err != nil { if err := engine.Index(ctx, id, publicID, version, note); err != nil {
return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil return mcp.NewToolResultError(fmt.Sprintf("search engine 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) 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( mcpServer.AddTool(
mcp.NewTool("delete_flashcard", 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.WithNumber("public_id",
mcp.Required(), mcp.Required(),
mcp.Description("Public identifier of the flashcard to delete"), 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) { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
args, ok := request.Params.Arguments.(map[string]any) args, ok := request.Params.Arguments.(map[string]any)
if !ok { 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 return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", err)), nil
} }
if err := sqlite.deleteFTS(ctx, latest.ID); err != nil { if err := engine.Delete(ctx, latest.ID, publicID); err != nil {
logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err)) logger.Warn("search engine 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 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) { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
args, ok := request.Params.Arguments.(map[string]any) args, ok := request.Params.Arguments.(map[string]any)
if !ok { if !ok {
@@ -364,13 +360,8 @@ 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 { if err := engine.Index(ctx, id, publicID, 1, note); err != nil {
return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil return mcp.NewToolResultError(fmt.Sprintf("search engine 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) 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) { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
args, ok := request.Params.Arguments.(map[string]any) args, ok := request.Params.Arguments.(map[string]any)
if !ok { if !ok {
@@ -411,23 +402,9 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool
} }
} }
var results []searchResult results, err := engine.Search(ctx, query, limit)
var err error if err != nil {
if qdrant != nil { return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), 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
}
} }
output, err := json.MarshalIndent(results, "", " ") output, err := json.MarshalIndent(results, "", " ")
+14 -27
View File
@@ -6,36 +6,20 @@ import (
"strconv" "strconv"
"github.com/qdrant/go-client/qdrant" "github.com/qdrant/go-client/qdrant"
"go.uber.org/zap"
) )
const collectionName = "flashcard" 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 { type qdrantStore struct {
client *qdrant.Client client *qdrant.Client
embedder *embedder embedder *embedder
sqlite *sqliteStore
logger *zap.Logger
} }
func newQdrantStore(client *qdrant.Client, embedder *embedder) *qdrantStore { func newQdrantStore(client *qdrant.Client, embedder *embedder, sqlite *sqliteStore, logger *zap.Logger) *qdrantStore {
return &qdrantStore{client: client, embedder: embedder} return &qdrantStore{client: client, embedder: embedder, sqlite: sqlite, logger: logger}
} }
func (s *qdrantStore) ensureCollection(ctx context.Context) error { func (s *qdrantStore) ensureCollection(ctx context.Context) error {
@@ -61,7 +45,7 @@ func (s *qdrantStore) ensureCollection(ctx context.Context) error {
return nil 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) vec, err := s.embedder.embedText(ctx, note)
if err != nil { if err != nil {
return fmt.Errorf("embed text: %w", err) 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 { 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 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 wait := true
_, err := s.client.Delete(ctx, &qdrant.DeletePoints{ _, err := s.client.Delete(ctx, &qdrant.DeletePoints{
CollectionName: collectionName, CollectionName: collectionName,
@@ -101,12 +84,12 @@ func (s *qdrantStore) delete(ctx context.Context, publicID int64) error {
), ),
}) })
if err != nil { 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 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) vec, err := s.embedder.embedText(ctx, query)
if err != nil { if err != nil {
return nil, fmt.Errorf("embed query: %w", err) 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) searchResults = append(searchResults, r)
} }
+9
View File
@@ -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)
}
+20
View File
@@ -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"`
}