Update main.py

main.py

@@ -21,8 +21,6 @@ FALLBACK_MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mistral-7B-Instruct-v0.1"
 ]
 
-MAX_RETRIES = 3  # Maximum number of retries
-
 class InputData(BaseModel):
     model: str
     system_prompt_template: str
@@ -53,48 +51,45 @@ async def generate_response(data: InputData) -> Any:
     seed = random.randint(0, 2**32 - 1)
 
     models_to_try = [data.model] + FALLBACK_MODELS
-    retries = 0
-
-    while retries < MAX_RETRIES:
-        for model in models_to_try:
-            try:
-                response = client.text_generation(inputs,
-                                                  temperature=1.0,
-                                                  max_new_tokens=1000,
-                                                  seed=seed)
-
-                strict_response = str(response)
-                
-                repaired_response = repair_json(strict_response,
-                                                return_objects=True)
-
-                if isinstance(repaired_response, str):
-                    raise HTTPException(status_code=500, detail="Invalid response from model")
-                else:
-                    cleaned_response = {}
-                    for key, value in repaired_response.items():
-                        cleaned_key = key.replace("###", "")
-                        cleaned_response[cleaned_key] = value
-
-                    for i, text in enumerate(cleaned_response["New response"]):
-                        if i <= 2:
-                            sentences = tokenizer.tokenize(text)
-                            if sentences:
-                                cleaned_response["New response"][i] = sentences[0]
-                        else:
-                            del cleaned_response["New response"][i]
-                    if cleaned_response.get("Sentence count"):
-                        if cleaned_response["Sentence count"] > 3:
-                            cleaned_response["Sentence count"] = 3
+
+    for model in models_to_try:
+        try:
+            response = client.text_generation(inputs,
+                                              temperature=1.0,
+                                              max_new_tokens=1000,
+                                              seed=seed)
+
+            strict_response = str(response)
+            
+            repaired_response = repair_json(strict_response,
+                                            return_objects=True)
+
+            if isinstance(repaired_response, str):
+                raise HTTPException(status_code=500, detail="Invalid response from model")
+            else:
+                cleaned_response = {}
+                for key, value in repaired_response.items():
+                    cleaned_key = key.replace("###", "")
+                    cleaned_response[cleaned_key] = value
+
+                for i, text in enumerate(cleaned_response["New response"]):
+                    if i <= 2:
+                        sentences = tokenizer.tokenize(text)
+                        if sentences:
+                            cleaned_response["New response"][i] = sentences[0]
                     else:
-                        cleaned_response["Sentence count"] = len(cleaned_response["New response"])
+                        del cleaned_response["New response"][i]
+                if cleaned_response.get("Sentence count"):
+                    if cleaned_response["Sentence count"] > 3:
+                        cleaned_response["Sentence count"] = 3
+                else:
+                    cleaned_response["Sentence count"] = len(cleaned_response["New response"])
 
-                    data.history += str(cleaned_response)
+                data.history += str(cleaned_response)
 
-                    return cleaned_response
+                return cleaned_response
 
-            except Exception as e:
-                print(f"Model {model} failed with error: {e}")
-                retries += 1
+        except Exception as e:
+            print(f"Model {model} failed with error: {e}")
 
-    raise HTTPException(status_code=500, detail="All models failed to generate response after maximum retries")
+    raise HTTPException(status_code=500, detail="All models failed to generate response")