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"])