1use 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 {}