fuzzing_engine.py

import time
from groq import Groq
from tqdm import tqdm
import json
from config import *
from data.benchmark_prompts import get_all_prompts

class LLMFuzzingEngine:
    def __init__(self):
        if not GROQ_API_KEY:
            raise ValueError("GROQ_API_KEY not found in environment")
        
        self.client = Groq(api_key=GROQ_API_KEY)
        self.model = MODEL_NAME
        
    def run_single_inference(self, prompt, temperature):
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a helpful assistant. Provide concise, accurate answers."
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=temperature,
                max_tokens=MAX_TOKENS,
                top_p=1.0,
            )
            
            return response.choices[0].message.content
        
        except Exception as e:
            print(f"\n Error during inference: {e}")
            return None
    
    def fuzz_single_prompt(self, prompt_data, temperature, run_id):
        output = self.run_single_inference(prompt_data['prompt'], temperature)
        
        return {
            'prompt_id': prompt_data['id'],
            'category': prompt_data['category'],
            'prompt': prompt_data['prompt'],
            'temperature': temperature,
            'run_id': run_id,
            'output': output,
            'timestamp': time.time()
        }
    
    def run_full_experiment(self):
        prompts = get_all_prompts()
        total_experiments = len(prompts) * len(TEMPERATURES) * RUNS_PER_CONDITION
        
        print(f"\n{'='*60}")
        print(f"LLM FUZZ-BENCH: Systematic Consistency Analysis")
        print(f"{'='*60}")
        print(f"Model: {self.model}")
        print(f"Prompts: {len(prompts)} ({len(set(p['category'] for p in prompts))} categories)")
        print(f"Temperatures: {TEMPERATURES}")
        print(f"Runs per condition: {RUNS_PER_CONDITION}")
        print(f"Total experiments: {total_experiments}")
        print(f"{'='*60}\n")
        
        results = []
        
        with tqdm(total=total_experiments, desc="Running experiments") as pbar:
            for prompt_data in prompts:
                for temperature in TEMPERATURES:
                    for run_id in range(RUNS_PER_CONDITION):
                        result = self.fuzz_single_prompt(prompt_data, temperature, run_id)
                        results.append(result)
                        pbar.update(1)
                        
                        time.sleep(0.1)
        
        output_file = f"{RESULTS_DIR}/raw_results.json"
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        
        print(f"\n Experiment complete! Results saved to {output_file}")
        return results

if __name__ == "__main__":
    engine = LLMFuzzingEngine()
    results = engine.run_full_experiment()