gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: PR comments fixes

Browse source
This commit is contained in:
Andrej Milicevic 2026-01-15 20:00:17 +01:00
parent 2cdbc02b35
commit c6ff0eca3d
4 changed files with 84 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
.github/workflows/vector_db_tests.yml vendored
View file

@ -103,6 +103,55 @@ jobs:
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
run: uv run python ./cognee/tests/test_pgvector.py
run-pgvector-multi-user-tests:
name: PGVector Multi-User Tests
runs-on: ubuntu-22.04
if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'postgres') }}
services:
postgres:
image: pgvector/pgvector:pg17
env:
POSTGRES_USER: cognee
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
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@v4
with:
fetch-depth: 0
- name: Cognee Setup
uses: ./.github/actions/cognee_setup
with:
python-version: ${{ inputs.python-version }}
extra-dependencies: "postgres"
- name: Run PGVector Permissions Tests
env:
ENV: 'dev'
ENABLE_BACKEND_ACCESS_CONTROL: 'true'
LLM_MODEL: ${{ secrets.LLM_MODEL }}
LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
VECTOR_DB_URL: 127.0.0.1
VECTOR_DB_PORT: 5432
VECTOR_DB_USERNAME: cognee
VECTOR_DB_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
VECTOR_DATASET_DATABASE_HANDLER: pgvector
run: uv run python ./cognee/tests/test_permissions.py
run-lancedb-tests:
name: LanceDB Tests
runs-on: ubuntu-22.04

2
cognee/infrastructure/databases/dataset_database_handler/supported_dataset_database_handlers.py
View file

@ -17,7 +17,7 @@ supported_dataset_database_handlers = {
"handler_provider": "neo4j",
},
"lancedb": {"handler_instance": LanceDBDatasetDatabaseHandler, "handler_provider": "lancedb"},
"pgvector": {
"pgvector_local": {
"handler_instance": PGVectorDatasetDatabaseHandler,
"handler_provider": "pgvector",
},

33
cognee/infrastructure/databases/vector/pgvector/PGVectorDatasetDatabaseHandler.py
View file

@ -13,35 +13,6 @@ class PGVectorDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
Handler for interacting with PGVector Dataset databases.
"""
@classmethod
async def _create_pg_database(cls, vector_config):
"""
Create the necessary Postgres database, and the PGVector extension on it.
This is defined here because the creation needs the latest vector config,
which is not yet saved in the vector config context variable here.
"""
from cognee.infrastructure.databases.relational.create_relational_engine import (
create_relational_engine,
)
from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from sqlalchemy import text
pg_relational_engine = create_relational_engine(
db_path="",
db_host=vector_config["vector_db_url"],
db_name=vector_config["vector_db_name"],
db_port=vector_config["vector_db_port"],
db_username=vector_config["vector_db_username"],
db_password=vector_config["vector_db_password"],
db_provider="postgres",
)
await pg_relational_engine.create_database()
vector_engine = create_vector_engine(**vector_config)
async with vector_engine.engine.begin() as connection:
await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
@classmethod
async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
vector_config = get_vectordb_config()
@ -65,7 +36,9 @@ class PGVectorDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
"vector_dataset_database_handler": "pgvector",
}
await cls._create_pg_database(
from .create_db_and_tables import create_pg_database
await create_pg_database(
{
"vector_db_provider": new_vector_config["vector_database_provider"],
"vector_db_url": new_vector_config["vector_database_url"],

31
cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py
View file

@ -10,3 +10,34 @@ async def create_db_and_tables():
if vector_config["vector_db_provider"] == "pgvector":
async with vector_engine.engine.begin() as connection:
await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
async def create_pg_database(vector_config):
"""
Create the necessary Postgres database, and the PGVector extension on it.
This is defined separately because the creation needs the latest vector config,
which is not yet saved in the vector config context variable.
TODO: We can maybe merge this with create_db_and_tables(), but it seemed simpler to separate them for now
"""
from cognee.infrastructure.databases.relational.create_relational_engine import (
create_relational_engine,
)
from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from sqlalchemy import text
pg_relational_engine = create_relational_engine(
db_path="",
db_host=vector_config["vector_db_url"],
db_name=vector_config["vector_db_name"],
db_port=vector_config["vector_db_port"],
db_username=vector_config["vector_db_username"],
db_password=vector_config["vector_db_password"],
db_provider="postgres",
)
await pg_relational_engine.create_database()
vector_engine = create_vector_engine(**vector_config)
async with vector_engine.engine.begin() as connection:
await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))