Add Couchbase VectorDB Support (#3525)

* add couchbase vector store * add unit test * cleanup * add couchbase to website * fix code formatting --------- Co-authored-by: Jack Gerrits <[email protected]>
microsoft · Sep 25, 2024 · 2e1f788 · 2e1f788
1 parent 5eb1b0b
commit 2e1f788
Show file tree

Hide file tree

Showing 6 changed files with 569 additions and 3 deletions.
diff --git a/.github/workflows/contrib-tests.yml b/.github/workflows/contrib-tests.yml
@@ -91,6 +91,17 @@ jobs:
         image: mongodb/mongodb-atlas-local:latest
         ports:
           - 27017:27017
+      couchbase:
+        image: couchbase:enterprise-7.6.3
+        ports:
+          - "8091-8095:8091-8095"
+          - "11210:11210"
+          - "9102:9102"
+        healthcheck: # checks couchbase server is up
+          test: ["CMD", "curl", "-v", "http://localhost:8091/pools"]
+          interval: 20s
+          timeout: 20s
+          retries: 5
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
@@ -111,6 +122,9 @@ jobs:
       - name: Install mongodb when on linux
         run: |
           pip install -e .[retrievechat-mongodb]
+      - name: Install couchbase when on linux
+        run: |
+          pip install -e .[retrievechat-couchbase]
       - name: Install unstructured when python-version is 3.9 and on linux
         if: matrix.python-version == '3.9'
         run: |

diff --git a/autogen/agentchat/contrib/vectordb/base.py b/autogen/agentchat/contrib/vectordb/base.py
@@ -201,7 +201,7 @@ class VectorDBFactory:
     Factory class for creating vector databases.
     """
 
-    PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant"]
+    PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant", "couchbase"]
 
     @staticmethod
     def create_vector_db(db_type: str, **kwargs) -> VectorDB:
@@ -231,6 +231,10 @@ def create_vector_db(db_type: str, **kwargs) -> VectorDB:
             from .qdrant import QdrantVectorDB
 
             return QdrantVectorDB(**kwargs)
+        if db_type.lower() in ["couchbase", "couchbasedb", "capella"]:
+            from .couchbase import CouchbaseVectorDB
+
+            return CouchbaseVectorDB(**kwargs)
         else:
             raise ValueError(
                 f"Unsupported vector database type: {db_type}. Valid types are {VectorDBFactory.PREDEFINED_VECTOR_DB}."