diff --git a/src/core/database.py b/src/core/database.py index c8dd43d8..ccdc834a 100644 --- a/src/core/database.py +++ b/src/core/database.py @@ -77,6 +77,12 @@ class DataBaseManager: "graph": data["graph"] } + # 检查所有文件,如果出现状态是 processing 的,那么设置为 failed + for db in self.data["databases"]: + for file in db.files: + if file["status"] == "processing" or file["status"] == "waiting": + file["status"] = "failed" + def _save_databases(self): """将数据库的信息保存到本地的文件里面""" self._update_database() @@ -129,29 +135,36 @@ class DataBaseManager: # filenames = [f["filename"] for f in db.files] # if os.path.basename(file) in filenames: # continue - db.files.append({ + new_file = { "file_id": "file_" + hashstr(file + str(time.time())), "filename": os.path.basename(file), "path": file, "type": file.split(".")[-1], "status": "waiting", "created_at": time.time() - }) - new_files.append((len(db.files) - 1, file)) + } + db.files.append(new_file) + new_files.append(new_file) - for idx, file in new_files: + for new_file in new_files: + file_id = new_file["file_id"] + idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0] db.files[idx]["status"] = "processing" try: - text = self.read_text(file) + text = self.read_text(new_file["path"]) chunks = self.chunking(text) self.knowledge_base.add_documents( docs=chunks, collection_name=db.metaname, - file_id=db.files[idx]["file_id"]) + file_id=file_id) + + idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0] db.files[idx]["status"] = "done" + except Exception as e: logger.error(f"Failed to add documents to collection {db.metaname}, {e}") + idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0] db.files[idx]["status"] = "failed" self._save_databases() diff --git a/src/models/embedding.py b/src/models/embedding.py index 64c6e173..c6246724 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -8,7 +8,7 @@ logger = setup_logger("EmbeddingModel") SUPPORT_LIST = { "bge-large-zh-v1.5": "BAAI/bge-large-zh-v1.5", - "zhipu": "embedding-2", + "zhipu": "embedding-3", } RERANKER_LIST = { @@ -57,11 +57,19 @@ class ZhipuEmbedding: self.query_instruction_for_retrieval = "为这个句子生成表示以用于检索相关文章:" def predict(self, message): - response = self.client.embeddings.create( - model=SUPPORT_LIST[self.config.embed_model], - input=message - ) - return [a["embedding"] for a in response["data"]] + + data = [] + + for i in range(0, len(message), 10): + group_msg = message[i:i+10] + response = self.client.embeddings.create( + model=SUPPORT_LIST[self.config.embed_model], + input=group_msg + ) + + data.extend([a.embedding for a in response.data]) + + return data def encode(self, message): return self.predict(message) diff --git a/src/views/database_view.py b/src/views/database_view.py index 8516cdb9..29ec83f2 100644 --- a/src/views/database_view.py +++ b/src/views/database_view.py @@ -103,7 +103,7 @@ def upload_file(): if file: upload_dir = os.path.join(startup.config.save_dir, "data/uploads") os.makedirs(upload_dir, exist_ok=True) - filename = f"{hashstr(file.filename, 6, with_salt=True)}_{file.filename}" + filename = f"{hashstr(file.filename, 4, with_salt=True)}_{file.filename}" file_path = os.path.join(upload_dir, filename) file.save(file_path) return jsonify({'message': 'File successfully uploaded', 'file_path': file_path}), 200 diff --git a/web/src/components/ChatComponent.vue b/web/src/components/ChatComponent.vue index 34ed5156..f2cf4348 100644 --- a/web/src/components/ChatComponent.vue +++ b/web/src/components/ChatComponent.vue @@ -330,7 +330,8 @@ const appendAiMessage = (message, refs=null) => { role: 'received', text: message, refs, - status: "querying" + status: "querying", + model_name: configStore.config.model_name }) scrollToBottom() } @@ -641,11 +642,13 @@ watch( .message-box.received { color: initial; width: fit-content; - padding-top: 16px; text-align: left; word-wrap: break-word; margin: 0; padding-bottom: 0; + padding-top: 16px; + padding-left: 0; + padding-right: 0; text-align: justify; }