Jack Amadeo
676 lines
21 KiB
Rust
676 lines
21 KiB
Rust
use anyhow::Result;
|
|
use futures::Stream;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use super::errors::ProviderError;
|
|
use super::retry::RetryConfig;
|
|
use crate::config::base::ConfigValue;
|
|
use crate::conversation::message::Message;
|
|
use crate::conversation::Conversation;
|
|
use crate::model::ModelConfig;
|
|
use crate::utils::safe_truncate;
|
|
use rmcp::model::Tool;
|
|
use utoipa::ToSchema;
|
|
|
|
use once_cell::sync::Lazy;
|
|
use std::ops::{Add, AddAssign};
|
|
use std::pin::Pin;
|
|
use std::sync::Mutex;
|
|
|
|
/// A global store for the current model being used, we use this as when a provider returns, it tells us the real model, not an alias
|
|
pub static CURRENT_MODEL: Lazy<Mutex<Option<String>>> = Lazy::new(|| Mutex::new(None));
|
|
|
|
/// Set the current model in the global store
|
|
pub fn set_current_model(model: &str) {
|
|
if let Ok(mut current_model) = CURRENT_MODEL.lock() {
|
|
*current_model = Some(model.to_string());
|
|
}
|
|
}
|
|
|
|
/// Get the current model from the global store, the real model, not an alias
|
|
pub fn get_current_model() -> Option<String> {
|
|
CURRENT_MODEL.lock().ok().and_then(|model| model.clone())
|
|
}
|
|
|
|
pub static MSG_COUNT_FOR_SESSION_NAME_GENERATION: usize = 3;
|
|
|
|
/// Information about a model's capabilities
|
|
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)]
|
|
pub struct ModelInfo {
|
|
/// The name of the model
|
|
pub name: String,
|
|
/// The maximum context length this model supports
|
|
pub context_limit: usize,
|
|
/// Cost per token for input (optional)
|
|
pub input_token_cost: Option<f64>,
|
|
/// Cost per token for output (optional)
|
|
pub output_token_cost: Option<f64>,
|
|
/// Currency for the costs (default: "$")
|
|
pub currency: Option<String>,
|
|
/// Whether this model supports cache control
|
|
pub supports_cache_control: Option<bool>,
|
|
}
|
|
|
|
impl ModelInfo {
|
|
/// Create a new ModelInfo with just name and context limit
|
|
pub fn new(name: impl Into<String>, context_limit: usize) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
context_limit,
|
|
input_token_cost: None,
|
|
output_token_cost: None,
|
|
currency: None,
|
|
supports_cache_control: None,
|
|
}
|
|
}
|
|
|
|
/// Create a new ModelInfo with cost information (per token)
|
|
pub fn with_cost(
|
|
name: impl Into<String>,
|
|
context_limit: usize,
|
|
input_cost: f64,
|
|
output_cost: f64,
|
|
) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
context_limit,
|
|
input_token_cost: Some(input_cost),
|
|
output_token_cost: Some(output_cost),
|
|
currency: Some("$".to_string()),
|
|
supports_cache_control: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
|
|
pub enum ProviderType {
|
|
Preferred,
|
|
Builtin,
|
|
Declarative,
|
|
Custom,
|
|
}
|
|
|
|
/// Metadata about a provider's configuration requirements and capabilities
|
|
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
|
pub struct ProviderMetadata {
|
|
/// The unique identifier for this provider
|
|
pub name: String,
|
|
/// Display name for the provider in UIs
|
|
pub display_name: String,
|
|
/// Description of the provider's capabilities
|
|
pub description: String,
|
|
/// The default/recommended model for this provider
|
|
pub default_model: String,
|
|
/// A list of currently known models with their capabilities
|
|
pub known_models: Vec<ModelInfo>,
|
|
/// Link to the docs where models can be found
|
|
pub model_doc_link: String,
|
|
/// Required configuration keys
|
|
pub config_keys: Vec<ConfigKey>,
|
|
}
|
|
|
|
impl ProviderMetadata {
|
|
pub fn new(
|
|
name: &str,
|
|
display_name: &str,
|
|
description: &str,
|
|
default_model: &str,
|
|
model_names: Vec<&str>,
|
|
model_doc_link: &str,
|
|
config_keys: Vec<ConfigKey>,
|
|
) -> Self {
|
|
Self {
|
|
name: name.to_string(),
|
|
display_name: display_name.to_string(),
|
|
description: description.to_string(),
|
|
default_model: default_model.to_string(),
|
|
known_models: model_names
|
|
.iter()
|
|
.map(|&name| ModelInfo {
|
|
name: name.to_string(),
|
|
context_limit: ModelConfig::new_or_fail(name).context_limit(),
|
|
input_token_cost: None,
|
|
output_token_cost: None,
|
|
currency: None,
|
|
supports_cache_control: None,
|
|
})
|
|
.collect(),
|
|
model_doc_link: model_doc_link.to_string(),
|
|
config_keys,
|
|
}
|
|
}
|
|
|
|
pub fn with_models(
|
|
name: &str,
|
|
display_name: &str,
|
|
description: &str,
|
|
default_model: &str,
|
|
models: Vec<ModelInfo>,
|
|
model_doc_link: &str,
|
|
config_keys: Vec<ConfigKey>,
|
|
) -> Self {
|
|
Self {
|
|
name: name.to_string(),
|
|
display_name: display_name.to_string(),
|
|
description: description.to_string(),
|
|
default_model: default_model.to_string(),
|
|
known_models: models,
|
|
model_doc_link: model_doc_link.to_string(),
|
|
config_keys,
|
|
}
|
|
}
|
|
|
|
pub fn empty() -> Self {
|
|
Self {
|
|
name: "".to_string(),
|
|
display_name: "".to_string(),
|
|
description: "".to_string(),
|
|
default_model: "".to_string(),
|
|
known_models: vec![],
|
|
model_doc_link: "".to_string(),
|
|
config_keys: vec![],
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration key metadata for provider setup
|
|
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
|
pub struct ConfigKey {
|
|
/// The name of the configuration key (e.g., "API_KEY")
|
|
pub name: String,
|
|
/// Whether this key is required for the provider to function
|
|
pub required: bool,
|
|
/// Whether this key should be stored securely (e.g., in keychain)
|
|
pub secret: bool,
|
|
/// Optional default value for the key
|
|
pub default: Option<String>,
|
|
/// Whether this key should be configured using OAuth device code flow
|
|
/// When true, the provider's configure_oauth() method will be called instead of prompting for manual input
|
|
pub oauth_flow: bool,
|
|
}
|
|
|
|
impl ConfigKey {
|
|
/// Create a new ConfigKey
|
|
pub fn new(name: &str, required: bool, secret: bool, default: Option<&str>) -> Self {
|
|
Self {
|
|
name: name.to_string(),
|
|
required,
|
|
secret,
|
|
default: default.map(|s| s.to_string()),
|
|
oauth_flow: false,
|
|
}
|
|
}
|
|
|
|
pub fn from_value_type<T: ConfigValue>(required: bool, secret: bool) -> Self {
|
|
Self {
|
|
name: T::KEY.to_string(),
|
|
required,
|
|
secret,
|
|
default: Some(T::DEFAULT.to_string()),
|
|
oauth_flow: false,
|
|
}
|
|
}
|
|
|
|
/// Create a new ConfigKey that uses OAuth device code flow for configuration
|
|
///
|
|
/// This is used for providers that support OAuth authentication instead of manual API key entry.
|
|
/// When oauth_flow is true, the configuration system will call the provider's configure_oauth() method.
|
|
pub fn new_oauth(name: &str, required: bool, secret: bool, default: Option<&str>) -> Self {
|
|
Self {
|
|
name: name.to_string(),
|
|
required,
|
|
secret,
|
|
default: default.map(|s| s.to_string()),
|
|
oauth_flow: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ProviderUsage {
|
|
pub model: String,
|
|
pub usage: Usage,
|
|
}
|
|
|
|
impl ProviderUsage {
|
|
pub fn new(model: String, usage: Usage) -> Self {
|
|
Self { model, usage }
|
|
}
|
|
|
|
/// Ensures this ProviderUsage has token counts, estimating them if necessary
|
|
pub async fn ensure_tokens(
|
|
&mut self,
|
|
system_prompt: &str,
|
|
request_messages: &[Message],
|
|
response: &Message,
|
|
tools: &[Tool],
|
|
) -> Result<(), ProviderError> {
|
|
crate::providers::usage_estimator::ensure_usage_tokens(
|
|
self,
|
|
system_prompt,
|
|
request_messages,
|
|
response,
|
|
tools,
|
|
)
|
|
.await
|
|
.map_err(|e| ProviderError::ExecutionError(format!("Failed to ensure usage tokens: {}", e)))
|
|
}
|
|
|
|
/// Combine this ProviderUsage with another, adding their token counts
|
|
/// Uses the model from this ProviderUsage
|
|
pub fn combine_with(&self, other: &ProviderUsage) -> ProviderUsage {
|
|
ProviderUsage {
|
|
model: self.model.clone(),
|
|
usage: self.usage + other.usage,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)]
|
|
pub struct Usage {
|
|
pub input_tokens: Option<i32>,
|
|
pub output_tokens: Option<i32>,
|
|
pub total_tokens: Option<i32>,
|
|
}
|
|
|
|
fn sum_optionals<T>(a: Option<T>, b: Option<T>) -> Option<T>
|
|
where
|
|
T: Add<Output = T> + Default,
|
|
{
|
|
match (a, b) {
|
|
(Some(x), Some(y)) => Some(x + y),
|
|
(Some(x), None) => Some(x + T::default()),
|
|
(None, Some(y)) => Some(T::default() + y),
|
|
(None, None) => None,
|
|
}
|
|
}
|
|
|
|
impl Add for Usage {
|
|
type Output = Self;
|
|
|
|
fn add(self, other: Self) -> Self {
|
|
Self::new(
|
|
sum_optionals(self.input_tokens, other.input_tokens),
|
|
sum_optionals(self.output_tokens, other.output_tokens),
|
|
sum_optionals(self.total_tokens, other.total_tokens),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl AddAssign for Usage {
|
|
fn add_assign(&mut self, rhs: Self) {
|
|
*self = *self + rhs;
|
|
}
|
|
}
|
|
|
|
impl Usage {
|
|
pub fn new(
|
|
input_tokens: Option<i32>,
|
|
output_tokens: Option<i32>,
|
|
total_tokens: Option<i32>,
|
|
) -> Self {
|
|
let calculated_total = if total_tokens.is_none() {
|
|
match (input_tokens, output_tokens) {
|
|
(Some(input), Some(output)) => Some(input + output),
|
|
(Some(input), None) => Some(input),
|
|
(None, Some(output)) => Some(output),
|
|
(None, None) => None,
|
|
}
|
|
} else {
|
|
total_tokens
|
|
};
|
|
|
|
Self {
|
|
input_tokens,
|
|
output_tokens,
|
|
total_tokens: calculated_total,
|
|
}
|
|
}
|
|
}
|
|
|
|
use async_trait::async_trait;
|
|
|
|
/// Trait for LeadWorkerProvider-specific functionality
|
|
pub trait LeadWorkerProviderTrait {
|
|
/// Get information about the lead and worker models for logging
|
|
fn get_model_info(&self) -> (String, String);
|
|
|
|
/// Get the currently active model name
|
|
fn get_active_model(&self) -> String;
|
|
}
|
|
|
|
/// Base trait for AI providers (OpenAI, Anthropic, etc)
|
|
#[async_trait]
|
|
pub trait Provider: Send + Sync {
|
|
/// Get the metadata for this provider type
|
|
fn metadata() -> ProviderMetadata
|
|
where
|
|
Self: Sized;
|
|
|
|
/// Get the name of this provider instance
|
|
fn get_name(&self) -> &str;
|
|
|
|
// Internal implementation of complete, used by complete_fast and complete
|
|
// Providers should override this to implement their actual completion logic
|
|
async fn complete_with_model(
|
|
&self,
|
|
model_config: &ModelConfig,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError>;
|
|
|
|
// Default implementation: use the provider's configured model
|
|
async fn complete(
|
|
&self,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
let model_config = self.get_model_config();
|
|
self.complete_with_model(&model_config, system, messages, tools)
|
|
.await
|
|
}
|
|
|
|
// Check if a fast model is configured, otherwise fall back to regular model
|
|
async fn complete_fast(
|
|
&self,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
let model_config = self.get_model_config();
|
|
let fast_config = model_config.use_fast_model();
|
|
|
|
match self
|
|
.complete_with_model(&fast_config, system, messages, tools)
|
|
.await
|
|
{
|
|
Ok(result) => Ok(result),
|
|
Err(e) => {
|
|
if fast_config.model_name != model_config.model_name {
|
|
tracing::warn!(
|
|
"Fast model {} failed with error: {}. Falling back to regular model {}",
|
|
fast_config.model_name,
|
|
e,
|
|
model_config.model_name
|
|
);
|
|
self.complete_with_model(&model_config, system, messages, tools)
|
|
.await
|
|
} else {
|
|
Err(e)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get the model config from the provider
|
|
fn get_model_config(&self) -> ModelConfig;
|
|
|
|
fn retry_config(&self) -> RetryConfig {
|
|
RetryConfig::default()
|
|
}
|
|
|
|
async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
|
|
Ok(None)
|
|
}
|
|
|
|
fn supports_embeddings(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
async fn supports_cache_control(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
/// Create embeddings if supported. Default implementation returns an error.
|
|
async fn create_embeddings(&self, _texts: Vec<String>) -> Result<Vec<Vec<f32>>, ProviderError> {
|
|
Err(ProviderError::ExecutionError(
|
|
"This provider does not support embeddings".to_string(),
|
|
))
|
|
}
|
|
|
|
/// Check if this provider is a LeadWorkerProvider
|
|
/// This is used for logging model information at startup
|
|
fn as_lead_worker(&self) -> Option<&dyn LeadWorkerProviderTrait> {
|
|
None
|
|
}
|
|
|
|
async fn stream(
|
|
&self,
|
|
_system: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
Err(ProviderError::NotImplemented(
|
|
"streaming not implemented".to_string(),
|
|
))
|
|
}
|
|
|
|
fn supports_streaming(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
/// Get the currently active model name
|
|
/// For regular providers, this returns the configured model
|
|
/// For LeadWorkerProvider, this returns the currently active model (lead or worker)
|
|
fn get_active_model_name(&self) -> String {
|
|
if let Some(lead_worker) = self.as_lead_worker() {
|
|
lead_worker.get_active_model()
|
|
} else {
|
|
self.get_model_config().model_name
|
|
}
|
|
}
|
|
|
|
/// Returns the first 3 user messages as strings for session naming
|
|
fn get_initial_user_messages(&self, messages: &Conversation) -> Vec<String> {
|
|
messages
|
|
.iter()
|
|
.filter(|m| m.role == rmcp::model::Role::User)
|
|
.take(MSG_COUNT_FOR_SESSION_NAME_GENERATION)
|
|
.map(|m| m.as_concat_text())
|
|
.collect()
|
|
}
|
|
|
|
/// Generate a session name/description based on the conversation history
|
|
/// Creates a prompt asking for a concise description in 4 words or less.
|
|
async fn generate_session_name(
|
|
&self,
|
|
messages: &Conversation,
|
|
) -> Result<String, ProviderError> {
|
|
let context = self.get_initial_user_messages(messages);
|
|
let prompt = self.create_session_name_prompt(&context);
|
|
let message = Message::user().with_text(&prompt);
|
|
let result = self
|
|
.complete_fast(
|
|
"Reply with only a description in four words or less",
|
|
&[message],
|
|
&[],
|
|
)
|
|
.await?;
|
|
|
|
let description = result
|
|
.0
|
|
.as_concat_text()
|
|
.split_whitespace()
|
|
.collect::<Vec<_>>()
|
|
.join(" ");
|
|
|
|
Ok(safe_truncate(&description, 100))
|
|
}
|
|
|
|
// Generate a prompt for a session name based on the conversation history
|
|
fn create_session_name_prompt(&self, context: &[String]) -> String {
|
|
// Create a prompt for a concise description
|
|
let mut prompt = "Based on the conversation so far, provide a concise description of this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the description".to_string();
|
|
|
|
if !context.is_empty() {
|
|
prompt = format!(
|
|
"Here are the first few user messages:\n{}\n\n{}",
|
|
context.join("\n"),
|
|
prompt
|
|
);
|
|
}
|
|
prompt
|
|
}
|
|
|
|
/// Configure OAuth authentication for this provider
|
|
///
|
|
/// This method is called when a provider has configuration keys marked with oauth_flow = true.
|
|
/// Providers that support OAuth should override this method to implement their specific OAuth flow.
|
|
///
|
|
/// # Returns
|
|
/// * `Ok(())` if OAuth configuration succeeds and credentials are saved
|
|
/// * `Err(ProviderError)` if OAuth fails or is not supported by this provider
|
|
///
|
|
/// # Default Implementation
|
|
/// The default implementation returns an error indicating OAuth is not supported.
|
|
async fn configure_oauth(&self) -> Result<(), ProviderError> {
|
|
Err(ProviderError::ExecutionError(
|
|
"OAuth configuration not supported by this provider".to_string(),
|
|
))
|
|
}
|
|
}
|
|
|
|
/// A message stream yields partial text content but complete tool calls, all within the Message object
|
|
/// So a message with text will contain potentially just a word of a longer response, but tool calls
|
|
/// messages will only be yielded once concatenated.
|
|
pub type MessageStream = Pin<
|
|
Box<dyn Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), ProviderError>> + Send>,
|
|
>;
|
|
|
|
pub fn stream_from_single_message(message: Message, usage: ProviderUsage) -> MessageStream {
|
|
let stream = futures::stream::once(async move { Ok((Some(message), Some(usage))) });
|
|
Box::pin(stream)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::collections::HashMap;
|
|
|
|
use serde_json::json;
|
|
#[test]
|
|
fn test_usage_creation() {
|
|
let usage = Usage::new(Some(10), Some(20), Some(30));
|
|
assert_eq!(usage.input_tokens, Some(10));
|
|
assert_eq!(usage.output_tokens, Some(20));
|
|
assert_eq!(usage.total_tokens, Some(30));
|
|
}
|
|
|
|
#[test]
|
|
fn test_usage_serialization() -> Result<()> {
|
|
let usage = Usage::new(Some(10), Some(20), Some(30));
|
|
let serialized = serde_json::to_string(&usage)?;
|
|
let deserialized: Usage = serde_json::from_str(&serialized)?;
|
|
|
|
assert_eq!(usage.input_tokens, deserialized.input_tokens);
|
|
assert_eq!(usage.output_tokens, deserialized.output_tokens);
|
|
assert_eq!(usage.total_tokens, deserialized.total_tokens);
|
|
|
|
// Test JSON structure
|
|
let json_value: serde_json::Value = serde_json::from_str(&serialized)?;
|
|
assert_eq!(json_value["input_tokens"], json!(10));
|
|
assert_eq!(json_value["output_tokens"], json!(20));
|
|
assert_eq!(json_value["total_tokens"], json!(30));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_set_and_get_current_model() {
|
|
// Set the model
|
|
set_current_model("gpt-4o");
|
|
|
|
// Get the model and verify
|
|
let model = get_current_model();
|
|
assert_eq!(model, Some("gpt-4o".to_string()));
|
|
|
|
// Change the model
|
|
set_current_model("claude-sonnet-4-20250514");
|
|
|
|
// Get the updated model and verify
|
|
let model = get_current_model();
|
|
assert_eq!(model, Some("claude-sonnet-4-20250514".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_provider_metadata_context_limits() {
|
|
// Test that ProviderMetadata::new correctly sets context limits
|
|
let test_models = vec!["gpt-4o", "claude-sonnet-4-20250514", "unknown-model"];
|
|
let metadata = ProviderMetadata::new(
|
|
"test",
|
|
"Test Provider",
|
|
"Test Description",
|
|
"gpt-4o",
|
|
test_models,
|
|
"https://example.com",
|
|
vec![],
|
|
);
|
|
|
|
let model_info: HashMap<String, usize> = metadata
|
|
.known_models
|
|
.into_iter()
|
|
.map(|m| (m.name, m.context_limit))
|
|
.collect();
|
|
|
|
// gpt-4o should have 128k limit
|
|
assert_eq!(*model_info.get("gpt-4o").unwrap(), 128_000);
|
|
|
|
// claude-sonnet-4-20250514 should have 200k limit
|
|
assert_eq!(
|
|
*model_info.get("claude-sonnet-4-20250514").unwrap(),
|
|
200_000
|
|
);
|
|
|
|
// unknown model should have default limit (128k)
|
|
assert_eq!(*model_info.get("unknown-model").unwrap(), 128_000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_creation() {
|
|
// Test direct ModelInfo creation
|
|
let info = ModelInfo {
|
|
name: "test-model".to_string(),
|
|
context_limit: 1000,
|
|
input_token_cost: None,
|
|
output_token_cost: None,
|
|
currency: None,
|
|
supports_cache_control: None,
|
|
};
|
|
assert_eq!(info.context_limit, 1000);
|
|
|
|
// Test equality
|
|
let info2 = ModelInfo {
|
|
name: "test-model".to_string(),
|
|
context_limit: 1000,
|
|
input_token_cost: None,
|
|
output_token_cost: None,
|
|
currency: None,
|
|
supports_cache_control: None,
|
|
};
|
|
assert_eq!(info, info2);
|
|
|
|
// Test inequality
|
|
let info3 = ModelInfo {
|
|
name: "test-model".to_string(),
|
|
context_limit: 2000,
|
|
input_token_cost: None,
|
|
output_token_cost: None,
|
|
currency: None,
|
|
supports_cache_control: None,
|
|
};
|
|
assert_ne!(info, info3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_with_cost() {
|
|
let info = ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001);
|
|
assert_eq!(info.name, "gpt-4o");
|
|
assert_eq!(info.context_limit, 128000);
|
|
assert_eq!(info.input_token_cost, Some(0.0000025));
|
|
assert_eq!(info.output_token_cost, Some(0.00001));
|
|
assert_eq!(info.currency, Some("$".to_string()));
|
|
}
|
|
}
|