cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py (1)

187-189: Exception chaining missing (raise … from e)

You wrapped the internal failure in CollectionNotFoundError but the original stack-trace is lost.
Either re-raise the caught exception (raise) or chain it:

-                raise CollectionNotFoundError(
-                    f"Collection '{collection_name}' not found!", log_level="DEBUG"
-                )
+                raise CollectionNotFoundError(
+                    f"Collection '{collection_name}' not found!", log_level="DEBUG"
+                ) from None

This satisfies Ruff B904 and keeps signal/noise balanced.

cognee-mcp/src/server.py (3)

42-46: Preserve root cause when re-raising

cognify_task swallows the original exception context, hindering debugging.

-            except Exception as e:
-                logger.error("Cognify process failed.")
-                raise ValueError(f"Failed to cognify: {str(e)}")
+            except Exception as e:
+                logger.error("Cognify process failed.")
+                raise ValueError(f"Failed to cognify: {e}") from e

Same pattern applies to other tasks that may raise.

🧰 Tools

🪛 Ruff (0.8.2)

46-46: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling

(B904)

---

56-60: log_file may be None

get_log_file_location() returns None if no FileHandler is configured, so the user sees an
unhelpful message (…file at: None).
Fallback to a default string or guard against None.

---

171-191: Graceful shutdown and signal handling

asyncio.run() inside a try/except logs startup errors but does not handle
KeyboardInterrupt or SIGTERM, which matter for containerised deployments.
Consider using asyncio.run(main(), debug=False) inside a proper signal handler
or rely on FastMCP’s own graceful shutdown utilities.

cognee/infrastructure/databases/graph/graph_db_interface.py (1)

61-61: 🛠️ Refactor suggestion

Reduced log noise by changing error messages to debug level.

Changed the logging level from ERROR to DEBUG for relationship errors and session commit errors in the record_graph_changes decorator. This reduces log noise since these errors might be expected in certain scenarios and don't necessarily indicate critical problems.

Consider making this configurable like the recent changes to the exception classes. This would provide consistent error handling across the codebase:

-logger.debug(f"Error adding relationship: {e}")
+logger.log(log_level, f"Error adding relationship: {e}")

Where log_level could be a parameter of the decorator or determined contextually.

Also applies to: 81-81, 88-88

🤖 Prompt for AI Agents (early access)

In cognee/infrastructure/databases/graph/graph_db_interface.py at lines 61, 81, and 88, the logging level for errors in the record_graph_changes decorator was changed from ERROR to DEBUG to reduce log noise. To improve consistency and flexibility, refactor the decorator to accept a configurable log_level parameter that defaults to DEBUG but can be set to other levels as needed. Update the logging calls to use this parameter instead of a hardcoded level.

cognee/tests/test_neo4j.py (1)

77-86: 🛠️ Refactor suggestion

Use unittest.skip instead of commenting-out flaky test

Commenting code hides drift and makes future maintenance harder.
Mark the flaky case as skipped so it still shows up in test reports and can be re-enabled later.

-    # NOTE: Due to the test failing often on weak LLM models we've removed this test for now
-    # search_results = await cognee.search(
-    #     query_type=SearchType.NATURAL_LANGUAGE,
-    #     query_text=f"Find nodes connected to node with name {random_node_name}",
-    # )
-    # assert len(search_results) != 0, "Query related natural language don't exist."
-    # print("\nExtracted results are:\n")
-    # for result in search_results:
-    #     print(f"{result}\n")
+    @unittest.skip("Temporarily disabled – flaky on weaker LLM models")
+    async def _natural_language_search():
+        search_results = await cognee.search(
+            query_type=SearchType.NATURAL_LANGUAGE,
+            query_text=f"Find nodes connected to node with name {random_node_name}",
+        )
+        self.assertTrue(search_results, "Natural-language search returned no results.")

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents (early access)

In cognee/tests/test_neo4j.py around lines 77 to 86, the flaky test is currently commented out, which hides it from test reports and complicates maintenance. Replace the commented-out test code with a proper test function decorated with unittest.skip, providing a reason for skipping. This way, the test remains visible in reports and can be easily re-enabled later.

cognee-mcp/src/server.py (2)

78-86: 🛠️ Refactor suggestion

Background task errors are lost

asyncio.create_task() is fire-and-forget; unhandled exceptions in codify_task
are only logged via the default loop exception handler.
Capture the task and add task.add_done_callback to log/propagate failures:

-    asyncio.create_task(codify_task(repo_path))
+    task = asyncio.create_task(codify_task(repo_path))
+    task.add_done_callback(
+        lambda t: logger.error(f"Codify task failed: {t.exception()}") if t.exception() else None
+    )

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents (early access)

In cognee-mcp/src/server.py around lines 78 to 86, the current use of asyncio.create_task() for the codify_task is fire-and-forget, causing any exceptions in the background task to be lost and only logged by the default loop handler. To fix this, capture the created task in a variable and add a done callback function that checks for exceptions and logs or propagates them appropriately, ensuring background task errors are not silently ignored.

---

111-120: ⚠️ Potential issue

Potential IndexError for empty result set

search_results[0] assumes at least one item. If the query misses you’ll raise
IndexError, which is mapped to an unhelpful MCP error.

