From c3387ec4a9e6a1e71deebfb11e2bb9a8c141ec3a Mon Sep 17 00:00:00 2001 From: mlogclub Date: Sun, 7 Jun 2026 17:32:29 +0800 Subject: [PATCH] feat: add LanceDB vector database provider - Implemented LanceDBProvider for managing vector collections, including methods for creating, deleting, and searching collections. - Added support for LanceDB configuration in VectorDBConfig. - Introduced stub for LanceDB provider when not built with the appropriate tag. - Updated provider initialization to include LanceDB as a supported type. - Created types and interfaces for vector operations, including Vector, SearchRequest, and SearchResult. - Added tests for VectorDBConfig unmarshalling and for ensuring proper error handling when LanceDB is not built. - Updated enums to include LanceDB type and labels. - Modified frontend enums to reflect the addition of LanceDB. --- config/config.example.yaml | 14 +- docker/agent-desk.yaml | 12 +- go.mod | 18 +- go.sum | 39 +- internal/ai/rag/vectordb/lancedb.go | 489 +++++++++++++++++++ internal/ai/rag/vectordb/lancedb_stub.go | 13 + internal/ai/rag/vectordb/provider.go | 4 +- internal/ai/rag/vectordb/provider_test.go | 23 + internal/ai/rag/vectordb/qdrant.go | 47 +- internal/ai/rag/vectordb/types.go | 48 ++ internal/pkg/config/config.go | 14 +- internal/pkg/config/vector_db_config_test.go | 45 ++ internal/pkg/enums/knowledge.go | 6 +- internal/pkg/enums/knowledge_test.go | 12 + web/lib/generated/enums.ts | 4 + 15 files changed, 711 insertions(+), 77 deletions(-) create mode 100644 internal/ai/rag/vectordb/lancedb.go create mode 100644 internal/ai/rag/vectordb/lancedb_stub.go create mode 100644 internal/ai/rag/vectordb/provider_test.go create mode 100644 internal/ai/rag/vectordb/types.go create mode 100644 internal/pkg/config/vector_db_config_test.go create mode 100644 internal/pkg/enums/knowledge_test.go diff --git a/config/config.example.yaml b/config/config.example.yaml index c5b696e..2f13af3 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -77,11 +77,15 @@ storage: signedUrlExpireSeconds: 600 vectorDB: - type: qdrant - host: 127.0.0.1 - grpcPort: 6334 - apiKey: "" - useTls: false + type: qdrant # qdrant, lancedb + qdrant: + host: 127.0.0.1 + grpcPort: 6334 + apiKey: "" + useTls: false + # LanceDB requires building the backend with -tags lancedb and LanceDB native libraries. + lancedb: + path: data/lancedb mcp: # Global switch for MCP tool integration. diff --git a/docker/agent-desk.yaml b/docker/agent-desk.yaml index 45f2380..137b512 100644 --- a/docker/agent-desk.yaml +++ b/docker/agent-desk.yaml @@ -46,10 +46,14 @@ storage: vectorDB: type: qdrant - host: qdrant - grpcPort: 6334 - apiKey: "" - useTls: false + qdrant: + host: qdrant + grpcPort: 6334 + apiKey: "" + useTls: false + # LanceDB requires building the backend with -tags lancedb and LanceDB native libraries. + lancedb: + path: /app/data/lancedb mcp: enabled: true diff --git a/go.mod b/go.mod index d5857a7..1ca07f3 100644 --- a/go.mod +++ b/go.mod @@ -40,6 +40,7 @@ require ( ) require ( + github.com/apache/arrow/go/v17 v17.0.0 github.com/aymerick/douceur v0.2.0 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/bmatcuk/doublestar/v4 v4.10.0 // indirect @@ -65,6 +66,7 @@ require ( github.com/go-sql-driver/mysql v1.7.0 // indirect github.com/goccy/go-json v0.10.5 // indirect github.com/goccy/go-yaml v1.19.2 // indirect + github.com/google/flatbuffers v24.3.25+incompatible // indirect github.com/google/jsonschema-go v0.4.2 // indirect github.com/goph/emperror v0.17.2 // indirect github.com/gorilla/css v1.0.0 // indirect @@ -72,17 +74,20 @@ require ( github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/compress v1.18.4 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/lancedb/lancedb-go v0.1.2 github.com/leodido/go-urn v1.4.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect - github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/meguminnnnnnnnn/go-openai v0.1.2 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/ncruces/go-strftime v0.1.9 // indirect github.com/nikolalohinski/gonja v1.5.3 // indirect github.com/nxadm/tail v1.4.11 // indirect github.com/pelletier/go-toml/v2 v2.2.4 // indirect + github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.59.0 // indirect @@ -104,6 +109,7 @@ require ( github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 // indirect github.com/yargevad/filepathx v1.0.0 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + github.com/zeebo/xxh3 v1.0.2 // indirect go.mongodb.org/mongo-driver/v2 v2.5.0 // indirect go.opentelemetry.io/otel v1.42.0 // indirect golang.org/x/arch v0.22.0 // indirect @@ -111,12 +117,14 @@ require ( golang.org/x/mod v0.33.0 // indirect golang.org/x/sync v0.20.0 // indirect golang.org/x/sys v0.42.0 // indirect + golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 // indirect golang.org/x/time v0.14.0 // indirect + golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260209200024-4cfbd4190f57 // indirect google.golang.org/grpc v1.78.0 // indirect google.golang.org/protobuf v1.36.11 // indirect - modernc.org/libc v1.22.5 // indirect - modernc.org/mathutil v1.5.0 // indirect - modernc.org/memory v1.5.0 // indirect - modernc.org/sqlite v1.23.1 // indirect + modernc.org/libc v1.41.0 // indirect + modernc.org/mathutil v1.6.0 // indirect + modernc.org/memory v1.7.2 // indirect + modernc.org/sqlite v1.29.6 // indirect ) diff --git a/go.sum b/go.sum index 056b82e..48c5922 100644 --- a/go.sum +++ b/go.sum @@ -5,6 +5,8 @@ github.com/alicebob/miniredis/v2 v2.30.0 h1:uA3uhDbCxfO9+DI/DuGeAMr9qI+noVWwGPNT github.com/alicebob/miniredis/v2 v2.30.0/go.mod h1:84TWKZlxYkfgMucPBf5SOQBYJceZeQRFIaQgNMiCX6Q= github.com/aliyun/aliyun-oss-go-sdk v3.0.2+incompatible h1:8psS8a+wKfiLt1iVDX79F7Y6wUM49Lcha2FMXt4UM8g= github.com/aliyun/aliyun-oss-go-sdk v3.0.2+incompatible/go.mod h1:T/Aws4fEfogEE9v+HPhhw+CntffsBHJ8nXQCwKr0/g8= +github.com/apache/arrow/go/v17 v17.0.0 h1:RRR2bdqKcdbss9Gxy2NS/hK8i4LDMh23L6BbkN5+F54= +github.com/apache/arrow/go/v17 v17.0.0/go.mod h1:jR7QHkODl15PfYyjM2nU+yTLScZ/qfj7OSUZmJ8putc= github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= @@ -115,6 +117,8 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/gomarkdown/markdown v0.0.0-20240328165702-4d01890c35c0 h1:4gjrh/PN2MuWCCElk8/I4OCKRKWCCo2zEct3VKCbibU= github.com/gomarkdown/markdown v0.0.0-20240328165702-4d01890c35c0/go.mod h1:JDGcbDT52eL4fju3sZ4TeHGsQwhG9nbDV21aMyhwPoA= +github.com/google/flatbuffers v24.3.25+incompatible h1:CX395cjN9Kke9mmalRoL3d81AtFUxJM+yDthflgJGkI= +github.com/google/flatbuffers v24.3.25+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= @@ -156,6 +160,8 @@ github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHm github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/compress v1.18.4 h1:RPhnKRAQ4Fh8zU2FY/6ZFDwTVTxgJ/EMydqSTzE9a2c= +github.com/klauspost/compress v1.18.4/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= @@ -166,13 +172,14 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/lancedb/lancedb-go v0.1.2 h1:ucM+KNN5J886OilSh4MRdyBa1sinHyrisoaswNISNFk= +github.com/lancedb/lancedb-go v0.1.2/go.mod h1:HzleylKfuw2HgfBBfrE3tb4LMKNdJ3/TQ1Ziyd+CLZk= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= -github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= @@ -195,6 +202,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uYLpLIr5fm8diHn0JbqRycJi6w0Ms= +github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= +github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= @@ -221,6 +230,8 @@ github.com/panjf2000/ants/v2 v2.12.0 h1:u9JhESo83i/GkZnhfTNuFMMWcNt7mnV1bGJ6FT4w github.com/panjf2000/ants/v2 v2.12.0/go.mod h1:tSQuaNQ6r6NRhPt+IZVUevvDyFMTs+eS4ztZc52uJTY= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ= +github.com/pierrec/lz4/v4 v4.1.21/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -232,7 +243,6 @@ github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw= github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= -github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= 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/richardlehane/mscfb v1.0.6 h1:eN3bvvZCp00bs7Zf52bxNwAx5lJDBK1tCuH19qq5aC8= @@ -314,6 +324,10 @@ github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yuin/gopher-lua v0.0.0-20220504180219-658193537a64 h1:5mLPGnFdSsevFRFc9q3yYbBkB6tsm4aCwwQV/j1JQAQ= github.com/yuin/gopher-lua v0.0.0-20220504180219-658193537a64/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0= +github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA= go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE= go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= @@ -379,11 +393,12 @@ golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 h1:bTLqdHv7xrGlFbvf5/TXNxy/iUwwdkjhqQTJDjW7aj0= +golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4/go.mod h1:g5NllXBEermZrmR51cJDQxmJUHUOfRAaNyWBM+R+548= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.41.0 h1:QCgPso/Q3RTJx2Th4bDLqML4W6iJiaXFq2/ftQF13YU= golang.org/x/term v0.41.0/go.mod h1:3pfBgksrReYfZ5lvYM0kSO0LIkAl4Yl2bXOkKP7Ec2A= @@ -403,6 +418,8 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 h1:+cNy6SZtPcJQH3LJVLOSmiC7MMxXNOb3PU/VUEz+EhU= +golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= google.golang.org/genproto/googleapis/rpc v0.0.0-20260209200024-4cfbd4190f57 h1:mWPCjDEyshlQYzBpMNHaEof6UX1PmHcaUODUywQ0uac= @@ -442,11 +459,11 @@ gorm.io/driver/sqlite v1.5.7/go.mod h1:U+J8craQU6Fzkcvu8oLeAQmi50TkwPEhHDEjQZXDa gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= gorm.io/gorm v1.25.12 h1:I0u8i2hWQItBq1WfE0o2+WuL9+8L21K9e2HHSTE/0f8= gorm.io/gorm v1.25.12/go.mod h1:xh7N7RHfYlNc5EmcI/El95gXusucDrQnHXe0+CgWcLQ= -modernc.org/libc v1.22.5 h1:91BNch/e5B0uPbJFgqbxXuOnxBQjlS//icfQEGmvyjE= -modernc.org/libc v1.22.5/go.mod h1:jj+Z7dTNX8fBScMVNRAYZ/jF91K8fdT2hYMThc3YjBY= -modernc.org/mathutil v1.5.0 h1:rV0Ko/6SfM+8G+yKiyI830l3Wuz1zRutdslNoQ0kfiQ= -modernc.org/mathutil v1.5.0/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= -modernc.org/memory v1.5.0 h1:N+/8c5rE6EqugZwHii4IFsaJ7MUhoWX07J5tC/iI5Ds= -modernc.org/memory v1.5.0/go.mod h1:PkUhL0Mugw21sHPeskwZW4D6VscE/GQJOnIpCnW6pSU= -modernc.org/sqlite v1.23.1 h1:nrSBg4aRQQwq59JpvGEQ15tNxoO5pX/kUjcRNwSAGQM= -modernc.org/sqlite v1.23.1/go.mod h1:OrDj17Mggn6MhE+iPbBNf7RGKODDE9NFT0f3EwDzJqk= +modernc.org/libc v1.41.0 h1:g9YAc6BkKlgORsUWj+JwqoB1wU3o4DE3bM3yvA3k+Gk= +modernc.org/libc v1.41.0/go.mod h1:w0eszPsiXoOnoMJgrXjglgLuDy/bt5RR4y3QzUUeodY= +modernc.org/mathutil v1.6.0 h1:fRe9+AmYlaej+64JsEEhoWuAYBkOtQiMEU7n/XgfYi4= +modernc.org/mathutil v1.6.0/go.mod h1:Ui5Q9q1TR2gFm0AQRqQUaBWFLAhQpCwNcuhBOSedWPo= +modernc.org/memory v1.7.2 h1:Klh90S215mmH8c9gO98QxQFsY+W451E8AnzjoE2ee1E= +modernc.org/memory v1.7.2/go.mod h1:NO4NVCQy0N7ln+T9ngWqOQfi7ley4vpwvARR+Hjw95E= +modernc.org/sqlite v1.29.6 h1:0lOXGrycJPptfHDuohfYgNqoe4hu+gYuN/pKgY5XjS4= +modernc.org/sqlite v1.29.6/go.mod h1:S02dvcmm7TnTRvGhv8IGYyLnIt7AS2KPaB1F/71p75U= diff --git a/internal/ai/rag/vectordb/lancedb.go b/internal/ai/rag/vectordb/lancedb.go new file mode 100644 index 0000000..9293a52 --- /dev/null +++ b/internal/ai/rag/vectordb/lancedb.go @@ -0,0 +1,489 @@ +//go:build lancedb + +package vectordb + +import ( + "context" + "fmt" + "math" + "os" + "strconv" + "strings" + + "agent-desk/internal/pkg/config" + + "github.com/apache/arrow/go/v17/arrow" + "github.com/apache/arrow/go/v17/arrow/array" + "github.com/apache/arrow/go/v17/arrow/memory" + "github.com/lancedb/lancedb-go/pkg/contracts" + "github.com/lancedb/lancedb-go/pkg/lancedb" +) + +const lanceDBVectorColumn = "vector" + +type LanceDBProvider struct { + conn contracts.IConnection +} + +func NewLanceDBProvider(cfg *config.LanceDBVectorDBConfig) (Provider, error) { + if cfg == nil { + return nil, fmt.Errorf("lancedb config is nil") + } + path := strings.TrimSpace(cfg.Path) + if path == "" { + path = "data/lancedb" + } + if err := os.MkdirAll(path, 0o755); err != nil { + return nil, fmt.Errorf("failed to create lancedb directory %s: %w", path, err) + } + conn, err := lancedb.Connect(context.Background(), path, nil) + if err != nil { + return nil, err + } + return &LanceDBProvider{conn: conn}, nil +} + +func (p *LanceDBProvider) Close() error { + if p.conn == nil || p.conn.IsClosed() { + return nil + } + return p.conn.Close() +} + +func (p *LanceDBProvider) CreateCollection(ctx context.Context, name string, dimension int) error { + if dimension <= 0 { + return fmt.Errorf("invalid lancedb vector dimension: %d", dimension) + } + if err := p.ensureOpen(); err != nil { + return err + } + schema, err := newLanceDBSchema(dimension) + if err != nil { + return err + } + table, err := p.conn.CreateTable(ctx, name, schema) + if err != nil { + return fmt.Errorf("failed to create lancedb table %s: %w", name, err) + } + return table.Close() +} + +func (p *LanceDBProvider) DeleteCollection(ctx context.Context, name string) error { + if err := p.ensureOpen(); err != nil { + return err + } + if err := p.conn.DropTable(ctx, name); err != nil { + return fmt.Errorf("failed to delete lancedb table %s: %w", name, err) + } + return nil +} + +func (p *LanceDBProvider) GetCollection(ctx context.Context, name string) (*CollectionInfo, error) { + table, err := p.openTable(ctx, name) + if err != nil { + return nil, err + } + defer table.Close() + + schema, err := table.Schema(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get lancedb table schema %s: %w", name, err) + } + count, err := table.Count(ctx) + if err != nil { + return nil, fmt.Errorf("failed to count lancedb table %s: %w", name, err) + } + return &CollectionInfo{ + Name: name, + Dimension: lanceDBVectorDimension(schema), + PointCount: int(count), + Status: "ok", + }, nil +} + +func (p *LanceDBProvider) ListCollections(ctx context.Context) ([]string, error) { + if err := p.ensureOpen(); err != nil { + return nil, err + } + names, err := p.conn.TableNames(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list lancedb tables: %w", err) + } + return names, nil +} + +func (p *LanceDBProvider) UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error { + if len(vectors) == 0 { + return nil + } + table, err := p.openTable(ctx, collectionName) + if err != nil { + return err + } + defer table.Close() + + ids := make([]string, 0, len(vectors)) + for _, vector := range vectors { + if strings.TrimSpace(vector.ID) != "" { + ids = append(ids, vector.ID) + } + } + if len(ids) > 0 { + if err := table.Delete(ctx, lanceDBStringInFilter("id", ids)); err != nil { + return fmt.Errorf("failed to delete existing lancedb vectors from %s: %w", collectionName, err) + } + } + + record, release, err := newLanceDBVectorRecord(vectors) + if err != nil { + return err + } + defer release() + + if err := table.AddRecords(ctx, []arrow.Record{record}, nil); err != nil { + return fmt.Errorf("failed to add lancedb vectors to %s: %w", collectionName, err) + } + return nil +} + +func (p *LanceDBProvider) DeleteVectors(ctx context.Context, collectionName string, ids []string) error { + if len(ids) == 0 { + return nil + } + table, err := p.openTable(ctx, collectionName) + if err != nil { + return err + } + defer table.Close() + + if err := table.Delete(ctx, lanceDBStringInFilter("id", ids)); err != nil { + return fmt.Errorf("failed to delete lancedb vectors from %s: %w", collectionName, err) + } + return nil +} + +func (p *LanceDBProvider) Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) { + table, err := p.openTable(ctx, req.CollectionName) + if err != nil { + return nil, err + } + defer table.Close() + + filter := lanceDBSearchFilter(req.Filter) + var rows []map[string]interface{} + if filter == "" { + rows, err = table.VectorSearch(ctx, lanceDBVectorColumn, req.Vector, req.TopK) + } else { + rows, err = table.VectorSearchWithFilter(ctx, lanceDBVectorColumn, req.Vector, req.TopK, filter) + } + if err != nil { + return nil, fmt.Errorf("failed to search lancedb table %s: %w", req.CollectionName, err) + } + + results := make([]SearchResult, 0, len(rows)) + for _, row := range rows { + score := lanceDBScoreFromRow(row) + if req.ScoreThreshold > 0 && score < req.ScoreThreshold { + continue + } + results = append(results, SearchResult{ + ID: valueToString(row["id"]), + Score: score, + Payload: lanceDBPayloadFromRow(row), + }) + } + return results, nil +} + +func (p *LanceDBProvider) ensureOpen() error { + if p == nil || p.conn == nil || p.conn.IsClosed() { + return fmt.Errorf("lancedb provider is closed") + } + return nil +} + +func (p *LanceDBProvider) openTable(ctx context.Context, name string) (contracts.ITable, error) { + if err := p.ensureOpen(); err != nil { + return nil, err + } + table, err := p.conn.OpenTable(ctx, name) + if err != nil { + return nil, fmt.Errorf("failed to open lancedb table %s: %w", name, err) + } + return table, nil +} + +func newLanceDBSchema(dimension int) (contracts.ISchema, error) { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: lanceDBVectorColumn, Type: arrow.FixedSizeListOf(int32(dimension), arrow.PrimitiveTypes.Float32), Nullable: false}, + {Name: "knowledge_base_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "document_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "document_title", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "faq_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "faq_question", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "chunk_no", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "chunk_type", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "section_path", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "title", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "content", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "provider", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + return lancedb.NewSchema(schema) +} + +func newLanceDBVectorRecord(vectors []Vector) (arrow.Record, func(), error) { + dimension := 0 + for _, item := range vectors { + if len(item.Vector) > 0 { + dimension = len(item.Vector) + break + } + } + if dimension <= 0 { + return nil, nil, fmt.Errorf("lancedb vector dimension is empty") + } + for _, item := range vectors { + if len(item.Vector) != dimension { + return nil, nil, fmt.Errorf("inconsistent lancedb vector dimension for %s: got %d, want %d", item.ID, len(item.Vector), dimension) + } + } + + pool := memory.NewGoAllocator() + idBuilder := array.NewStringBuilder(pool) + kbIDBuilder := array.NewInt64Builder(pool) + documentIDBuilder := array.NewInt64Builder(pool) + documentTitleBuilder := array.NewStringBuilder(pool) + faqIDBuilder := array.NewInt64Builder(pool) + faqQuestionBuilder := array.NewStringBuilder(pool) + chunkNoBuilder := array.NewInt32Builder(pool) + chunkTypeBuilder := array.NewStringBuilder(pool) + sectionPathBuilder := array.NewStringBuilder(pool) + titleBuilder := array.NewStringBuilder(pool) + contentBuilder := array.NewStringBuilder(pool) + providerBuilder := array.NewStringBuilder(pool) + vectorBuilder := array.NewFloat32Builder(pool) + + for _, item := range vectors { + payload := item.Payload + idBuilder.Append(item.ID) + vectorBuilder.AppendValues(item.Vector, nil) + kbIDBuilder.Append(payload.KnowledgeBaseID) + documentIDBuilder.Append(payload.DocumentID) + documentTitleBuilder.Append(payload.DocumentTitle) + faqIDBuilder.Append(payload.FaqID) + faqQuestionBuilder.Append(payload.FaqQuestion) + chunkNoBuilder.Append(int32(payload.ChunkNo)) + chunkTypeBuilder.Append(payload.ChunkType) + sectionPathBuilder.Append(payload.SectionPath) + titleBuilder.Append(payload.Title) + contentBuilder.Append(payload.Content) + providerBuilder.Append(payload.Provider) + } + + idArray := idBuilder.NewArray() + vectorValues := vectorBuilder.NewArray() + kbIDArray := kbIDBuilder.NewArray() + documentIDArray := documentIDBuilder.NewArray() + documentTitleArray := documentTitleBuilder.NewArray() + faqIDArray := faqIDBuilder.NewArray() + faqQuestionArray := faqQuestionBuilder.NewArray() + chunkNoArray := chunkNoBuilder.NewArray() + chunkTypeArray := chunkTypeBuilder.NewArray() + sectionPathArray := sectionPathBuilder.NewArray() + titleArray := titleBuilder.NewArray() + contentArray := contentBuilder.NewArray() + providerArray := providerBuilder.NewArray() + + vectorType := arrow.FixedSizeListOf(int32(dimension), arrow.PrimitiveTypes.Float32) + vectorArray := array.NewFixedSizeListData( + array.NewData(vectorType, len(vectors), []*memory.Buffer{nil}, []arrow.ArrayData{vectorValues.Data()}, 0, 0), + ) + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: lanceDBVectorColumn, Type: vectorType, Nullable: false}, + {Name: "knowledge_base_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "document_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "document_title", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "faq_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "faq_question", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "chunk_no", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "chunk_type", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "section_path", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "title", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "content", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "provider", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + columns := []arrow.Array{ + idArray, + vectorArray, + kbIDArray, + documentIDArray, + documentTitleArray, + faqIDArray, + faqQuestionArray, + chunkNoArray, + chunkTypeArray, + sectionPathArray, + titleArray, + contentArray, + providerArray, + } + record := array.NewRecord(schema, columns, int64(len(vectors))) + release := func() { + record.Release() + for _, column := range columns { + column.Release() + } + vectorValues.Release() + } + return record, release, nil +} + +func lanceDBVectorDimension(schema *arrow.Schema) int { + if schema == nil { + return 0 + } + for i := 0; i < schema.NumFields(); i++ { + field := schema.Field(i) + if field.Name != lanceDBVectorColumn { + continue + } + listType, ok := field.Type.(*arrow.FixedSizeListType) + if !ok { + return 0 + } + return int(listType.Len()) + } + return 0 +} + +func lanceDBSearchFilter(filter *SearchFilter) string { + if filter == nil { + return "" + } + parts := make([]string, 0, 2) + if len(filter.KnowledgeBaseIDs) > 0 { + parts = append(parts, lanceDBIntInFilter("knowledge_base_id", filter.KnowledgeBaseIDs)) + } + if len(filter.DocumentIDs) > 0 { + parts = append(parts, lanceDBIntInFilter("document_id", filter.DocumentIDs)) + } + return strings.Join(parts, " AND ") +} + +func lanceDBIntInFilter(column string, values []int64) string { + items := make([]string, 0, len(values)) + for _, value := range values { + items = append(items, strconv.FormatInt(value, 10)) + } + return fmt.Sprintf("%s IN (%s)", column, strings.Join(items, ",")) +} + +func lanceDBStringInFilter(column string, values []string) string { + items := make([]string, 0, len(values)) + for _, value := range values { + items = append(items, "'"+strings.ReplaceAll(value, "'", "''")+"'") + } + return fmt.Sprintf("%s IN (%s)", column, strings.Join(items, ",")) +} + +func lanceDBScoreFromRow(row map[string]interface{}) float32 { + for _, key := range []string{"_distance", "distance"} { + if value, ok := row[key]; ok { + distance := valueToFloat64(value) + if math.IsNaN(distance) { + break + } + score := 1 - distance + if score < 0 { + return 0 + } + if score > 1 { + return 1 + } + return float32(score) + } + } + for _, key := range []string{"_score", "score"} { + if value, ok := row[key]; ok { + score := valueToFloat64(value) + if !math.IsNaN(score) { + return float32(score) + } + } + } + return 0 +} + +func lanceDBPayloadFromRow(row map[string]interface{}) ChunkPayload { + return ChunkPayload{ + KnowledgeBaseID: valueToInt64(row["knowledge_base_id"]), + DocumentID: valueToInt64(row["document_id"]), + DocumentTitle: valueToString(row["document_title"]), + FaqID: valueToInt64(row["faq_id"]), + FaqQuestion: valueToString(row["faq_question"]), + ChunkNo: int(valueToInt64(row["chunk_no"])), + ChunkType: valueToString(row["chunk_type"]), + SectionPath: valueToString(row["section_path"]), + Title: valueToString(row["title"]), + Content: valueToString(row["content"]), + Provider: valueToString(row["provider"]), + } +} + +func valueToString(value interface{}) string { + switch v := value.(type) { + case nil: + return "" + case string: + return v + case []byte: + return string(v) + default: + return fmt.Sprint(value) + } +} + +func valueToInt64(value interface{}) int64 { + switch v := value.(type) { + case int: + return int64(v) + case int32: + return int64(v) + case int64: + return v + case uint64: + return int64(v) + case float32: + return int64(v) + case float64: + return int64(v) + case string: + ret, _ := strconv.ParseInt(v, 10, 64) + return ret + default: + return 0 + } +} + +func valueToFloat64(value interface{}) float64 { + switch v := value.(type) { + case float32: + return float64(v) + case float64: + return v + case int: + return float64(v) + case int32: + return float64(v) + case int64: + return float64(v) + case string: + ret, err := strconv.ParseFloat(v, 64) + if err == nil { + return ret + } + } + return math.NaN() +} diff --git a/internal/ai/rag/vectordb/lancedb_stub.go b/internal/ai/rag/vectordb/lancedb_stub.go new file mode 100644 index 0000000..33224d4 --- /dev/null +++ b/internal/ai/rag/vectordb/lancedb_stub.go @@ -0,0 +1,13 @@ +//go:build !lancedb + +package vectordb + +import ( + "fmt" + + "agent-desk/internal/pkg/config" +) + +func NewLanceDBProvider(_ *config.LanceDBVectorDBConfig) (Provider, error) { + return nil, fmt.Errorf("LanceDB provider is not built. Rebuild with -tags lancedb and configure LanceDB native libraries") +} diff --git a/internal/ai/rag/vectordb/provider.go b/internal/ai/rag/vectordb/provider.go index 6c6f835..d60fcb3 100644 --- a/internal/ai/rag/vectordb/provider.go +++ b/internal/ai/rag/vectordb/provider.go @@ -18,7 +18,9 @@ func Init(cfg *config.VectorDBConfig) error { var err error switch enums.VectorDBType(cfg.Type) { case enums.VectorDBTypeQdrant: - defaultProvider, err = NewQdrantProvider(cfg) + defaultProvider, err = NewQdrantProvider(&cfg.Qdrant) + case enums.VectorDBTypeLanceDB: + defaultProvider, err = NewLanceDBProvider(&cfg.LanceDB) default: return fmt.Errorf("unsupported vectordb type: %s", cfg.Type) } diff --git a/internal/ai/rag/vectordb/provider_test.go b/internal/ai/rag/vectordb/provider_test.go new file mode 100644 index 0000000..66a223e --- /dev/null +++ b/internal/ai/rag/vectordb/provider_test.go @@ -0,0 +1,23 @@ +package vectordb + +import ( + "strings" + "testing" + + "agent-desk/internal/pkg/config" +) + +func TestInitLanceDBWithoutBuildTagReturnsActionableError(t *testing.T) { + err := Init(&config.VectorDBConfig{ + Type: "lancedb", + LanceDB: config.LanceDBVectorDBConfig{ + Path: "data/lancedb", + }, + }) + if err == nil { + t.Fatal("Init(lancedb) error = nil, want actionable build tag error") + } + if !strings.Contains(err.Error(), "LanceDB provider is not built") { + t.Fatalf("Init(lancedb) error = %q, want build tag guidance", err.Error()) + } +} diff --git a/internal/ai/rag/vectordb/qdrant.go b/internal/ai/rag/vectordb/qdrant.go index 3f7850c..5c005c7 100644 --- a/internal/ai/rag/vectordb/qdrant.go +++ b/internal/ai/rag/vectordb/qdrant.go @@ -9,56 +9,11 @@ import ( "agent-desk/internal/pkg/config" ) -type Vector struct { - ID string `json:"id"` - Vector []float32 `json:"vector"` - Payload ChunkPayload `json:"payload"` -} - -type SearchRequest struct { - CollectionName string `json:"collectionName"` - Vector []float32 `json:"vector"` - TopK int `json:"topK"` - ScoreThreshold float32 `json:"scoreThreshold"` - Filter *SearchFilter `json:"filter,omitempty"` -} - -type SearchFilter struct { - KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds,omitempty"` - DocumentIDs []int64 `json:"documentIds,omitempty"` -} - -type SearchResult struct { - ID string `json:"id"` - Score float32 `json:"score"` - Payload ChunkPayload `json:"payload"` -} - -type CollectionInfo struct { - Name string `json:"name"` - Dimension int `json:"dimension"` - PointCount int `json:"pointCount"` - Status string `json:"status"` -} - -type Provider interface { - CreateCollection(ctx context.Context, name string, dimension int) error - DeleteCollection(ctx context.Context, name string) error - GetCollection(ctx context.Context, name string) (*CollectionInfo, error) - ListCollections(ctx context.Context) ([]string, error) - - UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error - DeleteVectors(ctx context.Context, collectionName string, ids []string) error - - Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) - Close() error -} - type QdrantProvider struct { client *qdrant.Client } -func NewQdrantProvider(cfg *config.VectorDBConfig) (*QdrantProvider, error) { +func NewQdrantProvider(cfg *config.QdrantVectorDBConfig) (*QdrantProvider, error) { if cfg == nil { return nil, fmt.Errorf("vectordb config is nil") } diff --git a/internal/ai/rag/vectordb/types.go b/internal/ai/rag/vectordb/types.go new file mode 100644 index 0000000..82e57b4 --- /dev/null +++ b/internal/ai/rag/vectordb/types.go @@ -0,0 +1,48 @@ +package vectordb + +import "context" + +type Vector struct { + ID string `json:"id"` + Vector []float32 `json:"vector"` + Payload ChunkPayload `json:"payload"` +} + +type SearchRequest struct { + CollectionName string `json:"collectionName"` + Vector []float32 `json:"vector"` + TopK int `json:"topK"` + ScoreThreshold float32 `json:"scoreThreshold"` + Filter *SearchFilter `json:"filter,omitempty"` +} + +type SearchFilter struct { + KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds,omitempty"` + DocumentIDs []int64 `json:"documentIds,omitempty"` +} + +type SearchResult struct { + ID string `json:"id"` + Score float32 `json:"score"` + Payload ChunkPayload `json:"payload"` +} + +type CollectionInfo struct { + Name string `json:"name"` + Dimension int `json:"dimension"` + PointCount int `json:"pointCount"` + Status string `json:"status"` +} + +type Provider interface { + CreateCollection(ctx context.Context, name string, dimension int) error + DeleteCollection(ctx context.Context, name string) error + GetCollection(ctx context.Context, name string) (*CollectionInfo, error) + ListCollections(ctx context.Context) ([]string, error) + + UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error + DeleteVectors(ctx context.Context, collectionName string, ids []string) error + + Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) + Close() error +} diff --git a/internal/pkg/config/config.go b/internal/pkg/config/config.go index 0602467..e2d37bd 100644 --- a/internal/pkg/config/config.go +++ b/internal/pkg/config/config.go @@ -123,14 +123,22 @@ type OSSStorageConfig struct { } type VectorDBConfig struct { - Type string `yaml:"type"` + Type string `yaml:"type"` + Qdrant QdrantVectorDBConfig `yaml:"qdrant"` + LanceDB LanceDBVectorDBConfig `yaml:"lancedb"` +} + +type QdrantVectorDBConfig struct { Host string `yaml:"host"` - Port int `yaml:"port"` - APIKey string `yaml:"apiKey"` GrpcPort int `yaml:"grpcPort"` + APIKey string `yaml:"apiKey"` UseTLS bool `yaml:"useTls"` } +type LanceDBVectorDBConfig struct { + Path string `yaml:"path"` +} + type MCPConfig struct { Enabled bool `yaml:"enabled"` Servers map[string]MCPServerConfig `yaml:"servers"` diff --git a/internal/pkg/config/vector_db_config_test.go b/internal/pkg/config/vector_db_config_test.go new file mode 100644 index 0000000..30acf93 --- /dev/null +++ b/internal/pkg/config/vector_db_config_test.go @@ -0,0 +1,45 @@ +package config + +import ( + "testing" + + "gopkg.in/yaml.v3" +) + +func TestVectorDBConfigUnmarshalNestedProviders(t *testing.T) { + raw := []byte(` +vectorDB: + type: lancedb + qdrant: + host: 127.0.0.1 + grpcPort: 6334 + apiKey: secret + useTls: true + lancedb: + path: data/lancedb +`) + + var cfg Config + if err := yaml.Unmarshal(raw, &cfg); err != nil { + t.Fatalf("yaml.Unmarshal() error = %v", err) + } + + if cfg.VectorDB.Type != "lancedb" { + t.Fatalf("VectorDB.Type = %q, want %q", cfg.VectorDB.Type, "lancedb") + } + if cfg.VectorDB.Qdrant.Host != "127.0.0.1" { + t.Fatalf("VectorDB.Qdrant.Host = %q, want %q", cfg.VectorDB.Qdrant.Host, "127.0.0.1") + } + if cfg.VectorDB.Qdrant.GrpcPort != 6334 { + t.Fatalf("VectorDB.Qdrant.GrpcPort = %d, want %d", cfg.VectorDB.Qdrant.GrpcPort, 6334) + } + if cfg.VectorDB.Qdrant.APIKey != "secret" { + t.Fatalf("VectorDB.Qdrant.APIKey = %q, want %q", cfg.VectorDB.Qdrant.APIKey, "secret") + } + if !cfg.VectorDB.Qdrant.UseTLS { + t.Fatal("VectorDB.Qdrant.UseTLS = false, want true") + } + if cfg.VectorDB.LanceDB.Path != "data/lancedb" { + t.Fatalf("VectorDB.LanceDB.Path = %q, want %q", cfg.VectorDB.LanceDB.Path, "data/lancedb") + } +} diff --git a/internal/pkg/enums/knowledge.go b/internal/pkg/enums/knowledge.go index 0ab43c8..837ea25 100644 --- a/internal/pkg/enums/knowledge.go +++ b/internal/pkg/enums/knowledge.go @@ -3,11 +3,13 @@ package enums type VectorDBType string const ( - VectorDBTypeQdrant VectorDBType = "qdrant" + VectorDBTypeQdrant VectorDBType = "qdrant" + VectorDBTypeLanceDB VectorDBType = "lancedb" ) var vectorDBTypeLabelMap = map[VectorDBType]string{ - VectorDBTypeQdrant: "Qdrant", + VectorDBTypeQdrant: "Qdrant", + VectorDBTypeLanceDB: "LanceDB", } func GetVectorDBTypeLabel(dbType VectorDBType) string { diff --git a/internal/pkg/enums/knowledge_test.go b/internal/pkg/enums/knowledge_test.go new file mode 100644 index 0000000..b47ef7c --- /dev/null +++ b/internal/pkg/enums/knowledge_test.go @@ -0,0 +1,12 @@ +package enums + +import "testing" + +func TestVectorDBTypeLabelIncludesLanceDB(t *testing.T) { + if VectorDBTypeLanceDB != "lancedb" { + t.Fatalf("VectorDBTypeLanceDB = %q, want %q", VectorDBTypeLanceDB, "lancedb") + } + if got := GetVectorDBTypeLabel(VectorDBTypeLanceDB); got != "LanceDB" { + t.Fatalf("GetVectorDBTypeLabel(VectorDBTypeLanceDB) = %q, want %q", got, "LanceDB") + } +} diff --git a/web/lib/generated/enums.ts b/web/lib/generated/enums.ts index a8ba935..0790ce4 100644 --- a/web/lib/generated/enums.ts +++ b/web/lib/generated/enums.ts @@ -347,10 +347,12 @@ export const StatusLabels: Record = { export enum ThirdProvider { WxWork = "wxwork", Dingtalk = "dingtalk", + OIDC = "oidc", } export const ThirdProviderLabels: Record = { [ThirdProvider.WxWork]: "企业微信", [ThirdProvider.Dingtalk]: "钉钉", + [ThirdProvider.OIDC]: "OIDC", } export enum TicketStatus { @@ -366,7 +368,9 @@ export const TicketStatusLabels: Record = { export enum VectorDBType { Qdrant = "qdrant", + LanceDB = "lancedb", } export const VectorDBTypeLabels: Record = { [VectorDBType.Qdrant]: "Qdrant", + [VectorDBType.LanceDB]: "LanceDB", }