| import re |
|
|
| import chainlit as cl |
| import tiktoken |
| from langchain.callbacks.base import BaseCallbackHandler |
|
|
|
|
| def format_docs(documents, max_context_size=100000, separator="\n\n"): |
| context = "" |
| encoder = tiktoken.get_encoding("cl100k_base") |
| i = 0 |
| for doc in documents: |
| i += 1 |
| if len(encoder.encode(context)) < max_context_size: |
| source = doc.metadata["link"] |
| title = doc.metadata["title"] |
| context += ( |
| f"Article: {title}\n" |
| + doc.page_content |
| + f"\nSource: {source}" |
| + separator |
| ) |
| return context |
|
|
|
|
| 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 = [] |
|
|
| def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs): |
| for d in documents: |
| source_doc = d.page_content + "\nSource: " + d.metadata["link"] |
| self.sources.append(source_doc) |
|
|
| def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs): |
| if len(self.sources): |
| |
| sources_element = [ |
| cl.Text(name=f"source_{idx+1}", content=content) |
| for idx, content in enumerate(self.sources) |
| ] |
| source_names = [el.name for el in sources_element] |
| self.msg.elements += sources_element |
| self.msg.content += f"\nSources: {', '.join(source_names)}" |
|
|
| def clean_text(text): |
| text = re.sub("[Tt]weet", "", text) |
| text = re.sub(r"\ +", " ", text) |
| text = re.sub(r"\n+", "\n", text) |
| return text.strip() |
|
|
