cognee 0.5.3.dev1__tar.gz → 0.5.4.dev0__tar.gz

cognee-0.5.4.dev0/.dockerignore.ci

@@ -0,0 +1,4 @@
+*
+!pyproject.toml
+!uv.lock
+!README.md

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/.env.template

@@ -98,7 +98,18 @@ DB_NAME=cognee_db
 #     DATABASE_CONNECT_ARGS='{"sslmode": "require", "connect_timeout": 10}'
 #   For SQLite with custom timeout:
 #     DATABASE_CONNECT_ARGS='{"timeout": 60}'
-#DATABASE_CONNECT_ARGS='{}'
+# DATABASE_CONNECT_ARGS='{}'
+
+# -- Advanced: Custom database connection pool arguments (optional) ---------------
+# Pass additional connection pool parameters as JSON. Useful for managing the number of connections and their lifetimes.
+# Examples:
+#   SqlAlchemy Defaults:
+#     POOL_ARGS='{"pool_size": 5, "max_overflow": 10, "pool_recycle": -1, "pool_timeout": 30}'
+#   For heavy and slow workloads (eg: using Cognee with large datasets):
+#     POOL_ARGS='{"pool_size": 20, "max_overflow": 20, "pool_recycle": 280, "pool_timeout": 280, "pool_pre_ping": true}'
+# 		Note: Be careful using so many connections per instance in a multi-client environment as you can easily overwhelm a DB server this way.
+#		The suggested value for max_overflow is the number of CPU cores available, based on the SqlAlchemy and Postgresql recommendations.
+# POOL_ARGS='{}'
 
 ################################################################################
 # 🕸️ Graph Database settings

cognee-0.5.4.dev0/Dockerfile.ci

@@ -0,0 +1,26 @@
+FROM ghcr.io/astral-sh/uv:python3.11-bookworm-slim
+WORKDIR /app
+ENV UV_LINK_MODE=copy
+
+# System deps (superset of all CI needs)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    gcc g++ libpq-dev git curl cmake clang build-essential \
+    postgresql-client libgl1 libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY pyproject.toml uv.lock README.md ./
+
+# ALL extras used across CI - no-install-project so project code isn't baked in
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --frozen --no-install-project \
+    --extra api --extra docs --extra evals --extra codegraph --extra ollama \
+    --extra dev --extra neo4j --extra redis --extra postgres --extra aws \
+    --extra baml --extra distributed --extra deepeval --extra scraping \
+    --extra notebook --extra docling
+
+# Marker file for container detection in cognee_setup
+RUN echo "test-machine" > /app/.anon_id
+
+ENV PATH="/app/.venv/bin:$PATH"
+ENV PYTHONPATH=/app
+ENV PYTHONUNBUFFERED=1

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.5.3.dev1
+Version: 0.5.4.dev0
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/alembic/versions/18f01b0a0b4c_create_principal_configuration_table.py

