Skip to main content

hawdb_core/graph_rag/query_generation/
model.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::{PropertyType, Value};
16use std::collections::BTreeMap;
17use std::fmt::{Display, Formatter};
18
19pub const MAX_GRAPH_RAG_QUERY_LIMIT: usize = 100;
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub enum GraphRagQueryPattern {
23    Node {
24        label: String,
25    },
26    Route {
27        source_label: String,
28        relationship_type: String,
29        target_label: String,
30    },
31    TwoHopRoute {
32        source_label: String,
33        first_relationship_type: String,
34        intermediate_label: String,
35        second_relationship_type: String,
36        target_label: String,
37    },
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
41pub enum GraphRagQueryBinding {
42    Source,
43    Relationship,
44    Intermediate,
45    SecondRelationship,
46    Target,
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum GraphRagQueryPredicateOperator {
51    Eq,
52    NotEq,
53    Lt,
54    Lte,
55    Gt,
56    Gte,
57    In,
58    Contains,
59    StartsWith,
60    EndsWith,
61    IsNull,
62    IsNotNull,
63}
64
65impl GraphRagQueryPredicateOperator {
66    pub(super) const fn token(self) -> &'static str {
67        match self {
68            Self::Eq => "=",
69            Self::NotEq => "<>",
70            Self::Lt => "<",
71            Self::Lte => "<=",
72            Self::Gt => ">",
73            Self::Gte => ">=",
74            Self::In => "IN",
75            Self::Contains => "CONTAINS",
76            Self::StartsWith => "STARTS WITH",
77            Self::EndsWith => "ENDS WITH",
78            Self::IsNull => "IS NULL",
79            Self::IsNotNull => "IS NOT NULL",
80        }
81    }
82
83    pub(super) const fn requires_parameter(self) -> bool {
84        !matches!(self, Self::IsNull | Self::IsNotNull)
85    }
86
87    pub(super) const fn requires_string_property(self) -> bool {
88        matches!(self, Self::Contains | Self::StartsWith | Self::EndsWith)
89    }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct GraphRagQueryPredicate {
94    pub binding: GraphRagQueryBinding,
95    pub property: String,
96    pub operator: GraphRagQueryPredicateOperator,
97    pub parameter: Option<String>,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct GraphRagQueryProjection {
102    pub binding: GraphRagQueryBinding,
103    pub property: String,
104    pub alias: String,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct GraphRagQueryDraft {
109    pub schema_fingerprint: u64,
110    pub pattern: GraphRagQueryPattern,
111    pub predicates: Vec<GraphRagQueryPredicate>,
112    pub projections: Vec<GraphRagQueryProjection>,
113    pub limit: usize,
114}
115
116#[derive(Debug, Clone, PartialEq, Eq)]
117pub struct GraphRagGeneratedQuery {
118    pub(super) cypher: String,
119    pub(super) schema_fingerprint: u64,
120    pub(super) context_commit_epoch: u64,
121    pub(super) required_parameters: Vec<String>,
122    pub(super) parameter_requirements: Vec<GraphRagQueryParameterRequirement>,
123}
124
125impl GraphRagGeneratedQuery {
126    pub fn cypher(&self) -> &str {
127        &self.cypher
128    }
129
130    pub const fn schema_fingerprint(&self) -> u64 {
131        self.schema_fingerprint
132    }
133
134    pub const fn context_commit_epoch(&self) -> u64 {
135        self.context_commit_epoch
136    }
137
138    pub fn required_parameters(&self) -> &[String] {
139        &self.required_parameters
140    }
141
142    pub fn parameter_requirements(&self) -> &[GraphRagQueryParameterRequirement] {
143        &self.parameter_requirements
144    }
145
146    pub fn validate_parameters(
147        &self,
148        parameters: &BTreeMap<String, Value>,
149    ) -> Result<(), GraphRagQueryParameterError> {
150        for requirement in &self.parameter_requirements {
151            let value = parameters.get(&requirement.name).ok_or_else(|| {
152                GraphRagQueryParameterError::Missing {
153                    parameter: requirement.name.clone(),
154                }
155            })?;
156            if !requirement.accepts(value) {
157                return Err(GraphRagQueryParameterError::TypeMismatch {
158                    parameter: requirement.name.clone(),
159                    expected: requirement.clone(),
160                    actual: value_type_name(value),
161                });
162            }
163        }
164        if let Some(parameter) = parameters
165            .keys()
166            .find(|parameter| !self.required_parameters.contains(parameter))
167        {
168            return Err(GraphRagQueryParameterError::Unexpected {
169                parameter: parameter.clone(),
170            });
171        }
172        Ok(())
173    }
174}
175
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177pub enum GraphRagQueryParameterCardinality {
178    Scalar,
179    List,
180}
181
182#[derive(Debug, Clone, PartialEq, Eq)]
183pub struct GraphRagQueryParameterRequirement {
184    pub name: String,
185    pub value_type: PropertyType,
186    pub cardinality: GraphRagQueryParameterCardinality,
187}
188
189impl GraphRagQueryParameterRequirement {
190    fn accepts(&self, value: &Value) -> bool {
191        match self.cardinality {
192            GraphRagQueryParameterCardinality::Scalar => {
193                value_matches_property_type(value, self.value_type)
194            }
195            GraphRagQueryParameterCardinality::List => match value {
196                Value::List(values) => values
197                    .iter()
198                    .all(|value| value_matches_property_type(value, self.value_type)),
199                _ => false,
200            },
201        }
202    }
203}
204
205#[derive(Debug, Clone, PartialEq, Eq)]
206pub enum GraphRagQueryParameterError {
207    Missing {
208        parameter: String,
209    },
210    Unexpected {
211        parameter: String,
212    },
213    TypeMismatch {
214        parameter: String,
215        expected: GraphRagQueryParameterRequirement,
216        actual: &'static str,
217    },
218}
219
220impl Display for GraphRagQueryParameterError {
221    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
222        match self {
223            Self::Missing { parameter } => {
224                write!(formatter, "missing GraphRAG query parameter: {parameter}")
225            }
226            Self::Unexpected { parameter } => {
227                write!(
228                    formatter,
229                    "unexpected GraphRAG query parameter: {parameter}"
230                )
231            }
232            Self::TypeMismatch {
233                parameter,
234                expected,
235                actual,
236            } => write!(
237                formatter,
238                "GraphRAG query parameter {parameter} must be {}, got {actual}",
239                parameter_requirement_name(expected)
240            ),
241        }
242    }
243}
244
245impl std::error::Error for GraphRagQueryParameterError {}
246
247fn value_matches_property_type(value: &Value, value_type: PropertyType) -> bool {
248    matches!(value, Value::Null)
249        || match value_type {
250            PropertyType::Any => true,
251            PropertyType::Bool => matches!(value, Value::Bool(_)),
252            PropertyType::Int => matches!(value, Value::Int(_)),
253            PropertyType::Float => matches!(value, Value::Int(_) | Value::Float(_)),
254            PropertyType::String | PropertyType::Text => matches!(value, Value::String(_)),
255            PropertyType::List => matches!(value, Value::List(_)),
256            PropertyType::Vector(dimension) => {
257                matches!(value, Value::Vector(value) if value.dimension() == dimension)
258            }
259        }
260}
261
262fn value_type_name(value: &Value) -> &'static str {
263    match value {
264        Value::Null => "null",
265        Value::Bool(_) => "bool",
266        Value::Int(_) => "int",
267        Value::Float(_) => "float",
268        Value::String(_) => "string",
269        Value::Binary(_) => "binary",
270        Value::Vector(_) => "vector",
271        Value::Uuid(_) => "uuid",
272        Value::List(_) => "list",
273        Value::Map(_) => "map",
274    }
275}
276
277fn parameter_requirement_name(requirement: &GraphRagQueryParameterRequirement) -> String {
278    let value_type = match requirement.value_type {
279        PropertyType::Any => "any",
280        PropertyType::Bool => "bool",
281        PropertyType::Int => "int",
282        PropertyType::Float => "float",
283        PropertyType::String => "string",
284        PropertyType::Text => "text",
285        PropertyType::List => "list",
286        PropertyType::Vector(_) => "vector",
287    };
288    match requirement.cardinality {
289        GraphRagQueryParameterCardinality::Scalar => value_type.to_string(),
290        GraphRagQueryParameterCardinality::List => format!("list<{value_type}>"),
291    }
292}
293
294#[derive(Debug, Clone, PartialEq, Eq)]
295pub enum GraphRagQueryGenerationError {
296    SchemaContextIntegrityMismatch {
297        expected: u64,
298        actual: u64,
299    },
300    SchemaFingerprintMismatch {
301        expected: u64,
302        actual: u64,
303    },
304    InvalidIdentifier {
305        kind: &'static str,
306        value: String,
307    },
308    UnknownLabel(String),
309    UnknownRoute {
310        source_label: String,
311        relationship_type: String,
312        target_label: String,
313    },
314    BindingUnavailable(GraphRagQueryBinding),
315    PropertyUnavailable {
316        binding: GraphRagQueryBinding,
317        property: String,
318    },
319    OperatorRequiresStringProperty {
320        binding: GraphRagQueryBinding,
321        property: String,
322    },
323    MissingParameter {
324        property: String,
325    },
326    UnexpectedParameter {
327        property: String,
328    },
329    ConflictingParameterRequirement {
330        parameter: String,
331        first: GraphRagQueryParameterRequirement,
332        second: GraphRagQueryParameterRequirement,
333    },
334    EmptyProjection,
335    InvalidLimit {
336        limit: usize,
337        maximum: usize,
338    },
339}
340
341impl Display for GraphRagQueryGenerationError {
342    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
343        match self {
344            Self::SchemaContextIntegrityMismatch { expected, actual } => write!(
345                formatter,
346                "schema context integrity mismatch: expected {expected:016x}, got {actual:016x}"
347            ),
348            Self::SchemaFingerprintMismatch { expected, actual } => write!(
349                formatter,
350                "schema fingerprint mismatch: expected {expected:016x}, got {actual:016x}"
351            ),
352            Self::InvalidIdentifier { kind, value } => {
353                write!(formatter, "invalid {kind} identifier: {value}")
354            }
355            Self::UnknownLabel(label) => write!(formatter, "label is not in schema context: {label}"),
356            Self::UnknownRoute {
357                source_label,
358                relationship_type,
359                target_label,
360            } => write!(
361                formatter,
362                "route is not in schema context: ({source_label})-[:{relationship_type}]->({target_label})"
363            ),
364            Self::BindingUnavailable(binding) => {
365                write!(formatter, "query binding is unavailable: {binding:?}")
366            }
367            Self::PropertyUnavailable { binding, property } => write!(
368                formatter,
369                "property is not in schema context for {binding:?}: {property}"
370            ),
371            Self::OperatorRequiresStringProperty { binding, property } => write!(
372                formatter,
373                "predicate requires a string property for {binding:?}: {property}"
374            ),
375            Self::MissingParameter { property } => {
376                write!(formatter, "predicate parameter is required for: {property}")
377            }
378            Self::UnexpectedParameter { property } => {
379                write!(formatter, "predicate parameter is not allowed for: {property}")
380            }
381            Self::ConflictingParameterRequirement {
382                parameter,
383                first,
384                second,
385            } => write!(
386                formatter,
387                "parameter {parameter} has conflicting requirements: {} and {}",
388                parameter_requirement_name(first),
389                parameter_requirement_name(second)
390            ),
391            Self::EmptyProjection => formatter.write_str("at least one projection is required"),
392            Self::InvalidLimit { limit, maximum } => {
393                write!(formatter, "query limit must be between 1 and {maximum}, got {limit}")
394            }
395        }
396    }
397}
398
399impl std::error::Error for GraphRagQueryGenerationError {}