import os from llama_index import GPTSimpleVectorIndex, SimpleDirectoryReader, download_loader from llama_index import Document, LLMPredictor, PromptHelper, QuestionAnswerPrompt, JSONReader from langchain.llms import OpenAIChat, OpenAI from zipfile import ZipFile from googlesearch import search as google_search from baidusearch.baidusearch import search as baidu_search import traceback import openai from utils import * def save_index(index, index_name, exist_ok=False): file_path = f"./index/{index_name}.json" if not os.path.exists(file_path) or exist_ok: index.save_to_disk(file_path) print(f'Saved file "{file_path}".') else: i = 1 while True: new_file_path = f'{os.path.splitext(file_path)[0]}_{i}{os.path.splitext(file_path)[1]}' if not os.path.exists(new_file_path): index.save_to_disk(new_file_path) print(f'Saved file "{new_file_path}".') break i += 1 def construct_index(api_key, file_list, index_name, max_input_size=4096, num_outputs=512, max_chunk_overlap=20, raw=False): documents = [] if not raw: txt_set = [] for file in file_list: if os.path.splitext(file.name)[1] == '.pdf': CJKPDFReader = download_loader("CJKPDFReader") loader = CJKPDFReader() documents += loader.load_data(file=file.name) elif os.path.splitext(file.name)[1] == '.docx': DocxReader = download_loader("DocxReader") loader = DocxReader() documents += loader.load_data(file=file.name) elif os.path.splitext(file.name)[1] == '.epub': EpubReader = download_loader("EpubReader") loader = EpubReader() documents += loader.load_data(file=file.name) else: with open(file.name, 'r', encoding="utf-8") as f: txt_set.append(f.read()) documents += [Document(k) for k in txt_set] else: documents += [Document(k.text.encode("UTF-8", errors="strict").decode()) for k in file_list] # Customizing LLM llm_predictor = LLMPredictor(llm=OpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key)) prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap) index = GPTSimpleVectorIndex(documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper) if not raw: save_index(index, index_name) newlist = refresh_json_list(plain=True) return newlist, newlist else: save_index(index, index_name, exist_ok=True) return index def chat_ai(api_key, index_select, question, prompt_tmpl, sim_k, chat_tone ,context, chatbot, search_mode=[], suggested_user_question = ""): os.environ["OPENAI_API_KEY"] = api_key print(f"Question: {question}") if question=="": question = suggested_user_question if chat_tone == 0: temprature = 2 elif chat_tone == 1: temprature = 1 else: temprature = 0.5 if not search_mode: response = ask_ai(api_key, index_select, question, prompt_tmpl, sim_k, temprature, context) else: print(f"You asked: {question}") BeautifulSoupWebReader = download_loader("BeautifulSoupWebReader") loader = BeautifulSoupWebReader() chat = OpenAI(model_name="gpt-3.5-turbo", openai_api_key=api_key) search_terms = chat.generate([f"Please extract search terms from the user’s question. The search terms is a concise sentence, which will be searched on Google to obtain relevant information to answer the user’s question, too generalized search terms doesn’t help. Please provide no more than two search terms. Please provide the most relevant search terms only, the search terms should directly correspond to the user’s question. Please separate different search items with commas, with no quote marks. The user’s question is: {question}"]).generations[0][0].text.strip() search_terms = search_terms.replace('"', '') search_terms = search_terms.replace(".", "") links = [] for keywords in search_terms.split(","): keywords = keywords.strip() for search_engine in search_mode: if "Google" in search_engine: print(f"Googling: {keywords}") search_iter = google_search(keywords, num_results=5) links += [next(search_iter) for _ in range(10)] if "Baidu" in search_engine: print(f"Baiduing: {keywords}") search_results = baidu_search(keywords, num_results=5) links += [i["url"] for i in search_results if i["url"].startswith("http") and (not "@" in i["url"])] if "Manual" in search_engine: print(f"Searching manually: {keywords}") print("Please input links manually. (Enter 'q' to quit.)") while True: link = input("请手动输入一个链接：\n") if link == "q": break else: links.append(link) links = list(set(links)) if len(links) == 0: msg = "No links found." print(msg) chatbot.append((question, msg)) return context, chatbot, gr.Dropdown.update(choices=[]) print("Extracting data from links...") print('\n'.join(links)) documents = loader.load_data(urls=links) # convert to utf-8 encoding index = construct_index(api_key, documents, " ".join(search_terms.split(",")), raw=True) print("Generating response...") response = ask_ai(api_key, index_select, question, prompt_tmpl, sim_k, temprature, context, raw = index) response = response.split("\n") suggested_next_turns = [] for index, line in enumerate(response): if "next user turn" in line: suggested_next_turns = response[index+1:] response = response[:index] break suggested_next_turns = [i.split()[1] for i in suggested_next_turns] response = "\n".join(response) response = parse_text(response) context.append({"role": "user", "content": question}) context.append({"role": "assistant", "content": response}) chatbot.append((question, response)) os.environ["OPENAI_API_KEY"] = "" return context, chatbot, gr.Dropdown.update(choices=suggested_next_turns) def ask_ai(api_key, index_select, question, prompt_tmpl, sim_k=1, temprature=0, prefix_messages=[], raw = None): os.environ["OPENAI_API_KEY"] = api_key if raw is not None: index = raw else: index = load_index(index_select) prompt = QuestionAnswerPrompt(prompt_tmpl) llm_predictor = LLMPredictor(llm=OpenAI(temperature=temprature, model_name="gpt-3.5-turbo", openai_api_key=api_key, prefix_messages=prefix_messages)) try: response = index.query(question, llm_predictor=llm_predictor, similarity_top_k=sim_k, text_qa_template=prompt) except: traceback.print_exc() return "" print(f"Response: {response.response}") os.environ["OPENAI_API_KEY"] = "" return response.response def load_index(index_name): index_path = f"./index/{index_name}.json" if not os.path.exists(index_path): return None index = GPTSimpleVectorIndex.load_from_disk(index_path) return index def display_json(json_select): json_path = f"./index/{json_select}.json" if not os.path.exists(json_path): return None documents = JSONReader().load_data(f"./index/{json_select}.json") return documents[0]