final

inference.py

@@ -1,46 +1,97 @@
 import os
+import sys
 import pandas as pd
+import httpx
 from openai import OpenAI
 from env.environment import DataCleaningEnv
 
 # ------------------ ENV VARIABLES ------------------
-API_BASE_URL = os.getenv("API_BASE_URL")
-API_KEY = os.getenv("API_KEY")           # ✅ their injected key
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
 
-# ------------------ DEBUG: print env at startup ------------------
-print(f"[DEBUG] API_BASE_URL={API_BASE_URL}")
-print(f"[DEBUG] API_KEY={'set' if API_KEY else 'MISSING'}")
-print(f"[DEBUG] MODEL_NAME={MODEL_NAME}")
-
 # ------------------ OPENAI CLIENT ------------------
-# ✅ Fail hard if env vars are missing — don't silently skip
-if not API_BASE_URL:
-    raise RuntimeError("API_BASE_URL environment variable is not set!")
-if not API_KEY:
-    raise RuntimeError("API_KEY environment variable is not set!")
+client = None
+
+try:
+    if not HF_TOKEN:
+        raise RuntimeError("API_KEY or HF_TOKEN environment variable is not set")
 
-client = OpenAI(
-    base_url=API_BASE_URL,
-    api_key=API_KEY
-)
+    client = OpenAI(
+        base_url=API_BASE_URL,
+        api_key=HF_TOKEN,
+        http_client=httpx.Client(timeout=30.0, trust_env=False)
+    )
+except Exception as e:
+    print(f"[WARN] OpenAI client init failed: {e}", file=sys.stderr)
+    client = None
 
 MAX_STEPS = 6
 
 # ------------------ LOGGING ------------------
+def clean_line(value):
+    return str(value).replace("\r", " ").replace("\n", " ")
+
 def log_start(task, env, model):
-    print(f"[START] task={task} env={env} model={model}")
+    print(f"[START] task={clean_line(task)} env={clean_line(env)} model={clean_line(model)}")
 
 def log_step(step, action, reward, done, error):
-    error_val = error if error else "null"
-    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}")
+    error_val = clean_line(error) if error else "null"
+    print(f"[STEP] step={step} action={clean_line(action)} reward={reward:.2f} done={str(done).lower()} error={error_val}")
 
-def log_end(success, steps, score, rewards):
+def log_end(success, steps, rewards, score=None):
     rewards_str = ",".join(f"{r:.2f}" for r in rewards)
-    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}")
+
+    if score is None:
+        print(f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str}")
+    else:
+        print(f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards={rewards_str}")
 
 # ------------------ LLM DECISION ------------------
+def fallback_action(history):
+    actions = [
+        {"type": "fill_nulls", "column": "city"},
+        {"type": "deduplicate", "column": "customer_id"},
+        {"type": "convert_types", "column": "age"},
+        {"type": "trim_whitespace", "column": "city"},
+        {"type": "normalize", "column": "income"},
+        {"type": "remove_nulls", "column": "customer_id"},
+    ]
+
+    for action in actions:
+        if str(action) not in history:
+            return action
+
+    return {"type": "deduplicate", "column": "customer_id"}
+
+def normalize_action(action, df, history):
+    valid_actions = {
+        "fill_nulls",
+        "remove_nulls",
+        "deduplicate",
+        "convert_types",
+        "trim_whitespace",
+        "normalize",
+    }
+
+    if not isinstance(action, dict):
+        return fallback_action(history)
+
+    action_type = action.get("type")
+    column = action.get("column")
+
+    if action_type not in valid_actions:
+        return fallback_action(history)
+
+    if action_type != "deduplicate" and column not in df.columns:
+        return fallback_action(history)
+
+    return action
+
 def get_action_from_llm(dataset, history):
+    if client is None:
+        return fallback_action(history)
+
     prompt = f"""
 You are an intelligent data cleaning agent.
 
@@ -57,78 +108,107 @@ Return ONLY:
 action_type,column_name
 """
 
-    # ✅ No try/except — let it raise so we know if LLM call fails
-    response = client.chat.completions.create(
-        model=MODEL_NAME,
-        messages=[{"role": "user", "content": prompt}],
-        temperature=0.3,
-        max_tokens=50
-    )
+    try:
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.3,
+            max_tokens=50
+        )
 
