meeting_attendee_recommendation.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. import time
  2. from loguru import logger
  3. from app.services.milvus import Milvus
  4. from app.services.milvus_helpers import my_milvus
  5. from app.services.my_redis import MyRedis
  6. @logger.catch()
  7. async def build_recommendations(company_id: str, user_id: str) -> list[str]:
  8. MILVUS_CLI = my_milvus.get("cli")
  9. myredis = MyRedis()
  10. r = myredis.get_client()
  11. user_pk = r.hget(f"meeting_user_id_to_embedding_id_map:{company_id}", user_id)
  12. user_id_list = []
  13. if user_pk:
  14. user_emb = await MILVUS_CLI.query_vector_by_pk(
  15. f"meeting_attendee_rec_{company_id}",
  16. int(user_pk)
  17. )
  18. pk_list = await MILVUS_CLI.search_vectors(
  19. f"meeting_attendee_rec_{company_id}",
  20. [user_emb],
  21. 10
  22. )
  23. pk_list = [str(item) for item in pk_list]
  24. user_id_list = r.hmget(
  25. f"embedding_id_to_meeting_user_id_map:{company_id}", pk_list
  26. )
  27. return user_id_list
  28. @logger.catch()
  29. async def build_recommendations_v2(company_id: str, user_id: str) -> list[str]:
  30. old_milvus = Milvus()
  31. myredis = MyRedis()
  32. r = myredis.get_client()
  33. user_pk = r.hget(f"meeting_user_id_to_embedding_id_map:{company_id}", user_id)
  34. user_id_list = []
  35. if user_pk:
  36. start = time.perf_counter()
  37. user_emb = await old_milvus.query_embedding_by_pk(
  38. f"meeting_attendee_rec_{company_id}",
  39. "pk",
  40. int(user_pk)
  41. )
  42. logger.debug(time.perf_counter() - start)
  43. start = time.perf_counter()
  44. pk_list = await old_milvus.search(
  45. [user_emb],
  46. f"meeting_attendee_rec_{company_id}",
  47. "embeddings",
  48. 10
  49. )
  50. logger.debug(time.perf_counter() - start)
  51. pk_list = [str(item) for item in pk_list]
  52. user_id_list = r.hmget(
  53. f"embedding_id_to_meeting_user_id_map:{company_id}", pk_list
  54. )
  55. return user_id_list