Serjesh/rag_pure_python

1# You can find this code for Chainlit python streaming here (https://docs.chainlit.io/concepts/streaming/python)2import sys3import os4sys.path.append('../../lutil')5import openai  # importing openai for API usage6import chainlit as cl  # importing chainlit for our app7from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools8from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools9from dotenv import load_dotenv10from aimakerspace.text_utils import TextFileLoader, CharacterTextSplitter11from aimakerspace.vectordatabase import VectorDatabase12import asyncio13from raq_qa_reterieval_wandb import RetrievalAugmentedQAPipeline,raqa_prompt,user_prompt14from aimakerspace.openai_utils.chatmodel import ChatOpenAI15import wandb16 17 18 19 20load_dotenv()21openai.api_key = os.environ["OPENAI_API_KEY"]22os.environ["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"]23 24 25 26@cl.on_chat_start  # marks a function that will be executed at the start of a user session27async def start_chat():28    msg = cl.Message(29        content=f"Loading Dataset ...", disable_human_feedback=True30    )31    await msg.send()32    text_loader = TextFileLoader("../../data/KingLear.txt")33    documents = text_loader.load_documents()34    text_splitter = CharacterTextSplitter()35    split_documents = text_splitter.split_texts(documents)36    vector_db = VectorDatabase()37    vector_db = asyncio.run(vector_db.abuild_from_list(split_documents))38    chat_openai = ChatOpenAI()39    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(40    vector_db_retriever=vector_db,41    llm=chat_openai,42    wandb_project="RAQ in pure python HF")43    msg.content = f"Dataset loading is done. You can now ask questions!"44    await msg.update()45    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)46 47@cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user48async def main(message: str):49    # settings = cl.user_session.get("settings")50    chain = cl.user_session.get("chain")  51 52    output = chain.run_pipeline(message)53    print(output)54    msg = cl.Message(content=f"{output}")55    # msg.prompt = output56    await msg.send()57