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

agent_air_runtime/client/
models.rs
