app.py

import gradio as gr
import os
import time
import requests
import json
import plotly.express as px
import pandas as pd
import re

# Model options for dropdown with both Together AI and Anthropic models
together_models = [
    "Qwen/Qwen2.5-Coder-32B-Instruct",
    "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
    "deepseek-ai/DeepSeek-R1-Distill-Llama-70B",
    "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
]

anthropic_models = [
    "claude-3-7-sonnet-20250219",
    "claude-3-haiku-20240307"
]

all_models = together_models + anthropic_models

VALIDATION_SCHEMA = {
    "quality_rating": "int (1–10)",
    "accuracy": "float (0.0–1.0)",
    "completeness": "float (0.0–1.0)",
    "best_practices_alignment": "float (0.0–1.0)",
    "explanations": {
        "quality_rating": "string",
        "accuracy": "string",
        "completeness": "string",
        "best_practices_alignment": "string"
    }
}

def get_api_key(provider):
    if provider == "together":
        api_key = os.getenv("TOGETHER_API_KEY")
        if not api_key:
            raise ValueError("TOGETHER_API_KEY not set. Please add it in Space secrets.")
        return api_key
    elif provider == "anthropic":
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY not set. Please add it in Space secrets.")
        return api_key
    else:
        raise ValueError(f"Unknown provider: {provider}")

def get_provider(model):
    if model in together_models:
        return "together"
    elif model in anthropic_models:
        return "anthropic"
    else:
        raise ValueError(f"Unknown model: {model}")

def call_together_api(model, prompt, temperature=0.7, max_tokens=1500):
    api_key = get_api_key("together")
    system_message = (
        "You are a Salesforce development expert specializing in B2B Commerce migrations,"
        " CloudCraze to B2B Lightning Experience conversions, and Apex code optimization."
    )
    try:
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_message},
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
            "top_p": 0.9
        }
        resp = requests.post(
            "https://api.together.xyz/v1/chat/completions",
            headers=headers,
            json=payload
        )
        if resp.status_code != 200:
            return f"Error: Status {resp.status_code}: {resp.text}"
        data = resp.json()
        text = data["choices"][0]["message"]["content"]
        return text
    except Exception as e:
        return f"Error calling Together AI API: {e}"

def call_anthropic_api(model, prompt, temperature=0.7, max_tokens=1500):
    api_key = get_api_key("anthropic")
    system_message = (
        "You are a Salesforce development expert specializing in B2B Commerce migrations,"
        " CloudCraze to B2B Lightning Experience conversions, and Apex code optimization."
    )
    try:
        headers = {
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json"
        }
        payload = {
            "model": model,
            "system": system_message,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature,
            "max_tokens": max_tokens
        }
        resp = requests.post(
            "https://api.anthropic.com/v1/messages",
            headers=headers,
            json=payload
        )
        if resp.status_code != 200:
            return f"Error: Status {resp.status_code}: {resp.text}"
        data = resp.json()
        text = data["content"][0]["text"]
        return text
    except Exception as e:
        return f"Error calling Anthropic API: {e}"

def call_llm(model, prompt, temperature=0.7, max_tokens=1500):
    provider = get_provider(model)
    if provider == "together":
        return call_together_api(model, prompt, temperature, max_tokens)
    elif provider == "anthropic":
        return call_anthropic_api(model, prompt, temperature, max_tokens)
    else:
        return f"Error: Unknown provider for model {model}"

def extract_code_blocks(text):
    """Extract code blocks from the model's response"""
    # Pattern to match code blocks with ```apex or ```java or just ```
    pattern = r"```(?:apex|java)?(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    
    # Clean up extracted code blocks
    code_blocks = []
    for block in matches:
        # Remove leading/trailing whitespace
        cleaned_block = block.strip()
        if cleaned_block:
            code_blocks.append(cleaned_block)
    
    # If no code blocks found but text contains code-like content
    if not code_blocks and ("public" in text or "private" in text or "trigger" in text):
        # Try to extract the most code-like part
        lines = text.split('\n')
        potential_code = []
        in_code_section = False
        
        for line in lines:
            if any(keyword in line for keyword in ["public", "private", "class", "trigger", "@"]):
                in_code_section = True
            
            if in_code_section:
                potential_code.append(line)
        
        if potential_code:
            code_blocks.append('\n'.join(potential_code))
    
    return '\n\n'.join(code_blocks) if code_blocks else ""

