Spaces:
Sleeping
Sleeping
File size: 5,565 Bytes
a617acd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from __future__ import annotations
import argparse
import json
import os
from typing import Any
import httpx
from openai import OpenAI
SYSTEM_PROMPT = (
"You are an audit agent. Return strict JSON with keys: action_type, violation_type, confidence, note. "
"Choose action_type from submit_finding, flag_human_review, noop."
)
def _build_action(task_id: str, observation: dict[str, Any], client: OpenAI, model: str) -> dict[str, Any]:
"""Build an action using the OpenAI Chat Completions API."""
documents = observation.get("documents", [])
doc_id = documents[0]["id"] if documents else "UNKNOWN"
user_prompt = (
"Task: " + task_id + "\n"
"Given this sample document id, propose one conservative action.\n"
f"document_id: {doc_id}\n"
"Return JSON only."
)
completion = client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
temperature=0,
max_tokens=200,
)
text = (completion.choices[0].message.content or "").strip()
# Strip markdown fences if present
if text.startswith("```"):
lines = text.split("\n")
lines = [l for l in lines if not l.strip().startswith("```")]
text = "\n".join(lines).strip()
# Safe fallback if model output is not parseable JSON.
if not text.startswith("{"):
return {"action_type": "noop", "task_id": task_id, "note": "fallback_no_json"}
try:
payload = json.loads(text)
except Exception:
return {"action_type": "noop", "task_id": task_id, "note": "fallback_parse_error"}
action_type = payload.get("action_type", "noop")
if action_type not in {"submit_finding", "flag_human_review", "noop"}:
action_type = "noop"
if action_type != "submit_finding":
return {"action_type": action_type, "task_id": task_id, "note": payload.get("note", "")}
violation_type = payload.get("violation_type", "duplicate_receipt")
confidence = float(payload.get("confidence", 0.5))
confidence = max(0.0, min(1.0, confidence))
return {
"action_type": "submit_finding",
"task_id": task_id,
"finding": {
"document_id": doc_id,
"violation_type": violation_type,
"evidence": [doc_id],
"confidence": confidence,
},
"note": payload.get("note", "baseline_action"),
}
def _build_heuristic_action(task_id: str, observation: dict[str, Any]) -> dict[str, Any]:
"""Free fallback policy for local validation when API credits are unavailable."""
documents = observation.get("documents", [])
doc_id = documents[0]["id"] if documents else "UNKNOWN"
violation_map = {
"easy": "duplicate_receipt",
"medium": "sod_conflict",
"hard": "shell_company",
}
return {
"action_type": "submit_finding",
"task_id": task_id,
"finding": {
"document_id": doc_id,
"violation_type": violation_map.get(task_id, "duplicate_receipt"),
"evidence": [doc_id],
"confidence": 0.5,
},
"note": "heuristic_fallback_policy",
}
def run_task(
env_url: str,
task_id: str,
client: OpenAI | None,
model: str,
seed: int,
policy: str,
) -> float:
with httpx.Client(timeout=20.0) as http:
obs = http.post(f"{env_url}/reset", json={"task_id": task_id, "seed": seed}).json()
total = 0.0
steps = 0
done = False
while not done:
if policy == "heuristic":
action = _build_heuristic_action(task_id=task_id, observation=obs)
else:
if client is None:
raise RuntimeError("OPENAI_API_KEY is required for policy=openai")
action = _build_action(task_id=task_id, observation=obs, client=client, model=model)
result = http.post(f"{env_url}/step", json=action).json()
total += float(result["reward"]["normalized"])
steps += 1
done = bool(result["done"])
obs = result["observation"]
# Mean normalized reward per step (bounded [0,1] by construction)
return round(total / steps, 6)
def main() -> None:
parser = argparse.ArgumentParser(description="Run reproducible baseline scores on all AuditEnv tasks.")
parser.add_argument("--env-url", default=os.getenv("AUDITENV_BASE_URL", "http://127.0.0.1:8000"))
parser.add_argument("--model", default=os.getenv("AUDITENV_BASELINE_MODEL", "gpt-4.1-mini"))
parser.add_argument(
"--policy",
choices=["openai", "heuristic"],
default="openai",
help="Action policy: 'openai' uses API key, 'heuristic' is free local fallback.",
)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
client: OpenAI | None = None
if args.policy == "openai":
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise RuntimeError("OPENAI_API_KEY is required for --policy openai")
client = OpenAI(api_key=api_key)
scores = {}
for task_id in ["easy", "medium", "hard"]:
scores[task_id] = run_task(args.env_url, task_id, client, args.model, args.seed, args.policy)
print("Baseline scores (normalized):")
for task_id in ["easy", "medium", "hard"]:
print(f"- {task_id}: {scores[task_id]:.6f}")
if __name__ == "__main__":
main()
|