code_updater.py

from typing import Dict
import openai
from contextlib import contextmanager
import re
import json

@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 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 analyze_files(config_content: str, project_content: str, current_handlers: str) -> dict:
    """First analyze all files and extract relevant information"""
    analysis_prompt = f"""   
You are an expert software engineer performing a comprehensive analysis. You must understand these key relationships:

1. Config Output to UI Components Mapping:
   - A single output in config (e.g., 'generated_bd_SOW') can map to multiple UI components
   - UI components typically come in pairs:
     * _text: For textbox display
     * _markdown: For markdown display
   - As for mandays , it only maps to one UI components which is in *_dataframe
   Example:
   Config output 'generated_bd_SOW' maps to:
   - generated_bd_SOW_text
   - generated_bd_SOW_markdown
   
2. Config Input to UI Component Mapping : 
   - The inputs should take the _text or _markdown of the component

3. Component Name and Config Key Validation:
   - Names in handlers must exactly match config
   - Both section names and component names must match 
   - Pay special attention to component name changes
   - Pay attention to case sensitivity

4. Prompt Name Validation : 
   - Prompt names in config are the source of truth 
   - Any references to these names in Project.py and event_handlers.py must match EXACTLY
   - Check for:
         * Function names in event_handlers.py
         * Component references in all_components dictionary
         * Input/output mappings in step handlers
         * Method names in Project.py
   
4. Input/Output Name Validation:
   - Config keys must match exactly in handlers
   - Check all_components dictionary keys in handlers against config
   - Example: if config uses 'generate_development_mandays', handlers must use the same key
   - Flag any mismatches between config keys and handler references

4. Function Parameter Mapping:
   - Function parameters in Project.py define the actual data flow
   - UI components must map to these parameters correctly

Context Files:

1. page_prompts_config.py (SOURCE OF TRUTH FOR NAMES & UI):
{config_content}

2. Project.py (SOURCE OF TRUTH FOR FUNCTION PARAMETERS):
{project_content}

3. event_handlers.py (NEEDS VALIDATION):
{current_handlers}

Analysis Requirements:

1. UI Component Validation:
   - Check each config output's UI components
   - Verify both _text and _markdown variants exist where specified
   - Ensure case matches exactly

2. Component Name Validation:
   - Compare against config definitions
   - Check both section names and component names
   - Flag any case mismatches

3. Function Parameter Validation:
   - Verify against Project.py signatures
   - Check parameter order and types
   - Ensure all required parameters are included

   
Return your analysis in this EXACT JSON format:
{{
    "step_1": {{
        "current_code": "exact current handler code",
        "needs_update": false,
        "updates": []
    }},
    "step_2": {{
        "current_code": "exact current handler code",
        "needs_update": true,
        "updates": [
            {{
                "type": "component_name",
                "line": "all_components['generate_BD_SOW']['generated_BD_SOW_text']",
                "replacement": "all_components['generate_bd_SOW']['generated_bd_SOW_text']",
                "reason": "Case mismatch with config definition"
            }}
        ]
    }},
    "step_3": {{
        "current_code": "exact current handler code",
        "needs_update": true,
        "updates": [
            {{
                "type": "component_name",
                "line": "all_components['generate_BD_SOW']['generated_BD_SOW_text']",
                "replacement": "all_components['generate_bd_SOW']['generated_bd_SOW_text']",
                "reason": "Case mismatch with config definition"
            }}
        ]
    }}
}}

Requirements:
1. Include ALL steps (1-3)
2. For each step:
   - Include exact current code
   - Set needs_update to true if changes needed
   - List ALL required updates
3. For each update:
   - Provide exact current line
   - Provide exact replacement
   - Include clear reason
4. Check:
   - Case sensitivity
   - Component names
   - Function parameters
   - UI component mappings
   
Every discrepancy must be reported as an issue with:
- Clear description of the problem
- Exact current implementation
- Exact expected implementation from config
- Specific location in the code

If no issues are found for a step, include an empty issues array.

Remember: 
- Component names must match config
- Function parameters must match Project.py signatures
- Some components may be correct in function but need case fixes from config
- DO NOT suggest changes unless there is a 100% exact mismatch
- If a component or parameter does not have any changes and it does exist and works, DO NOT suggest changes
"""
    
    analysis = call_o1_mini(analysis_prompt)
    return analysis