-            elif (
-                search_type.upper() == "GRAPH_COMPLETION" or search_type.upper() == "RAG_COMPLETION"
-            ):
-                return search_results[0]
+            elif search_type.upper() in {"GRAPH_COMPLETION", "RAG_COMPLETION"}:
+                return search_results[0] if search_results else "No results found."

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

                query_type=SearchType[search_type.upper()], query_text=search_query
            )

            if search_type.upper() == "CODE":
                return json.dumps(search_results, cls=JSONEncoder)
            elif search_type.upper() in {"GRAPH_COMPLETION", "RAG_COMPLETION"}:
                return search_results[0] if search_results else "No results found."
            elif search_type.upper() == "CHUNKS":
                # ...

🤖 Prompt for AI Agents (early access)

In cognee-mcp/src/server.py around lines 111 to 120, the code accesses search_results[0] without checking if search_results is empty, risking an IndexError. To fix this, add a condition to verify that search_results contains at least one item before accessing the first element. If it is empty, handle the case gracefully by returning an appropriate response or error message instead of directly accessing search_results[0].

cognee/infrastructure/databases/graph/networkx/adapter.py (1)

45-45: 💡 Verification agent

🧩 Analysis chain

Consistent type updates for node identifiers

The type annotations have been updated from str to UUID for node identifiers across multiple methods in the NetworkXAdapter class. This change improves type safety and enforces consistent use of UUID objects throughout the graph operations.

However, I noticed some inconsistency in the codebase - some methods still use str type annotations for node identifiers, such as:

• add_edge (lines 75-81)
• add_edges (line 93)
• has_edge (line 62)

Consider updating the type annotations in these methods as well to maintain consistency across the entire adapter class:

-    async def has_edge(self, from_node: str, to_node: str, edge_label: str) -> bool:
+    async def has_edge(self, from_node: UUID, to_node: UUID, edge_label: str) -> bool:

-    async def add_edge(
-        self,
-        from_node: str,
-        to_node: str,
-        relationship_name: str,
-        edge_properties: Dict[str, Any] = {},
-    ) -> None:
+    async def add_edge(
+        self,
+        from_node: UUID,
+        to_node: UUID,
+        relationship_name: str,
+        edge_properties: Dict[str, Any] = {},
+    ) -> None:

-    async def add_edges(self, edges: list[tuple[str, str, str, dict]]) -> None:
+    async def add_edges(self, edges: list[tuple[UUID, UUID, str, dict]]) -> None:

Also applies to: 139-139, 177-177, 183-183, 218-218, 267-267, 278-278, 624-624, 629-629

---

🏁 Script executed:

#!/bin/bash
# Check for mixed usage of string and UUID node identifiers in the codebase

echo "Checking for NetworkXAdapter method calls with mixed node ID types..."

# Find places where the adapter methods are called
rg -A 1 "networkx_adapter\.(has_node|get_edges|extract_node|extract_nodes|get_neighbors|get_node|get_nodes)" --type py

echo -e "\nChecking for UUID imports and conversions..."
rg -A 3 "from uuid import UUID" --type py | head -20

echo -e "\nChecking for ID parsing or conversion..."
rg -A 1 "parse_id|UUID\(" --type py | head -20

Length of output: 4059

---

🏁 Script executed:

#!/bin/bash
# Inspect NetworkXAdapter method signatures for edge-related methods
echo "has_edge signature:"
rg -n -A0 "async def has_edge" cognee/infrastructure/databases/graph/networkx/adapter.py || echo "has_edge not found"

echo -e "\nadd_edge signature:"
rg -n -A3 "async def add_edge" cognee/infrastructure/databases/graph/networkx/adapter.py || echo "add_edge not found"

echo -e "\nadd_edges signature:"
rg -n -A0 "async def add_edges" cognee/infrastructure/databases/graph/networkx/adapter.py || echo "add_edges not found"

Length of output: 922

---

Ensure UUID Consistency for Edge-Related Methods

The has_node method now correctly uses UUID for identifiers, but the following methods in cognee/infrastructure/databases/graph/networkx/adapter.py still accept str:

• Line 62 – has_edge(self, from_node: str, to_node: str, edge_label: str) -> bool
• Lines 75–81 – add_edge(self, from_node: str, to_node: str, relationship_name: str, edge_properties: Dict[str, Any] = {}) -> None
• Line 93 – add_edges(self, edges: list[tuple[str, str, str, dict]]) -> None

Please update their signatures to use UUID for node identifiers:

-    async def has_edge(self, from_node: str, to_node: str, edge_label: str) -> bool:
+    async def has_edge(self, from_node: UUID, to_node: UUID, edge_label: str) -> bool:

-    async def add_edge(
-        self,
-        from_node: str,
-        to_node: str,
-        relationship_name: str,
-        edge_properties: Dict[str, Any] = {},
-    ) -> None:
+    async def add_edge(
+        self,
+        from_node: UUID,
+        to_node: UUID,
+        relationship_name: str,
+        edge_properties: Dict[str, Any] = {},
+    ) -> None:

-    async def add_edges(self, edges: list[tuple[str, str, str, dict]]) -> None:
+    async def add_edges(self, edges: list[tuple[UUID, UUID, str, dict]]) -> None:

This will enforce consistent use of UUID across all graph operations.

🤖 Prompt for AI Agents (early access)

In cognee/infrastructure/databases/graph/networkx/adapter.py around lines 62, 75-81, and 93, the methods has_edge, add_edge, and add_edges currently use str type annotations for node identifiers. Update these method signatures to use UUID instead of str for all node identifier parameters to maintain consistent type usage across the class. This involves changing from_node and to_node parameters in has_edge and add_edge, and the corresponding tuple elements in add_edges, to UUID type annotations.