code_modifier.py

from typing import Dict, List, Set, Optional
from dataclasses import dataclass
from enum import Enum
import openai
from contextlib import contextmanager
import re
import json

@dataclass
class UIComponent:
    name: str
    type: str
    config: Dict

@dataclass
class PromptConfig:
    name: str
    inputs: List[str]
    outputs: List[str]
    ui_components: Dict[str, UIComponent]
    step: Optional[str] = None

class ChangeType(Enum):
    ADD_PROMPT = "add_prompt"
    REMOVE_PROMPT = "remove_prompt"
    RENAME_PROMPT = "rename_prompt"
    ADD_INPUT = "add_input"
    REMOVE_INPUT = "remove_input"
    RENAME_INPUT = "rename_input"
    ADD_OUTPUT = "add_output"
    REMOVE_OUTPUT = "remove_output"
    RENAME_OUTPUT = "rename_output"
    ADD_UI = "add_ui"
    REMOVE_UI = "remove_ui"
    RENAME_UI = "rename_ui"
    CHANGE_STEP = "change_step"

@dataclass
class ConfigChange:
    type: ChangeType
    prompt_name: str
    old_value: Optional[str] = None
    new_value: Optional[str] = None
    details: Optional[Dict] = None

@dataclass
class AnalysisResult:
    changes: List[ConfigChange]
    affected_files: Set[str]
    required_updates: Dict[str, List[Dict]]
    
@contextmanager
def openai_session():
    """Context manager to properly handle OpenAI API sessions"""
    try:
        client = openai.OpenAI()
        yield client
    finally:
        if hasattr(client, 'close'):
            client.close()