@@ -27,11 +27,11 @@ def upgrade() -> None:
     if table_name not in inspector.get_table_names():
         op.create_table(
             table_name,
+            sa.Column("id", sa.UUID(), primary_key=True, index=True),
             sa.Column(
                 "owner_id",
                 sa.UUID(),
                 sa.ForeignKey("principals.id", ondelete="CASCADE"),
-                primary_key=True,
                 index=True,
             ),
             sa.Column("name", sa.String(), unique=False, nullable=False),

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/alembic/versions/b9274c27a25a_kuzu_11_migration.py

@@ -31,7 +31,7 @@ def upgrade() -> None:
 
         databases_root = os.path.join(base_config.system_root_directory, "databases")
         if not os.path.isdir(databases_root):
-            raise FileNotFoundError(f"Directory not found: {databases_root}")
+            return
 
         for current_path, dirnames, _ in os.walk(databases_root):
             # If file is kuzu graph database

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/api/v1/config/config.py

@@ -84,17 +84,25 @@ class config:
         llm_config.llm_api_key = llm_api_key
 
     @staticmethod
-    def set_llm_config(config_dict: dict):
+    def _update_config(config_obj, config_dict: dict):
         """
-        Updates the llm config with values from config_dict.
+        Updates a config object with values from config_dict after attribute validation.
         """
-        llm_config = get_llm_config()
         for key, value in config_dict.items():
-            if hasattr(llm_config, key):
-                object.__setattr__(llm_config, key, value)
+            if hasattr(config_obj, key):
+                object.__setattr__(config_obj, key, value)
             else:
                 raise InvalidConfigAttributeError(attribute=key)
 
+        return config_obj
+
+    @staticmethod
+    def set_llm_config(config_dict: dict):
+        """
+        Updates the llm config with values from config_dict.
+        """
+        config._update_config(get_llm_config(), config_dict)
+
     @staticmethod
     def set_chunk_strategy(chunk_strategy: object):
         chunk_config = get_chunk_config()
@@ -125,48 +133,28 @@ class config:
         """
         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 InvalidConfigAttributeError(attribute=key)
+        config._update_config(get_relational_config(), config_dict)
 
     @staticmethod
     def set_migration_db_config(config_dict: dict):
         """
         Updates the relational db config with values from config_dict.
         """
-        migration_db_config = get_migration_config()
-        for key, value in config_dict.items():
-            if hasattr(migration_db_config, key):
-                object.__setattr__(migration_db_config, key, value)
-            else:
-                raise InvalidConfigAttributeError(attribute=key)
+        config._update_config(get_migration_config(), config_dict)
 
     @staticmethod
     def set_graph_db_config(config_dict: dict) -> None:
         """
         Updates the graph db config with values from config_dict.
         """
-        graph_db_config = get_graph_config()
-        for key, value in config_dict.items():
-            if hasattr(graph_db_config, key):
-                object.__setattr__(graph_db_config, key, value)
-            else:
-                raise AttributeError(f"'{key}' is not a valid attribute of the config.")
+        config._update_config(get_graph_config(), config_dict)
 
     @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 InvalidConfigAttributeError(attribute=key)
+        config._update_config(get_vectordb_config(), config_dict)
 
     @staticmethod
     def set_vector_db_key(db_key: str):
@@ -197,13 +185,9 @@ class config:
         """
         Updates the translation config with values from config_dict.
         """
-        translation_config = get_translation_config()
-        for key, value in config_dict.items():
-            if hasattr(translation_config, key):
-                object.__setattr__(translation_config, key, value)
-            else:
-                raise InvalidConfigAttributeError(attribute=key)
+        config._update_config(get_translation_config(), config_dict)
 
+    @staticmethod
     def set(key: str, value):
         """
         Generic setter that maps configuration keys to their specific setter methods.

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/api/v1/datasets/datasets.py

@@ -84,7 +84,7 @@ class datasets:
         # even if some fail.
         if dataset_data:
             results = await asyncio.gather(
-                *[delete_data(data) for data in dataset_data],
+                *[delete_data(data, dataset_id) for data in dataset_data],
                 return_exceptions=True,
             )
             deletion_errors = [r for r in results if isinstance(r, Exception)]
@@ -145,7 +145,7 @@ class datasets:
         else:
             await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
 
-        await delete_data(data)
+        await delete_data(data, dataset_id)
 
         dataset_data = await get_dataset_data(dataset.id)
         if not dataset_data and delete_dataset_if_empty:

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/api/v1/health/health.py

@@ -249,13 +249,8 @@ class HealthChecker:
             ("vector_db", self.check_vector_db()),
             ("graph_db", self.check_graph_db()),
             ("file_storage", self.check_file_storage()),
-            ("llm_provider", self.check_llm_provider()),
-            ("embedding_service", self.check_embedding_service()),
         ]
 
