1use std::collections::BTreeSet;
2
3use serde::Serialize;
4
5use plexus_serde::{Op, Plan};
6
7use crate::capabilities::{required_capabilities, EngineCapabilities, ExprKind, OpKind};
8
9pub const EMBEDDED_GRAPH_RAG_PROFILE: &str = "embedded-graph-rag-v0";
10pub const EMBEDDED_PHASE_1_RELEASE_CANDIDATE: &str = "embedded-phase-1";
11pub const PLEXUS_CAPABILITY_REJECTION_SCENARIO: &str = "sr.plexus.capability-rejection.v1";
12
13fn allowed_ops() -> BTreeSet<OpKind> {
14 BTreeSet::from([
15 OpKind::ConstRow,
16 OpKind::Filter,
17 OpKind::Project,
18 OpKind::Limit,
19 OpKind::VectorScan,
20 OpKind::Rerank,
21 OpKind::Return,
22 ])
23}
24
25fn required_ops() -> BTreeSet<OpKind> {
26 BTreeSet::from([OpKind::VectorScan, OpKind::Rerank, OpKind::Return])
27}
28
29fn allowed_exprs() -> BTreeSet<ExprKind> {
30 BTreeSet::from([
31 ExprKind::ColRef,
32 ExprKind::PropAccess,
33 ExprKind::IntLiteral,
34 ExprKind::FloatLiteral,
35 ExprKind::BoolLiteral,
36 ExprKind::StringLiteral,
37 ExprKind::NullLiteral,
38 ExprKind::Cmp,
39 ExprKind::And,
40 ExprKind::Or,
41 ExprKind::Not,
42 ExprKind::IsNull,
43 ExprKind::IsNotNull,
44 ExprKind::StartsWith,
45 ExprKind::EndsWith,
46 ExprKind::Contains,
47 ExprKind::ListLiteral,
48 ExprKind::Arith,
49 ExprKind::Param,
50 ExprKind::Case,
51 ExprKind::VectorSimilarity,
52 ])
53}
54
55fn has_graph_ref(plan: &Plan) -> bool {
56 plan.ops.iter().any(|op| match op {
57 Op::ScanNodes { graph_ref, .. }
58 | Op::Expand { graph_ref, .. }
59 | Op::OptionalExpand { graph_ref, .. }
60 | Op::ExpandVarLen { graph_ref, .. } => graph_ref.is_some(),
61 _ => false,
62 })
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
66pub struct EmbeddedGraphRagAssessment {
67 pub profile: String,
68 pub supported: bool,
69 pub required_ops: Vec<String>,
70 pub allowed_ops: Vec<String>,
71 pub allowed_exprs: Vec<String>,
72 pub missing_required_ops: Vec<String>,
73 pub unsupported_ops: Vec<String>,
74 pub unsupported_exprs: Vec<String>,
75 pub graph_ref_unsupported: bool,
76}
77
78impl EmbeddedGraphRagAssessment {
79 pub fn rejection_summary(&self) -> String {
80 let mut parts = Vec::new();
81 if !self.missing_required_ops.is_empty() {
82 parts.push(format!(
83 "missing required ops [{}]",
84 self.missing_required_ops.join(", ")
85 ));
86 }
87 if !self.unsupported_ops.is_empty() {
88 parts.push(format!(
89 "unsupported ops [{}]",
90 self.unsupported_ops.join(", ")
91 ));
92 }
93 if !self.unsupported_exprs.is_empty() {
94 parts.push(format!(
95 "unsupported exprs [{}]",
96 self.unsupported_exprs.join(", ")
97 ));
98 }
99 if self.graph_ref_unsupported {
100 parts.push("graph_ref is not supported".to_string());
101 }
102 parts.join("; ")
103 }
104}
105
106#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
107#[error(
108 "plan is outside {profile}: {reason}",
109 profile = .assessment.profile,
110 reason = .assessment.rejection_summary()
111)]
112pub struct EmbeddedGraphRagProfileError {
113 pub assessment: EmbeddedGraphRagAssessment,
114}
115
116#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
117pub struct EmbeddedGraphRagCapabilityGap {
118 pub profile: String,
119 pub supported: bool,
120 pub missing_ops: Vec<String>,
121 pub missing_exprs: Vec<String>,
122}
123
124pub fn assess_embedded_graph_rag_plan(plan: &Plan) -> EmbeddedGraphRagAssessment {
125 let required = required_capabilities(plan);
126 let allowed_ops = allowed_ops();
127 let allowed_exprs = allowed_exprs();
128 let profile_required_ops = required_ops();
129
130 let missing_required_ops: Vec<_> = profile_required_ops
131 .difference(&required.required_ops)
132 .copied()
133 .map(OpKind::as_str)
134 .map(str::to_string)
135 .collect();
136 let unsupported_ops: Vec<_> = required
137 .required_ops
138 .difference(&allowed_ops)
139 .copied()
140 .map(OpKind::as_str)
141 .map(str::to_string)
142 .collect();
143 let unsupported_exprs: Vec<_> = required
144 .required_exprs
145 .difference(&allowed_exprs)
146 .copied()
147 .map(ExprKind::as_str)
148 .map(str::to_string)
149 .collect();
150 let graph_ref_unsupported = has_graph_ref(plan);
151 let supported = missing_required_ops.is_empty()
152 && unsupported_ops.is_empty()
153 && unsupported_exprs.is_empty()
154 && !graph_ref_unsupported;
155
156 EmbeddedGraphRagAssessment {
157 profile: EMBEDDED_GRAPH_RAG_PROFILE.to_string(),
158 supported,
159 required_ops: profile_required_ops
160 .iter()
161 .copied()
162 .map(OpKind::as_str)
163 .map(str::to_string)
164 .collect(),
165 allowed_ops: allowed_ops
166 .iter()
167 .copied()
168 .map(OpKind::as_str)
169 .map(str::to_string)
170 .collect(),
171 allowed_exprs: allowed_exprs
172 .iter()
173 .copied()
174 .map(ExprKind::as_str)
175 .map(str::to_string)
176 .collect(),
177 missing_required_ops,
178 unsupported_ops,
179 unsupported_exprs,
180 graph_ref_unsupported,
181 }
182}
183
184#[allow(clippy::result_large_err)]
185pub fn validate_embedded_graph_rag_plan(plan: &Plan) -> Result<(), EmbeddedGraphRagProfileError> {
186 let assessment = assess_embedded_graph_rag_plan(plan);
187 if assessment.supported {
188 Ok(())
189 } else {
190 Err(EmbeddedGraphRagProfileError { assessment })
191 }
192}
193
194pub fn assess_embedded_graph_rag_capabilities(
195 capabilities: &EngineCapabilities,
196) -> EmbeddedGraphRagCapabilityGap {
197 let missing_ops: Vec<_> = required_ops()
198 .difference(&capabilities.supported_ops)
199 .copied()
200 .map(OpKind::as_str)
201 .map(str::to_string)
202 .collect();
203 let missing_exprs = if capabilities
204 .supported_exprs
205 .contains(&ExprKind::VectorSimilarity)
206 {
207 Vec::new()
208 } else {
209 vec![ExprKind::VectorSimilarity.as_str().to_string()]
210 };
211
212 EmbeddedGraphRagCapabilityGap {
213 profile: EMBEDDED_GRAPH_RAG_PROFILE.to_string(),
214 supported: missing_ops.is_empty() && missing_exprs.is_empty(),
215 missing_ops,
216 missing_exprs,
217 }
218}