khushmagrawal/devsecops_env

1import os2import json3import asyncio4from typing import List, Optional5from openai import OpenAI6 7import sys8sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))9 10try:11    from devsecops_env.server.devsecops_env_environment import DevSecOpsEnvironment12    from devsecops_env.models import DevsecopsAction13    from devsecops_env.server.graders import compute_reward14except ImportError:15    from server.devsecops_env_environment import DevSecOpsEnvironment16    from models import DevsecopsAction17    from server.graders import compute_reward18 19# Setup Configuration20API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")21MODEL_NAME = os.getenv("MODEL_NAME", "HuggingFaceH4/zephyr-7b-beta")22HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")23 24BENCHMARK = "devsecops_env"25MAX_STEPS = 1026SUCCESS_SCORE_THRESHOLD = 0.527 28 29def log_start(task: str, env: str, model: str) -> None:30    print(f"[START] task={task} env={env} model={model}", flush=True)31 32 33def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:34    error_val = error if error else "null"35    done_val = str(done).lower()36    # Ensuring no internal quotes mess up line structure37    escaped_action = action.replace("\n", " ").replace('"', "'")38    print(f"[STEP] step={step} action={escaped_action} reward={reward:.2f} done={done_val} error={error_val}", flush=True)39 40 41def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:42    rewards_str = ",".join(f"{r:.2f}" for r in rewards)43    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)44 45 46def get_mock_action(task_id: str, step: int) -> str:47    """Mock agent behavior for when HF_TOKEN is not provided."""48    if task_id == "task1":49        if step == 1: return '{"tool_name": "inspect_diff"}'50        if step == 2: return '{"tool_name": "run_ci", "scope": "full"}'51        return '{"tool_name": "make_decision", "verdict": "MERGE", "justification": "Looks good"}'52 53    elif task_id == "task2":54        if step == 1: return '{"tool_name": "inspect_diff"}'55        if step == 2: return '{"tool_name": "run_ci", "scope": "full"}'56        return '{"tool_name": "make_decision", "verdict": "REQUEST_CHANGES", "justification": "Tests fail"}'57 58    elif task_id == "task3":59        if step == 1: return '{"tool_name": "inspect_diff"}'60        if step == 2: return '{"tool_name": "query_package_registry", "pkg": "cryptoutils", "version": "2.1.5"}'61        return '{"tool_name": "make_decision", "verdict": "BLOCK", "justification": "Malware detected"}'62        63    return '{"tool_name": "inspect_diff"}'64 65 66import textwrap67 68def build_prompt(obs) -> str:69    history = "\n".join([f"- Step {t.step}: {t.tool_name} -> {t.result[:100]}" for t in obs.pipeline_history])70    71    prompt = textwrap.dedent(f"""72        # TASK: Review Pull Request73        PR Title: {obs.pr.title}74        Task Objective: {obs.task_id}75        76        # CURRENT STATE77        Step Count: {obs.step_count}/1078        CI Budget: {obs.budget.ci_runs} runs remaining79        80        # HISTORY81        {history if history else "No actions taken yet."}82        83        # LAST TOOL OUTPUT84        {obs.last_tool_output if obs.last_tool_output else "None"}85        86        # AVAILABLE TOOLS87        - inspect_diff (No params) -> View the code changes88        - run_ci (scope: "unit_only" or "full") -> Run tests89        - query_package_registry (pkg: str, version: str) -> Check package safety90        - search_vuln_db (pkg: str, version: str) -> Check for CVEs91        - patch_code (file: str, old_code: str, new_code: str) -> Fix a bug92        - make_decision (verdict: "MERGE"|"BLOCK"|"REQUEST_CHANGES", justification: str) -> FINISH TASK93        94        # REQUIREMENT95        You MUST output a FLAT JSON object. Do not nest parameters inside 'required_parameters'.96        Example: {{"tool_name": "run_ci", "scope": "unit_only"}}97        98        If you have enough info, use 'make_decision' to end the episode.99    """).strip()100    return prompt101 102 103def get_model_action(client, obs, step: int) -> str:104    105    try:106        user_prompt = build_prompt(obs)107        response = client.chat.completions.create(108            model=MODEL_NAME,109            messages=[110                {"role": "system", "content": "You are a DevSecOps Expert AI. You only reply with functional JSON actions."},111                {"role": "user", "content": user_prompt}112            ],113            temperature=0.1,114            max_tokens=512,115        )116        return response.choices[0].message.content.strip()117    except Exception as e:118        print(f"[ERROR] Request failed: {e}")119        # Return mock to allow testing environment logic when API is down120        return get_mock_action(obs.task_id, step)121 122 123async def run_episode(client, env, task_id: str) -> None:124    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)125    126    obs = env.reset(options={"task": task_id})127    128    rewards = []129    steps_taken = 0130    success = False131    132    for step in range(1, MAX_STEPS + 1):133        if obs.done:134            break135            136        action_json = get_model_action(client, obs, step)137        138        try:139            import re140            141            if not action_json:142                raise ValueError("Model returned None (API call failed)")143                144            cleaned_json = str(action_json).strip()145            146            # Try to find a JSON block in markdown147            match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', cleaned_json, re.DOTALL)148            if match:149                cleaned_json = match.group(1)150            else:151                # Fallback: try to find the first { and last }152                start_idx = cleaned_json.find('{')153                end_idx = cleaned_json.rfind('}')154                if start_idx != -1 and end_idx != -1 and end_idx > start_idx:155                    cleaned_json = cleaned_json[start_idx:end_idx+1]156            157            action_dict = json.loads(cleaned_json)158            action = DevsecopsAction(**action_dict)159            error_msg = None160        except Exception as e:161            action = DevsecopsAction(tool_name="inspect_diff")162            raw_text = str(action_json)[:40] if action_json else "None"163            error_msg = f"Parse err: {str(e)[:40]} | Raw: {raw_text}"164            action_json = '{"tool_name": "inspect_diff"}'165            166        obs = env.step(action)167        168        reward = obs.reward169        done = obs.done170        171        rewards.append(reward)172        steps_taken = step173        174        log_step(step=step, action=action_json, reward=reward, done=done, error=error_msg)175        176        if done:177            break178            179    last_verdict = None180    for record in reversed(obs.pipeline_history):181        if record.tool_name == "make_decision":182            if "merge" in record.result.lower():183                last_verdict = "MERGE"184            elif "block" in record.result.lower():185                last_verdict = "BLOCK"186            elif "request" in record.result.lower():187                last_verdict = "REQUEST_CHANGES"188            break189    190    score = obs.reward191    score = min(max(score, 0.001), 0.999)192    success = score >= SUCCESS_SCORE_THRESHOLD193    194    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)195 196 197async def main() -> None:198    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "dummy_key")199    200    env = DevSecOpsEnvironment()201    202    tasks = ["task1", "task2", "task3"]203    for task_id in tasks:204        await run_episode(client, env, task_id)205        206 207if __name__ == "__main__":208    asyncio.run(main())209