def extract_explanations(text, code_blocks):
    """Extract explanations from the model's response by removing code blocks"""
    # Replace all code blocks with placeholders
    explanation = text
    for block in code_blocks.split('\n\n'):
        # Escape special regex characters in the block
        escaped_block = re.escape(block)
        explanation = re.sub(escaped_block, "[CODE BLOCK REMOVED]", explanation, flags=re.DOTALL)
    
    # Remove code block markers
    explanation = re.sub(r"```(?:apex|java)?.*?```", "", explanation, flags=re.DOTALL)
    
    # Clean up the explanation
    explanation = explanation.replace("[CODE BLOCK REMOVED]", "*Code has been moved to the Code Output section*")
    
    return explanation.strip()

def correct_apex_trigger(model, trigger_code):
    if not trigger_code.strip():
        return "Please provide Apex Trigger code to correct.", "", ""
    prompt = f"""
Please analyze and correct the following Apex Trigger code for migration from CloudCraze to B2B Lightning Experience.
Identify any issues, optimize it according to best practices, and provide the corrected code.

```apex
{trigger_code}
```

Please return the corrected code with explanations of changes. Put the corrected code in ```apex code blocks.
"""
    response = call_llm(model, prompt)
    
    # Extract code blocks and explanations
    code_output = extract_code_blocks(response)
    explanation = extract_explanations(response, code_output)
    
    return response, code_output, explanation

def convert_cc_object(model, cc_object_code):
    if not cc_object_code.strip():
        return "Please provide CloudCraze Object code to convert.", "", ""
    prompt = f"""
Please convert the following CloudCraze Object code to B2B Lightning Experience format.
Identify the corresponding B2B LEx system object, map fields, and provide the full definition.

```
{cc_object_code}
```

Return:
1. Equivalent B2B LEx object name
2. Field mappings
3. Full implementation code in ```apex code blocks
4. Additional steps if needed
"""
    response = call_llm(model, prompt)
    
    # Extract code blocks and explanations
    code_output = extract_code_blocks(response)
    explanation = extract_explanations(response, code_output)
    
    return response, code_output, explanation

def validate_apex_trigger(validation_model, original_code, corrected_code):
    if not validation_model or not original_code.strip() or not corrected_code.strip():
        return "Please provide all required inputs for validation."
    
    prompt = f"""
I need you to validate and review the following Apex code migration. This involves both CloudCraze to B2B Lightning Experience conversions AND general Apex code optimization.

ORIGINAL CODE:
```apex
{original_code}
```

CORRECTED CODE:
```apex
{corrected_code}
```

Please review the corrected code and provide:
1. Is the correction accurate and complete? Highlight any missed issues.
2. Are there any potential runtime errors or performance concerns?
3. Does it follow Salesforce best practices for B2B Lightning Experience?
4. Are there optimization opportunities that were missed?
5. Evaluate both the B2B Commerce conversion aspects AND the general Apex optimization.
6. Rate the quality of the correction on a scale of 1-10 with explanation.

Additionally, provide a structured assessment in JSON format following this schema:
{json.dumps(VALIDATION_SCHEMA, indent=2)}

Be thorough and detailed in your assessment, addressing both the conversion and optimization aspects.
"""
    return call_llm(validation_model, prompt)

