diff --git a/graphs_trees/graph/graph.py b/graphs_trees/graph/graph.py
index 755b4bb..0681db3 100644
--- a/graphs_trees/graph/graph.py
+++ b/graphs_trees/graph/graph.py
@@ -47,20 +47,18 @@ class Graph:
     def add_node(self, key):
         if key is None:
             raise TypeError('key cannot be None')
-        if key in self.nodes:
-            return self.nodes[key]
-        self.nodes[key] = Node(key)
+        if key not in self.nodes:
+            self.nodes[key] = Node(key)
         return self.nodes[key]
 
     def add_edge(self, source_key, dest_key, weight=0):
         if source_key is None or dest_key is None:
-            raise Exception('Invalid key')
+            raise KeyError('Invalid key')
         if source_key not in self.nodes:
             self.add_node(source_key)
         if dest_key not in self.nodes:
             self.add_node(dest_key)
-        self.nodes[source_key].add_neighbor(self.nodes[dest_key],
-                                            weight)
+        self.nodes[source_key].add_neighbor(self.nodes[dest_key], weight)
 
     def add_undirected_edge(self, src_key, dst_key, weight=0):
         if src_key is None or dst_key is None:
diff --git a/graphs_trees/graph/graph_solution.ipynb b/graphs_trees/graph/graph_solution.ipynb
index 64ec832..0d392d6 100644
--- a/graphs_trees/graph/graph_solution.ipynb
+++ b/graphs_trees/graph/graph_solution.ipynb
@@ -215,20 +215,18 @@
     "    def add_node(self, key):\n",
     "        if key is None:\n",
     "            raise TypeError('key cannot be None')\n",
-    "        if key in self.nodes:\n",
-    "            return self.nodes[key]\n",
-    "        self.nodes[key] = Node(key)\n",
+    "        if key not in self.nodes:\n",
+    "            self.nodes[key] = Node(key)\n",
     "        return self.nodes[key]\n",
     "\n",
     "    def add_edge(self, source_key, dest_key, weight=0):\n",
     "        if source_key is None or dest_key is None:\n",
-    "            raise Exception('Invalid key')\n",
+    "            raise KeyError('Invalid key')\n",
     "        if source_key not in self.nodes:\n",
     "            self.add_node(source_key)\n",
     "        if dest_key not in self.nodes:\n",
     "            self.add_node(dest_key)\n",
-    "        self.nodes[source_key].add_neighbor(self.nodes[dest_key],\n",
-    "                                            weight)\n",
+    "        self.nodes[source_key].add_neighbor(self.nodes[dest_key], weight)\n",
     "\n",
     "    def add_undirected_edge(self, src_key, dst_key, weight=0):\n",
     "        if src_key is None or dst_key is None:\n",