cortex_sdk/
aql.rs

1use crate::aql_support::{
2    quote_string, validate_identifier, validate_optional_decimal, validate_optional_fragment,
3};
4pub use crate::aql_support::{AqlBuildError, AqlRetrievalMode};
5
6#[derive(Clone, Debug, PartialEq, Eq)]
7pub struct RetrieveContextBuilder {
8    task: String,
9    brain: String,
10    mode: Option<AqlRetrievalMode>,
11    budget_tokens: Option<u64>,
12    candidate_limit: Option<u32>,
13    where_clause: Option<String>,
14    require_citations: bool,
15    min_confidence: Option<String>,
16    source_trust: Option<String>,
17    freshness_seconds: Option<u64>,
18    explain: bool,
19}
20
21impl RetrieveContextBuilder {
22    pub fn new(task: impl Into<String>, brain: impl Into<String>) -> Self {
23        Self {
24            task: task.into(),
25            brain: brain.into(),
26            mode: None,
27            budget_tokens: None,
28            candidate_limit: None,
29            where_clause: None,
30            require_citations: false,
31            min_confidence: None,
32            source_trust: None,
33            freshness_seconds: None,
34            explain: false,
35        }
36    }
37
38    pub fn mode(mut self, mode: AqlRetrievalMode) -> Self {
39        self.mode = Some(mode);
40        self
41    }
42
43    pub fn budget_tokens(mut self, budget_tokens: u64) -> Self {
44        self.budget_tokens = Some(budget_tokens);
45        self
46    }
47
48    pub fn limit_candidates(mut self, candidate_limit: u32) -> Self {
49        self.candidate_limit = Some(candidate_limit);
50        self
51    }
52
53    pub fn where_clause(mut self, where_clause: impl Into<String>) -> Self {
54        self.where_clause = Some(where_clause.into());
55        self
56    }
57
58    pub fn require_citations(mut self) -> Self {
59        self.require_citations = true;
60        self
61    }
62
63    pub fn require_confidence_at_least(mut self, decimal: impl Into<String>) -> Self {
64        self.min_confidence = Some(decimal.into());
65        self
66    }
67
68    pub fn require_source_trust_at_least(mut self, decimal: impl Into<String>) -> Self {
69        self.source_trust = Some(decimal.into());
70        self
71    }
72
73    pub fn require_freshness_seconds(mut self, seconds: u64) -> Self {
74        self.freshness_seconds = Some(seconds);
75        self
76    }
77
78    pub fn explain(mut self) -> Self {
79        self.explain = true;
80        self
81    }
82
83    pub fn build(&self) -> Result<String, AqlBuildError> {
84        validate_identifier("brain", &self.brain)?;
85        validate_optional_fragment("where_clause", self.where_clause.as_deref())?;
86        validate_optional_decimal("min_confidence", self.min_confidence.as_deref())?;
87        validate_optional_decimal("source_trust", self.source_trust.as_deref())?;
88
89        let mut statement = String::new();
90        if self.explain {
91            statement.push_str("EXPLAIN ");
92        }
93        statement.push_str("RETRIEVE CONTEXT FOR TASK ");
94        statement.push_str(&quote_string(&self.task));
95        statement.push_str(" IN BRAIN ");
96        statement.push_str(&self.brain);
97        if let Some(mode) = self.mode {
98            statement.push_str(" USING MODE ");
99            statement.push_str(mode.as_str());
100        }
101        if let Some(budget) = self.budget_tokens {
102            statement.push_str(" BUDGET ");
103            statement.push_str(&budget.to_string());
104            statement.push_str(" TOKENS");
105        }
106        if let Some(limit) = self.candidate_limit {
107            statement.push_str(" LIMIT ");
108            statement.push_str(&limit.to_string());
109            statement.push_str(" CANDIDATES");
110        }
111        if let Some(where_clause) = self.where_clause.as_deref() {
112            statement.push_str(" WHERE ");
113            statement.push_str(where_clause.trim());
114        }
115        if self.require_citations {
116            statement.push_str(" REQUIRE citations");
117        }
118        if let Some(decimal) = self.min_confidence.as_deref() {
119            statement.push_str(" REQUIRE confidence >= ");
120            statement.push_str(decimal);
121        }
122        if let Some(decimal) = self.source_trust.as_deref() {
123            statement.push_str(" REQUIRE source_trust >= ");
124            statement.push_str(decimal);
125        }
126        if let Some(seconds) = self.freshness_seconds {
127            statement.push_str(" REQUIRE freshness <= ");
128            statement.push_str(&seconds.to_string());
129            statement.push_str(" SECONDS");
130        }
131        statement.push(';');
132        Ok(statement)
133    }
134}
135
136#[derive(Clone, Debug, PartialEq, Eq)]
137pub struct VerifyFactBuilder {
138    fact: String,
139    brain: String,
140}
141
142impl VerifyFactBuilder {
143    pub fn new(fact: impl Into<String>, brain: impl Into<String>) -> Self {
144        Self {
145            fact: fact.into(),
146            brain: brain.into(),
147        }
148    }
149
150    pub fn build(&self) -> Result<String, AqlBuildError> {
151        validate_identifier("brain", &self.brain)?;
152        Ok(format!(
153            "VERIFY FACT {} IN BRAIN {};",
154            quote_string(&self.fact),
155            self.brain
156        ))
157    }
158}
159
160#[derive(Clone, Debug, PartialEq, Eq)]
161pub struct RememberBuilder {
162    content: String,
163    scope: String,
164    memory_type: String,
165    ttl_seconds: Option<u64>,
166}
167
168impl RememberBuilder {
169    pub fn new(
170        content: impl Into<String>,
171        scope: impl Into<String>,
172        memory_type: impl Into<String>,
173    ) -> Self {
174        Self {
175            content: content.into(),
176            scope: scope.into(),
177            memory_type: memory_type.into(),
178            ttl_seconds: None,
179        }
180    }
181
182    pub fn ttl_seconds(mut self, seconds: u64) -> Self {
183        self.ttl_seconds = Some(seconds);
184        self
185    }
186
187    pub fn build(&self) -> Result<String, AqlBuildError> {
188        validate_identifier("scope", &self.scope)?;
189        validate_identifier("memory_type", &self.memory_type)?;
190        let mut statement = format!(
191            "REMEMBER {} IN SCOPE {} AS TYPE {}",
192            quote_string(&self.content),
193            self.scope,
194            self.memory_type
195        );
196        if let Some(seconds) = self.ttl_seconds {
197            statement.push_str(" TTL ");
198            statement.push_str(&seconds.to_string());
199            statement.push_str(" SECONDS");
200        }
201        statement.push(';');
202        Ok(statement)
203    }
204}
205
206#[derive(Clone, Copy, Debug, PartialEq, Eq)]
207pub struct Aql;
208
209impl Aql {
210    pub fn retrieve_context(
211        task: impl Into<String>,
212        brain: impl Into<String>,
213    ) -> RetrieveContextBuilder {
214        RetrieveContextBuilder::new(task, brain)
215    }
216
217    pub fn verify_fact(fact: impl Into<String>, brain: impl Into<String>) -> VerifyFactBuilder {
218        VerifyFactBuilder::new(fact, brain)
219    }
220
221    pub fn remember(
222        content: impl Into<String>,
223        scope: impl Into<String>,
224        memory_type: impl Into<String>,
225    ) -> RememberBuilder {
226        RememberBuilder::new(content, scope, memory_type)
227    }
228}
cortex_sdk/aql.rs

cortex_sdk/
aql.rs
