coding_agent.py

# Copyright (c) 2025 Andrey Volkov

# This work is a derivative work based on the original by Mahtab Syed.
# Original author copyright notice:
# MIT License
# Copyright (c) 2025 Mahtab Syed
# https://www.linkedin.com/in/mahtabsyed/

import os
import random
import re
import sys
import json
import time
import subprocess
import tempfile
import argparse
from pathlib import Path
from google import genai
from google.genai import errors
from patch import patch_code, is_unified_diff
from sandbox_execution import execute_sandboxed
from token_tracker import TokenUsageTracker
from utils import *

# Initialize Gemini LLM key
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    raise ValueError("GEMINI_API_KEY environment variable not set")

default_llm_model = "gemini-2.5-flash"
print(f"📡 Initializing Gemini LLM ...")
llm = genai.Client(api_key=api_key)

class Iteration:
    def __init__(self):
        self.code = None
        self.feedback = None
        self.flags = set()
        self.program_output = None
        self.score = None

    def add_flag(self, name: str):
        self.flags.add(name)
    
    def get_score(self) -> int:
        """Returns the score, or 0 if not set"""
        return self.score if self.score is not None else 0

class Context:
    def __init__(self, filename, use_case, goals):
        self.filename = filename
        self.use_case = use_case
        self.goals = goals
        self.research_summary = ""
        self._iterations = []
        self.current_iteration = None

    @property
    def iterations(self):
        """Returns a copy of the iterations list"""
        return self._iterations.copy()
    
    @property
    def previous(self):
        if len(self._iterations):
            return self._iterations[-1]
        return None

    @property
    def current(self):
        if not self.current_iteration:
            raise RuntimeError("Start an iteration before accessing current")
        return self.current_iteration

    @property
    def iter_no(self):
        return len(self._iterations) + 1

    def start_iteration(self):
        if self.current_iteration:
            self._iterations.append(self.current_iteration)
        self.current_iteration = Iteration()

    def erase_iteration(self):
        self.current_iteration = None

    def trim_iterations(self, limit_by):
        """Trim iterations to keep only the first limit_by iterations.
        The last kept iteration becomes the current iteration."""
        self._iterations = self._iterations[:limit_by]
        if self._iterations:
            self.current_iteration = self._iterations.pop()
        else:
            self.current_iteration = None

    def save_to(self, filename_template, content, content_name=None):
        """
            Saves content to a file with a name based on the template.
            Template parameters are:
                - {name}, which is a solution base name
                - {iter} which is the iteration number.
        """
        try:
            fn = filename_template.format(name=self.filename, iter=self.iter_no)
            save_to_file(fn, content, content_name)
        except KeyError as ke:
            print(f"Error creating file name: key {ke} in the template is invalid: {filename_template}")
            sys.exit(1)
            

llm_config_coder = genai.types.GenerateContentConfig(
    temperature=0.3,
    tools=[
        genai.types.Tool(code_execution=genai.types.ToolCodeExecution), 
    ],
)

llm_config_reviewer = genai.types.GenerateContentConfig(
    temperature=0.5,
    response_modalities=["TEXT"],  # Force text output
)

llm_config_research = genai.types.GenerateContentConfig(
    tools=[
        {"url_context": {}}
    ],
    response_modalities=["TEXT"],  # Force text output
    max_output_tokens=50000,  # limit the total output tokens
    thinkingConfig=genai.types.ThinkingConfig(
        thinking_budget=20000,
    )
)

llm_config_goals_check = genai.types.GenerateContentConfig(
    temperature=0.3,
    responseMimeType="text/x.enum",
    responseSchema={
        "type": "object",
        "properties": {
            "result": {
                "type": "string",
                "enum": ["Yes", "No"],
                "description": "Whether the goals have been met"
            },
            "score": {
                "type": "number",
                "minimum": 0,
                "maximum": 100,
                "description": "Completion score (0-100) of the solution against the goals"
            }
        }
    }
)

