extract.py

import json
import os
import re
from openai import AsyncAzureOpenAI
from pydantic import BaseModel
import asyncio
import aiofiles
import time
import datetime
import argparse

class TestExtraction(BaseModel):
    test: str

prompt_template = """You are given an LLM-generated response to a test-input-generation task.

**Your job:** 
Extract the *complete, valid, and effective* test input from the response, even if the response does not strictly follow the expected ` ```plaintext ... ``` ` format.
The test input is usually the actual data or values to be given to the algorithm, not code or explanations.
Sometimes the LLM response may have formatting errors, be missing code blocks, or present the test input in plain text.
Please do your best to identify and extract the correct test input, ignoring code generation code, explanations, or markdown formatting.

**Common formats to look for:**
1. Content within ` ```plaintext ... ``` ` or ` ``` ... ``` ` code blocks
2. Content following **Test Input:** heading until the next heading (like **Explanation:** or **Output:**)
3. Plain text that represents test data/input values

**Requirements:**
- If you cannot find any valid test input in the response, return "None" as the test value.
- The extracted test input should be as complete, valid, and precise as possible, even if the response is imperfectly formatted.
- If there are multiple possible candidates, choose the one that most likely represents the actual input expected by the algorithm problem.
- For **Test Input:** format, extract everything from after the heading until the next markdown heading or explanation section.
- Remove any trailing explanations or comments, keep only the raw input data.

**LLM Response:**
{response}
"""

client = AsyncAzureOpenAI(
)

# concurrent control parameter
MAX_REQUESTS_PER_MINUTE = 20
REQUEST_INTERVAL = 60 / MAX_REQUESTS_PER_MINUTE

# runtime control
MAX_RUNTIME_MINUTES = 600
program_start_time = time.time()
stop_event = asyncio.Event()

# statistics variable
processed_count = 0
error_count = 0

def extract_test_regex(response):
    answer_pattern = r'<answer>\s*(.*?)\s*</answer>'
    answer_match = re.search(answer_pattern, response, re.DOTALL)
    
    if answer_match:
        answer_content = answer_match.group(1).strip()
        
        pattern1 = r'```plaintext\s*(.*?)```'
        match1 = re.search(pattern1, answer_content, re.DOTALL)
        if match1:
            return match1.group(1).strip()
        
        pattern2 = r'```\s*(.*?)```'
        match2 = re.search(pattern2, answer_content, re.DOTALL)
        if match2:
            return match2.group(1).strip()
    
    pattern1 = r'```plaintext\s*(.*?)```'
    match1 = re.search(pattern1, response, re.DOTALL)
    if match1:
        return match1.group(1).strip()
    
    pattern2 = r'```\s*(.*?)```'
    match2 = re.search(pattern2, response, re.DOTALL)
    if match2:
        return match2.group(1).strip()
    
    return None


async def check_runtime():
    while True:
        current_runtime_minutes = (time.time() - program_start_time) / 60
        if current_runtime_minutes >= MAX_RUNTIME_MINUTES:
            print(f"Reach the maximum runtime {MAX_RUNTIME_MINUTES} minutes, stop all new requests")
            stop_event.set()
            break
        await asyncio.sleep(10)

async def extract_test_llm(response, semaphore):
    global processed_count, error_count
    
    if stop_event.is_set():
        return "ERROR"
    
    async with semaphore:
        if stop_event.is_set():
            return "ERROR"
            
        start_time = time.time()
        
        try:
            prompt = prompt_template.format(response=response)
            
            response_obj = await client.responses.parse(
                model="gpt-4.1-mini",
                input=[
                    {"role": "user", "content": prompt}
                ],
                text_format=TestExtraction,
                max_output_tokens=1024
            )
            
            extracted_test = response_obj.output_parsed.test
            
            processed_count += 1
            return extracted_test
                
        except Exception as e:
            error_count += 1
            error_type = type(e).__name__
            
            if "Content Exists Risk" in str(e):
                print("Content Exists Risk. Skipping.")
            else:
                print(f"Error ({error_type}): {e}")
            
            return "ERROR"
        
        finally:
            elapsed = time.time() - start_time
            if elapsed < REQUEST_INTERVAL:
                await asyncio.sleep(REQUEST_INTERVAL - elapsed)

async def extract_test(response, semaphore):
    regex_result = extract_test_regex(response)
    if regex_result:
        return regex_result
    
    return await extract_test_llm(response, semaphore)

async def process_record(record, response_idx, response, model_name, existing_responses, semaphore):
    try:
        new_record = record.copy()
        new_record['extra_info']['model_name'] = model_name
        new_record['extra_info']['response_id'] = response_idx
        
        temp_extra_info = new_record['extra_info'].copy()
        temp_extra_info.pop('response_id', None)
        base_key = create_extra_info_key(temp_extra_info)
        response_key = (base_key, response_idx)
        
        if response_key in existing_responses:
            return None, None
        
        new_record['response'] = response
        del new_record['responses']
        
        extracted_test = await extract_test(response, semaphore)
        
        new_record['test'] = extracted_test
        
        return new_record, response_key
        
    except Exception as e:
        print(f"Error processing record: {e}")
        return None, None

