Spaces:
Running
Running
File size: 1,586 Bytes
c7426d8 1031c5b c7426d8 4fe2f95 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
"""Module containing utility functions for the chatbot application."""
import json
from chain import RAGChain, FollowUpChain
from schema import ReqData
from retriever import DocRetriever
followUpChain = FollowUpChain()
async def generate(req: ReqData):
"""
Asynchronously generates responses based on the provided request data.
This function uses different processing chains depending on the `web` attribute of the request.
It streams chunks of data and yields server-sent events (SSE) for answers and contexts.
Additionally, it generates follow-up questions and updates citations.
Args:
req (ReqData): Request data containing user and chat info, query, and other parameters.
Yields:
str: Server-sent events (SSE) for answers, contexts, and follow-up questions in JSON format.
"""
chain = RAGChain(DocRetriever(req=req))
session_id = "/".join([req.user_id, req.chat_id])
contexts = []
for chunk in chain.stream({"input": req.query},
config={"configurable": {"session_id": session_id}}):
if 'answer' in chunk:
yield "event: answer\n"
yield f"data: {json.dumps(chunk)}\n\n"
print(chunk['answer'], end="", flush=True)
elif 'context' in chunk:
for context in chunk['context']:
yield "event: context\n"
yield f"data: {json.dumps({'context': context.metadata})}\n\n"
yield "event: questions\n"
yield f"data: {json.dumps({'questions': followUpChain.invoke(req.query, contexts)})}\n\n"
|