Update app.py

app.py

@@ -26,6 +26,49 @@ top_p = st.sidebar.slider("Top P:", 0.0, 1.0, 1.0, step=0.01)
 system_message = st.sidebar.text_area("System Message:")
 followup_message = st.sidebar.text_area("Edit Message:")
 row_count = st.session_state.get("row_count", 1)
+generate_all_edits = st.sidebar.button("Generate All Edits")
+
+if generate_all_edits:
+    # Creating a list of messages
+    messages = [st.session_state.get(f"message{i}", "") for i in range(row_count)]
+
+    jobs = Queue()
+    results = Queue()
+
+    # Create workers
+    workers = [WorkerThread(jobs, results) for _ in range(num_concurrent_calls)]
+
+    # Start all workers
+    for worker in workers:
+        worker.start()
+
+    # Put all the jobs into the queue
+    for i, message in enumerate(messages):
+        jobs.put((i, [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": message},
+            {"role": "user", "content": followup_message}
+        ]))
+
+    # Put a None for each worker to indicate the end of jobs
+    for _ in range(num_concurrent_calls):
+        jobs.put(None)
+
+    # Wait for all of the tasks to finish
+    for worker in workers:
+        worker.join()
+
+    # Collect all the results
+    while not results.empty():
+        i, response, prompt_tokens, response_tokens, word_count, error_message = results.get()
+        if error_message is not None:
+            st.write(f"Error on row {i}: {error_message}")
+        st.session_state[f"followup_response{i}"] = response
+        st.session_state[f"prompt_tokens{i}"] = prompt_tokens
+        st.session_state[f"response_tokens{i}"] = response_tokens
+        st.session_state[f"word_count{i}"] = word_count
+
+# ...
 
 if add_row:
     row_count += 1