meeting_attendee_recommendation.py 1.9 KB

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