File size: 2,434 Bytes
b9c7274 ed5def4 b9c7274 ed5def4 b9c7274 667b788 b9c7274 ed5def4 b9c7274 ed5def4 b9c7274 |
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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import chainlit as cl
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import StrOutputParser
from langchain.schema.runnable import RunnableConfig, RunnablePassthrough
from langchain_openai import ChatOpenAI
import configs
from prompts import prompt
from utils import format_docs, process_documents
doc_search = process_documents(configs.DOCS_STORAGE_PATH)
model = ChatOpenAI(name=configs.CHAT_MODEL, streaming=True)
@cl.on_chat_start
async def on_chat_start():
retriever = doc_search.as_retriever(search_kwargs={"k": 10})
runnable = (
{"context": retriever | format_docs, "question": RunnablePassthrough()}
| prompt
| model
| StrOutputParser()
)
cl.user_session.set("runnable", runnable)
@cl.on_message
async def on_message(message: cl.Message):
runnable = cl.user_session.get("runnable")
msg = cl.Message(content="")
class PostMessageHandler(BaseCallbackHandler):
"""
Callback handler for handling the retriever and LLM processes.
Used to post the sources of the retrieved documents as a Chainlit element.
"""
def __init__(self, msg: cl.Message):
BaseCallbackHandler.__init__(self)
self.msg = msg
self.sources = set() # To store unique pairs
def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
for d in documents:
source_page_pair = d.metadata["source"]
self.sources.add(source_page_pair) # Add unique pairs to the set
def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
if len(self.sources):
sources_text = "\n".join(
[f"{source}#page={page}" for source, page in self.sources]
)
self.msg.elements.append(
cl.Text(name="Sources", content=sources_text, display="inline")
)
async with cl.Step(type="run", name="QA Assistant"):
async for chunk in runnable.astream(
message.content,
config=RunnableConfig(
callbacks=[cl.LangchainCallbackHandler(), PostMessageHandler(msg)]
),
):
await msg.stream_token(chunk)
with open(configs.HIISTORY_FILE, "a") as f:
f.write(f"""{message.content}[SEP]{msg.content}[END]\n\n""")
await msg.send()
|