import os
from concurrent import futures
from langchain.embeddings import OpenAIEmbeddings
import ai_pb2
import ai_pb2_grpc
import grpc


class AIServer(ai_pb2_grpc.LLMQueryServicer):
    def AddDocument(self, request, context):
        print("AddDocument called with", request.text)

        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()