Update app.py

app.py

@@ -81,3 +81,57 @@ async def refine_code(websocket: WebSocket):
 
     except WebSocketDisconnect:
         print("Client disconnected")
+@app.websocket("/ws/optimize_code")
+async def optimize_code(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        data = await websocket.receive_text()
+        request_data = json.loads(data)
+        code = request_data.get("code")
+        language = request_data.get("language")
+        
+        prompt = f"Analyze the following {language} code for syntax, linting, and optimization suggestions:\n\n{code}"
+        messages = [{"role": "user", "content": prompt}]
+        
+        stream = client.chat.completions.create(
+            model="Qwen/Qwen2.5-Coder-32B-Instruct",
+            messages=messages,
+            temperature=0.5,
+            max_tokens=512,
+            top_p=0.7,
+            stream=True
+        )
+        
+        for chunk in stream:
+            await websocket.send_text(chunk.choices[0].delta.content)
+
+    except WebSocketDisconnect:
+        print("Client disconnected")
+
+# WebSocket for Failure Mode Analysis and Test Bench Generation
+@app.websocket("/ws/generate_test_bench")
+async def generate_test_bench(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        data = await websocket.receive_text()
+        request_data = json.loads(data)
+        code = request_data.get("code")
+        language = request_data.get("language")
+        
+        prompt = f"Generate test benches and identify potential failure modes for the following {language} code:\n\n{code}"
+        messages = [{"role": "user", "content": prompt}]
+        
+        stream = client.chat.completions.create(
+            model="Qwen/Qwen2.5-Coder-32B-Instruct",
+            messages=messages,
+            temperature=0.5,
+            max_tokens=512,
+            top_p=0.7,
+            stream=True
+        )
+        
+        for chunk in stream:
+            await websocket.send_text(chunk.choices[0].delta.content)
+
+    except WebSocketDisconnect:
+        print("Client disconnected")