39 lines
897 B
Python
39 lines
897 B
Python
|
import os
|
||
|
|
from concurrent import futures
|
||
|
|
import ai_pb2
|
||
|
|
import ai_pb2_grpc
|
||
|
|
import grpc
|
||
|
|
import init
|
||
|
|
|
||
|
|
|
||
|
|
class AIServer(ai_pb2_grpc.LLMQueryServicer):
|
||
|
|
def QueryDocumentRequest(self, request, context):
|
||
|
|
_text = request.text
|
||
|
|
_user_id = request.user_id
|
||
|
|
_database = request.database
|
||
|
|
_collection = request.collection
|
||
|
|
|
||
|
|
vector = init.text_to_vector(_text)
|
||
|
|
data = init.insert_document(_user_id, vector, _database)
|
||
|
|
|
||
|
|
|
||
|
|
return ai_pb2.AddDocumentReply(
|
||
|
|
id=request.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))
|
||
|
|
ai_pb2_grpc.add_LLMQueryServicer_to_server(AIServer(), server)
|
||
|
|
server.add_insecure_port(_ADDR)
|
||
|
|
server.start()
|
||
|
|
server.wait_for_termination()
|
||
|
|
|
||
|
|
|
||
|
|
serve()
|