Update app.py

app.py

@@ -1,11 +1,8 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-import openai  # Can be replaced with another AI API
+import gradio as gr
+import openai  # Can be replaced with another AI model API
 import os
 import subprocess
 
-app = FastAPI()
-
 # AI Model Selection
 MODEL_MAP = {
     "gpt": "gpt-4",
@@ -13,40 +10,48 @@ MODEL_MAP = {
     "custom": "http://custom-llm-api"
 }
 
-class TestRequest(BaseModel):
-    model: str
-    prompt: str
-
-@app.post("/generate_test")
-def generate_test(request: TestRequest):
+def generate_test(model, prompt):
     """Generate test script using the selected AI model."""
-    model_choice = request.model
-    prompt = request.prompt
-    
-    model_api = MODEL_MAP.get(model_choice, "gpt-4")
+    model_api = MODEL_MAP.get(model, "gpt-4")
     try:
         response = openai.ChatCompletion.create(
             model=model_api,
             messages=[{"role": "system", "content": prompt}]
         )
         test_script = response["choices"][0]["message"]["content"]
-        
+
         # Save the generated test script
         with open("generated_test.py", "w") as f:
             f.write(test_script)
 
-        return {"message": "Test generated", "script": test_script}
-    
+        return test_script
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return f"Error: {str(e)}"
 
-@app.post("/run_tests")
 def run_tests():
     """Execute AI-generated test scripts and return results."""
     try:
         result = subprocess.run(["pytest", "generated_test.py", "--json-report"], capture_output=True, text=True)
-        return {"output": result.stdout, "error": result.stderr}
+        return result.stdout + "\n" + result.stderr
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return f"Error: {str(e)}"
+
+# Gradio UI
+with gr.Blocks() as app:
+    gr.Markdown("# AI Test Framework")
+    
+    with gr.Row():
+        model = gr.Dropdown(["gpt", "gemini", "custom"], value="gpt", label="Choose AI Model")
+        prompt = gr.Textbox(label="Test Prompt", placeholder="Describe the test case...")
+
+    generate_button = gr.Button("Generate Test Script")
+    test_script_output = gr.Textbox(label="Generated Test Script", interactive=False)
+
+    generate_button.click(generate_test, inputs=[model, prompt], outputs=test_script_output)
+
+    run_tests_button = gr.Button("Run Tests")
+    test_results_output = gr.Textbox(label="Test Execution Results", interactive=False)
+
+    run_tests_button.click(run_tests, inputs=[], outputs=test_results_output)
 
-# Run with: uvicorn app:app --host 0.0.0.0 --port 7860
+app.launch()