llm_interface_enhanced.py

import json
import re
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

# Initialize the LLM pipeline (using a free model from Hugging Face)
model_name = "microsoft/DialoGPT-medium"  # Fallback to a smaller model if needed
try:
    # Try to use a more capable model if available
    model_name = "microsoft/DialoGPT-large"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    # Temporarily disable LLM pipeline to use improved fallback logic
    llm_pipeline = None  # pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
except:
    # Fallback to a simpler approach
    try:
        llm_pipeline = None  # pipeline("text-generation", model="gpt2", device=0 if torch.cuda.is_available() else -1)
    except:
        llm_pipeline = None

def generate_initial_robot_design_prompt():
    """Generate the initial prompt for LLM robot design"""
    prompt = """You are an expert robot design AI. Your task is to design a robot that can successfully pass a predefined obstacle.

Obstacle Description: A rectangular block: 0.5m wide, 0.1m deep, 0.05m (5cm) high, located at X=0.75m.

Robot Task: Start at X=0m, drive forward, and cross the obstacle completely.

Available Robot Parameters for your design (provide in JSON format within a 'robot_specs' key):
- "wheel_type": Choose one from ["small_high_grip", "large_smooth", "tracked_base"]
- "body_clearance_cm": An integer between 1 and 10 (cm)
- "approach_sensor_enabled": true or false (For MVP, its effect is conceptual)
- "main_material": Choose one from ["light_plastic", "sturdy_metal_alloy"]

Success Criteria (these are fixed and how your design will be judged):
- robot_crossed_obstacle: Robot's center X-coordinate > 0.8m
- no_significant_collision_with_obstacle: Minimal or no impeding contacts with the obstacle during the pass
- robot_remains_upright: Robot does not fall over

Output Format: Provide your design as a JSON object with keys: robot_design_iteration (start with 1), design_reasoning (your brief explanation), and robot_specs (containing the parameters above).

Example robot_specs: {"wheel_type": "large_smooth", "body_clearance_cm": 7, "approach_sensor_enabled": true, "main_material": "light_plastic"}

Please provide your robot design now:"""
    
    return prompt

def generate_initial_drone_design_prompt():
    """Generate the initial prompt for LLM drone design"""
    prompt = """You are an expert drone design AI. Your task is to design a drone that can successfully fly over a predefined obstacle.

Obstacle Description: A rectangular block: 0.5m wide, 0.1m deep, 0.05m (5cm) high, located at X=0.75m.

Drone Task: Start at X=0m, fly forward, and cross the obstacle completely by flying over it.

Available Drone Parameters for your design (provide in JSON format within a 'robot_specs' key):
- "propeller_size": Choose one from ["small_agile", "medium", "large_stable"]
- "flight_height_cm": An integer between 10 and 50 (cm) - target altitude for crossing
- "stability_mode": Choose one from ["auto_hover", "manual_control"]
- "main_material": Choose one from ["light_carbon_fiber", "sturdy_aluminum"]

Success Criteria (these are fixed and how your design will be judged):
- robot_crossed_obstacle: Drone's center X-coordinate > 0.8m (same key name for compatibility)
- no_significant_collision_with_obstacle: Minimal or no contacts with the obstacle during flight
- robot_remains_upright: Drone maintains stable flight orientation (same key name for compatibility)

Output Format: Provide your design as a JSON object with keys: robot_design_iteration (start with 1), design_reasoning (your brief explanation), and robot_specs (containing the parameters above).

Example robot_specs: {"propeller_size": "medium", "flight_height_cm": 20, "stability_mode": "auto_hover", "main_material": "light_carbon_fiber"}

Please provide your drone design now:"""
    
    return prompt

def generate_iterative_robot_design_prompt(previous_attempt_details, iteration_count):
    """Generate iterative prompt based on previous robot attempt feedback"""
    prompt = f"""Your previous robot design attempt (iteration {iteration_count-1}) had the following specs and outcome:

Previous Specs: {previous_attempt_details['robot_specs']}
Previous Reasoning: {previous_attempt_details['design_reasoning']}
Simulation Feedback: {previous_attempt_details['feedback_from_simulation']}

Please refine your design to meet the success criteria:
- robot_crossed_obstacle: True
- no_significant_collision_with_obstacle: True  
- robot_remains_upright: True

Reminder of available parameters:
- "wheel_type": ["small_high_grip", "large_smooth", "tracked_base"]
- "body_clearance_cm": integer between 1 and 10
- "approach_sensor_enabled": true or false
- "main_material": ["light_plastic", "sturdy_metal_alloy"]

Output Format: Provide your new design as a JSON object with keys: robot_design_iteration (should be {iteration_count}), design_reasoning, and robot_specs.

Please provide your improved robot design now:"""
    
    return prompt

