Drac0528 commited on
Commit
100988d
·
verified ·
1 Parent(s): ad262b3

Delete inference.py

Browse files
Files changed (1) hide show
  1. inference.py +0 -251
inference.py DELETED
@@ -1,251 +0,0 @@
1
- #!/usr/bin/env python3
2
- from __future__ import annotations
3
-
4
- import asyncio
5
- import json
6
- import os
7
- from typing import Any, Dict, List, Optional
8
-
9
- from openai import OpenAI
10
-
11
- try:
12
- from code_security_auditor_env import CodeSecurityAction, CodeSecurityAuditorEnv
13
- except ImportError:
14
- from client import CodeSecurityAuditorEnv
15
- from models import CodeSecurityAction
16
-
17
- API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
18
- MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
19
- API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
20
- LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
21
- ENV_BASE_URL = os.getenv("ENV_BASE_URL")
22
- DEFAULT_ENV_BASE_URL = os.getenv("DEFAULT_ENV_BASE_URL", "http://127.0.0.1:8000")
23
- DEFAULT_LOCAL_IMAGE_NAME = os.getenv("DEFAULT_LOCAL_IMAGE_NAME", "code-security-auditor-env:latest")
24
- TASK_IDS = [t.strip() for t in os.getenv("TASK_IDS", "easy,medium,hard").split(",") if t.strip()]
25
- MAX_STEPS = int(os.getenv("MAX_STEPS", "12"))
26
- TEMPERATURE = 0.0
27
- MAX_TOKENS = 260
28
- BENCHMARK = "code_security_auditor_env"
29
-
30
- SYSTEM_PROMPT = (
31
- "You are a senior application security reviewer. Produce strictly valid JSON for the next action. "
32
- "Allowed action_type values: inspect_file, submit_finding, submit_final_report. "
33
- "Do not include markdown fences. Keep fields concise and accurate."
34
- )
35
-
36
-
37
- def log_start(task: str, env: str, model: str) -> None:
38
- print(f"[START] task={task} env={env} model={model}", flush=True)
39
-
40
-
41
- def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
42
- err = error if error else "null"
43
- print(
44
- f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={err}",
45
- flush=True,
46
- )
47
-
48
-
49
- def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
50
- rewards_str = ",".join(f"{r:.2f}" for r in rewards)
51
- print(
52
- f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
53
- flush=True,
54
- )
55
-
56
-
57
- def _compact_action_str(action: Dict[str, Any]) -> str:
58
- return json.dumps(action, separators=(",", ":"), ensure_ascii=True)
59
-
60
-
61
- def _default_action() -> Dict[str, Any]:
62
- return {
63
- "action_type": "submit_final_report",
64
- "confidence": 0.5,
65
- "summary": "fallback-finalize",
66
- "evidence": "fallback-finalize",
67
- }
68
-
69
-
70
- def _safe_error(exc: Exception) -> str:
71
- msg = str(exc).strip()
72
- if not msg:
73
- msg = exc.__class__.__name__
74
- return msg.replace("\n", " ")[:240]
75
-
76
-
77
- def _parse_action(raw: str, available_files: List[str]) -> Dict[str, Any]:
78
- try:
79
- parsed = json.loads(raw)
80
- if not isinstance(parsed, dict):
81
- return _default_action()
82
- except Exception:
83
- return _default_action()
84
-
85
- action_type = parsed.get("action_type")
86
- if action_type not in {"inspect_file", "submit_finding", "submit_final_report"}:
87
- return _default_action()
88
-
89
- action: Dict[str, Any] = {
90
- "action_type": action_type,
91
- "confidence": float(parsed.get("confidence", 0.5)),
92
- "summary": str(parsed.get("summary", ""))[:400],
93
- "evidence": str(parsed.get("evidence", ""))[:700],
94
- }
95
-
96
- if parsed.get("filename"):
97
- filename = str(parsed["filename"])
98
- if filename in available_files:
99
- action["filename"] = filename
100
- if parsed.get("line_start") is not None:
101
- try:
102
- action["line_start"] = max(1, int(parsed["line_start"]))
103
- except Exception:
104
- pass
105
- if parsed.get("line_end") is not None:
106
- try:
107
- action["line_end"] = max(1, int(parsed["line_end"]))
108
- except Exception:
109
- pass
110
- if parsed.get("vuln_type") is not None:
111
- action["vuln_type"] = str(parsed["vuln_type"])
112
- if parsed.get("severity") is not None:
113
- action["severity"] = str(parsed["severity"])
114
-
115
- action["confidence"] = min(1.0, max(0.0, action["confidence"]))
116
-
117
- return action
118
-
119
-
120
- def _build_prompt(obs: Any, step: int) -> str:
121
- findings = obs.findings_so_far[-4:] if obs.findings_so_far else []
122
- snippet = obs.file_excerpt[:1800] if obs.file_excerpt else ""
123
- return (
124
- f"Task: {obs.task_id} ({obs.difficulty})\\n"
125
- f"Objective: {obs.objective}\\n"
126
- f"Step: {step}\\n"
127
- f"Steps remaining: {obs.steps_remaining}\\n"
128
- f"Files: {', '.join(obs.available_files)}\\n"
129
- f"Last feedback: {obs.last_feedback}\\n"
130
- f"Focused file: {obs.focused_file}\\n"
131
- f"Recent findings: {json.dumps(findings)}\\n"
132
- f"Visible snippet:\\n{snippet}\\n"
133
- "Return one JSON object with action_type and required fields."
134
- )
135
-
136
-
137
- def _query_model(client: OpenAI, obs: Any, step: int) -> Dict[str, Any]:
138
- user_prompt = _build_prompt(obs, step)
139
- try:
140
- resp = client.chat.completions.create(
141
- model=MODEL_NAME,
142
- messages=[
143
- {"role": "system", "content": SYSTEM_PROMPT},
144
- {"role": "user", "content": user_prompt},
145
- ],
146
- temperature=TEMPERATURE,
147
- max_tokens=MAX_TOKENS,
148
- stream=False,
149
- )
150
- content = (resp.choices[0].message.content or "").strip()
151
- return _parse_action(content, obs.available_files)
152
- except Exception:
153
- return _default_action()
154
-
155
-
156
- async def _create_env() -> CodeSecurityAuditorEnv:
157
- # Prefer explicit configuration, then fall back to common local defaults.
158
- if ENV_BASE_URL:
159
- return CodeSecurityAuditorEnv(base_url=ENV_BASE_URL)
160
-
161
- if LOCAL_IMAGE_NAME:
162
- return await CodeSecurityAuditorEnv.from_docker_image(LOCAL_IMAGE_NAME)
163
-
164
- try:
165
- return CodeSecurityAuditorEnv(base_url=DEFAULT_ENV_BASE_URL)
166
- except Exception:
167
- return await CodeSecurityAuditorEnv.from_docker_image(DEFAULT_LOCAL_IMAGE_NAME)
168
-
169
-
170
- async def run_task(env: CodeSecurityAuditorEnv, client: OpenAI, task_id: str) -> float:
171
- log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
172
-
173
- rewards: List[float] = []
174
- steps_taken = 0
175
- score = 0.0
176
- success = False
177
-
178
- try:
179
- result = await env.reset(task_id=task_id)
180
- obs = result.observation
181
-
182
- for step in range(1, MAX_STEPS + 1):
183
- if result.done:
184
- break
185
-
186
- action_dict = _query_model(client, obs, step)
187
- action_str = _compact_action_str(action_dict)
188
-
189
- action = CodeSecurityAction(**action_dict)
190
- result = await env.step(action)
191
- obs = result.observation
192
-
193
- reward = float(result.reward or 0.0)
194
- done = bool(result.done)
195
- error = obs.metadata.get("last_action_error")
196
-
197
- rewards.append(reward)
198
- steps_taken = step
199
- log_step(step=step, action=action_str, reward=reward, done=done, error=error)
200
-
201
- if done:
202
- break
203
-
204
- score = float(obs.reward or 0.0)
205
- score = min(max(score, 0.0), 1.0)
206
- success = score >= 0.6
207
- except Exception as exc:
208
- # Keep evaluator contract: do not crash inference.py on transient/runtime errors.
209
- log_step(step=max(1, steps_taken), action="{}", reward=0.0, done=True, error=_safe_error(exc))
210
- if not rewards:
211
- rewards.append(0.0)
212
- steps_taken = max(1, steps_taken)
213
- score = 0.0
214
- success = False
215
- finally:
216
- log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
217
-
218
- return score
219
-
220
-
221
- async def main() -> None:
222
- # Keep script resilient in validators even if a key is temporarily unavailable.
223
- api_key = API_KEY or "missing"
224
-
225
- client = OpenAI(base_url=API_BASE_URL, api_key=api_key)
226
-
227
- try:
228
- env = await _create_env()
229
- except Exception as exc:
230
- # Emit structured logs for each task and exit cleanly.
231
- err = _safe_error(exc)
232
- for task_id in TASK_IDS:
233
- log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
234
- log_step(step=1, action="{}", reward=0.0, done=True, error=err)
235
- log_end(success=False, steps=1, score=0.0, rewards=[0.0])
236
- return
237
-
238
- try:
239
- scores: List[float] = []
240
- for task_id in TASK_IDS:
241
- score = await run_task(env, client, task_id)
242
- scores.append(score)
243
-
244
- # Keep strict output format requirement: no extra structured tags beyond START/STEP/END.
245
- _ = scores
246
- finally:
247
- await env.close()
248
-
249
-
250
- if __name__ == "__main__":
251
- asyncio.run(main())