feat: added return types

brickllm/graphs/brickschema_graph.py

@@ -1,4 +1,3 @@
-import json
 from langgraph.graph import START, END, StateGraph
 from .. import State, GraphConfig
 from ..nodes import (
@@ -51,7 +50,7 @@ def _compiled_graph(self):
             raise ValueError("Graph is not compiled yet. Please compile the graph first.")
         return self.graph
 
-    def display(self, filename="graph.png"):
+    def display(self, filename="graph.png") -> None:
         """Display the compiled graph as an image.
         
         Args:
@@ -80,10 +79,21 @@ def run(self, prompt, stream=False):
         input_data = {"user_prompt": prompt}
 
         if stream:
+            events = []
             # Stream the content of the graph state at each node
             for event in self.graph.stream(input_data, self.config, stream_mode="values"):
-                print(json.dumps(event, indent=2))
+                events.append(event)
+            return events
         else:
             # Invoke the graph without streaming
             result = self.graph.invoke(input_data, self.config)
-            print(json.dumps(result, indent=2))
+            return result
+
+    def get_state_snapshots(self) -> list:
+        """Get all the state snapshots from the graph execution."""
+
+        all_states = []
+        for state in self.graph.get_state_history(self.config):
+            all_states.append(state)
+
+        return all_states

examples/brickschema_ttl.py

@@ -1,8 +1,5 @@
 from brickllm.graphs import BrickSchemaGraph
 
-# Create an instance of BrickSchemaGraph
-brick_graph = BrickSchemaGraph()
-
 # Specify the user prompt
 building_description = """
 I have a building located in Bolzano.
@@ -13,11 +10,16 @@
 - CO sensor.
 """
 
+# Create an instance of BrickSchemaGraph
+brick_graph = BrickSchemaGraph()
+
 # Display the graph
 brick_graph.display()
 
 # Run the graph without streaming
-brick_graph.run(
+result = brick_graph.run(
     prompt=building_description,
     stream=False
-)
+)
+
+print(result)