Skip to main content

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}