raphtory-graphql 0.17.0

Raphtory GraphQL server
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

use crate::rayon::blocking_compute;

use super::vector_selection::GqlVectorSelection;
use dynamic_graphql::{InputObject, ResolvedObject, ResolvedObjectFields};
use raphtory::{
    db::api::view::MaterializedGraph, errors::GraphResult,
    vectors::vectorised_graph::VectorisedGraph,
};

#[derive(InputObject)]
pub(super) struct VectorisedGraphWindow {
    /// Start time.
    start: i64,
    /// End time.
    end: i64,
}

pub(super) trait IntoWindowTuple {
    fn into_window_tuple(self) -> Option<(i64, i64)>;
}

impl IntoWindowTuple for Option<VectorisedGraphWindow> {
    fn into_window_tuple(self) -> Option<(i64, i64)> {
        self.map(|window| (window.start, window.end))
    }
}

#[derive(ResolvedObject)]
#[graphql(name = "VectorisedGraph")]
pub(crate) struct GqlVectorisedGraph(VectorisedGraph<MaterializedGraph>);

impl From<VectorisedGraph<MaterializedGraph>> for GqlVectorisedGraph {
    fn from(value: VectorisedGraph<MaterializedGraph>) -> Self {
        Self(value.clone())
    }
}

#[ResolvedObjectFields]
impl GqlVectorisedGraph {
    /// Optmize the vector index
    async fn optimize_index(&self) -> GraphResult<bool> {
        self.0.optimize_index().await?;
        Ok(true)
    }

    /// Returns an empty selection of documents.
    async fn empty_selection(&self) -> GqlVectorSelection {
        self.0.empty_selection().into()
    }

    /// Search the top scoring entities according to a specified query returning no more than a specified limit of entities.
    async fn entities_by_similarity(
        &self,
        query: String,
        limit: usize,
        window: Option<VectorisedGraphWindow>,
    ) -> GraphResult<GqlVectorSelection> {
        let vector = self.0.embed_text(query).await?;
        let w = window.into_window_tuple();
        let cloned = self.0.clone();
        let query =
            blocking_compute(move || cloned.entities_by_similarity(&vector, limit, w)).await;
        Ok(query.execute().await?.into())
    }

    /// Search the top scoring nodes according to a specified query returning no more than a specified limit of nodes.
    async fn nodes_by_similarity(
        &self,
        query: String,
        limit: usize,
        window: Option<VectorisedGraphWindow>,
    ) -> GraphResult<GqlVectorSelection> {
        let vector = self.0.embed_text(query).await?;
        let w = window.into_window_tuple();
        let cloned = self.0.clone();
        let query = blocking_compute(move || cloned.nodes_by_similarity(&vector, limit, w)).await;
        Ok(query.execute().await?.into())
    }

    /// Search the top scoring edges according to a specified query returning no more than a specified limit of edges.
    async fn edges_by_similarity(
        &self,
        query: String,
        limit: usize,
        window: Option<VectorisedGraphWindow>,
    ) -> GraphResult<GqlVectorSelection> {
        let vector = self.0.embed_text(query).await?;
        let w = window.into_window_tuple();
        let cloned = self.0.clone();
        let query = blocking_compute(move || cloned.edges_by_similarity(&vector, limit, w)).await;
        Ok(query.execute().await?.into())
    }
}