open-douyin/grpc_m/send_data_to_vector.py

import asyncio
import grpc
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from config import logger,GRPC_VECTOR
from grpc_m import vector_service_pb2, vector_service_pb2_grpc  
from db.docs import DocumentsRepository,DocumentCategories,Documents,DocStatus

def send_to_grpc_vetcor(category_id,doc_model:Documents):
    loop = asyncio.get_running_loop()
    loop.create_task(load_user_docs_async(category_id,doc_model))
    
async def load_user_docs_async(category_id: str, doc_model:Documents):  
    user_doc_relative_path = DocumentsRepository.get_user_file_relpath_from_docmodel(doc_model)
    try:
        async with grpc.aio.insecure_channel(GRPC_VECTOR) as channel:  
            stub = vector_service_pb2_grpc.VectorServiceStub(channel)  
            logger.info(f"send to vector: category_id={category_id} user_doc_relative_path={user_doc_relative_path}")
            request = vector_service_pb2.SaveDocToVectorRequest(  
                user_doc_relative_path=user_doc_relative_path,
                category_id=str(category_id) # 将 uuid 类型转为字符串，使其符合 proto 文件定义的协议
            )
            response:vector_service_pb2.SaveDocToVectorResponse = await stub.SaveDocToVector(request)  
    
            # 输出响应结果  
            if response.status == vector_service_pb2.ErrorCode.SUCCESS:
                doc_model.status = DocStatus.COMPLETED
                DocumentsRepository().update(doc_model)
                logger.info(f"Document conver to vector sucess. category_id:{category_id}  doc_id:{doc_model.id} path:{doc_model.path}")
            else:
                logger.error(f"Response: status = {response.status}, {doc_model}")
            return response.status
    except Exception:
        logger.exception(f"vector server {GRPC_VECTOR} error:{response.status}")

async def main():
    dm = Documents(open_id="user123", path="思维链-文档说明.md")
    send_to_grpc_vetcor("123", dm)
    while True:
        await asyncio.sleep(1)
if __name__ == "__main__":
    asyncio.run(main())