Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import time
 import requests
 
-# Model options for dropdown
+# 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",
@@ -11,16 +11,40 @@ together_models = [
     "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
 ]
 
-# Helper to fetch API key from environment
-def get_together_client():
-    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
+anthropic_models = [
+    "claude-3-7-sonnet-20250219",
+    "claude-3-haiku-20240307"
+]
+
+all_models = together_models + anthropic_models
+
+# Helper to fetch API keys from environment
+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}")
+
+# Determine the provider based on model name
+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}")
 
 # Call Together AI via REST API
-def call_llm(model, prompt, temperature=0.7, max_tokens=1500):
-    api_key = get_together_client()
+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."
@@ -52,7 +76,53 @@ def call_llm(model, prompt, temperature=0.7, max_tokens=1500):
         text = data["choices"][0]["message"]["content"]
         return text
     except Exception as e:
-        return f"Error calling API: {e}"
+        return f"Error calling Together AI API: {e}"
+
+# Call Anthropic API
+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."
+    )
+    start = time.time()
+    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}"
+
+# Unified function to call the appropriate API based on the model
+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}"
 
 # Build prompt for trigger correction
 def correct_apex_trigger(model, trigger_code):
@@ -96,9 +166,10 @@ def main():
         gr.Markdown("# Salesforce B2B Commerce Migration Assistant")
         gr.Markdown("This tool helps migrate CloudCraze code to B2B Lightning Experience.")
 
+        # Create a dropdown with model type labels
         model_dropdown = gr.Dropdown(
-            choices=together_models,
-            value=together_models[0],
+            choices=all_models,
+            value=anthropic_models[0],  # Default to Claude 3.7 Sonnet
             label="Select AI Model"
         )
 
@@ -155,6 +226,7 @@ def main():
             """
 - **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.
 
 Always review AI-generated code before production use.
             """
@@ -163,4 +235,4 @@ Always review AI-generated code before production use.
     app.launch()
 
 if __name__ == "__main__":
-    main()
+    main()