-        # Non-critical services (only for detailed checks)
-        non_critical_checks = []
-
         # Run critical checks
         critical_results = await asyncio.gather(
             *[check for _, check in critical_checks], return_exceptions=True
@@ -273,7 +268,13 @@ class HealthChecker:
                 components[name] = result
 
         # Run non-critical checks if detailed (currently none)
-        if detailed and non_critical_checks:
+        if detailed:
+            # Non-critical services (only for detailed checks)
+            non_critical_checks = [
+                ("llm_provider", self.check_llm_provider()),
+                ("embedding_service", self.check_embedding_service()),
+            ]
+
             non_critical_results = await asyncio.gather(
                 *[check for _, check in non_critical_checks], return_exceptions=True
             )

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/api/v1/users/routers/get_configuration_router.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter
 from fastapi import Form, Depends
+from uuid import UUID
 
 from cognee.api.DTO import InDTO
 from cognee.modules.users.models import User
@@ -11,6 +12,7 @@ from cognee.modules.users.methods import (
 )
 from cognee.modules.users.methods import (
     get_principal_configuration as method_get_principal_configuration,
+    get_principal_all_configuration as method_get_principal_all_configuration,
 )
 
 logger = get_logger()
@@ -36,11 +38,17 @@ def get_configuration_router() -> APIRouter:
             principal_id=user.id, name=payload.name, configuration=payload.config
         )
 
-    @router.get("/get_user_configuration/{name}", response_model=dict)
+    @router.get("/get_user_configuration/{config_id}", response_model=dict)
     async def get_user_configuration(
-        name: str,
+        config_id: UUID,
         user: User = Depends(get_authenticated_user),
     ):
-        return await method_get_principal_configuration(principal_id=user.id, name=name)
+        return await method_get_principal_configuration(config_id=config_id)
+
+    @router.get("/get_user_configuration/", response_model=list)
+    async def get_user_all_configuration(
+        user: User = Depends(get_authenticated_user),
+    ):
+        return await method_get_principal_all_configuration(principal_id=user.id)
 
     return router

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/context_global_variables.py

@@ -103,12 +103,13 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
     """
     If backend access control is enabled this function will ensure all datasets have their own databases,
     access to which will be enforced by given permissions.
-    Database name will be determined by dataset_id and LanceDB and KuzuDB use will be enforced.
+    Database name will be determined by dataset and the appropriate vector and
+    graph database handlers will be enforced.
 
     Note: This is only currently supported by the following databases:
           Relational: SQLite, Postgres
-          Vector: LanceDB
-          Graph: KuzuDB
+          Vector: LanceDB, pgvector
+          Graph: KuzuDB, neo4j_aura_dev
 
     Args:
         dataset: Cognee dataset name or id

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/cache/cache_db_interface.py

@@ -74,10 +74,12 @@ class CacheDBInterface(ABC):
         qa_id: str,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ):
         """
         Add a Q/A/context triplet to a cache session.
         Uses the same QA fields as update_qa_entry for consistent structure.
+        used_graph_element_ids: Optional dict with keys "node_ids" and "edge_ids" (lists of str).
         """
         pass
 
@@ -114,6 +116,7 @@ class CacheDBInterface(ABC):
         answer: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> bool:
         """
         Update a QA entry by qa_id. Same QA fields as create_qa_entry.

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/cache/config.py

@@ -15,10 +15,12 @@ class CacheConfig(BaseSettings):
     - agentic_lock_timeout: Maximum time (in seconds) to wait for the lock release.
     - usage_logging: Enable/disable usage logging for API endpoints and MCP tools.
     - usage_logging_ttl: Time-to-live for usage logs in seconds (default: 7 days).
