feat: added fts5 backend
This commit is contained in:
didier committed
1 parent
d5cae6efa9
commit
98aa7dd68c
6 files changed
+141
-37
No files matched your search
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user