utils.py

"""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"