Update write_story.py

write_story.py

@@ -11,8 +11,8 @@ from config import save_novel_chapter
 from config import generate_uuid
 
 def print_step_costs(response, model):
-    input = response['usage']['prompt_tokens']
-    output = response['usage']['completion_tokens']
+    input = response.usage.prompt_tokens
+    output = response.usage.completion_tokens
 
     if model == "gpt-4" or model == "gpt-4":
         input_per_token = 0.00003
@@ -60,7 +60,7 @@ def generate_plots(prompt):
 
     print_step_costs(response, llm_model_name)
 
-    return response['choices'][0]['message']['content'].split('\n')
+    return response.choices[0].message.content.split('\n')
 
 
 def select_most_engaging(plots):
@@ -79,7 +79,7 @@ def select_most_engaging(plots):
 
     print_step_costs(response, llm_model_name)
 
-    return response['choices'][0]['message']['content']
+    return response.choices[0].message.content
 
 
 def improve_plot(plot):
@@ -93,7 +93,7 @@ def improve_plot(plot):
 
     print_step_costs(response, llm_model_name)
 
-    return response['choices'][0]['message']['content']
+    return response.choices[0].message.content
 
 
 def get_title(plot):
@@ -109,7 +109,7 @@ def get_title(plot):
 
     print_step_costs(response, llm_model_name)
 
-    return response['choices'][0]['message']['content']
+    return response.choices[0].message.content
 
 
 def write_first_chapter(plot, first_chapter_title, writing_style, claude=True):
@@ -200,13 +200,13 @@ def write_first_chapter(plot, first_chapter_title, writing_style, claude=True):
                 {"role": "system",
                  "content": "You are a world-class fantasy writer. Your job is to take your student's rough initial draft of the first chapter of their fantasy novel, and rewrite it to be significantly better, with much more detail."},
                 {"role": "user",
-                 "content": f"Here is the high-level plot you asked your student to follow: {plot}\n\nHere is the first chapter they wrote: {response['choices'][0]['message']['content']}\n\nNow, rewrite the first chapter of this novel, in a way that is far superior to your student's chapter. It should still follow the exact same plot, but it should be far more detailed, much longer, and more engaging. Here is a description of the writing style you should use: `{writing_style}`.Please respond in Chinese."}
+                 "content": f"Here is the high-level plot you asked your student to follow: {plot}\n\nHere is the first chapter they wrote: {response.choices[0].message.content}\n\nNow, rewrite the first chapter of this novel, in a way that is far superior to your student's chapter. It should still follow the exact same plot, but it should be far more detailed, much longer, and more engaging. Here is a description of the writing style you should use: `{writing_style}`.Please respond in Chinese."}
             ]
         )
 
         print_step_costs(response, llm_model_name)
 
-        return improved_response['choices'][0]['message']['content']
+        return improved_response.choices[0].message.content
 
 
 def write_chapter(previous_chapters, plot, chapter_title, claude=True):
@@ -246,7 +246,7 @@ def write_chapter(previous_chapters, plot, chapter_title, claude=True):
 
             print_step_costs(response, llm_model_name)
 
-            return response['choices'][0]['message']['content']
+            return response.choices[0].message.content
         except:
             response = completion_with_backoff(
                 model=llm_model_name,
@@ -259,7 +259,7 @@ def write_chapter(previous_chapters, plot, chapter_title, claude=True):
 
             print_step_costs(response, llm_model_name)
 
-            return response['choices'][0]['message']['content']
+            return response.choices[0].message.content
 
 
 def generate_storyline(prompt, num_chapters):
@@ -291,7 +291,7 @@ def generate_storyline(prompt, num_chapters):
                         "Your job is to take your student's rough initial draft of the storyline of a fantasy novel in Chinese, "
                         "and rewrite it to be significantly better. Please respond in Chinese"},
             {"role": "user",
-             "content": f"Here is the draft storyline they wrote: {response['choices'][0]['message']['content']}\n\nNow, "
+             "content": f"Here is the draft storyline they wrote: {response.choices[0].message.content}\n\nNow, "
                         f"rewrite the storyline in Chinese, in a way that is far superior to your student's version. "
                         f"It should have the same number of chapters, "
                         f"but it should be much improved in as many ways as possible. "
@@ -302,7 +302,7 @@ def generate_storyline(prompt, num_chapters):
 
     print_step_costs(improved_response, llm_model_name)
 
-    return improved_response['choices'][0]['message']['content']
+    return improved_response.choices[0].message.content
 
 
 def write_to_file(prompt, content):