Update main.py

main.py

@@ -7,7 +7,6 @@ import random
 from json_repair import repair_json
 import nltk
 import sys
-sys.setrecursionlimit(10000)  # Set a higher recursion limit
 
 app = FastAPI()
 
@@ -22,6 +21,8 @@ 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
@@ -52,45 +53,48 @@ async def generate_response(data: InputData) -> Any:
     seed = random.randint(0, 2**32 - 1)
 
     models_to_try = [data.model] + FALLBACK_MODELS
-
-    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
+    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["Sentence count"] = len(cleaned_response["New response"])
+                    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
+                    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}")
+            except Exception as e:
+                print(f"Model {model} failed with error: {e}")
+                retries += 1
 
-    raise HTTPException(status_code=500, detail="All models failed to generate response")
+    raise HTTPException(status_code=500, detail="All models failed to generate response after maximum retries")