exec-ops-env / inference.py
saaheerpurav's picture
first commit
1bb18a0
Raw
History Blame Contribute Delete
9.39 kB
"""
inference.py - AI Executive Operations Manager
Runs all 3 tasks using an LLM agent and prints structured logs.
Environment variables:
API_BASE_URL - OpenAI-compatible API base URL
MODEL_NAME - Model to use (e.g., "gpt-4o-mini", "meta-llama/Llama-3.1-8B-Instruct")
HF_TOKEN - API key / HuggingFace token
STDOUT format (strict):
[START] task=<task_id> env=exec-ops model=<model>
[STEP] step=<n> action=<type>('<id>') reward=<0.00> done=<true|false> error=<msg|null>
[END] success=<true|false> steps=<n> score=<0.00> rewards=<r1,r2,...>
"""
import os
import json
import sys
import dotenv
# Force UTF-8 output on Windows
if sys.stdout.encoding != "utf-8":
sys.stdout.reconfigure(encoding="utf-8")
from openai import OpenAI
from env import ExecOpsEnv, Action
from env.grader import grade
# -------------------------------------------------------
# Configuration from environment
# -------------------------------------------------------
dotenv.load_dotenv()
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
SYSTEM_PROMPT = """You are an expert AI assistant helping a startup CEO (Alex Rivera at NovaTech AI) manage their inbox efficiently.
Your job: Given the current environment state, choose ONE action for ONE email.
Triage principles:
1. Handle the highest PRIORITY * URGENCY emails first
2. Reply to items requiring personal CEO attention (investor relations, critical incidents, legal deadlines)
3. Delegate items that can be handled by team members (HR, finance, operations)
4. Schedule items requiring a meeting
5. Ignore truly trivial items only when critical items remain
You MUST respond with ONLY valid JSON, no explanation:
{"type": "reply|schedule|delegate|ignore", "email_id": "<id>"}
IMPORTANT: email_id MUST be one of the ids from the "unhandled_emails" list (e.g. "e1", "e2").
Do NOT use goal_id values (e.g. "g1", "g2") as the email_id.
Choose wisely - you have limited steps. Ignoring P4-P5 items is heavily penalized."""
def _log(msg: str) -> None:
"""Print a structured log line to stdout."""
print(msg, flush=True)
def _err(msg: str) -> None:
"""Print diagnostic info to stderr - never pollutes stdout."""
print(msg, file=sys.stderr, flush=True)
def _clean_error(msg: str) -> str:
"""Collapse error message to a single line safe for the log format."""
return msg.replace("\n", " ").replace("\r", "").strip()[:120]
def parse_action_from_response(text: str) -> dict:
"""Extract JSON action from LLM response, handling markdown code blocks."""
text = text.strip()
if "```json" in text:
text = text.split("```json")[1].split("```")[0].strip()
elif "```" in text:
parts = text.split("```")
for part in parts[1::2]:
part = part.strip()
if part.startswith("{"):
text = part
break
start = text.find("{")
end = text.rfind("}") + 1
if start >= 0 and end > start:
text = text[start:end]
return json.loads(text)
def get_fallback_action(obs: dict) -> Action:
"""Fallback: pick highest priority * urgency unhandled email and reply or delegate."""
inbox = obs.get("inbox", [])
unhandled = [e for e in inbox if not e.get("handled", False)]
if not unhandled:
return None
best = max(unhandled, key=lambda e: e.get("priority", 1) * e.get("urgency", 0.5))
action_type = "delegate" if best.get("priority", 1) <= 2 else "reply"
return Action(type=action_type, email_id=best["id"])
def run_task(task_id: str) -> float:
"""Run a single task with the LLM agent. Returns final grade in [0, 1]."""
_log(f"[START] task={task_id} env=exec-ops model={MODEL_NAME}")
step_count = 0
rewards: list[float] = []
pending_step_log = None # deferred so we can force done=true on the last step
final_score = 0.0
env = None
try:
env = ExecOpsEnv(task_id)
obs = env.reset()
max_steps = obs.get("max_steps", 10)
while not obs.get("done", False) and step_count < max_steps:
unhandled = [e for e in obs.get("inbox", []) if not e.get("handled", False)]
if not unhandled:
break
# Flush previous buffered step - it is not the last, so done=false
if pending_step_log is not None:
_log(pending_step_log.replace("_DONE_", "false"))
pending_step_log = None
prompt_state = {
"time": obs.get("time"),
"steps_remaining": obs.get("steps_remaining", max_steps - step_count),
"unhandled_emails": [
{
"id": e["id"],
"sender": e["sender"].split("<")[0].strip(),
"subject": e["subject"],
"priority": e["priority"],
"urgency": round(e["urgency"], 2),
}
for e in sorted(unhandled, key=lambda x: -(x["priority"] * x["urgency"]))
],
"pending_goals": [
{"goal_id": g["id"], "description": g["description"], "priority": g["priority"]}
for g in obs.get("pending_goals", [])
],
}
# --- Get action from LLM (with fallback) ---
step_error = None
action = None
try:
response = client.chat.completions.create(
model=MODEL_NAME,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": f"Current state:\n{json.dumps(prompt_state, indent=2)}\n\nChoose your action (JSON only):"}
],
temperature=0.1,
max_tokens=80,
)
action_text = response.choices[0].message.content
action_data = parse_action_from_response(action_text)
action = Action(type=action_data["type"], email_id=action_data["email_id"])
valid_ids = {e["id"] for e in unhandled}
if action.email_id not in valid_ids:
raise ValueError(f"Invalid email_id: {action.email_id}")
except Exception as e:
step_error = _clean_error(str(e))
_err(f"[fallback] LLM error ({task_id} step {step_count+1}): {type(e).__name__}: {e}")
action = get_fallback_action(obs)
if action is None:
break
# --- Execute step; retry with fallback on failure ---
candidates = [action]
fb = get_fallback_action(obs)
if fb is not None and (action is None or fb.email_id != action.email_id):
candidates.append(fb)
executed = False
for attempt in candidates:
try:
result = env.step(attempt)
reward = result.get("reward", 0.0)
done = result.get("done", False)
obs = result.get("observation", result)
if "done" not in obs:
obs["done"] = done
step_count += 1
rewards.append(reward)
action_str = f"{attempt.type}('{attempt.email_id}')"
error_field = step_error if step_error else "null"
pending_step_log = (
f"[STEP] step={step_count} action={action_str} "
f"reward={reward:.2f} done=_DONE_ error={error_field}"
)
executed = True
break
except Exception as e:
step_error = _clean_error(str(e))
_err(f"[fallback] Step error ({task_id} step {step_count+1}): {type(e).__name__}: {e}")
if not executed:
break
except Exception as e:
_err(f"[error] Unexpected error in task '{task_id}': {e}")
finally:
# Flush the final buffered step - always done=true
if pending_step_log is not None:
_log(pending_step_log.replace("_DONE_", "true"))
try:
final_score = grade(env._state) if env is not None else 0.0
except Exception:
final_score = 0.0
success = "true" if final_score >= 0.5 else "false"
rewards_str = ",".join(f"{r:.2f}" for r in rewards) if rewards else "0.00"
_log(f"[END] success={success} steps={step_count} score={final_score:.3f} rewards={rewards_str}")
return final_score
def main():
tasks = ["easy", "medium", "hard"]
scores = {}
for task_id in tasks:
try:
score = run_task(task_id)
scores[task_id] = score
except Exception as e:
_err(f"[error] Task '{task_id}' failed: {e}")
scores[task_id] = 0.0
_log(f"[END] success=false steps=0 score=0.000 rewards=0.00")
return sum(scores.values()) / len(scores) if scores else 0.0
if __name__ == "__main__":
avg = main()
sys.exit(0)