app.py

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

# 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 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.
"""
    return call_llm(model, prompt)

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
4. Additional steps if needed
"""
    return call_llm(model, prompt)

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]
            data = json.loads(json_str)
            
            # Extract the required metrics
            metrics = {
                "quality_rating": data.get("quality_rating", 0),
                "accuracy": data.get("accuracy", 0.0),
                "completeness": data.get("completeness", 0.0),
                "best_practices_alignment": data.get("best_practices_alignment", 0.0)
            }
            
            return metrics
        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 main():
    with gr.Blocks(title="Salesforce B2B Commerce Migration Assistant") as app:
        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"
            )
            trigger_output = gr.Textbox(
                lines=15,
                label="Corrected Apex Trigger",
                interactive=True
            )
            trigger_button = gr.Button("Correct Apex Trigger")
            trigger_button.click(
                fn=correct_apex_trigger,
                inputs=[primary_model_dropdown, trigger_input],
                outputs=trigger_output,
                show_progress=True
            )
            
            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
            
            validate_trigger_button.click(
                fn=validate_and_chart_trigger,
                inputs=[validation_model_dropdown, trigger_input, trigger_output],
                outputs=[trigger_validation_output, trigger_chart],
                show_progress=True
            )
            
            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_output, 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"
            )
            object_output = gr.Textbox(
                lines=15,
                label="Converted B2B LEx Object",
                interactive=True
            )
            object_button = gr.Button("Convert Object")
            object_button.click(
                fn=convert_cc_object,
                inputs=[primary_model_dropdown, object_input],
                outputs=object_output,
                show_progress=True
            )
            
            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
            
            validate_object_button.click(
                fn=validate_and_chart_object,
                inputs=[validation_model_dropdown, object_input, object_output],
                outputs=[object_validation_output, object_chart],
                show_progress=True
            )
            
            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_output, object_validation_output, object_chart]
                )

        with gr.Accordion("About This Tool", open=False):
    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.
        
        Validation now outputs four key metrics (quality, accuracy, completeness, best practices) as both JSON and a fun chart.
        Always review AI-generated code before production use.
        """)

    app.launch()

if __name__ == "__main__":
    main()