mod.py

import requests
import json
import random
import time
import os
import ast
import sys

# --- Configuration ---
OLLAMA_HOST = "http://localhost:11434"
OLLAMA_API_URL = f"{OLLAMA_HOST}/api/generate"

# --- ANSI Color Codes for Rich CLI Output ---
class Color:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# --- Pre-defined Task Library ---
TaskManager = {
    'sort_list': {
        "description": "Write a Python function `sort_list(numbers)` that sorts a list of numbers in ascending order without using `sorted()` or `.sort()`.",
        "problem_type": "function",
        "name": "sort_list",
        "test_cases": [
            {'input': [[3, 1, 4, 1, 5, 9]], 'expected': [1, 1, 3, 4, 5, 9]}, # Corrected expected output
            {'input': [[]], 'expected': []},
            {'input': [[5, 4, 3, 2, 1]], 'expected': [1, 2, 3, 4, 5]}
        ]
    },
    'find_primes': {
        "description": "Write a Python function `find_primes(n)` that returns a list of all prime numbers up to `n`.",
        "problem_type": "function",
        "name": "find_primes",
        "test_cases": [
            {'input': [20], 'expected': [2, 3, 5, 7, 11, 13, 17, 19]},
            {'input': [2], 'expected': [2]},
            {'input': [1], 'expected': []}
        ]
    },
}

# --- Evaluation Engine ---

def evaluate_function_solution(code, function_name, test_cases):
    """Evaluates a single function solution based on passed test cases."""
    if not test_cases: return 0.0
    passed_cases = 0
    for tc in test_cases:
        try:
            local_scope = {}
            exec(code, globals(), local_scope)
            func = local_scope.get(function_name)
            if not callable(func): continue

            result = func(*tc['input'])
            
            is_list_of_lists = isinstance(result, list) and all(isinstance(elem, list) for elem in result)
            if is_list_of_lists:
                # Order-independent comparison for lists of lists
                sorted_result = sorted(map(sorted, result))
                sorted_expected = sorted(map(sorted, tc['expected']))
                if sorted_result == sorted_expected:
                    passed_cases += 1
            elif result == tc['expected']:
                passed_cases += 1
        except Exception:
            pass # A failed execution means the test case is not passed
    return passed_cases / len(test_cases)

def evaluate_class_solution(code, class_name, test_cases):
    """Evaluates a class-based solution by running sequences of method calls."""
    if not test_cases: return 0.0
    passed_scenarios = 0
    for tc in test_cases:
        try:
            local_scope = {}
            exec(code, globals(), local_scope)
            SolutionClass = local_scope.get(class_name)
            if not SolutionClass: continue

            instance = SolutionClass(*tc['constructor_args'])
            
            scenario_correct = True
            for op in tc['operations']:
                method = getattr(instance, op['method'])
                result = method(*op['args'])
                if result != op['expected']:
                    scenario_correct = False
                    break # One wrong operation fails the whole scenario
            
            if scenario_correct:
                passed_scenarios += 1
        except Exception:
            # Any exception during a scenario (init or method call) means it fails.
            pass
            
    return passed_scenarios / len(test_cases)

def master_evaluator(code, task_details):
    """Delegates evaluation based on the problem type."""
    problem_type = task_details.get("problem_type")
    name = task_details.get("name")
    test_cases = task_details.get("test_cases")
    
    if not all([problem_type, name, test_cases, code]):
        return 0.0

    if problem_type == "function":
        return evaluate_function_solution(code, name, test_cases)
    elif problem_type == "class":
        return evaluate_class_solution(code, name, test_cases)
    return 0.0

# --- Ollama API Call ---
def call_ollama(prompt, model_name, output_format=""):
    max_retries = 3
    payload = {"model": model_name, "prompt": prompt, "stream": False, "options": {"temperature": 0.5}}
    if output_format == "json":
        payload["format"] = "json"

    for attempt in range(max_retries):
        try:
            response = requests.post(OLLAMA_API_URL, json=payload, timeout=180)
            response.raise_for_status()
            response_text = response.json().get('response', '')
            if output_format == "json":
                return response_text
            if '```python' in response_text:
                return response_text.split('```python')[1].split('```')[0].strip()
            if '```' in response_text:
                return response_text.split('```')[1].split('```')[0].strip()
            return response_text.strip()
        except requests.exceptions.RequestException as e:
            if attempt >= max_retries - 1: raise e
            time.sleep((2 ** attempt) + random.uniform(0, 1))
    return ""

# --- Evolutionary Algorithm Core ---
def generate_solution(task, model_name):
    problem_type = task.get("problem_type", "function")
    prompt = f"You are an expert Python programmer. Solve the following problem by writing a single, complete Python {problem_type}. Provide only the raw Python code. Do not include explanations, examples, or markdown.\n\nProblem: {task['description']}\n{problem_type.capitalize()} Name: `{task['name']}`"
    return call_ollama(prompt, model_name)

