cortex_sdk/
grounded_answer.rs1use std::collections::BTreeSet;
2
3use crate::aql::{Aql, AqlBuildError, AqlRetrievalMode};
4use crate::types::{
5 AnswerGroundingOptionsResponse, AnswerGroundingReportResponse, ContextPackResponse,
6 VerificationReportResponse,
7};
8
9#[derive(Clone, Debug, PartialEq, Eq)]
10pub struct GroundedAnswerRequest {
11 scope: String,
12 brain: String,
13 question: String,
14 mode: Option<AqlRetrievalMode>,
15 budget_tokens: Option<u64>,
16 candidate_limit: Option<u32>,
17 where_clause: Option<String>,
18 require_citations: bool,
19 grounding_options: AnswerGroundingOptionsResponse,
20 verify_answer: bool,
21}
22
23impl GroundedAnswerRequest {
24 pub fn new(
25 scope: impl Into<String>,
26 brain: impl Into<String>,
27 question: impl Into<String>,
28 ) -> Self {
29 Self {
30 scope: scope.into(),
31 brain: brain.into(),
32 question: question.into(),
33 mode: None,
34 budget_tokens: None,
35 candidate_limit: None,
36 where_clause: None,
37 require_citations: true,
38 grounding_options: AnswerGroundingOptionsResponse {
39 require_citations: true,
40 ..AnswerGroundingOptionsResponse::default()
41 },
42 verify_answer: true,
43 }
44 }
45
46 pub fn mode(mut self, mode: AqlRetrievalMode) -> Self {
47 self.mode = Some(mode);
48 self
49 }
50
51 pub fn budget_tokens(mut self, budget_tokens: u64) -> Self {
52 self.budget_tokens = Some(budget_tokens);
53 self
54 }
55
56 pub fn limit_candidates(mut self, candidate_limit: u32) -> Self {
57 self.candidate_limit = Some(candidate_limit);
58 self
59 }
60
61 pub fn where_clause(mut self, where_clause: impl Into<String>) -> Self {
62 self.where_clause = Some(where_clause.into());
63 self
64 }
65
66 pub fn require_citations(mut self, require: bool) -> Self {
67 self.require_citations = require;
68 self.grounding_options.require_citations = require;
69 self
70 }
71
72 pub fn grounding_options(mut self, options: AnswerGroundingOptionsResponse) -> Self {
73 self.grounding_options = options;
74 self.require_citations = options.require_citations;
75 self
76 }
77
78 pub fn verify_answer(mut self, verify: bool) -> Self {
79 self.verify_answer = verify;
80 self
81 }
82
83 pub fn scope(&self) -> &str {
84 &self.scope
85 }
86
87 pub fn brain(&self) -> &str {
88 &self.brain
89 }
90
91 pub fn question(&self) -> &str {
92 &self.question
93 }
94
95 pub fn should_verify_answer(&self) -> bool {
96 self.verify_answer
97 }
98
99 pub fn grounding_options_value(&self) -> AnswerGroundingOptionsResponse {
100 self.grounding_options
101 }
102
103 pub fn retrieve_statement(&self) -> Result<String, AqlBuildError> {
104 let mut builder = Aql::retrieve_context(&self.question, &self.brain);
105 if let Some(mode) = self.mode {
106 builder = builder.mode(mode);
107 }
108 if let Some(budget_tokens) = self.budget_tokens {
109 builder = builder.budget_tokens(budget_tokens);
110 }
111 if let Some(candidate_limit) = self.candidate_limit {
112 builder = builder.limit_candidates(candidate_limit);
113 }
114 if let Some(where_clause) = self.where_clause.as_deref() {
115 builder = builder.where_clause(where_clause);
116 }
117 if self.require_citations {
118 builder = builder.require_citations();
119 }
120 builder.build()
121 }
122
123 pub fn verify_statement(&self, answer: &str) -> Result<Option<String>, AqlBuildError> {
124 if !self.verify_answer || answer.trim().is_empty() {
125 return Ok(None);
126 }
127 Aql::verify_fact(answer, &self.brain).build().map(Some)
128 }
129}
130
131#[derive(Clone, Debug, PartialEq, Eq)]
132pub struct GroundedAnswerResponse {
133 pub question: String,
134 pub answer: String,
135 pub retrieve_statement: String,
136 pub verify_statement: Option<String>,
137 pub context: ContextPackResponse,
138 pub grounding: AnswerGroundingReportResponse,
139 pub verification: Option<VerificationReportResponse>,
140 pub citations: Vec<String>,
141 pub used_context_cell_ids: Vec<u64>,
142 pub rejected: bool,
143}
144
145impl GroundedAnswerResponse {
146 pub fn from_context_answer(
147 request: &GroundedAnswerRequest,
148 retrieve_statement: String,
149 verify_statement: Option<String>,
150 context: ContextPackResponse,
151 answer: String,
152 verification: Option<VerificationReportResponse>,
153 ) -> Self {
154 let grounding =
155 context.ground_answer_with_options(&answer, request.grounding_options_value());
156 let citations = collect_citations(&context, &grounding);
157 let used_context_cell_ids = collect_used_cell_ids(&context, &grounding);
158 let rejected = grounding.rejected;
159 Self {
160 question: request.question().to_owned(),
161 answer,
162 retrieve_statement,
163 verify_statement,
164 context,
165 grounding,
166 verification,
167 citations,
168 used_context_cell_ids,
169 rejected,
170 }
171 }
172
173 pub fn answer_supported(&self) -> bool {
174 self.grounding.answer_supported
175 }
176
177 pub fn verification_status(&self) -> Option<&str> {
178 self.verification
179 .as_ref()
180 .map(|value| value.status.as_str())
181 }
182}
183
184fn collect_citations(
185 context: &ContextPackResponse,
186 grounding: &AnswerGroundingReportResponse,
187) -> Vec<String> {
188 let mut seen = BTreeSet::new();
189 let mut citations = Vec::new();
190 for citation in grounding
191 .spans
192 .iter()
193 .flat_map(|span| span.citations.iter())
194 .chain(
195 context
196 .cells
197 .iter()
198 .filter_map(|cell| cell.citation.as_ref()),
199 )
200 {
201 if !citation.trim().is_empty() && seen.insert(citation.clone()) {
202 citations.push(citation.clone());
203 }
204 }
205 citations
206}
207
208fn collect_used_cell_ids(
209 context: &ContextPackResponse,
210 grounding: &AnswerGroundingReportResponse,
211) -> Vec<u64> {
212 let mut seen = BTreeSet::new();
213 let mut ids = Vec::new();
214 for cell_id in grounding
215 .spans
216 .iter()
217 .flat_map(|span| span.supported_by_cell_ids.iter().copied())
218 .chain(context.cells.iter().map(|cell| cell.cell_id))
219 {
220 if seen.insert(cell_id) {
221 ids.push(cell_id);
222 }
223 }
224 ids
225}