queries.rs 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. use heed3::RoTxn;
  2. use get_routes::handler;
  3. use helix_db::{field_remapping, identifier_remapping, traversal_remapping, exclude_field, value_remapping, embed};
  4. use helix_db::helix_engine::vector_core::vector::HVector;
  5. use helix_db::{
  6. helix_engine::graph_core::ops::{
  7. g::G,
  8. in_::{in_::InAdapter, in_e::InEdgesAdapter, to_n::ToNAdapter, to_v::ToVAdapter},
  9. out::{from_n::FromNAdapter, from_v::FromVAdapter, out::OutAdapter, out_e::OutEdgesAdapter},
  10. source::{
  11. add_e::{AddEAdapter, EdgeType},
  12. add_n::AddNAdapter,
  13. e_from_id::EFromIdAdapter,
  14. e_from_type::EFromTypeAdapter,
  15. n_from_id::NFromIdAdapter,
  16. n_from_type::NFromTypeAdapter,
  17. n_from_index::NFromIndexAdapter,
  18. },
  19. tr_val::{Traversable, TraversalVal},
  20. util::{
  21. dedup::DedupAdapter, filter_mut::FilterMut,
  22. filter_ref::FilterRefAdapter, range::RangeAdapter, update::UpdateAdapter,
  23. map::MapAdapter, paths::ShortestPathAdapter, props::PropsAdapter, drop::Drop,
  24. },
  25. vectors::{insert::InsertVAdapter, search::SearchVAdapter, brute_force_search::BruteForceSearchVAdapter},
  26. bm25::search_bm25::SearchBM25Adapter,
  27. },
  28. helix_engine::types::GraphError,
  29. helix_gateway::router::router::HandlerInput,
  30. node_matches, props,
  31. protocol::count::Count,
  32. protocol::remapping::{RemappingMap, ResponseRemapping},
  33. protocol::response::Response,
  34. protocol::{
  35. filterable::Filterable, remapping::Remapping, return_values::ReturnValue, value::Value, id::ID,
  36. },
  37. providers::embedding_providers::get_embedding_model,
  38. };
  39. use sonic_rs::{Deserialize, Serialize};
  40. use std::collections::{HashMap, HashSet};
  41. use std::sync::Arc;
  42. use std::time::Instant;
  43. use std::cell::RefCell;
  44. use chrono::{DateTime, Utc};
  45. pub struct Professor {
  46. pub name: String,
  47. pub title: String,
  48. pub page: String,
  49. pub bio: String,
  50. }
  51. pub struct HasResearchAreaAndDescriptionEmbedding {
  52. pub from: Professor,
  53. pub to: ResearchAreaAndDescriptionEmbedding,
  54. pub areas_and_descriptions: String,
  55. }
  56. pub struct ResearchAreaAndDescriptionEmbedding {
  57. pub areas_and_descriptions: String,
  58. }
  59. #[derive(Serialize, Deserialize)]
  60. pub struct get_professor_research_areas_with_descriptions_v2Input {
  61. pub professor_id: ID
  62. }
  63. #[handler]
  64. pub fn get_professor_research_areas_with_descriptions_v2 (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  65. let data: get_professor_research_areas_with_descriptions_v2Input = match sonic_rs::from_slice(&input.request.body) {
  66. Ok(data) => data,
  67. Err(err) => return Err(GraphError::from(err)),
  68. };
  69. let mut remapping_vals = RemappingMap::new();
  70. let db = Arc::clone(&input.graph.storage);
  71. let txn = db.graph_env.read_txn().unwrap();
  72. let research_areas = G::new(Arc::clone(&db), &txn)
  73. .n_from_id(&data.professor_id)
  74. .out("HasResearchAreaAndDescriptionEmbedding",&EdgeType::Vec)
  75. .check_property("areas_and_descriptions").collect_to_obj();
  76. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  77. return_vals.insert("research_areas".to_string(), ReturnValue::from_traversal_value_array_with_mixin(G::new_from(Arc::clone(&db), &txn, research_areas.clone())
  78. .map_traversal(|item, txn| { traversal_remapping!(remapping_vals, item.clone(), "areas_and_descriptions" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  79. .check_property("areas_and_descriptions").collect_to_obj())?;
  80. Ok(item) }).collect_to::<Vec<_>>().clone(), remapping_vals.borrow_mut()));
  81. txn.commit().unwrap();
  82. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  83. Ok(())
  84. }
  85. #[derive(Serialize, Deserialize)]
  86. pub struct get_professor_research_areas_with_descriptions_v1Input {
  87. pub professor_id: ID
  88. }
  89. #[handler]
  90. pub fn get_professor_research_areas_with_descriptions_v1 (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  91. let data: get_professor_research_areas_with_descriptions_v1Input = match sonic_rs::from_slice(&input.request.body) {
  92. Ok(data) => data,
  93. Err(err) => return Err(GraphError::from(err)),
  94. };
  95. let mut remapping_vals = RemappingMap::new();
  96. let db = Arc::clone(&input.graph.storage);
  97. let txn = db.graph_env.read_txn().unwrap();
  98. let research_areas = G::new(Arc::clone(&db), &txn)
  99. .n_from_id(&data.professor_id)
  100. .out("HasResearchAreaAndDescriptionEmbedding",&EdgeType::Vec).collect_to::<Vec<_>>();
  101. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  102. return_vals.insert("research_areas".to_string(), ReturnValue::from_traversal_value_array_with_mixin(G::new_from(Arc::clone(&db), &txn, research_areas.clone())
  103. .check_property("areas_and_descriptions").collect_to_obj().clone(), remapping_vals.borrow_mut()));
  104. txn.commit().unwrap();
  105. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  106. Ok(())
  107. }