Update llama_func.py

llama_func.py

@@ -1,6 +1,7 @@
 import os
 import logging
 
+from llama_index import GPTSimpleVectorIndex
 from llama_index import download_loader
 from llama_index import (
     Document,
@@ -9,146 +10,181 @@ from llama_index import (
     QuestionAnswerPrompt,
     RefinePrompt,
 )
+from langchain.llms import OpenAI
 import colorama
-# import PyPDF2
-from tqdm import tqdm
+
 
 from presets import *
 from utils import *
-from config import local_embedding
-
-
-def get_index_name(file_src):
-    file_paths = [x.name for x in file_src]
-    file_paths.sort(key=lambda x: os.path.basename(x))
-
-    md5_hash = hashlib.md5()
-    for file_path in file_paths:
-        with open(file_path, "rb") as f:
-            while chunk := f.read(8192):
-                md5_hash.update(chunk)
-
-    return md5_hash.hexdigest()
-
-
-def block_split(text):
-    blocks = []
-    while len(text) > 0:
-        blocks.append(Document(text[:1000]))
-        text = text[1000:]
-    return blocks
 
 
 def get_documents(file_src):
     documents = []
+    index_name = ""
     logging.debug("Loading documents...")
     logging.debug(f"file_src: {file_src}")
     for file in file_src:
-        filepath = file.name
-        filename = os.path.basename(filepath)
-        file_type = os.path.splitext(filepath)[1]
-        logging.info(f"loading file: {filename}")
-        try:
-            if file_type == ".pdf":
-                logging.debug("Loading PDF...")
-                CJKPDFReader = download_loader("CJKPDFReader")
-                loader = CJKPDFReader()
-                text_raw = loader.load_data(file=filepath)[0].text
-            elif file_type == ".docx":
-                logging.debug("Loading Word...")
-                DocxReader = download_loader("DocxReader")
-                loader = DocxReader()
-                text_raw = loader.load_data(file=filepath)[0].text
-            elif file_type == ".epub":
-                logging.debug("Loading EPUB...")
-                EpubReader = download_loader("EpubReader")
-                loader = EpubReader()
-                text_raw = loader.load_data(file=filepath)[0].text
-            elif file_type == ".xlsx":
-                logging.debug("Loading Excel...")
-                text_list = excel_to_string(filepath)
-                for elem in text_list:
-                    documents.append(Document(elem))
-                continue
-            else:
-                logging.debug("Loading text file...")
-                with open(filepath, "r", encoding="utf-8") as f:
-                    text_raw = f.read()
-        except Exception as e:
-            logging.error(f"Error loading file: {filename}")
-            pass
-        text = add_space(text_raw)
-        # text = block_split(text)
-        # documents += text
-        documents += [Document(text)]
-    logging.debug("Documents loaded.")
-    return documents
+        logging.debug(f"file: {file.name}")
+        index_name += file.name
+        if os.path.splitext(file.name)[1] == ".pdf":
+            logging.debug("Loading PDF...")
+            CJKPDFReader = download_loader("CJKPDFReader")
+            loader = CJKPDFReader()
+            documents += loader.load_data(file=file.name)
+        elif os.path.splitext(file.name)[1] == ".docx":
+            logging.debug("Loading DOCX...")
+            DocxReader = download_loader("DocxReader")
+            loader = DocxReader()
+            documents += loader.load_data(file=file.name)
+        elif os.path.splitext(file.name)[1] == ".epub":
+            logging.debug("Loading EPUB...")
+            EpubReader = download_loader("EpubReader")
+            loader = EpubReader()
+            documents += loader.load_data(file=file.name)
+        else:
+            logging.debug("Loading text file...")
+            with open(file.name, "r", encoding="utf-8") as f:
+                text = add_space(f.read())
+                documents += [Document(text)]
+    index_name = sha1sum(index_name)
+    return documents, index_name
 
 
 def construct_index(
     api_key,
     file_src,
     max_input_size=4096,
-    num_outputs=5,
+    num_outputs=1,
     max_chunk_overlap=20,
     chunk_size_limit=600,
     embedding_limit=None,
     separator=" ",
+    num_children=10,
+    max_keywords_per_chunk=10,
 ):
-    from langchain.chat_models import ChatOpenAI
-    from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-    from llama_index import GPTSimpleVectorIndex, ServiceContext, LangchainEmbedding, OpenAIEmbedding
-
-    if api_key:
-        os.environ["OPENAI_API_KEY"] = api_key
-    else:
-        # 由于一个依赖的愚蠢的设计，这里必须要有一个API KEY
-        os.environ["OPENAI_API_KEY"] = "sk-xxxxxxx"
+    os.environ["OPENAI_API_KEY"] = api_key
     chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
     embedding_limit = None if embedding_limit == 0 else embedding_limit
     separator = " " if separator == "" else separator
 
+    llm_predictor = LLMPredictor(
+        llm=OpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)
+    )
     prompt_helper = PromptHelper(
-        max_input_size=max_input_size,
-        num_output=num_outputs,
-        max_chunk_overlap=max_chunk_overlap,
-        embedding_limit=embedding_limit,
-        chunk_size_limit=600,
+        max_input_size,
+        num_outputs,
+        max_chunk_overlap,
+        embedding_limit,
+        chunk_size_limit,
         separator=separator,
     )