def generate_iterative_drone_design_prompt(previous_attempt_details, iteration_count):
    """Generate iterative prompt based on previous drone attempt feedback"""
    prompt = f"""Your previous drone design attempt (iteration {iteration_count-1}) had the following specs and outcome:

Previous Specs: {previous_attempt_details['robot_specs']}
Previous Reasoning: {previous_attempt_details['design_reasoning']}
Simulation Feedback: {previous_attempt_details['feedback_from_simulation']}

Please refine your design to meet the success criteria:
- robot_crossed_obstacle: True (drone crosses x > 0.8m)
- no_significant_collision_with_obstacle: True (minimal contact during flight)
- robot_remains_upright: True (stable flight orientation)

Reminder of available parameters:
- "propeller_size": ["small_agile", "medium", "large_stable"]
- "flight_height_cm": integer between 10 and 50
- "stability_mode": ["auto_hover", "manual_control"]
- "main_material": ["light_carbon_fiber", "sturdy_aluminum"]

Output Format: Provide your new design as a JSON object with keys: robot_design_iteration (should be {iteration_count}), design_reasoning, and robot_specs.

Please provide your improved drone design now:"""
    
    return prompt

def call_llm_api(prompt_text):
    """Call LLM API and parse response"""
    try:
        # If we have a working LLM pipeline, use it
        if llm_pipeline is not None:
            # Generate response
            response = llm_pipeline(
                prompt_text,
                max_length=len(prompt_text.split()) + 200,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=llm_pipeline.tokenizer.eos_token_id
            )
            
            generated_text = response[0]['generated_text']
            # Extract the part after the prompt
            response_text = generated_text[len(prompt_text):].strip()
            
            # Try to extract JSON from the response
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                try:
                    parsed_response = json.loads(json_str)
                    
                    # Validate required keys
                    required_keys = ['robot_design_iteration', 'design_reasoning', 'robot_specs']
                    if all(key in parsed_response for key in required_keys):
                        return parsed_response
                    else:
                        return generate_fallback_design_response(prompt_text)
                except json.JSONDecodeError:
                    return generate_fallback_design_response(prompt_text)
            else:
                return generate_fallback_design_response(prompt_text)
        else:
            # Use fallback response when LLM is not available
            return generate_fallback_response(prompt_text)
    
    except Exception as e:
        print(f"LLM API call failed: {e}")
        return generate_fallback_response(prompt_text)

