fix: reuse PostgreSQL database connections (#574)

<!-- .github/pull_request_template.md -->

## Description
Fix PostgreSQL database connection problems

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Refactor**
- Improved the system’s database connection process to enhance
compatibility across multiple relational databases. The application now
dynamically selects the optimal connection method—reusing established
connections when possible—to ensure improved stability and performance
without affecting the public interface.
- Streamlined the creation of the embedding engine by removing it as a
parameter and generating it internally.
- Removed dependency on the embedding engine in the vector engine
retrieval process.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
Igor Ilic 2025-02-24 20:35:40 +01:00 committed by GitHub
parent 6e567445b5
commit 4f354ba534
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 17 additions and 5 deletions

View file

@ -1,14 +1,17 @@
from .embeddings import get_embedding_engine
from functools import lru_cache
@lru_cache
def create_vector_engine(
embedding_engine,
vector_db_url: str,
vector_db_port: str,
vector_db_key: str,
vector_db_provider: str,
):
embedding_engine = get_embedding_engine()
if vector_db_provider == "weaviate":
from .weaviate_db import WeaviateAdapter

View file

@ -1,7 +1,6 @@
from .config import get_vectordb_config
from .embeddings import get_embedding_engine
from .create_vector_engine import create_vector_engine
def get_vector_engine():
return create_vector_engine(get_embedding_engine(), **get_vectordb_config().to_dict())
return create_vector_engine(**get_vectordb_config().to_dict())

View file

@ -10,6 +10,7 @@ from cognee.exceptions import InvalidValueError
from cognee.infrastructure.databases.exceptions import EntityNotFoundError
from cognee.infrastructure.engine import DataPoint
from cognee.infrastructure.engine.utils import parse_id
from cognee.infrastructure.databases.relational import get_relational_engine
from ...relational.ModelBase import Base
from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
@ -36,8 +37,17 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
self.api_key = api_key
self.embedding_engine = embedding_engine
self.db_uri: str = connection_string
self.engine = create_async_engine(self.db_uri)
self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
relational_db = get_relational_engine()
# If postgreSQL is used we must use the same engine and sessionmaker
if relational_db.engine.dialect.name == "postgresql":
self.engine = relational_db.engine
self.sessionmaker = relational_db.sessionmaker
else:
# If not create new instances of engine and sessionmaker
self.engine = create_async_engine(self.db_uri)
self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
# Has to be imported at class level
# Functions reading tables from database need to know what a Vector column type is