def mutate(solution, task, model_name):
    problem_type = task.get("problem_type", "function")
    prompt = f"You are an expert Python code optimizer. Improve the following Python {problem_type} by applying a significant mutation (e.g., fix a bug, refactor, try a new algorithm). Provide only the raw, mutated Python code.\n\nProblem: {task['description']}\n\n{problem_type.capitalize()} to Mutate:\n```python\n{solution}\n```"
    return call_ollama(prompt, model_name)

def crossover(p1, p2, task, model_name):
    problem_type = task.get("problem_type", "function")
    prompt = f"You are an expert Python code synthesizer. Merge the two parent solutions below into a single, superior child solution, combining their best aspects. Provide only the raw Python code for the new {problem_type}.\n\nProblem: {task['description']}\n\nParent 1:\n```python\n{p1}\n```\n\nParent 2:\n```python\n{p2}\n```"
    return call_ollama(prompt, model_name)

# --- Main Evolution Loop ---
def run_evolution(task, model_name):
    POPULATION_SIZE = 8
    NUM_GENERATIONS = 5
    TOURNAMENT_SIZE = 3
    CROSSOVER_CHANCE = 0.5

    print(f"{Color.OKBLUE}🌱 Generating initial population of {POPULATION_SIZE} solutions...{Color.ENDC}")
    population = [sol for sol in [generate_solution(task, model_name) for _ in range(POPULATION_SIZE)] if sol]

    best_overall_solution = {"code": "# No solution found.", "fitness": 0.0}

    for gen in range(NUM_GENERATIONS):
        print(f"\n{Color.HEADER}--- 🧬 Generation {gen + 1}/{NUM_GENERATIONS} ---{Color.ENDC}")
        
        fitness_scores = {code: master_evaluator(code, task) for code in population}
        
        if not fitness_scores:
            print(f"{Color.WARNING}  Population is empty or failed evaluation. Ending evolution.{Color.ENDC}")
            break

        best_this_gen_code = max(fitness_scores, key=fitness_scores.get)
        best_this_gen_fitness = fitness_scores[best_this_gen_code]
        
        print(f"  Best fitness this generation: {Color.OKGREEN}{best_this_gen_fitness:.2f}{Color.ENDC}")

        if best_this_gen_fitness > best_overall_solution["fitness"]:
            best_overall_solution = {"code": best_this_gen_code, "fitness": best_this_gen_fitness}
            print(f"  {Color.OKGREEN}🏆 New overall best solution found!{Color.ENDC}")

        if best_overall_solution["fitness"] == 1.0:
            print(f"\n{Color.OKGREEN}{Color.BOLD}✅ Perfect solution found! Evolution complete.{Color.ENDC}")
            break
        
        next_generation = [best_this_gen_code] # Elitism

        while len(next_generation) < POPULATION_SIZE:
            p1 = max(random.sample(population, k=TOURNAMENT_SIZE), key=lambda i: fitness_scores.get(i, 0.0))
            p2 = max(random.sample(population, k=TOURNAMENT_SIZE), key=lambda i: fitness_scores.get(i, 0.0))
            
            child = crossover(p1, p2, task, model_name) if random.random() < CROSSOVER_CHANCE else mutate(p1, task, model_name)
            next_generation.append(child if child else p1)
        
        population = next_generation

    print(f"\n{Color.HEADER}--- 🏁 Evolution Finished ---{Color.ENDC}")
    print(f"{Color.BOLD}Best overall fitness score: {best_overall_solution['fitness']:.2f}{Color.ENDC}")
    print(f"{Color.BOLD}Best solution found:{Color.ENDC}")
    print(f"{Color.OKCYAN}{'-'*50}\n{best_overall_solution['code']}\n{'-'*50}{Color.ENDC}")

