Spaces:
Sleeping
Sleeping
| """Baseline inference script — OpenAI API against all 3 tasks.""" | |
| import os, json, asyncio | |
| import httpx | |
| BASE_URL = os.environ.get("ENV_BASE_URL", "http://localhost:7860") | |
| OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "") | |
| OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") | |
| MODEL = os.environ.get("BASELINE_MODEL", "gpt-4o-mini") | |
| SYSTEM = """You are an expert email manager AI agent. Process emails by: | |
| 1. Correct category: urgent/normal/low/spam/internal/external | |
| 2. Accurate priority 0.0-1.0 | |
| 3. Right action: triage/respond/escalate/archive/flag/skip | |
| 4. Quality responses when needed | |
| RULES: Urgent/security/legal → escalate first. Spam → archive. Premium customers with issues → respond with empathy and resolution. Always output valid JSON only.""" | |
| async def call_llm(messages): | |
| if not OPENAI_API_KEY: | |
| return json.dumps({"action_type": "archive", "email_id": "unknown"}) | |
| async with httpx.AsyncClient(timeout=60) as c: | |
| r = await c.post(f"{OPENAI_BASE_URL}/chat/completions", | |
| headers={"Authorization": f"Bearer {OPENAI_API_KEY}", "Content-Type": "application/json"}, | |
| json={"model": MODEL, "messages": messages, "temperature": 0.1, | |
| "max_tokens": 500, "response_format": {"type": "json_object"}}) | |
| r.raise_for_status() | |
| return r.json()["choices"][0]["message"]["content"] | |
| def parse_action(text, email_id): | |
| try: | |
| a = json.loads(text) | |
| if "email_id" not in a: a["email_id"] = email_id | |
| if "action_type" not in a: a["action_type"] = "archive" | |
| return a | |
| except: return {"action_type": "archive", "email_id": email_id} | |
| def priority_heuristic(email): | |
| s = email.get("subject","").lower(); sn = email.get("sender","").lower() | |
| score = 0.5 | |
| if any(w in s for w in ["critical","p0","down","breach","security","expired","gdpr"]): score=0.95 | |
| elif any(w in s for w in ["urgent","asap","today","deadline","final","board","legal","aws bill"]): score=0.80 | |
| elif any(w in s for w in ["complaint","dispute","churn","billing","renewal","nps"]): score=0.70 | |
| elif any(w in s for w in ["newsletter","lunch","cafeteria","reminder","digest","offer"]): score=0.05 | |
| if any(w in sn for w in ["pagerduty","promo","domains-r","crypto","paypa1"]): score=0.0 | |
| score += email.get("thread_depth",0)*0.03 | |
| return min(1.0,max(0.0,score)) | |
| async def run_task(task_id): | |
| print(f"\n{'='*55}\nTask: {task_id}\n{'='*55}") | |
| if not OPENAI_API_KEY: | |
| return {"task_id": task_id, "score": 0.0, "error": "OPENAI_API_KEY not set"} | |
| async with httpx.AsyncClient(base_url=BASE_URL, timeout=30) as c: | |
| obs = (await c.post(f"/reset?task_id={task_id}")).json()["observation"] | |
| print(f"Emails: {len(obs['inbox'])} | Max steps: {obs['context']['max_steps']}") | |
| processed = set(); step = 0 | |
| while step < obs['context']['max_steps'] + 2: | |
| pending = [e for e in obs['inbox'] if e['id'] not in processed] | |
| if not pending: break | |
| email = sorted(pending, key=lambda e: priority_heuristic(e), reverse=True)[0] | |
| body = obs.get('current_email_body','') if obs.get('current_email',{}) and obs['current_email']['id']==email['id'] else email.get('preview','') | |
| msg = [{"role":"system","content":SYSTEM}, | |
| {"role":"user","content":f"Task: {obs['task_description'][:120]}\n\nEmail:\nID: {email['id']}\nSubject: {email['subject']}\nFrom: {email['sender']}\nBody: {body[:400]}\n\nOutput JSON action for email_id: {email['id']}"}] | |
| llm = await call_llm(msg); action = parse_action(llm, email['id']) | |
| processed.add(email['id']) | |
| sr = await c.post(f"/step?task_id={task_id}", json=action) | |
| if sr.status_code != 200: step+=1; continue | |
| sd = sr.json(); obs = sd['observation']; step+=1 | |
| print(f" [{email['id']}] {action.get('action_type')} cat={action.get('category')} pri={action.get('priority_score')} r={sd['reward']:+.3f}") | |
| if sd['done']: break | |
| grade = (await c.post(f"/grader?task_id={task_id}")).json() | |
| print(f"\n Score: {grade['episode_score']:.4f} | {'PASSED' if grade['passed'] else 'FAILED'} | {grade['summary']}") | |
| return {"task_id": task_id, "score": grade['episode_score'], "passed": grade['passed'], "breakdown": grade.get('breakdown',{})} | |
| async def run_baseline_all(): | |
| if not OPENAI_API_KEY: | |
| return {"task_1":{"score":0.0,"error":"No API key"},"task_2":{"score":0.0,"error":"No API key"},"task_3":{"score":0.0,"error":"No API key"},"model":MODEL,"note":"Set OPENAI_API_KEY"} | |
| results = {} | |
| for tid in ["task_1","task_2","task_3"]: | |
| try: results[tid] = await run_task(tid) | |
| except Exception as e: results[tid] = {"score":0.0,"error":str(e)} | |
| scores = [r["score"] for r in results.values() if "error" not in r] | |
| results["average_score"] = round(sum(scores)/len(scores),4) if scores else 0.0 | |
| results["model"] = MODEL | |
| print(f"\n{'='*55}\nBASELINE SUMMARY\n{'='*55}") | |
| for t in ["task_1","task_2","task_3"]: print(f" {t}: {results[t].get('score',0):.4f} {'✅' if results[t].get('passed') else '❌'}") | |
| print(f" Average: {results['average_score']:.4f} | Model: {MODEL}") | |
| return results | |
| if __name__ == "__main__": | |
| asyncio.run(run_baseline_all()) | |