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"