Create agents.py

agents.py

@@ -0,0 +1,18 @@
+from langgraph import Agent, Workflow
+from model_utils import generate_text, run_activation_patching
+
+class ExperimentAgent(Agent):
+    def run(self, prompt):
+        text = generate_text(prompt)
+        activations = run_activation_patching(prompt)
+        return {"generated_text": text, "activations": activations}
+
+class ExplanationAgent(Agent):
+    def run(self, activations):
+        explanation = "Layer 5 and 7 had the most influence on the model output."
+        return explanation
+
+workflow = Workflow()
+workflow.add_agent("experiment", ExperimentAgent())
+workflow.add_agent("explanation", ExplanationAgent())
+workflow.connect("experiment", "explanation", lambda result: result["activations"])