Rabi Mishra
438 lines
15 KiB
Rust
438 lines
15 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use serde_json::{json, Value};
|
|
|
|
use super::api_client::{ApiClient, AuthMethod};
|
|
use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage};
|
|
use super::errors::ProviderError;
|
|
use super::retry::ProviderRetry;
|
|
use super::utils::{
|
|
get_model, handle_response_openai_compat, handle_status_openai_compat, stream_openai_compat,
|
|
RequestLog,
|
|
};
|
|
use crate::conversation::message::Message;
|
|
|
|
use crate::model::ModelConfig;
|
|
use crate::providers::formats::openai::{create_request, get_usage};
|
|
use crate::providers::formats::openrouter as openrouter_format;
|
|
use rmcp::model::Tool;
|
|
|
|
pub const OPENROUTER_DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4";
|
|
pub const OPENROUTER_DEFAULT_FAST_MODEL: &str = "google/gemini-2.5-flash";
|
|
pub const OPENROUTER_MODEL_PREFIX_ANTHROPIC: &str = "anthropic";
|
|
|
|
// OpenRouter can run many models, we suggest the default
|
|
pub const OPENROUTER_KNOWN_MODELS: &[&str] = &[
|
|
"x-ai/grok-code-fast-1",
|
|
"anthropic/claude-sonnet-4.5",
|
|
"anthropic/claude-sonnet-4",
|
|
"anthropic/claude-opus-4.1",
|
|
"anthropic/claude-opus-4",
|
|
"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,
|
|
supports_streaming: bool,
|
|
#[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,
|
|
supports_streaming: true,
|
|
name: Self::metadata().name,
|
|
})
|
|
}
|
|
|
|
async fn post(
|
|
&self,
|
|
session_id: Option<&str>,
|
|
payload: &Value,
|
|
) -> Result<Value, ProviderError> {
|
|
let response = self
|
|
.api_client
|
|
.response_post(session_id, "api/v1/chat/completions", payload)
|
|
.await?;
|
|
|
|
let response_body = handle_response_openai_compat(response)
|
|
.await
|
|
.map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response: {e}")))?;
|
|
|
|
if let Some(error_obj) = response_body.get("error") {
|
|
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);
|
|
|
|
if error_code == 400 && error_message.contains("maximum context length") {
|
|
return Err(ProviderError::ContextLengthExceeded(
|
|
error_message.to_string(),
|
|
));
|
|
}
|
|
|
|
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())),
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
fn is_gemini_model(model_name: &str) -> bool {
|
|
model_name.starts_with("google/")
|
|
}
|
|
|
|
async fn create_request_based_on_model(
|
|
provider: &OpenRouterProvider,
|
|
session_id: Option<&str>,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<Value> {
|
|
let mut payload = create_request(
|
|
&provider.model,
|
|
system,
|
|
messages,
|
|
tools,
|
|
&super::utils::ImageFormat::OpenAi,
|
|
false,
|
|
)?;
|
|
|
|
if let Some(session_id) = session_id.filter(|id| !id.is_empty()) {
|
|
if let Some(obj) = payload.as_object_mut() {
|
|
obj.insert("user".to_string(), Value::String(session_id.to_string()));
|
|
}
|
|
}
|
|
|
|
if provider.supports_cache_control().await {
|
|
payload = update_request_for_anthropic(&payload);
|
|
}
|
|
|
|
if is_gemini_model(&provider.model.model_name) {
|
|
openrouter_format::add_reasoning_details_to_request(&mut payload, messages);
|
|
}
|
|
|
|
if let Some(obj) = payload.as_object_mut() {
|
|
obj.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"),
|
|
),
|
|
],
|
|
)
|
|
.with_unlisted_models()
|
|
}
|
|
|
|
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,
|
|
session_id: Option<&str>,
|
|
model_config: &ModelConfig,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
let payload =
|
|
create_request_based_on_model(self, session_id, system, messages, tools).await?;
|
|
let mut log = RequestLog::start(model_config, &payload)?;
|
|
|
|
let response = self
|
|
.with_retry(|| async {
|
|
let payload_clone = payload.clone();
|
|
self.post(session_id, &payload_clone).await
|
|
})
|
|
.await?;
|
|
|
|
let response_model = get_model(&response);
|
|
let message = if is_gemini_model(&self.model.model_name) {
|
|
openrouter_format::response_to_message(&response)?
|
|
} else {
|
|
crate::providers::formats::openai::response_to_message(&response)?
|
|
};
|
|
|
|
let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
|
|
tracing::debug!("Failed to get usage data");
|
|
Usage::default()
|
|
});
|
|
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
|
|
.request(None, "api/v1/models")
|
|
.response_get()
|
|
.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)
|
|
}
|
|
|
|
fn supports_streaming(&self) -> bool {
|
|
self.supports_streaming
|
|
}
|
|
|
|
async fn stream(
|
|
&self,
|
|
session_id: &str,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
let mut payload = create_request(
|
|
&self.model,
|
|
system,
|
|
messages,
|
|
tools,
|
|
&super::utils::ImageFormat::OpenAi,
|
|
true,
|
|
)?;
|
|
|
|
if self.supports_cache_control().await {
|
|
payload = update_request_for_anthropic(&payload);
|
|
}
|
|
|
|
if is_gemini_model(&self.model.model_name) {
|
|
openrouter_format::add_reasoning_details_to_request(&mut payload, messages);
|
|
}
|
|
|
|
if let Some(obj) = payload.as_object_mut() {
|
|
obj.insert("transforms".to_string(), json!(["middle-out"]));
|
|
}
|
|
|
|
let mut log = RequestLog::start(&self.model, &payload)?;
|
|
|
|
let response = self
|
|
.with_retry(|| async {
|
|
let resp = self
|
|
.api_client
|
|
.response_post(Some(session_id), "api/v1/chat/completions", &payload)
|
|
.await?;
|
|
handle_status_openai_compat(resp).await
|
|
})
|
|
.await
|
|
.inspect_err(|e| {
|
|
let _ = log.error(e);
|
|
})?;
|
|
|
|
stream_openai_compat(response, log)
|
|
}
|
|
}
|