Update main.py

main.py

@@ -39,10 +39,13 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
         for user_input in data.user_inputs:
             data.history += data.prompt_template.replace("{Prompt}", user_input)
         return {
-            "response": "",
+            "response": [],
+            "sentence_count": None,
             "history": data.history + data.end_token
         }
 
+    responses = []
+
     if data.segment:
         for user_input in data.user_inputs:
             user_sentences = tokenizer.tokenize(user_input)
@@ -74,19 +77,17 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
             ai_sentences = tokenizer.tokenize(response_str)
             if data.max_sentences is not None:
                 ai_sentences = ai_sentences[:data.max_sentences]
-            ai_response_str = "\n".join(ai_sentences)
+            responses = ai_sentences
+            sentence_count = len(ai_sentences)
         else:
-            ai_response_str = response_str
-
-        data.history += ai_response_str + "\n"
+            responses = [response_str]
+            sentence_count = None
 
-        cleaned_response = {
-            "New response": ai_sentences if data.segment else [response_str],
-            "Sentence count": len(ai_sentences) if data.segment else 1
-        }
+        data.history += response_str + "\n"
 
         return {
-            "response": cleaned_response,
+            "response": responses,
+            "sentence_count": sentence_count,
             "history": data.history + data.end_token
         }