giesAIexperiments/makerlab-bot
0
1import os2import pickle3import re4import time5from typing import List, Union6from urllib.parse import urlparse, urljoin7 8import faiss9import requests10from PyPDF2 import PdfReader11from bs4 import BeautifulSoup12from langchain import OpenAI, LLMChain13from langchain.agents import ConversationalAgent14from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser15from langchain.prompts import BaseChatPromptTemplate16from langchain.chains import ConversationalRetrievalChain17from langchain.docstore.document import Document18from langchain.embeddings import OpenAIEmbeddings19from langchain.memory import ConversationBufferWindowMemory20from langchain.schema import AgentAction, AgentFinish, HumanMessage21from langchain.text_splitter import CharacterTextSplitter22from langchain.vectorstores.faiss import FAISS23 24book_url = 'https://g.co/kgs/2VFC7u'25book_file = "Book.pdf"26url = 'https://makerlab.illinois.edu/'27 28pickle_file = "open_ai.pkl"29index_file = "open_ai.index"30 31gpt_3_5 = OpenAI(model_name='gpt-3.5-turbo',temperature=0)32 33embeddings = OpenAIEmbeddings()34 35chat_history = []36 37memory = ConversationBufferWindowMemory(memory_key="chat_history")38 39gpt_3_5_index = None40 41class CustomOutputParser(AgentOutputParser):42 43 def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:44 # Check if agent replied without using tools45 if "AI:" in llm_output:46 return AgentFinish(return_values={"output": llm_output.split("AI:")[-1].strip()},47 log=llm_output)48 # Check if agent should finish49 if "Final Answer:" in llm_output:50 return AgentFinish(51 # Return values is generally always a dictionary with a single `output` key52 # It is not recommended to try anything else at the moment :)53 return_values={"output": llm_output.split("Final Answer:")[-1].strip()},54 log=llm_output,55 )56 # Parse out the action and action input57 regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)"58 match = re.search(regex, llm_output, re.DOTALL)59 if not match:60 raise ValueError(f"Could not parse LLM output: `{llm_output}`")61 action = match.group(1).strip()62 action_input = match.group(2)63 # Return the action and action input64 return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)65 66# Set up a prompt template67class CustomPromptTemplate(BaseChatPromptTemplate):68 # The template to use69 template: str70 # The list of tools available71 tools: List[Tool]72 73 def format_messages(self, **kwargs) -> str:74 # Get the intermediate steps (AgentAction, Observation tuples)75 # Format them in a particular way76 intermediate_steps = kwargs.pop("intermediate_steps")77 thoughts = ""78 for action, observation in intermediate_steps:79 thoughts += action.log80 thoughts += f"\nObservation: {observation}\nThought: "81 # Set the agent_scratchpad variable to that value82 kwargs["agent_scratchpad"] = thoughts83 # Create a tools variable from the list of tools provided84 kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])85 # Create a list of tool names for the tools provided86 kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])87 formatted = self.template.format(**kwargs)88 return [HumanMessage(content=formatted)]89 90def get_search_index():91 global gpt_3_5_index92 if os.path.isfile(pickle_file) and os.path.isfile(index_file) and os.path.getsize(pickle_file) > 0:93 # Load index from pickle file94 with open(pickle_file, "rb") as f:95 search_index = pickle.load(f)96 else:97 search_index = create_index()98 99 gpt_3_5_index = search_index100 101 102def create_index():103 source_chunks = create_chunk_documents()104 search_index = search_index_from_docs(source_chunks)105 faiss.write_index(search_index.index, index_file)106 # Save index to pickle file107 with open(pickle_file, "wb") as f:108 pickle.dump(search_index, f)109 return search_index110 111 112def create_chunk_documents():113 sources = fetch_data_for_embeddings(url, book_file, book_url)114 # print("sources" + str(len(sources)))115 116 splitter = CharacterTextSplitter(separator=" ", chunk_size=800, chunk_overlap=0)117 118 source_chunks = splitter.split_documents(sources)119 120 for chunk in source_chunks:121 print("Size of chunk: " + str(len(chunk.page_content) + len(chunk.metadata)))122 if chunk.page_content is None or chunk.page_content == '':123 print("removing chunk: "+ chunk.page_content)124 source_chunks.remove(chunk)125 elif len(chunk.page_content) >=1000:126 print("splitting document")127 source_chunks.extend(splitter.split_documents([chunk]))128 # print("Chunks: " + str(len(source_chunks)) + "and type " + str(type(source_chunks)))129 return source_chunks130 131 132def fetch_data_for_embeddings(url, book_file, book_url):133 sources = get_website_data(url)134 sources.extend(get_document_data(book_file, book_url))135 return sources136 137def get_website_data(index_url):138 # Get all page paths from index139 paths = get_paths(index_url)140 141 # Filter out invalid links and join them with the base URL142 links = get_links(index_url, paths)143 144 return get_content_from_links(links, index_url)145 146 147def get_content_from_links(links, index_url):148 content_list = []149 for link in set(links):150 if link.startswith(index_url):151 page_data = requests.get(link).content152 soup = BeautifulSoup(page_data, "html.parser")153 154 # Get page content155 content = soup.get_text(separator="\n")156 # print(link)157 158 # Get page metadata159 metadata = {"source": link}160 161 content_list.append(Document(page_content=content, metadata=metadata))162 time.sleep(1)163 # print("content list" + str(len(content_list)))164 return content_list165 166 167def get_paths(index_url):168 index_data = requests.get(index_url).content169 soup = BeautifulSoup(index_data, "html.parser")170 paths = set([a.get('href') for a in soup.find_all('a', href=True)])171 return paths172 173 174def get_links(index_url, paths):175 links = []176 for path in paths:177 url = urljoin(index_url, path)178 parsed_url = urlparse(url)179 if parsed_url.scheme in ["http", "https"] and "squarespace" not in parsed_url.netloc:180 links.append(url)181 return links182 183 184def get_document_data(book_file, book_url):185 document_list = []186 if os.path.isfile(book_file):187 with open(book_file, 'rb') as f:188 pdf_reader = PdfReader(f)189 for i in range(len(pdf_reader.pages)):190 page_text = pdf_reader.pages[i].extract_text()191 metadata = {"source": book_url}192 document_list.append(Document(page_content=page_text, metadata=metadata))193 194 # print("document list" + str(len(document_list)))195 return document_list196 197def search_index_from_docs(source_chunks):198 # Create index from chunk documents199 # print("Size of chunk" + str(len(source_chunks)))200 search_index = FAISS.from_texts([doc.page_content for doc in source_chunks], embeddings, metadatas=[doc.metadata for doc in source_chunks])201 return search_index202 203 204def get_qa_chain(gpt_3_5_index):205 global gpt_3_5206 print("index: " + str(gpt_3_5_index))207 return ConversationalRetrievalChain.from_llm(gpt_3_5, chain_type="stuff", get_chat_history=get_chat_history,208 retriever=gpt_3_5_index.as_retriever(), return_source_documents=True, verbose=True)209 210def get_chat_history(inputs) -> str:211 res = []212 for human, ai in inputs:213 res.append(f"Human:{human}\nAI:{ai}")214 return "\n".join(res)215 216 217def generate_answer(question) -> str:218 global chat_history, gpt_3_5_index219 gpt_3_5_chain = get_qa_chain(gpt_3_5_index)220 result = gpt_3_5_chain(221 {"question": question, "chat_history": chat_history,"vectordbkwargs": {"search_distance": 0.8}})222 print("REsult: " + str(result))223 chat_history = [(question, result["answer"])]224 sources = []225 226 for document in result['source_documents']:227 source = document.metadata['source']228 sources.append(source)229 230 source = ',\n'.join(set(sources))231 return result['answer'] + '\nSOURCES: ' + source232 233 234def get_agent_chain(prompt, tools):235 global gpt_3_5236 # output_parser = CustomOutputParser()237 llm_chain = LLMChain(llm=gpt_3_5, prompt=prompt)238 agent = ConversationalAgent(llm_chain=llm_chain, tools=tools, verbose=True)239 agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory,240 intermediate_steps=True)241 return agent_chain242 243 244def get_prompt_and_tools():245 tools = get_tools()246 247 prefix = """Have a conversation with a human, answering the following questions as best you can. 248 Always try to use Vectorstore first. 249 Your name is Makerlab Bot because you are a personal assistant of Makerlab. You have access to the following tools:"""250 suffix = """Begin! If you use any tool, ALWAYS return a "SOURCES" part in your answer"251 252 {chat_history}253 Question: {input}254 {agent_scratchpad}255 SOURCES:"""256 prompt = ConversationalAgent.create_prompt(257 tools,258 prefix=prefix,259 suffix=suffix,260 input_variables=["input", "chat_history", "agent_scratchpad"]261 )262 # print("Template: " + prompt.template)263 return prompt, tools264 265 266def get_tools():267 tools = [268 Tool(269 name="Vectorstore",270 func=generate_answer,271 description="useful for when you need to answer questions about the Makerlab or 3D Printing.",272 return_direct=True273 )]274 return tools275 276def get_custom_agent(prompt, tools):277 278 llm_chain = LLMChain(llm=gpt_3_5, prompt=prompt)279 280 output_parser = CustomOutputParser()281 tool_names = [tool.name for tool in tools]282 agent = LLMSingleActionAgent(283 llm_chain=llm_chain,284 output_parser=output_parser,285 stop=["\nObservation:"],286 allowed_tools=tool_names287 )288 agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory,289 intermediate_steps=True)290 return agent_executor291 292def get_prompt_and_tools_for_custom_agent():293 template = """294 Have a conversation with a human, answering the following questions as best you can. 295 Always try to use Vectorstore first. 296 Your name is Makerlab Bot because you are a personal assistant of Makerlab. You have access to the following tools:297 298 {tools}299 300 To answer for the new input, use the following format:301 302 New Input: the input question you must answer303 Thought: Do I need to use a tool? Yes304 Action: the action to take, should be one of [{tool_names}]305 Action Input: the input to the action306 Observation: the result of the action307 ... (this Thought/Action/Action Input/Observation can repeat N times)308 Thought: I now know the final answer309 Final Answer: the final answer to the original input question. SOURCES: the sources referred to find the final answer310 311 312 When you have a response to say to the Human and DO NOT need to use a tool:313 1. DO NOT return "SOURCES" if you did not use any tool.314 2. You MUST use this format:315 ```316 Thought: Do I need to use a tool? No317 AI: [your response here]318 ```319 320 Begin! Remember to speak as a personal assistant when giving your final answer.321 ALWAYS return a "SOURCES" part in your answer, if you used any tool. 322 323 Previous conversation history:324 {chat_history}325 New input: {input}326 {agent_scratchpad}327 SOURCES:"""328 tools = get_tools()329 prompt = CustomPromptTemplate(330 template=template,331 tools=tools,332 # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically333 # This includes the `intermediate_steps` variable because that is needed334 input_variables=["input", "intermediate_steps", "chat_history"]335 )336 return prompt, tools