From 3ff885d0d38d6645f55e060504d45dce0c7aca26 Mon Sep 17 00:00:00 2001
From: imClumsyPanda <littlepanda0716@gmail.com>
Date: Sat, 13 May 2023 21:35:17 +0800
Subject: [PATCH] update api.py

---
 api.py | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/api.py b/api.py
index 00d8600..c675543 100644
--- a/api.py
+++ b/api.py
@@ -181,7 +181,7 @@ async def delete_docs(
         if os.path.exists(doc_path):
             os.remove(doc_path)
         else:
-            return {"code": 1, "msg": f"document {doc_name} not found"}
+            BaseResponse(code=1, msg=f"document {doc_name} not found")
 
         remain_docs = await list_docs(knowledge_base_id)
         if remain_docs["code"] != 0 or len(remain_docs["data"]) == 0:
@@ -211,24 +211,30 @@ async def local_doc_chat(
 ):
     vs_path = os.path.join(VS_ROOT_PATH, knowledge_base_id)
     if not os.path.exists(vs_path):
-        raise ValueError(f"Knowledge base {knowledge_base_id} not found")
+        # return BaseResponse(code=1, msg=f"Knowledge base {knowledge_base_id} not found")
+        return ChatMessage(
+            question=question,
+            response=f"Knowledge base {knowledge_base_id} not found",
+            history=history,
+            source_documents=[],
+        )
+    else:
+        for resp, history in local_doc_qa.get_knowledge_based_answer(
+                query=question, vs_path=vs_path, chat_history=history, streaming=True
+        ):
+            pass
+        source_documents = [
+            f"""出处 [{inum + 1}] {os.path.split(doc.metadata['source'])[-1]}：\n\n{doc.page_content}\n\n"""
+            f"""相关度：{doc.metadata['score']}\n\n"""
+            for inum, doc in enumerate(resp["source_documents"])
+        ]
 
-    for resp, history in local_doc_qa.get_knowledge_based_answer(
-            query=question, vs_path=vs_path, chat_history=history, streaming=True
-    ):
-        pass
-    source_documents = [
-        f"""出处 [{inum + 1}] {os.path.split(doc.metadata['source'])[-1]}：\n\n{doc.page_content}\n\n"""
-        f"""相关度：{doc.metadata['score']}\n\n"""
-        for inum, doc in enumerate(resp["source_documents"])
-    ]
-
-    return ChatMessage(
-        question=question,
-        response=resp["result"],
-        history=history,
-        source_documents=source_documents,
-    )
+        return ChatMessage(
+            question=question,
+            response=resp["result"],
+            history=history,
+            source_documents=source_documents,
+        )
 
 
 async def chat(