Rahmath1/self_improving_agent
0
1"""2inference.py — Self-Improving EDA Agent3===================================4MANDATORY environment variables:5 API_BASE_URL The API endpoint for the LLM6 MODEL_NAME The model identifier7 HF_TOKEN Your Hugging Face API key8 9Self-improvement loop:10 1. Curriculum picks task + dataset based on current skill level11 2. Memory provides reflection from past episodes12 3. Expert provides dynamic requirements13 4. Agent runs episode with full context14 5. Results recorded → curriculum may promote → memory updated → expert evaluates15 6. Repeat with harder challenges16 17STDOUT FORMAT:18 [START] task=<task> env=eda_openenv model=<model>19 [STEP] step=<n> action=<action> reward=<0.00> done=<true|false> error=<msg|null>20 [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...>21"""22 23import os24import re25import json26import argparse27from typing import List, Optional28 29from openai import OpenAI30 31 32# from server.Advance_agent_environment import AdvanceAgentEnvironment as EDAEnv33# # TASKS, TASK_ACTION_MAP34# from models import AdvanceAgentAction as Action35# #Reward36# from pipeline import validate_action, apply_order_bonus, PIPELINE, get_completed_actions37# from Curriculum import CurriculumManager38# from Self_Agent.Advance_agent.self_imrovement.memory import MemoryBank39# from Self_Agent.Advance_agent.self_imrovement.expert import SimulatedExpert40 41 42# To this:43from self_improvement.curriculum import CurriculumManager44from self_improvement.memory import MemoryBank45from self_improvement.expert import SimulatedExpert46from self_improvement.injection_detector import InjectionDetector47 48# ─────────────────────────────────────────49# Config50# ─────────────────────────────────────────51API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")52API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")53MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")54BENCHMARK = "eda_openenv_selfplay"55 56MAX_STEPS = 1057TEMPERATURE = 0.058MAX_TOKENS = 20059SUCCESS_SCORE_THRESHOLD = 0.3560 61VALID_ACTIONS = [62 "clean_data", "eda", "feature_engineering", "train_model",63 "missing", "correlation", "insight",64]65 66SYSTEM_PROMPT = """You are a self-improving EDA agent. You have access to your past performance history and expert feedback.67 68Your goal is to learn from mistakes and improve across episodes.69 70Pipeline order (must follow strictly):711. clean_data722. eda733. feature_engineering744. train_model755. Then: task-specific action76 77Task actions:78- detect_missing → missing79- find_correlation → correlation80- generate_insight → insight81 82Read your memory and expert requirements carefully before deciding.83Respond ONLY with JSON: {"action": "<action_name>", "reason": "<one sentence>"}"""84 85 86# ─────────────────────────────────────────87# Structured logging88# ─────────────────────────────────────────89def log_start(task: str, env: str, model: str) -> None:90 print(f"[START] task={task} env={env} model={model}", flush=True)91 92def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:93 print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error or 'null'}", flush=True)94 95def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:96 rewards_str = ",".join(f"{r:.2f}" for r in rewards)97 print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)98 99 100def _safe(value) -> float:101 try:102 f = float(value) if value is not None else 0.5103 except (TypeError, ValueError):104 f = 0.5105 return round(max(0.02, min(0.98, f)), 4)106 107 108# ─────────────────────────────────────────109# Self-Improving LLM Agent110# ─────────────────────────────────────────111class SelfImprovingAgent:112 113 def __init__(self, memory: MemoryBank, expert: SimulatedExpert):114 if not API_KEY:115 print("[warn] HF_TOKEN not set", flush=True)116 self.client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "dummy")117 self.model = MODEL_NAME118 self.memory = memory119 self.expert = expert120 121 def select_action(self, obs, history: list, episode: int) -> tuple[str, str]:122 completed = get_completed_actions(history)123 next_pipeline_step = next((s for s in PIPELINE if s not in completed), "pipeline complete")124 125 # Build rich context with memory + expert126 reflection = self.memory.reflection_prompt(obs.task, 1)127 expert_context = self.expert.expert_prompt(obs.task, episode)128 129 user_message = f"""{reflection}130 131{expert_context}132 133## Current Episode Observation134Task : {obs.task}135Columns : {obs.columns}136History : {obs.history}137Completed: {completed}138Next step: {next_pipeline_step}139 140Dataset preview:141{json.dumps(obs.dataset_head[:3], indent=2)}142 143Based on your memory and expert requirements, what is the best action?"""144 145 try:146 completion = self.client.chat.completions.create(147 model=self.model,148 temperature=TEMPERATURE,149 max_tokens=MAX_TOKENS,150 stream=False,151 messages=[152 {"role": "system", "content": SYSTEM_PROMPT},153 {"role": "user", "content": user_message},154 ],155 )156 raw = completion.choices[0].message.content or ""157 except Exception as exc:158 print(f"[DEBUG] Model error: {exc}", flush=True)159 if next_pipeline_step != "pipeline complete":160 return next_pipeline_step, "fallback"161 return TASK_ACTION_MAP.get(obs.task, "missing"), "fallback"162 163 try:164 clean = re.sub(r"```json|```", "", raw).strip()165 parsed = json.loads(clean)166 action = parsed.get("action", "").strip()167 reason = parsed.get("reason", "")168 if action not in VALID_ACTIONS:169 action = next_pipeline_step if next_pipeline_step != "pipeline complete" else TASK_ACTION_MAP.get(obs.task, "missing")170 reason = "fallback"171 except json.JSONDecodeError:172 action = next_pipeline_step if next_pipeline_step != "pipeline complete" else TASK_ACTION_MAP.get(obs.task, "missing")173 reason = "fallback"174 175 return action, reason176 177 178# ─────────────────────────────────────────179# Episode Runner180# ─────────────────────────────────────────181def run_episode(env: EDAEnv, agent: SelfImprovingAgent,182 task_override: dict, episode: int) -> dict:183 obs = env.reset()184 env._task = task_override.copy()185 obs = env._get_obs()186 187 history = []188 rewards = []189 step = 0190 done = False191 success = False192 193 log_start(task=obs.task, env=BENCHMARK, model=MODEL_NAME)194 195 try:196 for step_num in range(1, MAX_STEPS + 1):197 if done:198 break199 200 action_type, reason = agent.select_action(obs, history, episode)201 penalty = validate_action(action_type, history)202 203 if penalty:204 raw_reward = _safe(penalty.score)205 error = "out-of-order"206 else:207 action = Action(action_type=action_type)208 obs = env.step(action)209 done = obs.done210 r = apply_order_bonus(action_type, history,211 Reward(score=_safe(obs.reward), feedback="", is_penalty=False))212 raw_reward = _safe(r.score)213 error = None214 215 rewards.append(raw_reward)216 step = step_num217 log_step(step=step_num, action=action_type, reward=raw_reward, done=done, error=error)218 219 history.append({220 "action": action_type,221 "reward": raw_reward,222 "is_penalty": penalty is not None,223 "done": done,224 })225 226 score = _safe(sum(rewards) / len(rewards)) if rewards else 0.5227 success = score >= SUCCESS_SCORE_THRESHOLD228 229 finally:230 log_end(success=success, steps=step, score=score, rewards=rewards)231 232 return {233 "task": env._task["name"],234 "score": score,235 "steps": step,236 "success": success,237 "rewards": rewards,238 "penalties": sum(1 for h in history if h["is_penalty"]),239 "actions": [h["action"] for h in history],240 }241 242 243# ─────────────────────────────────────────244# Main Self-Play Loop245# ─────────────────────────────────────────246def main() -> None:247 parser = argparse.ArgumentParser(description="EDA OpenEnv — Self-Improving Agent")248 parser.add_argument("--episodes", type=int, default=9, help="Total episodes to run (default: 9 = 3 per task)")249 parser.add_argument("--steps", type=int, default=10, help="Max steps per episode")250 parser.add_argument("--reset", action="store_true", help="Reset curriculum and memory")251 args = parser.parse_args()252 253 # Init self-improvement components254 curriculum = CurriculumManager()255 memory = MemoryBank()256 expert = SimulatedExpert()257 258 if args.reset:259 curriculum.reset()260 memory.clear()261 expert.reset()262 print("[INFO] Reset curriculum, memory and expert state.", flush=True)263 264 # Init env and agent265 env = EDAEnv(max_steps=args.steps)266 agent = SelfImprovingAgent(memory=memory, expert=expert)267 268 print(f"\n{'═'*60}", flush=True)269 print(f" EDA OpenEnv — Self-Improving Agent", flush=True)270 print(f" Model : {MODEL_NAME}", flush=True)271 print(f" Episodes: {args.episodes}", flush=True)272 print(f" Start : {curriculum.current_level_name()}", flush=True)273 print(f"{'═'*60}\n", flush=True)274 275 all_results = []276 episode_num = 0277 278 # Run episodes cycling through all 3 tasks279 for i in range(args.episodes):280 task_dict = TASKS[i % len(TASKS)]281 episode_num += 1282 283 # Update dataset based on curriculum level284 df = curriculum.next_dataset()285 env.df = df286 287 print(f"\n{'─'*60}", flush=True)288 print(f" Episode {episode_num} | {curriculum.current_level_name()} | Task: {task_dict['name']}", flush=True)289 print(f" {curriculum.summary()}", flush=True)290 print(f"{'─'*60}", flush=True)291 292 # Run episode293 result = run_episode(env, agent, task_dict, episode_num)294 all_results.append(result)295 296 # Get expert evaluation297 expert_fb = expert.evaluate(298 episode = episode_num,299 task = result["task"],300 actions = result["actions"],301 score = result["score"],302 df = df,303 )304 305 # Store in memory with expert feedback306 promoted = curriculum.record(307 score = result["score"],308 steps = result["steps"],309 penalties= result["penalties"],310 task = result["task"],311 )312 313 memory.store(314 episode = episode_num,315 task = result["task"],316 level = curriculum.state.current_level,317 score = result["score"],318 steps = result["steps"],319 penalties = result["penalties"],320 actions = result["actions"],321 expert_feedback= expert_fb.feedback,322 )323 324 # Print episode summary325 icon = "🎉" if promoted else ("✅" if result["success"] else "📉")326 print(f"\n {icon} Score: {result['score']:.4f} | "327 f"Expert: {expert_fb.feedback[:60]}...", flush=True)328 if promoted:329 print(f" 🚀 PROMOTED to {curriculum.current_level_name()}!", flush=True)330 331 # ── Final Summary ─────────────────────────────────────────332 print(f"\n{'═'*60}", flush=True)333 print(f" SELF-IMPROVEMENT SUMMARY", flush=True)334 print(f"{'═'*60}", flush=True)335 print(f" Final Level : {curriculum.current_level_name()}", flush=True)336 print(f" Final Avg : {curriculum.rolling_avg():.4f}", flush=True)337 338 scores_by_episode = [r["score"] for r in all_results]339 first_3_avg = sum(scores_by_episode[:3]) / 3 if len(scores_by_episode) >= 3 else 0340 last_3_avg = sum(scores_by_episode[-3:]) / 3 if len(scores_by_episode) >= 3 else 0341 improvement = last_3_avg - first_3_avg342 343 print(f" First 3 avg : {first_3_avg:.4f}", flush=True)344 print(f" Last 3 avg : {last_3_avg:.4f}", flush=True)345 print(f" Improvement : {improvement:+.4f} {'📈' if improvement > 0 else '📉'}", flush=True)346 print(f"{'═'*60}\n", flush=True)347 348 # Save results349 output = {350 "model": MODEL_NAME,351 "total_episodes": episode_num,352 "final_level": curriculum.state.current_level,353 "final_level_name": curriculum.current_level_name(),354 "first_3_avg": round(first_3_avg, 4),355 "last_3_avg": round(last_3_avg, 4),356 "improvement": round(improvement, 4),357 "episode_scores": scores_by_episode,358 "curriculum_stats": curriculum.get_stats(),359 }360 with open("selfplay_results.json", "w") as f:361 json.dump(output, f, indent=2)362 print("Results saved → selfplay_results.json", flush=True)363 364 365if __name__ == "__main__":366 main()367 368 369# ─────────────────────────────────────────370# Run with injection detection + comparison371# ─────────────────────────────────────────372def run_with_injection_comparison(episodes_per_phase: int = 6, steps: int = 10):373 """374 Runs two phases:375 Phase 1 (Base) — no fine-tuning, tracks injection performance376 Phase 2 (Fine-tuned) — after curriculum learning, compare metrics377 378 Produces the BASE vs FINE-TUNED comparison table.379 """380 from injection_detector import InjectionDetector381 382 curriculum = CurriculumManager("curriculum_compare.json")383 memory = MemoryBank("memory_compare.json")384 expert = SimulatedExpert("expert_compare.json")385 detector = InjectionDetector()386 387 env = EDAEnv(max_steps=steps)388 agent = SelfImprovingAgent(memory=memory, expert=expert)389 390 print(f"\n{'═'*60}", flush=True)391 print(f" BASE vs FINE-TUNED Evaluation with Injection Detection", flush=True)392 print(f"{'═'*60}\n", flush=True)393 394 def run_phase(phase_name: str, mode: str, n_episodes: int) -> float:395 """Run a phase and return average reward."""396 detector.set_mode(mode)397 phase_scores = []398 399 print(f"\n{'─'*60}", flush=True)400 print(f" Phase: {phase_name}", flush=True)401 print(f"{'─'*60}", flush=True)402 403 for i in range(n_episodes):404 task_dict = TASKS[i % len(TASKS)]405 episode_num = curriculum.state.episode_count + 1406 407 # Get curriculum dataset and inject adversarial content408 df = curriculum.next_dataset()409 df, injected_at = detector.inject_dataset(df, n_injections=1)410 env.df = df411 412 # Scan dataset for injections BEFORE running413 injection_events = detector.scan_dataset(df, episode_num)414 if injection_events:415 print(f" 🛡️ Detected {len(injection_events)} injection(s) in dataset", flush=True)416 417 result = run_episode(env, agent, task_dict, episode_num)418 phase_scores.append(result["score"])419 420 # Record metrics421 expert_fb = expert.evaluate(422 episode=episode_num, task=result["task"],423 actions=result["actions"], score=result["score"], df=df,424 )425 curriculum.record(426 score=result["score"], steps=result["steps"],427 penalties=result["penalties"], task=result["task"],428 )429 memory.store(430 episode=episode_num, task=result["task"],431 level=curriculum.state.current_level,432 score=result["score"], steps=result["steps"],433 penalties=result["penalties"], actions=result["actions"],434 expert_feedback=expert_fb.feedback,435 )436 437 print(f" Ep {episode_num} | {task_dict['name']:<22} | "438 f"score={result['score']:.4f} | "439 f"injections_caught={len(injection_events)}", flush=True)440 441 return round(sum(phase_scores) / len(phase_scores), 4) if phase_scores else 0.5442 443 # Run base phase444 base_avg = run_phase("BASE MODEL", "base", episodes_per_phase)445 446 # Run fine-tuned phase (after curriculum learning)447 ft_avg = run_phase("FINE-TUNED MODEL", "finetuned", episodes_per_phase)448 449 # Print comparison report450 report = detector.comparison_report(base_avg, ft_avg)451 print(f"\n{report}", flush=True)452 453 # Save full comparison454 stats = detector.get_stats_dict(base_avg, ft_avg)455 with open("comparison_results.json", "w") as f:456 json.dump(stats, f, indent=2)457 print("\nComparison saved → comparison_results.json", flush=True)458 459 return stats460 461 462if __name__ == "__main__" and False:463 # Run comparison mode directly464 run_with_injection_comparison()