hjgfg

agent.py

@@ -264,8 +264,31 @@ def build_graph(provider: str = "huggingface"):
     builder.add_edge("tools", "assistant")
 
     # Compile graph
-    return builder.compile()
+    compiled_graph = builder.compile() # This line should already be there or be the next line
 
+    # --- START: Add this visualization code ---
+    try:
+        print("Attempting to generate graph visualization...")
+        image_filename = "langgraph_state_diagram.png"
+        # Using draw_mermaid_png as it's often more robust
+        image_bytes = compiled_graph.get_graph().draw_mermaid_png()
+        with open(image_filename, "wb") as f:
+            f.write(image_bytes)
+        print(f"SUCCESS: Graph visualization saved to '{image_filename}'")
+
+    except ImportError as e:
+        print(f"WARNING: Could not generate graph image due to missing package: {e}. "
+              "Ensure 'pygraphviz' and 'graphviz' (system) are installed, or Mermaid components are available.")
+    except Exception as e:
+        print(f"WARNING: An error occurred while generating the graph image: {e}")
+        try:
+            print("\nGraph (DOT format as fallback):\n", compiled_graph.get_graph().to_string())
+        except Exception as dot_e:
+            print(f"Could not even get DOT string: {dot_e}")
+    # --- END: Visualization code ---
+
+    return compiled_graph # This should be the last line of the function
+    
 # test
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"

app.py

@@ -24,7 +24,6 @@ class BasicAgent:
         self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
-        print(f"**--**")
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         # Wrap the question in a HumanMessage from langchain_core
         messages = [HumanMessage(content=question)]