feat: added fts5 backend
This commit is contained in:
1 parent
d5cae6efa9
commit
98aa7dd68c
6 files changed
+118
-14
No files matched your search
@@ -10,3 +10,4 @@ result
|
|||||||
# IDE
|
# IDE
|
||||||
.idea/
|
.idea/
|
||||||
.vscode/
|
.vscode/
|
||||||
|
*.db
|
||||||
@@ -14,6 +14,7 @@ require (
|
|||||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||||
github.com/google/jsonschema-go v0.4.2 // indirect
|
github.com/google/jsonschema-go v0.4.2 // indirect
|
||||||
github.com/google/uuid v1.6.0 // 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/mattn/go-isatty v0.0.21 // indirect
|
||||||
github.com/mfridman/interpolate v0.0.2 // indirect
|
github.com/mfridman/interpolate v0.0.2 // indirect
|
||||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||||
|
|||||||
@@ -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/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 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
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 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ import (
|
|||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/joho/godotenv"
|
||||||
"github.com/mark3labs/mcp-go/mcp"
|
"github.com/mark3labs/mcp-go/mcp"
|
||||||
"github.com/mark3labs/mcp-go/server"
|
"github.com/mark3labs/mcp-go/server"
|
||||||
"github.com/qdrant/go-client/qdrant"
|
"github.com/qdrant/go-client/qdrant"
|
||||||
@@ -21,13 +22,15 @@ import (
|
|||||||
var version = "0.1.0"
|
var version = "0.1.0"
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
|
godotenv.Load()
|
||||||
|
|
||||||
logger, err := zap.NewProduction()
|
logger, err := zap.NewProduction()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
panic(fmt.Sprintf("failed to init logger: %v", err))
|
panic(fmt.Sprintf("failed to init logger: %v", err))
|
||||||
}
|
}
|
||||||
defer logger.Sync()
|
defer logger.Sync()
|
||||||
|
|
||||||
qdrantURL := requireEnv(logger, "QDRANT_URL")
|
qdrantURL := os.Getenv("QDRANT_URL")
|
||||||
openAIKey := requireEnv(logger, "OPENAI_API_KEY")
|
openAIKey := requireEnv(logger, "OPENAI_API_KEY")
|
||||||
token := requireEnv(logger, "MCP_TOKEN")
|
token := requireEnv(logger, "MCP_TOKEN")
|
||||||
embeddingModel := os.Getenv("EMBEDDING_MODEL")
|
embeddingModel := os.Getenv("EMBEDDING_MODEL")
|
||||||
@@ -41,6 +44,8 @@ func main() {
|
|||||||
qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
|
qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
|
||||||
qdrantTLS := os.Getenv("QDRANT_TLS") == "true"
|
qdrantTLS := os.Getenv("QDRANT_TLS") == "true"
|
||||||
|
|
||||||
|
var qdrantStr *qdrantStore
|
||||||
|
if qdrantURL != "" {
|
||||||
qdrantClient, err := qdrant.NewClient(&qdrant.Config{
|
qdrantClient, err := qdrant.NewClient(&qdrant.Config{
|
||||||
Host: qdrantURL,
|
Host: qdrantURL,
|
||||||
Port: qdrantPort,
|
Port: qdrantPort,
|
||||||
@@ -53,7 +58,14 @@ func main() {
|
|||||||
defer qdrantClient.Close()
|
defer qdrantClient.Close()
|
||||||
|
|
||||||
emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
|
emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
|
||||||
qdrantStore := newQdrantStore(qdrantClient, emb)
|
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)
|
||||||
@@ -62,10 +74,6 @@ func main() {
|
|||||||
}
|
}
|
||||||
defer sqliteStore.close()
|
defer sqliteStore.close()
|
||||||
|
|
||||||
if err := qdrantStore.ensureCollection(context.Background()); err != nil {
|
|
||||||
logger.Fatal("failed to ensure collection", zap.Error(err))
|
|
||||||
}
|
|
||||||
|
|
||||||
mcpServer := server.NewMCPServer(
|
mcpServer := server.NewMCPServer(
|
||||||
"flashcard-mcp",
|
"flashcard-mcp",
|
||||||
version,
|
version,
|
||||||
@@ -73,10 +81,10 @@ func main() {
|
|||||||
server.WithLogging(),
|
server.WithLogging(),
|
||||||
)
|
)
|
||||||
|
|
||||||
registerStoreFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
|
||||||
registerSearchFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr)
|
||||||
registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
|
||||||
registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
|
||||||
|
|
||||||
streamableServer := server.NewStreamableHTTPServer(mcpServer)
|
streamableServer := server.NewStreamableHTTPServer(mcpServer)
|
||||||
|
|
||||||
@@ -111,7 +119,6 @@ func main() {
|
|||||||
if err := streamableServer.Shutdown(shutdownCtx); err != nil {
|
if err := streamableServer.Shutdown(shutdownCtx); err != nil {
|
||||||
logger.Error("shutdown error", zap.Error(err))
|
logger.Error("shutdown error", zap.Error(err))
|
||||||
}
|
}
|
||||||
qdrantClient.Close()
|
|
||||||
sqliteStore.close()
|
sqliteStore.close()
|
||||||
os.Exit(0)
|
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(
|
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."),
|
||||||
@@ -251,9 +258,17 @@ 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 {
|
||||||
|
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 {
|
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))
|
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)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -309,9 +324,14 @@ 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 {
|
||||||
|
logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err))
|
||||||
|
}
|
||||||
|
if qdrant != nil {
|
||||||
if err := qdrant.delete(ctx, publicID); 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))
|
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
|
||||||
}
|
}
|
||||||
@@ -344,9 +364,14 @@ 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 {
|
||||||
|
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 {
|
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))
|
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)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -386,16 +411,24 @@ func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.Tool
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
results, err := qdrant.search(ctx, query, limit)
|
var results []searchResult
|
||||||
|
var err error
|
||||||
|
if qdrant != nil {
|
||||||
|
results, err = qdrant.search(ctx, query, limit)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
|
return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
for i, r := range results {
|
for i, r := range results {
|
||||||
if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
|
if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
|
||||||
results[i].Entry = *cached
|
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, "", " ")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -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;
|
||||||
@@ -184,6 +184,68 @@ func (s *sqliteStore) tombstone(ctx context.Context, publicID int64, version int
|
|||||||
return id, nil
|
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 {
|
func (s *sqliteStore) close() error {
|
||||||
return s.db.Close()
|
return s.db.Close()
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in new issue
Block a user