def validate_cc_object_conversion(validation_model, original_object, converted_object):
    if not validation_model or not original_object.strip() or not converted_object.strip():
        return "Please provide all required inputs for validation."
    
    prompt = f"""
I need you to validate and review the following CloudCraze Object conversion to B2B Lightning Experience format. This involves both CloudCraze to B2B Lightning Experience conversions AND optimization of the resulting code.

ORIGINAL CLOUDCRAZE OBJECT:
```
{original_object}
```

CONVERTED B2B LEX OBJECT:
```
{converted_object}
```

Please review the conversion and provide:
1. Is the object mapping correct and complete? Identify any missed fields or relationships.
2. Are there any potential data migration issues or concerns?
3. Does the converted object follow B2B Lightning Experience best practices?
4. Are there optimization opportunities that were missed in the conversion?
5. Evaluate both the B2B Commerce conversion aspects AND the optimization of the resulting code.
6. Rate the quality of the conversion on a scale of 1-10 with explanation.

Additionally, provide a structured assessment in JSON format following this schema:
{json.dumps(VALIDATION_SCHEMA, indent=2)}

Be thorough and detailed in your assessment, addressing both the conversion and optimization aspects.
"""
    return call_llm(validation_model, prompt)

def extract_validation_metrics(validation_text):
    try:
        # Try to find JSON in the response
        start_idx = validation_text.find('{')
        end_idx = validation_text.rfind('}') + 1
        
        if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
            json_str = validation_text[start_idx:end_idx]
            # Print for debugging
            print("Extracted JSON:", json_str)
            
            try:
                data = json.loads(json_str)
                
                # Extract the required metrics with default values
                metrics = {
                    "quality_rating": float(data.get("quality_rating", 0)),
                    "accuracy": float(data.get("accuracy", 0.0)),
                    "completeness": float(data.get("completeness", 0.0)),
                    "best_practices_alignment": float(data.get("best_practices_alignment", 0.0))
                }
                
                return metrics
            except json.JSONDecodeError as je:
                print(f"JSON decode error: {je}")
                return None
        else:
            print("No JSON found in validation text")
            return None
    except Exception as e:
        print(f"Error extracting metrics: {e}")
        return None

def create_radar_chart(metrics):
    if not metrics:
        return None
    
    # Create data for the radar chart
    categories = ["Quality", "Accuracy", "Completeness", "Best Practices"]
    values = [
        metrics["quality_rating"] / 10,  # Normalize to 0-1 scale
        metrics["accuracy"],
        metrics["completeness"],
        metrics["best_practices_alignment"]
    ]
    
    # Create a DataFrame for plotting
    df = pd.DataFrame({
        'Category': categories,
        'Value': values
    })
    
    # Create the radar chart
    fig = px.line_polar(
        df, r='Value', theta='Category', line_close=True,
        range_r=[0, 1], title="Validation Assessment"
    )
    fig.update_traces(fill='toself')
    
    return fig

