Cog 170 pgvector adapter

.env.template

@@ -7,24 +7,26 @@ GRAPHISTRY_PASSWORD=
 
 SENTRY_REPORTING_URL=
 
-GRAPH_DATABASE_PROVIDER="neo4j" # or "networkx"
+# "neo4j" or "networkx"
+GRAPH_DATABASE_PROVIDER="neo4j"
 # Not needed if using networkx
 GRAPH_DATABASE_URL=
 GRAPH_DATABASE_USERNAME=
 GRAPH_DATABASE_PASSWORD=
 
-VECTOR_DB_PROVIDER="qdrant" # or "weaviate" or "lancedb"
-# Not needed if using "lancedb"
+# "qdrant", "pgvector", "weaviate" or "lancedb"
+VECTOR_DB_PROVIDER="qdrant"
+# Not needed if using "lancedb" or "pgvector"
 VECTOR_DB_URL=
 VECTOR_DB_KEY=
 
-# Database provider
-DB_PROVIDER="sqlite" # or "postgres"
+# Relational Database provider "sqlite" or "postgres"
+DB_PROVIDER="sqlite"
 
 # Database name
 DB_NAME=cognee_db
 
-# Postgres specific parameters (Only if Postgres is run)
+# Postgres specific parameters (Only if Postgres or PGVector is used)
 DB_HOST=127.0.0.1
 DB_PORT=5432
 DB_USERNAME=cognee

.github/workflows/test_pgvector.yml

@@ -0,0 +1,67 @@
+name: test | pgvector
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  RUNTIME__LOG_LEVEL: ERROR
+
+jobs:
+  get_docs_changes:
+    name: docs changes
+    uses: ./.github/workflows/get_docs_changes.yml
+
+  run_pgvector_integration_test:
+    name: test
+    needs: get_docs_changes
+    if: needs.get_docs_changes.outputs.changes_outside_docs == 'true'
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+    services:
+      postgres:
+        image: pgvector/pgvector:pg17
+        env:
+          POSTGRES_USER: cognee
+          POSTGRES_PASSWORD: cognee
+          POSTGRES_DB: cognee_db
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          - 5432:5432
+
+    steps:
+      - name: Check out
+        uses: actions/checkout@master
+
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11.x'
+
+      - name: Install Poetry
+        uses: snok/[email protected]
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      - name: Install dependencies
+        run: poetry install -E postgres --no-interaction
+
+      - name: Run default PGVector
+        env:
+          ENV: 'dev'
+          LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+        run: poetry run python ./cognee/tests/test_pgvector.py

README.md

@@ -190,11 +190,11 @@ Cognee supports a variety of tools and services for different operations:
 
 - **Local Setup**: By default, LanceDB runs locally with NetworkX and OpenAI.
 
-- **Vector Stores**: Cognee supports Qdrant and Weaviate for vector storage.
+- **Vector Stores**: Cognee supports LanceDB, Qdrant, PGVector and Weaviate for vector storage.
 
 - **Language Models (LLMs)**: You can use either Anyscale or Ollama as your LLM provider.
 
-- **Graph Stores**: In addition to LanceDB, Neo4j is also supported for graph storage.
+- **Graph Stores**: In addition to NetworkX, Neo4j is also supported for graph storage.
 
 - **User management**: Create individual user graphs and manage permissions
 

cognee/api/client.py

@@ -374,7 +374,7 @@ class LLMConfigDTO(InDTO):
     api_key: str
 
 class VectorDBConfigDTO(InDTO):
-    provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"]]
+    provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
     url: str
     api_key: str
 

cognee/api/v1/add/add.py

@@ -8,15 +8,18 @@
 from cognee.modules.ingestion import get_matched_datasets, save_data_to_file
 from cognee.shared.utils import send_telemetry
 from cognee.base_config import get_base_config
-from cognee.infrastructure.databases.relational import get_relational_engine, create_db_and_tables
+from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.users.methods import get_default_user
 from cognee.tasks.ingestion import get_dlt_destination
 from cognee.modules.users.permissions.methods import give_permission_on_document
 from cognee.modules.users.models import User
 from cognee.modules.data.methods import create_dataset
+from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables
+from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables
 
 async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_name: str = "main_dataset", user: User = None):
-    await create_db_and_tables()
+    await create_relational_db_and_tables()
+    await create_pgvector_db_and_tables()
 
     if isinstance(data, str):
         if "data://" in data:

cognee/api/v1/add/add_v2.py

@@ -3,10 +3,12 @@
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.pipelines import run_tasks, Task
 from cognee.tasks.ingestion import save_data_to_storage, ingest_data
-from cognee.infrastructure.databases.relational import create_db_and_tables
+from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables
+from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables
 
 async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_name: str = "main_dataset", user: User = None):
-    await create_db_and_tables()
+    await create_relational_db_and_tables()
+    await create_pgvector_db_and_tables()
 
     if user is None:
         user = await get_default_user()

cognee/api/v1/config/config.py

@@ -95,6 +95,30 @@ def set_vector_db_provider(vector_db_provider: str):
         vector_db_config = get_vectordb_config()
         vector_db_config.vector_db_provider = vector_db_provider
 
+    @staticmethod
+    def set_relational_db_config(config_dict: dict):
+        """
+            Updates the relational db config with values from config_dict.
+        """
+        relational_db_config = get_relational_config()
+        for key, value in config_dict.items():
+            if hasattr(relational_db_config, key):
+                object.__setattr__(relational_db_config, key, value)
+            else:
+                raise AttributeError(f"'{key}' is not a valid attribute of the config.")
+
+    @staticmethod
+    def set_vector_db_config(config_dict: dict):
+        """
+            Updates the vector db config with values from config_dict.
+        """
+        vector_db_config = get_vectordb_config()
+        for key, value in config_dict.items():
+            if hasattr(vector_db_config, key):
+                object.__setattr__(vector_db_config, key, value)
+            else:
+                raise AttributeError(f"'{key}' is not a valid attribute of the config.")
+
     @staticmethod
     def set_vector_db_key(db_key: str):
         vector_db_config = get_vectordb_config()

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -119,6 +119,8 @@ async def delete_database(self):
                 self.db_path = None
             else:
                 async with self.engine.begin() as connection:
+                    # Load the schema information into the MetaData object
+                    await connection.run_sync(Base.metadata.reflect)
                     for table in Base.metadata.sorted_tables:
                         drop_table_query = text(f"DROP TABLE IF EXISTS {table.name} CASCADE")
                         await connection.execute(drop_table_query)

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -1,5 +1,7 @@
 from typing import Dict
 
+from ..relational.config import get_relational_config
+
 class VectorConfig(Dict):
     vector_db_url: str
     vector_db_key: str
@@ -26,6 +28,25 @@ def create_vector_engine(config: VectorConfig, embedding_engine):
                 api_key = config["vector_db_key"],
                 embedding_engine = embedding_engine
             )
+    elif config["vector_db_provider"] == "pgvector":
+        from .pgvector.PGVectorAdapter import PGVectorAdapter
+
+        # Get configuration for postgres database
+        relational_config = get_relational_config()
+        db_username = relational_config.db_username
+        db_password = relational_config.db_password
+        db_host = relational_config.db_host
+        db_port = relational_config.db_port
+        db_name = relational_config.db_name
+
+        connection_string: str = (
+                            f"postgresql+asyncpg://{db_username}:{db_password}@{db_host}:{db_port}/{db_name}"
+        )
+
+        return PGVectorAdapter(connection_string, 
+                        config["vector_db_key"], 
+                        embedding_engine
+        )
     else:
         from .lancedb.LanceDBAdapter import LanceDBAdapter