-    index_name = get_index_name(file_src)
+    documents, index_name = get_documents(file_src)
     if os.path.exists(f"./index/{index_name}.json"):
         logging.info("找到了缓存的索引文件，加载中……")
         return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")
     else:
         try:
-            documents = get_documents(file_src)
-            if local_embedding:
-                embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name = "sentence-transformers/distiluse-base-multilingual-cased-v2"))
-            else:
-                embed_model = OpenAIEmbedding()
-            logging.info("构建索引中……")
-            with retrieve_proxy():
-                service_context = ServiceContext.from_defaults(
-                    prompt_helper=prompt_helper,
-                    chunk_size_limit=chunk_size_limit,
-                    embed_model=embed_model,
-                )
-                index = GPTSimpleVectorIndex.from_documents(
-                    documents, service_context=service_context
-                )
-            logging.debug("索引构建完成！")
+            logging.debug("构建索引中……")
+            index = GPTSimpleVectorIndex(
+                documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper
+            )
             os.makedirs("./index", exist_ok=True)
             index.save_to_disk(f"./index/{index_name}.json")
-            logging.debug("索引已保存至本地!")
             return index
-
         except Exception as e:
-            logging.error("索引构建失败！", e)
             print(e)
             return None
 
 
+def chat_ai(
+    api_key,
+    index,
+    question,
+    context,
+    chatbot,
+):
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    logging.info(f"Question: {question}")
+
+    response, chatbot_display, status_text = ask_ai(
+        api_key,
+        index,
+        question,
+        replace_today(PROMPT_TEMPLATE),
+        REFINE_TEMPLATE,
+        SIM_K,
+        INDEX_QUERY_TEMPRATURE,
+        context,
+    )
+    if response is None:
+        status_text = "查询失败，请换个问法试试"
+        return context, chatbot
+    response = response
+
+    context.append({"role": "user", "content": question})
+    context.append({"role": "assistant", "content": response})
+    chatbot.append((question, chatbot_display))
+
+    os.environ["OPENAI_API_KEY"] = ""
+    return context, chatbot, status_text
+
+
+def ask_ai(
+    api_key,
+    index,
+    question,
+    prompt_tmpl,
+    refine_tmpl,
+    sim_k=1,
+    temprature=0,
+    prefix_messages=[],
+):
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    logging.debug("Index file found")
+    logging.debug("Querying index...")
+    llm_predictor = LLMPredictor(
+        llm=OpenAI(
+            temperature=temprature,
+            model_name="gpt-3.5-turbo-0301",
+            prefix_messages=prefix_messages,
+        )
+    )
+
+    response = None  # Initialize response variable to avoid UnboundLocalError
+    qa_prompt = QuestionAnswerPrompt(prompt_tmpl)
+    rf_prompt = RefinePrompt(refine_tmpl)
+    response = index.query(
+        question,
+        llm_predictor=llm_predictor,
+        similarity_top_k=sim_k,
+        text_qa_template=qa_prompt,
+        refine_template=rf_prompt,
+        response_mode="compact",
+    )
+
+    if response is not None:
+        logging.info(f"Response: {response}")
+        ret_text = response.response
+        nodes = []
+        for index, node in enumerate(response.source_nodes):
+            brief = node.source_text[:25].replace("\n", "")
+            nodes.append(
+                f"<details><summary>[{index+1}]\t{brief}...</summary><p>{node.source_text}</p></details>"
+            )
+        new_response = ret_text + "\n----------\n" + "\n\n".join(nodes)
+        logging.info(
+            f"Response: {colorama.Fore.BLUE}{ret_text}{colorama.Style.RESET_ALL}"
+        )
+        os.environ["OPENAI_API_KEY"] = ""
+        return ret_text, new_response, f"查询消耗了{llm_predictor.last_token_usage} tokens"
+    else:
+        logging.warning("No response found, returning None")
+        os.environ["OPENAI_API_KEY"] = ""
+        return None
+
+
 def add_space(text):
     punctuations = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}
     for cn_punc, en_punc in punctuations.items():