853 lines · 26.9 KB · Rust
RawDownload
use anyhow::{Result, anyhow};use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};use http_client::{    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, http,};use serde::{Deserialize, Serialize};use serde_json::Value;pub use settings::DataCollection;pub use settings::ModelMode;pub use settings::OpenRouterAvailableModel as AvailableModel;pub use settings::OpenRouterProvider as Provider;use std::{convert::TryFrom, io, time::Duration};use strum::EnumString;use thiserror::Error; pub const OPEN_ROUTER_API_URL: &str = "https://openrouter.ai/api/v1"; fn extract_retry_after(headers: &http::HeaderMap) -> Option<std::time::Duration> {    if let Some(reset) = headers.get("X-RateLimit-Reset") {        if let Ok(s) = reset.to_str() {            if let Ok(epoch_ms) = s.parse::<u64>() {                let now = std::time::SystemTime::now()                    .duration_since(std::time::UNIX_EPOCH)                    .unwrap_or_default()                    .as_millis() as u64;                if epoch_ms > now {                    return Some(std::time::Duration::from_millis(epoch_ms - now));                }            }        }    }    None} fn is_none_or_empty<T: AsRef<[U]>, U>(opt: &Option<T>) -> bool {    opt.as_ref().is_none_or(|v| v.as_ref().is_empty())} #[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]#[serde(rename_all = "lowercase")]pub enum Role {    User,    Assistant,    System,    Tool,} impl TryFrom<String> for Role {    type Error = anyhow::Error;     fn try_from(value: String) -> Result<Self> {        match value.as_str() {            "user" => Ok(Self::User),            "assistant" => Ok(Self::Assistant),            "system" => Ok(Self::System),            "tool" => Ok(Self::Tool),            _ => Err(anyhow!("invalid role '{value}'")),        }    }} impl From<Role> for String {    fn from(val: Role) -> Self {        match val {            Role::User => "user".to_owned(),            Role::Assistant => "assistant".to_owned(),            Role::System => "system".to_owned(),            Role::Tool => "tool".to_owned(),        }    }} #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]pub struct Model {    pub name: String,    pub display_name: Option<String>,    pub max_tokens: u64,    pub supports_tools: Option<bool>,    pub supports_images: Option<bool>,    #[serde(default)]    pub mode: ModelMode,    pub provider: Option<Provider>,} impl Model {    pub fn default() -> Self {        Self::new(            "openrouter/auto",            Some("Auto Router"),            Some(2000000),            Some(true),            Some(false),            Some(ModelMode::Default),            None,        )    }     pub fn new(        name: &str,        display_name: Option<&str>,        max_tokens: Option<u64>,        supports_tools: Option<bool>,        supports_images: Option<bool>,        mode: Option<ModelMode>,        provider: Option<Provider>,    ) -> Self {        Self {            name: name.to_owned(),            display_name: display_name.map(|s| s.to_owned()),            max_tokens: max_tokens.unwrap_or(2000000),            supports_tools,            supports_images,            mode: mode.unwrap_or(ModelMode::Default),            provider,        }    }     pub fn id(&self) -> &str {        &self.name    }     pub fn display_name(&self) -> &str {        self.display_name.as_ref().unwrap_or(&self.name)    }     pub fn max_token_count(&self) -> u64 {        self.max_tokens    }     pub fn max_output_tokens(&self) -> Option<u64> {        None    }     pub fn supports_tool_calls(&self) -> bool {        self.supports_tools.unwrap_or(false)    }     pub fn supports_parallel_tool_calls(&self) -> bool {        false    }} #[derive(Debug, Serialize, Deserialize)]pub struct Request {    pub model: String,    pub messages: Vec<RequestMessage>,    pub stream: bool,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub session_id: Option<String>,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub max_tokens: Option<u64>,    #[serde(default, skip_serializing_if = "Vec::is_empty")]    pub stop: Vec<String>,    pub temperature: f32,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub tool_choice: Option<ToolChoice>,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub parallel_tool_calls: Option<bool>,    #[serde(default, skip_serializing_if = "Vec::is_empty")]    pub tools: Vec<ToolDefinition>,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub reasoning: Option<Reasoning>,    pub usage: RequestUsage,    pub provider: Option<Provider>,} #[derive(Debug, Default, Serialize, Deserialize)]pub struct RequestUsage {    pub include: bool,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum ToolChoice {    Auto,    Required,    None,    #[serde(untagged)]    Other(ToolDefinition),} #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]#[derive(Clone, Deserialize, Serialize, Debug)]#[serde(tag = "type", rename_all = "snake_case")]pub enum ToolDefinition {    #[allow(dead_code)]    Function { function: FunctionDefinition },} #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]#[derive(Clone, Debug, Serialize, Deserialize)]pub struct FunctionDefinition {    pub name: String,    pub description: Option<String>,    pub parameters: Option<Value>,} #[derive(Debug, Serialize, Deserialize)]pub struct Reasoning {    #[serde(skip_serializing_if = "Option::is_none")]    pub effort: Option<String>,    #[serde(skip_serializing_if = "Option::is_none")]    pub max_tokens: Option<u32>,    #[serde(skip_serializing_if = "Option::is_none")]    pub exclude: Option<bool>,    #[serde(skip_serializing_if = "Option::is_none")]    pub enabled: Option<bool>,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]#[serde(tag = "role", rename_all = "lowercase")]pub enum RequestMessage {    Assistant {        content: Option<MessageContent>,        #[serde(default, skip_serializing_if = "Vec::is_empty")]        tool_calls: Vec<ToolCall>,        #[serde(default, skip_serializing_if = "Option::is_none")]        reasoning_details: Option<std::sync::Arc<serde_json::Value>>,    },    User {        content: MessageContent,    },    System {        content: MessageContent,    },    Tool {        content: MessageContent,        tool_call_id: String,    },} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]#[serde(untagged)]pub enum MessageContent {    Plain(String),    Multipart(Vec<MessagePart>),} impl MessageContent {    pub fn empty() -> Self {        Self::Plain(String::new())    }     pub fn push_part(&mut self, part: MessagePart) {        match self {            Self::Plain(text) if text.is_empty() => {                *self = Self::Multipart(vec![part]);            }            Self::Plain(text) => {                let text_part = MessagePart::Text {                    text: std::mem::take(text),                    cache_control: None,                };                *self = Self::Multipart(vec![text_part, part]);            }            Self::Multipart(parts) => parts.push(part),        }    }} impl From<Vec<MessagePart>> for MessageContent {    fn from(parts: Vec<MessagePart>) -> Self {        if parts.len() == 1            && let MessagePart::Text {                text,                cache_control,            } = &parts[0]            && cache_control.is_none()        {            return Self::Plain(text.clone());        }        Self::Multipart(parts)    }} impl From<String> for MessageContent {    fn from(text: String) -> Self {        Self::Plain(text)    }} impl From<&str> for MessageContent {    fn from(text: &str) -> Self {        Self::Plain(text.to_string())    }} impl MessageContent {    pub fn as_text(&self) -> Option<&str> {        match self {            Self::Plain(text) => Some(text),            Self::Multipart(parts) if parts.len() == 1 => {                if let MessagePart::Text { text, .. } = &parts[0] {                    Some(text)                } else {                    None                }            }            _ => None,        }    }     pub fn to_text(&self) -> String {        match self {            Self::Plain(text) => text.clone(),            Self::Multipart(parts) => parts                .iter()                .filter_map(|part| {                    if let MessagePart::Text { text, .. } = part {                        Some(text.as_str())                    } else {                        None                    }                })                .collect::<Vec<_>>()                .concat(),        }    }} #[derive(Debug, Serialize, Deserialize, Copy, Clone, Eq, PartialEq)]#[serde(rename_all = "lowercase")]pub enum CacheControlType {    Ephemeral,} #[derive(Debug, Serialize, Deserialize, Copy, Clone, Eq, PartialEq)]pub enum CacheTtl {    /// Anthropic's default ephemeral TTL (currently 5 minutes). Refreshes for    /// free on every cache hit.    #[serde(rename = "5m")]    FiveMinutes,    /// Anthropic's extended ephemeral TTL (currently 1 hour). Costs 2x base    /// input tokens to write, but persists across longer idle gaps.    #[serde(rename = "1h")]    OneHour,} #[derive(Debug, Serialize, Deserialize, Copy, Clone, Eq, PartialEq)]pub struct CacheControl {    #[serde(rename = "type")]    pub cache_type: CacheControlType,    /// Omitting this field uses the API's default 5-minute TTL.    #[serde(default, skip_serializing_if = "Option::is_none")]    pub ttl: Option<CacheTtl>,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]#[serde(tag = "type", rename_all = "snake_case")]pub enum MessagePart {    Text {        text: String,        #[serde(default, skip_serializing_if = "Option::is_none")]        cache_control: Option<CacheControl>,    },    #[serde(rename = "image_url")]    Image { image_url: String },} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]pub struct ToolCall {    pub id: String,    #[serde(flatten)]    pub content: ToolCallContent,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]#[serde(tag = "type", rename_all = "lowercase")]pub enum ToolCallContent {    Function { function: FunctionContent },} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]pub struct FunctionContent {    pub name: String,    pub arguments: String,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub thought_signature: Option<String>,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]pub struct ResponseMessageDelta {    pub role: Option<Role>,    pub content: Option<String>,    pub reasoning: Option<String>,    #[serde(default, skip_serializing_if = "is_none_or_empty")]    pub tool_calls: Option<Vec<ToolCallChunk>>,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub reasoning_details: Option<serde_json::Value>,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]pub struct ToolCallChunk {    pub index: usize,    pub id: Option<String>,    pub function: Option<FunctionChunk>,} #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]pub struct FunctionChunk {    pub name: Option<String>,    pub arguments: Option<String>,    #[serde(default)]    pub thought_signature: Option<String>,} #[derive(Serialize, Deserialize, Debug, Default)]pub struct PromptTokensDetails {    #[serde(default)]    pub cached_tokens: u64,    #[serde(default)]    pub cache_write_tokens: u64,} #[derive(Serialize, Deserialize, Debug)]pub struct Usage {    pub prompt_tokens: u64,    pub completion_tokens: u64,    pub total_tokens: u64,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub prompt_tokens_details: Option<PromptTokensDetails>,} #[derive(Serialize, Deserialize, Debug)]pub struct ChoiceDelta {    pub index: u32,    pub delta: ResponseMessageDelta,    pub finish_reason: Option<String>,} #[derive(Serialize, Deserialize, Debug)]pub struct ResponseStreamEvent {    #[serde(default, skip_serializing_if = "Option::is_none")]    pub id: Option<String>,    pub created: u32,    pub model: String,    pub choices: Vec<ChoiceDelta>,    pub usage: Option<Usage>,} #[derive(Serialize, Deserialize, Debug)]pub struct Response {    pub id: String,    pub object: String,    pub created: u64,    pub model: String,    pub choices: Vec<Choice>,    pub usage: Usage,} #[derive(Serialize, Deserialize, Debug)]pub struct Choice {    pub index: u32,    pub message: RequestMessage,    pub finish_reason: Option<String>,} #[derive(Default, Debug, Clone, PartialEq, Deserialize)]pub struct ListModelsResponse {    pub data: Vec<ModelEntry>,} #[derive(Default, Debug, Clone, PartialEq, Deserialize)]pub struct ModelEntry {    pub id: String,    pub name: String,    pub created: usize,    pub description: String,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub context_length: Option<u64>,    #[serde(default, skip_serializing_if = "Vec::is_empty")]    pub supported_parameters: Vec<String>,    #[serde(default, skip_serializing_if = "Option::is_none")]    pub architecture: Option<ModelArchitecture>,} #[derive(Default, Debug, Clone, PartialEq, Deserialize)]pub struct ModelArchitecture {    #[serde(default, skip_serializing_if = "Vec::is_empty")]    pub input_modalities: Vec<String>,} pub async fn stream_completion(    client: &dyn HttpClient,    api_url: &str,    api_key: &str,    request: Request,    extra_headers: &CustomHeaders,) -> Result<BoxStream<'static, Result<ResponseStreamEvent, OpenRouterError>>, OpenRouterError> {    let uri = format!("{api_url}/chat/completions");    let request = HttpRequest::builder()        .method(Method::POST)        .uri(uri)        .header("Content-Type", "application/json")        .header("Authorization", format!("Bearer {}", api_key))        .header("HTTP-Referer", "https://zed.dev")        .header("X-Title", "Zed Editor")        .extra_headers(extra_headers)        .body(AsyncBody::from(

Lines 1–500 of 853.