fix

.gitignore

@@ -5,3 +5,13 @@ __pycache__/
 files/input/ir2023_ashare.pdf
 knowledge_bases/中国移动.faiss
 knowledge_bases/中国移动.pkl
+files/input/ClientList.csv
+db.json
+files/input/history_Example_20240109T225059.json
+files/input/ir2023_ashare.txt
+knowledge_bases/json.faiss
+knowledge_bases/json.pkl
+knowledge_bases/list.faiss
+knowledge_bases/list.pkl
+knowledge_bases/text.faiss
+knowledge_bases/text.pkl

app.py

@@ -8,6 +8,10 @@ from sam_everything import SamAnything
 from ocr import do_ocr
 from retriever import knowledgeBase
 import time
+from pathlib import Path
+
+current_file_path = Path(__file__).resolve()
+absolute_path = (current_file_path.parent / "files" / "input").resolve()
 
 components = {}
 
@@ -137,6 +141,14 @@ def create_ui():
                                                     col_count=(1, "fixed"),
                                                     interactive=False
                         )
+                        components["file_expr"] = gr.FileExplorer(
+                            scale=1,
+                            value=[],
+                            file_count="single",
+                            root=absolute_path,
+                            # ignore_glob="**/__init__.py",
+                            elem_id="file_expr",
+                        )
                 with gr.Column(scale=2):
                     with gr.Row():
                         with gr.Column(scale=2):
@@ -246,6 +258,10 @@ def create_event_handlers():
         llm, gradio('ak','sk','llm_client'), None
     )
 
+    components['db_view'].select(
+        db_expr, gradio('db_view'), gradio('file_expr')
+    )
+
 def do_refernce(algo_type,input_image):
 # def do_refernce():
     print("input image",input_image)
@@ -422,6 +438,14 @@ def file_handler(file_objs,name):
     dfs = knowledgeBase.get_df_bases()
     return dfs,gr.CheckboxGroup(dbs,label="知识库", info="可选择1个或多个知识库"),gr.Dropdown(dbs,multiselect=True, label="知识库选择")
 
+def db_expr(selected_index: gr.SelectData, dataframe_origin):
+    print("db_expr",selected_index.index)
+    
+    dbname = dataframe_origin.iloc[selected_index.index[0],selected_index.index[1]]
+    print("db_expr",dbname)
+    
+    return knowledgeBase.get_db_files(dbname)
+
 def do_search(selected_dbs,user_input):
     print("do_search:",selected_dbs,user_input)
     context = knowledgeBase.retrieve_documents(selected_dbs,user_input)

retriever.py

@@ -5,6 +5,7 @@ from langchain_community.docstore.in_memory import InMemoryDocstore
 import faiss
 import os
 import glob
+import json
 from typing import Any,List,Dict
 from embedding import Embedding
 
@@ -16,6 +17,7 @@ class KnowledgeBaseManager:
         self.batch_size = batch_size
         self.embeddings = Embedding()
         self.knowledge_bases: Dict[str, FAISS] = {}
+        self.db_files_map: Dict[str, list] = {}
         os.makedirs(self.base_path, exist_ok=True)
 
         faiss_files = glob.glob(os.path.join(base_path, '*.faiss'))
@@ -33,12 +35,14 @@ class KnowledgeBaseManager:
             return
         
         self.knowledge_bases[name] = kb
+        self.db_files_map[name] = []
         self.save_knowledge_base(name)
         print(f"Knowledge base '{name}' created.")
 
     def delete_knowledge_base(self, name: str):
         if name in self.knowledge_bases:
             del self.knowledge_bases[name]
+            del self.db_files_map[name]
             os.remove(os.path.join(self.base_path, f"{name}.faiss"))
             print(f"Knowledge base '{name}' deleted.")
         else:
@@ -48,6 +52,15 @@ class KnowledgeBaseManager:
         kb_path = os.path.join(self.base_path, f"{name}.faiss")
         if os.path.exists(kb_path):
             self.knowledge_bases[name] = FAISS.load_local(self.base_path, self.embeddings, name, allow_dangerous_deserialization=True)
+            # 加载文件中的数据
+            try:
+                with open('db.json', 'r+') as f:
+                    self.db_files_map = json.load(f)
+            except FileNotFoundError:
+                # 如果文件不存在，则创建一个空的文件并初始化 self.db_files_map
+                with open('db.json', 'w+') as f:
+                    self.db_files_map = {}
+                    json.dump(self.db_files_map, f)
             print(f"Knowledge base '{name}' loaded.")
         else:
             print(f"Knowledge base '{name}' does not exist.")
@@ -55,6 +68,8 @@ class KnowledgeBaseManager:
     def save_knowledge_base(self, name: str):
         if name in self.knowledge_bases:
             self.knowledge_bases[name].save_local(self.base_path, name)
+            with open('db.json', 'w') as f:
+                json.dump(self.db_files_map, f)
             print(f"Knowledge base '{name}' saved.")
         else:
             print(f"Knowledge base '{name}' does not exist.")
@@ -72,6 +87,7 @@ class KnowledgeBaseManager:
             self.create_knowledge_base(name)
         
         kb = self.knowledge_bases[name]
+        self.db_files_map[name].extend([os.path.basename(file_path) for file_path in file_paths])
         documents = self.load_documents(file_paths)
         print(f"Loaded {len(documents)} documents.")
         print(documents)
@@ -138,6 +154,10 @@ class KnowledgeBaseManager:
         
         return results
     
+    def get_db_files(self,name):
+        data = self.db_files_map[name]
+        return data
+    
     def get_bases(self):
         data = self.knowledge_bases.keys()
         return list(data)