Update app.py

app.py

@@ -1,16 +1,57 @@
 import gradio as gr
-from openai import OpenAI
+import time
 
-def generate_solution(api_key, problem_statement):
-    # Check if inputs are provided
-    if not api_key or not problem_statement:
-        return "Please provide both an API key and a problem statement."
+# Try importing OpenAI with error handling
+try:
+    from openai import OpenAI
+except ImportError:
+    print("OpenAI package not installed. Installing now...")
+    import subprocess
+    import sys
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "openai"])
+    from openai import OpenAI
+
+def test_api_connection(api_key):
+    """Test the API connection and return a status message"""
+    if not api_key or api_key.strip() == "":
+        return False, "API key is required"
+    
+    try:
+        client = OpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            api_key=api_key.strip(),
+        )
+        # Simple request to test connection
+        client.models.list()
+        return True, "API connection successful"
+    except Exception as e:
+        return False, f"API connection failed: {str(e)}"
+
+def generate_solution(api_key, problem_statement, progress=gr.Progress()):
+    """Generate solution with progress updates"""
+    progress(0, desc="Starting...")
+    
+    # Input validation
+    if not api_key or api_key.strip() == "":
+        return "Error: OpenRouter API key is required"
+    
+    if not problem_statement or problem_statement.strip() == "":
+        return "Error: Please provide a problem statement"
     
+    progress(0.1, desc="Validating API key...")
+    # Test API connection first
+    success, message = test_api_connection(api_key)
+    if not success:
+        return f"Error: {message}"
+    
+    progress(0.3, desc="Sending request to AI model...")
     try:
         client = OpenAI(
             base_url="https://openrouter.ai/api/v1",
             api_key=api_key.strip(),
         )