def generate_fallback_response(prompt_text):
    """Generate a reasonable fallback response when LLM is not available"""
    # Determine if this is robot or drone based on prompt
    is_drone = "drone" in prompt_text.lower() or "flight" in prompt_text.lower() or "propeller" in prompt_text.lower()
    
    if "iteration 1" in prompt_text.lower() or "previous" not in prompt_text.lower():
        # Initial design
        if is_drone:
            return {
                "robot_design_iteration": 1,
                "design_reasoning": "For the initial drone design, I'm choosing medium propellers for balanced thrust and maneuverability, a flight height of 20cm to safely clear the 5cm obstacle, auto-hover stability mode for better control, and light carbon fiber material for optimal weight-to-strength ratio.",
                "robot_specs": {
                    "propeller_size": "medium",
                    "flight_height_cm": 20,
                    "stability_mode": "auto_hover",
                    "main_material": "light_carbon_fiber"
                }
            }
        else:  # robot
            return {
                "robot_design_iteration": 1,
                "design_reasoning": "For the initial design, I'm choosing large smooth wheels for better obstacle climbing ability, moderate body clearance of 6cm to clear the 5cm obstacle, light plastic material for better mobility, and enabling approach sensors for potential future enhancements.",
                "robot_specs": {
                    "wheel_type": "large_smooth",
                    "body_clearance_cm": 6,
                    "approach_sensor_enabled": True,
                    "main_material": "light_plastic"
                }
            }
    else:
        # Iterative design - analyze feedback and improve
        if "fell_over" in prompt_text.lower() or "not.*stable" in prompt_text.lower():
            # Vehicle fell over or became unstable
            if is_drone:
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Previous drone became unstable. Switching to large stable propellers for better stability, increasing flight height to 25cm for more clearance, maintaining auto-hover mode, and using sturdy aluminum for better stability in flight.",
                    "robot_specs": {
                        "propeller_size": "large_stable",
                        "flight_height_cm": 25,
                        "stability_mode": "auto_hover",
                        "main_material": "sturdy_aluminum"
                    }
                }
            else:  # robot
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Previous robot fell over. Switching to tracked base for better stability, increasing body clearance to 8cm for better obstacle clearance, and using sturdy metal alloy for lower center of mass and better stability.",
                    "robot_specs": {
                        "wheel_type": "tracked_base",
                        "body_clearance_cm": 8,
                        "approach_sensor_enabled": True,
                        "main_material": "sturdy_metal_alloy"
                    }
                }
        elif "failed_to_reach" in prompt_text.lower() or "collided_and_stuck" in prompt_text.lower():
            # Vehicle couldn't reach or got stuck
            if is_drone:
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Previous drone failed to reach or had collision issues. Increasing flight height to 30cm for better obstacle clearance, using small agile propellers for better maneuverability, and sturdy aluminum for durability.",
                    "robot_specs": {
                        "propeller_size": "small_agile",
                        "flight_height_cm": 30,
                        "stability_mode": "auto_hover",
                        "main_material": "sturdy_aluminum"
                    }
                }
            else:  # robot
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Previous robot couldn't reach the obstacle or got stuck. Increasing body clearance to 9cm to ensure obstacle clearance, using small high-grip wheels for better traction, and sturdy material for durability.",
                    "robot_specs": {
                        "wheel_type": "small_high_grip",
                        "body_clearance_cm": 9,
                        "approach_sensor_enabled": True,
                        "main_material": "sturdy_metal_alloy"
                    }
                }
        else:
            # Default iterative improvement
            if is_drone:
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Based on the previous attempt's feedback, I'm adjusting to medium propellers for balanced performance, increasing flight height to 25cm for better obstacle clearance, maintaining auto-hover for stability, and using light carbon fiber for optimal performance.",
                    "robot_specs": {
                        "propeller_size": "medium",
                        "flight_height_cm": 25,
                        "stability_mode": "auto_hover",
                        "main_material": "light_carbon_fiber"
                    }
                }
            else:  # robot
                return {
                    "robot_design_iteration": 2,
                    "design_reasoning": "Based on the previous attempt's feedback, I'm increasing the body clearance to 8cm to ensure better obstacle clearance, switching to tracked base for better traction and stability, and using sturdy metal alloy for better durability during obstacle crossing.",
                    "robot_specs": {
                        "wheel_type": "tracked_base",
                        "body_clearance_cm": 8,
                        "approach_sensor_enabled": True,
                        "main_material": "sturdy_metal_alloy"
                    }
                }

def generate_fallback_design_response(prompt_text=""):
    """Generate a reasonable fallback response when LLM is not available"""
    # Determine if this is robot or drone based on prompt
    is_drone = "drone" in prompt_text.lower() or "flight" in prompt_text.lower() or "propeller" in prompt_text.lower()
    
    if "iteration 1" in prompt_text.lower() or "previous" not in prompt_text.lower():
        # Initial design
        if is_drone:
            return {
                "robot_design_iteration": 1,
                "design_reasoning": "For the initial drone design, I'm choosing medium propellers for balanced thrust and maneuverability, a flight height of 20cm to safely clear the 5cm obstacle, auto-hover stability mode for better control, and light carbon fiber material for optimal weight-to-strength ratio.",
                "llm_interpreted_success_conditions": [
                    "Successfully fly over the obstacle without collision",
                    "Maintain stable flight throughout the process",
                    "Reach the target position beyond the obstacle",
                    "Land safely if required"
                ],
                "robot_specs": {
                    "propeller_size": "medium",
                    "flight_height_cm": 20,
                    "stability_mode": "auto_hover",
                    "main_material": "light_carbon_fiber"
                }
            }
        else:  # robot
            return {
                "robot_design_iteration": 1,
                "design_reasoning": "For the initial design, I'm choosing large smooth wheels for better obstacle climbing ability, moderate body clearance of 6cm to clear the 5cm obstacle, light plastic material for better mobility, and enabling approach sensors for potential future enhancements.",
                "llm_interpreted_success_conditions": [
                    "Successfully cross the obstacle without getting stuck",
                    "Maintain stability and not fall over",
                    "Reach the target position beyond the obstacle",
                    "Complete the crossing efficiently"
                ],
                "robot_specs": {
                    "wheel_type": "large_smooth",
                    "body_clearance_cm": 6,
                    "approach_sensor_enabled": True,
                    "main_material": "light_plastic"
                }
            }
    else:
        # Iterative design - analyze feedback and improve
        return generate_iterative_fallback_response(prompt_text)

