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    /// Add a custom `TurnExtractor` implementation.
190    pub fn extractor(mut self, extractor: Arc<dyn TurnExtractor>) -> Self {
191        // Auto-enable transcription
192        self.config = self
193            .config
194            .input_transcription(true)
195            .output_transcription(true);
196        self.extractors.push(extractor);
197        self
198    }
199
200    /// Called when a TurnExtractor produces a result.
201    ///
202    /// The callback receives the extractor name and the extracted JSON value.
203    pub fn on_extracted<F, Fut>(mut self, f: F) -> Self
204    where
205        F: Fn(String, serde_json::Value) -> Fut + Send + Sync + 'static,
206        Fut: Future<Output = ()> + Send + 'static,
207    {
208        self.callbacks.on_extracted = Some(Arc::new(move |name, value| Box::pin(f(name, value))));
209        self
210    }
211
212    /// Called when a TurnExtractor fails.
213    ///
214    /// The callback receives the extractor name and error message.
215    /// Use this for custom error handling (alerting, retry logic, etc.).
216    pub fn on_extraction_error<F, Fut>(mut self, f: F) -> Self
217    where
218        F: Fn(String, String) -> Fut + Send + Sync + 'static,
219        Fut: Future<Output = ()> + Send + 'static,
220    {
221        self.callbacks.on_extraction_error =
222            Some(Arc::new(move |name, error| Box::pin(f(name, error))));
223        self
224    }
225}