Build and Tune AI Agents in Colab with Microsoft Agent-Lightning

!pip -q install agentlightning openai nest_asyncio python-dotenv > /dev/null
import os, threading, time, asyncio, nest_asyncio, random
from getpass import getpass
from agentlightning.litagent import LitAgent
from agentlightning.trainer import Trainer
from agentlightning.server import AgentLightningServer
from agentlightning.types import PromptTemplate
import openai
if not os.getenv("OPENAI_API_KEY"):
   try:
       os.environ["OPENAI_API_KEY"] = getpass(" Enter OPENAI_API_KEY (leave blank if using a local/proxy base): ") or ""
   except Exception:
       pass
MODEL = os.getenv("MODEL", "gpt-4o-mini")

class QAAgent(LitAgent):
   def training_rollout(self, task, rollout_id, resources):
       """Given a task {'prompt':..., 'answer':...}, ask LLM using the server-provided system prompt and return a reward in [0,1]."""
       sys_prompt = resources["system_prompt"].template
       user = task["prompt"]; gold = task.get("answer","").strip().lower()
       try:
           r = openai.chat.completions.create(
               model=MODEL,
               messages=[{"role":"system","content":sys_prompt},
                         {"role":"user","content":user}],
               temperature=0.2,
           )
           pred = r.choices[0].message.content.strip()
       except Exception as e:
           pred = f"[error]{e}"
       def score(pred, gold):
           P = pred.lower()
           base = 1.0 if gold and gold in P else 0.0
           gt = set(gold.split()); pr = set(P.split());
           inter = len(gt & pr); denom = (len(gt)+len(pr)) or 1
           overlap = 2*inter/denom
           brevity = 0.2 if base==1.0 and len(P.split())<=8 else 0.0
           return max(0.0, min(1.0, 0.7*base + 0.25*overlap + brevity))
       return float(score(pred, gold))

TASKS = [
   {"prompt":"Capital of France?","answer":"Paris"},
   {"prompt":"Who wrote Pride and Prejudice?","answer":"Jane Austen"},
   {"prompt":"2+2 = ?","answer":"4"},
]
PROMPTS = [
   "You are a terse expert. Answer with only the final fact, no sentences.",
   "You are a helpful, knowledgeable AI. Prefer concise, correct answers.",
   "Answer as a rigorous evaluator; return only the canonical fact.",
   "Be a friendly tutor. Give the one-word answer if obvious."
]
nest_asyncio.apply()
HOST, PORT = "127.0.0.1", 9997

async def run_server_and_search():
   server = AgentLightningServer(host=HOST, port=PORT)
   await server.start()
   print(" Server started")
   await asyncio.sleep(1.5)
   results = []
   for sp in PROMPTS:
       await server.update_resources({"system_prompt": PromptTemplate(template=sp, engine="f-string")})
       scores = []
       for t in TASKS:
           tid = await server.queue_task(sample=t, mode="train")
           rollout = await server.poll_completed_rollout(tid, timeout=40)  # waits for a worker
           if rollout is None:
               print(" Timeout waiting for rollout; continuing...")
               continue
           scores.append(float(getattr(rollout, "final_reward", 0.0)))
       avg = sum(scores)/len(scores) if scores else 0.0
       print(f" Prompt avg: {avg:.3f}  |  {sp}")
       results.append((sp, avg))
   best = max(results, key=lambda x: x[1]) if results else ("<none>",0)
   print("n BEST PROMPT:", best[0], " | score:", f"{best[1]:.3f}")
   await server.stop()

def run_client_in_thread():
   agent = QAAgent()
   trainer = Trainer(n_workers=2)    
   trainer.fit(agent, backend=f"http://{HOST}:{PORT}")
client_thr = threading.Thread(target=run_client_in_thread, daemon=True)
client_thr.start()
asyncio.run(run_server_and_search())