# --- CLI Interaction ---
def get_custom_task(model_name):
    print(f"\n{Color.HEADER}--- ✍️  Define a Custom Task ---{Color.ENDC}")
    print(f"{Color.OKCYAN}Paste your full task description. The model will extract the details.{Color.ENDC}")
    print(f"On a new line, type {Color.BOLD}'DONE'{Color.ENDC} when you are finished.")

    lines = []
    while True:
        try:
            line = input()
            if line.strip().upper() == 'DONE':
                break
            lines.append(line)
        except EOFError:
            break
    full_text = "\n".join(lines)

    if not full_text.strip(): return None

    parsing_prompt = f"""
You are an expert text analysis tool for programming problems. Your task is to parse a problem description into a structured JSON object.

1.  **Determine Problem Type**: Analyze if the solution is a single function or a class with methods. Set "problem_type" to "function" or "class".
2.  **Extract Name**: Find the function or class name (e.g., `findLadders`, `LRUCache`).
3.  **Extract Test Cases**: This is the most critical step.
    * For **functions**, the "input" value must be a list of arguments. Example: `{{"input": [arg1, arg2], "expected": "value"}}`.
    * For **classes**, parse the sequence of operations. The "constructor_args" should be a list. Each step in "operations" must have a "method", "args" (as a list), and "expected" output.
    * **IMPORTANT**: The JSON `null` should be used for expected outputs that are `null` in the problem description.

Analyze the following problem description:
---
{full_text}
---

Respond with a single, valid JSON object and nothing else.

**Function Example Format**:
`{{
  "problem_type": "function",
  "name": "functionName",
  "test_cases": [
    {{"input": ["arg1", "arg2"], "expected": "some_value"}}
  ]
}}`

**Class Example Format**:
`{{
  "problem_type": "class",
  "name": "ClassName",
  "test_cases": [
    {{
      "constructor_args": [2],
      "operations": [
        {{"method": "put", "args": [1, 1], "expected": null}},
        {{"method": "get", "args": [1], "expected": 1}}
      ]
    }}
  ]
}}`
"""
    
    try:
        print(f"\n{Color.OKBLUE}🤖 Parsing the task description with '{model_name}'...{Color.ENDC}")
        json_response_str = call_ollama(parsing_prompt, model_name, output_format="json")
        parsed_data = json.loads(json_response_str)

        if "name" not in parsed_data or "problem_type" not in parsed_data:
            raise ValueError("Parsed JSON is missing required fields.")

        print(f"{Color.OKGREEN}✅ Successfully parsed custom task!{Color.ENDC}")
        parsed_data["description"] = full_text
        return parsed_data

    except Exception as e:
        print(f"{Color.FAIL}Error: The LLM failed to parse the task description correctly. Details: {e}{Color.ENDC}")
        return None

def main():
    os.system('cls' if os.name == 'nt' else 'clear')
    print(f"{Color.HEADER}{Color.BOLD}Welcome to AlphaEvolve CLI{Color.ENDC}")
    print("An evolutionary coding agent powered by local Ollama models.")
    
    # --- MODIFICATION START ---
    # Automatically detect available Ollama models
    available_models = []
    try:
        # Check connection and fetch models by querying the tags endpoint
        response = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        response.raise_for_status()
        print(f"{Color.OKGREEN}🔗 Ollama server connection successful.{Color.ENDC}")
        data = response.json()
        # Get the list of model names from the response
        available_models = sorted([model['name'] for model in data.get('models', [])])
    except requests.exceptions.RequestException:
        print(f"{Color.FAIL}Error: Could not connect to Ollama server at {OLLAMA_HOST}. Please ensure it's running.{Color.ENDC}")
        return
    except json.JSONDecodeError:
        print(f"{Color.FAIL}Error: Received an invalid response when fetching models from Ollama.{Color.ENDC}")
        return

    if not available_models:
        print(f"{Color.WARNING}Warning: No models found on the Ollama server.{Color.ENDC}")
        print(f"{Color.OKCYAN}Please pull a model first (e.g., `ollama pull llama3:8b`){Color.ENDC}")
        return

    print(f"\n{Color.BOLD}1. Select an Ollama Model:{Color.ENDC}")
    for i, model_name in enumerate(available_models):
        print(f"  [{i+1}] {model_name}")
        
    selected_model = None
    while selected_model is None:
        try:
            model_choice_str = input("Enter model number: ")
            choice_idx = int(model_choice_str) - 1
            if 0 <= choice_idx < len(available_models):
                selected_model = available_models[choice_idx]
            else:
                print(f"{Color.FAIL}Invalid number. Please select from the list.{Color.ENDC}")
        except ValueError:
            print(f"{Color.FAIL}Invalid input. Please enter a number.{Color.ENDC}")
    # --- MODIFICATION END ---

    print(f"\n{Color.BOLD}2. Select a Task:{Color.ENDC}")
    task_keys = list(TaskManager.keys())
    for i, key in enumerate(task_keys):
        print(f"  [{i+1}] {key.replace('_', ' ').title()}")
    print(f"  [{len(task_keys) + 1}] Custom Task...")
    
    task_details = None
    while not task_details:
        task_choice_str = input("Enter task number: ")
        try:
            choice = int(task_choice_str)
            if 1 <= choice <= len(task_keys):
                task_details = TaskManager[task_keys[choice - 1]]
            elif choice == len(task_keys) + 1:
                task_details = get_custom_task(selected_model)
            else:
                print(f"{Color.FAIL}Invalid choice. Please try again.{Color.ENDC}")
        except (ValueError, IndexError):
            print(f"{Color.FAIL}Invalid input. Please try again.{Color.ENDC}")

    task_details["evaluator"] = lambda code: master_evaluator(code, task_details)
    
    print(f"\n{Color.OKCYAN}🚀 Starting evolution for task '{task_details['name']}' using model '{selected_model}'...{Color.ENDC}")
    run_evolution(task_details, selected_model)

if __name__ == "__main__":
    main()