Navigam/jira-to-code

1# inference.py — ReAct Agent for Jira-to-Code Environment2#3# Architecture:4#   Phase 1: Episodic Memory — persistent messages[] across the episode5#   Phase 2: ReAct Pattern — "thought" key forces reasoning before action6#   Phase 3: Robust Parsing — JSON extraction with markdown-fence stripping7#   Phase 4: Self-Correction — negative rewards inject corrective prompts8#   Phase 5: Multi-Task Loop — evaluates all 6 tasks in one run9 10import argparse11import json12import os13import re14import textwrap15import time16from typing import List, Optional17 18from openai import OpenAI19from dotenv import load_dotenv20 21load_dotenv()22 23# Our environment for local/direct testing24from server.env import JiraToCodeEnv25from src.jira_to_code.models import JiraCodeAction26 27# --- HACKATHON MANDATORY CONFIGURATION ---28API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"29MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"30HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")31 32BENCHMARK = "jira-to-code"33# MAX_STEPS is now dynamic based on task level34SUCCESS_SCORE_THRESHOLD = 0.9  # Account for step penalties35ALL_TASKS = list(JiraToCodeEnv.TASKS.keys())36MAX_HISTORY_MESSAGES = 30  # Context-window safety: trim if exceeded37MAX_RETRIES = 5            # Rate limit retry attempts38RETRY_BASE_DELAY = 2       # Base delay in seconds for exponential backoff39 40# --- SYSTEM PROMPT (ReAct + Reward-Aware) ---41SYSTEM_PROMPT = textwrap.dedent("""\42You are an expert software engineer resolving Jira tickets.43You operate in a sandboxed workspace. You can read files, write code, list files, run tests, and submit your solution.44 45## Rules461. ALWAYS respond with ONLY a valid JSON object. No markdown fences, no explanations outside JSON.472. You MUST include a "thought" key FIRST to reason about your plan before acting.483. Work step-by-step: list files, read the code, understand the bug/requirement, write a fix, run tests, then submit.494. If tests fail, carefully read the traceback and fix your code before re-submitting.505. Only use "submit" when you are confident all tests will pass.516. Be efficient — each step has a small penalty. Aim to solve in the fewest steps possible.527. Read the test file to understand exactly what is expected before writing code.53 54## Valid action_types55- "list_files" — List all files in the workspace (file_path and content should be null)56- "read_file" — Read a file's contents (requires file_path, content should be null)57- "write_file" — Write/overwrite a file (requires file_path and content)58- "run_tests" — Run pytest on the workspace (file_path and content should be null)59- "submit" — Final submission, runs tests and ends the episode (file_path and content should be null)60 61## Reward Structure62- list_files / read_file: 0.01 (initial exploration)63- write_file: +0.05 (reward for taking action)64- run_tests (all pass): +0.5 | run_tests (partial): proportional | run_tests (crash): 0.0165- submit (all pass): +1.0 | submit (partial): proportional66- Every step: 0.01 minimum reward (be efficient!)67 68## JSON Schema69{70  "thought": "Your reasoning about what to do next and why",71  "action_type": "one of: list_files, read_file, write_file, run_tests, submit",72  "file_path": "string or null",73  "content": "string or null"74}75 76## Strategy Guide771. First, list_files to see the workspace structure.782. Read the test file to understand the exact expected behavior.793. Read the source file to understand the current (buggy/incomplete) code.804. Write the fix/implementation.815. Run tests to verify.826. If tests pass, submit. If not, read the error, fix, and retry.83""").strip()84 85 86# --- MANDATORY LOGGING FUNCTIONS ---87def log_start(task: str, env: str, model: str) -> None:88    print(f"[START] task={task} env={env} model={model}", flush=True)89 90 91def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:92    error_val = error if error else "null"93    done_val = str(done).lower()94    print(95        f"[STEP] step={step} action={action} reward={reward:.2f} "96        f"done={done_val} error={error_val}",97        flush=True,98    )99 100 101def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:102    rewards_str = ",".join(f"{r:.2f}" for r in rewards)103    print(104        f"[END] success={str(success).lower()} steps={steps} "105        f"score={score:.3f} rewards={rewards_str}",106        flush=True,107    )108 109 110# --- PHASE 3: ROBUST JSON PARSING ---111def extract_json(raw_text: str) -> dict:112    """113    Extract a JSON object from LLM output, handling:114    - Markdown code fences (```json ... ```)115    - Leading/trailing whitespace and text116    - Nested braces via brace-counting117    """118    cleaned = raw_text.strip()119    cleaned = re.sub(r'^```(?:json)?\s*', '', cleaned)120    cleaned = re.sub(r'\s*```\s*$', '', cleaned)121    cleaned = cleaned.strip()122 123    # Try direct parse first124    try:125        return json.loads(cleaned)126    except json.JSONDecodeError:127        pass128 129    # Fallback: find the first balanced {...} block via brace counting130    start = cleaned.find('{')131    if start == -1:132        raise ValueError("No JSON object found in response")133 134    depth = 0135    in_string = False136    escape_next = False137    for i in range(start, len(cleaned)):138        c = cleaned[i]139        if escape_next:140            escape_next = False141            continue142        if c == '\\' and in_string:143            escape_next = True144            continue145        if c == '"' and not escape_next:146            in_string = not in_string147            continue148        if in_string:149            continue150        if c == '{':151            depth += 1152        elif c == '}':153            depth -= 1154            if depth == 0:155                return json.loads(cleaned[start:i + 1])156 157    raise ValueError("Unbalanced braces in JSON")158 159 160def parse_action(raw_text: str) -> JiraCodeAction:161    """Parse LLM output into a JiraCodeAction, extracting JSON robustly."""162    action_dict = extract_json(raw_text)163    # Remove the 'thought' key — it's for reasoning only, not part of the action model164    action_dict.pop("thought", None)165    return JiraCodeAction(**action_dict)166 167 168# --- PHASE 1 & 2: BUILD OBSERVATION MESSAGE ---169def build_observation_message(step: int, obs, reward: float) -> str:170    """Format environment observation as a user message for the conversation history."""171    parts = [172        f"--- Step {step} Observation ---",173        f"Ticket: {obs.jira_ticket}",174        f"Files in workspace: {', '.join(obs.file_tree) if obs.file_tree else 'None'}",175    ]176    if obs.current_file_content is not None:177        parts.append(f"File Content:\n```\n{obs.current_file_content}\n```")178    if obs.test_output:179        parts.append(f"Test Output:\n```\n{obs.test_output}\n```")180    if obs.error:181        parts.append(f"Error: {obs.error}")182    parts.append(f"Reward: {reward:.2f}")183    parts.append("Respond with your next action as JSON.")184    return "\n".join(parts)185 186 187def trim_history(messages: list, max_messages: int = MAX_HISTORY_MESSAGES) -> None:188    """Trim oldest non-system messages if history exceeds max to avoid context overflow."""189    while len(messages) > max_messages:190        # Keep index 0 (system prompt), remove index 1191        messages.pop(1)192 193 194# --- MAIN AGENT LOOP FOR ONE TASK ---195def run_agent_episode(client: OpenAI, task_name: str) -> tuple:196    """197    Run a full agent episode for one task.198    Returns: (score, steps_taken, rewards, success)199    """200    os.environ["JIRA_TASK_LEVEL"] = task_name201    env = JiraToCodeEnv()202 203    rewards: List[float] = []204    steps_taken = 0205    score = 0.0206    success = False207 208    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)209 210    try:211        obs = env.reset()212 213        task_max_steps = 10 if "easy" in task_name else 20214 215        # Phase 1: Episodic memory — persistent conversation history216        messages = [217            {"role": "system", "content": SYSTEM_PROMPT},218            {"role": "user", "content": build_observation_message(0, obs, 0.0)},219        ]220 221        for step in range(1, task_max_steps + 1):222            trim_history(messages)223 224            # Call the LLM with rate-limit retry + exponential backoff225            raw_text = None226            for attempt in range(MAX_RETRIES):227                try:228                    completion = client.chat.completions.create(229                        model=MODEL_NAME,230                        messages=messages,231                        temperature=0.2,232                        max_tokens=2048,233                    )234                    raw_text = (completion.choices[0].message.content or "").strip()235                    break  # Success236                except Exception as exc:237                    exc_str = str(exc)238                    is_rate_limit = "429" in exc_str or "rate" in exc_str.lower()239                    if is_rate_limit and attempt < MAX_RETRIES - 1:240                        delay = RETRY_BASE_DELAY * (2 ** attempt)241                        print(f"  [RATE LIMIT] Retry {attempt + 1}/{MAX_RETRIES} in {delay}s...", flush=True)242                        time.sleep(delay)243                        continue244                    # Non-rate-limit error or final attempt — give up245                    messages.append({246                        "role": "user",247                        "content": f"API ERROR: {exc}. Please try again with a valid JSON action.",248                    })249                    log_step(step=step, action=f"API_ERROR: {exc}", reward=0.0, done=False, error=exc_str)250                    rewards.append(0.0)251                    steps_taken = step252                    break253 254            if raw_text is None:255                continue  # Skip to next step if all retries failed256 257            # Phase 1: Append assistant response to history258            messages.append({"role": "assistant", "content": raw_text})259 260            # Phase 3: Robust parsing with safe fallback261            try:262                action = parse_action(raw_text)263                action_log = action.model_dump_json()264            except Exception as exc:265                # Parse failure — No-Op fallback + corrective injection266                action = JiraCodeAction(action_type="list_files")267                action_log = f"PARSE_ERROR: {exc}"268 269                # Phase 4: Inject corrective message270                messages.append({271                    "role": "user",272                    "content": (273                        f"ERROR: Your last response was not valid JSON.\n"274                        f"Parse error: {exc}\n"275                        f"You MUST respond with ONLY a valid JSON object. "276                        f"No markdown, no explanations.\nTry again."277                    ),278                })279 280            # Take step in environment281            obs, reward, done, _ = env.step(action)282            error = obs.error283 284            # Ensure individual step rewards are strictly positive (min 0.01)285            reward = max(reward, 0.01)286 287            rewards.append(reward)288            steps_taken = step289 290            # Escape newlines for single-line logging291            safe_action_str = action_log.replace('\n', '\\n').replace('\r', '')292            log_step(step=step, action=safe_action_str, reward=reward, done=done, error=error)293 294            if done:295                break296 297            # Phase 1: Append observation to conversation history298            obs_message = build_observation_message(step, obs, reward)299 300            # Phase 4: Self-correction prompt injection on low/negative reward or error301            if reward <= 0.01 or obs.error:302                obs_message += (303                    f"\n\nLOW/NEGATIVE RESULT (reward={reward:.2f})."304                    f"\nCarefully analyze the error/test output above."305                    f"\nIdentify the root cause and write a fix."306                    f"\nDo NOT repeat the same action that just failed."307                )308            elif reward >= 0.4:309                obs_message += (310                    "\n\nTests are passing! If all tests pass, use 'submit' to finalize."311                )312 313            messages.append({"role": "user", "content": obs_message})314 315        # Calculate final score (clamp strictly between 0 and 1)316        score = min(max(sum(rewards), 0.01), 0.99)317        success = score >= SUCCESS_SCORE_THRESHOLD318 319    finally:320        env.close()321        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)322 323    return score, steps_taken, rewards, success324 325 326# --- PHASE 5: MULTI-TASK EVALUATION ---327def main() -> None:328    parser = argparse.ArgumentParser(description="Jira-to-Code ReAct Agent")329    parser.add_argument(330        "--tasks",331        type=str,332        default=None,333        help=(334            "Comma-separated list of tasks to run. "335            f"Available: {', '.join(ALL_TASKS)}. "336            "Default: all tasks."337        ),338    )339    args = parser.parse_args()340 341    import random342 343    # Determine which tasks to run344    if args.tasks:345        tasks = [t.strip() for t in args.tasks.split(",")]346        invalid = [t for t in tasks if t not in ALL_TASKS]347        if invalid:348            print(f"ERROR: Unknown tasks: {invalid}", flush=True)349            print(f"Available: {ALL_TASKS}", flush=True)350            return351    else:352        # Baseline inference: 1 easy, 1 medium, 1 hard randomly sampled353        easies = [t for t in ALL_TASKS if "easy" in t]354        mediums = [t for t in ALL_TASKS if "medium" in t]355        hards = [t for t in ALL_TASKS if "hard" in t]356        357        tasks = []358        if easies: tasks.append(random.choice(easies))359        if mediums: tasks.append(random.choice(mediums))360        if hards: tasks.append(random.choice(hards))361 362    print(f"Running tasks: {tasks}", flush=True)363 364    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)365 366    total_score = 0.0367    results = []368 369    for task in tasks:370        score, steps, rewards, success = run_agent_episode(client, task)371        results.append({372            "task": task,373            "score": score,374            "steps": steps,375            "success": success,376        })377        total_score += score378 379        print("Waiting 20 seconds before next task to respect API limits...", flush=True)380        time.sleep(20)381 382    # Summary383    print("\n" + "=" * 50, flush=True)384    print("EVALUATION SUMMARY", flush=True)385    print("=" * 50, flush=True)386    for r in results:387        status = "PASS" if r["success"] else "FAIL"388        print(389            f"  {r['task']:10s} | score={r['score']:.3f} | "390            f"steps={r['steps']:2d} | {status}",391            flush=True,392        )393    avg_score = total_score / len(tasks)394    print(f"  {'AVERAGE':10s} | score={avg_score:.3f}", flush=True)395    print(f"  {'TOTAL':10s} | score={total_score:.3f} / {len(tasks):.1f}", flush=True)396    print("=" * 50, flush=True)397 398 399if __name__ == "__main__":400    main()