douyin/grpc_m/send_data_to_vector.py

import asyncio
from grpclib.client import Channel
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from config import logger,GRPC_VECTOR_HOST,GRPC_VECTOR_PORT
from grpc_m.vector import vector_pb2, vector_grpc
from db.docs import DocumentsRepository,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 Channel(GRPC_VECTOR_HOST, GRPC_VECTOR_PORT) as channel:
            stub = vector_grpc.VectorServiceStub(channel)  
            logger.info(f"send to vector: category_id={category_id} user_doc_relative_path={user_doc_relative_path}")

            request = vector_pb2.SaveDocToVectorRequest(  
                user_doc_relative_path=user_doc_relative_path,
                category_id=str(category_id) # 将 uuid 类型转为字符串，使其符合 proto 文件定义的协议
            )
            response:vector_pb2.SaveDocToVectorResponse = await stub.SaveDocToVector(request)  
            # send_to_grpc_vetcor(category_id, doc_model)
            # 输出响应结果  
            if response.status == vector_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_HOST, GRPC_VECTOR_PORT} error:{response.status}")

async def langchain_chat(collection_name:str, prompt:str):
    async with Channel(GRPC_VECTOR_HOST, GRPC_VECTOR_PORT) as channel:
        vector = vector_grpc.VectorServiceStub(channel)  
        request = vector_pb2.DocChatRequest(collection_name=collection_name, prompt=prompt)
        res = await vector.DocChat(request)
        logger.debug(f"{collection_name, res.reply, prompt}")
        return res.reply
    
async def simarity_search():
    pass
    
async def main():
    dm = Documents(open_id="user123", path="思维链-文档说明.md")
    # send_to_grpc_vetcor("123", dm)
    await langchain_chat("123")
if __name__ == "__main__":
    asyncio.run(main())