Upload app.py with huggingface_hub

app.py

@@ -1,53 +1,121 @@
 # -*- coding: utf-8 -*-
 import os
-import gradio as gr
 import warnings
-from starlette.exceptions import StarletteDeprecationWarning
-warnings.filterwarnings("ignore", category=StarletteDeprecationWarning)
 import asyncio
-asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
-import openai
 
-# Load OpenAI API key from environment
+import gradio as gr
+
+# Suppress Starlette deprecation warnings that appear in Spaces
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# ----------------------------------------------------------------------
+# LLM client setup
+# ----------------------------------------------------------------------
+# Try Groq first (preferred). If the GROQ_API_KEY env var is not set,
+# fall back to OpenAI (requires OPENAI_API_KEY).
+groq_api_key = os.getenv("GROQ_API_KEY")
 openai_api_key = os.getenv("OPENAI_API_KEY")
-if not openai_api_key:
-    openai_api_key = ""
-openai.api_key = openai_api_key
 
-def improve_code(user_code, language):
+if groq_api_key:
+    try:
+        from groq import Groq
+        groq_client = Groq(api_key=groq_api_key)
+    except Exception as e:
+        groq_client = None
+        print(f"Failed to init Groq client: {e}")
+else:
+    groq_client = None
+
+if not groq_api_key and openai_api_key:
+    try:
+        import openai
+        openai.api_key = openai_api_key
+    except Exception as e:
+        openai = None
+        print(f"Failed to init OpenAI client: {e}")
+else:
+    openai = None
 
+# ----------------------------------------------------------------------
+# Core logic
+# ----------------------------------------------------------------------
+def improve_code(user_code: str, language: str) -> str:
     """
-    Sends the user code to OpenAI GPT model and asks for an improved version.
+    Sends the user code to a LLM (Groq preferred, OpenAI fallback) and asks
+    for an improved version. Returns the improved code or an error message.
     """
-    system_prompt = f"You are an expert {language} developer. Refactor and improve the following code. Keep the same functionality but make it cleaner, more efficient and well commented."
+    system_prompt = (
+        f"You are an expert {language} developer. Refactor and improve the "
+        "following code. Keep the same functionality but make it cleaner, "
+        "more efficient and well commented."
+    )
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_code}
+        {"role": "user", "content": user_code},
     ]
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=messages,
-            temperature=0.2,
-            max_tokens=1024,
-            n=1,
-            stop=None,
-        )
-        improved = response.choices[0].message.content.strip()
-        return improved
-    except Exception as e:
-        return f"Error: {str(e)}"
 
-# Build Gradio interface
+    # ------------------------------------------------------------------
+    # Try Groq
+    # ------------------------------------------------------------------
+    if groq_client:
+        try:
+            response = groq_client.chat.completions.create(
+                model="llama3-70b-8192",
+                messages=messages,
+                temperature=0.2,
+                max_tokens=1024,
+            )
+            improved = response.choices[0].message.content.strip()
+            return improved
+        except Exception as e:
+            return f"Groq error: {str(e)}"
+
+    # ------------------------------------------------------------------
+    # Fallback to OpenAI
+    # ------------------------------------------------------------------
+    if openai:
+        try:
+            response = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=messages,
+                temperature=0.2,
+                max_tokens=1024,
+            )
+            improved = response.choices[0].message.content.strip()
+            return improved
+        except Exception as e:
+            return f"OpenAI error: {str(e)}"
+
+    return "No LLM credentials found. Set GROQ_API_KEY or OPENAI_API_KEY."
+
+# ----------------------------------------------------------------------
+# Gradio UI
+# ----------------------------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# Code Improver")
-    gr.Markdown("Enter your code and select the programming language. The model will return a cleaner version.")
+    gr.Markdown(
+        "Enter your code, select the programming language, and click **Improve**. "
+        "The model will return a cleaner version."
+    )
     with gr.Row():
-        code_input = gr.Textbox(label="Original Code", lines=15, placeholder="Paste your code here")
+        code_input = gr.Textbox(
+            label="Original Code",
+            lines=15,
+            placeholder="Paste your code here",
+        )
         language_dropdown = gr.Dropdown(
-            choices=["Python", "JavaScript", "Java", "C++", "C#", "Go", "Ruby", "PHP"],
+            choices=[
+                "Python",
+                "JavaScript",
+                "Java",
+                "C++",
+                "C#",
+                "Go",
+                "Ruby",
+                "PHP",
+            ],
             value="Python",
-            label="Language"
+            label="Language",
         )
     improve_button = gr.Button("Improve")
     output_box = gr.Textbox(label="Improved Code", lines=15)
@@ -55,7 +123,7 @@ with gr.Blocks() as demo:
     improve_button.click(
         fn=improve_code,
         inputs=[code_input, language_dropdown],
-        outputs=output_box
+        outputs=output_box,
     )
 
 demo.launch(server_name="0.0.0.0")