Change outputs to use streaming

app.py

@@ -2,15 +2,20 @@
 import collections
 import os
 from itertools import islice
+from queue import Queue
 
+from anyio.from_thread import start_blocking_portal
 import gradio as gr
 from diff_match_patch import diff_match_patch
 from langchain.chains import LLMChain
-from langchain.chat_models import PromptLayerChatOpenAI
+from langchain.chat_models import PromptLayerChatOpenAI, ChatOpenAI
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
 from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
 from langchain.schema import HumanMessage
+from langchain.callbacks.manager import AsyncCallbackManager
+
+from util import StreamingLLMCallbackHandler, concatenate_generators
 
 GRAMMAR_PROMPT = "Proofread for grammar and spelling without adding new paragraphs:\n{content}"
 
@@ -71,6 +76,7 @@ def load_chain(api_key, api_type):
         "model_name": "gpt-3.5-turbo",
         "api_key": api_key, # deliberately not use "openai_api_key" and other openai args since those apply globally
         "pl_tags": ["grammar"],
+        "streaming": True,
     }
     if api_type == "OpenAI":
       llm = PromptLayerChatOpenAI(**shared_args)
@@ -106,16 +112,35 @@ def load_chain(api_key, api_type):
     return chain, llm, chain_intro, chain_body1
 
 
-def run_diff(content, chain):
+def run_diff(content, chain: LLMChain):
   chain.memory.clear()
-  edited = "\n".join([(chain.run(x) if should_check else x) for x, should_check in split_paragraphs(content)])
+  edited = chain.run(content)
   return diff_words(content, edited) + (edited,)
 
-def run(content, chain):
+# https://github.com/hwchase17/langchain/issues/2428#issuecomment-1512280045
+def run(content, chain: LLMChain):
   chain.memory.clear()
-  return chain.run(content)
 
-def run_followup(followup_question, input_vars, chain, chat):
+  q = Queue()
+  job_done = object()
+  async def task():
+    result = await chain.arun(content, callbacks=[StreamingLLMCallbackHandler(q)])
+    q.put(job_done)
+    return result
+
+  with start_blocking_portal() as portal:
+    portal.start_task_soon(task)
+
+    content = ""
+    while True:
+      next_token = q.get(True, timeout=10)
+      if next_token is job_done:
+        break
+      content += next_token
+      yield content
+
+# TODO share code with above
+def run_followup(followup_question, input_vars, chain, chat: ChatOpenAI):
 
   history = [HumanMessage(content=chain.prompt.format(content=m.content)) if isinstance(m, HumanMessage) else m
              for m in chain.memory.chat_memory.messages]
@@ -123,16 +148,37 @@ def run_followup(followup_question, input_vars, chain, chat):
       *history,
       HumanMessagePromptTemplate.from_template(followup_question)])
   messages = prompt.format_prompt(**input_vars).to_messages()
-  return chat(messages).content
+
+  q = Queue()
+  job_done = object()
+  async def task():
+    result = await chat.agenerate([messages], callbacks=[StreamingLLMCallbackHandler(q)])
+    q.put(job_done)
+    return result.generations[0][0].message.content
+
+  with start_blocking_portal() as portal:
+    portal.start_task_soon(task)
+
+    content = ""
+    while True:
+      next_token = q.get(True, timeout=10)
+      if next_token is job_done:
+        break
+      content += next_token
+      yield content
+
 
 def run_body(content, title, chain, llm):
   if not title:
     return "Please enter the book title."
-  output1 = run(content, chain)
-  output2 = run_followup(BODY_PROMPT2, {}, chain, llm)
-  output3 = run_followup(BODY_PROMPT3, {"title": title}, chain, llm)
-  output3 = output3.split("----")[-1]
-  return output1 + "\n\n" + output2 + "\n\n7. Whether supporting evidence is from the book:" + output3
+
+  yield from concatenate_generators(
+    run(content, chain),
+    "\n\n",
+    run_followup(BODY_PROMPT2, {}, chain, llm),
+    "\n\n7. Whether supporting evidence is from the book:",
+    (output.split("----")[-1] for output in run_followup(BODY_PROMPT3, {"title": title}, chain, llm))
+    )
 
 def run_custom(content, llm, prompt):
   chain = LLMChain(llm=llm,
@@ -143,9 +189,9 @@ def run_custom(content, llm, prompt):
                   ))
   return chain.run(content), chain
 
+# not currently used
 def split_paragraphs(text):
-  # return [(x, x != "" and not x.startswith("#") and not x.isspace()) for x in text.split("\n")]
-  return [(text, True)]
+  return [(x, x != "" and not x.startswith("#") and not x.isspace()) for x in text.split("\n")]
 
 def sliding_window(iterable, n):
     # sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
@@ -167,7 +213,7 @@ def diff_words(content, edited):
   diff = dmp.diff_main(content, edited)
   dmp.diff_cleanupSemantic(diff)
   diff += [(None, None)]
-  # print(diff)
+
   for [(change, text), (next_change, next_text)] in sliding_window(diff, 2):
     if change == 0:
       before.append((text, None))
@@ -364,4 +410,5 @@ with demo:
 port = os.environ.get("SERVER_PORT", None)
 if port:
   port = int(port)
+demo.queue()
 demo.launch(debug=True, server_port=port)

util.py

@@ -0,0 +1,25 @@
+from typing import Any
+from types import GeneratorType
+from langchain.callbacks.base import AsyncCallbackHandler
+
+class StreamingLLMCallbackHandler(AsyncCallbackHandler):
+    """Callback handler for streaming LLM responses to a queue."""
+
+    def __init__(self, q):
+        self.q = q
+
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        self.q.put(token)
+
+
+def concatenate_generators(*args):
+    final_outputs = ""
+    for g in args:
+        if isinstance(g, GeneratorType):
+            for v in g:
+                yield final_outputs + v
+            result = v
+        else:
+            yield final_outputs + g
+            result = g
+        final_outputs += result