Fix couchbase vectordb support

.github/workflows/contrib-test.yml

@@ -94,6 +94,12 @@ 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"
     steps:
       - uses: actions/checkout@v4
       - uses: astral-sh/setup-uv@v5
@@ -116,6 +122,9 @@ jobs:
       - name: Install pgvector when on linux
         run: |
           uv pip install --system -e .[test,retrievechat-pgvector,retrievechat-mongodb]
+      - name: Install couchbase when on linux
+        run: |
+          uv pip install --system -e .[retrievechat-couchbase]
       - name: Set AUTOGEN_USE_DOCKER based on OS
         shell: bash
         run: |
@@ -130,7 +139,7 @@ jobs:
         uses: codecov/codecov-action@v5
         with:
           files: ./coverage.xml
-          flags: retrievechat, retrievechat-qdrant, retrievechat-pgvector, retrievechat-mongodb, ubuntu-latest, ${{ matrix.python-version }}
+          flags: retrievechat, retrievechat-qdrant, retrievechat-pgvector, retrievechat-mongodb, retrievechat-couchbase, ubuntu-latest, ${{ matrix.python-version }}
           fail_ci_if_error: true
           token: ${{ secrets.CODECOV_TOKEN }}
 

.github/workflows/test-with-optional-deps.yml

@@ -65,6 +65,7 @@ jobs:
           - "commsagent-telegram"
           - "jupyter-executor"
           - "retrievechat"
+          - "retrievechat-couchbase"
           - "retrievechat-pgvector"
           - "retrievechat-mongodb"
           - "retrievechat-qdrant"

autogen/agentchat/contrib/llamaindex_conversable_agent.py

@@ -9,7 +9,7 @@
 from ... import OpenAIWrapper
 from ...import_utils import optional_import_block, require_optional_import
 from .. import Agent, ConversableAgent
-from .vectordb.utils import get_logger
+from .vectordb import get_logger
 
 logger = get_logger(__name__)
 

autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py

@@ -10,7 +10,7 @@
 from ...import_utils import optional_import_block, require_optional_import
 from ...retrieve_utils import TEXT_FORMATS, get_files_from_dir, split_files_to_chunks
 from ..contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
-from ..contrib.vectordb.utils import (
+from ..contrib.vectordb import (
     chroma_results_to_query_results,
     filter_results_by_distance,
     get_logger,

autogen/agentchat/contrib/retrieve_user_proxy_agent.py

@@ -23,8 +23,11 @@
 from ...token_count_utils import count_token
 from .. import UserProxyAgent
 from ..agent import Agent
-from ..contrib.vectordb.base import Document, QueryResults, VectorDB, VectorDBFactory
-from ..contrib.vectordb.utils import (
+from ..contrib.vectordb import (
+    Document,
+    QueryResults,
+    VectorDB,
+    VectorDBFactory,
     chroma_results_to_query_results,
     filter_results_by_distance,
     get_logger,

autogen/agentchat/contrib/vectordb/__init__.py

@@ -2,4 +2,19 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__all__: list[str] = []
+from .base import Document, QueryResults, VectorDB, VectorDBFactory
+from .utils import (
+    chroma_results_to_query_results,
+    filter_results_by_distance,
+    get_logger,
+)
+
+__all__ = [
+    "Document",
+    "QueryResults",
+    "VectorDB",
+    "VectorDBFactory",
+    "chroma_results_to_query_results",
+    "filter_results_by_distance",
+    "get_logger",
+]

autogen/agentchat/contrib/vectordb/base.py

@@ -15,11 +15,14 @@
     runtime_checkable,
 )
 
+from ....doc_utils import export_module
+
 Metadata = Union[Mapping[str, Any], None]
 Vector = Union[Sequence[float], Sequence[int]]
 ItemID = Union[str, int]  # chromadb doesn't support int ids, VikingDB does
 
 
+@export_module("agentchat.contrib.vectordb")
 class Document(TypedDict):
     """A Document is a record in the vector database.
 
@@ -42,6 +45,7 @@ class Document(TypedDict):
 QueryResults = list[list[tuple[Document, float]]]
 
 
+@export_module("agentchat.contrib.vectordb")
 @runtime_checkable
 class VectorDB(Protocol):
     """Abstract class for vector database. A vector database is responsible for storing and retrieving documents.
@@ -63,9 +67,7 @@ class VectorDB(Protocol):
 
     active_collection: Any = None
     type: str = ""
-    embedding_function: Optional[Callable[[list[str]], list[list[float]]]] = (
-        None  # embeddings = embedding_function(sentences)
-    )
+    embedding_function: Optional[Callable[[list[str]], list[list[float]]]] = None
 
     def create_collection(self, collection_name: str, overwrite: bool = False, get_or_create: bool = True) -> Any:
         """Create a collection in the vector database.
@@ -190,10 +192,11 @@ def get_docs_by_ids(
         ...
 
 
+@export_module("agentchat.contrib.vectordb")
 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:
@@ -222,6 +225,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}."