diff --git a/.env.template b/.env.template index 7dcd4f346..f1a549c86 100644 --- a/.env.template +++ b/.env.template @@ -142,6 +142,28 @@ MIGRATION_DB_PROVIDER="sqlite" # MIGRATION_DB_HOST="127.0.0.1" # MIGRATION_DB_PORT=5432 +################################################################################ +# ⚡ Performance & Caching Settings +################################################################################ + +# Cache Control: Toggle database engine caching behavior +# By default, database engines (vector, graph, relational) are cached using lru_cache +# to improve performance. However, this can cause issues with event loops in certain +# scenarios (e.g., when using threading + asyncio patterns, Jupyter notebooks, etc.) + +# -- Global cache control (disables ALL caches) ------------------------------ +# Uncomment to disable all engine caching globally: +# COGNEE_DISABLE_ALL_CACHES=true + +# -- Per-engine cache control (fine-grained control) ------------------------- +# Uncomment to disable specific engine caches: +# COGNEE_DISABLE_VECTOR_ENGINE_CACHE=true +# COGNEE_DISABLE_GRAPH_ENGINE_CACHE=true +# COGNEE_DISABLE_RELATIONAL_ENGINE_CACHE=true + +# Note: If you're experiencing "RuntimeError: asyncio.Lock is bound to a different +# event loop" errors, try setting COGNEE_DISABLE_ALL_CACHES=true + ################################################################################ # 🔒 Security Settings ################################################################################ diff --git a/cognee/infrastructure/databases/graph/get_graph_engine.py b/cognee/infrastructure/databases/graph/get_graph_engine.py index 1ea61d29f..f2893aa75 100644 --- a/cognee/infrastructure/databases/graph/get_graph_engine.py +++ b/cognee/infrastructure/databases/graph/get_graph_engine.py @@ -1,6 +1,6 @@ """Factory function to get the appropriate graph client based on the graph type.""" -from functools import lru_cache +from cognee.shared.cache_utils import cacheable from .config import get_graph_context_config from .graph_db_interface import GraphDBInterface @@ -24,7 +24,7 @@ async def get_graph_engine() -> GraphDBInterface: return graph_client -@lru_cache +@cacheable(key="GRAPH_ENGINE") def create_graph_engine( graph_database_provider, graph_file_path, diff --git a/cognee/infrastructure/databases/relational/create_relational_engine.py b/cognee/infrastructure/databases/relational/create_relational_engine.py index deaeaa2da..b37a1f320 100644 --- a/cognee/infrastructure/databases/relational/create_relational_engine.py +++ b/cognee/infrastructure/databases/relational/create_relational_engine.py @@ -1,8 +1,8 @@ from .sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter -from functools import lru_cache +from cognee.shared.cache_utils import cacheable -@lru_cache +@cacheable(key="RELATIONAL_ENGINE") def create_relational_engine( db_path: str, db_name: str, diff --git a/cognee/infrastructure/databases/vector/create_vector_engine.py b/cognee/infrastructure/databases/vector/create_vector_engine.py index ce66d1381..bcd9f6c66 100644 --- a/cognee/infrastructure/databases/vector/create_vector_engine.py +++ b/cognee/infrastructure/databases/vector/create_vector_engine.py @@ -1,6 +1,9 @@ from .supported_databases import supported_databases from .embeddings import get_embedding_engine +from cognee.shared.cache_utils import cacheable + +@cacheable(key="VECTOR_ENGINE") def create_vector_engine( vector_db_provider: str, vector_db_url: str, diff --git a/cognee/shared/cache_utils.py b/cognee/shared/cache_utils.py new file mode 100644 index 000000000..063380776 --- /dev/null +++ b/cognee/shared/cache_utils.py @@ -0,0 +1,46 @@ +"""Cache utilities for cognee.""" + +import os +from functools import lru_cache + + +def cacheable(func=None, *, key=None): + """ + Decorator to optionally cache function results based on environment variables. + + Usage: + @cacheable + def my_func(): ... + + @cacheable(key="VECTOR_ENGINE") + def my_func(): ... + + Configuration: + COGNEE_DISABLE_ALL_CACHES=true: Disables caching for all functions decorated with @cacheable + COGNEE_DISABLE_{KEY}_CACHE=true: Disables caching for specific function if key is provided + + Example: + # Disable all caches globally + COGNEE_DISABLE_ALL_CACHES=true + + # Disable only vector engine cache + COGNEE_DISABLE_VECTOR_ENGINE_CACHE=true + """ + + def decorator(f): + # Check global kill switch first + if os.getenv("COGNEE_DISABLE_ALL_CACHES", "false").lower() == "true": + return f + + # Check specific feature kill switch if key is provided + if key and os.getenv(f"COGNEE_DISABLE_{key}_CACHE", "false").lower() == "true": + return f + + return lru_cache(f) + + if func is None: + # Called with arguments: @cacheable(key="SOMETHING") + return decorator + else: + # Called without arguments: @cacheable + return decorator(func)