From 6be6b3d222c41699b52c374e55c180cec16fec5e Mon Sep 17 00:00:00 2001
From: Igor Ilic <30923996+dexters1@users.noreply.github.com>
Date: Sat, 8 Feb 2025 00:46:46 +0100
Subject: [PATCH] =?UTF-8?q?fix:=20resolve=20key=20error=20for=20visualizat?=
 =?UTF-8?q?ion,=20handle=20ValueError=20for=20jedi=20=E2=80=A6=20(#509)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…and move duplicate edge information to debug log

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

## Description
Fix visualization bug
Handle ValueError for CodeGraph
Move debug information from print to debug logs

## 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

- **Bug Fixes**
- Enhanced error handling across several modules to ensure smoother
operation when unexpected conditions occur.
- Updated diagnostic and logging mechanisms to provide more robust
system feedback and reduce potential disruptions.
- Improved robustness in the deletion of properties to prevent runtime
errors related to missing keys.
- Added additional exception handling for better analysis of code
entities.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 .../modules/graph/cognee_graph/CogneeGraph.py   |  8 ++++----
 .../cognee_network_visualization.py             | 17 +++++++++++++----
 .../repo_processor/get_local_dependencies.py    |  3 +++
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py
index dab8371b9..b53367e8d 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py
@@ -1,5 +1,4 @@
-import numpy as np
-
+import logging
 from typing import List, Dict, Union
 
 from cognee.exceptions import InvalidValueError
@@ -8,7 +7,8 @@ from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInte
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
 from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph
 import heapq
-import asyncio
+
+logger = logging.getLogger(__name__)
 
 
 class CogneeGraph(CogneeAbstractGraph):
@@ -40,7 +40,7 @@ class CogneeGraph(CogneeAbstractGraph):
             edge.node1.add_skeleton_edge(edge)
             edge.node2.add_skeleton_edge(edge)
         else:
-            print(f"Edge {edge} already exists in the graph.")
+            logger.debug(f"Edge {edge} already exists in the graph.")
 
     def get_node(self, node_id: str) -> Node:
         return self.nodes.get(node_id, None)
diff --git a/cognee/modules/visualization/cognee_network_visualization.py b/cognee/modules/visualization/cognee_network_visualization.py
index df06b2094..e92f7783f 100644
--- a/cognee/modules/visualization/cognee_network_visualization.py
+++ b/cognee/modules/visualization/cognee_network_visualization.py
@@ -21,10 +21,19 @@ async def cognee_network_visualization(graph_data):
         node_info["id"] = str(node_id)
         node_info["color"] = color_map.get(node_info.get("pydantic_type", "default"), "#D3D3D3")
         node_info["name"] = node_info.get("name", str(node_id))
-        del node_info[
-            "updated_at"
-        ]  #:TODO: We should decide what properties to show on the nodes and edges, we dont necessarily need all.
-        del node_info["created_at"]
+
+        try:
+            del node_info[
+                "updated_at"
+            ]  #:TODO: We should decide what properties to show on the nodes and edges, we dont necessarily need all.
+        except KeyError:
+            pass
+
+        try:
+            del node_info["created_at"]
+        except KeyError:
+            pass
+
         nodes_list.append(node_info)
         G.add_node(node_id, **node_info)
 
diff --git a/cognee/tasks/repo_processor/get_local_dependencies.py b/cognee/tasks/repo_processor/get_local_dependencies.py
index 600ff25a5..c04bdbaf2 100644
--- a/cognee/tasks/repo_processor/get_local_dependencies.py
+++ b/cognee/tasks/repo_processor/get_local_dependencies.py
@@ -79,6 +79,9 @@ def _update_code_entity(script: jedi.Script, code_entity: Dict[str, any]) -> Non
     except AssertionError as e:
         # TODO: See if there is a way to handle AttributeError properly
         logger.error(f"Failed to analyze code entity {code_entity['name']}: {e}")
+    except ValueError as e:
+        # TODO: See if there is a way to handle ValueError properly
+        logger.error(f"Failed to analyze code entity {code_entity['name']}: {e}")
     except Exception as e:
         # logging.warning(f"Failed to analyze code entity {code_entity['name']}: {e}")
         logger.error(f"Failed to analyze code entity {code_entity['name']}: {e}")