gemini_adk_fluent_rs/live/extraction.rs
1//! Extraction pipeline configuration methods for `Live`.
2
3use std::future::Future;
4use std::sync::Arc;
5
6use serde::Serialize;
7use serde::de::DeserializeOwned;
8
9use gemini_adk_rs::live::extractor::{ExtractionTrigger, LlmExtractor, TurnExtractor};
10use gemini_adk_rs::llm::BaseLlm;
11
12use super::Live;
13
14impl Live {
15 // -- Turn Extraction Pipeline --
16
17 /// Add a turn extractor that runs an OOB LLM after each turn to extract
18 /// structured data from the transcript window.
19 ///
20 /// Automatically enables both input and output transcription.
21 /// The extraction result is stored in `State` under the type name
22 /// (e.g., `"OrderState"`) and can be read via `handle.extracted::<T>(name)`.
23 ///
24 /// The type `T` must implement `JsonSchema` for schema-guided extraction.
25 /// The window size defaults to 3 turns.
26 pub fn extract_turns<T>(self, llm: Arc<dyn BaseLlm>, prompt: impl Into<String>) -> Self
27 where
28 T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
29 {
30 self.extract_turns_windowed::<T>(llm, prompt, 3)
31 }
32
33 /// Register a deterministic [`Extract`](gemini_adk_rs::extract::Extract)
34 /// record — CPU recognizers over the transcript, no model, no network. The
35 /// recognized fields are promoted into `State`, where `Flow` guards
36 /// (`done(captured([...]))`) and repair read them. Composes with
37 /// `extract_turns` (LLM) on the same session for a cheap-first cascade.
38 pub fn extract_record(mut self, spec: gemini_adk_rs::extract::Extract) -> Self {
39 self.config = self.config.input_transcription(true);
40 self.extractors.push(spec.into_extractor());
41 self
42 }
43
44 /// Like `extract_turns` but with a custom window size.
45 pub fn extract_turns_windowed<T>(
46 mut self,
47 llm: Arc<dyn BaseLlm>,
48 prompt: impl Into<String>,
49 window_size: usize,
50 ) -> Self
51 where
52 T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
53 {
54 // Auto-enable transcription
55 self.config = self
56 .config
57 .input_transcription(true)
58 .output_transcription(true);
59
60 // Derive name from type
61 let name = std::any::type_name::<T>()
62 .rsplit("::")
63 .next()
64 .unwrap_or("Extraction")
65 .to_string();
66
67 let schema = gemini_adk_rs::tool::wire_schema::<T>();
68
69 // Auto-register LLM for connection warming
70 self.warm_up_llms.push(llm.clone());
71
72 let extractor = LlmExtractor::new(name, llm, prompt, window_size)
73 .with_schema(schema)
74 .with_min_words(3);
75 self.extractors.push(Arc::new(extractor));
76 self
77 }
78
79 /// Like `extract_turns_windowed` but with a custom extraction trigger.
80 ///
81 /// Use `ExtractionTrigger::AfterToolCall` when tool calls are the primary
82 /// state source, `ExtractionTrigger::Interval(n)` to reduce extraction
83 /// frequency, or `ExtractionTrigger::OnPhaseChange` for phase-entry extraction.
84 pub fn extract_turns_triggered<T>(
85 mut self,
86 llm: Arc<dyn BaseLlm>,
87 prompt: impl Into<String>,
88 window_size: usize,
89 trigger: ExtractionTrigger,
90 ) -> Self
91 where
92 T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
93 {
94 // Auto-enable transcription
95 self.config = self
96 .config
97 .input_transcription(true)
98 .output_transcription(true);
99
100 let name = std::any::type_name::<T>()
101 .rsplit("::")
102 .next()
103 .unwrap_or("Extraction")
104 .to_string();
105
106 let schema = gemini_adk_rs::tool::wire_schema::<T>();
107
108 self.warm_up_llms.push(llm.clone());
109
110 let extractor = LlmExtractor::new(name, llm, prompt, window_size)
111 .with_schema(schema)
112 .with_min_words(3)
113 .with_trigger(trigger);
114 self.extractors.push(Arc::new(extractor));
115 self
116 }
117
118 /// Like [`extract_turns_triggered`](Self::extract_turns_triggered), but lets
119 /// callers configure the underlying [`LlmExtractor`] before registration.
120 ///
121 /// Use this for field promotion rules, custom minimum word counts, or other
122 /// extraction policies that should live at the SDK layer instead of app
123 /// callback glue.
124 pub fn extract_turns_configured<T>(
125 mut self,
126 llm: Arc<dyn BaseLlm>,
127 prompt: impl Into<String>,
128 window_size: usize,
129 trigger: ExtractionTrigger,
130 configure: impl FnOnce(LlmExtractor) -> LlmExtractor,
131 ) -> Self
132 where
133 T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
134 {
135 self.config = self
136 .config
137 .input_transcription(true)
138 .output_transcription(true);
139
140 let name = std::any::type_name::<T>()
141 .rsplit("::")
142 .next()
143 .unwrap_or("Extraction")
144 .to_string();
145
146 let schema = gemini_adk_rs::tool::wire_schema::<T>();
147
148 self.warm_up_llms.push(llm.clone());
149
150 let extractor = LlmExtractor::new(name, llm, prompt, window_size)
151 .with_schema(schema)
152 .with_min_words(3)
153 .with_trigger(trigger);
154 self.extractors.push(Arc::new(configure(extractor)));
155 self
156 }
157
158 /// Like [`extract_turns`](Self::extract_turns), but schema-as-data: no
159 /// Rust type required. The extraction result is stored in `State` under
160 /// `name`; pair with
161 /// [`FieldPromotion`](gemini_adk_rs::live::extractor::FieldPromotion)
162 /// rules via [`extractor`](Self::extractor) (or a
163 /// [`SessionSpec`](crate::spec::SessionSpec) `extract` entry, which wires
164 /// promotions declaratively) to land individual fields in the bare keys
165 /// flow guards read.
166 ///
167 /// This is the piece that lets a JSON-authored flow advance from speech
168 /// alone: extraction fills the state that `captured`/`is_true` guards
169 /// latch on, with no tool call anywhere.
170 pub fn extract_json(
171 mut self,
172 llm: Arc<dyn BaseLlm>,
173 name: impl Into<String>,
174 schema: serde_json::Value,
175 prompt: impl Into<String>,
176 ) -> Self {
177 self.config = self
178 .config
179 .input_transcription(true)
180 .output_transcription(true);
181 self.warm_up_llms.push(llm.clone());
182 let extractor = LlmExtractor::new(name.into(), llm, prompt.into(), 3)
183 .with_schema(schema)
184 .with_min_words(3);
185 self.extractors.push(Arc::new(extractor));
186 self
187 }
188
189 /// The questions a decision model answers about the conversation, for
190 /// guards that use the `decided` atom (`Guard::decided("confirmed")`,
191 /// `{"decided": "confirmed"}` in a spec).
192 ///
193 /// At each decision point (the caller's turn ends, or the model calls a
194 /// tool) the runtime asks the questions the flow can act on, in one
195 /// request, against the rolling conversation. Turns on input and output
196 /// transcription, which the conversation is built from.
197 ///
198 /// ```rust,ignore
199 /// use gemini_adk_rs::decision::{Decision, Decisions, GatewayDecisionModel, Question};
200 ///
201 /// let decisions = Decisions::new(Arc::new(GatewayDecisionModel::from_env()?)).question(
202 /// "confirmed",
203 /// Decision::new(Question::boolean(
204 /// "In their last turn, did the caller agree to the booking that was read back?",
205 /// )),
206 /// );
207 /// Live::builder()
208 /// .decisions(decisions)
209 /// .converse(&booking) // .commit("book", Guard::decided("confirmed"))
210 /// ```
211 pub fn decisions(mut self, decisions: gemini_adk_rs::decision::Decisions) -> Self {
212 self.config = self
213 .config
214 .input_transcription(true)
215 .output_transcription(true);
216 self.decisions = Some(Arc::new(decisions));
217 self
218 }
219
220 /// Add a custom `TurnExtractor` implementation.
221 pub fn extractor(mut self, extractor: Arc<dyn TurnExtractor>) -> Self {
222 // Auto-enable transcription
223 self.config = self
224 .config
225 .input_transcription(true)
226 .output_transcription(true);
227 self.extractors.push(extractor);
228 self
229 }
230
231 /// Called when a TurnExtractor produces a result.
232 ///
233 /// The callback receives the extractor name and the extracted JSON value.
234 pub fn on_extracted<F, Fut>(mut self, f: F) -> Self
235 where
236 F: Fn(String, serde_json::Value) -> Fut + Send + Sync + 'static,
237 Fut: Future<Output = ()> + Send + 'static,
238 {
239 self.callbacks.on_extracted = Some(Arc::new(move |name, value| Box::pin(f(name, value))));
240 self
241 }
242
243 /// Called when a TurnExtractor fails.
244 ///
245 /// The callback receives the extractor name and error message.
246 /// Use this for custom error handling (alerting, retry logic, etc.).
247 pub fn on_extraction_error<F, Fut>(mut self, f: F) -> Self
248 where
249 F: Fn(String, String) -> Fut + Send + Sync + 'static,
250 Fut: Future<Output = ()> + Send + 'static,
251 {
252 self.callbacks.on_extraction_error =
253 Some(Arc::new(move |name, error| Box::pin(f(name, error))));
254 self
255 }
256}