def generate_iterative_fallback_response(prompt_text):
    """Generate iterative fallback response based on feedback analysis"""
    is_drone = "drone" in prompt_text.lower() or "flight" in prompt_text.lower() or "propeller" in prompt_text.lower()
    
    # Extract iteration number
    iteration_match = re.search(r"iteration (\d+)", prompt_text.lower())
    iteration = int(iteration_match.group(1)) if iteration_match else 2
    
    # Analyze feedback for improvements
    improvements = []
    reasoning_parts = []
    
    if "failed_to_reach" in prompt_text or "didn't cross" in prompt_text:
        if is_drone:
            improvements.append("Increase flight height for better clearance")
            reasoning_parts.append("increasing flight height to ensure obstacle clearance")
        else:
            improvements.append("Use larger wheels or higher clearance")
            reasoning_parts.append("using larger wheels for better obstacle traversal")
    
    if "fell_over" in prompt_text or "not upright" in prompt_text:
        if is_drone:
            improvements.append("Switch to more stable configuration")
            reasoning_parts.append("using large stable propellers for better stability")
        else:
            improvements.append("Lower center of gravity with heavier material")
            reasoning_parts.append("switching to sturdy metal alloy for better stability")
    
    if "collided" in prompt_text or "stuck" in prompt_text:
        if is_drone:
            improvements.append("Increase flight altitude")
            reasoning_parts.append("flying higher to avoid collision")
        else:
            improvements.append("Increase ground clearance")
            reasoning_parts.append("increasing body clearance to avoid getting stuck")
    
    # Generate improved specs
    if is_drone:
        specs = {
            "propeller_size": "large_stable" if "stability" in prompt_text else "medium",
            "flight_height_cm": min(50, 15 + (iteration * 5)),  # Incrementally increase height
            "stability_mode": "auto_hover",
            "main_material": "sturdy_aluminum" if "stability" in prompt_text else "light_carbon_fiber"
        }
        reasoning = f"For iteration {iteration}, I'm " + ", ".join(reasoning_parts) + " to address the previous failure."
    else:
        specs = {
            "wheel_type": "large_smooth" if iteration <= 3 else "tracked_base",
            "body_clearance_cm": min(10, 4 + iteration),  # Incrementally increase clearance
            "approach_sensor_enabled": True,
            "main_material": "sturdy_metal_alloy" if "stability" in prompt_text else "light_plastic"
        }
        reasoning = f"For iteration {iteration}, I'm " + ", ".join(reasoning_parts) + " to overcome the obstacle."
    
    return {
        "robot_design_iteration": iteration,
        "design_reasoning": reasoning,
        "llm_interpreted_success_conditions": [
            "Successfully cross the obstacle without failure",
            "Maintain stability throughout the process",
            "Reach the target position efficiently",
            "Avoid collision or getting stuck"
        ],
        "robot_specs": specs
    }

def generate_initial_robot_design_prompt_with_criteria(task_description, success_criteria):
    """Generate initial robot design prompt with user-defined criteria"""
    criteria_text = "\n".join([f"- {criterion}" for criterion in success_criteria])
    
    prompt = f"""You are an expert robot design AI. Your task is to design a robot based on the following user requirements:

USER TASK: {task_description}

USER SUCCESS CRITERIA (as interpreted by the system):
{criteria_text}

ENVIRONMENT:
Obstacle: Rectangular block (5cm high, 50cm wide, 10cm deep) at x=0.75m
Robot starts at x=0m and must traverse forward

AVAILABLE ROBOT PARAMETERS (provide in JSON format within 'robot_specs'):
- "wheel_type": ["small_high_grip", "large_smooth", "tracked_base"]
- "body_clearance_cm": integer 1-10 (ground clearance in cm)
- "approach_sensor_enabled": true/false
- "main_material": ["light_plastic", "sturdy_metal_alloy"]

REQUIRED OUTPUT FORMAT:
{{
  "robot_design_iteration": 1,
  "design_reasoning": "Your detailed explanation of design choices",
  "llm_interpreted_success_conditions": ["condition 1", "condition 2", ...],
  "robot_specs": {{
    "wheel_type": "your_choice",
    "body_clearance_cm": your_number,
    "approach_sensor_enabled": your_boolean,
    "main_material": "your_choice"
  }}
}}

Please provide your robot design now:"""
    
    return prompt