-    print(f"[DEBUG] LLM raw response: {response}")   # ✅ confirm call went through
+        output = response.choices[0].message.content.strip()
 
-    output = response.choices[0].message.content.strip()
-    print(f"[DEBUG] LLM output: {output}")
+        first_line = output.splitlines()[0].strip()
+        parts = first_line.split(",", 1)
 
-    parts = output.split(",")
-    if len(parts) < 2:
-        return {"type": "deduplicate", "column": "customer_id"}
+        if len(parts) < 2:
+            return fallback_action(history)
 
-    action_type, column = parts[0], parts[1]
-    return {"type": action_type.strip(), "column": column.strip()}
+        action_type, column = parts[0], parts[1]
+        return {"type": action_type.strip(), "column": column.strip()}
+
+    except Exception as e:
+        print(f"[WARN] LLM call failed: {e}", file=sys.stderr)
+        return fallback_action(history)
 
 # ------------------ MAIN ------------------
 def main():
-    env = DataCleaningEnv(task=1)
-    obs = env.reset()
-
+    env = None
     rewards = []
     steps_taken = 0
     history = []
+    success = False
+    score = None
 
     log_start("task1", "data_cleaning", MODEL_NAME)
 
-    for step in range(1, MAX_STEPS + 1):
+    try:
+        env = DataCleaningEnv(task=1)
+        obs = env.reset()
 
-        df = pd.DataFrame(obs["dataset"])
-        col_info = {}
+        for step in range(1, MAX_STEPS + 1):
 
-        for col in df.columns:
-            col_info[col] = {
-                "nulls": float(df[col].isnull().mean()),
-                "dtype": str(df[col].dtype),
-                "unique": int(df[col].nunique())
-            }
+            df = pd.DataFrame(obs["dataset"])
+            col_info = {}
 
-        summary = f"""
+            for col in df.columns:
+                col_info[col] = {
+                    "nulls": float(df[col].isnull().mean()),
+                    "dtype": str(df[col].dtype),
+                    "unique": int(df[col].nunique())
+                }
+
+            summary = f"""
 Columns: {list(df.columns)}
 Column Info: {col_info}
 Duplicates: {df.duplicated().sum()}
 Sample: {df.head(3).to_dict()}
 """
 
-        action = get_action_from_llm(summary, history)
+            action = get_action_from_llm(summary, history)
+            action = normalize_action(action, df, history)
+
+            if str(action) in history:
+                action = fallback_action(history)
+                action = normalize_action(action, df, history)
+
+            history.append(str(action))
+
+            error = None
+
+            try:
+                obs, reward, done, _ = env.step(action)
+            except Exception as e:
+                error = str(e)
+                reward = 0.0
+                done = True
 
-        if str(action) in history:
-            action = {"type": "deduplicate", "column": "customer_id"}
+            rewards.append(reward)
+            steps_taken = step
 
-        history.append(str(action))
+            log_step(step, str(action), reward, done, error)
 
-        obs, reward, done, _ = env.step(action)
+            if done:
+                break
 
-        rewards.append(reward)
-        steps_taken = step
+        final = env.submit_cleaned_data(env.dirty_df)
+        score = final["final_score"]
+        success = score > 0.3
 
-        log_step(step, str(action), reward, done, None)
+    except Exception as e:
+        print(f"[WARN] inference failed: {e}", file=sys.stderr)
+        success = False
+        score = None
 
-        if done:
-            break
+    finally:
+        if env is not None and hasattr(env, "close"):
+            try:
+                env.close()
+            except Exception as e:
+                print(f"[WARN] env.close failed: {e}", file=sys.stderr)
 
-    final = env.submit_cleaned_data(env.dirty_df)
-    score = final["final_score"]
-    success = score > 0.3
-    log_end(success, steps_taken, score, rewards)
+        log_end(success, steps_taken, rewards, score)
 
 # ------------------ RUN ------------------
 if __name__ == "__main__":
-    main()   # ✅ removed try/except so real errors surface
+    main()