def executor_agent(step: Dict[str, Any], context: Dict[str, Any]) -> StepResult:
   step_id = int(step.get("id", 0))
   title = step.get("title", f"Step {step_id}")
   tool = step.get("tool", "llm")


   ctx_compact = {
       "goal": context.get("goal"),
       "assumptions": context.get("assumptions", []),
       "prior_results": [
           {"step_id": r.step_id, "title": r.title, "tool": r.tool, "output": r.output[:1500]}
           for r in context.get("results", [])
       ],
   }


   if tool == "python":
       code = llm_chat(
           EXECUTOR_SYSTEM,
           user=(
               f"Step:\n{json.dumps(step, indent=2)}\n\n"
               f"Context:\n{json.dumps(ctx_compact, indent=2)}\n\n"
               f"Write Python code that completes the step. Output ONLY code."
           ),
           max_new_tokens=700,
           temperature=0.2,
       )
       py = run_python(code)
       out = []
       out.append("PYTHON_CODE:\n" + code)
       out.append("\nEXECUTION_OK: " + str(py["ok"]))
       if py["stdout"]:
           out.append("\nSTDOUT:\n" + py["stdout"])
       if py["error"]:
           out.append("\nERROR:\n" + py["error"])
       return StepResult(step_id=step_id, title=title, tool=tool, output="\n".join(out))


   result_text = llm_chat(
       EXECUTOR_SYSTEM,
       user=(
           f"Step:\n{json.dumps(step, indent=2)}\n\n"
           f"Context:\n{json.dumps(ctx_compact, indent=2)}\n\n"
           f"Return the step result."
       ),
       max_new_tokens=700,
       temperature=0.3,
   )
   return StepResult(step_id=step_id, title=title, tool=tool, output=result_text)




def aggregator_agent(task: str, plan: Dict[str, Any], results: List[StepResult]) -> str:
   payload = {
       "task": task,
       "plan": plan,
       "results": [{"step_id": r.step_id, "title": r.title, "tool": r.tool, "output": r.output[:2500]} for r in results],
   }
   return llm_chat(
       AGGREGATOR_SYSTEM,
       user=f"Combine everything into the final answer.\n\nINPUT:\n{json.dumps(payload, indent=2)}",
       max_new_tokens=900,
       temperature=0.2,
   )




def run_hierarchical_agent(task: str, verbose: bool = True) -> Dict[str, Any]:
   plan = planner_agent(task)


   if verbose:
       print("\n====================")
       print("PLAN (from Planner)")
       print("====================")
       print(json.dumps(plan, indent=2))


   context = {
       "goal": plan.get("goal", task),
       "assumptions": plan.get("assumptions", []),
       "results": [],
   }


   results: List[StepResult] = []
   for step in plan.get("steps", []):
       res = executor_agent(step, context)
       results.append(res)
       context["results"].append(res)


       if verbose:
           print("\n--------------------")
           print(f"STEP {res.step_id}: {res.title}  [tool={res.tool}]")
           print("--------------------")
           print(res.output)


   final = aggregator_agent(task, plan, results)
   if verbose:
       print("\n====================")
       print("FINAL (from Aggregator)")
       print("====================")
       print(final)


   return {"task": task, "plan": plan, "results": results, "final": final}




demo_task = """
Create a practical checklist to launch a small multi-agent system in Python for coordinating logistics:
- One planner agent that decomposes tasks
- Two executor agents (routing + inventory)
- A simple memory store for past decisions
Keep it lightweight and runnable in Colab.
"""


_ = run_hierarchical_agent(demo_task, verbose=True)


print("\n\nType your own task (or press Enter to skip):")
user_task = input().strip()
if user_task:
   _ = run_hierarchical_agent(user_task, verbose=True)