Skip to main content

agent_core_runtime/client/
models.rs

1/// Message role in a conversation.
2#[derive(Debug, Clone, PartialEq)]
3pub enum Role {
4    /// System message providing instructions or context.
5    System,
6    /// User message containing requests or responses.
7    User,
8    /// Assistant message containing responses or tool use.
9    Assistant,
10}
11
12/// Image source types supported by LLM providers.
13#[derive(Debug, Clone, PartialEq)]
14pub enum ImageSource {
15    /// Base64-encoded image data with media type (e.g., "image/jpeg", "image/png").
16    Base64 {
17        media_type: String,
18        data: String,
19    },
20    /// URL reference to an image.
21    Url(String),
22}
23
24/// Tool invocation requested by the assistant.
25#[derive(Debug, Clone, PartialEq)]
26pub struct ToolUse {
27    /// Unique identifier for this tool use (used to match with ToolResult).
28    pub id: String,
29    /// Name of the tool to invoke.
30    pub name: String,
31    /// Input arguments as a JSON string.
32    pub input: String,
33}
34
35/// Result of a tool invocation, sent back to the assistant.
36#[derive(Debug, Clone, PartialEq)]
37pub struct ToolResult {
38    /// ID of the tool use this result corresponds to.
39    pub tool_use_id: String,
40    /// Result content (can be text or error message).
41    pub content: String,
42    /// Whether this result represents an error.
43    pub is_error: bool,
44}
45
46/// Content block types that can appear in messages.
47#[derive(Debug, Clone, PartialEq)]
48pub enum Content {
49    /// Plain text content.
50    Text(String),
51    /// Image content with source data.
52    Image(ImageSource),
53    /// Tool invocation from the assistant.
54    ToolUse(ToolUse),
55    /// Tool result from the user.
56    ToolResult(ToolResult),
57}
58
59/// Conversation message with role and content blocks.
60#[derive(Debug, Clone, PartialEq)]
61pub struct Message {
62    /// Role of the message sender.
63    pub role: Role,
64    /// Content blocks (text, images, tool use, tool results).
65    pub content: Vec<Content>,
66    /// Provider-specific response metadata (safety ratings, grounding, citations).
67    pub response_metadata: Option<ResponseMetadata>,
68}
69
70/// Provider-specific metadata returned with responses.
71#[derive(Debug, Clone, PartialEq, Default)]
72pub struct ResponseMetadata {
73    /// Content safety ratings (Gemini).
74    pub safety_ratings: Option<Vec<SafetyRating>>,
75    /// Grounding/citation metadata (Gemini).
76    pub grounding: Option<GroundingMetadata>,
77}
78
79/// Content safety rating from the model.
80#[derive(Debug, Clone, PartialEq)]
81pub struct SafetyRating {
82    /// Harm category (e.g., "HARM_CATEGORY_HARASSMENT").
83    pub category: String,
84    /// Probability level (e.g., "NEGLIGIBLE", "LOW", "MEDIUM", "HIGH").
85    pub probability: String,
86    /// Whether this category was blocked.
87    pub blocked: bool,
88}
89
90/// Grounding and citation metadata.
91#[derive(Debug, Clone, PartialEq, Default)]
92pub struct GroundingMetadata {
93    /// Web search queries used for grounding.
94    pub web_search_queries: Vec<String>,
95    /// Grounding chunks with source information.
96    pub grounding_chunks: Vec<GroundingChunk>,
97    /// Grounding supports linking content to sources.
98    pub grounding_supports: Vec<GroundingSupport>,
99}
100
101/// A source chunk used for grounding.
102#[derive(Debug, Clone, PartialEq)]
103pub struct GroundingChunk {
104    /// Source type (e.g., "web").
105    pub source_type: String,
106    /// URI of the source.
107    pub uri: Option<String>,
108    /// Title of the source.
109    pub title: Option<String>,
110}
111
112/// Links a segment of generated content to grounding sources.
113#[derive(Debug, Clone, PartialEq)]
114pub struct GroundingSupport {
115    /// Start index in the generated text.
116    pub start_index: usize,
117    /// End index in the generated text.
118    pub end_index: usize,
119    /// Indices into grounding_chunks that support this segment.
120    pub chunk_indices: Vec<usize>,
121    /// Confidence scores for each supporting chunk.
122    pub confidence_scores: Vec<f32>,
123}
124
125impl Message {
126    /// Create a new message with a single text content block.
127    pub fn new(role: Role, text: impl Into<String>) -> Self {
128        Self {
129            role,
130            content: vec![Content::Text(text.into())],
131            response_metadata: None,
132        }
133    }
134
135    /// Create a message with multiple content blocks.
136    pub fn with_content(role: Role, content: Vec<Content>) -> Self {
137        Self {
138            role,
139            content,
140            response_metadata: None,
141        }
142    }
143
144    /// Create a message with content and metadata.
145    pub fn with_metadata(role: Role, content: Vec<Content>, metadata: ResponseMetadata) -> Self {
146        Self {
147            role,
148            content,
149            response_metadata: Some(metadata),
150        }
151    }
152
153    /// Create a system message with text content.
154    pub fn system(text: impl Into<String>) -> Self {
155        Self::new(Role::System, text)
156    }
157
158    /// Create a user message with text content.
159    pub fn user(text: impl Into<String>) -> Self {
160        Self::new(Role::User, text)
161    }
162
163    /// Create an assistant message with text content.
164    pub fn assistant(text: impl Into<String>) -> Self {
165        Self::new(Role::Assistant, text)
166    }
167
168    /// Create a user message with a tool result.
169    pub fn tool_result(tool_use_id: impl Into<String>, content: impl Into<String>, is_error: bool) -> Self {
170        Self {
171            role: Role::User,
172            content: vec![Content::ToolResult(ToolResult {
173                tool_use_id: tool_use_id.into(),
174                content: content.into(),
175                is_error,
176            })],
177            response_metadata: None,
178        }
179    }
180}
181
182/// Tool definition for function calling.
183#[derive(Debug, Clone, PartialEq)]
184pub struct Tool {
185    /// Name of the tool.
186    pub name: String,
187    /// Description of what the tool does.
188    pub description: String,
189    /// JSON Schema for the tool's input parameters.
190    pub input_schema: String,
191}
192
193impl Tool {
194    /// Create a new tool definition.
195    pub fn new(
196        name: impl Into<String>,
197        description: impl Into<String>,
198        input_schema: impl Into<String>,
199    ) -> Self {
200        Self {
201            name: name.into(),
202            description: description.into(),
203            input_schema: input_schema.into(),
204        }
205    }
206}
207
208/// Controls how the model uses tools.
209#[derive(Debug, Clone, PartialEq)]
210pub enum ToolChoice {
211    /// Model decides whether to use tools.
212    Auto,
213    /// Model must use at least one tool.
214    Any,
215    /// Model must use the specified tool.
216    Tool(String),
217    /// Model cannot use any tools.
218    None,
219}
220
221impl Default for ToolChoice {
222    fn default() -> Self {
223        Self::Auto
224    }
225}
226
227/// Metadata for request tracking.
228#[derive(Debug, Clone, PartialEq, Default)]
229pub struct Metadata {
230    /// User identifier for tracking/billing.
231    pub user_id: Option<String>,
232}
233
234/// Options for LLM message requests.
235#[derive(Debug, Clone, Default)]
236pub struct MessageOptions {
237    /// Sampling temperature (0.0-1.0).
238    pub temperature: Option<f32>,
239    /// Maximum tokens to generate.
240    pub max_tokens: Option<u32>,
241    /// Model to use.
242    pub model: Option<String>,
243    /// Tools available for the model to use.
244    pub tools: Option<Vec<Tool>>,
245    /// How the model should use tools.
246    pub tool_choice: Option<ToolChoice>,
247    /// Custom stop sequences.
248    pub stop_sequences: Option<Vec<String>>,
249    /// Nucleus sampling parameter.
250    pub top_p: Option<f32>,
251    /// Top-K sampling parameter.
252    pub top_k: Option<u32>,
253    /// Request metadata.
254    pub metadata: Option<Metadata>,
255}
256
257// ============================================================================
258// Streaming Types
259// ============================================================================
260
261/// Events emitted during streaming responses.
262#[derive(Debug, Clone, PartialEq)]
263pub enum StreamEvent {
264    /// Stream started, contains message metadata.
265    MessageStart {
266        message_id: String,
267        model: String,
268    },
269    /// A content block is starting.
270    ContentBlockStart {
271        index: usize,
272        block_type: ContentBlockType,
273    },
274    /// Incremental text content.
275    TextDelta {
276        index: usize,
277        text: String,
278    },
279    /// Incremental JSON for tool input.
280    InputJsonDelta {
281        index: usize,
282        json: String,
283    },
284    /// A content block has finished.
285    ContentBlockStop {
286        index: usize,
287    },
288    /// Message-level updates (stop reason, usage).
289    MessageDelta {
290        stop_reason: Option<String>,
291        usage: Option<Usage>,
292    },
293    /// Stream has ended.
294    MessageStop,
295    /// Keep-alive ping.
296    Ping,
297}
298
299/// Type of content block in streaming.
300#[derive(Debug, Clone, PartialEq)]
301pub enum ContentBlockType {
302    Text,
303    ToolUse { id: String, name: String },
304}
305
306/// Token usage statistics.
307#[derive(Debug, Clone, PartialEq, Default)]
308pub struct Usage {
309    /// Input tokens used.
310    pub input_tokens: u32,
311    /// Output tokens generated.
312    pub output_tokens: u32,
313}