Add refresh_ontology() method to kg

graphrag_sdk/kg.py

@@ -61,18 +61,18 @@ def __init__(
         # Connect to database
         self.db = FalkorDB(host=host, port=port, username=username, password=password)
         self.graph = self.db.select_graph(name)
-        ontology_graph = self.db.select_graph("{" + name + "}" + "_schema")
+        self.ontology_graph = self.db.select_graph("{" + name + "}" + "_schema")
 
         # Load / Save ontology to database
         if ontology is None:
             # Load ontology from DB
-            ontology = Ontology.from_schema_graph(ontology_graph)
+            ontology = Ontology.from_schema_graph(self.ontology_graph)
 
             if len(ontology.entities) == 0:
                 raise Exception("The ontology is empty. Load a valid ontology or create one using the ontology module.")
         else:
             # Save ontology to DB
-            ontology.save_to_graph(ontology_graph)
+            ontology.save_to_graph(self.ontology_graph)
 
         self._ontology = ontology
         self._name = name
@@ -210,6 +210,22 @@ def chat_session(self) -> ChatSession:
         chat_session = ChatSession(self._model_config, self.ontology, self.graph, self.cypher_system_instruction,
                                    self.qa_system_instruction, self.cypher_gen_prompt, self.qa_prompt, self.cypher_gen_prompt_history)
         return chat_session
+
+    def refresh_ontology(self) -> None:
+        """
+        Refresh the ontology by reloading it from the database.
+        This is useful when the schema has been updated.
+
+        Raises:
+            Exception: If the refreshed ontology is empty and no fallback is available.
+        """
+        # Reload ontology from database
+        refreshed_ontology = Ontology.from_schema_graph(self.ontology_graph)
+
+        # Always update the ontology, even if it's empty
+        # This allows users to intentionally clear the ontology if needed
+        self._ontology = refreshed_ontology
+
     def add_node(self, entity: str, attributes: dict) -> None:
         """
         Add a node to the knowledge graph, checking if it matches the ontology