app.py

import gradio as gr
import os
import time
import json
import requests

# Set environment variables
# Note: For Hugging Face Spaces, you'll set these in the Secrets tab
# os.environ["TOGETHER_API_KEY"] = "your-api-key-here"

# Model options for dropdown
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"
]

# Initialize Together client
def get_together_client():
    api_key = os.environ.get("TOGETHER_API_KEY")
    if not api_key:
        raise ValueError("TOGETHER_API_KEY not found. Please set it in the environment variables.")
    return api_key

# Function to call LLM using REST API directly
def call_llm(model, prompt, temperature=0.7, max_tokens=1500):
    api_key = get_together_client()
    
    # System message for Salesforce expertise
    system_message = "You are a Salesforce development expert specializing in B2B Commerce migrations, CloudCraze to B2B Lightning Experience conversions, and Apex code optimization. Provide clear, step-by-step guidance and corrected code."
    
    start_time = time.time()
    
    try:
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_message},
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
            "top_p": 0.9
        }
        
        response = requests.post(
            "https://api.together.xyz/v1/chat/completions", 
            headers=headers, 
            json=data
        )
        
        if response.status_code != 200:
            return f"Error: API returned status code {response.status_code}: {response.text}"
        
        response_data = response.json()
        completion_text = response_data["choices"][0]["message"]["content"]
        
        end_time = time.time()
        duration = end_time - start_time
        
        # Add duration to the completion text if needed
        # completion_text += f"\n\nResponse time: {duration:.2f} seconds"
        
        return completion_text
        
    except Exception as e:
        return f"Error calling API: {str(e)}"

# For Apex Trigger correction
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 what was changed and why.
"""
    
    return call_llm(model, prompt)

# For CloudCraze Object conversion
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 appropriate B2B LEx system object that corresponds to this CloudCraze object,
map the fields correctly, and provide the complete definition for the B2B LEx version.

```
{cc_object_code}
```

Please return:
1. The equivalent B2B LEx system object name
2. Field mapping between CloudCraze and B2B LEx
3. The complete definition for the B2B LEx implementation
4. Any additional steps needed for proper migration
"""
    
    return call_llm(model, prompt)

# Gradio Interface
with gr.Blocks(title="Salesforce B2B Commerce Migration Assistant") as app:
    gr.Markdown("# Salesforce B2B Commerce Migration Assistant")
    gr.Markdown("This tool helps with migrating from CloudCraze to B2B Lightning Experience")
    
    model_dropdown = gr.Dropdown(
        choices=TOGETHER_MODELS,
        value=TOGETHER_MODELS[0],
        label="Select AI Model"
    )
    
    with gr.Tab("Apex Trigger Correction"):
        gr.Markdown("### Apex Trigger Correction")
        gr.Markdown("Paste your Apex Trigger code below to correct issues for B2B LEx compatibility")
        
        trigger_input = gr.Textbox(
            lines=15,
            placeholder="Paste your Apex Trigger code here...",
            label="Apex Trigger Code"
        )
        
        trigger_output = gr.Textbox(
            lines=20,
            label="Corrected Apex Trigger"
        )
        
        trigger_button = gr.Button("Correct Apex Trigger")
        trigger_button.click(
            fn=correct_apex_trigger,
            inputs=[model_dropdown, trigger_input],
            outputs=trigger_output
        )
    
    with gr.Tab("CloudCraze Object Conversion"):
        gr.Markdown("### CloudCraze to B2B LEx Object Conversion")
        gr.Markdown("Paste your CloudCraze Object code to convert it to B2B Lightning Experience format")
        
        object_input = gr.Textbox(
            lines=15,
            placeholder="Paste your CloudCraze Object code here...",
            label="CloudCraze Object Code"
        )
        
        object_output = gr.Textbox(
            lines=20,
            label="Converted B2B LEx Object"
        )
        
        object_button = gr.Button("Convert Object")
        object_button.click(
            fn=convert_cc_object,
            inputs=[model_dropdown, object_input],
            outputs=object_output
        )
    
    gr.Markdown("### About This Tool")
    gr.Markdown("""
    This tool uses advanced AI models to assist with Salesforce B2B Commerce migrations.
    - The Apex Trigger Correction tool helps fix trigger code for B2B Lightning Experience compatibility
    - The CloudCraze Object Conversion tool maps CC objects to their B2B LEx equivalents
    
    **Note**: Always review AI-generated output before implementing in production environments.
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()