Files
tkmind_go/crates/goose/src/providers/openrouter.rs
T
2025-10-28 14:21:24 -04:00

375 lines
13 KiB
Rust

use anyhow::{Error, Result};
use async_trait::async_trait;
use serde_json::{json, Value};
use super::api_client::{ApiClient, AuthMethod};
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{
get_model, handle_response_google_compat, handle_response_openai_compat, is_google_model,
RequestLog,
};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::formats::openai::{create_request, get_usage, response_to_message};
use rmcp::model::Tool;
pub const OPENROUTER_DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4";
pub const OPENROUTER_DEFAULT_FAST_MODEL: &str = "google/gemini-flash-2.5";
pub const OPENROUTER_MODEL_PREFIX_ANTHROPIC: &str = "anthropic";
// OpenRouter can run many models, we suggest the default
pub const OPENROUTER_KNOWN_MODELS: &[&str] = &[
"anthropic/claude-sonnet-4.5",
"anthropic/claude-sonnet-4",
"anthropic/claude-opus-4.1",
"anthropic/claude-opus-4",
"anthropic/claude-3.7-sonnet",
"google/gemini-2.5-pro",
"google/gemini-2.5-flash",
"deepseek/deepseek-r1-0528",
"qwen/qwen3-coder",
"moonshotai/kimi-k2",
];
pub const OPENROUTER_DOC_URL: &str = "https://openrouter.ai/models";
#[derive(serde::Serialize)]
pub struct OpenRouterProvider {
#[serde(skip)]
api_client: ApiClient,
model: ModelConfig,
#[serde(skip)]
name: String,
}
impl OpenRouterProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let model = model.with_fast(OPENROUTER_DEFAULT_FAST_MODEL.to_string());
let config = crate::config::Config::global();
let api_key: String = config.get_secret("OPENROUTER_API_KEY")?;
let host: String = config
.get_param("OPENROUTER_HOST")
.unwrap_or_else(|_| "https://openrouter.ai".to_string());
let auth = AuthMethod::BearerToken(api_key);
let api_client = ApiClient::new(host, auth)?
.with_header("HTTP-Referer", "https://block.github.io/goose")?
.with_header("X-Title", "goose")?;
Ok(Self {
api_client,
model,
name: Self::metadata().name,
})
}
async fn post(&self, payload: &Value) -> Result<Value, ProviderError> {
let response = self
.api_client
.response_post("api/v1/chat/completions", payload)
.await?;
// Handle Google-compatible model responses differently
if is_google_model(payload) {
return handle_response_google_compat(response).await;
}
// For OpenAI-compatible models, parse the response body to JSON
let response_body = handle_response_openai_compat(response)
.await
.map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response: {e}")))?;
let _debug = format!(
"OpenRouter request with payload: {} and response: {}",
serde_json::to_string_pretty(payload).unwrap_or_else(|_| "Invalid JSON".to_string()),
serde_json::to_string_pretty(&response_body)
.unwrap_or_else(|_| "Invalid JSON".to_string())
);
// OpenRouter can return errors in 200 OK responses, so we have to check for errors explicitly
// https://openrouter.ai/docs/api-reference/errors
if let Some(error_obj) = response_body.get("error") {
// If there's an error object, extract the error message and code
let error_message = error_obj
.get("message")
.and_then(|m| m.as_str())
.unwrap_or("Unknown OpenRouter error");
let error_code = error_obj.get("code").and_then(|c| c.as_u64()).unwrap_or(0);
// Check for context length errors in the error message
if error_code == 400 && error_message.contains("maximum context length") {
return Err(ProviderError::ContextLengthExceeded(
error_message.to_string(),
));
}
// Return appropriate error based on the OpenRouter error code
match error_code {
401 | 403 => return Err(ProviderError::Authentication(error_message.to_string())),
429 => {
return Err(ProviderError::RateLimitExceeded {
details: error_message.to_string(),
retry_delay: None,
})
}
500 | 503 => return Err(ProviderError::ServerError(error_message.to_string())),
_ => return Err(ProviderError::RequestFailed(error_message.to_string())),
}
}
// No error detected, return the response body
Ok(response_body)
}
}
/// Update the request when using anthropic model.
/// For anthropic model, we can enable prompt caching to save cost. Since openrouter is the OpenAI compatible
/// endpoint, we need to modify the open ai request to have anthropic cache control field.
fn update_request_for_anthropic(original_payload: &Value) -> Value {
let mut payload = original_payload.clone();
if let Some(messages_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("messages"))
.and_then(|messages| messages.as_array_mut())
{
// Add "cache_control" to the last and second-to-last "user" messages.
// During each turn, we mark the final message with cache_control so the conversation can be
// incrementally cached. The second-to-last user message is also marked for caching with the
// cache_control parameter, so that this checkpoint can read from the previous cache.
let mut user_count = 0;
for message in messages_spec.iter_mut().rev() {
if message.get("role") == Some(&json!("user")) {
if let Some(content) = message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*content = json!([{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]);
}
}
user_count += 1;
if user_count >= 2 {
break;
}
}
}
// Update the system message to have cache_control field.
if let Some(system_message) = messages_spec
.iter_mut()
.find(|msg| msg.get("role") == Some(&json!("system")))
{
if let Some(content) = system_message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*system_message = json!({
"role": "system",
"content": [{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]
});
}
}
}
}
if let Some(tools_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("tools"))
.and_then(|tools| tools.as_array_mut())
{
// Add "cache_control" to the last tool spec, if any. This means that all tool definitions,
// will be cached as a single prefix.
if let Some(last_tool) = tools_spec.last_mut() {
if let Some(function) = last_tool.get_mut("function") {
function
.as_object_mut()
.unwrap()
.insert("cache_control".to_string(), json!({ "type": "ephemeral" }));
}
}
}
payload
}
async fn create_request_based_on_model(
provider: &OpenRouterProvider,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> anyhow::Result<Value, Error> {
let mut payload = create_request(
&provider.model,
system,
messages,
tools,
&super::utils::ImageFormat::OpenAi,
)?;
if provider.supports_cache_control().await {
payload = update_request_for_anthropic(&payload);
}
// Always add transforms: ["middle-out"] for OpenRouter to handle prompts > context size
payload
.as_object_mut()
.unwrap()
.insert("transforms".to_string(), json!(["middle-out"]));
Ok(payload)
}
#[async_trait]
impl Provider for OpenRouterProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::new(
"openrouter",
"OpenRouter",
"Router for many model providers",
OPENROUTER_DEFAULT_MODEL,
OPENROUTER_KNOWN_MODELS.to_vec(),
OPENROUTER_DOC_URL,
vec![
ConfigKey::new("OPENROUTER_API_KEY", true, true, None),
ConfigKey::new(
"OPENROUTER_HOST",
false,
false,
Some("https://openrouter.ai"),
),
],
)
}
fn get_name(&self) -> &str {
&self.name
}
fn get_model_config(&self) -> ModelConfig {
self.model.clone()
}
#[tracing::instrument(
skip(self, model_config, system, messages, tools),
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
)]
async fn complete_with_model(
&self,
model_config: &ModelConfig,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request_based_on_model(self, system, messages, tools).await?;
let mut log = RequestLog::start(model_config, &payload)?;
// Make request
let response = self
.with_retry(|| async {
let payload_clone = payload.clone();
self.post(&payload_clone).await
})
.await?;
// Parse response
let message = response_to_message(&response)?;
let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
tracing::debug!("Failed to get usage data");
Usage::default()
});
let response_model = get_model(&response);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
/// Fetch supported models from OpenRouter API (only models with tool support)
async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
// Handle request failures gracefully
// If the request fails, fall back to manual entry
let response = match self.api_client.response_get("api/v1/models").await {
Ok(response) => response,
Err(e) => {
tracing::warn!("Failed to fetch models from OpenRouter API: {}, falling back to manual model entry", e);
return Ok(None);
}
};
// Handle JSON parsing failures gracefully
let json: serde_json::Value = match response.json().await {
Ok(json) => json,
Err(e) => {
tracing::warn!("Failed to parse OpenRouter API response as JSON: {}, falling back to manual model entry", e);
return Ok(None);
}
};
// Check for error in response
if let Some(err_obj) = json.get("error") {
let msg = err_obj
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("unknown error");
tracing::warn!("OpenRouter API returned an error: {}", msg);
return Ok(None);
}
let data = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
ProviderError::UsageError("Missing data field in JSON response".into())
})?;
let mut models: Vec<String> = data
.iter()
.filter_map(|model| {
// Get the model ID
let id = model.get("id").and_then(|v| v.as_str())?;
// Check if the model supports tools
let supported_params =
match model.get("supported_parameters").and_then(|v| v.as_array()) {
Some(params) => params,
None => {
// If supported_parameters is missing, skip this model (assume no tool support)
tracing::debug!(
"Model '{}' missing supported_parameters field, skipping",
id
);
return None;
}
};
let has_tool_support = supported_params
.iter()
.any(|param| param.as_str() == Some("tools"));
if has_tool_support {
Some(id.to_string())
} else {
None
}
})
.collect();
// If no models with tool support were found, fall back to manual entry
if models.is_empty() {
tracing::warn!("No models with tool support found in OpenRouter API response, falling back to manual model entry");
return Ok(None);
}
models.sort();
Ok(Some(models))
}
async fn supports_cache_control(&self) -> bool {
self.model
.model_name
.starts_with(OPENROUTER_MODEL_PREFIX_ANTHROPIC)
}
}