81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
import os
|
|
from concurrent import futures
|
|
import document_query_pb2
|
|
import document_query_pb2_grpc
|
|
import grpc
|
|
import documents_pb2
|
|
import init
|
|
import doc_client
|
|
from langchain.llms.openai import OpenAI
|
|
from langchain.schema.document import Document
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
|
|
|
|
|
|
class AIServer(document_query_pb2_grpc.DocumentQuery):
|
|
def Query(self, request, context):
|
|
vec = init.text_to_vector(request.question)
|
|
|
|
question = request.question + "(必须使用中文回复)"
|
|
|
|
search_param = {
|
|
"data": [vec],
|
|
"anns_field": "vector",
|
|
"param": {"metric_type": "L2"},
|
|
"limit": 10,
|
|
"expr": "user_id == " + str(request.user_id),
|
|
"output_fields": ["document_id", "user_id"],
|
|
}
|
|
|
|
res = init.collection.search(**search_param)
|
|
|
|
document_ids = []
|
|
real_document = []
|
|
|
|
for i in range(len(res[0])):
|
|
_doc_id = res[0][i].id
|
|
print("正在获取 " + str(_doc_id) + " 的内容...")
|
|
|
|
try:
|
|
_doc_content = doc_client.stub.GetDocumentById(documents_pb2.GetDocumentByIdRequest(
|
|
id=_doc_id
|
|
))
|
|
_doc_content_full = _doc_content.title + "\n" + _doc_content.content
|
|
|
|
# real_document.append(_doc_content)
|
|
doc_obj = Document(page_content=_doc_content_full, metadata={"source": _doc_content.title})
|
|
|
|
real_document.append(doc_obj)
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
print(real_document)
|
|
|
|
print("正在调用 LLM...")
|
|
chain = load_qa_with_sources_chain(OpenAI(temperature=0), chain_type="map_reduce",
|
|
return_intermediate_steps=True,
|
|
verbose=True)
|
|
output = chain({"input_documents": real_document, "question": question}, return_only_outputs=False)
|
|
print("回复：" + output["output_text"])
|
|
|
|
return document_query_pb2.QueryResponse(
|
|
text=output["output_text"]
|
|
)
|
|
|
|
|
|
def serve():
|
|
_ADDR = os.getenv("BIND")
|
|
if _ADDR is None:
|
|
_ADDR = "[::]:50051"
|
|
print("Listening on", _ADDR)
|
|
|
|
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
|
|
document_query_pb2_grpc.add_DocumentQueryServicer_to_server(AIServer(), server)
|
|
server.add_insecure_port(_ADDR)
|
|
server.start()
|
|
server.wait_for_termination()
|
|
|
|
|
|
serve()
|