adk_core/
model.rs

1use crate::schema_adapter::{GenericSchemaAdapter, SchemaAdapter};
2use crate::{Result, types::Content};
3use async_trait::async_trait;
4use futures::stream::Stream;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::pin::Pin;
8
9/// A pinned, boxed stream of [`LlmResponse`] results from a model.
10pub type LlmResponseStream = Pin<Box<dyn Stream<Item = Result<LlmResponse>> + Send>>;
11
12/// The core trait for all LLM providers.
13///
14/// Implementations wrap a specific model API (Gemini, OpenAI, Anthropic, etc.)
15/// and produce a stream of responses for a given request.
16#[async_trait]
17pub trait Llm: Send + Sync {
18    /// Returns the model identifier (e.g., "gemini-2.5-flash").
19    fn name(&self) -> &str;
20    /// Generates content from the given request, optionally streaming.
21    async fn generate_content(&self, req: LlmRequest, stream: bool) -> Result<LlmResponseStream>;
22
23    /// Returns the schema adapter for this provider.
24    ///
25    /// The schema adapter normalizes raw JSON Schema from MCP tools into the
26    /// format accepted by this provider's function-calling API.
27    ///
28    /// Default implementation returns [`GenericSchemaAdapter`], which applies
29    /// safe transforms suitable for most providers. Override this method to
30    /// return a provider-specific adapter (e.g., `GeminiSchemaAdapter`,
31    /// `OpenAiStrictSchemaAdapter`).
32    fn schema_adapter(&self) -> &dyn SchemaAdapter {
33        &GenericSchemaAdapter
34    }
35
36    /// Returns `true` if this model is configured to use a server-managed
37    /// environment (e.g., Gemini Interactions API) where the provider owns
38    /// the tool-calling loop and filesystem.
39    ///
40    /// This is used at agent build time to detect conflicts with client-side
41    /// sandbox tools that would produce competing filesystems.
42    ///
43    /// Default implementation returns `false`. Override in providers that
44    /// support server-managed environments.
45    fn uses_interactions_api(&self) -> bool {
46        false
47    }
48}
49
50/// A request to an LLM provider.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct LlmRequest {
53    /// The model identifier to use for generation.
54    pub model: String,
55    /// The conversation contents (system, user, model messages).
56    pub contents: Vec<Content>,
57    /// Optional generation configuration (temperature, tokens, etc.).
58    pub config: Option<GenerateContentConfig>,
59    /// Tool declarations keyed by tool name.
60    #[serde(skip)]
61    pub tools: HashMap<String, serde_json::Value>,
62    /// Provider-neutral identifier of the previous response to continue from.
63    /// The agent populates this from the most recent event's `interaction_id`.
64    /// Maps to `previous_interaction_id` for the Gemini Interactions transport;
65    /// ignored by transports that do not support response chaining.
66    #[serde(skip_serializing_if = "Option::is_none", default)]
67    pub previous_response_id: Option<String>,
68}
69
70/// Configuration for LLM content generation.
71#[derive(Debug, Clone, Default, Serialize, Deserialize)]
72pub struct GenerateContentConfig {
73    /// Sampling temperature (0.0 = deterministic, higher = more random).
74    pub temperature: Option<f32>,
75    /// Nucleus sampling threshold.
76    pub top_p: Option<f32>,
77    /// Top-k sampling parameter.
78    pub top_k: Option<i32>,
79    /// Frequency penalty to reduce repetition.
80    #[serde(skip_serializing_if = "Option::is_none", default)]
81    pub frequency_penalty: Option<f32>,
82    /// Presence penalty to encourage topic diversity.
83    #[serde(skip_serializing_if = "Option::is_none", default)]
84    pub presence_penalty: Option<f32>,
85    /// Maximum number of output tokens to generate.
86    pub max_output_tokens: Option<i32>,
87    /// Random seed for reproducible generation.
88    #[serde(skip_serializing_if = "Option::is_none", default)]
89    pub seed: Option<i64>,
90    /// Number of top log probabilities to return per token.
91    #[serde(skip_serializing_if = "Option::is_none", default)]
92    pub top_logprobs: Option<u8>,
93    /// Sequences that stop generation when encountered.
94    #[serde(default, skip_serializing_if = "Vec::is_empty")]
95    pub stop_sequences: Vec<String>,
96    /// JSON Schema for structured output.
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub response_schema: Option<serde_json::Value>,
99
100    /// Optional cached content name for Gemini provider.
101    /// When set, the Gemini provider attaches this to the generation request.
102    #[serde(skip_serializing_if = "Option::is_none", default)]
103    pub cached_content: Option<String>,
104
105    /// Provider-specific request options keyed by provider namespace.
106    #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
107    pub extensions: serde_json::Map<String, serde_json::Value>,
108}
109
110/// A response from an LLM provider.
111#[derive(Debug, Clone, Default, Serialize, Deserialize)]
112pub struct LlmResponse {
113    /// The generated content (text, function calls, etc.).
114    pub content: Option<Content>,
115    /// Token usage statistics.
116    pub usage_metadata: Option<UsageMetadata>,
117    /// Reason the model stopped generating.
118    pub finish_reason: Option<FinishReason>,
119    /// Citation sources referenced in the response.
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub citation_metadata: Option<CitationMetadata>,
122    /// Whether this is a partial streaming chunk.
123    pub partial: bool,
124    /// Whether the model has finished its turn.
125    pub turn_complete: bool,
126    /// Whether the response was interrupted.
127    pub interrupted: bool,
128    /// Error code from the provider, if any.
129    pub error_code: Option<String>,
130    /// Error message from the provider, if any.
131    pub error_message: Option<String>,
132    /// Provider-specific metadata (e.g., response IDs, routing info).
133    ///
134    /// `LlmAgent` recognizes two optional boolean fields: `content_complete`
135    /// on a non-partial response replaces accumulated deltas with authoritative
136    /// content; `continue_turn` requests another bounded model iteration for
137    /// a provider-native tool pause without a client function call.
138    #[serde(skip_serializing_if = "Option::is_none", default)]
139    pub provider_metadata: Option<serde_json::Value>,
140    /// Interactions API: the server-assigned interaction id for this turn, used
141    /// to continue the conversation via `previous_interaction_id`. `None` for
142    /// the generateContent transport and non-Gemini providers.
143    #[serde(skip_serializing_if = "Option::is_none", default)]
144    pub interaction_id: Option<String>,
145}
146
147/// Trait for LLM providers that support prompt caching.
148///
149/// Providers implementing this trait can create and delete cached content
150/// resources, enabling automatic prompt caching lifecycle management by the
151/// runner. The runner stores an `Option<Arc<dyn CacheCapable>>` alongside the
152/// primary `Arc<dyn Llm>` and calls these methods when [`ContextCacheConfig`]
153/// is active.
154///
155/// # Example
156///
157/// ```rust,ignore
158/// use adk_core::CacheCapable;
159///
160/// let cache_name = model
161///     .create_cache("You are a helpful assistant.", &tools, 600)
162///     .await?;
163/// // ... use cache_name in generation requests ...
164/// model.delete_cache(&cache_name).await?;
165/// ```
166#[async_trait]
167pub trait CacheCapable: Send + Sync {
168    /// Create a cached content resource from the given system instruction,
169    /// tool definitions, and TTL.
170    ///
171    /// Returns the provider-specific cache name (e.g. `"cachedContents/abc123"`
172    /// for Gemini) that can be attached to subsequent generation requests via
173    /// [`GenerateContentConfig::cached_content`].
174    async fn create_cache(
175        &self,
176        system_instruction: &str,
177        tools: &HashMap<String, serde_json::Value>,
178        ttl_seconds: u32,
179    ) -> Result<String>;
180
181    /// Delete a previously created cached content resource by name.
182    async fn delete_cache(&self, name: &str) -> Result<()>;
183}
184
185/// Configuration for automatic prompt caching lifecycle management.
186///
187/// When set on runner configuration, the runner will automatically create and manage
188/// cached content resources for supported providers (currently Gemini).
189///
190/// # Example
191///
192/// ```rust
193/// use adk_core::ContextCacheConfig;
194///
195/// let config = ContextCacheConfig {
196///     min_tokens: 4096,
197///     ttl_seconds: 600,
198///     cache_intervals: 3,
199/// };
200/// assert_eq!(config.min_tokens, 4096);
201/// ```
202#[derive(Debug, Clone, Serialize, Deserialize)]
203pub struct ContextCacheConfig {
204    /// Minimum system instruction + tool token count to trigger caching.
205    /// Set to 0 to disable caching.
206    pub min_tokens: u32,
207
208    /// Cache time-to-live in seconds.
209    /// Set to 0 to disable caching.
210    pub ttl_seconds: u32,
211
212    /// Maximum number of LLM invocations before cache refresh.
213    /// After this many invocations, the runner creates a new cache
214    /// and deletes the old one.
215    pub cache_intervals: u32,
216}
217
218impl Default for ContextCacheConfig {
219    fn default() -> Self {
220        Self { min_tokens: 4096, ttl_seconds: 600, cache_intervals: 3 }
221    }
222}
223
224/// Token usage statistics from an LLM response.
225#[derive(Debug, Clone, Default, Serialize, Deserialize)]
226pub struct UsageMetadata {
227    /// Number of tokens in the prompt.
228    pub prompt_token_count: i32,
229    /// Number of tokens in the generated response.
230    pub candidates_token_count: i32,
231    /// Total tokens (prompt + response).
232    pub total_token_count: i32,
233
234    /// Tokens read from cache (Gemini/Anthropic).
235    #[serde(skip_serializing_if = "Option::is_none", default)]
236    pub cache_read_input_token_count: Option<i32>,
237
238    /// Tokens written to cache during this request.
239    #[serde(skip_serializing_if = "Option::is_none", default)]
240    pub cache_creation_input_token_count: Option<i32>,
241
242    /// Tokens used for thinking/reasoning (thinking models).
243    #[serde(skip_serializing_if = "Option::is_none", default)]
244    pub thinking_token_count: Option<i32>,
245
246    /// Audio input tokens (multimodal models).
247    #[serde(skip_serializing_if = "Option::is_none", default)]
248    pub audio_input_token_count: Option<i32>,
249
250    /// Audio output tokens (multimodal models).
251    #[serde(skip_serializing_if = "Option::is_none", default)]
252    pub audio_output_token_count: Option<i32>,
253
254    /// Estimated cost in USD for this request.
255    #[serde(skip_serializing_if = "Option::is_none", default)]
256    pub cost: Option<f64>,
257
258    /// Whether this request used a bring-your-own-key provider.
259    #[serde(skip_serializing_if = "Option::is_none", default)]
260    pub is_byok: Option<bool>,
261
262    /// Provider-native usage data retained when available, including fields
263    /// that are not yet projected into the normalized counters above.
264    #[serde(skip_serializing_if = "Option::is_none", default)]
265    pub provider_usage: Option<serde_json::Value>,
266}
267
268/// Citation metadata emitted by model providers for source attribution.
269#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
270#[serde(rename_all = "camelCase")]
271pub struct CitationMetadata {
272    /// The list of citation sources in this response.
273    #[serde(default)]
274    pub citation_sources: Vec<CitationSource>,
275}
276
277/// One citation source with optional offsets.
278#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
279#[serde(rename_all = "camelCase")]
280pub struct CitationSource {
281    /// URI of the cited source.
282    pub uri: Option<String>,
283    /// Title of the cited source.
284    pub title: Option<String>,
285    /// Start character index in the response text.
286    pub start_index: Option<i32>,
287    /// End character index in the response text.
288    pub end_index: Option<i32>,
289    /// License of the cited source.
290    pub license: Option<String>,
291    /// Publication date of the cited source.
292    pub publication_date: Option<String>,
293}
294
295/// Reason the model stopped generating content.
296#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
297pub enum FinishReason {
298    /// Natural stop (end of response).
299    Stop,
300    /// Hit the maximum token limit.
301    MaxTokens,
302    /// Content filtered for safety.
303    Safety,
304    /// Content blocked due to recitation/copyright.
305    Recitation,
306    /// Other/unknown reason.
307    Other,
308}
309
310impl LlmRequest {
311    /// Creates a new request with the given model and contents.
312    pub fn new(model: impl Into<String>, contents: Vec<Content>) -> Self {
313        Self {
314            model: model.into(),
315            contents,
316            config: None,
317            tools: HashMap::new(),
318            previous_response_id: None,
319        }
320    }
321
322    /// Set the response schema for structured output.
323    pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self {
324        let config = self.config.get_or_insert(GenerateContentConfig::default());
325        config.response_schema = Some(schema);
326        self
327    }
328
329    /// Set the generation config.
330    pub fn with_config(mut self, config: GenerateContentConfig) -> Self {
331        self.config = Some(config);
332        self
333    }
334}
335
336impl LlmResponse {
337    /// Creates a complete (non-streaming) response with the given content.
338    pub fn new(content: Content) -> Self {
339        Self {
340            content: Some(content),
341            usage_metadata: None,
342            finish_reason: Some(FinishReason::Stop),
343            citation_metadata: None,
344            partial: false,
345            turn_complete: true,
346            interrupted: false,
347            error_code: None,
348            error_message: None,
349            provider_metadata: None,
350            interaction_id: None,
351        }
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    #[test]
360    fn test_llm_request_creation() {
361        let req = LlmRequest::new("test-model", vec![]);
362        assert_eq!(req.model, "test-model");
363        assert!(req.contents.is_empty());
364    }
365
366    #[test]
367    fn test_llm_request_with_response_schema() {
368        let schema = serde_json::json!({
369            "type": "object",
370            "properties": {
371                "name": { "type": "string" }
372            }
373        });
374        let req = LlmRequest::new("test-model", vec![]).with_response_schema(schema.clone());
375
376        assert!(req.config.is_some());
377        let config = req.config.unwrap();
378        assert!(config.response_schema.is_some());
379        assert_eq!(config.response_schema.unwrap(), schema);
380    }
381
382    #[test]
383    fn test_llm_request_with_config() {
384        let config = GenerateContentConfig {
385            temperature: Some(0.7),
386            top_p: Some(0.9),
387            top_k: Some(40),
388            frequency_penalty: Some(0.2),
389            presence_penalty: Some(-0.3),
390            max_output_tokens: Some(1024),
391            seed: Some(42),
392            top_logprobs: Some(5),
393            stop_sequences: vec!["END".to_string()],
394            ..Default::default()
395        };
396        let req = LlmRequest::new("test-model", vec![]).with_config(config);
397
398        assert!(req.config.is_some());
399        let config = req.config.unwrap();
400        assert_eq!(config.temperature, Some(0.7));
401        assert_eq!(config.max_output_tokens, Some(1024));
402        assert_eq!(config.frequency_penalty, Some(0.2));
403        assert_eq!(config.presence_penalty, Some(-0.3));
404        assert_eq!(config.seed, Some(42));
405        assert_eq!(config.top_logprobs, Some(5));
406        assert_eq!(config.stop_sequences, vec!["END"]);
407    }
408
409    #[test]
410    fn test_llm_response_creation() {
411        let content = Content::new("assistant");
412        let resp = LlmResponse::new(content);
413        assert!(resp.content.is_some());
414        assert!(resp.turn_complete);
415        assert!(!resp.partial);
416        assert_eq!(resp.finish_reason, Some(FinishReason::Stop));
417        assert!(resp.citation_metadata.is_none());
418        assert!(resp.provider_metadata.is_none());
419    }
420
421    #[test]
422    fn test_llm_response_deserialize_without_citations() {
423        let json = serde_json::json!({
424            "content": {
425                "role": "model",
426                "parts": [{"text": "hello"}]
427            },
428            "partial": false,
429            "turn_complete": true,
430            "interrupted": false
431        });
432
433        let response: LlmResponse = serde_json::from_value(json).expect("should deserialize");
434        assert!(response.citation_metadata.is_none());
435    }
436
437    #[test]
438    fn test_llm_response_roundtrip_with_citations() {
439        let response = LlmResponse {
440            content: Some(Content::new("model").with_text("hello")),
441            usage_metadata: None,
442            finish_reason: Some(FinishReason::Stop),
443            citation_metadata: Some(CitationMetadata {
444                citation_sources: vec![CitationSource {
445                    uri: Some("https://example.com".to_string()),
446                    title: Some("Example".to_string()),
447                    start_index: Some(0),
448                    end_index: Some(5),
449                    license: None,
450                    publication_date: Some("2026-01-01T00:00:00Z".to_string()),
451                }],
452            }),
453            partial: false,
454            turn_complete: true,
455            interrupted: false,
456            error_code: None,
457            error_message: None,
458            provider_metadata: None,
459            interaction_id: None,
460        };
461
462        let encoded = serde_json::to_string(&response).expect("serialize");
463        let decoded: LlmResponse = serde_json::from_str(&encoded).expect("deserialize");
464        assert_eq!(decoded.citation_metadata, response.citation_metadata);
465    }
466
467    #[test]
468    fn test_generate_content_config_roundtrip_with_extensions() {
469        let mut extensions = serde_json::Map::new();
470        extensions.insert(
471            "openrouter".to_string(),
472            serde_json::json!({
473                "provider": {
474                    "zdr": true,
475                    "order": ["openai", "anthropic"]
476                },
477                "plugins": [
478                    { "id": "web", "enabled": true }
479                ]
480            }),
481        );
482
483        let config = GenerateContentConfig {
484            temperature: Some(0.4),
485            top_p: Some(0.8),
486            top_k: Some(12),
487            frequency_penalty: Some(0.1),
488            presence_penalty: Some(0.2),
489            max_output_tokens: Some(512),
490            seed: Some(7),
491            top_logprobs: Some(3),
492            stop_sequences: vec!["STOP".to_string(), "DONE".to_string()],
493            response_schema: Some(serde_json::json!({
494                "type": "object",
495                "properties": { "answer": { "type": "string" } },
496                "required": ["answer"]
497            })),
498            cached_content: Some("cachedContents/abc123".to_string()),
499            extensions,
500        };
501
502        let encoded = serde_json::to_string(&config).expect("serialize");
503        let decoded: GenerateContentConfig = serde_json::from_str(&encoded).expect("deserialize");
504
505        assert_eq!(decoded.temperature, config.temperature);
506        assert_eq!(decoded.top_p, config.top_p);
507        assert_eq!(decoded.top_k, config.top_k);
508        assert_eq!(decoded.frequency_penalty, config.frequency_penalty);
509        assert_eq!(decoded.presence_penalty, config.presence_penalty);
510        assert_eq!(decoded.max_output_tokens, config.max_output_tokens);
511        assert_eq!(decoded.seed, config.seed);
512        assert_eq!(decoded.top_logprobs, config.top_logprobs);
513        assert_eq!(decoded.stop_sequences, config.stop_sequences);
514        assert_eq!(decoded.response_schema, config.response_schema);
515        assert_eq!(decoded.cached_content, config.cached_content);
516        assert_eq!(decoded.extensions, config.extensions);
517    }
518
519    #[test]
520    fn test_llm_response_and_usage_roundtrip_with_provider_metadata() {
521        let response = LlmResponse {
522            content: Some(Content::new("model").with_text("hello")),
523            usage_metadata: Some(UsageMetadata {
524                prompt_token_count: 10,
525                candidates_token_count: 20,
526                total_token_count: 30,
527                cache_read_input_token_count: Some(5),
528                cache_creation_input_token_count: Some(2),
529                thinking_token_count: Some(3),
530                audio_input_token_count: Some(4),
531                audio_output_token_count: Some(6),
532                cost: Some(0.0125),
533                is_byok: Some(true),
534                provider_usage: Some(serde_json::json!({
535                    "server_tool_use": {
536                        "web_search_requests": 1
537                    },
538                    "prompt_tokens_details": {
539                        "video_tokens": 8
540                    }
541                })),
542            }),
543            finish_reason: Some(FinishReason::Stop),
544            citation_metadata: None,
545            partial: false,
546            turn_complete: true,
547            interrupted: false,
548            error_code: None,
549            error_message: None,
550            provider_metadata: Some(serde_json::json!({
551                "openrouter": {
552                    "responseId": "resp_123",
553                    "outputItems": 2
554                }
555            })),
556            interaction_id: None,
557        };
558
559        let encoded = serde_json::to_string(&response).expect("serialize");
560        let decoded: LlmResponse = serde_json::from_str(&encoded).expect("deserialize");
561
562        assert_eq!(decoded.provider_metadata, response.provider_metadata);
563        assert_eq!(
564            decoded.usage_metadata.as_ref().and_then(|u| u.cost),
565            response.usage_metadata.as_ref().and_then(|u| u.cost),
566        );
567        assert_eq!(
568            decoded.usage_metadata.as_ref().and_then(|u| u.is_byok),
569            response.usage_metadata.as_ref().and_then(|u| u.is_byok),
570        );
571        assert_eq!(
572            decoded.usage_metadata.as_ref().and_then(|u| u.provider_usage.clone()),
573            response.usage_metadata.as_ref().and_then(|u| u.provider_usage.clone()),
574        );
575    }
576
577    #[test]
578    fn test_finish_reason() {
579        assert_eq!(FinishReason::Stop, FinishReason::Stop);
580        assert_ne!(FinishReason::Stop, FinishReason::MaxTokens);
581    }
582}
adk_core/model.rs

adk_core/
model.rs
