This commit is contained in:
Wenjie Zhang 2024-08-25 12:34:35 +08:00
parent 4083fa6938
commit def68e8a3e
4 changed files with 39 additions and 15 deletions

View File

@ -77,6 +77,12 @@ class DataBaseManager:
"graph": data["graph"] "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): def _save_databases(self):
"""将数据库的信息保存到本地的文件里面""" """将数据库的信息保存到本地的文件里面"""
self._update_database() self._update_database()
@ -129,29 +135,36 @@ class DataBaseManager:
# filenames = [f["filename"] for f in db.files] # filenames = [f["filename"] for f in db.files]
# if os.path.basename(file) in filenames: # if os.path.basename(file) in filenames:
# continue # continue
db.files.append({ new_file = {
"file_id": "file_" + hashstr(file + str(time.time())), "file_id": "file_" + hashstr(file + str(time.time())),
"filename": os.path.basename(file), "filename": os.path.basename(file),
"path": file, "path": file,
"type": file.split(".")[-1], "type": file.split(".")[-1],
"status": "waiting", "status": "waiting",
"created_at": time.time() "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" db.files[idx]["status"] = "processing"
try: try:
text = self.read_text(file) text = self.read_text(new_file["path"])
chunks = self.chunking(text) chunks = self.chunking(text)
self.knowledge_base.add_documents( self.knowledge_base.add_documents(
docs=chunks, docs=chunks,
collection_name=db.metaname, 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" db.files[idx]["status"] = "done"
except Exception as e: except Exception as e:
logger.error(f"Failed to add documents to collection {db.metaname}, {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" db.files[idx]["status"] = "failed"
self._save_databases() self._save_databases()

View File

@ -8,7 +8,7 @@ logger = setup_logger("EmbeddingModel")
SUPPORT_LIST = { SUPPORT_LIST = {
"bge-large-zh-v1.5": "BAAI/bge-large-zh-v1.5", "bge-large-zh-v1.5": "BAAI/bge-large-zh-v1.5",
"zhipu": "embedding-2", "zhipu": "embedding-3",
} }
RERANKER_LIST = { RERANKER_LIST = {
@ -57,11 +57,19 @@ class ZhipuEmbedding:
self.query_instruction_for_retrieval = "为这个句子生成表示以用于检索相关文章:" self.query_instruction_for_retrieval = "为这个句子生成表示以用于检索相关文章:"
def predict(self, message): def predict(self, message):
response = self.client.embeddings.create(
model=SUPPORT_LIST[self.config.embed_model], data = []
input=message
) for i in range(0, len(message), 10):
return [a["embedding"] for a in response["data"]] 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): def encode(self, message):
return self.predict(message) return self.predict(message)

View File

@ -103,7 +103,7 @@ def upload_file():
if file: if file:
upload_dir = os.path.join(startup.config.save_dir, "data/uploads") upload_dir = os.path.join(startup.config.save_dir, "data/uploads")
os.makedirs(upload_dir, exist_ok=True) 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_path = os.path.join(upload_dir, filename)
file.save(file_path) file.save(file_path)
return jsonify({'message': 'File successfully uploaded', 'file_path': file_path}), 200 return jsonify({'message': 'File successfully uploaded', 'file_path': file_path}), 200

View File

@ -330,7 +330,8 @@ const appendAiMessage = (message, refs=null) => {
role: 'received', role: 'received',
text: message, text: message,
refs, refs,
status: "querying" status: "querying",
model_name: configStore.config.model_name
}) })
scrollToBottom() scrollToBottom()
} }
@ -641,11 +642,13 @@ watch(
.message-box.received { .message-box.received {
color: initial; color: initial;
width: fit-content; width: fit-content;
padding-top: 16px;
text-align: left; text-align: left;
word-wrap: break-word; word-wrap: break-word;
margin: 0; margin: 0;
padding-bottom: 0; padding-bottom: 0;
padding-top: 16px;
padding-left: 0;
padding-right: 0;
text-align: justify; text-align: justify;
} }