Update app.py

app.py

@@ -1,6 +1,5 @@
 from llama_cpp import Llama
 from concurrent.futures import ThreadPoolExecutor, as_completed
-import re
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
@@ -51,6 +50,7 @@ for config in model_configs:
 
 class ChatRequest(BaseModel):
     message: str
+    max_tokens_per_part: int = 256
 
 def normalize_input(input_text):
     return input_text.strip()
@@ -66,15 +66,28 @@ def remove_duplicates(text):
             seen_lines.add(line)
     return '\n'.join(unique_lines)
 
-def generate_model_response(model, inputs):
+def generate_model_response(model, inputs, max_tokens_per_part):
     try:
         if model is None:
-            return ""
-        response = model(inputs)
-        return remove_duplicates(response['choices'][0]['text'])
+            return []
+        full_response = ""
+        responses = []
+        tokens_generated = 0
+        while True:
+            response_part = model(inputs, max_tokens=max_tokens_per_part, stop=["\n\n"])
+            text = response_part['choices'][0]['text']
+            if not text.strip():
+                break
+            full_response += text
+            tokens_generated += len(response_part['choices'][0]['token'])
+            responses.append(remove_duplicates(text))
+            if "eos_token" in response_part['choices'][0]['token']:
+                break
+            inputs = ""
+        return responses
     except Exception as e:
         print(f"Error generating response: {e}")
-        return f"Error: {e}"
+        return [f"Error: {e}"]
 
 app = FastAPI()
 origins = ["*"]
@@ -90,17 +103,23 @@ app.add_middleware(
 async def generate(request: ChatRequest):
     inputs = normalize_input(request.message)
     with ThreadPoolExecutor() as executor:
-        futures = [executor.submit(generate_model_response, model, inputs) for model in models.values()]
+        futures = [executor.submit(generate_model_response, model, inputs, request.max_tokens_per_part) for model in models.values()]
         responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(models.keys(), as_completed(futures))]
 
     unique_responses = {}
-    for response in responses:
-        if response['model'] not in unique_responses and response['response']:
-            unique_responses[response['model']] = response['response']
+    for response_set in responses:
+        model_name = response_set['model']
+        if model_name not in unique_responses:
+            unique_responses[model_name] = []
+        unique_responses[model_name].extend(response_set['response'])
+
 
     formatted_response = ""
-    for model, response in unique_responses.items():
-        formatted_response += f"**{model}:**\n{response}\n\n"
+    for model, response_parts in unique_responses.items():
+        formatted_response += f"**{model}:**\n"
+        for i, part in enumerate(response_parts):
+            formatted_response += f"Part {i+1}:\n{part}\n\n"
+
 
     return {"response": formatted_response}