[Re-Opened] Support for PGVector Database in Autogen (#2439)

* PGVector Contrib Initial Commit - KnucklesTeam:autogen:pgvector_contrib fork

* Update website/docs/ecosystem/pgvector.md

Co-authored-by: Chi Wang <[email protected]>

* Updated qdrant installation instructions.

* Fixed openai version.

* Added dependencies to install for qdrant and pgvector in contrib tests.

* Added dependencies to install for qdrant and pgvector in contrib tests.

* Cleaned up dependencies.

* Removed flaml out of setup.py. Used only for notebook example.

* Added PGVector notebook link

---------

Co-authored-by: Chi Wang <[email protected]>

.github/workflows/contrib-openai.yml

@@ -41,19 +41,23 @@ jobs:
           pip install -e .
           python -c "import autogen"
           pip install coverage pytest-asyncio
+      - name: Install PostgreSQL
+        run: |
+          sudo apt install postgresql -y
+      - name: Start PostgreSQL service
+        run: sudo service postgresql start
       - name: Install packages for test when needed
         run: |
           pip install docker
-          pip install qdrant_client[fastembed]
-          pip install -e .[retrievechat]
+          pip install -e .[retrievechat-qdrant,retrievechat-pgvector]
       - name: Coverage
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY }}
           AZURE_OPENAI_API_BASE: ${{ secrets.AZURE_OPENAI_API_BASE }}
           OAI_CONFIG_LIST: ${{ secrets.OAI_CONFIG_LIST }}
         run: |
-          coverage run -a -m pytest test/agentchat/contrib/test_retrievechat.py::test_retrievechat test/agentchat/contrib/test_qdrant_retrievechat.py::test_retrievechat
+          coverage run -a -m pytest test/agentchat/contrib/test_retrievechat.py::test_retrievechat test/agentchat/contrib/test_qdrant_retrievechat.py::test_retrievechat test/agentchat/contrib/test_pgvector_retrievechat.py::test_retrievechat
           coverage xml
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3

.github/workflows/contrib-tests.yml

@@ -42,16 +42,20 @@ jobs:
       - name: Install qdrant_client when python-version is 3.10
         if: matrix.python-version == '3.10'
         run: |
-          pip install qdrant_client[fastembed]
+          pip install .[retrievechat-qdrant]
       - name: Install unstructured when python-version is 3.9 and on linux
-        if: matrix.python-version == '3.9' && matrix.os == 'ubuntu-latest'
         run: |
           sudo apt-get update
           sudo apt-get install -y tesseract-ocr poppler-utils
           pip install unstructured[all-docs]==0.13.0
-      - name: Install packages and dependencies for RetrieveChat
+      - name: Install and Start PostgreSQL
+        runs-on: ubuntu-latest
         run: |
-          pip install -e .[retrievechat]
+          sudo apt install postgresql -y
+          sudo service postgresql start
+      - name: Install packages and dependencies for PGVector
+        run: |
+          pip install -e .[retrievechat-pgvector]
       - name: Set AUTOGEN_USE_DOCKER based on OS
         shell: bash
         run: |

autogen/agentchat/contrib/vectordb/base.py

@@ -185,7 +185,7 @@ class VectorDBFactory:
     Factory class for creating vector databases.
     """
 
-    PREDEFINED_VECTOR_DB = ["chroma"]
+    PREDEFINED_VECTOR_DB = ["chroma", "pgvector"]
 
     @staticmethod
     def create_vector_db(db_type: str, **kwargs) -> VectorDB:
@@ -203,6 +203,10 @@ def create_vector_db(db_type: str, **kwargs) -> VectorDB:
             from .chromadb import ChromaVectorDB
 
             return ChromaVectorDB(**kwargs)
+        if db_type.lower() in ["pgvector", "pgvectordb"]:
+            from .pgvectordb import PGVectorDB
+
+            return PGVectorDB(**kwargs)
         else:
             raise ValueError(
                 f"Unsupported vector database type: {db_type}. Valid types are {VectorDBFactory.PREDEFINED_VECTOR_DB}."