rohan13's picture
Duplicate from rohan13/coursera-qa-bot
9e35a98
import os
import pickle
import re
from typing import List, Union
import faiss
from langchain import OpenAI, LLMChain
from langchain.agents import ConversationalAgent
from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import DirectoryLoader, TextLoader, UnstructuredHTMLLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import BaseChatPromptTemplate
from langchain.schema import AgentAction, AgentFinish, HumanMessage
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
os.environ['OPENAI_API_KEY'] = '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'
pickle_file = "open_ai.pkl"
index_file = "open_ai.index"
gpt_3_5 = OpenAI(model_name='gpt-4',temperature=0)
embeddings = OpenAIEmbeddings()
chat_history = []
memory = ConversationBufferWindowMemory(memory_key="chat_history")
gpt_3_5_index = None
class CustomOutputParser(AgentOutputParser):
def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
# Check if agent replied without using tools
if "AI:" in llm_output:
return AgentFinish(return_values={"output": llm_output.split("AI:")[-1].strip()},
log=llm_output)
# Check if agent should finish
if "Final Answer:" in llm_output:
return AgentFinish(
# Return values is generally always a dictionary with a single `output` key
# It is not recommended to try anything else at the moment :)
return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
log=llm_output,
)
# Parse out the action and action input
regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)"
match = re.search(regex, llm_output, re.DOTALL)
if not match:
raise ValueError(f"Could not parse LLM output: `{llm_output}`")
action = match.group(1).strip()
action_input = match.group(2)
# Return the action and action input
return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
# Set up a prompt template
class CustomPromptTemplate(BaseChatPromptTemplate):
# The template to use
template: str
# The list of tools available
tools: List[Tool]
def format_messages(self, **kwargs) -> str:
# Get the intermediate steps (AgentAction, Observation tuples)
# Format them in a particular way
intermediate_steps = kwargs.pop("intermediate_steps")
thoughts = ""
for action, observation in intermediate_steps:
thoughts += action.log
thoughts += f"\nObservation: {observation}\nThought: "
# Set the agent_scratchpad variable to that value
kwargs["agent_scratchpad"] = thoughts
# Create a tools variable from the list of tools provided
kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
# Create a list of tool names for the tools provided
kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
formatted = self.template.format(**kwargs)
return [HumanMessage(content=formatted)]
def get_search_index():
global gpt_3_5_index
if os.path.isfile(pickle_file) and os.path.isfile(index_file) and os.path.getsize(pickle_file) > 0:
# Load index from pickle file
with open(pickle_file, "rb") as f:
search_index = pickle.load(f)
else:
search_index = create_index()
gpt_3_5_index = search_index
def create_index():
source_chunks = create_chunk_documents()
search_index = search_index_from_docs(source_chunks)
faiss.write_index(search_index.index, index_file)
# Save index to pickle file
with open(pickle_file, "wb") as f:
pickle.dump(search_index, f)
return search_index
def search_index_from_docs(source_chunks):
# print("source chunks: " + str(len(source_chunks)))
# print("embeddings: " + str(embeddings))
search_index = FAISS.from_documents(source_chunks, embeddings)
return search_index
def get_html_files():
loader = DirectoryLoader('docs', glob="**/*.html", loader_cls=UnstructuredHTMLLoader, recursive=True)
document_list = loader.load()
return document_list
def fetch_data_for_embeddings():
document_list = get_text_files()
document_list.extend(get_html_files())
print("document list" + str(len(document_list)))
return document_list
def get_text_files():
loader = DirectoryLoader('docs', glob="**/*.txt", loader_cls=TextLoader, recursive=True)
document_list = loader.load()
return document_list
def create_chunk_documents():
sources = fetch_data_for_embeddings()
splitter = CharacterTextSplitter(separator=" ", chunk_size=800, chunk_overlap=0)
source_chunks = splitter.split_documents(sources)
print("sources" + str(len(source_chunks)))
return source_chunks
def get_qa_chain(gpt_3_5_index):
global gpt_3_5
return ConversationalRetrievalChain.from_llm(gpt_3_5, chain_type="stuff", get_chat_history=get_chat_history,
retriever=gpt_3_5_index.as_retriever(), return_source_documents=True, verbose=True)
def get_chat_history(inputs) -> str:
res = []
for human, ai in inputs:
res.append(f"Human:{human}\nAI:{ai}")
return "\n".join(res)
def generate_answer(question) -> str:
global chat_history, gpt_3_5_index
gpt_3_5_chain = get_qa_chain(gpt_3_5_index)
result = gpt_3_5_chain(
{"question": question, "chat_history": chat_history, "vectordbkwargs": {"search_distance": 0.6}})
chat_history = [(question, result["answer"])]
sources = []
for document in result['source_documents']:
source = document.metadata['source']
sources.append(source.split('/')[-1].split('.')[0])
source = ',\n'.join(set(sources))
return result['answer'] + '\nSOURCES: ' + source
def get_agent_chain(prompt, tools):
global gpt_3_5
llm_chain = LLMChain(llm=gpt_3_5, prompt=prompt)
agent = ConversationalAgent(llm_chain=llm_chain, tools=tools, verbose=True)
agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory,
intermediate_steps=True)
return agent_chain
def get_prompt_and_tools():
tools = get_tools()
prefix = """Have a conversation with a human, answering the following questions as best you can. Always try to use Vectorstore first. Your name is Coursera Bot because your knowledge base is Coursera course. You have access to the following tools:"""
suffix = """Begin! If you used vectorstore tool, ALWAYS return a "SOURCES" part in your answer"
{chat_history}
Question: {input}
{agent_scratchpad}
sources:"""
prompt = ConversationalAgent.create_prompt(
tools,
prefix=prefix,
suffix=suffix,
input_variables=["input", "chat_history", "agent_scratchpad"]
)
return prompt, tools
def get_tools():
tools = [
Tool(
name="Vectorstore",
func=generate_answer,
description="useful for when you need to answer questions about the coursera course on 3D Printing.",
return_direct=True
)]
return tools
def get_custom_agent(prompt, tools):
llm_chain = LLMChain(llm=gpt_3_5, prompt=prompt)
output_parser = CustomOutputParser()
tool_names = [tool.name for tool in tools]
agent = LLMSingleActionAgent(
llm_chain=llm_chain,
output_parser=output_parser,
stop=["\nObservation:"],
allowed_tools=tool_names
)
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory,
intermediate_steps=True)
return agent_executor
def get_prompt_and_tools_for_custom_agent():
template = """
Have a conversation with a human, answering the following questions as best you can.
ALWAYS try to use Vectorstore first.
You are a teaching assistant for a Coursera Course: The 3D Printing Evolution and can answer any question about that using vectorstore . You have access to the following tools:
{tools}
ALWAYS use one of the 2 formats listed below to respond.
To answer for the new input, use the following format:
New Input: the input question you must answer
Thought: Do I need to use a tool? Yes
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question. SOURCES: the sources referred to find the final answer
When you have a response to say to the Human and DO NOT need to use a tool:
1. DO NOT return "SOURCES" if you did not use any tool.
2. You MUST use this format:
```
Thought: Do I need to use a tool? No
AI: [your response here]
```
Begin! Remember to speak as a personal assistant when giving your final answer.
ALWAYS return a "SOURCES" part in your answer, if you used any tool.
Previous conversation history:
{chat_history}
New input: {input}
{agent_scratchpad}
SOURCES:"""
tools = get_tools()
prompt = CustomPromptTemplate(
template=template,
tools=tools,
# This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically
# This includes the `intermediate_steps` variable because that is needed
input_variables=["input", "intermediate_steps", "chat_history"]
)
return prompt, tools