Update inference.py

inference.py

@@ -1,61 +1,96 @@
+import asyncio
 import os
 import json
-import asyncio
+import textwrap
+from typing import List, Optional
 from openai import OpenAI
+from my_env_v4 import MyEnvV4Action, MyEnvV4Env
 
-# 1. MANDATORY START TAG
-# Use os.getenv only for the model name to allow for validator flexibility
+# 1. INFRASTRUCTURE CONFIGURATION
+# Prioritize injected variables to pass Phase 2 Deep Validation
+API_KEY = os.getenv("API_KEY") or os.getenv("HF_TOKEN")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-print(f"[START] task=ecommerce_lifecycle_v1 env=scaler_ecommerce_bot model={MODEL_NAME}", flush=True)
+IMAGE_NAME = os.getenv("IMAGE_NAME")
 
-def run_inference(data):
-    """Bridge for the FastAPI server"""
-    return asyncio.run(execute_agent_logic(data))
+# Task Metadata
+TASK_NAME = "ecommerce_lifecycle_v1"
+BENCHMARK = "scaler_ecommerce_bot"
+MAX_STEPS = 1 
+TEMPERATURE = 0.7
+MAX_TOKENS = 150
 
-async def execute_agent_logic(data):
-    # 2. STRICT ENVIRONMENT VARIABLE USE
-    # Do not use fallbacks like "https://router.huggingface.co/v1" here.
-    # The validator MUST inject these for the proxy check to work.
-    try:
-        api_base = os.environ["API_BASE_URL"]
-        api_key = os.environ["API_KEY"]
-    except KeyError as e:
-        # This will show up in your HF Container logs if the validator fails to inject
-        print(f"[DEBUG] Missing infrastructure variable: {e}", flush=True)
-        # Fallback to HF_TOKEN only as a last resort for local testing
-        api_base = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-        api_key = os.getenv("API_KEY") or os.getenv("HF_TOKEN")
-
-    # 3. INITIALIZE OPENAI CLIENT
-    client = OpenAI(
-        base_url=api_base,
-        api_key=api_key
-    )
-
-    query = data.get("query", "No query provided")
-    
+# 2. LOGGING HELPERS (Mandatory for Scaler 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 INTERACTION LOGIC
+def get_model_message(client: OpenAI, history: List[str]) -> str:
     try:
-        # 4. THE PROXY CALL
-        # Using the client this way ensures traffic hits the LiteLLM proxy
+        # Mandatory Proxy Call
         completion = client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
-                {"role": "system", "content": "You are an ecommerce agent. Respond in JSON: {'reply': str, 'action': str}"},
-                {"role": "user", "content": query}
+                {"role": "system", "content": "You are an e-commerce assistant. Respond with the action to take."},
+                {"role": "user", "content": "Handle this request: refund for order 101."}
             ],
-            response_format={ "type": "json_object" }
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
         )
+        return (completion.choices[0].message.content or "").strip()
+    except Exception as exc:
+        print(f"[DEBUG] Model request failed: {exc}", flush=True)
+        return "escalate_to_human"
+
+# 4. AGENT LOOP
+async def main() -> None:
+    # Initialize client using injected infrastructure variables
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    env = await MyEnvV4Env.from_docker_image(IMAGE_NAME)
+
+    history: List[str] = []
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        result = await env.reset() # Handled by server /reset
+
+        for step in range(1, MAX_STEPS + 1):
+            if result.done:
+                break
+
+            message = get_model_message(client, history)
+            result = await env.step(MyEnvV4Action(message=message))
+
+            reward = result.reward or 0.0
+            log_step(step=step, action=message, reward=reward, done=result.done, error=None)
+
+            rewards.append(reward)
+            steps_taken = step
+            history.append(f"Step {step}: {message} -> {reward}")
+
+            if result.done:
+                break
 
-        decision = json.loads(completion.choices[0].message.content)
-        action = decision.get("action", "none")
-        
-        # 5. EMIT MANDATORY LOGS FOR SCORING
-        print(f"[STEP] step=1 action={action} reward=1.00 done=true error=null", flush=True)
-        print(f"[END] success=true steps=1 score=1.00 rewards=1.00", flush=True)
+        score = sum(rewards) / len(rewards) if rewards else 0.0
+        success = score >= 1.0
 
-        return decision
+    finally:
+        await env.close()
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
-    except Exception as e:
-        print(f"[STEP] step=1 action=error reward=0.00 done=true error={str(e)}", flush=True)
-        print(f"[END] success=false steps=1 score=0.00 rewards=0.00", flush=True)
-        return {"reply": "Connection failed", "status": "error"}
+if __name__ == "__main__":
+    asyncio.run(main())