+    - auto_feedback: When caching is True, run automatic feedback detection on each query (default False).
     """
 
     cache_backend: Literal["redis", "fs"] = "fs"
     caching: bool = False
+    auto_feedback: bool = False
     shared_kuzu_lock: bool = False
     cache_host: str = "localhost"
     cache_port: int = 6379
@@ -35,6 +37,7 @@ class CacheConfig(BaseSettings):
         return {
             "cache_backend": self.cache_backend,
             "caching": self.caching,
+            "auto_feedback": self.auto_feedback,
             "shared_kuzu_lock": self.shared_kuzu_lock,
             "cache_host": self.cache_host,
             "cache_port": self.cache_port,

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/cache/fscache/FsCacheAdapter.py

@@ -44,6 +44,7 @@ class FSCacheAdapter(CacheDBInterface):
         qa_id: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> dict:
         entry = SessionQAEntry(
             time=datetime.utcnow().isoformat(),
@@ -53,6 +54,7 @@ class FSCacheAdapter(CacheDBInterface):
             qa_id=qa_id or str(uuid.uuid4()),
             feedback_text=feedback_text,
             feedback_score=feedback_score,
+            used_graph_element_ids=used_graph_element_ids,
         )
         return entry.model_dump()
 
@@ -76,6 +78,7 @@ class FSCacheAdapter(CacheDBInterface):
         answer: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> dict:
         merged = {**entry}
         if question is not None:
@@ -88,6 +91,8 @@ class FSCacheAdapter(CacheDBInterface):
             merged["feedback_text"] = feedback_text
         if feedback_score is not None:
             merged["feedback_score"] = feedback_score
+        if used_graph_element_ids is not None:
+            merged["used_graph_element_ids"] = used_graph_element_ids
         return merged
 
     @staticmethod
@@ -132,11 +137,18 @@ class FSCacheAdapter(CacheDBInterface):
         qa_id: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ):
         try:
             session_key = self._session_key(user_id, session_id)
             qa_entry = self._build_qa_entry_dump(
-                question, context, answer, qa_id, feedback_text, feedback_score
+                question,
+                context,
+                answer,
+                qa_id,
+                feedback_text,
+                feedback_score,
+                used_graph_element_ids=used_graph_element_ids,
             )
             entries = self._load_entries(session_key)
             entries.append(qa_entry)
@@ -167,6 +179,7 @@ class FSCacheAdapter(CacheDBInterface):
         answer: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> bool:
         """
         Update a QA entry by qa_id. Same QA fields as create_qa_entry.
@@ -180,7 +193,13 @@ class FSCacheAdapter(CacheDBInterface):
             if idx is None:
                 return False
             merged = self._merge_entry_update(
-                entries[idx], question, context, answer, feedback_text, feedback_score
+                entries[idx],
+                question,
+                context,
+                answer,
+                feedback_text,
+                feedback_score,
+                used_graph_element_ids=used_graph_element_ids,
             )
             entries[idx] = self._validate_entry_dict(merged)
             self._save_entries(session_key, entries)

cognee-0.5.4.dev0/cognee/infrastructure/databases/cache/models.py