def call_o1_mini(prompt: str) -> str:
    """Call the o1-mini model with the given prompt"""
    with openai_session() as client:
        try:
            response = client.chat.completions.create(
                model="o1-mini",
                messages=[{"role": "user", "content": prompt}]
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error generating output: {str(e)}"
        
def clean_ai_response(response: str) -> str:
    """Clean the AI response to ensure valid JSON"""
    try:
        # Remove any potential markdown code block markers
        response = re.sub(r'```json\s*', '', response)
        response = re.sub(r'```\s*$', '', response)
        
        # Remove any leading/trailing whitespace
        response = response.strip()
        
        # Attempt to parse JSON to validate
        json.loads(response)
        
        return response
    except json.JSONDecodeError as e:
        print(f"Error cleaning response: {str(e)}")
        print("Raw response:", response)
        return "{}"

def analyze_config_structure(config_content: str, project_content: str, current_handlers: str) -> Dict:
    """First AI: Analyzes the current config structure and determines required changes"""
    analysis_prompt = f"""
    You are an expert software engineer performing a comprehensive analysis.You must analyze the current state of the code WITHOUT relying on any previous context or memory.
    
    IMPORTANT RULES:
    1. ONLY analyze the code provided in this prompt
    2. DO NOT reference or recall any previous analyses
    3. DO NOT suggest changes based on previous versions
    4. ONLY compare against the current files provided
    5. IGNORE any remembered patterns or previous suggestions
    6. The config file (page_prompts_config.py) is the source of truth for naming conventions
    7. Component names MUST match their corresponding config names EXACTLY
    
    Your task is to analyze and focus on the configuration structure and validate the relationships between components.

    1. OUTPUT AND INPUT VALIDATION RULES:
       - Not all outputs/inputs require UI component mappings
       - Special outputs like 'quotation_cost' can be internal values used as inputs and outputs in other steps
       - Internal values can be passed between steps without UI components , so just DO NOT suggest any changes for them.
       - Examples of internal values:
         * quotation_cost
         * project_detail
         * gathered_project_input
         * reviewed_project_input
         * client_follow_up_questions 
         * follow_up_questions
         * client_followup_question_sample_answers

    2. OUTPUT TO UI COMPONENT MAPPING RULES:
       - Most outputs in config map to specific UI components
       - Standard Output Pattern:
         * Base output name MUST map to TWO components:
           - base_name_text: For textbox display
           - base_name_markdown: For markdown display
       - Special Case - Mandays Pattern:
         * Mandays outputs MUST map to ONE component:
           - base_name_dataframe: For dataframe display
           
    3. INPUT TO UI COMPONENT MAPPING RULES:
       - Only UI-interactive inputs need component mappings
       - When needed, inputs MUST reference either the _text or _markdown variant

    4. NAMING CONVENTION (Source of Truth: Config File)
       MUST CHECK:
       - Prompt names 
       - Component names 
       - Input/Output references
       
       WHERE TO CHECK:
       - event_handlers.py: Function names, component references
       - Project.py: Method names, mappings
       - all_components dictionary: Component keys
       
    5. HANDLER VALIDATION RULES:
       - Each step handler must correctly map:
         * Inputs to component _text or _markdown
         * Outputs to both _text and _markdown (or _dataframe for mandays)
       - all_components dictionary keys must exactly match config naming
       - Function parameters should align with Project.py signatures
    
    6. PROMPTCONFIG CHANGES:
       A. REMOVED PROMPTCONFIGS:
          - When a PromptConfig is removed from config:
            * ALL references to that prompt must be removed from Project.py
            * ALL UI components associated with that prompt must be removed from event_handlers.py
            * ALL method signatures and return values must be updated accordingly
            * ALL dependent code that referenced the removed components must be updated
          - Check for:
            * Orphaned UI components
            * Unused function parameters
            * Outdated return values
            * Broken event handler mappings

       B. NEW PROMPTCONFIGS:
          - When a new PromptConfig is added to config:
            * Add corresponding execute_prompt() call in Project.py
            * Add UI component references in event_handlers.py
            * Add new step button handler if step is specified
            * Add necessary input/output mappings
          - Check for:
            * Required method signatures
            * Input/output component creation
            * Event handler connections
            * Step button integration
            * State management updates


    Context Files:
    1. Current Config:
    {config_content}

    2. Project.py:
    {project_content}

    3. event_handlers.py:
    {current_handlers}
    
    CRITICAL ANALYSIS RULES:
    1. RENAMED PROMPTS:
       - First check if a component missing from config exists with a similar name
       - Compare component names using similarity matching (e.g., generate_Tech_SOW vs generate_Technical_SOW)
       - If a similar name exists in config, treat it as a rename operation
       - Add it to the 'renamed_prompts' section with old and new names
       - Do NOT mark it as a removed prompt
       
    2. REMOVED PROMPTS:
       - Only mark a prompt as removed if NO similar name exists in config
       - If unsure, prefer rename over removal
       - For confirmed removals:
         * List ALL code locations that need cleanup
         * Include exact file paths and line numbers if possible
         * Specify what code needs to be removed

    3. NAMING UPDATES:
       - When a component name changes in config:
         * Update ALL references to match exactly
         * Maintain consistent casing and formatting
         * Update both component names and method references
    
    Return a JSON analysis with this EXACT structure as shown as an example below:
    {{
        "component_mappings": {{
            "outputs": [
                {{
                    "config_name": "base_output_name",
                    "required_components": ["component1", "component2"],
                    "current_components": ["existing1", "existing2"],
                    "is_valid": false,
                    "issues": ["detailed_issue_description"]
                }}
            ],
            "inputs": [
                {{
                    "config_name": "input_name",
                    "required_component": "required_component_name",
                    "current_component": "current_component_name",
                    "is_valid": false,
                    "issues": ["detailed_issue_description"]
                }}
            ],
            "buttons": [
                {{
                    "button_id": "button_name",
                    "type": "step|action|recalculate|upload",
                    "handler": "handler_function_name",
                    "inputs": ["input1", "input2"],
                    "outputs": ["output1", "output2"],
                    "state_updates": ["state1", "state2"],
                    "is_valid": false,
                    "issues": ["detailed_issue_description"]
                }}
            ]
        }},
       "event_handlers": {{
            "handlers": [
                {{
                    "button_id": "button_name",
                    "type": "step|recalculate|upload",
                    "input_mappings": {{
                        "is_valid": false,
                        "issues": ["detailed_issue_description"]
                    }},
                    "output_mappings": {{
                        "is_valid": false,
                        "issues": ["detailed_issue_description"]
                    }},
                    "state_mappings": {{
                        "is_valid": false,
                        "issues": ["detailed_issue_description"]
                    }}
                }}
            ]
        }},
        "prompt_changes": {{
            "removed_prompts": [
                {{
                    "prompt_name": "name_of_removed_prompt",
                    "affected_components": ["component1", "component2"],
                    "affected_handlers": ["handler1", "handler2"],
                    "cleanup_required": {{
                        "Project.py": [
                            {{
                                "type": "method_removal|signature_update",
                                "location": "exact_location",
                                "code_to_remove": "code_snippet",
                                "reason": "detailed_explanation"
                            }}
                        ],
                        "event_handlers.py": [
                            {{
                                "type": "component_removal|handler_update",
                                "location": "exact_location",
                                "code_to_remove": "code_snippet",
                                "reason": "detailed_explanation"
                            }}
                        ]
                    }}
                }}
            ],
              "renamed_prompts": [
                {{
                    "old_name": "old_prompt_name",
                    "new_name": "new_prompt_name",
                    "affected_files": {{
                        "Project.py": [
                            {{
                                "type": "method_rename",
                                "location": "method references",
                                "old_code": "old_code_snippet",
                                "new_code": "new_code_snippet",
                                "reason": "Prompt name updated in config"
                            }}
                        ],
                        "event_handlers.py": [
                            {{
                                "type": "component_rename",
                                "location": "component references",
                                "old_code": "old_code_snippet",
                                "new_code": "new_code_snippet",
                                "reason": "Prompt name updated in config"
                            }}
                        ]
                    }}
                }}
            ]
              "new_prompts": [
                {{
                    "prompt_name": "name_of_new_prompt",
                    "required_components": ["component1", "component2"],
                    "required_handlers": ["handler1", "handler2"],
                    "additions_required": {{
                        "Project.py": [
                            {{
                                "type": "method_addition|signature_update",
                                "location": "exact_location",
                                "code_to_add": "code_snippet",
                                "reason": "detailed_explanation"
                            }}
                        ],
                        "event_handlers.py": [
                            {{
                                "type": "component_addition|handler_creation",
                                "location": "exact_location",
                                "code_to_add": "code_snippet",
                                "reason": "detailed_explanation"
                            }}
                        ]
                    }}
                }}
            ]
        }},
        "required_updates": {{
            "Project.py": [
                {{
                    "type": "method|mapping",
                    "location": "exact_location",
                    "reason": "detailed_explanation"
                }}
            ],
            "event_handlers.py": [
                {{
                    "button_id": "button_name",
                    "type": "step|recalculate|upload",
                    "current": "current_code",
                    "required": "required_code",
                    "reason": "detailed_explanation"
                }}
            ]
        }}
    }}

    IMPORTANT:
    1. Validate EVERY output has correct component mappings
    2. Check EVERY input references correct component variant
    3. Verify ALL component names match exactly
    4. Flag ANY case mismatches or naming inconsistencies
    5. Return ONLY the JSON object, no additional text
    """

    try:
        result = call_o1_mini(analysis_prompt)
        cleaned_result = clean_ai_response(result)
        return json.loads(cleaned_result)
    except Exception as e:
        print(f"Error in analysis: {str(e)}")
        print("Raw response:", result)
        return {
            "component_mappings": {"outputs": [], "inputs": []},
            "event_handlers": {"steps": []},
            "required_updates": {}
        }


def generate_code_updates(analysis_result: Dict) -> Dict:
    """Second AI: Generates specific code updates based on the analysis"""
    code_generation_prompt = f"""
    You are an expert code generator. Based on this analysis, generate EXACT code updates.
    
    Analysis Result:
    {json.dumps(analysis_result, indent=2)}

    CRITICAL RULES:
    1. REMOVED PROMPTS:
       When a prompt is removed from config:
       - REMOVE all related component references
       - REMOVE all related handler code
       - DO NOT suggest adding new components
       - DO NOT suggest new mappings
       
    2. NEW PROMPTS:
       When a new prompt is added to config:
       - ADD corresponding execute_prompt() method in Project.py
       - ADD UI component references in event_handlers.py
       - ADD new step button handler if step is specified
       - ADD necessary input/output mappings
       - FOLLOW exact naming from config
       - ENSURE all required components are created
       
    3. Processing Order:
       a) Process removals FIRST (prompt_changes.removed_prompts)
       b) Then process additions (prompt_changes.new_prompts)
       c) Finally process other updates

    Generate a JSON response with this EXACT structure:
    {{
        "Project.py": {{
            "updates": [
                {{
                    "type": "removal|addition|update",  # Must specify type
                    "location": "exact_location",
                    "explanation": "Removing/Adding/Updating for prompt 'prompt_name'",
                    "old_code": "existing code to remove or update",
                    "new_code": "new code to add or update with",
                    "template": {{  # Only for additions
                        "method_signature": "def method_name(self, ...)",
                        "docstring": "\"\"\"Method documentation\"\"\"",
                        "implementation": "method implementation",
                        "return_statement": "return statement"
                    }}
                }}
            ]
        }},
        "event_handlers.py": {{
            "updates": [
                {{
                    "type": "removal|addition|update",
                    "location": "event_handlers.py - Component References",
                    "explanation": "Adding/Removing components for prompt 'prompt_name'",
                    "old_code": "code to remove or update",
                    "new_code": "new code or empty string for removals",
                    "components": [  # Only for additions
                        {{
                            "name": "component_name",
                            "type": "text|markdown|dataframe",
                            "initial_value": "initial value"
                        }}
                    ],
                    "handler": {{  # Only for additions requiring handlers
                        "name": "handler_name",
                        "inputs": ["input1", "input2"],
                        "outputs": ["output1", "output2"],
                        "implementation": "handler implementation"
                    }}
                }}
            ]
        }}
    }}

    Requirements:
    1. Process changes in correct order (removals -> additions -> updates)
    2. Each update MUST specify correct type (removal|addition|update)
    3. For removals, new_code should be empty string
    4. For additions:
       - Include complete method templates
       - Include all required components
       - Include handler implementations if needed
       - Follow exact naming from config
    5. Maintain exact function signatures
    6. Preserve existing code structure
    7. Include clear comments explaining changes
    8. Each file must have an 'updates' array, even if empty

    IMPORTANT: Return ONLY the JSON object, without any markdown formatting or explanation.
    """
    
    try:
        result = call_o1_mini(code_generation_prompt)
        cleaned_result = clean_ai_response(result)
        parsed_result = json.loads(cleaned_result)
        
        # Ensure correct structure and validate updates
        for file_name in ["Project.py", "event_handlers.py"]:
            if file_name not in parsed_result:
                parsed_result[file_name] = {"updates": []}
            
            # Ensure updates array exists
            if "updates" not in parsed_result[file_name]:
                parsed_result[file_name]["updates"] = []
                
            # Validate each update
            for update in parsed_result[file_name]["updates"]:
                if "type" not in update:
                    update["type"] = "update"  # Default type if missing
                if update["type"] == "removal" and update.get("new_code"):
                    update["new_code"] = ""  # Ensure removals have empty new_code
                if update["type"] == "addition":
                    # Ensure template exists for Project.py additions
                    if file_name == "Project.py" and "template" not in update:
                        update["template"] = {
                            "method_signature": "",
                            "docstring": "",
                            "implementation": "",
                            "return_statement": ""
                        }
                    # Ensure components and handler exist for event_handlers.py additions
                    if file_name == "event_handlers.py":
                        if "components" not in update:
                            update["components"] = []
                        if "handler" not in update:
                            update["handler"] = {
                                "name": "",
                                "inputs": [],
                                "outputs": [],
                                "implementation": ""
                            }
                    
        return parsed_result
        
    except Exception as e:
        print(f"Error generating code updates: {str(e)}")
        print("Raw response:", result)
        return {
            "Project.py": {"updates": []},
            "event_handlers.py": {"updates": []}
        }

def apply_updates_with_ai(file_path: str, updates: List[Dict]) -> bool:
    """Third AI: Intelligently applies code updates while preserving structure"""
    try:
        with open(file_path, 'r') as f:
            current_code = f.read()
            
        # Store original empty lines and their positions
        original_lines = current_code.splitlines(keepends=True)
        empty_line_indices = [i for i, line in enumerate(original_lines) if not line.strip()]
        trailing_spaces = [len(line) - len(line.rstrip()) for line in original_lines]
            
        update_prompt = f"""
        You are an expert code updater. Your task is to apply updates to the code while perfectly preserving its structure and formatting.
        
        CRITICAL RULES:
        1. DO NOT modify any line breaks or indentation
        2. DO NOT add or remove any lines
        3. DO NOT change any code structure or formatting
        4. ONLY replace the exact old_code with new_code
        5. Return the COMPLETE file with ALL original formatting preserved
        6. Keep ALL whitespace, commas, and brackets exactly as they are
        7. Return ONLY the code, no explanations or markdown
        8. Preserve ALL empty lines in their exact positions
        9. Keep ALL trailing spaces at the end of each line
        
        Current Code:
        {current_code}
        
        Required Updates:
        {json.dumps(updates, indent=2)}
        
        IMPORTANT: 
        - Empty lines are at these positions: {empty_line_indices}
        - Each line must maintain its exact trailing spaces
        - Return the complete file with perfect structure preservation
        """
        
        # Get AI-generated update
        result = call_o1_mini(update_prompt)
        updated_code = clean_code_response(result)
        
        # Split into lines while preserving endings
        updated_lines = updated_code.splitlines(keepends=True)
        
        # Handle line structure preservation
        updated_code = preserve_line_structure(original_lines, updated_lines, empty_line_indices, trailing_spaces)
        
        # Verify updates were applied using flexible pattern matching
        for update in updates:
            old_code = update.get('old_code', '').strip()
            new_code = update.get('new_code', '').strip()
            
            # Skip empty updates
            if not old_code and not new_code:
                continue
                
            # For renames, extract the core identifiers
            if update.get('type') == 'rename':
                old_identifier = extract_identifier(old_code)
                new_identifier = extract_identifier(new_code)
                
                if old_identifier in updated_code:
                    print(f"Error: Old identifier '{old_identifier}' still present")
                    return False
                if new_identifier not in updated_code:
                    print(f"Error: New identifier '{new_identifier}' not found")
                    return False
            else:
                # For other updates, check if the change was applied
                if old_code and old_code in updated_code:
                    print(f"Error: Old code still present")
                    return False
                if new_code and new_code not in updated_code:
                    print(f"Error: New code not found")
                    return False
        
        # Write the updated code
        with open(file_path, 'w') as f:
            f.write(updated_code)
            
        return True
        
    except Exception as e:
        print(f"Error applying AI updates: {str(e)}")
        return False

def clean_code_response(result: str) -> str:
    """Clean up the AI response code"""
    if '```python' in result:
        result = result.split('```python', 1)[1]
    if '```' in result:
        result = result.split('```', 1)[0]
    return result.strip()

def preserve_line_structure(original_lines: List[str], updated_lines: List[str], 
                          empty_indices: List[int], trailing_spaces: List[int]) -> str:
    """Preserve the original code structure"""
    if len(original_lines) != len(updated_lines):
        fixed_lines = []
        for i in range(len(original_lines)):
            if i in empty_indices:
                fixed_lines.append('\n')
            else:
                line = updated_lines[i].rstrip() if i < len(updated_lines) else ''
                line = line + ' ' * trailing_spaces[i]
                fixed_lines.append(line + '\n')
        return ''.join(fixed_lines)
    return ''.join(updated_lines)

def extract_identifier(code: str) -> str:
    """Extract the core identifier from a code snippet"""
    # Remove common code patterns
    code = code.replace('self.', '')
    code = code.replace('execute_prompt(', '')
    code = code.replace('generated_', '')
    code = code.replace('all_components[', '')
    code = code.replace(']', '')
    code = code.replace('"', '')
    code = code.replace("'", '')
    
    # Return the cleaned identifier
    return code.strip()

def extract_config_without_prompts() -> str:
    """Extract config file content without prompts"""
    try:
        with open('page_prompts_config.py', 'r') as f:
            config_lines = []
            in_prompt_block = False
            in_multiline_string = False
            prompt_indent = 0
            
            for line in f:
                stripped_line = line.strip()
                
                # Skip empty lines
                if not stripped_line:
                    continue
                
                # Detect start of multi-line string
                if '"""' in line:
                    if not in_multiline_string:
                        in_multiline_string = True
                        in_prompt_block = True
                        continue
                    else:
                        in_multiline_string = False
                        in_prompt_block = False
                        continue
                
                # Skip lines while in a prompt block
                if in_prompt_block or in_multiline_string:
                    continue
                
                # Detect single-line prompt assignments
                if 'prompt=' in line and not in_multiline_string:
                    continue
                
                # Keep all other lines
                config_lines.append(line)
            
            return ''.join(config_lines)
    except Exception as e:
        return f"Error reading config file: {str(e)}"

def main():
    """Main function to handle config changes"""
    try:
        # Read current files
        config_content = extract_config_without_prompts()
        if config_content.startswith("Error"):
            print(config_content)
            return
            
        with open('Project.py', 'r') as f:
            project_content = f.read()
        with open('event_handlers.py', 'r') as f:
            current_handlers = f.read()

        # First AI: Analyze current structure
        print("\nAnalyzing configuration structure...")
        analysis_result = analyze_config_structure(config_content, project_content, current_handlers)

        if not analysis_result.get('component_mappings'):
            print("Error: Invalid analysis result format")
            return

        # Display analysis results
        print("\nConfiguration Analysis:")
        print("1. Output Mappings:", len(analysis_result['component_mappings']['outputs']))
        print("2. Input Mappings:", len(analysis_result['component_mappings']['inputs']))
        print("3. Event Handlers:", len(analysis_result['event_handlers']['handlers']))

        if analysis_result.get('required_updates'):
            print("\nRequired Updates:")
            for file, updates in analysis_result['required_updates'].items():
                print(f"\n{file}:")
                for update in updates:
                    # Safely access update information with fallbacks
                    update_type = update.get('type', 'Unknown')
                    explanation = update.get('reason', update.get('explanation', 'No explanation provided'))
                    print(f"- {update_type}: {explanation}")

        if analysis_result.get('validation_issues'):
            print("\nValidation Issues:")
            for issue in analysis_result['validation_issues']:
                severity = issue.get('severity', 'unknown')
                description = issue.get('description', 'No description provided')
                issue_type = issue.get('type', 'Unknown')
                print(f"- {issue_type} ({severity}): {description}")

        # Get user confirmation
        confirm = input("\nProceed with generating code updates? (y/n): ")
        if confirm.lower() != 'y':
            print("Analysis complete. Update cancelled.")
            return

        # Second AI: Generate code updates
        print("\nGenerating code updates...")
        code_updates = generate_code_updates(analysis_result)

        # Display proposed changes
        print("\nProposed Code Updates:")
        for file, updates in code_updates.items():
            print(f"\n{file}:")
            for update in updates['updates']:
                print(f"- Location: {update.get('location', 'Unknown location')}")
                print(f"  Explanation: {update.get('explanation', 'No explanation provided')}")
                print("  Old code:")
                print(f"```\n{update.get('old_code', 'No old code provided')}\n```")
                print("  New code:")
                print(f"```\n{update.get('new_code', 'No new code provided')}\n```")

        # Final confirmation
        confirm = input("\nApply these code updates? (y/n): ")
        if confirm.lower() != 'y':
            print("Updates cancelled.")
            return

        # Apply updates
        success = True
        if 'Project.py' in code_updates:
            success &= apply_updates_with_ai('Project.py' ,code_updates['Project.py']['updates'])
        if 'event_handlers.py' in code_updates:
            success &= apply_updates_with_ai('event_handlers.py' , code_updates['event_handlers.py']['updates'])

        if success:
            print("Successfully updated all files")
        else:
            print("Some updates failed - please check the files manually")

    except Exception as e:
        print(f"Error in update process: {str(e)}")
        import traceback
        print("Traceback:", traceback.format_exc())

if __name__ == "__main__":
    main()