Update aAzel.py

aAzel.py

@@ -35,12 +35,11 @@ def generate_chapter_part(prompt, previous_output=None, max_tokens=2048):
         prompt = prompt + "\n\n" + previous_output
 
     response = ""
-    while True:
-        partial_response = model.generate_content([author_persona, author_guidelines, prompt], max_output_tokens=max_tokens)
+    output_length = 0
+    while output_length < max_tokens:
+        partial_response = model.generate_content([author_persona, author_guidelines, prompt], max_output_length=max_tokens - output_length)
         response += partial_response.text
-        if len(model.count_tokens(response)) >= max_tokens:
-            break
-        prompt = partial_response.text
+        output_length = len(model.count_tokens(response))
 
     return response