add country authority

lesson_graph.py

@@ -17,6 +17,8 @@ class NodeType(Enum):
     ASSESSMENT = "Assessment"
     RESOURCE = "Resource"
     SCHOOL_BOARD = "School Board"
+    COUNTRY_AUTHORITY = "Country Authority"  # New NodeType
+
 
 class ResetState(NamedTuple):
     teacher_name: str
@@ -27,6 +29,7 @@ class ResetState(NamedTuple):
     assessment: str
     resource: str
     school_board: str
+    country_authority: str  # New field
     message: str
 
 class LessonGraph:
@@ -38,7 +41,8 @@ class LessonGraph:
         "activity": "",
         "assessment": "",
         "resource": "",
-        "school_board": ""
+        "school_board": "",
+        "country_authority": ""  # New field
     }
 
     REQUIRED_FIELDS = ["teacher_name", "subject", "grade_level"]
@@ -51,7 +55,8 @@ class LessonGraph:
         NodeType.ACTIVITY: "#FF99FF",
         NodeType.ASSESSMENT: "#FFFF99",
         NodeType.RESOURCE: "#99FFFF",
-        NodeType.SCHOOL_BOARD: "#CCCCCC"
+        NodeType.SCHOOL_BOARD: "#CCCCCC",
+        NodeType.COUNTRY_AUTHORITY: "#FFA07A"  # New color for Country Authority
     }
 
     def __init__(self):
@@ -72,20 +77,25 @@ class LessonGraph:
         Add nodes and edges to the lesson plan graph for the given inputs.
         Returns a search string and the graph image.
         """
-        self.graph.clear()  # Clear previous graph
+        self.graph.clear()
         self.inputs.update(kwargs)
-        
-        # Validate only required fields
         self.validate_required_fields()
-
-        # Define node details and add them to the graph
+    
+        # Define required nodes
         nodes = {
             self.inputs["teacher_name"]: {"type": NodeType.USER, "role": "Teacher"},
             self.inputs["subject"]: {"type": NodeType.SUBJECT, "description": "Core subject area"},
             self.inputs["grade_level"]: {"type": NodeType.GRADE_LEVEL, "description": "Target grade for the lesson"},
         }
-
-        # Add optional nodes if they exist
+    
+        # Include country authority if provided
+        if self.inputs.get("country_authority"):
+            nodes[self.inputs["country_authority"]] = {
+                "type": NodeType.COUNTRY_AUTHORITY,
+                "description": "Sets national curriculum standards"
+            }
+    
+        # Optional nodes
         optional_nodes = {
             "learning_objective": NodeType.LEARNING_OBJECTIVE,
             "activity": NodeType.ACTIVITY,
@@ -93,88 +103,58 @@ class LessonGraph:
             "resource": NodeType.RESOURCE,
             "school_board": NodeType.SCHOOL_BOARD
         }
-
+    
         for field, node_type in optional_nodes.items():
             if self.inputs.get(field):
                 nodes[self.inputs[field]] = {"type": node_type, "description": f"{node_type.value}"}
-
+    
         # Add nodes to the graph
         for node, attributes in nodes.items():
             self.graph.add_node(node, **attributes)
-
-        # Define the relationships between nodes
+    
+        # Define relationships between nodes
         edges = [
             (self.inputs["teacher_name"], self.inputs["subject"], {"relationship": "TEACHES"}),
             (self.inputs["subject"], self.inputs["grade_level"], {"relationship": "HAS_GRADE"})
         ]
-
-        # Add optional edges
+    
+        # Relationships involving country authority
+        if self.inputs.get("country_authority"):
+            if self.inputs.get("learning_objective"):
+                edges.append((self.inputs["country_authority"], self.inputs["learning_objective"], {"relationship": "DEFINES"}))
+            if self.inputs.get("school_board"):
+                edges.append((self.inputs["country_authority"], self.inputs["school_board"], {"relationship": "OVERSEES"}))
+    
+        # Existing optional edges
         if self.inputs.get("learning_objective"):
-            edges.extend([
-                (self.inputs["subject"], self.inputs["learning_objective"], {"relationship": "COVERS"}),
-                (self.inputs["learning_objective"], self.inputs["school_board"], {"relationship": "ALIGNS_WITH"}) if self.inputs.get("school_board") else None
-            ])
-
+            edges.append((self.inputs["subject"], self.inputs["learning_objective"], {"relationship": "COVERS"}))
+            if self.inputs.get("school_board"):
+                edges.append((self.inputs["learning_objective"], self.inputs["school_board"], {"relationship": "ALIGNS_WITH"}))
+    
         if self.inputs.get("activity") and self.inputs.get("learning_objective"):
             edges.append((self.inputs["activity"], self.inputs["learning_objective"], {"relationship": "ACHIEVES"}))
-
+    
         if self.inputs.get("activity") and self.inputs.get("resource"):
             edges.append((self.inputs["activity"], self.inputs["resource"], {"relationship": "REQUIRES"}))
-
+    
         if self.inputs.get("learning_objective") and self.inputs.get("assessment"):
             edges.append((self.inputs["learning_objective"], self.inputs["assessment"], {"relationship": "EVALUATED_BY"}))
-
+    
         if self.inputs.get("school_board"):
             edges.append((self.inputs["teacher_name"], self.inputs["school_board"], {"relationship": "BELONGS_TO"}))
-
+    
         # Remove None entries from edges list
         edges = [edge for edge in edges if edge is not None]
         self.graph.add_edges_from(edges)
-
+    
         # Generate the search string for content discovery
         search_string = f"{self.inputs['subject']} {self.inputs['grade_level']} {self.inputs.get('learning_objective', '')} {self.inputs.get('activity', '')} {self.inputs.get('resource', '')}".strip()
-        
+    
         # Get the graph image
         image = self.draw_graph()
-
-        # Return the search string and the graph image
+    
         return search_string, image
 
-    def draw_graph(self) -> Image.Image:
-        """
-        Visualize the graph using Matplotlib, handling layout, labels, and interactivity.
-        """
-        fig, ax = plt.subplots(figsize=(14, 10))
-        pos = nx.spring_layout(self.graph, k=0.9, iterations=50)
-
-        self._draw_nodes(ax, pos)
-        self._draw_edges(ax, pos)
-        self._add_legend(ax)
-        
-        plt.title("Your Educational Landscape", fontsize=16)
-        plt.axis('off')
-        plt.tight_layout()
-        
-        self._add_interactivity()
-        
-        # Save the plot to a BytesIO object
-        buf = BytesIO()
-        plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
-        buf.seek(0)
-        plt.close(fig)
-
-        return Image.open(buf)
-
-    def _draw_nodes(self, ax, pos):
-        node_colors = [self.COLOR_MAP[self.graph.nodes[node]['type']] for node in self.graph.nodes()]
-        nx.draw_networkx_nodes(self.graph, pos, node_color=node_colors, node_size=3000, alpha=0.8, ax=ax)
-        nx.draw_networkx_labels(self.graph, pos, font_size=10, font_weight="bold", ax=ax)
-
-    def _draw_edges(self, ax, pos):
-        nx.draw_networkx_edges(self.graph, pos, edge_color='gray', arrows=True, arrowsize=20, ax=ax)
-        edge_labels = nx.get_edge_attributes(self.graph, 'relationship')
-        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels, font_size=8, ax=ax)
-
     def _add_legend(self, ax):
         legend_elements = [mpatches.Patch(color=color, label=node_type.value) for node_type, color in self.COLOR_MAP.items()]
         ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1, 1), title="Node Types")