def toggle_dark_mode(dark_mode, elements):
    """Toggle between light and dark mode for interface elements"""
    if dark_mode:
        style = "background-color: #2e2e2e; color: #ffffff;"
        code_style = "background-color: #1e1e1e; color: #e0e0e0; font-family: 'Courier New', monospace;"
    else:
        style = "background-color: #ffffff; color: #333333;"
        code_style = "background-color: #f5f5f5; color: #333333; font-family: 'Courier New', monospace;"
    
    return [style for _ in range(len(elements) // 2)] + [code_style for _ in range(len(elements) // 2)]

def main():
    with gr.Blocks(title="Salesforce B2B Commerce Migration Assistant", theme=gr.themes.Soft(primary_hue="blue")) as app:
        # Initialize dark mode state
        dark_mode = gr.State(False)
        
        gr.Markdown("# Salesforce B2B Commerce Migration Assistant")
        gr.Markdown("This tool helps migrate CloudCraze code to B2B Lightning Experience.")

        # Create model dropdowns
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Primary Model")
                primary_model_dropdown = gr.Dropdown(
                    choices=all_models,
                    value=anthropic_models[0],  # Default to Claude 3.7 Sonnet
                    label="Select Primary AI Model for Conversion"
                )
            
            with gr.Column():
                gr.Markdown("### Validation Model")
                validation_model_dropdown = gr.Dropdown(
                    choices=all_models,
                    value=anthropic_models[1],  # Default to Claude 3 Haiku
                    label="Select Validation AI Model for Review",
                    info="This model will validate and review the output from the primary model"
                )

        with gr.Tab("Apex Trigger Correction"):
            gr.Markdown("### Apex Trigger Correction")
            gr.Markdown("Paste your Apex Trigger code below:")

            trigger_input = gr.Textbox(
                lines=12,
                placeholder="Paste Apex Trigger code here...",
                label="Apex Trigger Code"
            )
            
            with gr.Row():
                trigger_button = gr.Button("Correct Apex Trigger", variant="primary")
                copy_code_button = gr.Button("Copy Code", variant="secondary")
                
            with gr.Accordion("Full Model Response (Hidden by Default)", open=False):
                trigger_full_response = gr.Textbox(
                    lines=15,
                    label="Full Model Response",
                    interactive=False
                )
            
            with gr.Row():
                with gr.Column():
                    trigger_explanation = gr.Textbox(
                        lines=10,
                        label="Explanation",
                        placeholder="Explanation will appear here after correction",
                        interactive=False,
                        elem_id="trigger_explanation"
                    )
                
                with gr.Column():
                    trigger_code_output = gr.Code(
                        language="python",  # Using Python syntax highlighting as it's supported
                        label="Corrected Code (Apex)",
                        value="# Corrected Apex code will appear here",
                        elem_id="trigger_code_output"
                    )
            
            gr.Markdown("### Validation Results")
            with gr.Row():
                with gr.Column(scale=2):
                    trigger_validation_output = gr.Textbox(
                        lines=15,
                        label="Validation Assessment",
                        placeholder="Validation results will appear here after clicking 'Validate Correction'",
                        interactive=True
                    )
                with gr.Column(scale=1):
                    trigger_chart = gr.Plot(label="Validation Metrics")
            
            validate_trigger_button = gr.Button("Validate Correction")
            
            def validate_and_chart_trigger(model, original, corrected):
                validation_text = validate_apex_trigger(model, original, corrected)
                metrics = extract_validation_metrics(validation_text)
                chart = create_radar_chart(metrics) if metrics else None
                return validation_text, chart
            
            # Handle button clicks
            trigger_button.click(
                fn=correct_apex_trigger,
                inputs=[primary_model_dropdown, trigger_input],
                outputs=[trigger_full_response, trigger_code_output, trigger_explanation],
                show_progress=True
            )
            
            validate_trigger_button.click(
                fn=validate_and_chart_trigger,
                inputs=[validation_model_dropdown, trigger_input, trigger_code_output],
                outputs=[trigger_validation_output, trigger_chart],
                show_progress=True
            )
            
            # Copy code button functionality
            def trigger_copy_to_clipboard():
                # This function doesn't need to return anything - the copy happens in the frontend
                return None
            
            copy_code_button.click(
                fn=trigger_copy_to_clipboard,
                inputs=[],
                outputs=[],
                _js="""() => {
                    const codeElem = document.getElementById('trigger_code_output');
                    if (codeElem) {
                        navigator.clipboard.writeText(codeElem.textContent);
                        return [];
                    }
                }"""
            )
            
            with gr.Row():
                trigger_clear = gr.Button("Clear Input")
                trigger_clear.click(lambda: "", [], trigger_input)
                
                results_clear = gr.Button("Clear Results")
                results_clear.click(
                    lambda: ["", "", "", "", None],
                    [],
                    [trigger_full_response, trigger_code_output, trigger_explanation, trigger_validation_output, trigger_chart]
                )

        with gr.Tab("CloudCraze Object Conversion"):
            gr.Markdown("### CloudCraze Object Conversion")
            gr.Markdown("Paste your CloudCraze Object code below:")

            object_input = gr.Textbox(
                lines=12,
                placeholder="Paste CC object code here...",
                label="CloudCraze Object Code"
            )
            
            with gr.Row():
                object_button = gr.Button("Convert Object", variant="primary")
                object_copy_code_button = gr.Button("Copy Code", variant="secondary")
                
            with gr.Accordion("Full Model Response (Hidden by Default)", open=False):
                object_full_response = gr.Textbox(
                    lines=15,
                    label="Full Model Response",
                    interactive=False
                )
            
            with gr.Row():
                with gr.Column():
                    object_explanation = gr.Textbox(
                        lines=10,
                        label="Explanation",
                        placeholder="Explanation will appear here after conversion",
                        interactive=False,
                        elem_id="object_explanation"
                    )
                
                with gr.Column():
                    object_code_output = gr.Code(
                        language="python",  # Using Python syntax highlighting as it's supported
                        label="Converted Code (Apex)",
                        value="# Converted Apex code will appear here",
                        elem_id="object_code_output"
                    )
            
            gr.Markdown("### Validation Results")
            with gr.Row():
                with gr.Column(scale=2):
                    object_validation_output = gr.Textbox(
                        lines=15,
                        label="Validation Assessment",
                        placeholder="Validation results will appear here after clicking 'Validate Conversion'",
                        interactive=True
                    )
                with gr.Column(scale=1):
                    object_chart = gr.Plot(label="Validation Metrics")
            
            validate_object_button = gr.Button("Validate Conversion")
            
            def validate_and_chart_object(model, original, converted):
                validation_text = validate_cc_object_conversion(model, original, converted)
                metrics = extract_validation_metrics(validation_text)
                chart = create_radar_chart(metrics) if metrics else None
                return validation_text, chart
            
            # Handle button clicks
            object_button.click(
                fn=convert_cc_object,
                inputs=[primary_model_dropdown, object_input],
                outputs=[object_full_response, object_code_output, object_explanation],
                show_progress=True
            )
            
            validate_object_button.click(
                fn=validate_and_chart_object,
                inputs=[validation_model_dropdown, object_input, object_code_output],
                outputs=[object_validation_output, object_chart],
                show_progress=True
            )
            
            # Copy code button functionality
            def object_copy_to_clipboard():
                # This function doesn't need to return anything - the copy happens in the frontend
                return None
            
            object_copy_code_button.click(
                fn=object_copy_to_clipboard,
                inputs=[],
                outputs=[],
                _js="""() => {
                    const codeElem = document.getElementById('object_code_output');
                    if (codeElem) {
                        navigator.clipboard.writeText(codeElem.textContent);
                        return [];
                    }
                }"""
            )
            
            with gr.Row():
                object_clear = gr.Button("Clear Input")
                object_clear.click(lambda: "", [], object_input)
                
                object_results_clear = gr.Button("Clear Results")
                object_results_clear.click(
                    lambda: ["", "", "", "", None],
                    [],
                    [object_full_response, object_code_output, object_explanation, object_validation_output, object_chart]
                )

        # UI Preferences
        with gr.Accordion("UI Preferences", open=False):
            dark_mode_toggle = gr.Checkbox(label="Dark Mode", value=False)
            
            # Create a list of UI elements to update
            ui_elements = [
                trigger_explanation, trigger_code_output,
                object_explanation, object_code_output
            ]
            
            dark_mode_toggle.change(
                fn=toggle_dark_mode,
                inputs=[dark_mode_toggle, gr.State(ui_elements)],
                outputs=ui_elements
            )

        gr.Markdown("### About This Tool")
        gr.Markdown(
            """
                 **Primary Model**: Performs the initial code conversion or correction.
                **Validation Model**: Reviews and validates the output from the primary model, identifying potential issues or improvements.
                **Trigger Correction**: Fixes Apex Triggers for B2B LEx compatibility.
                **Object Conversion**: Maps and converts CloudCraze object definitions to B2B LEx.
                **Model Selection**: Choose from Together AI models or Anthropic's Claude models.
                
                **New Interface Features**:
                - Separated code output in dedicated code editor interface with syntax highlighting
                - Clear separation between explanations and code
                - "Copy Code" button for easy copying
                - Dark mode option for comfortable viewing
                - Hidden full model response to reduce clutter
                
                **Validation** outputs four key metrics (quality, accuracy, completeness, best practices) as both JSON and a radar chart.
                Always review AI-generated code before production use.
                """
        )

    app.launch()

if __name__ == "__main__":
    main()