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

gemini_adk_rs/live/
extractor.rs