def generate_handler_code(step_number: int, analysis_result: dict) -> str:
    """Generate the handler code based on the analysis"""
    step_key = f"step_{step_number}"
    step_data = analysis_result.get(step_key, {})
    
    if not step_data.get("needs_update", False):
        return ""
        
    current_code = step_data["current_code"]
    updates = step_data.get("updates", [])
    
    # Apply all updates
    updated_code = current_code
    for update in updates:
        updated_code = updated_code.replace(update["line"], update["replacement"])
    
    return updated_code

def analyze_prompts_and_generate_handlers() -> Dict[int, str]:
    """Generate step handlers separately"""
    try:
        # Read all files once
        config_content = extract_config_without_prompts()
        print("\nConfig content loaded")
        
        with open('Project.py', 'r') as f:
            project_content = f.read()
        print("Project.py content loaded")
            
        with open('event_handlers.py', 'r') as f:
            current_handlers = f.read()
        print("event_handlers.py content loaded")
        
        # Do one analysis for all files
        print("\nAnalyzing all files...")
        analysis_result = analyze_files(config_content, project_content, current_handlers)
        
        # Parse the analysis result from string to JSON
        try:
            if isinstance(analysis_result, str):
                # Remove ```json and ``` markers if present
                analysis_result = analysis_result.replace('```json\n', '').replace('\n```', '').strip()
                analysis_result = json.loads(analysis_result)
            
            print("\nRaw analysis result:")
            print(json.dumps(analysis_result, indent=2))

            handlers = {}
            for step in range(1, 4):
                step_key = f"step_{step}"
                if step_key not in analysis_result:
                    print(f"\nNo analysis results found for {step_key}")
                    continue
                    
                step_code = generate_handler_code(step, analysis_result)
                
                # Clean up and validate the response
                if step_code:
                    step_code = step_code.strip()
                    
                    if not step_code.startswith(f"step_buttons['Step {step}"):
                        print(f"\nInvalid handler code format for Step {step}")
                        # Extract current handler as fallback
                        pattern = rf"step_buttons\['Step {step} : [^']*'\]\.click\([^)]+\)"
                        current_handler = re.search(pattern, current_handlers, re.DOTALL)
                        handlers[step] = current_handler.group(0) if current_handler else None
                    else:
                        handlers[step] = step_code
                        print(f"Generated handler for Step {step}")
            
            return handlers

        except json.JSONDecodeError as e:
            print(f"Error parsing analysis result: {str(e)}")
            print("Raw analysis result:", analysis_result)
            return {}

    except Exception as e:
        print(f"Error generating handlers: {str(e)}")
        import traceback
        print("Traceback:", traceback.format_exc())
        return {}

def main():
    """Generate and save the step handlers"""
    handlers = analyze_prompts_and_generate_handlers()
    print("Generated handlers for steps:", list(handlers.keys()))
    
    # Let the user review the changes
    for step, code in handlers.items():
        if not code:  
            continue
            
        print(f"\nStep {step} Handler:")
        print(code)
        update = input(f"Do you want to update Step {step} handler? (y/n): ")
        
        if update.lower() == 'y':
            try:
                # Read current content
                with open('event_handlers.py', 'r') as f:
                    content = f.read()
                
                # Find and replace the specific step handler
                pattern = rf"step_buttons\['Step {step} : [^']*'\]\.click\([^)]+\)"
                match = re.search(pattern, content, re.DOTALL)
                
                if not match:
                    print(f"Warning: Couldn't find Step {step} handler in file")
                    continue
                
                # Replace the matched content with the new code
                updated_content = content[:match.start()] + code.strip() + content[match.end():]
                
                # Write back to file
                with open('event_handlers.py', 'w') as f:
                    f.write(updated_content)
                    
                print(f"Successfully updated Step {step} handler in event_handlers.py")
                
                # Verify the write
                with open('event_handlers.py', 'r') as f:
                    new_content = f.read()
                if code.strip() in new_content:
                    print("Verified: Update successful")
                else:
                    print("Warning: Update may not have been successful")
                    
            except Exception as e:
                print(f"Error updating file: {str(e)}")
                print(f"Error details: {type(e).__name__}")

if __name__ == "__main__":
    main()