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        // Generate JSON schema from the type
68        let root_schema = schemars::schema_for!(T);
69        let schema = serde_json::to_value(root_schema).unwrap_or(serde_json::Value::Null);
70
71        // Auto-register LLM for connection warming
72        self.warm_up_llms.push(llm.clone());
73
74        let extractor = LlmExtractor::new(name, llm, prompt, window_size)
75            .with_schema(schema)
76            .with_min_words(3);
77        self.extractors.push(Arc::new(extractor));
78        self
79    }
80
81    /// Like `extract_turns_windowed` but with a custom extraction trigger.
82    ///
83    /// Use `ExtractionTrigger::AfterToolCall` when tool calls are the primary
84    /// state source, `ExtractionTrigger::Interval(n)` to reduce extraction
85    /// frequency, or `ExtractionTrigger::OnPhaseChange` for phase-entry extraction.
86    pub fn extract_turns_triggered<T>(
87        mut self,
88        llm: Arc<dyn BaseLlm>,
89        prompt: impl Into<String>,
90        window_size: usize,
91        trigger: ExtractionTrigger,
92    ) -> Self
93    where
94        T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
95    {
96        // Auto-enable transcription
97        self.config = self
98            .config
99            .input_transcription(true)
100            .output_transcription(true);
101
102        let name = std::any::type_name::<T>()
103            .rsplit("::")
104            .next()
105            .unwrap_or("Extraction")
106            .to_string();
107
108        let root_schema = schemars::schema_for!(T);
109        let schema = serde_json::to_value(root_schema).unwrap_or(serde_json::Value::Null);
110
111        self.warm_up_llms.push(llm.clone());
112
113        let extractor = LlmExtractor::new(name, llm, prompt, window_size)
114            .with_schema(schema)
115            .with_min_words(3)
116            .with_trigger(trigger);
117        self.extractors.push(Arc::new(extractor));
118        self
119    }
120
121    /// Like [`extract_turns_triggered`](Self::extract_turns_triggered), but lets
122    /// callers configure the underlying [`LlmExtractor`] before registration.
123    ///
124    /// Use this for field promotion rules, custom minimum word counts, or other
125    /// extraction policies that should live at the SDK layer instead of app
126    /// callback glue.
127    pub fn extract_turns_configured<T>(
128        mut self,
129        llm: Arc<dyn BaseLlm>,
130        prompt: impl Into<String>,
131        window_size: usize,
132        trigger: ExtractionTrigger,
133        configure: impl FnOnce(LlmExtractor) -> LlmExtractor,
134    ) -> Self
135    where
136        T: DeserializeOwned + Serialize + schemars::JsonSchema + Send + Sync + 'static,
137    {
138        self.config = self
139            .config
140            .input_transcription(true)
141            .output_transcription(true);
142
143        let name = std::any::type_name::<T>()
144            .rsplit("::")
145            .next()
146            .unwrap_or("Extraction")
147            .to_string();
148
149        let root_schema = schemars::schema_for!(T);
150        let schema = serde_json::to_value(root_schema).unwrap_or(serde_json::Value::Null);
151
152        self.warm_up_llms.push(llm.clone());
153
154        let extractor = LlmExtractor::new(name, llm, prompt, window_size)
155            .with_schema(schema)
156            .with_min_words(3)
157            .with_trigger(trigger);
158        self.extractors.push(Arc::new(configure(extractor)));
159        self
160    }
161
162    /// Like [`extract_turns`](Self::extract_turns), but schema-as-data: no
163    /// Rust type required. The extraction result is stored in `State` under
164    /// `name`; pair with
165    /// [`FieldPromotion`](gemini_adk_rs::live::extractor::FieldPromotion)
166    /// rules via [`extractor`](Self::extractor) (or a
167    /// [`SessionSpec`](crate::spec::SessionSpec) `extract` entry, which wires
168    /// promotions declaratively) to land individual fields in the bare keys
169    /// flow guards read.
170    ///
171    /// This is the piece that lets a JSON-authored flow advance from speech
172    /// alone: extraction fills the state that `captured`/`is_true` guards
173    /// latch on, with no tool call anywhere.
174    pub fn extract_json(
175        mut self,
176        llm: Arc<dyn BaseLlm>,
177        name: impl Into<String>,
178        schema: serde_json::Value,
179        prompt: impl Into<String>,
180    ) -> Self {
181        self.config = self
182            .config
183            .input_transcription(true)
184            .output_transcription(true);
185        self.warm_up_llms.push(llm.clone());
186        let extractor = LlmExtractor::new(name.into(), llm, prompt.into(), 3)
187            .with_schema(schema)
188            .with_min_words(3);
189        self.extractors.push(Arc::new(extractor));
190        self
191    }
192
193    /// Add a custom `TurnExtractor` implementation.
194    pub fn extractor(mut self, extractor: Arc<dyn TurnExtractor>) -> Self {
195        // Auto-enable transcription
196        self.config = self
197            .config
198            .input_transcription(true)
199            .output_transcription(true);
200        self.extractors.push(extractor);
201        self
202    }
203
204    /// Called when a TurnExtractor produces a result.
205    ///
206    /// The callback receives the extractor name and the extracted JSON value.
207    pub fn on_extracted<F, Fut>(mut self, f: F) -> Self
208    where
209        F: Fn(String, serde_json::Value) -> Fut + Send + Sync + 'static,
210        Fut: Future<Output = ()> + Send + 'static,
211    {
212        self.callbacks.on_extracted = Some(Arc::new(move |name, value| Box::pin(f(name, value))));
213        self
214    }
215
216    /// Called when a TurnExtractor fails.
217    ///
218    /// The callback receives the extractor name and error message.
219    /// Use this for custom error handling (alerting, retry logic, etc.).
220    pub fn on_extraction_error<F, Fut>(mut self, f: F) -> Self
221    where
222        F: Fn(String, String) -> Fut + Send + Sync + 'static,
223        Fut: Future<Output = ()> + Send + 'static,
224    {
225        self.callbacks.on_extraction_error =
226            Some(Arc::new(move |name, error| Box::pin(f(name, error))));
227        self
228    }
229}
gemini_adk_fluent_rs/live/extraction.rs

gemini_adk_fluent_rs/live/
extraction.rs
