Update app.py

app.py

@@ -11,6 +11,7 @@ load_dotenv()
 
 app = FastAPI()
 
+# Global data storage
 global_data = {
     'models': {},
     'tokens': {
@@ -75,7 +76,7 @@ class ModelManager:
             self.loaded = True
             return models
         except Exception as e:
-            print(f"Error loading all models: {e}")
+            print(f"Error loading models: {e}")
             return []
 
 model_manager = ModelManager()
@@ -113,49 +114,34 @@ def remove_repetitive_responses(responses):
             unique_responses.append(response)
     return unique_responses
 
-def generate_chat_response(request, model_data):
-    model = model_data['model']
-    try:
-        user_input = normalize_input(request.message)
-        response = model(user_input, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
-        return response
-    except Exception as e:
-        print(f"Error generating response with model {model_data['name']}: {e}")
-        return None
-
-@app.post("/generate")
 @spaces.GPU(duration=0)
+@app.post("/generate")
 async def generate(request: ChatRequest):
     try:
-        responses = []
-        models = global_data['models']
-        for model_data in models:
-            response = generate_chat_response(request, model_data)
-            if response:
-                responses.append({
-                    "model": model_data['name'],
-                    "response": response
-                })
-        
-        if not responses:
-            raise HTTPException(status_code=500, detail="Error: No responses generated.")
+        if not global_data['models']:
+            raise HTTPException(status_code=500, detail="Models not loaded")
         
-        responses = remove_repetitive_responses(responses)
+        model = global_data['models'][0]['model']
+        inputs = normalize_input(request.message)
+        responses = model.generate(inputs, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
         best_response = responses[0] if responses else {}
+        unique_responses = remove_repetitive_responses(responses)
         return {
             "best_response": best_response,
-            "all_responses": responses
+            "all_responses": unique_responses
         }
-    except Exception:
-        pass
+    except Exception as e:
+        print(f"Error in generate endpoint: {e}")
+        return {"error": str(e)}
 
 @app.api_route("/{method_name:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
 async def handle_request(method_name: str, request: Request):
     try:
         body = await request.json()
         return {"message": "Request handled successfully", "body": body}
-    except Exception:
-        pass
+    except Exception as e:
+        print(f"Error handling request: {e}")
+        return {"error": str(e)}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)