feat: couchbase integration (langgenius#6165)

Co-authored-by: crazywoola <[email protected]> Co-authored-by: Elliot Scribner <[email protected]> Co-authored-by: crazywoola <[email protected]> Co-authored-by: Bowen Liang <[email protected]>
luckylhb90 · Oct 29, 2024 · c8ef922 · c8ef922
1 parent fc37e65
commit c8ef922
Show file tree

Hide file tree

Showing 21 changed files with 639 additions and 7 deletions.
diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml
@@ -78,14 +78,15 @@ jobs:
       - name: Run Workflow
         run: poetry run -C api bash dev/pytest/pytest_workflow.sh
 
-      - name: Set up Vector Stores (Weaviate, Qdrant, PGVector, Milvus, PgVecto-RS, Chroma, MyScale, ElasticSearch)
+      - name: Set up Vector Stores (Weaviate, Qdrant, PGVector, Milvus, PgVecto-RS, Chroma, MyScale, ElasticSearch, Couchbase)
         uses: hoverkraft-tech/[email protected]
         with:
           compose-file: |
             docker/docker-compose.yaml
           services: |
             weaviate
             qdrant
+            couchbase-server
             etcd
             minio
             milvus-standalone

diff --git a/.github/workflows/expose_service_ports.sh b/.github/workflows/expose_service_ports.sh
@@ -7,5 +7,7 @@ yq eval '.services["milvus-standalone"].ports += ["19530:19530"]' -i docker/dock
 yq eval '.services.pgvector.ports += ["5433:5432"]' -i docker/docker-compose.yaml
 yq eval '.services["pgvecto-rs"].ports += ["5431:5432"]' -i docker/docker-compose.yaml
 yq eval '.services["elasticsearch"].ports += ["9200:9200"]' -i docker/docker-compose.yaml
+yq eval '.services.couchbase-server.ports += ["8091-8096:8091-8096"]' -i docker/docker-compose.yaml
+yq eval '.services.couchbase-server.ports += ["11210:11210"]' -i docker/docker-compose.yaml
 
-echo "Ports exposed for sandbox, weaviate, qdrant, chroma, milvus, pgvector, pgvecto-rs, elasticsearch"
+echo "Ports exposed for sandbox, weaviate, qdrant, chroma, milvus, pgvector, pgvecto-rs, elasticsearch, couchbase"
diff --git a/.gitignore b/.gitignore
@@ -173,6 +173,7 @@ docker/volumes/myscale/log/*
 docker/volumes/unstructured/*
 docker/volumes/pgvector/data/*
 docker/volumes/pgvecto_rs/data/*
+docker/volumes/couchbase/*
 
 docker/nginx/conf.d/default.conf
 docker/nginx/ssl/*
@@ -189,4 +190,4 @@ pyrightconfig.json
 api/.vscode
 
 .idea/
-.vscode
+.vscode
diff --git a/api/.env.example b/api/.env.example
@@ -120,7 +120,7 @@ SUPABASE_URL=your-server-url
 WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 
-# Vector database configuration, support: weaviate, qdrant, milvus, myscale, relyt, pgvecto_rs, pgvector, pgvector, chroma, opensearch, tidb_vector, vikingdb, upstash
+# Vector database configuration, support: weaviate, qdrant, milvus, myscale, relyt, pgvecto_rs, pgvector, pgvector, chroma, opensearch, tidb_vector, couchbase, vikingdb, upstash
 VECTOR_STORE=weaviate
 
 # Weaviate configuration
@@ -136,6 +136,13 @@ QDRANT_CLIENT_TIMEOUT=20
 QDRANT_GRPC_ENABLED=false
 QDRANT_GRPC_PORT=6334
 
+#Couchbase configuration
+COUCHBASE_CONNECTION_STRING=127.0.0.1
+COUCHBASE_USER=Administrator
+COUCHBASE_PASSWORD=password
+COUCHBASE_BUCKET_NAME=Embeddings
+COUCHBASE_SCOPE_NAME=_default
+
 # Milvus configuration
 MILVUS_URI=http://127.0.0.1:19530
 MILVUS_TOKEN=

diff --git a/api/commands.py b/api/commands.py
@@ -278,6 +278,7 @@ def migrate_knowledge_vector_database():
         VectorType.BAIDU,
         VectorType.VIKINGDB,
         VectorType.UPSTASH,
+        VectorType.COUCHBASE,
     }
     page = 1
     while True:

diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py
@@ -17,6 +17,7 @@
 from configs.middleware.storage.volcengine_tos_storage_config import VolcengineTOSStorageConfig
 from configs.middleware.vdb.analyticdb_config import AnalyticdbConfig
 from configs.middleware.vdb.chroma_config import ChromaConfig
+from configs.middleware.vdb.couchbase_config import CouchbaseConfig
 from configs.middleware.vdb.elasticsearch_config import ElasticsearchConfig
 from configs.middleware.vdb.milvus_config import MilvusConfig
 from configs.middleware.vdb.myscale_config import MyScaleConfig
@@ -251,6 +252,7 @@ class MiddlewareConfig(
     TiDBVectorConfig,
     WeaviateConfig,
     ElasticsearchConfig,
+    CouchbaseConfig,
     InternalTestConfig,
     VikingDBConfig,
     UpstashConfig,

diff --git a/api/configs/middleware/vdb/couchbase_config.py b/api/configs/middleware/vdb/couchbase_config.py
@@ -0,0 +1,34 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class CouchbaseConfig(BaseModel):
+    """
+    Couchbase configs
+    """
+
+    COUCHBASE_CONNECTION_STRING: Optional[str] = Field(
+        description="COUCHBASE connection string",
+        default=None,
+    )
+
+    COUCHBASE_USER: Optional[str] = Field(
+        description="COUCHBASE user",
+        default=None,
+    )
+
+    COUCHBASE_PASSWORD: Optional[str] = Field(
+        description="COUCHBASE password",
+        default=None,
+    )
+
+    COUCHBASE_BUCKET_NAME: Optional[str] = Field(
+        description="COUCHBASE bucket name",
+        default=None,
+    )
+
+    COUCHBASE_SCOPE_NAME: Optional[str] = Field(
+        description="COUCHBASE scope name",
+        default=None,
+    )
diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py
@@ -640,6 +640,7 @@ def get(self):
                 | VectorType.ELASTICSEARCH
                 | VectorType.PGVECTOR
                 | VectorType.TIDB_ON_QDRANT
+                | VectorType.COUCHBASE
             ):
                 return {
                     "retrieval_method": [
@@ -678,6 +679,7 @@ def get(self, vector_type):
                 | VectorType.MYSCALE
                 | VectorType.ORACLE
                 | VectorType.ELASTICSEARCH
+                | VectorType.COUCHBASE
                 | VectorType.PGVECTOR
             ):
                 return {

diff --git a/api/core/rag/datasource/vdb/couchbase/__init__.py b/api/core/rag/datasource/vdb/couchbase/__init__.py