增加 多线程

document_ai/run.py

@@ -0,0 +1,20 @@
+from threading import Thread
+
+import worker
+import server
+
+if __name__ == '__main__':
+    # Start the worker thread
+    worker_thread = Thread(target=worker.sync_documents, args=())
+    worker_thread.start()
+
+    # Start the server thread
+    server_thread = Thread(target=server.serve, args=())
+    server_thread.start()
+
+    # Wait for the threads to finish
+    worker_thread.join()
+    server_thread.join()
+
+
+

document_ai/server.py

@@ -52,10 +52,10 @@ class AIServer(document_query_pb2_grpc.DocumentQuery):
 
         print(real_document)
 
-        print("正在调用 LLM...")
+        print("正在调用 LLM: " + question + "...")
         chain = load_qa_with_sources_chain(OpenAI(temperature=0), chain_type="map_reduce",
-                                           return_intermediate_steps=True,
+                                           return_intermediate_steps=False,
-                                           verbose=True)
+                                           verbose=False)
         output = chain({"input_documents": real_document, "question": question}, return_only_outputs=False)
         print("回复：" + output["output_text"])
 
@@ -77,4 +77,3 @@ def serve():
     server.wait_for_termination()
 
 
-serve()

document_ai/worker.py

@@ -1,27 +1,35 @@
+import time
+
 import documents_pb2_grpc
 import documents_pb2
 import init
 import doc_client
+import sys
+import signal
 
-print("获取需要更新的数据...")
-documents_response = doc_client.stub.GetNoVectorDocuments(documents_pb2.GetNoVectorDocumentsRequest()).documents
 
-# # get all documents with no vector
+def sync_documents():
-for document in documents_response:
+    while True:
-    docContent = document.title + "\n" + document.content
+        documents_response = doc_client.stub.GetNoVectorDocuments(documents_pb2.GetNoVectorDocumentsRequest()).documents
 
-    print("正在更新向量...")
+        # # get all documents with no vector
-    text_vector = init.text_to_vector(docContent)
+        for document in documents_response:
+            docContent = document.title + "\n" + document.content
 
-    # update vector
+            print("正在更新向量...")
-    update_vector_response = init.insert_document(document.id, document.user_id, text_vector)
+            text_vector = init.text_to_vector(docContent)
-    print(update_vector_response)
 
-    # update vector_id
+            # update vector
-    update_vector_id_response = doc_client.stub.UpdateDocument(documents_pb2.UpdateDocumentRequest(
+            update_vector_response = init.insert_document(document.id, document.user_id, text_vector)
-        id=document.id,
+            print(update_vector_response)
-        vector_id=update_vector_response
-    ))
 
-    print(update_vector_id_response)
+            # update vector_id
-    print("更新向量完成")
+            update_vector_id_response = doc_client.stub.UpdateDocument(documents_pb2.UpdateDocumentRequest(
+                id=document.id,
+                vector_id=update_vector_response
+            ))
+
+            print(update_vector_id_response)
+            print("更新向量完成")
+
+            time.sleep(1 * 60)