def generate_initial_drone_design_prompt_with_criteria(task_description, success_criteria):
    """Generate initial drone design prompt with user-defined criteria"""
    criteria_text = "\n".join([f"- {criterion}" for criterion in success_criteria])
    
    prompt = f"""You are an expert drone design AI. Your task is to design a drone based on the following user requirements:

USER TASK: {task_description}

USER SUCCESS CRITERIA (as interpreted by the system):
{criteria_text}

ENVIRONMENT:
Obstacle: Rectangular block (5cm high, 50cm wide, 10cm deep) at x=0.75m
Drone starts at x=0m and must fly over/around the obstacle

AVAILABLE DRONE PARAMETERS (provide in JSON format within 'robot_specs'):
- "propeller_size": ["small_agile", "medium", "large_stable"]
- "flight_height_cm": integer 10-50 (target flight altitude)
- "stability_mode": ["auto_hover", "manual_control"]
- "main_material": ["light_carbon_fiber", "sturdy_aluminum"]

REQUIRED OUTPUT FORMAT:
{{
  "robot_design_iteration": 1,
  "design_reasoning": "Your detailed explanation of design choices",
  "llm_interpreted_success_conditions": ["condition 1", "condition 2", ...],
  "robot_specs": {{
    "propeller_size": "your_choice",
    "flight_height_cm": your_number,
    "stability_mode": "your_choice",
    "main_material": "your_choice"
  }}
}}

Please provide your drone design now:"""
    
    return prompt

def generate_iterative_robot_design_prompt_with_criteria(previous_attempt_details, iteration_count, success_criteria):
    """Generate iterative robot design prompt with user-defined criteria"""
    criteria_text = "\n".join([f"- {criterion}" for criterion in success_criteria])
    
    prompt = f"""Your previous robot design attempt (iteration {iteration_count-1}) had the following specs and outcome:

Previous Specs: {previous_attempt_details['vehicle_specs']}
Previous Reasoning: {previous_attempt_details['design_reasoning']}
Simulation Feedback: {previous_attempt_details['feedback_from_simulation']}
LLM Success Conditions: {previous_attempt_details.get('llm_success_conditions', [])}

USER SUCCESS CRITERIA (as interpreted by the system):
{criteria_text}

Please refine your design to meet these criteria. Focus on addressing the specific failures from the previous attempt.

Reminder of available parameters:
- "wheel_type": ["small_high_grip", "large_smooth", "tracked_base"]
- "body_clearance_cm": integer between 1 and 10
- "approach_sensor_enabled": true or false
- "main_material": ["light_plastic", "sturdy_metal_alloy"]

Output Format: Provide your new design as a JSON object with keys: 
- robot_design_iteration (should be {iteration_count})
- design_reasoning (explain your improvements)
- llm_interpreted_success_conditions (your understanding of what success means)
- robot_specs (the parameters)

Please provide your improved robot design now:"""
    
    return prompt

def generate_iterative_drone_design_prompt_with_criteria(previous_attempt_details, iteration_count, success_criteria):
    """Generate iterative drone design prompt with user-defined criteria"""
    criteria_text = "\n".join([f"- {criterion}" for criterion in success_criteria])
    
    prompt = f"""Your previous drone design attempt (iteration {iteration_count-1}) had the following specs and outcome:

Previous Specs: {previous_attempt_details['vehicle_specs']}
Previous Reasoning: {previous_attempt_details['design_reasoning']}
Simulation Feedback: {previous_attempt_details['feedback_from_simulation']}
LLM Success Conditions: {previous_attempt_details.get('llm_success_conditions', [])}

USER SUCCESS CRITERIA (as interpreted by the system):
{criteria_text}

Please refine your design to meet these criteria. Focus on addressing the specific failures from the previous attempt.

Reminder of available parameters:
- "propeller_size": ["small_agile", "medium", "large_stable"]
- "flight_height_cm": integer between 10 and 50
- "stability_mode": ["auto_hover", "manual_control"]
- "main_material": ["light_carbon_fiber", "sturdy_aluminum"]

Output Format: Provide your new design as a JSON object with keys:
- robot_design_iteration (should be {iteration_count})
- design_reasoning (explain your improvements)
- llm_interpreted_success_conditions (your understanding of what success means)
- robot_specs (the parameters)

Please provide your improved drone design now:"""
    
    return prompt