refactor: SearchEngine interface, remove nilable qdrantStore
This commit is contained in:
1 parent
98aa7dd68c
commit
ffdc023065
5 files changed
+132
-116
No files matched your search
@@ -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)
|
||||
}
|
||||
@@ -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 := 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.Delete(ctx, latest.ID, publicID); err != nil {
|
||||
logger.Warn("search engine delete failed", zap.Int64("id", latest.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,24 +402,10 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool
|
||||
}
|
||||
}
|
||||
|
||||
var results []searchResult
|
||||
var err error
|
||||
if qdrant != nil {
|
||||
results, err = qdrant.search(ctx, query, limit)
|
||||
results, err := engine.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, "", " ")
|
||||
if err != nil {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
Reference in new issue
Block a user