feat: add streaming handler

mongulu-cm · Dec 26, 2023 · a7a616d · a7a616d
1 parent 8325a9a
commit a7a616d
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/etl/experiments/ui.py b/etl/experiments/ui.py
@@ -13,7 +13,9 @@
     SystemMessagePromptTemplate,
 )
 from langchain.vectorstores import FAISS
-
+from langchain.callbacks.streaming_stdout_final_only import (
+    StreamingStdOutCallbackHandler,
+)
 
 sentry_sdk.init(
     dsn="https://a38e91a66c70912c38406fef32d86809@o4504301629407232.ingest.sentry.io/4506436450844672",
@@ -48,7 +50,7 @@
     vectors = FAISS.from_documents(data, embeddings)
     vectors.save_local(embedding_pth)
 
-llm = ChatOpenAI(max_tokens=500, temperature=0, model_name="gpt-3.5-turbo",streaming=True)
+llm = ChatOpenAI(max_tokens=500, temperature=0, model_name="gpt-3.5-turbo",streaming=True,callbacks=[StreamingStdOutCallbackHandler()])
 chain_type_kwargs = {"prompt": CHAT_PROMPT}