Update inference.py

inference.py

@@ -1,66 +1,113 @@
+import asyncio
 import os
-import sys
+import textwrap
 import json
+from typing import List, Optional
 from openai import OpenAI
 
-# 1. MANDATORY START TAG
-# Updated to match the model we are actually calling
-print("[START] task=ecommerce_lifecycle_v1 env=ecommerce_bot model=Qwen2.5-72B-Instruct", flush=True)
+# Import the environment classes - ensure these files are in your root!
+# Based on your previous screenshots, you might need to adjust the import name
+try:
+    from my_env_v4 import MyEnvV4Action, MyEnvV4Env
+except ImportError:
+    # Fallback if the naming in your project is different
+    from env import MyEnvV4Action, MyEnvV4Env 
 
-def run_inference(data):
-    # 2. CONFIGURATION FROM YOUR SCREENSHOT (image_42204f.png)
-    # This ensures we hit the exact proxy the validator expects
-    api_base = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
-    
-    # In HF Spaces, the 'API_KEY' for the router is usually your HF_TOKEN secret
-    api_key = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-    model_name = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
-
-    # Initialize OpenAI client pointed at the Hugging Face Router
-    client = OpenAI(
-        base_url=api_base,
-        api_key=api_key
-    )
-
-    query = data.get("query", "")
-    order_id = data.get("order_id", "Unknown")
-    
-    # 3. LOG THE STEP - Include model name to prove we are using the correct router
-    print(f"[STEP] step=1 action=hf_router_call model={model_name} reward=0.5 done=False", flush=True)
+# 1. CONFIGURATION (Environment Variables)
+IMAGE_NAME = os.getenv("IMAGE_NAME") 
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
+MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
+
+TASK_NAME = "ecommerce_lifecycle_v1"
+BENCHMARK = "scaler_ecommerce_bot"
+MAX_STEPS = 5  # Adjusted based on your previous logic
+TEMPERATURE = 0.7
+MAX_TOKENS = 150
+
+# 2. LOGGING HELPERS (Mandatory for Scaler Stdout Parsing)
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+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()
+    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True)
 
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True)
+
+# 3. LLM INFERENCE LOGIC
+def get_model_decision(client: OpenAI, query: str, history: List[str]) -> dict:
     try:
-        # 4. THE ACTUAL PROXY CALL
-        # Note: We use the parameters (temp, max_tokens) observed in your screenshot
         completion = client.chat.completions.create(
-            model=model_name, 
+            model=MODEL_NAME,
             messages=[
-                {"role": "system", "content": "You are a professional e-commerce agent. Respond ONLY in valid JSON with fields: 'reply', 'action', and 'status'. Actions: issue_refund, update_address, escalate_to_hitl."},
-                {"role": "user", "content": f"User Query: {query} | Order ID: {order_id}"}
+                {"role": "system", "content": "You are an e-commerce agent. Respond ONLY in JSON with 'reply', 'action' (issue_refund, update_address, escalate_to_hitl), and 'status'."},
+                {"role": "user", "content": f"Query: {query}\nHistory: {history}"}
             ],
-            temperature=0.7,
-            max_tokens=150
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
         )
-        
-        # 5. PARSE & CLEAN
-        raw_content = completion.choices[0].message.content
-        # Remove any markdown backticks if the model includes them
-        clean_content = raw_content.replace("```json", "").replace("```", "").strip()
-        llm_json = json.loads(clean_content)
-
-        # 6. LOG FINAL SUCCESS
-        print(f"[STEP] step=2 action={llm_json.get('action', 'complete')} reward=1.0 done=True", flush=True)
-        print(f"[END] task=ecommerce_lifecycle_v1 score=1.0 steps=2", flush=True)
-        
-        return llm_json
-
-    except Exception as e:
-        # Detailed error logging for your HF Container tab
-        print(f"HF ROUTER FAILURE: {str(e)}", flush=True)
-        print(f"[STEP] step=1 action=error reward=0.0 done=True", flush=True)
-        print(f"[END] task=ecommerce_lifecycle_v1 score=0.0 steps=1", flush=True)
-        
-        return {
-            "reply": "The fulfillment system is temporarily unreachable via the HF Router.",
-            "action": "error",
-            "status": "failed"
-        }
+        content = completion.choices[0].message.content.strip()
+        # Clean potential markdown
+        clean_content = content.replace("```json", "").replace("```", "").strip()
+        return json.loads(clean_content)
+    except Exception as exc:
+        print(f"[DEBUG] Proxy call failed: {exc}", flush=True)
+        return {"reply": "Error", "action": "escalate_to_hitl", "status": "error"}
+
+# 4. MAIN AGENT LOOP
+async def main() -> None:
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    
+    # Initialize the containerized environment
+    env = await MyEnvV4Env.from_docker_image(IMAGE_NAME)
+
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+    history = []
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        # Reset environment and get initial observation
+        result = await env.reset()
+        current_query = result.observation.get("query", "No query provided")
+
+        for step in range(1, MAX_STEPS + 1):
+            if result.done:
+                break
+
+            # Get decision from LLM via Proxy
+            decision = get_model_decision(client, current_query, history)
+            action_str = decision.get("action", "none")
+
+            # Execute step in environment
+            result = await env.step(MyEnvV4Action(action=action_str, metadata=decision))
+            
+            reward = result.reward or 0.0
+            done = result.done
+            
+            rewards.append(reward)
+            steps_taken = step
+            log_step(step=step, action=action_str, reward=reward, done=done, error=None)
+
+            history.append(f"Step {step}: {action_str} -> reward {reward}")
+
+            if done:
+                break
+
+        # Calculate final score (normalized to 1.0)
+        score = sum(rewards) / len(rewards) if rewards else 0.0
+        success = score > 0.5
+
+    finally:
+        await env.close()
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+if __name__ == "__main__":
+    asyncio.run(main())