queries.rs 3.1 KB

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