conversation_api/execution/
llm.rs1use crate::execution::{ExternalError, InvocationContext};
3use async_trait::async_trait;
4pub use llm_api::{
5 Completion, CompletionRequest, ContentPart, Continuation, FinishReason, INPUT_AUDIO,
6 INPUT_FILE, INPUT_IMAGE, INPUT_VIDEO, Message, MessageRole, ModelCapabilities,
7 ModelConstraints, ModelMode, ModelProfile, TokenUsage, ToolCall, UseCase,
8 input_modality_for_kind, payload_input_modalities,
9};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use thiserror::Error;
13
14#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
16pub struct ToolDefinition {
17 pub name: String,
19 pub description: String,
21 pub input_schema: Value,
23 pub policy: ToolPolicy,
25}
26
27#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
29#[serde(rename_all = "snake_case")]
30pub enum ToolEffect {
31 ReadOnly,
33 Mutating,
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
39#[serde(rename_all = "snake_case")]
40pub enum ApprovalRequirement {
41 Never,
43 WhenInteractive,
45 Always,
47}
48
49impl ApprovalRequirement {
50 #[must_use]
52 pub const fn rank(self) -> u8 {
53 match self {
54 Self::Never => 0,
55 Self::WhenInteractive => 1,
56 Self::Always => 2,
57 }
58 }
59}
60
61#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
63#[serde(tag = "type", rename_all = "snake_case")]
64pub enum ApprovalPolicy {
65 Fixed {
67 requirement: ApprovalRequirement,
69 },
70 PerInvocation {
72 minimum: ApprovalRequirement,
74 },
75}
76
77impl ApprovalPolicy {
78 #[must_use]
80 pub const fn permits(self, requirement: ApprovalRequirement) -> bool {
81 match self {
82 Self::Fixed { requirement: fixed } => fixed.rank() == requirement.rank(),
83 Self::PerInvocation { minimum } => requirement.rank() >= minimum.rank(),
84 }
85 }
86
87 #[must_use]
89 pub const fn is_fixed_never(self) -> bool {
90 matches!(
91 self,
92 Self::Fixed {
93 requirement: ApprovalRequirement::Never
94 }
95 )
96 }
97}
98
99#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
101pub struct ToolPolicy {
102 pub effect: ToolEffect,
104 pub approval: ApprovalPolicy,
106 pub parallel_safe: bool,
108}
109
110#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
112pub struct UsageSummary {
113 pub model_requests: u32,
115 pub reported_model_requests: u32,
117 pub input_tokens: u64,
119 pub output_tokens: u64,
121 pub cached_input_tokens: Option<u64>,
123 pub reasoning_output_tokens: Option<u64>,
125 pub peak_input_tokens: u64,
127 pub credits: Option<u64>,
129 pub tool_calls: u32,
131 pub window_truncations: u32,
133 pub trimmed_messages: u64,
135}
136
137impl UsageSummary {
138 pub fn record_model_attempt(&mut self) {
140 self.model_requests = self.model_requests.saturating_add(1);
141 }
142
143 pub fn record_usage(&mut self, usage: TokenUsage, reported: bool) {
145 if reported {
146 self.reported_model_requests = self.reported_model_requests.saturating_add(1);
147 }
148 self.input_tokens = self.input_tokens.saturating_add(usage.input_tokens);
149 self.output_tokens = self.output_tokens.saturating_add(usage.output_tokens);
150 self.peak_input_tokens = self.peak_input_tokens.max(usage.input_tokens);
151 add_optional(&mut self.cached_input_tokens, usage.cached_input_tokens);
152 add_optional(
153 &mut self.reasoning_output_tokens,
154 usage.reasoning_output_tokens,
155 );
156 add_optional(&mut self.credits, usage.credits);
157 }
158
159 pub fn record_model_request(&mut self, usage: Option<TokenUsage>) {
161 self.record_model_attempt();
162 let Some(usage) = usage else {
163 return;
164 };
165 self.record_usage(usage, true);
166 }
167
168 pub fn record_window_truncation(&mut self, removed_messages: usize) {
170 self.window_truncations = self.window_truncations.saturating_add(1);
171 self.trimmed_messages = self
172 .trimmed_messages
173 .saturating_add(u64::try_from(removed_messages).unwrap_or(u64::MAX));
174 }
175}
176
177fn add_optional(total: &mut Option<u64>, value: Option<u64>) {
178 if let Some(value) = value {
179 *total = Some(total.unwrap_or(0).saturating_add(value));
180 }
181}
182
183#[derive(Clone, Debug, Error, Eq, PartialEq)]
185#[error("{error}")]
186pub struct LlmFailure {
187 pub error: ExternalError,
189 pub usage: Option<TokenUsage>,
191}
192
193impl From<ExternalError> for LlmFailure {
194 fn from(error: ExternalError) -> Self {
195 Self { error, usage: None }
196 }
197}
198
199#[async_trait]
204pub trait Llm: Send + Sync {
205 async fn model_profile(
207 &self,
208 context: &InvocationContext,
209 use_case: &UseCase,
210 model_mode: &ModelMode,
211 ) -> Result<ModelProfile, ExternalError>;
212
213 async fn complete(
215 &self,
216 context: &InvocationContext,
217 request: CompletionRequest,
218 ) -> Result<Completion, LlmFailure>;
219}
220
221impl ToolDefinition {
222 pub fn model_definition(&self) -> llm_api::ToolDefinition {
223 llm_api::ToolDefinition {
224 name: self.name.clone(),
225 description: self.description.clone(),
226 input_schema: self.input_schema.clone(),
227 }
228 }
229}