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