Update app/main.py

app/main.py

@@ -24,10 +24,14 @@ def get_db():
         db.close()
 
 
-async def generate_stream(input_data: schemas.BaseModel, runnable: Runnable, callbacks: List[BaseCallbackHandler]=[]):
-    for output in runnable.stream(input_data.dict(), config={"callbacks": callbacks}): 
-        data = WellKnownLCSerializer().dumps(output).decode("utf-8")
-        yield {'data': data, "event": "data"} 
+async def generate_stream(input_data: schemas.BaseModel, callbacks):
+    stream = rag_chain.stream(
+        input_data.dict(),
+        config={"callbacks": [callback]}
+    )
+    for output in stream:
+        data = orjson.dumps(output).decode("utf-8")
+        yield {'data': data, "event": "data"}
     yield {"event": "end"}
 
 
@@ -82,12 +86,8 @@ async def rag_stream(request: Request, db: Session = Depends(get_db)):
         question=user_request.question,
         chat_history=format_chat_history(chat_history),
     )
-
-    return EventSourceResponse(generate_stream(
-        rag_input, 
-        rag_chain,
-        [LogResponseCallback(user_request, db)]
-    ))
+    callback = LogResponseCallback(user_request, db)
+    return EventSourceResponse(generate_stream(rag_input, callback))
 
 @app.post("/filtered_rag/stream")
 async def filtered_rag_stream(request: Request, db: Session = Depends(get_db)):