|
from dotenv import load_dotenv |
|
import openai |
|
import chainlit as cl |
|
from aimakerspace.vectordatabase import VectorDatabase |
|
from aimakerspace.vectordatabase import asyncio |
|
from aimakerspace.text_utils import TextFileLoader, CharacterTextSplitter |
|
import os |
|
import openai |
|
from getpass import getpass |
|
from aimakerspace.openai_utils.prompts import ( |
|
UserRolePrompt, |
|
SystemRolePrompt, |
|
AssistantRolePrompt, |
|
) |
|
from aimakerspace.openai_utils.chatmodel import ChatOpenAI |
|
|
|
|
|
load_dotenv() |
|
os.environ["OPENAI_API_KEY"] ="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" |
|
openai.api_key = os.environ["OPENAI_API_KEY"] |
|
|
|
def load(filename): |
|
text_loader = TextFileLoader(filename) |
|
documents = text_loader.load_documents() |
|
return documents |
|
|
|
model_name = "gpt-4" |
|
|
|
filename = "data/KingLear.txt" |
|
|
|
vector_db = VectorDatabase() |
|
documents = load(filename) |
|
text_splitter = CharacterTextSplitter() |
|
split_documents = text_splitter.split_texts(documents) |
|
vector_db = asyncio.run(vector_db.abuild_from_list(split_documents)) |
|
|
|
|
|
user_prompt_template = "{content}" |
|
user_role_prompt = UserRolePrompt(user_prompt_template) |
|
system_prompt_template = ( |
|
"You are an expert in {expertise}, you always answer in a kind way." |
|
) |
|
system_role_prompt = SystemRolePrompt(system_prompt_template) |
|
RAQA_PROMPT_TEMPLATE = """ |
|
Use the provided context to answer the user's query. |
|
|
|
You may not answer the user's query unless there is specific context in the following text. |
|
|
|
If you do not know the answer, or cannot answer, please respond with "I don't know". |
|
|
|
Context: |
|
{context} |
|
""" |
|
|
|
raqa_prompt = SystemRolePrompt(RAQA_PROMPT_TEMPLATE) |
|
|
|
USER_PROMPT_TEMPLATE = """ |
|
User Query: |
|
{user_query} |
|
""" |
|
|
|
user_prompt = UserRolePrompt(USER_PROMPT_TEMPLATE) |
|
class RetrievalAugmentedQAPipeline: |
|
def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None: |
|
self.llm = llm |
|
self.vector_db_retriever = vector_db_retriever |
|
|
|
def run_pipeline(self, user_query: str) -> str: |
|
context_list = self.vector_db_retriever.search_by_text(user_query, k=4) |
|
|
|
context_prompt = "" |
|
for context in context_list: |
|
context_prompt += context[0] + "\n" |
|
|
|
formatted_system_prompt = raqa_prompt.create_message(context=context_prompt) |
|
|
|
formatted_user_prompt = user_prompt.create_message(user_query=user_query) |
|
|
|
return self.llm.run([formatted_system_prompt, formatted_user_prompt]) |
|
|
|
async def stream_pipeline(self, user_query: str, message_history: [], msg: cl.Message) -> str: |
|
context_list = self.vector_db_retriever.search_by_text(user_query, k=4) |
|
|
|
context_prompt = "" |
|
for context in context_list: |
|
context_prompt += context[0] + "\n" |
|
|
|
formatted_system_prompt = raqa_prompt.create_message(context=context_prompt) |
|
|
|
formatted_user_prompt = user_prompt.create_message(user_query=user_query) |
|
|
|
message_history.append(formatted_system_prompt) |
|
message_history.append(formatted_user_prompt) |
|
|
|
await self.llm.stream_with_cl_message(message_history=message_history, chainlit_msg=msg) |
|
|
|
|
|
@cl.on_chat_start |
|
def start_chat(): |
|
cl.user_session.set( |
|
"message_history", |
|
[{"role": "system", "content": "You are a helpful assistant."}], |
|
) |
|
settings = { |
|
"temperature": 0.7, |
|
"max_tokens": 500, |
|
"top_p": 1, |
|
"frequency_penalty": 0, |
|
"presence_penalty": 0, |
|
} |
|
cl.user_session.set("settings", settings) |
|
|
|
|
|
@cl.on_message |
|
async def main(message: str): |
|
message_history = cl.user_session.get("message_history") |
|
|
|
qaPipeline = RetrievalAugmentedQAPipeline(vector_db_retriever=vector_db, llm=ChatOpenAI(model_name=model_name)) |
|
msg = cl.Message(content="") |
|
|
|
await qaPipeline.stream_pipeline(user_query=message, message_history=message_history, msg=msg) |
|
await msg.send() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
