Update main.py

main.py

@@ -30,7 +30,7 @@ class InputData(BaseModel):
     json_prompt: str
     history: str = ""
 
-@app.post("/generate-response/")
+@@app.post("/generate-response/")
 async def generate_response(data: InputData) -> Any:
     client = InferenceClient(model=data.model, token=HF_TOKEN)
 
@@ -50,10 +50,49 @@ async def generate_response(data: InputData) -> Any:
 
     seed = random.randint(0, 2**32 - 1)
 
-    models_to_try = [data.model] + FALLBACK_MODELS
+    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
+            else:
+                cleaned_response["Sentence count"] = len(cleaned_response["New response"])
+
+            data.history += str(cleaned_response)
+
+            return cleaned_response
+
+    except Exception as e:
+        print(f"Primary model {data.model} failed with error: {e}")
 
-    for model in models_to_try:
+    # If the primary model fails, try fallback models
+    for model in FALLBACK_MODELS:
         try:
+            client = InferenceClient(model=model, token=HF_TOKEN)
             response = client.text_generation(inputs,
                                               temperature=1.0,
                                               max_new_tokens=1000,