+        
+        progress(0.5, desc="Generating solution...")
         completion = client.chat.completions.create(
             model="open-r1/olympiccoder-7b:free",
             messages=[
@@ -26,60 +67,98 @@ def generate_solution(api_key, problem_statement):
             temperature=0.3,
             max_tokens=2048
         )
+        
+        progress(0.8, desc="Processing response...")
         response = completion.choices[0].message.content
-        return format_response(response)
+        
+        progress(0.9, desc="Formatting output...")
+        formatted_response = format_response(response)
+        
+        progress(1.0, desc="Done!")
+        return formatted_response
     
     except Exception as e:
-        return f"Error: {str(e)}"
+        error_message = str(e)
+        print(f"Error occurred: {error_message}")
+        return f"Error: {error_message}"
 
 def format_response(response):
+    """Format the AI response into markdown with code blocks"""
+    if not response:
+        return "Error: Received empty response from AI model"
+    
     formatted = []
+    lines = response.split('\n')
+    
+    # Find code blocks
     in_code = False
-    code_buffer = []
-    explanation_buffer = []
+    code_blocks = []
+    current_block = []
     
-    for line in response.split('\n'):
+    for line in lines:
         if line.strip().startswith('```'):
-            if in_code and code_buffer:
-                # Skip language marker if present
-                if len(code_buffer) > 0 and any(code_buffer[0].strip().lower() == lang for lang in ["python", "java", "c++", "cpp"]):
-                    code_content = '\n'.join(code_buffer[1:])
-                else:
-                    code_content = '\n'.join(code_buffer)
-                formatted.append(f"```python\n{code_content}\n```")
-                code_buffer = []
+            if in_code:
+                code_blocks.append(current_block)
+                current_block = []
+            else:
+                current_block = []
             in_code = not in_code
             continue
-            
+        
         if in_code:
-            code_buffer.append(line)
-        else:
-            explanation_buffer.append(line)
+            current_block.append(line)
     
-    # Handle case where code block wasn't properly closed
-    if in_code and code_buffer:
-        if len(code_buffer) > 0 and any(code_buffer[0].strip().lower() == lang for lang in ["python", "java", "c++", "cpp"]):
-            code_content = '\n'.join(code_buffer[1:])
-        else:
-            code_content = '\n'.join(code_buffer)
-        formatted.append(f"```python\n{code_content}\n```")
+    # Handle unclosed code block
+    if in_code and current_block:
+        code_blocks.append(current_block)
     
-    if explanation_buffer:
-        formatted.append("### Explanation\n" + '\n'.join(explanation_buffer))
+    # Process code blocks and explanations
+    processed_response = response
+    
+    # Replace each code block with properly formatted markdown
+    for i, block in enumerate(code_blocks):
+        # Skip language identifier if present
+        if block and any(block[0].strip().lower() == lang for lang in ["python", "java", "c++", "cpp", "javascript"]):
+            code_content = '\n'.join(block[1:])
+            lang = block[0].strip().lower()
+        else:
+            code_content = '\n'.join(block)
+            lang = "python"  # Default to Python
+        
+        # Create search pattern for the original block
+        if i < len(code_blocks) - 1 or not in_code:
+            block_text = '```' + '\n' + '\n'.join(block) + '\n' + '```'
+        else:
+            block_text = '```' + '\n' + '\n'.join(block)
         
-    return "\n\n".join(formatted)
+        # Replace with properly formatted block
+        formatted_block = f"```{lang}\n{code_content}\n```"
+        processed_response = processed_response.replace(block_text, formatted_block)
+    
+    return processed_response
 
-with gr.Blocks(title="Competitive Programming Assistant") as app:
+# Define the Gradio interface
+with gr.Blocks(title="Competitive Programming Assistant", theme=gr.themes.Soft()) as app:
     gr.Markdown("# 🏆 Competitive Programming Assistant")
     gr.Markdown("Powered by OlympicCoder-7B via OpenRouter AI")
     
     with gr.Row():
-        api_key = gr.Textbox(
-            label="OpenRouter API Key",
-            type="password",
-            placeholder="Enter your API key here...",
-            value=""  # Set default value to empty string
-        )
+        with gr.Column():
+            api_key = gr.Textbox(
+                label="OpenRouter API Key",
+                type="password",
+                placeholder="Enter your API key here...",
+                value=""
+            )
+            
+            test_btn = gr.Button("Test API Connection", variant="secondary")
+            api_status = gr.Textbox(label="API Status", interactive=False)
+            
+            test_btn.click(
+                fn=test_api_connection,
+                inputs=[api_key],
+                outputs=[gr.Checkbox(visible=False), api_status]
+            )
     
     with gr.Row():
         problem_input = gr.Textbox(
@@ -88,16 +167,15 @@ with gr.Blocks(title="Competitive Programming Assistant") as app:
             placeholder="Paste your programming problem here..."
         )
     
-    submit_btn = gr.Button("Generate Solution", variant="primary")
+    submit_btn = gr.Button("Generate Solution", variant="primary", size="lg")
     
-    output = gr.Markdown(label="Solution")
+    # Add a status indicator
+    status = gr.Markdown("Ready to generate solutions")
     
-    submit_btn.click(
-        fn=generate_solution,
-        inputs=[api_key, problem_input],
-        outputs=output
-    )
+    # Output area
+    solution_output = gr.Markdown(label="Solution")
     
+    # Example problems
     gr.Examples(
         examples=[
             [
@@ -109,6 +187,34 @@ with gr.Blocks(title="Competitive Programming Assistant") as app:
         ],
         inputs=[problem_input]
     )
+    
+    # Set up event handlers with additional status updates
+    def on_submit_click(api_key, problem):
+        return "Generating solution... Please wait."
+    
+    submit_btn.click(
+        fn=on_submit_click,
+        inputs=[api_key, problem_input],
+        outputs=[status],
+        queue=False
+    ).then(
+        fn=generate_solution,
+        inputs=[api_key, problem_input],
+        outputs=[solution_output],
+        queue=True
+    ).success(
+        fn=lambda: "Solution generated successfully!",
+        inputs=None,
+        outputs=[status]
+    ).error(
+        fn=lambda: "An error occurred. Please check your inputs and try again.",
+        inputs=None,
+        outputs=[status]
+    )
 
 if __name__ == "__main__":
-    app.launch(server_port=7860, share=True)  # Changed share to True to make it accessible externally
+    try:
+        print("Starting Competitive Programming Assistant...")
+        app.launch(server_port=7860, share=True, debug=True)
+    except Exception as e:
+        print(f"Error launching app: {str(e)}")