diff --git a/.env.example b/.env.example index 6096d2f..dd08071 100644 --- a/.env.example +++ b/.env.example @@ -7,3 +7,4 @@ OPENAI_API_KEY=sk-... EMBEDDING_MODEL=text-embedding-3-small MCP_TOKEN=change-me PORT=8080 +SQLITE_PATH=knowledge.db diff --git a/flake.nix b/flake.nix index cae191a..cbfa5f4 100644 --- a/flake.nix +++ b/flake.nix @@ -98,6 +98,12 @@ default = null; description = "MCP Bearer token directly. Prefer tokenFile for secrets."; }; + + sqlitePath = lib.mkOption { + type = lib.types.path; + default = "/var/lib/st-pierre/st-pierre.sqlite"; + description = "Path to the SQLite database file."; + }; }; config = lib.mkIf cfg.enable { @@ -119,6 +125,7 @@ Group = "st-pierre"; Restart = "always"; WorkingDirectory = "/var/lib/st-pierre"; + StateDirectory = "st-pierre"; ExecStart = "${st-pierre-pkg}/bin/st-pierre"; EnvironmentFile = lib.optional (cfg.openaiApiKeyFile != null) cfg.openaiApiKeyFile ++ lib.optional (cfg.tokenFile != null) cfg.tokenFile @@ -133,6 +140,7 @@ QDRANT_TLS = if cfg.qdrantTls then "true" else "false"; EMBEDDING_BASE_URL = cfg.embeddingBaseUrl; EMBEDDING_MODEL = cfg.embeddingModel; + SQLITE_PATH = builtins.toString cfg.sqlitePath; } // lib.optionalAttrs (cfg.openaiApiKey != null) { OPENAI_API_KEY = cfg.openaiApiKey; diff --git a/go.mod b/go.mod index a62471c..f22ef0d 100644 --- a/go.mod +++ b/go.mod @@ -3,20 +3,34 @@ module github.com/didier/st-pierre go 1.26.2 require ( - github.com/google/uuid v1.6.0 github.com/mark3labs/mcp-go v0.54.1 + github.com/pressly/goose/v3 v3.27.1 github.com/qdrant/go-client v1.18.2 + go.uber.org/zap v1.28.0 + modernc.org/sqlite v1.51.0 ) require ( + github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/jsonschema-go v0.4.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/mattn/go-isatty v0.0.21 // indirect + github.com/mfridman/interpolate v0.0.2 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect + github.com/sethvargo/go-retry v0.3.0 // indirect github.com/spf13/cast v1.7.1 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + go.uber.org/multierr v1.11.0 // indirect golang.org/x/net v0.53.0 // indirect + golang.org/x/sync v0.20.0 // indirect golang.org/x/sys v0.43.0 // indirect golang.org/x/text v0.36.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260427160629-7cedc36a6bc4 // indirect google.golang.org/grpc v1.80.0 // indirect google.golang.org/protobuf v1.36.11 // indirect + modernc.org/libc v1.72.3 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index aea0325..de8e8b8 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI= github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= @@ -16,22 +18,38 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8= github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= 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/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/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= 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/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mark3labs/mcp-go v0.54.1 h1:Ap/ptEB9FtWzFKM8NDsTA7QDxerQOC06eZigrTldVj0= github.com/mark3labs/mcp-go v0.54.1/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas= +github.com/mattn/go-isatty v0.0.21 h1:xYae+lCNBP7QuW4PUnNG61ffM4hVIfm+zUzDuSzYLGs= +github.com/mattn/go-isatty v0.0.21/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= +github.com/mfridman/interpolate v0.0.2 h1:pnuTK7MQIxxFz1Gr+rjSIx9u7qVjf5VOoM/u6BbAxPY= +github.com/mfridman/interpolate v0.0.2/go.mod h1:p+7uk6oE07mpE/Ik1b8EckO0O4ZXiGAfshKBWLUM9Xg= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pressly/goose/v3 v3.27.1 h1:6uEvcprBybDmW4hcz3gYujhARhye+GoWKhEWyzD5sh4= +github.com/pressly/goose/v3 v3.27.1/go.mod h1:maruOxsPnIG2yHHyo8UqKWXYKFcH7Q76csUV7+7KYoM= github.com/qdrant/go-client v1.18.2 h1:7ViiXB/fB4vfzdUtEYZ7g2vSG+yMU6EMu8CaNWF1g1c= github.com/qdrant/go-client v1.18.2/go.mod h1:Xkfp+r89uNOgSbvilVAhCZ3wKI4G+hB/r9Zr2m4zifI= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ= github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= +github.com/sethvargo/go-retry v0.3.0 h1:EEt31A35QhrcRZtrYFDTBg91cqZVnFL2navjDrah2SE= +github.com/sethvargo/go-retry v0.3.0/go.mod h1:mNX17F0C/HguQMyMyJxcnU471gOZGxCLyYaFyAZraas= github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= @@ -50,12 +68,26 @@ go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2W go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew= go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= +go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= +go.uber.org/zap v1.28.0 h1:IZzaP1Fv73/T/pBMLk4VutPl36uNC+OSUh3JLG3FIjo= +go.uber.org/zap v1.28.0/go.mod h1:rDLpOi171uODNm/mxFcuYWxDsqWSAVkFdX4XojSKg/Q= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI= +golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY= golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg= golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164= +golang.org/x/tools v0.43.0 h1:12BdW9CeB3Z+J/I/wj34VMl8X+fEXBxVR90JeMX5E7s= +golang.org/x/tools v0.43.0/go.mod h1:uHkMso649BX2cZK6+RpuIPXS3ho2hZo4FVwfoy1vIk0= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/genproto/googleapis/rpc v0.0.0-20260427160629-7cedc36a6bc4 h1:tEkOQcXgF6dH1G+MVKZrfpYvozGrzb91k6ha7jireSM= @@ -66,3 +98,31 @@ google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBN google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY= +modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ= +modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU= +modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.51.0 h1:aH/MMSoayAIhozZ7uJbVTT9QO/VhzBf0J9tymmmuC/U= +modernc.org/sqlite v1.51.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/main.go b/main.go index d520649..5520848 100644 --- a/main.go +++ b/main.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "fmt" - "log" "net/http" "os" "os/signal" @@ -16,14 +15,21 @@ import ( "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "github.com/qdrant/go-client/qdrant" + "go.uber.org/zap" ) var version = "0.1.0" func main() { - qdrantURL := requireEnv("QDRANT_URL") - openAIKey := requireEnv("OPENAI_API_KEY") - token := requireEnv("MCP_TOKEN") + logger, err := zap.NewProduction() + if err != nil { + panic(fmt.Sprintf("failed to init logger: %v", err)) + } + defer logger.Sync() + + qdrantURL := requireEnv(logger, "QDRANT_URL") + openAIKey := requireEnv(logger, "OPENAI_API_KEY") + token := requireEnv(logger, "MCP_TOKEN") embeddingModel := os.Getenv("EMBEDDING_MODEL") embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL") port := envOrDefault("PORT", "8080") @@ -42,15 +48,22 @@ func main() { UseTLS: qdrantTLS, }) if err != nil { - log.Fatalf("failed to create qdrant client: %v", err) + logger.Fatal("failed to create qdrant client", zap.Error(err)) } defer qdrantClient.Close() emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL) - store := newQdrantStore(qdrantClient, emb) + qdrantStore := newQdrantStore(qdrantClient, emb) - if err := store.ensureCollection(context.Background()); err != nil { - log.Fatalf("failed to ensure collection: %v", err) + sqlitePath := envOrDefault("SQLITE_PATH", "knowledge.db") + sqliteStore, err := newSqliteStore(sqlitePath) + if err != nil { + logger.Fatal("failed to init sqlite", zap.Error(err)) + } + defer sqliteStore.close() + + if err := qdrantStore.ensureCollection(context.Background()); err != nil { + logger.Fatal("failed to ensure collection", zap.Error(err)) } mcpServer := server.NewMCPServer( @@ -60,7 +73,8 @@ func main() { server.WithLogging(), ) - registerTools(mcpServer, store) + registerStoreKnowledge(mcpServer, sqliteStore, qdrantStore, logger) + registerSearchKnowledge(mcpServer, sqliteStore, qdrantStore, logger) streamableServer := server.NewStreamableHTTPServer(mcpServer) @@ -71,42 +85,57 @@ func main() { Handler: corsMiddleware(authHandler), } - log.Printf("MCP Knowledge Server starting on http://0.0.0.0:%s", port) + logger.Info("starting", zap.String("port", port)) go func() { sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) sig := <-sigChan - log.Printf("received signal %v, shutting down...", sig) + logger.Info("shutting down", zap.String("signal", sig.String())) shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if err := streamableServer.Shutdown(shutdownCtx); err != nil { - log.Printf("shutdown error: %v", err) + logger.Error("shutdown error", zap.Error(err)) } qdrantClient.Close() + sqliteStore.close() os.Exit(0) }() if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { - log.Fatalf("server stopped: %v", err) + logger.Fatal("server stopped", zap.Error(err)) } } -func registerTools(mcpServer *server.MCPServer, store *qdrantStore) { +func registerStoreKnowledge(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { mcpServer.AddTool( mcp.NewTool("store_knowledge", mcp.WithDescription("Store a piece of knowledge for later retrieval. Use this to persist useful information discovered during a session."), - mcp.WithString("content", + mcp.WithString("note", mcp.Required(), - mcp.Description("The text content to store"), + mcp.Description("The markdown body — keep it concise, one paragraph usually"), ), - mcp.WithObject("metadata", - mcp.Description("Optional metadata as key-value pairs (e.g. {\"source\": \"docs\", \"category\": \"golang\"})"), + mcp.WithArray("repos", + mcp.Description("Optional list of repository names"), + ), + mcp.WithString("host", + mcp.Description("Optional hostname or null if universal"), + ), + mcp.WithArray("tags", + mcp.Description("Optional list of tags"), + ), + mcp.WithString("user", + mcp.Description("Optional user identifier"), + ), + mcp.WithString("context", + mcp.Description("Optional context — what triggered this"), ), ), - handleStoreKnowledge(store), + handleStoreKnowledge(sqlite, qdrant, logger), ) +} +func registerSearchKnowledge(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) { mcpServer.AddTool( mcp.NewTool("search_knowledge", mcp.WithDescription("Search previously stored knowledge using semantic search. Returns the most relevant knowledge entries for the given query."), @@ -118,42 +147,52 @@ func registerTools(mcpServer *server.MCPServer, store *qdrantStore) { mcp.Description("Maximum number of results to return (default: 5)"), ), ), - handleSearchKnowledge(store), + handleSearchKnowledge(sqlite, qdrant), ) } -func handleStoreKnowledge(store *qdrantStore) server.ToolHandlerFunc { +func handleStoreKnowledge(sqlite *sqliteStore, qdrant *qdrantStore, 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 { return mcp.NewToolResultError("invalid arguments"), nil } - content, ok := args["content"].(string) - if !ok || content == "" { - return mcp.NewToolResultError("content is required and must be a non-empty string"), nil + note, ok := args["note"].(string) + if !ok || note == "" { + return mcp.NewToolResultError("note is required and must be a non-empty string"), nil } - var metadata map[string]any - if rawMeta, ok := args["metadata"]; ok && rawMeta != nil { - metaMap, ok := rawMeta.(map[string]any) - if !ok { - return mcp.NewToolResultError("metadata must be a key-value object"), nil - } - metadata = metaMap - } + repos := parseStringSlice(args["repos"]) + host := parseOptionalString(args["host"]) + tags := parseStringSlice(args["tags"]) + user := parseOptionalString(args["user"]) + context_ := parseOptionalString(args["context"]) - result, err := store.store(ctx, content, metadata) + id, err := sqlite.save(ctx, note, host, user, context_, repos, tags) if err != nil { - return mcp.NewToolResultError(fmt.Sprintf("failed to store knowledge: %v", err)), nil + return mcp.NewToolResultError(fmt.Sprintf("failed to persist: %v", err)), nil + } + + if err := qdrant.index(ctx, id, note); err != nil { + logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err)) + } + + entry, err := sqlite.getByID(ctx, id) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve: %v", err)), nil + } + + output, err := json.MarshalIndent(entry, "", " ") + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), nil } - output, _ := json.MarshalIndent(result, "", " ") return mcp.NewToolResultText(string(output)), nil } } -func handleSearchKnowledge(store *qdrantStore) server.ToolHandlerFunc { +func handleSearchKnowledge(sqlite *sqliteStore, qdrant *qdrantStore) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { args, ok := request.Params.Arguments.(map[string]any) if !ok { @@ -177,12 +216,22 @@ func handleSearchKnowledge(store *qdrantStore) server.ToolHandlerFunc { } } - results, err := store.search(ctx, query, limit) + results, err := qdrant.search(ctx, query, limit) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil } - output, _ := json.MarshalIndent(results, "", " ") + for i, r := range results { + if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil { + results[i].Entry = *cached + } + } + + output, err := json.MarshalIndent(results, "", " ") + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("failed to marshal results: %v", err)), nil + } + return mcp.NewToolResultText(string(output)), nil } } @@ -195,8 +244,8 @@ func authMiddleware(token string, next http.Handler) http.HandlerFunc { return } - authHeader := r.Header.Get("Authorization") - if !strings.HasPrefix(authHeader, "Bearer ") || strings.TrimPrefix(authHeader, "Bearer ") != token { + authHeader := strings.ToLower(r.Header.Get("Authorization")) + if !strings.HasPrefix(authHeader, "bearer ") || strings.TrimPrefix(authHeader, "bearer ") != token { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusUnauthorized) w.Write([]byte(`{"error":"unauthorized"}`)) @@ -222,10 +271,10 @@ func corsMiddleware(next http.Handler) http.HandlerFunc { } } -func requireEnv(key string) string { +func requireEnv(logger *zap.Logger, key string) string { val := os.Getenv(key) if val == "" { - log.Fatalf("environment variable %s is required", key) + logger.Fatal("missing required env var", zap.String("key", key)) } return val } diff --git a/migrations/000001_create_knowledge.sql b/migrations/000001_create_knowledge.sql new file mode 100644 index 0000000..ea1c2c3 --- /dev/null +++ b/migrations/000001_create_knowledge.sql @@ -0,0 +1,14 @@ +-- +goose Up +CREATE TABLE IF NOT EXISTS knowledge ( + id INTEGER PRIMARY KEY, + note TEXT NOT NULL, + repos TEXT, + host TEXT, + tags TEXT, + "user" TEXT, + context TEXT, + created TEXT NOT NULL DEFAULT (datetime('now')) +); + +-- +goose Down +DROP TABLE IF EXISTS knowledge; diff --git a/qdrant.go b/qdrant.go index 45ec1cd..8e91aa9 100644 --- a/qdrant.go +++ b/qdrant.go @@ -2,16 +2,30 @@ package main import ( "context" - "crypto/sha256" - "encoding/json" "fmt" + "strconv" - "github.com/google/uuid" "github.com/qdrant/go-client/qdrant" ) const collectionName = "knowledge" +type entry struct { + ID int64 `json:"id"` + 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"` + Created string `json:"created"` +} + +type searchResult struct { + Entry entry `json:",inline"` + Score float32 `json:"score"` +} + type qdrantStore struct { client *qdrant.Client embedder *embedder @@ -33,7 +47,7 @@ func (s *qdrantStore) ensureCollection(ctx context.Context) error { err = s.client.CreateCollection(ctx, &qdrant.CreateCollection{ CollectionName: collectionName, VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{ - Size: 1536, // text-embedding-3-small dimension + Size: 1536, Distance: qdrant.Distance_Cosine, }), }) @@ -44,63 +58,32 @@ func (s *qdrantStore) ensureCollection(ctx context.Context) error { return nil } -type storedKnowledge struct { - ID string `json:"id"` - Content string `json:"content"` - Metadata map[string]any `json:"metadata,omitempty"` -} - -type searchResult struct { - ID string `json:"id"` - Content string `json:"content"` - Score float32 `json:"score"` - Metadata map[string]any `json:"metadata,omitempty"` -} - -func (s *qdrantStore) store(ctx context.Context, content string, metadata map[string]any) (*storedKnowledge, error) { - vec, err := s.embedder.embedText(ctx, content) +func (s *qdrantStore) index(ctx context.Context, id int64, note string) error { + vec, err := s.embedder.embedText(ctx, note) if err != nil { - return nil, fmt.Errorf("embed text: %w", err) + return fmt.Errorf("embed text: %w", err) } - id := uuid.New().String() payload := map[string]*qdrant.Value{ - "content": qdrant.NewValueString(content), - } - - if metadata != nil { - metaBytes, err := json.Marshal(metadata) - if err != nil { - return nil, fmt.Errorf("marshal metadata: %w", err) - } - payload["_metadata"] = qdrant.NewValueString(string(metaBytes)) - } - - idHash := sha256.Sum256([]byte(id)) - idNum := uint64(0) - for i := 0; i < 8; i++ { - idNum = (idNum << 8) | uint64(idHash[i]) + "id": qdrant.NewValueString(strconv.FormatInt(id, 10)), + "content": qdrant.NewValueString(note), } _, err = s.client.Upsert(ctx, &qdrant.UpsertPoints{ CollectionName: collectionName, Points: []*qdrant.PointStruct{ { - Id: qdrant.NewIDNum(idNum), + Id: qdrant.NewIDNum(uint64(id)), Vectors: qdrant.NewVectors(vec...), Payload: payload, }, }, }) if err != nil { - return nil, fmt.Errorf("upsert point: %w", err) + return fmt.Errorf("upsert point: %w", err) } - return &storedKnowledge{ - ID: id, - Content: content, - Metadata: metadata, - }, nil + return nil } func (s *qdrantStore) search(ctx context.Context, query string, limit uint64) ([]searchResult, error) { @@ -130,14 +113,9 @@ func (s *qdrantStore) search(ctx context.Context, query string, limit uint64) ([ } if payload := pt.GetPayload(); payload != nil { - if v, ok := payload["content"]; ok { - r.ID = fmt.Sprintf("%d", pt.GetId().GetNum()) - r.Content = v.GetStringValue() - } - if v, ok := payload["_metadata"]; ok { - var meta map[string]any - if err := json.Unmarshal([]byte(v.GetStringValue()), &meta); err == nil { - r.Metadata = meta + if v, ok := payload["id"]; ok { + if n, err := strconv.ParseInt(v.GetStringValue(), 10, 64); err == nil { + r.Entry.ID = n } } } diff --git a/sqlite.go b/sqlite.go new file mode 100644 index 0000000..94deddf --- /dev/null +++ b/sqlite.go @@ -0,0 +1,148 @@ +package main + +import ( + "context" + "database/sql" + "embed" + "encoding/json" + "fmt" + "strings" + + "github.com/pressly/goose/v3" + _ "modernc.org/sqlite" +) + +//go:embed migrations/*.sql +var embedMigrations embed.FS + +type sqliteStore struct { + db *sql.DB +} + +func newSqliteStore(path string) (*sqliteStore, error) { + db, err := sql.Open("sqlite", path) + if err != nil { + return nil, fmt.Errorf("open sqlite: %w", err) + } + + db.SetMaxOpenConns(1) + + goose.SetBaseFS(embedMigrations) + + if err := goose.SetDialect("sqlite3"); err != nil { + db.Close() + return nil, fmt.Errorf("set dialect: %w", err) + } + + if err := goose.Up(db, "migrations"); err != nil { + db.Close() + return nil, fmt.Errorf("migrate: %w", err) + } + + return &sqliteStore{db: db}, nil +} + +func (s *sqliteStore) save(ctx context.Context, note, host, user, context_ string, repos, tags []string) (int64, error) { + reposJSON := marshalStringSlice(repos) + tagsJSON := marshalStringSlice(tags) + + result, err := s.db.ExecContext(ctx, + `INSERT INTO knowledge (note, repos, host, tags, "user", context) + VALUES (?, ?, ?, ?, ?, ?)`, + note, reposJSON, host, tagsJSON, user, context_, + ) + if err != nil { + return 0, fmt.Errorf("insert: %w", err) + } + + id, err := result.LastInsertId() + if err != nil { + return 0, fmt.Errorf("last insert id: %w", err) + } + + return id, nil +} + +func (s *sqliteStore) getByID(ctx context.Context, id int64) (*entry, error) { + row := s.db.QueryRowContext(ctx, + `SELECT id, note, repos, host, tags, "user", context, created + FROM knowledge WHERE id = ?`, id, + ) + + var e entry + var reposJSON, tagsJSON sql.NullString + var host, user, context_ sql.NullString + + err := row.Scan(&e.ID, &e.Note, &reposJSON, &host, &tagsJSON, &user, &context_, &e.Created) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("query: %w", err) + } + + if reposJSON.Valid { + e.Repos = unmarshalStringSlice(reposJSON.String) + } + if tagsJSON.Valid { + e.Tags = unmarshalStringSlice(tagsJSON.String) + } + if host.Valid { + e.Host = host.String + } + if user.Valid { + e.User = user.String + } + if context_.Valid { + e.Context = context_.String + } + + return &e, nil +} + +func (s *sqliteStore) close() error { + return s.db.Close() +} + +func marshalStringSlice(ss []string) string { + if len(ss) == 0 { + return "[]" + } + b, _ := json.Marshal(ss) + return string(b) +} + +func unmarshalStringSlice(raw string) []string { + if raw == "" || raw == "[]" || raw == "null" { + return nil + } + var ss []string + if err := json.Unmarshal([]byte(raw), &ss); err != nil { + return nil + } + return ss +} + +func parseStringSlice(raw any) []string { + if raw == nil { + return nil + } + arr, ok := raw.([]any) + if !ok { + return nil + } + result := make([]string, 0, len(arr)) + for _, v := range arr { + if s, ok := v.(string); ok { + result = append(result, s) + } + } + return result +} + +func parseOptionalString(raw any) string { + if s, ok := raw.(string); ok && strings.TrimSpace(s) != "" { + return s + } + return "" +}