def create_extra_info_key(extra_info):
    key_fields = [
        extra_info.get('problem_id', ''),
        str(extra_info.get('submission_id', '')),
        extra_info.get('task_type', ''),
        extra_info.get('model_name', ''),
        str(extra_info.get('response_id', ''))
    ]
    return '|'.join(key_fields)

def load_existing_records(jsonl_file):
    existing_responses = {}
    
    if os.path.exists(jsonl_file):
        try:
            with open(jsonl_file, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    if line.strip():
                        try:
                            record = json.loads(line.strip())
                            if 'extra_info' in record and record['extra_info']:
                                extra_info = record['extra_info'].copy()
                                response_id = extra_info.pop('response_id', None)
                                model_name = extra_info.pop('model_name', None)
                                
                                if model_name:
                                    extra_info['model_name'] = model_name
                                base_key = create_extra_info_key(extra_info)
                                
                                if base_key and response_id is not None:
                                    response_key = (base_key, response_id)
                                    existing_responses[response_key] = True
                        except json.JSONDecodeError as e:
                            print(f"Line {line_num} JSON parsing error: {e}")
                            continue
            print(f"Read {len(existing_responses)} existing response records")
        except Exception as e:
            print(f"Error reading existing jsonl file: {e}")
    
    return existing_responses


async def extract_responses_to_jsonl(model_name, task_id):
    """
    Asynchronously extract model responses to jsonl file
    """
    global processed_count, error_count
    processed_count = 0
    error_count = 0
    
    model_path = os.path.join(base_path, model_name)
    gen_file = os.path.join(model_path, f"task{task_id}_gen.jsonl")
    ext_file = os.path.join(model_path, f"task{task_id}_ext.jsonl")
    
    if not os.path.exists(gen_file):
        print(f"File not found: {gen_file}")
        return

    try:
        with open(gen_file, 'r', encoding='utf-8') as f:
            records = [json.loads(line.strip()) for line in f]
            
        print(f"Processing model: {model_name}, task: task{task_id}")
        
        os.makedirs(os.path.dirname(ext_file), exist_ok=True)
        
        existing_responses = load_existing_records(ext_file)
        
        semaphore = asyncio.Semaphore(MAX_REQUESTS_PER_MINUTE)
        
        runtime_check_task = asyncio.create_task(check_runtime())
        
        tasks = []
        task_data = []
        
        for record in records:
            responses = record['responses']

            for response_idx, response in enumerate(responses):
                temp_record = record.copy()
                temp_record['extra_info']['model_name'] = model_name
                base_key = create_extra_info_key(temp_record['extra_info'])
                
                response_key = (base_key, response_idx)
                if response_key in existing_responses:
                    continue
                
                task = asyncio.create_task(
                    process_record(record, response_idx, response, model_name, existing_responses, semaphore)
                )
                tasks.append(task)
                task_data.append((response_idx, response_idx))

        if len(tasks) == 0:
            print("No tasks to process")
            return

        print(f"Processing {len(tasks)} tasks")
        print(f"Start time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        
        _, pending = await asyncio.wait(
            tasks, 
            return_when=asyncio.FIRST_COMPLETED if stop_event.is_set() else asyncio.ALL_COMPLETED
        )
        
        completed_tasks = [task for task in tasks if task.done()]
        successful_records = 0
        
        for i, task in enumerate(completed_tasks):
            try:
                result = await task
                new_record, response_key = result
                
                if new_record and response_key:
                    existing_responses[response_key] = True
                    async with aiofiles.open(ext_file, 'a', encoding='utf-8') as f:
                        await f.write(json.dumps(new_record, ensure_ascii=False) + '\n')
                    
                    successful_records += 1
                    
                    if successful_records % 10 == 0:
                        current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        elapsed = time.time() - program_start_time
                        print(f"[{model_name}] [{current_time}] [elapsed: {elapsed:.2f}s] {successful_records} records processed")
                        
            except Exception as e:
                print(f"Error getting task result: {e}")
                
        if pending and stop_event.is_set():
            for task in pending:
                task.cancel()
            
            if pending:
                await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED)
        
        runtime_check_task.cancel()
        try:
            await runtime_check_task
        except asyncio.CancelledError:
            pass
        
        total_runtime = (time.time() - program_start_time) / 60
        print(f"Task completed, total runtime: {total_runtime:.2f} minutes")
        print(f"End time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"{successful_records} records processed, saved to {ext_file}")
        print(f"Success: {processed_count}, Error: {error_count}")

    except Exception as e:
        print(f"Error processing {model_name} task{task_id}: {e}")
        import traceback
        traceback.print_exc()

base_path = "results"

async def main(model):
    print(f"{'='*60}")
    print(f"Processing model: {model}")
    print(f"{'='*60}")
    
    await extract_responses_to_jsonl(model, 1)

    await extract_responses_to_jsonl(model, 2)
    
    print(f"Completed processing {model}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extract test cases from LLM responses")
    parser.add_argument("--model", type=str, required=True, help="Model name")

    args = parser.parse_args()

    asyncio.run(main(args.model))