Update main.py

main.py

@@ -1,12 +1,11 @@
 from fastapi import FastAPI, HTTPException
-from typing import Any
+from typing import Any, Dict
 from pydantic import BaseModel
 from os import getenv
 from huggingface_hub import InferenceClient
 import random
 from json_repair import repair_json
 import nltk
-import sys
 
 app = FastAPI()
 
@@ -26,7 +25,7 @@ class InputData(BaseModel):
     history: str = ""
 
 @app.post("/generate-response/")
-async def generate_response(data: InputData) -> Any:
+async def generate_response(data: InputData) -> Dict[str, Any]:
     client = InferenceClient(model=data.model, token=HF_TOKEN)
 
     sentences = tokenizer.tokenize(data.user_input)
@@ -38,27 +37,27 @@ async def generate_response(data: InputData) -> Any:
     data.history += data.prompt_template.replace("{Prompt}", str(data_dict))
 
     inputs = (
-        data.system_prompt_template.replace("{SystemPrompt}",
-                                       data.system_prompt) +
+        data.system_prompt_template.replace("{SystemPrompt}", data.system_prompt) +
         data.system_prompt_template.replace("{SystemPrompt}", data.json_prompt) +
-        data.history)
+        data.history
+    )
 
     seed = random.randint(0, 2**32 - 1)
 
     try:
-        response = client.text_generation(inputs,
-                                          temperature=1.0,
-                                          max_new_tokens=1000,
-                                          seed=seed)
+        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)
+        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():
@@ -72,6 +71,7 @@ async def generate_response(data: InputData) -> Any:
                         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
@@ -80,7 +80,10 @@ async def generate_response(data: InputData) -> Any:
 
             data.history += str(cleaned_response)
 
-            return cleaned_response
+            return {
+                "response": cleaned_response,
+                "history": data.history
+            }
 
     except Exception as e:
         print(f"Model {data.model} failed with error: {e}")