debug: (inference) yeah, given up fully part 5

Author: vishvam-cstk

inference.py

@@ -9,9 +9,6 @@
 from dotenv import load_dotenv
 from openai import OpenAI
 
-# ----------------------------
-# PATH SETUP
-# ----------------------------
 ROOT = Path(__file__).resolve().parent
 REPO_ROOT = ROOT.parent
 
@@ -20,9 +17,6 @@
 
 load_dotenv()
 
-# ----------------------------
-# CONFIG
-# ----------------------------
 TASK_DIR = ROOT / "tasks"
 SYSTEM_PROMPT = (ROOT / "system_prompt.txt").read_text().strip()
 
@@ -31,7 +25,7 @@
 
 BASE_URL = "https://vishvam10-cloudenv.hf.space"
 
-TASK_ORDER = ["easy-task"]
+TASK_ORDER = ["easy-task", "medium-task", "hard-task"]
 
 MAX_PARSE_RETRIES = 3
 MAX_STEPS = 5
@@ -46,7 +40,6 @@ def log_start(task: str, env: str, model: str) -> None:
 def log_step(
     step: int, action: str, reward: float, done: bool, error: Optional[str]
 ) -> None:
-
     error_val = error if error else "null"
     done_val = str(done).lower()
     action_safe = action.replace("\n", " ").replace("\r", "")[:120]
@@ -68,32 +61,38 @@ def log_end(
     )
 
 
-# ----------------------------
-# LOAD TASK
-# ----------------------------
 def load_task(name: str):
     return json.loads((TASK_DIR / f"{name}.json").read_text())
 
 
-# ----------------------------
-# RESET
-# ----------------------------
-async def env_reset(task_id: str, episode_id: str):
+async def register_pipeline(task: dict):
     async with httpx.AsyncClient(timeout=30) as http:
         r = await http.post(
-            f"{BASE_URL}/reset",
+            f"{BASE_URL}/api/pipelines",
             json={
-                "task_id": task_id,
-                "episode_id": episode_id,
+                "pipeline_id": task["pipeline_id"],
+                "broken_pipeline": task["broken_pipeline"],
+                "ideal_pipeline": task["ideal_pipeline"],
             },
         )
+
+        if r.status_code not in (200, 201, 400):
+            print(
+                f"[PIPELINE REGISTER ERROR] {r.status_code} {r.text}",
+                flush=True,
+            )
+
+
+async def env_reset(task_id: str, episode_id: str):
+    async with httpx.AsyncClient(timeout=30) as http:
+        r = await http.post(
+            f"{BASE_URL}/reset",
+            json={"task_id": task_id, "episode_id": episode_id},
+        )
         r.raise_for_status()
         return r.json()
 
 
-# ----------------------------
-# STEP
-# ----------------------------
 async def env_step(action: dict):
     async with httpx.AsyncClient(timeout=30) as http:
         r = await http.post(
@@ -104,9 +103,6 @@ async def env_step(action: dict):
         return r.json()
 
 
-# ----------------------------
-# SAFE JSON CLEANING
-# ----------------------------
 def clean_json(raw: str) -> str:
     raw = raw.strip()
     raw = re.sub(r"^```(?:json)?\s*", "", raw)
@@ -119,9 +115,6 @@ def clean_json(raw: str) -> str:
     return raw
 
 
-# ----------------------------
-# BUILD ACTION (DICT ONLY)
-# ----------------------------
 def build_action(raw: str, task_id: str, episode_id: str):
     try:
         raw = clean_json(raw)
@@ -162,9 +155,6 @@ def build_action(raw: str, task_id: str, episode_id: str):
         return None
 
 
-# ----------------------------
-# LLM CALL
-# ----------------------------
 def call_model(step, state, broken, ideal, history):
     prompt = f"""
 Return ONLY valid JSON.
@@ -189,15 +179,14 @@ def call_model(step, state, broken, ideal, history):
     return (res.choices[0].message.content or "").strip()
 
 
-# ----------------------------
-# EPISODE LOOP
-# ----------------------------
 async def run_episode(task):
     task_id = task["pipeline_id"]
     episode_id = f"{task_id}-episode-1"
 
     log_start(task=task_id, env="http", model=MODEL_NAME)
 
+    result = await env_reset(task_id, episode_id)
+
     await env_step(
         {
             "task_id": task_id,
@@ -208,15 +197,14 @@ async def run_episode(task):
             "payload": {},
         }
     )
-    result = await env_reset(task_id, episode_id)
-    
-    obs = result["observation"]
 
+    await register_pipeline(task)
+
+    obs = {}
     history = []
     rewards = []
     step_count = 0
 
-
     for step in range(MAX_STEPS):
         if obs.get("done"):
             break
@@ -226,7 +214,7 @@ async def run_episode(task):
         try:
             raw = call_model(
                 step,
-                obs["current_pipeline_state"],
+                obs,
                 task["broken_pipeline"],
                 task["ideal_pipeline"],
                 history,
@@ -260,36 +248,18 @@ async def run_episode(task):
         except Exception as e:
             error = str(e)
 
-            log_step(
-                step=step,
-                action="ERROR",
-                reward=0.0,
-                done=False,
-                error=error,
-            )
-
+            log_step(step, "ERROR", 0.0, False, error)
             break
 
         if obs.get("done"):
             break
 
-    # ----------------------------
-    # END LOG
-    # ----------------------------
     success = obs.get("done", False)
     score = sum(rewards) / (len(rewards) or 1)
 
-    log_end(
-        success=success,
-        steps=step_count,
-        score=score,
-        rewards=rewards,
-    )
+    log_end(success=success, steps=step_count, score=score, rewards=rewards)
 
 
-# ----------------------------
-# MAIN
-# ----------------------------
 async def main():
     for name in TASK_ORDER:
         task = load_task(name)