hawdb_core/graph_rag/
mod.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::schema::{Catalog, GraphStatistics, PropertyType, SchemaObjectState, TableKind};
16use std::collections::{BTreeMap, BTreeSet};
17use std::fmt::Write;
18
19mod fingerprint;
20mod query_generation;
21mod scoring;
22mod scoring_program;
23mod topology;
24
25use fingerprint::schema_context_fingerprint;
26pub use query_generation::{
27    GraphRagGeneratedQuery, GraphRagQueryBinding, GraphRagQueryDraft, GraphRagQueryGenerationError,
28    GraphRagQueryParameterCardinality, GraphRagQueryParameterError,
29    GraphRagQueryParameterRequirement, GraphRagQueryPattern, GraphRagQueryPredicate,
30    GraphRagQueryPredicateOperator, GraphRagQueryProjection, MAX_GRAPH_RAG_QUERY_LIMIT,
31};
32pub use scoring::{
33    DecayTerm, ScoreFeature, ScoringEvaluation, ScoringFeatureSource, ScoringSpec,
34    ScoringSpecError, ScoringTerm,
35};
36pub use scoring_program::{
37    MissingScoringFeature, ScoringCombination, ScoringProgram, ScoringProgramError,
38    ScoringProgramShape,
39};
40use topology::{common_path_summaries, route_summaries};
41
42pub const GRAPH_RAG_SCHEMA_CONTEXT_PROTOCOL: &str = "hawdb-graph-rag-schema-context-v1";
43pub const DEFAULT_GRAPH_RAG_MAX_LABELS: usize = 32;
44pub const DEFAULT_GRAPH_RAG_MAX_RELATIONSHIP_TYPES: usize = 32;
45pub const DEFAULT_GRAPH_RAG_MAX_PROPERTIES_PER_SUBJECT: usize = 12;
46pub const DEFAULT_GRAPH_RAG_MAX_ROUTES: usize = 64;
47pub const DEFAULT_GRAPH_RAG_MAX_COMMON_PATHS: usize = 32;
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct GraphRagSchemaContextOptions {
51    pub max_labels: usize,
52    pub max_relationship_types: usize,
53    pub max_properties_per_subject: usize,
54    pub max_routes: usize,
55    pub max_common_paths: usize,
56}
57
58impl Default for GraphRagSchemaContextOptions {
59    fn default() -> Self {
60        Self {
61            max_labels: DEFAULT_GRAPH_RAG_MAX_LABELS,
62            max_relationship_types: DEFAULT_GRAPH_RAG_MAX_RELATIONSHIP_TYPES,
63            max_properties_per_subject: DEFAULT_GRAPH_RAG_MAX_PROPERTIES_PER_SUBJECT,
64            max_routes: DEFAULT_GRAPH_RAG_MAX_ROUTES,
65            max_common_paths: DEFAULT_GRAPH_RAG_MAX_COMMON_PATHS,
66        }
67    }
68}
69
70#[derive(Debug, Clone, PartialEq, Eq)]
71pub struct GraphRagLabelSummary {
72    pub name: String,
73    pub node_count: u64,
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct GraphRagRelationshipTypeSummary {
78    pub name: String,
79    pub relationship_count: u64,
80    pub distinct_source_count: u64,
81    pub distinct_target_count: u64,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
85pub enum GraphRagPropertySubject {
86    Node,
87    Relationship,
88}
89
90impl GraphRagPropertySubject {
91    pub const fn as_str(self) -> &'static str {
92        match self {
93            Self::Node => "node",
94            Self::Relationship => "relationship",
95        }
96    }
97}
98
99#[derive(Debug, Clone, PartialEq, Eq)]
100pub struct GraphRagPropertySummary {
101    pub subject: GraphRagPropertySubject,
102    pub subject_name: String,
103    pub name: String,
104    pub value_type: PropertyType,
105    pub nullable: bool,
106    pub declared: bool,
107    pub distinct_count: Option<u64>,
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct GraphRagRouteSummary {
112    pub source_label: String,
113    pub relationship_type: String,
114    pub target_label: String,
115    pub observed_count: u64,
116    pub distinct_source_count: u64,
117    pub distinct_target_count: u64,
118}
119
120#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct GraphRagCommonPathSummary {
122    pub source_label: String,
123    pub relationship_type: String,
124    pub target_label: String,
125    pub hops: usize,
126    pub observed_count: u64,
127    pub distinct_source_count: u64,
128    pub distinct_target_count: u64,
129}
130
131#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
132pub struct GraphRagSchemaContextTruncation {
133    pub labels: bool,
134    pub relationship_types: bool,
135    pub properties: bool,
136    pub routes: bool,
137    pub common_paths: bool,
138}
139
140impl GraphRagSchemaContextTruncation {
141    pub const fn any(self) -> bool {
142        self.labels
143            || self.relationship_types
144            || self.properties
145            || self.routes
146            || self.common_paths
147    }
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
151pub struct GraphRagSchemaContext {
152    pub protocol: &'static str,
153    pub computed_at_commit_epoch: u64,
154    pub fingerprint: u64,
155    pub labels: Vec<GraphRagLabelSummary>,
156    pub relationship_types: Vec<GraphRagRelationshipTypeSummary>,
157    pub properties: Vec<GraphRagPropertySummary>,
158    pub routes: Vec<GraphRagRouteSummary>,
159    pub common_paths: Vec<GraphRagCommonPathSummary>,
160    pub truncation: GraphRagSchemaContextTruncation,
161}
162
163impl GraphRagSchemaContext {
164    pub fn generate_query(
165        &self,
166        draft: &GraphRagQueryDraft,
167    ) -> Result<GraphRagGeneratedQuery, GraphRagQueryGenerationError> {
168        query_generation::generate_query(self, draft)
169    }
170
171    pub fn render_compact_cypher_guidance(&self) -> String {
172        let mut output = String::new();
173        let _ = writeln!(
174            output,
175            "SCHEMA protocol={} epoch={} fingerprint={:016x}",
176            self.protocol, self.computed_at_commit_epoch, self.fingerprint
177        );
178        output.push_str(
179            "RULES read_only=true parameterize_values=true identifiers_from_schema_only=true max_hops=2\n",
180        );
181        for label in &self.labels {
182            let _ = writeln!(output, "LABEL {} count={}", label.name, label.node_count);
183        }
184        for relationship in &self.relationship_types {
185            let _ = writeln!(
186                output,
187                "RELATIONSHIP_TYPE {} count={} sources={} targets={}",
188                relationship.name,
189                relationship.relationship_count,
190                relationship.distinct_source_count,
191                relationship.distinct_target_count
192            );
193        }
194        for property in &self.properties {
195            let _ = writeln!(
196                output,
197                "PROPERTY {} {}.{} type={} nullable={} declared={} distinct={}",
198                property.subject.as_str(),
199                property.subject_name,
200                property.name,
201                property_type_name(property.value_type),
202                property.nullable,
203                property.declared,
204                optional_count(property.distinct_count)
205            );
206        }
207        for route in &self.routes {
208            let _ = writeln!(
209                output,
210                "ROUTE ({})-[:{}]->({}) count={} sources={} targets={}",
211                route.source_label,
212                route.relationship_type,
213                route.target_label,
214                route.observed_count,
215                route.distinct_source_count,
216                route.distinct_target_count
217            );
218        }
219        for path in &self.common_paths {
220            let _ = writeln!(
221                output,
222                "COMMON_PATH ({})-[:{}*{}]->({}) count={} sources={} targets={}",
223                path.source_label,
224                path.relationship_type,
225                path.hops,
226                path.target_label,
227                path.observed_count,
228                path.distinct_source_count,
229                path.distinct_target_count
230            );
231        }
232        if self.truncation.any() {
233            let _ = writeln!(
234                output,
235                "TRUNCATED labels={} relationship_types={} properties={} routes={} common_paths={}",
236                self.truncation.labels,
237                self.truncation.relationship_types,
238                self.truncation.properties,
239                self.truncation.routes,
240                self.truncation.common_paths
241            );
242        }
243        output
244    }
245}
246
247pub fn build_graph_rag_schema_context(
248    catalog: &Catalog,
249    statistics: &GraphStatistics,
250    options: GraphRagSchemaContextOptions,
251) -> GraphRagSchemaContext {
252    let mut labels = catalog
253        .labels()
254        .map(|label| GraphRagLabelSummary {
255            name: label.name.clone(),
256            node_count: statistics
257                .label_counts
258                .get(&label.id)
259                .copied()
260                .unwrap_or_default(),
261        })
262        .collect::<Vec<_>>();
263    labels.sort_by(|left, right| {
264        right
265            .node_count
266            .cmp(&left.node_count)
267            .then_with(|| left.name.cmp(&right.name))
268    });
269    let labels_truncated = truncate(&mut labels, options.max_labels);
270    let selected_labels = labels
271        .iter()
272        .map(|label| label.name.clone())
273        .collect::<BTreeSet<_>>();
274
275    let mut relationship_types = catalog
276        .rel_types()
277        .map(|rel_type| GraphRagRelationshipTypeSummary {
278            name: rel_type.name.clone(),
279            relationship_count: statistics
280                .rel_type_counts
281                .get(&rel_type.id)
282                .copied()
283                .unwrap_or_default(),
284            distinct_source_count: statistics
285                .rel_type_source_counts
286                .get(&rel_type.id)
287                .copied()
288                .unwrap_or_default(),
289            distinct_target_count: statistics
290                .rel_type_target_counts
291                .get(&rel_type.id)
292                .copied()
293                .unwrap_or_default(),
294        })
295        .collect::<Vec<_>>();
296    relationship_types.sort_by(|left, right| {
297        right
298            .relationship_count
299            .cmp(&left.relationship_count)
300            .then_with(|| left.name.cmp(&right.name))
301    });
302    let relationship_types_truncated =
303        truncate(&mut relationship_types, options.max_relationship_types);
304    let selected_relationship_types = relationship_types
305        .iter()
306        .map(|rel_type| rel_type.name.clone())
307        .collect::<BTreeSet<_>>();
308
309    let (properties, properties_truncated) = property_summaries(
310        catalog,
311        statistics,
312        &selected_labels,
313        &selected_relationship_types,
314        options.max_properties_per_subject,
315    );
316    let (routes, routes_truncated) = route_summaries(
317        catalog,
318        statistics,
319        &selected_labels,
320        &selected_relationship_types,
321        options.max_routes,
322    );
323    let (common_paths, common_paths_truncated) = common_path_summaries(
324        catalog,
325        statistics,
326        &selected_labels,
327        &selected_relationship_types,
328        options.max_common_paths,
329    );
330
331    let mut context = GraphRagSchemaContext {
332        protocol: GRAPH_RAG_SCHEMA_CONTEXT_PROTOCOL,
333        computed_at_commit_epoch: statistics.computed_at_commit_epoch,
334        fingerprint: 0,
335        labels,
336        relationship_types,
337        properties,
338        routes,
339        common_paths,
340        truncation: GraphRagSchemaContextTruncation {
341            labels: labels_truncated,
342            relationship_types: relationship_types_truncated,
343            properties: properties_truncated,
344            routes: routes_truncated,
345            common_paths: common_paths_truncated,
346        },
347    };
348    context.fingerprint = schema_context_fingerprint(&context);
349    context
350}
351
352fn property_summaries(
353    catalog: &Catalog,
354    statistics: &GraphStatistics,
355    selected_labels: &BTreeSet<String>,
356    selected_relationship_types: &BTreeSet<String>,
357    max_properties_per_subject: usize,
358) -> (Vec<GraphRagPropertySummary>, bool) {
359    let public_tables = catalog
360        .table_descriptors()
361        .filter(|table| table.state == SchemaObjectState::Public)
362        .map(|table| (table.id, (table.kind, table.name.clone())))
363        .collect::<BTreeMap<_, _>>();
364    let mut summaries =
365        BTreeMap::<(GraphRagPropertySubject, String, String), GraphRagPropertySummary>::new();
366
367    for property in catalog
368        .property_descriptors()
369        .filter(|property| property.state == SchemaObjectState::Public)
370    {
371        let Some((table_kind, table_name)) = public_tables.get(&property.table_id) else {
372            continue;
373        };
374        let subject = match table_kind {
375            TableKind::Node if selected_labels.contains(table_name) => {
376                GraphRagPropertySubject::Node
377            }
378            TableKind::Relationship if selected_relationship_types.contains(table_name) => {
379                GraphRagPropertySubject::Relationship
380            }
381            _ => continue,
382        };
383        let distinct_count =
384            property_distinct_count(catalog, statistics, subject, table_name, &property.name);
385        summaries.insert(
386            (subject, table_name.clone(), property.name.clone()),
387            GraphRagPropertySummary {
388                subject,
389                subject_name: table_name.clone(),
390                name: property.name.clone(),
391                value_type: property.value_type,
392                nullable: property.nullable,
393                declared: true,
394                distinct_count,
395            },
396        );
397    }
398
399    for ((label_id, property), distinct_count) in &statistics.property_distinct_counts {
400        let Some(label) = catalog.label_name(*label_id) else {
401            continue;
402        };
403        if !selected_labels.contains(label) {
404            continue;
405        }
406        summaries
407            .entry((
408                GraphRagPropertySubject::Node,
409                label.to_string(),
410                property.clone(),
411            ))
412            .or_insert_with(|| GraphRagPropertySummary {
413                subject: GraphRagPropertySubject::Node,
414                subject_name: label.to_string(),
415                name: property.clone(),
416                value_type: PropertyType::Any,
417                nullable: true,
418                declared: false,
419                distinct_count: Some(*distinct_count),
420            });
421    }
422    for ((rel_type_id, property), distinct_count) in &statistics.rel_property_distinct_counts {
423        let Some(rel_type) = catalog.rel_type_name(*rel_type_id) else {
424            continue;
425        };
426        if !selected_relationship_types.contains(rel_type) {
427            continue;
428        }
429        summaries
430            .entry((
431                GraphRagPropertySubject::Relationship,
432                rel_type.to_string(),
433                property.clone(),
434            ))
435            .or_insert_with(|| GraphRagPropertySummary {
436                subject: GraphRagPropertySubject::Relationship,
437                subject_name: rel_type.to_string(),
438                name: property.clone(),
439                value_type: PropertyType::Any,
440                nullable: true,
441                declared: false,
442                distinct_count: Some(*distinct_count),
443            });
444    }
445
446    let mut grouped =
447        BTreeMap::<(GraphRagPropertySubject, String), Vec<GraphRagPropertySummary>>::new();
448    for summary in summaries.into_values() {
449        grouped
450            .entry((summary.subject, summary.subject_name.clone()))
451            .or_default()
452            .push(summary);
453    }
454    let mut output = Vec::new();
455    let mut truncated = false;
456    for properties in grouped.values_mut() {
457        properties.sort_by(|left, right| {
458            right
459                .declared
460                .cmp(&left.declared)
461                .then_with(|| left.name.cmp(&right.name))
462        });
463        truncated |= truncate(properties, max_properties_per_subject);
464        output.append(properties);
465    }
466    (output, truncated)
467}
468
469fn property_distinct_count(
470    catalog: &Catalog,
471    statistics: &GraphStatistics,
472    subject: GraphRagPropertySubject,
473    subject_name: &str,
474    property: &str,
475) -> Option<u64> {
476    match subject {
477        GraphRagPropertySubject::Node => catalog.label_id(subject_name).and_then(|label_id| {
478            statistics
479                .property_distinct_counts
480                .get(&(label_id, property.to_string()))
481                .copied()
482        }),
483        GraphRagPropertySubject::Relationship => {
484            catalog.rel_type_id(subject_name).and_then(|rel_type_id| {
485                statistics
486                    .rel_property_distinct_counts
487                    .get(&(rel_type_id, property.to_string()))
488                    .copied()
489            })
490        }
491    }
492}
493
494fn truncate<T>(values: &mut Vec<T>, limit: usize) -> bool {
495    let truncated = values.len() > limit;
496    values.truncate(limit);
497    truncated
498}
499
500fn property_type_name(value_type: PropertyType) -> std::borrow::Cow<'static, str> {
501    match value_type {
502        PropertyType::Any => "any",
503        PropertyType::Bool => "bool",
504        PropertyType::Int => "int",
505        PropertyType::Float => "float",
506        PropertyType::String => "string",
507        PropertyType::Text => "text",
508        PropertyType::List => "list",
509        PropertyType::Vector(dimension) => return format!("vector({})", dimension.get()).into(),
510    }
511    .into()
512}
513
514fn optional_count(count: Option<u64>) -> String {
515    count
516        .map(|count| count.to_string())
517        .unwrap_or_else(|| "unknown".to_string())
518}
519
520#[cfg(test)]
521mod query_generation_tests;
522#[cfg(test)]
523mod tests;
hawdb_core/graph_rag/mod.rs

hawdb_core/graph_rag/
mod.rs
