queries.rs 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. use heed3::RoTxn;
  2. use get_routes::handler;
  3. use helix_db::{field_remapping, identifier_remapping, traversal_remapping, exclude_field};
  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},
  9. out::{from_n::FromNAdapter, 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::ResponseRemapping,
  33. protocol::response::Response,
  34. protocol::traversal_value::TraversalValue,
  35. protocol::{
  36. filterable::Filterable, remapping::Remapping, return_values::ReturnValue, value::Value, id::ID,
  37. },
  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 User {
  46. pub username: String,
  47. pub email: String,
  48. }
  49. pub struct Connection {
  50. pub from: User,
  51. pub to: Embedding,
  52. }
  53. pub struct Friend {
  54. pub from: User,
  55. pub to: User,
  56. }
  57. pub struct Embedding {
  58. pub created_at: DateTime<Utc>,
  59. }
  60. #[derive(Serialize, Deserialize)]
  61. pub struct search_vectorInput {
  62. pub query: Vec<f64>,
  63. pub k: i64
  64. }
  65. #[handler]
  66. pub fn search_vector (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  67. let data: search_vectorInput = match sonic_rs::from_slice(&input.request.body) {
  68. Ok(data) => data,
  69. Err(err) => return Err(GraphError::from(err)),
  70. };
  71. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  72. let db = Arc::clone(&input.graph.storage);
  73. let txn = db.graph_env.read_txn().unwrap();
  74. let result = G::new(Arc::clone(&db), &txn)
  75. .n_from_type("User")
  76. .out("Connection",&EdgeType::Vec)
  77. .brute_force_search_v(&data.query, data.k as usize).collect_to::<Vec<_>>();
  78. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  79. return_vals.insert("result".to_string(), ReturnValue::from_traversal_value_array_with_mixin(result.clone(), remapping_vals.borrow_mut()));
  80. txn.commit().unwrap();
  81. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  82. Ok(())
  83. }