d18bb6e512
Signed-off-by: DaeHee Lee <lee111dae11@proton.me>
304 lines
10 KiB
Rust
304 lines
10 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use futures::future::BoxFuture;
|
|
use serde_json::{json, Value};
|
|
|
|
use super::api_client::{ApiClient, AuthMethod};
|
|
use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata};
|
|
use super::errors::ProviderError;
|
|
use super::openai_compatible::{handle_status, stream_openai_compat};
|
|
use super::retry::ProviderRetry;
|
|
use super::utils::{ImageFormat, RequestLog};
|
|
use crate::conversation::message::Message;
|
|
use crate::model::ModelConfig;
|
|
use crate::providers::formats::openai::create_request;
|
|
use crate::providers::formats::openrouter as openrouter_format;
|
|
use rmcp::model::Tool;
|
|
|
|
const OPENROUTER_PROVIDER_NAME: &str = "openrouter";
|
|
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, OPENROUTER_PROVIDER_NAME)?;
|
|
|
|
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://goose-docs.ai")?
|
|
.with_header("X-Title", "goose")?;
|
|
|
|
Ok(Self {
|
|
api_client,
|
|
model,
|
|
supports_streaming: true,
|
|
name: OPENROUTER_PROVIDER_NAME.to_string(),
|
|
})
|
|
}
|
|
}
|
|
|
|
/// 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/")
|
|
}
|
|
|
|
impl ProviderDef for OpenRouterProvider {
|
|
type Provider = Self;
|
|
|
|
fn metadata() -> ProviderMetadata {
|
|
ProviderMetadata::new(
|
|
OPENROUTER_PROVIDER_NAME,
|
|
"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, true),
|
|
ConfigKey::new(
|
|
"OPENROUTER_HOST",
|
|
false,
|
|
false,
|
|
Some("https://openrouter.ai"),
|
|
false,
|
|
),
|
|
],
|
|
)
|
|
.with_setup_steps(vec![
|
|
"Go to https://openrouter.ai/settings/keys",
|
|
"Click 'Create' or use an existing API key",
|
|
"Copy the key and paste it above",
|
|
])
|
|
}
|
|
|
|
fn from_env(
|
|
model: ModelConfig,
|
|
_extensions: Vec<crate::config::ExtensionConfig>,
|
|
) -> BoxFuture<'static, Result<Self::Provider>> {
|
|
Box::pin(Self::from_env(model))
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for OpenRouterProvider {
|
|
fn get_name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model.clone()
|
|
}
|
|
|
|
/// Fetch supported models from OpenRouter API (only models with tool support)
|
|
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
|
let response = self
|
|
.api_client
|
|
.request(None, "api/v1/models")
|
|
.response_get()
|
|
.await
|
|
.map_err(|e| {
|
|
ProviderError::RequestFailed(format!(
|
|
"Failed to fetch models from OpenRouter API: {}",
|
|
e
|
|
))
|
|
})?;
|
|
|
|
let json: serde_json::Value = response.json().await.map_err(|e| {
|
|
ProviderError::RequestFailed(format!(
|
|
"Failed to parse OpenRouter API response as JSON: {}",
|
|
e
|
|
))
|
|
})?;
|
|
|
|
if let Some(err_obj) = json.get("error") {
|
|
let msg = err_obj
|
|
.get("message")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("unknown error");
|
|
return Err(ProviderError::RequestFailed(format!(
|
|
"OpenRouter API returned an error: {}",
|
|
msg
|
|
)));
|
|
}
|
|
|
|
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| {
|
|
let id = model.get("id").and_then(|v| v.as_str())?;
|
|
Some(id.to_string())
|
|
})
|
|
.collect();
|
|
|
|
models.sort();
|
|
Ok(models)
|
|
}
|
|
|
|
async fn supports_cache_control(&self) -> bool {
|
|
self.model
|
|
.model_name
|
|
.starts_with(OPENROUTER_MODEL_PREFIX_ANTHROPIC)
|
|
}
|
|
|
|
async fn stream(
|
|
&self,
|
|
model_config: &ModelConfig,
|
|
session_id: &str,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
let mut payload = create_request(
|
|
model_config,
|
|
system,
|
|
messages,
|
|
tools,
|
|
&ImageFormat::OpenAi,
|
|
true,
|
|
)?;
|
|
|
|
// Add user field for OpenRouter attribution/rate-limiting
|
|
if !session_id.is_empty() {
|
|
if let Some(obj) = payload.as_object_mut() {
|
|
obj.insert("user".to_string(), Value::String(session_id.to_string()));
|
|
}
|
|
}
|
|
|
|
if self.supports_cache_control().await {
|
|
payload = update_request_for_anthropic(&payload);
|
|
}
|
|
|
|
if is_gemini_model(&model_config.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(model_config, &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(resp).await
|
|
})
|
|
.await
|
|
.inspect_err(|e| {
|
|
let _ = log.error(e);
|
|
})?;
|
|
|
|
stream_openai_compat(response, log)
|
|
}
|
|
}
|