@@ -0,0 +1,62 @@
+from typing import Dict, List, Optional
+from pydantic import BaseModel, field_validator
+
+
+def _validate_list_of_str(value: object, key: str) -> List[str]:
+    if not isinstance(value, list):
+        raise ValueError(f"{key} must be a list")
+    for i, item in enumerate(value):
+        if not isinstance(item, str):
+            raise ValueError(f"{key}[{i}] must be a string")
+    return value
+
+
+class SessionQAEntry(BaseModel):
+    """
+    Canonical format for a QA entry stored in session cache.
+
+    Fields:
+        time: ISO format timestamp when the QA was created.
+        qa_id: Unique identifier for the entry (required for update/delete).
+        question: User's question.
+        context: Context used to generate the answer.
+        answer: Generated answer.
+        feedback_text: Optional user feedback text.
+        feedback_score: Optional feedback score 1-5.
+        used_graph_element_ids: Optional dict with only "node_ids" and "edge_ids" (lists of str).
+    """
+
+    time: str
+    question: str
+    context: str
+    answer: str
+    qa_id: Optional[str] = None
+    feedback_text: Optional[str] = None
+    feedback_score: Optional[int] = None
+    used_graph_element_ids: Optional[Dict[str, List[str]]] = None
+
+    @field_validator("used_graph_element_ids")
+    @classmethod
+    def used_graph_element_ids_only_node_and_edge_ids(
+        cls, v: Optional[Dict[str, List[str]]]
+    ) -> Optional[Dict[str, List[str]]]:
+        if v is None:
+            return None
+        if not isinstance(v, dict):
+            raise ValueError("used_graph_element_ids must be a dict or None")
+        allowed = {"node_ids", "edge_ids"}
+        if set(v.keys()) - allowed:
+            raise ValueError("used_graph_element_ids may only have keys 'node_ids' and 'edge_ids'")
+        out: Dict[str, List[str]] = {}
+        if "node_ids" in v:
+            out["node_ids"] = _validate_list_of_str(v["node_ids"], "node_ids")
+        if "edge_ids" in v:
+            out["edge_ids"] = _validate_list_of_str(v["edge_ids"], "edge_ids")
+        return out if out else None
+
+    @field_validator("feedback_score")
+    @classmethod
+    def feedback_score_range(cls, v: Optional[int]) -> Optional[int]:
+        if v is not None and (v < 1 or v > 5):
+            raise ValueError("feedback_score must be between 1 and 5")
+        return v

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/cache/redis/RedisAdapter.py

@@ -90,6 +90,7 @@ class RedisAdapter(CacheDBInterface):
         qa_id: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> dict:
         entry = SessionQAEntry(
             time=datetime.utcnow().isoformat(),
@@ -99,6 +100,7 @@ class RedisAdapter(CacheDBInterface):
             qa_id=qa_id or str(uuid.uuid4()),
             feedback_text=feedback_text,
             feedback_score=feedback_score,
+            used_graph_element_ids=used_graph_element_ids,
         )
         return entry.model_dump()
 
@@ -122,6 +124,7 @@ class RedisAdapter(CacheDBInterface):
         answer: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> dict:
         merged = {**entry}
         if question is not None:
@@ -134,6 +137,8 @@ class RedisAdapter(CacheDBInterface):
             merged["feedback_text"] = feedback_text
         if feedback_score is not None:
             merged["feedback_score"] = feedback_score
+        if used_graph_element_ids is not None:
+            merged["used_graph_element_ids"] = used_graph_element_ids
         return merged
 
     @staticmethod
@@ -204,6 +209,7 @@ class RedisAdapter(CacheDBInterface):
         qa_id: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ):
         """
         Add a Q/A/context triplet to a Redis list for this session.
@@ -212,7 +218,13 @@ class RedisAdapter(CacheDBInterface):
         try:
             session_key = self._session_key(user_id, session_id)
             qa_entry = self._build_qa_entry_dump(
-                question, context, answer, qa_id, feedback_text, feedback_score
+                question,
+                context,
+                answer,
+                qa_id,
+                feedback_text,
+                feedback_score,
+                used_graph_element_ids=used_graph_element_ids,
             )
             await self.async_redis.rpush(session_key, json.dumps(qa_entry))
         except (redis.ConnectionError, redis.TimeoutError) as e:
@@ -252,6 +264,7 @@ class RedisAdapter(CacheDBInterface):
         answer: str | None = None,
         feedback_text: str | None = None,
         feedback_score: int | None = None,
+        used_graph_element_ids: dict | None = None,
     ) -> bool:
         """
         Update a QA entry by qa_id. Same QA fields as create_qa_entry.
