gemini_adk_rs/live/
extractor.rs

1//! Turn-windowed extraction — OOB LLM structured data extraction between turns.
2//!
3//! A `TurnExtractor` runs after each turn completes, taking a window of recent
4//! transcript turns and producing a structured JSON value via an out-of-band
5//! LLM call.
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde_json::Value;
11
12use crate::llm::{BaseLlm, LlmError, LlmRequest};
13use crate::state::State;
14
15use super::phase::Phase;
16use super::transcript::TranscriptTurn;
17
18/// Controls WHEN an extractor runs.
19///
20/// The default is `EveryTurn`. Use `AfterToolCall` when tool calls are the primary state source,
21/// `Interval(n)` to reduce extraction frequency, or `OnPhaseChange`
22/// to extract only when entering a new conversation phase.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum ExtractionTrigger {
25    /// Run on every TurnComplete event (current default).
26    EveryTurn,
27    /// Run every N TurnComplete events.
28    Interval(u32),
29    /// Run after tool calls complete.
30    AfterToolCall,
31    /// Run when a phase transition occurs.
32    OnPhaseChange,
33    /// Run on GenerationComplete — before interruption truncation.
34    ///
35    /// Use this to extract from the model's full intended output, even if
36    /// the user barged in and the audio delivery was interrupted.
37    OnGenerationComplete,
38}
39
40/// How an extracted field should be merged into authoritative session state.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum MergePolicy {
43    /// Keep an existing state value; write only when the target key is absent.
44    KeepKnown,
45    /// Always overwrite the target state key with the extracted field value.
46    Overwrite,
47}
48
49/// Predicate used to decide whether an extracted field may be promoted.
50pub type PromotionPredicate = Arc<dyn Fn(&State, &Value) -> bool + Send + Sync>;
51
52/// Rule for promoting one raw extraction field into authoritative state.
53#[derive(Clone)]
54pub struct FieldPromotion {
55    /// Field name inside the extractor's JSON object.
56    pub field: String,
57    /// State key to write when the field is accepted.
58    pub state_key: String,
59    /// Merge behavior for the target state key.
60    pub merge: MergePolicy,
61    /// Optional acceptance predicate.
62    pub accept: Option<PromotionPredicate>,
63}
64
65impl FieldPromotion {
66    /// Promote `field` into the same state key using [`MergePolicy::KeepKnown`].
67    pub fn keep_known(field: impl Into<String>) -> Self {
68        let field = field.into();
69        Self {
70            state_key: field.clone(),
71            field,
72            merge: MergePolicy::KeepKnown,
73            accept: None,
74        }
75    }
76
77    /// Promote `field` into the same state key using [`MergePolicy::Overwrite`].
78    pub fn overwrite(field: impl Into<String>) -> Self {
79        let field = field.into();
80        Self {
81            state_key: field.clone(),
82            field,
83            merge: MergePolicy::Overwrite,
84            accept: None,
85        }
86    }
87
88    /// Promote a boolean field only when its extracted value is `true`.
89    pub fn true_only(field: impl Into<String>) -> Self {
90        Self::overwrite(field).accept_when(|_, value| value.as_bool() == Some(true))
91    }
92
93    /// Promote a string field only when its extracted value is non-empty.
94    pub fn non_empty(field: impl Into<String>) -> Self {
95        Self::overwrite(field)
96            .accept_when(|_, value| value.as_str().is_some_and(|s| !s.trim().is_empty()))
97    }
98
99    /// Promote into a custom target state key.
100    pub fn to(mut self, state_key: impl Into<String>) -> Self {
101        self.state_key = state_key.into();
102        self
103    }
104
105    /// Only accept this promotion when `predicate` returns true.
106    ///
107    /// This is the escape hatch for application-specific logic:
108    /// `FieldPromotion::overwrite("intent").accept_when(|state, value| ...)`.
109    pub fn accept_when(
110        mut self,
111        predicate: impl Fn(&State, &Value) -> bool + Send + Sync + 'static,
112    ) -> Self {
113        self.accept = Some(Arc::new(predicate));
114        self
115    }
116
117    /// Add an additional acceptance predicate, preserving any existing predicate.
118    pub fn and_accept_when(
119        mut self,
120        predicate: impl Fn(&State, &Value) -> bool + Send + Sync + 'static,
121    ) -> Self {
122        let previous = self.accept.take();
123        self.accept = Some(Arc::new(move |state, value| {
124            previous.as_ref().is_none_or(|accept| accept(state, value)) && predicate(state, value)
125        }));
126        self
127    }
128
129    /// Only promote after the named concept has been presented by a phase.
130    pub fn after_presented(self, concept: impl Into<String>) -> Self {
131        let concept = concept.into();
132        self.and_accept_when(move |state, _| Phase::is_presented(state, &concept))
133    }
134}
135
136/// Strip markdown code fences from LLM output.
137///
138/// Handles `` ```json\n...\n``` ``, `` ```\n...\n``` ``, and bare JSON.
139fn strip_code_fences(text: &str) -> &str {
140    let trimmed = text.trim();
141    if let Some(rest) = trimmed.strip_prefix("```") {
142        // Skip optional language tag (e.g., "json") on the first line
143        let rest = rest.trim_start_matches(|c: char| c != '\n');
144        let rest = rest.strip_prefix('\n').unwrap_or(rest);
145        // Strip trailing ```
146        let rest = rest.trim_end();
147        rest.strip_suffix("```").unwrap_or(rest).trim()
148    } else {
149        trimmed
150    }
151}
152
153/// Trait for between-turn extraction from transcript windows.
154///
155/// Implementations receive a window of recent transcript turns and produce
156/// a structured JSON value. The processor stores the result in `State`
157/// under the extractor's name.
158#[async_trait]
159pub trait TurnExtractor: Send + Sync {
160    /// Name of this extractor (used as the State key).
161    fn name(&self) -> &str;
162
163    /// How many recent turns this extractor needs.
164    fn window_size(&self) -> usize;
165
166    /// Whether this extractor should run for the current turn.
167    ///
168    /// Override to skip extraction on trivial turns (e.g., short utterances,
169    /// turns without user speech). Default returns `true` (always extract).
170    ///
171    /// This is checked before launching the async extraction, so returning
172    /// `false` avoids an LLM round-trip entirely.
173    fn should_extract(&self, window: &[TranscriptTurn]) -> bool {
174        let _ = window;
175        true
176    }
177
178    /// The trigger mode for this extractor.
179    ///
180    /// Controls when the extractor runs. Default is `EveryTurn`.
181    fn trigger(&self) -> ExtractionTrigger {
182        ExtractionTrigger::EveryTurn
183    }
184
185    /// Field promotion rules for this extractor.
186    ///
187    /// When empty, the runtime auto-flattens every top-level non-null field of
188    /// the extraction into state under the field's own name. When non-empty,
189    /// only these rules can promote raw extraction fields into authoritative
190    /// state.
191    fn promotion_rules(&self) -> &[FieldPromotion] {
192        &[]
193    }
194
195    /// Extract structured data from the transcript window.
196    async fn extract(&self, window: &[TranscriptTurn]) -> Result<Value, LlmError>;
197
198    /// Extract with access to session `State` — for extractors whose sources
199    /// bind arguments from `State` (e.g. async fetch/agent resolvers).
200    ///
201    /// The default delegates to [`extract`](Self::extract), so transcript-only
202    /// extractors need not implement it. The pipeline always calls this method.
203    async fn extract_with_state(
204        &self,
205        window: &[TranscriptTurn],
206        state: &State,
207    ) -> Result<Value, LlmError> {
208        let _ = state;
209        self.extract(window).await
210    }
211
212    /// An optional agent to run when this extractor's results land in state —
213    /// the `on_complete(dispatch(agent))` effect. Fired by the pipeline after
214    /// promotion, only when the extractor produced a non-empty object.
215    fn on_complete(&self) -> Option<OnComplete> {
216        None
217    }
218}
219
220/// A downstream agent fired when an extractor's results land in state.
221#[derive(Clone)]
222pub struct OnComplete {
223    /// The agent to run; it reads its inputs from `State`.
224    pub agent: Arc<dyn crate::text::TextAgent>,
225    /// How to run it (`Call` awaits inline; `Dispatch`/`Background` detached).
226    pub mode: crate::orchestration::AgentMode,
227}
228
229/// LLM-backed turn extractor that sends transcript windows to an OOB LLM
230/// with a structured extraction prompt.
231pub struct LlmExtractor {
232    name: String,
233    llm: Arc<dyn BaseLlm>,
234    prompt: String,
235    window_size: usize,
236    schema: Option<Value>,
237    /// Pre-rendered schema string (computed once at construction)
238    schema_str: Option<String>,
239    /// Minimum word count in the last user utterance to trigger extraction.
240    min_words: usize,
241    /// When this extractor should fire.
242    trigger: ExtractionTrigger,
243    /// Field promotion rules. Empty means every top-level non-null field is
244    /// auto-flattened into state under its own name.
245    promotion_rules: Vec<FieldPromotion>,
246}
247
248impl LlmExtractor {
249    /// Create a new LLM-backed extractor.
250    ///
251    /// - `name`: key for storing results in State
252    /// - `llm`: the out-of-band LLM to use for extraction
253    /// - `prompt`: system instruction describing what to extract
254    /// - `window_size`: how many recent turns to include
255    pub fn new(
256        name: impl Into<String>,
257        llm: Arc<dyn BaseLlm>,
258        prompt: impl Into<String>,
259        window_size: usize,
260    ) -> Self {
261        Self {
262            name: name.into(),
263            llm,
264            prompt: prompt.into(),
265            window_size,
266            schema: None,
267            schema_str: None,
268            min_words: 0,
269            trigger: ExtractionTrigger::EveryTurn,
270            promotion_rules: Vec::new(),
271        }
272    }
273
274    /// Set the minimum word count in the last user utterance to trigger extraction.
275    ///
276    /// Turns where the user said fewer than `n` words will skip the LLM call.
277    /// Useful for filtering out "uh huh", "ok", "yes" style responses.
278    pub fn with_min_words(mut self, n: usize) -> Self {
279        self.min_words = n;
280        self
281    }
282
283    /// Set a JSON Schema for structured output.
284    ///
285    /// When set, the schema is included in the prompt to guide the LLM
286    /// toward producing valid JSON matching the schema.
287    pub fn with_schema(mut self, schema: Value) -> Self {
288        self.schema_str = serde_json::to_string_pretty(&schema).ok();
289        self.schema = Some(schema);
290        self
291    }
292
293    /// Set the trigger mode for this extractor.
294    pub fn with_trigger(mut self, trigger: ExtractionTrigger) -> Self {
295        self.trigger = trigger;
296        self
297    }
298
299    /// Set explicit field promotion rules.
300    ///
301    /// Once promotion rules are present, top-level fields are no longer
302    /// automatically flattened into state; only accepted rules promote.
303    pub fn with_promotions(mut self, rules: Vec<FieldPromotion>) -> Self {
304        self.promotion_rules = rules;
305        self
306    }
307
308    /// Format transcript turns for the LLM prompt.
309    fn format_transcript(window: &[TranscriptTurn]) -> String {
310        let mut out = String::new();
311        for turn in window {
312            if !turn.user.is_empty() {
313                out.push_str("User: ");
314                out.push_str(turn.user.trim());
315                out.push('\n');
316            }
317            if !turn.model.is_empty() {
318                out.push_str("Assistant: ");
319                out.push_str(turn.model.trim());
320                out.push('\n');
321            }
322            out.push('\n');
323        }
324        out
325    }
326}
327
328#[async_trait]
329impl TurnExtractor for LlmExtractor {
330    fn name(&self) -> &str {
331        &self.name
332    }
333
334    fn window_size(&self) -> usize {
335        self.window_size
336    }
337
338    fn should_extract(&self, window: &[TranscriptTurn]) -> bool {
339        if self.min_words == 0 {
340            return true;
341        }
342        // Check the last user utterance
343        window
344            .iter()
345            .rev()
346            .find(|t| !t.user.is_empty())
347            .is_some_and(|t| t.user.split_whitespace().count() >= self.min_words)
348    }
349
350    fn trigger(&self) -> ExtractionTrigger {
351        self.trigger.clone()
352    }
353
354    fn promotion_rules(&self) -> &[FieldPromotion] {
355        &self.promotion_rules
356    }
357
358    async fn extract(&self, window: &[TranscriptTurn]) -> Result<Value, LlmError> {
359        let transcript = Self::format_transcript(window);
360
361        let mut request = LlmRequest::from_text(format!(
362            "Transcript:\n{transcript}\nExtract the requested information."
363        ));
364        request.system_instruction = Some(self.prompt.clone());
365
366        // Use native JSON mode when a schema is available — the API constrains
367        // the model to produce valid JSON matching the schema, eliminating
368        // markdown fences and malformed output.
369        if let Some(ref schema) = self.schema {
370            request.response_mime_type = Some("application/json".to_string());
371            request.response_json_schema = Some(schema.clone());
372        } else {
373            request.response_mime_type = Some("application/json".to_string());
374        }
375
376        let response = self.llm.generate(request).await?;
377        let text = response.text();
378
379        // Fallback: strip markdown code fences if the model still wraps output
380        let cleaned = strip_code_fences(&text);
381
382        serde_json::from_str(cleaned).map_err(|e| {
383            LlmError::Other(format!(
384                "Failed to parse extraction result as JSON: {e}. Raw: {text}"
385            ))
386        })
387    }
388}
389
390#[cfg(test)]
391mod tests {
392    use super::*;
393    use crate::llm::LlmResponse;
394    use gemini_genai_rs::prelude::{Content, Part, Role};
395    use std::time::Instant;
396
397    struct MockLlm {
398        response: String,
399    }
400
401    #[async_trait]
402    impl BaseLlm for MockLlm {
403        fn model_id(&self) -> &str {
404            "mock"
405        }
406        async fn generate(&self, _request: LlmRequest) -> Result<LlmResponse, LlmError> {
407            Ok(LlmResponse {
408                content: Content {
409                    role: Some(Role::Model),
410                    parts: vec![Part::Text {
411                        text: self.response.clone(),
412                    }],
413                },
414                finish_reason: Some("STOP".into()),
415                usage: None,
416            })
417        }
418    }
419
420    fn make_turns(pairs: &[(&str, &str)]) -> Vec<TranscriptTurn> {
421        pairs
422            .iter()
423            .enumerate()
424            .map(|(i, (user, model))| TranscriptTurn {
425                turn_number: i as u32,
426                user: user.to_string(),
427                model: model.to_string(),
428                tool_calls: Vec::new(),
429                timestamp: Instant::now(),
430            })
431            .collect()
432    }
433
434    #[tokio::test]
435    async fn llm_extractor_produces_json() {
436        let llm = Arc::new(MockLlm {
437            response: r#"{"phase": "ordering", "items": ["pizza"]}"#.to_string(),
438        });
439
440        let extractor = LlmExtractor::new("OrderState", llm, "Extract order state", 3);
441
442        let turns = make_turns(&[
443            ("I'd like a pizza", "Great! What size?"),
444            ("Large please", "Coming right up!"),
445        ]);
446
447        let result = extractor.extract(&turns).await.unwrap();
448        assert_eq!(result["phase"], "ordering");
449        assert_eq!(result["items"][0], "pizza");
450    }
451
452    #[tokio::test]
453    async fn llm_extractor_with_schema() {
454        let llm = Arc::new(MockLlm {
455            response: r#"{"sentiment": "positive", "score": 0.9}"#.to_string(),
456        });
457
458        let schema = serde_json::json!({
459            "type": "object",
460            "properties": {
461                "sentiment": {"type": "string", "enum": ["positive", "neutral", "negative"]},
462                "score": {"type": "number"}
463            }
464        });
465
466        let extractor =
467            LlmExtractor::new("Sentiment", llm, "Rate sentiment", 1).with_schema(schema);
468
469        let turns = make_turns(&[("This is great!", "Glad you think so!")]);
470        let result = extractor.extract(&turns).await.unwrap();
471        assert_eq!(result["sentiment"], "positive");
472    }
473
474    #[tokio::test]
475    async fn llm_extractor_invalid_json_returns_error() {
476        let llm = Arc::new(MockLlm {
477            response: "not json at all".to_string(),
478        });
479
480        let extractor = LlmExtractor::new("Bad", llm, "Extract", 1);
481        let turns = make_turns(&[("hi", "hello")]);
482        let result = extractor.extract(&turns).await;
483        assert!(result.is_err());
484    }
485
486    #[test]
487    fn format_transcript_readable() {
488        let turns = make_turns(&[("Hello", "Hi there!"), ("How are you?", "I'm doing well")]);
489        let formatted = LlmExtractor::format_transcript(&turns);
490        assert!(formatted.contains("User: Hello"));
491        assert!(formatted.contains("Assistant: Hi there!"));
492        assert!(formatted.contains("User: How are you?"));
493    }
494
495    #[tokio::test]
496    async fn llm_extractor_handles_markdown_fenced_json() {
497        let llm = Arc::new(MockLlm {
498            response: "```json\n{\"status\": \"ok\"}\n```".to_string(),
499        });
500
501        let extractor = LlmExtractor::new("Fenced", llm, "Extract", 1);
502        let turns = make_turns(&[("test", "reply")]);
503        let result = extractor.extract(&turns).await.unwrap();
504        assert_eq!(result["status"], "ok");
505    }
506
507    #[test]
508    fn strip_code_fences_variants() {
509        assert_eq!(super::strip_code_fences("```json\n{}\n```"), "{}");
510        assert_eq!(super::strip_code_fences("```\n{}\n```"), "{}");
511        assert_eq!(
512            super::strip_code_fences("  ```json\n{\"a\":1}\n```  "),
513            "{\"a\":1}"
514        );
515        assert_eq!(
516            super::strip_code_fences("{\"bare\":true}"),
517            "{\"bare\":true}"
518        );
519    }
520
521    #[test]
522    fn extractor_name_and_window_size() {
523        let llm = Arc::new(MockLlm {
524            response: "{}".to_string(),
525        });
526        let ext = LlmExtractor::new("TestExtractor", llm, "test", 5);
527        assert_eq!(ext.name(), "TestExtractor");
528        assert_eq!(ext.window_size(), 5);
529    }
530
531    #[test]
532    fn extractor_default_trigger_is_every_turn() {
533        let llm = Arc::new(MockLlm {
534            response: "{}".to_string(),
535        });
536        let ext = LlmExtractor::new("Test", llm, "test", 5);
537        assert_eq!(ext.trigger(), ExtractionTrigger::EveryTurn);
538    }
539
540    #[test]
541    fn extractor_with_trigger() {
542        let llm = Arc::new(MockLlm {
543            response: "{}".to_string(),
544        });
545        let ext = LlmExtractor::new("Test", llm, "test", 5)
546            .with_trigger(ExtractionTrigger::AfterToolCall);
547        assert_eq!(ext.trigger(), ExtractionTrigger::AfterToolCall);
548    }
549}
gemini_adk_rs/live/extractor.rs

gemini_adk_rs/live/
extractor.rs