llm_config_refine_task = genai.types.GenerateContentConfig(
    temperature=0.1,
    responseMimeType="application/json",
    responseSchema={
        "type": "object",
        "properties": {
            "refined_use_case": {
                "type": "string",
                "description": "The refined use case text"
            },
            "refined_goals": {
                "type": "array",
                "items": {
                    "type": "string"
                },
                "description": "List of refined goals as separate strings"
            }
        },
        "required": ["refined_use_case", "refined_goals"]
    },
)

# Default configuration structure
DEFAULT_TASK_CONFIG = {
    "coder_model": "gemini-2.5-pro",
    "reviewer_model": "gemini-2.5-pro", 
    "utility_model": "gemini-2.5-flash-lite",
    "max_rounds": 25,
    "basename": "code",
    "sandbox_method": "auto",  # Options: auto, firejail, docker, bubblewrap, subprocess
    "commandline_args": ""
}

# Initialize token usage tracker
token_tracker = TokenUsageTracker()

def llm_query(query, parts=None, config=llm_config_coder, model=default_llm_model):
    """
    Query the LLM with retries on server errors.
    Args:
        if parts is None:
            query: The input query string
        else:
            query: System prompt string (will be cached automatically by Gemini if unchanged)
            parts: List of (title, content) tuples to build the prompt
        config: LLM configuration
        model: LLM model name
    """
    max_retries = 10
    
    for attempt in range(max_retries):
        try:
            # mark start time
            start_time = time.monotonic()
            request_config = config
            
            if model.startswith("gemini-3"):
                request_config.temperature = 1.0 # For Gemini 3 it is important not to alter the default temperature

            if parts is None:
                response = llm.models.generate_content(
                    model=model, contents=query, config=request_config
                )
            else:
                # Use system instruction for caching - this gets cached automatically by Gemini
                request_config.system_instruction = query
                
                # Build parts as proper content structure (not concatenated strings)
                request_parts = []
                for title, content in parts:
                    request_parts.append({"text": f"\n\n# {title}\n{content}"})
                
                request_contents = [{"role": "user", "parts": request_parts}]
                
                response = llm.models.generate_content(
                    model=model, contents=request_contents, config=request_config
                )
            end_time = time.monotonic()
            # Calculate generation time in seconds
            generation_time = end_time - start_time
            text = response.text

            # Print usage info and record statistics
            token_tracker.print_call_info(response.usage_metadata, generation_time)
            token_tracker.record(model, response.usage_metadata, generation_time)

            return {"text": text, "full": response, "usage": response.usage_metadata, "response_time": generation_time}
        
        except errors.ServerError as e:
            if attempt < max_retries - 1:
                # 15 seconds for 503, 5 seconds for other 5xx errors
                delay = 15 if e.code == 503 else 5
                print(f"⚠️  Server error: {e}")
                print(f"🔄 Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
                time.sleep(delay)
            else:
                print(f"❌ Server error after {max_retries} retries: {e}")
                raise

# --- Agent-Specific Functions ---

def load_task_config(config_name: str) -> dict:
    """Load configuration from tasks/{config_name}/config.json"""
    config_path = Path(f"tasks/{config_name}/config.json")
    
    if not config_path.exists():
        print(f"⚠️  Config file not found: {config_path}")
        print("🔄 Using default configuration")
        return DEFAULT_TASK_CONFIG.copy()
    
    try:
        with open(config_path, "r") as f:
            config = json.load(f)
        
        # Start with defaults and update with loaded values
        final_config = DEFAULT_TASK_CONFIG.copy()
        final_config.update(config)
        
        print(f"📋 Loaded config: coder={final_config['coder_model']}, reviewer={final_config['reviewer_model']}, utility={final_config['utility_model']}, max_rounds={final_config['max_rounds']}")
        return final_config
        
    except (json.JSONDecodeError, KeyError) as e:
        print(f"❌ Error loading config from {config_path}: {e}")
        print("🔄 Using default configuration")
        return DEFAULT_TASK_CONFIG.copy()

def refine_goals(config: dict, context: Context):
    # Refines goals and use case in the context
    refine_prompt = load_file("scripts/refine task.md")
    refine_response = llm_query(refine_prompt.format_map({
        "use_case": context.use_case,
        "goals": context.goals
    }), config=llm_config_refine_task, model=config["reviewer_model"])

    # save the refined response for debugging
    refine_text = refine_response["text"]
    refine_json = json.loads(refine_text)
    context.save_to("{name}_refined_use_case.md", refine_json["refined_use_case"], content_name="refined use case")
    context.save_to("{name}_refined_goals.md", refine_json["refined_goals"], content_name="refined goals")
    context.use_case = refine_json["refined_use_case"]
    context.goals = refine_json["refined_goals"]  # Keep as list
    return True

def research(config: dict, context: Context):
    # Refines goals and use case in the context
    refine_prompt = load_file("scripts/research.md")
    if not "urls" in config:
        return False # Nothing to do   
    urls = ", ".join(config["urls"])
    
    response = llm_query(refine_prompt.format_map({
        "use_case": context.use_case,
        "goals": context.goals,
        "urls": urls
    }), config=llm_config_research, model=config["utility_model"])

    # save the refined response for debugging
    context.save_to("{name}_research_raw_{iter}.json", response["full"].model_dump_json(indent=2), content_name="research JSON response" )
    summary = response["text"]
    if not summary:
        print("⚠️  Research step returned empty summary.")
        exit(1)
    else:
        context.save_to("{name}_research_summary_{iter}.md", summary, content_name="research summary")
    context.research_summary = summary or "No research summary available."
    return True

def code(config: dict, context: Context, use_diffs: bool = True):

    if context.previous:
        print("🔄Creating code refinement prompt...")
        script_path = "scripts/coder fix.md"
    else:
        print("📝 Constructing code generation prompt...")
        script_path = "scripts/coder create.md"

    script = load_file(script_path)
    if use_diffs:
        system_prompt = to_string(select_variant(to_lines(script), "a"))
    else:
        system_prompt = to_string(select_variant(to_lines(script), "b"))

    system_parts = [
        ("Use Case", context.use_case),
        ("Goals", context.goals),
    ]
    user_parts = [
    ]

    if context.previous:
        system_parts.append(("Research Summary", context.research_summary))
        if context.previous.feedback:
            user_parts.append(("Feedback on the previous iteration", to_string(context.previous.feedback)))
        if context.previous.program_output:
            user_parts.append(("Previous iteration code execution output", to_string(context.previous.program_output)))
        if context.previous.code:
            user_parts.append(("Code from the previous iteration", format_code_block(context.previous.code)))
    else:
        # Append at least something to user parts to avoid empty user prompt
        user_parts.append(("Research Summary", context.research_summary))

    for title, content in system_parts:
        system_prompt += f"\n\n# {title}\n{content}"

    prompt_text = system_prompt
    for title, content in user_parts:
        prompt_text += f"\n\n# {title}\n{content}"
    context.save_to("{name}_coder_prompt_{iter}.md", prompt_text, content_name="coder prompt text")

    try:
        print("🚧 Generating code...")
        coder_config=llm_config_coder
        # if context.previous:
        #     coder_config.temperature=0.5
        # code_response = llm_query(prompt, config=coder_config, model=config["coder_model"])
        code_response = llm_query(system_prompt, parts=user_parts, config=coder_config, model=config["coder_model"])
        
        print("🧾 Processing LLM output...")
        # Save JSON response for debugging
        context.save_to("{name}_coder_raw_{iter}.json", code_response["full"].model_dump_json(indent=2), content_name="raw LLM JSON response" )
        context.save_to("{name}_coder_text_{iter}.md", code_response["text"], content_name="raw LLM text")

        # Check if LLM actually executed code
        response_obj = code_response["full"]
        if hasattr(response_obj, 'candidates') and response_obj.candidates:
            parts = response_obj.candidates[0].content.parts
            for part in parts:
                if hasattr(part, 'code_execution_result') and part.code_execution_result:
                    context.current.add_flag('llm_executed')
                    break
            if response_obj.candidates[0].url_context_metadata:
                print("🛈 LLM used URL context tool.")
        if not 'llm_executed' in context.current.flags:
            print("⚠️  WARNING: LLM did not execute code (the code may be non-runnable)")

        text = code_response["text"]
        code_blocks = find_code_blocks(text, delimiter="~~~", language="python")
        diff_blocks = find_code_blocks(text, delimiter="~~~", language="diff")
    except Exception as e:
        print(f"❌ Error during code generation: {e}")
        return False

    if code_blocks:
        context.save_to("{name}_coder_code_{iter}.py", code_blocks[0], content_name="code block")
    if diff_blocks:
        context.save_to("{name}_coder_diff_{iter}.patch", diff_blocks[0], content_name="diff patch")

    if diff_blocks and context.previous and context.previous.code:
        if not code_quality_gate(diff_blocks[0]):
            return False
        patch_lines = clean_code_block(diff_blocks[0])
        print("🛠️ Detected unified diff patch. Applying patch to previous code.")
        prev_code_lines = to_lines(context.previous.code)
        patch_code(prev_code_lines, patch_lines, fuzziness=2)
        context.current.code = prev_code_lines 
    elif code_blocks:
        if not code_quality_gate(code_blocks[0]):
            return False
        context.current.code = clean_code_block(code_blocks[0])  # Now a list
    else:
        context.current.code = []

    context.save_to("{name}_v{iter}.py", context.current.code, content_name="intermediate code")
    return True

def execute(config: dict, context: Context):
    # Execute code locally and get actual program output and/or errors
    sandbox_method = config.get("sandbox_method", "auto")
    commandline_args = config.get("commandline_args", "")

    venv_path = None
    python_packages = None
    if config.get("python_packages"):
        venv_path = f"solutions/venv/{context.filename}"
        python_packages = config.get("python_packages")

    print(f"🖥️  Executing code locally (sandbox: {sandbox_method}, args: {commandline_args if commandline_args else 'none'})...")
    local_exec_result = execute_sandboxed(to_string(context.current.code), method=sandbox_method, args=commandline_args, venv_path=venv_path, extra_packages=python_packages)
    local_exec_success = local_exec_result['success']

    if local_exec_success:
        context.current.add_flag("exec_success")
        actual_method = local_exec_result.get('method', sandbox_method)
        print(f"✅ Local execution successful using method: {actual_method}")
    else:
        print(f"❌ Local execution returned error: {local_exec_result['stderr']}")

    # Check if there were obvious syntax errors
    if "SyntaxError" in local_exec_result['stderr']  or "IndentationError" in local_exec_result['stderr']:
        context.current.add_flag("syntax_error")

    # Save local execution output
    program_output = ["Program exited with code " + str(local_exec_result['exit_code'])]
    program_output.extend(["", "Stdout:", "", "~~~shell"])
    program_output.extend(to_lines(local_exec_result['stdout']))
    program_output.extend(["~~~", "", "Stderr:", "", "~~~shell"])
    program_output.extend(to_lines(local_exec_result['stderr']))
    program_output.extend(["~~~"])

    context.save_to("{name}_v{iter}_output.txt", program_output, content_name="local execution output")
    context.current.program_output = program_output

def fix_syntax_errors(config: dict, context: Context):
    try:
        # Run syntax fix step. The model does not know anything about the goals, it has to merely has to fix syntax issues
        print("\n🚨 SyntaxError or IndentationError detected in program output. Running syntax fix iteration...")
        context.current.add_flag("syntax_fix")
        # load prompt for syntax fix
        syntax_fix_prompt = load_file("scripts/syntax fix.md")
        syntax_fix_prompt_formatted = syntax_fix_prompt.format_map({
            "previous_code": to_string(context.current.code),
            "program_output": to_string(context.current.program_output)
        })
        context.save_to("{name}_syntax_fix_prompt_v{iter}.md", syntax_fix_prompt_formatted, content_name="syntax fix prompt")
        syntax_fix_response = llm_query(syntax_fix_prompt_formatted, model=config["reviewer_model"]) # Coder or utility_model?
        context.save_to("{name}_syntax_fix_response_v{iter}.json", syntax_fix_response["full"].model_dump_json(indent=2), content_name="syntax fix response")
        syntax_fix_text = syntax_fix_response["text"]
        context.save_to("{name}_syntax_fix_response_v{iter}.md", syntax_fix_text, content_name="syntax fix response")
        diff_blocks = find_code_blocks(syntax_fix_text, delimiter="~~~", language="diff")
        if not diff_blocks:
            diff_blocks = find_code_blocks(syntax_fix_text, delimiter="```", language="diff")
        if not diff_blocks:
            print("❌ No diff block found in syntax fix response.")
            return False
    except Exception as e:
        print(f"❌ Error during syntax fix generation: {e}")
        return False

    if diff_blocks:
        print("🛠️ Applying syntax fix diff patch to current code.")
        patch_lines = clean_code_block(diff_blocks[0])
        code_lines = to_lines(context.current.code)
        patch_code(code_lines, patch_lines, fuzziness=2)
        context.current.code = code_lines 
        # Save fixed code
        context.save_to("{name}_v{iter}_syntax_fixed.py", context.current.code, content_name="syntax fixed code")
        return True
    else:
        return False

def feedback(config: dict, context: Context) -> str:
    print("🔍 Evaluating code against the goals...")

    script_path = "scripts/reviewer.md"
    system_prompt = load_file(script_path)

    system_parts = [
        ("Use Case", context.use_case),
        ("Research Summary", context.research_summary)
    ]
    user_parts = [
        ("Goals", context.goals)
    ]    

    if context.current.code:
        user_parts.append(("Code from this iteration", format_code_block(context.current.code)))
    if context.current.program_output:
        user_parts.append(("Code execution output", to_string(context.current.program_output)))
    if context.previous and context.previous.feedback:
        user_parts.append(("Your previous review", to_string(context.previous.feedback)))

    for title, content in system_parts:
        system_prompt += f"\n\n# {title}\n{content}"

    prompt_text = system_prompt
    for title, content in user_parts:
        prompt_text += f"\n\n# {title}\n{content}"
    context.save_to("{name}_review_prompt_{iter}.md", prompt_text, content_name="reviewer prompt text")

    context.current.feedback = llm_query(system_prompt, parts=user_parts, 
                                         config=llm_config_reviewer, model=config["reviewer_model"])["text"]
    if context.current.feedback:
        context.save_to("{name}_review_v{iter}.txt", context.current.feedback, content_name="code review")
        return True
    return False

def goals_met(config: dict, context: Context) -> tuple[bool, int]:
    """
    Uses the LLM to evaluate whether the goals have been met based on the feedback text.
    Returns tuple of (goals_met: bool, score: int).
    """
    script_path = "scripts/goals check.md"
    script = load_file(script_path)
    review_prompt = script.format_map({
        "goals": context.goals,
        "feedback_text": context.current.feedback
    })
    response_text = llm_query(review_prompt, config=llm_config_goals_check, model=config["utility_model"])["text"]
    
    # First try to parse as JSON, then fallback to extracting json code block
    try:
        json_block = to_string(clean_code_block(response_text))
        response_json = json.loads(json_block)
        result = response_json.get("result", "No").lower()
        score = response_json.get("score", 0)
        return (result == "yes", score)
    except json.JSONDecodeError:
        print(f"⚠️  Failed to parse goals check response as JSON, retrying with code block extraction...")

    try:
        json_blocks = find_code_blocks(response_text, delimiter="```", language="json")
        if json_blocks:
            json_block = to_string(clean_code_block(json_blocks[0]))
            response_json = json.loads(json_block)
            result = response_json.get("result", "No").lower()
            score = response_json.get("score", 0)
            return (result == "yes", score)
        else:
            print(f"⚠️  No code blocks found in response")
    except (json.JSONDecodeError, IndexError) as e:
        print(f"⚠️  Failed to parse goals check response as JSON: {response_text}")

    return (False, 0)

def progress_check(context: Context, reset_threshold: int) -> int:
    """ 
    Checks if there is progress in scores.
    If there is no progress over the last 3 iterations, returns the iteration number to return to.
    """
    # Calculate score sequence - use get_score() to handle None
    scores = [x.get_score() for x in context.iterations]
    scores.append(context.current.get_score())
    if len(scores) < reset_threshold:
        return None  # Not enough data to determine
    # Find last best score index (rightmost), and if it's older than reset_threshold iterations, return that index
    best_score = max(scores)
    # Find this score from the right
    best_index = len(scores) - 1 - scores[::-1].index(best_score)
    if best_index < len(scores) - reset_threshold:
        return best_index
    return None


def format_final_code(config: dict, context: Context, token_tracker: TokenUsageTracker) -> list:
    """
    Adds a comment header to the code.
    """
    comment = []
    comment.append(f"# Generated by AI Code Generation Agent")
    comment.append(f"# This Python program implements the following use case:")
    use_case_lines = to_lines(context.use_case)
    for line in use_case_lines:
        comment.append(f"# {line.strip()}")
    comment.append(f"# It shall meet the following goals:")
    goals_lines = to_lines(context.goals)
    for line in goals_lines:
        comment.append(f"# {line.strip()}")
    comment.append(f"# Models used: coder={config['coder_model']}, reviewer={config['reviewer_model']}, utility={config['utility_model']}")
    comment.append(f"# It required {len(context.iterations) + 1} coding rounds to complete.")
    comment.append(f"# Token usage summary:")
    for line in token_tracker.summary():
        comment.append(f"# {line}")
    comment.append("")
    
    code_lines = to_lines(context.current.code)
    return comment + code_lines

def create_filename(basename: str) -> str:
    # Create a filename by appending a random suffix to the basename
    random_suffix = str(random.randint(1000, 9999))
    return f"{basename}_{random_suffix}"

# --- Main Agent Function ---
def run_code_agent(task_config: dict, use_case: str, goals: str, flag_refine_goals: bool = True, flag_diffs: bool = True, reset_threshold: int = 3) -> str:
    max_iterations = task_config["max_rounds"]
    
    print("\n🎯 Use Case:")
    print(use_case)
    print("🎯 Goals:")
    print(goals)

    # Print the task configuration
    print(f"🛠️ Task Configuration: coder_model={task_config['coder_model']}, reviewer_model={task_config['reviewer_model']}, utility_model={task_config['utility_model']}, max_rounds={max_iterations}")

    filename = create_filename(task_config["basename"])
    print(f"🔁 Base name is {filename} for this run")
    context = Context(filename, use_case, goals)

    # Refine the use case and goals before starting (if enabled)
    if flag_refine_goals:
        print("\n🔍 Refining use case and goals before starting...")
        refine_goals(task_config, context)
    else:
        print("\n⏭️  Skipping goals refinement (using original goals)")
    
    # Format goals as a string with a bullet list inside
    context.goals = format_goals(context.goals)

    # Append URLs to the use case if provided in task config
    if "urls" in task_config:
        context.use_case += f"\n\nThe following URLs provide additional context:\n"
        for url in task_config["urls"]:
            context.use_case += f"- {url}\n"

    if "python_packages" in task_config:
        print(f"📦 Additional Python packages to install in sandbox: {task_config['python_packages']}")
        context.use_case += f"\n\nThe following extra Python packages will be available for use: {', '.join(task_config['python_packages'])}\n"

    # Call research step if URLs are provided
    if "urls" in task_config:
        print("\n🔬 Performing research using provided URLs...")
        research(task_config, context)

    for i in range(max_iterations):
        print(f"\n=== 🔁 Iteration {i + 1} of {max_iterations} ===")

        context.start_iteration()

        # Run coding stage
        if not code(task_config, context, use_diffs=flag_diffs):
            context.erase_iteration()
            print("❌ Model generated some bad output, repeating iteration")
            continue

        # Execute code
        execute(task_config, context)

        # If there were syntax errors, run one round of fixing them
        if "syntax_error" in context.current.flags:
            if fix_syntax_errors(task_config, context):
                # If there were successful changes, execute once more
                execute(task_config, context)


        print("\n📤 Submitting code for feedback review...")
        if not feedback(task_config, context):
            print("❌ No feedback received, repeating iteration...")
            context.erase_iteration()
            continue

        done_flag, score = goals_met(task_config, context)
        context.current.score = score

        if done_flag:
            print("✅ LLM confirms goals are met. Stopping iteration.")
            break

        print("🛠️ Goals not fully met. Preparing for next iteration...")
        # Create scores from context
        scores = [x.score for x in context.iterations]
        scores.append(context.current.score)
        print(f"📊 Completion score progression: {scores}")

        if reset_threshold > 0:
            return_to_iteration = progress_check(context, reset_threshold)
            if return_to_iteration is not None:
                context.trim_iterations(return_to_iteration+1)
                if "restarted_from_no_progress" not in context.current.flags:
                    print(f"🔄 No progress detected. Resetting to iteration {return_to_iteration + 1} and continuing from there.")
                    context.current.add_flag("restarted_from_no_progress")
                else:
                    print("⚠️  No progress detected again after restart. Restarting from step 1.")
                    context.trim_iterations(0)

    # Print token usage summary
    token_tracker.print_summary()

    final_code = format_final_code(task_config, context, token_tracker)
    code_filename = f"{filename}.py"
    return save_to_file(code_filename, final_code, content_name="final code")

def run_test():
    query = "Fetch this URL https://en.wikipedia.org/wiki/Code_128 and create a table matching ASCII codes to code sequences for the characters used in Code 128 barcode standard."
    response = llm_query(query, model="gemini-2.5-flash-lite", config=llm_config_coder)
    print("Response:")
    print(response["text"])
    response_obj = response["full"]
    if hasattr(response_obj, 'candidates') and response_obj.candidates:
        if response_obj.candidates[0].url_context_metadata:
            print("🛈 LLM used URL context tool.")


# --- CLI Test Run ---
if __name__ == "__main__":
    print("\n🧠 Welcome to the AI Code Generation Agent")

    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="AI Code Generation Agent")
    parser.add_argument("config_name", help="Configuration name (task directory in tasks/)")
    parser.add_argument("--refine-goals", dest="refine_goals", action="store_true", 
                        help="Refine use case and goals before starting (default)")
    parser.add_argument("--no-refine-goals", dest="refine_goals", action="store_false",
                        help="Skip goals refinement, use original goals as-is")
    parser.add_argument("--diffs", dest="diffs", action="store_true", 
                        help="Use unified diffs for coder output (default)")
    parser.add_argument("--no-diffs", dest="diffs", action="store_false",
                        help="Do not use unified diffs for coder output")
    parser.add_argument("--reset", type=int, help="Number of unsuccessful operations before resetting the the last successful iteration")
    parser.add_argument("--no-reset", dest="reset", action="store_const", const=0,
                        help="Disable resetting on no progress")
    parser.set_defaults(refine_goals=True, diffs=True)
    parser.set_defaults(reset=3)
    args = parser.parse_args()
    config_name = args.config_name

    if not os.path.exists(f"tasks/{config_name}/"):
        print(f"Configuration for '{config_name}' not found in 'tasks/{config_name}/'.")
        sys.exit(1)

    # Load task configuration
    task_config = load_task_config(config_name)
    
    use_case_input = load_file(f"tasks/{config_name}/hl_spec.md")
    goals_input = load_file(f"tasks/{config_name}/ac.md")
    run_code_agent(task_config, use_case_input, goals_input, flag_refine_goals=args.refine_goals, flag_diffs=args.diffs, reset_threshold=args.reset)
    # test = run_test()