@@ -265,7 +278,13 @@ class RedisAdapter(CacheDBInterface):
             if idx is None:
                 return False
             merged = self._merge_entry_update(
-                entries[idx], question, context, answer, feedback_text, feedback_score
+                entries[idx],
+                question,
+                context,
+                answer,
+                feedback_text,
+                feedback_score,
+                used_graph_element_ids=used_graph_element_ids,
             )
             entries[idx] = self._validate_entry_dict(merged)
             await self._write_entry_at(session_key, idx, entries[idx])

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -859,7 +859,22 @@ class KuzuAdapter(GraphDBInterface):
             - List[Dict[str, Any]]: A list of dictionaries representing neighboring nodes'
               properties.
         """
-        return await self.get_neighbours(node_id)
+        query_str = """
+        MATCH (n:Node)-[r]-(m:Node)
+        WHERE n.id = $id
+        RETURN DISTINCT {
+            id: m.id,
+            name: m.name,
+            type: m.type,
+            properties: m.properties
+        }
+        """
+        try:
+            result = await self.query(query_str, {"id": node_id})
+            return [self._parse_node_properties(row[0]) for row in result] if result else []
+        except Exception as e:
+            logger.error(f"Failed to get neighbours for node {node_id}: {e}")
+            return []
 
     async def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
         """
@@ -934,37 +949,6 @@ class KuzuAdapter(GraphDBInterface):
             logger.error(f"Failed to get nodes: {e}")
             return []
 
-    async def get_neighbours(self, node_id: str) -> List[Dict[str, Any]]:
-        """
-        Get all neighbouring nodes.
-
-        This method retrieves all neighboring nodes connected to a specified node and returns
-        them as a list of dictionaries. It may return an empty list if no neighbors exist or an
-        error occurs.
-
-        Parameters:
-        -----------
-
-            - node_id (str): The identifier of the node for which to find neighbors.
-
-        Returns:
-        --------
-
-            - List[Dict[str, Any]]: A list of dictionaries representing neighboring nodes'
-              properties.
-        """
-        query_str = """
-        MATCH (n)-[r]-(m)
-        WHERE n.id = $id
-        RETURN DISTINCT properties(m)
-        """
-        try:
-            result = await self.query(query_str, {"id": node_id})
-            return [row[0] for row in result] if result else []
-        except Exception as e:
-            logger.error(f"Failed to get neighbours for node {node_id}: {e}")
-            return []
-
     async def get_predecessors(
         self, node_id: Union[str, UUID], edge_label: Optional[str] = None
     ) -> List[Dict[str, Any]]:

{cognee-0.5.3.dev1 → cognee-0.5.4.dev0}/cognee/infrastructure/databases/relational/config.py

@@ -21,6 +21,7 @@ class RelationalConfig(BaseSettings):
     db_password: Union[str, None] = None  # "cognee"
     db_provider: str = "sqlite"
     database_connect_args: Union[str, None] = None
+    pool_args: Union[str, None] = None
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -37,12 +38,27 @@ class RelationalConfig(BaseSettings):
             try:
                 parsed_args = json.loads(self.database_connect_args)
                 if isinstance(parsed_args, dict):
-                    self.database_connect_args = parsed_args
+                    # Note: For caching purposes, database_connect_args is stored as a sorted tuple of key-value pairs in the config
+                    #       It is later returned to a dictionary format
+                    self.database_connect_args = tuple(sorted(parsed_args.items()))
                 else:
                     self.database_connect_args = {}
             except json.JSONDecodeError:
                 self.database_connect_args = {}
 
+        # Parse pool_args if provided as JSON string
+        if self.pool_args and isinstance(self.pool_args, str):
+            try:
+                parsed_args = json.loads(self.pool_args)
+                if isinstance(parsed_args, dict):
+                    # Note: For caching purposes, pool_args is stored as a sorted tuple of key-value pairs in the config
+                    #       It is later returned to a dictionary format
+                    self.pool_args = tuple(sorted(parsed_args.items()))
+                else:
+                    self.pool_args = {}
+            except json.JSONDecodeError:
+                self.pool_args = {}
+
         return self
 
     def to_dict(self) -> dict:
@@ -65,6 +81,7 @@ class RelationalConfig(BaseSettings):
             "db_password": self.db_password,
             "db_provider": self.db_provider,
             "database_connect_args": self.database_connect_args,
+            "pool_args": self.pool_args,
         }