1725 lines · 64.0 KB · Rust
RawDownload
use anyhow::{Result, anyhow};use collections::HashMap;use futures::{Stream, StreamExt};use language_model_core::{    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelImage,    LanguageModelRequest, LanguageModelRequestMessage, LanguageModelToolChoice,    LanguageModelToolResultContent, LanguageModelToolUse, LanguageModelToolUseId, MessageContent,    Role, StopReason, TokenUsage,    util::{fix_streamed_json, parse_tool_arguments},};use std::pin::Pin;use std::sync::Arc; use crate::responses::{    Request as ResponseRequest, ResponseFunctionCallItem, ResponseFunctionCallOutputContent,    ResponseFunctionCallOutputItem, ResponseInputContent, ResponseInputItem, ResponseMessageItem,    ResponseOutputItem, ResponseSummary as ResponsesSummary, ResponseUsage as ResponsesUsage,    StreamEvent as ResponsesStreamEvent,};use crate::{    FunctionContent, FunctionDefinition, ImageUrl, MessagePart, ReasoningEffort,    ResponseStreamEvent, ToolCall, ToolCallContent,}; pub fn into_open_ai(    request: LanguageModelRequest,    model_id: &str,    supports_parallel_tool_calls: bool,    supports_prompt_cache_key: bool,    max_output_tokens: Option<u64>,    reasoning_effort: Option<ReasoningEffort>,    interleaved_reasoning: bool,) -> crate::Request {    let stream = !model_id.starts_with("o1-");     let mut messages = Vec::new();    let mut current_reasoning: Option<String> = None;    for message in request.messages {        for content in message.content {            match content {                MessageContent::Thinking { text, .. } if interleaved_reasoning => {                    current_reasoning.get_or_insert_default().push_str(&text);                }                MessageContent::Text(text) | MessageContent::Thinking { text, .. } => {                    let should_add = if message.role == Role::User {                        // Including whitespace-only user messages can cause error with OpenAI compatible APIs                        // See https://github.com/zed-industries/zed/issues/40097                        !text.trim().is_empty()                    } else {                        !text.is_empty()                    };                    if should_add {                        add_message_content_part(                            MessagePart::Text { text },                            message.role,                            &mut messages,                        );                        if let Some(reasoning) = current_reasoning.take() {                            if let Some(crate::RequestMessage::Assistant {                                reasoning_content,                                ..                            }) = messages.last_mut()                            {                                *reasoning_content = Some(reasoning);                            }                        }                    }                }                MessageContent::RedactedThinking(_) => {}                MessageContent::Image(image) => {                    add_message_content_part(                        MessagePart::Image {                            image_url: ImageUrl {                                url: image.to_base64_url(),                                detail: None,                            },                        },                        message.role,                        &mut messages,                    );                }                MessageContent::ToolUse(tool_use) => {                    let tool_call = ToolCall {                        id: tool_use.id.to_string(),                        content: ToolCallContent::Function {                            function: FunctionContent {                                name: tool_use.name.to_string(),                                arguments: serde_json::to_string(&tool_use.input)                                    .unwrap_or_default(),                            },                        },                    };                     if let Some(crate::RequestMessage::Assistant { tool_calls, .. }) =                        messages.last_mut()                    {                        tool_calls.push(tool_call);                    } else {                        messages.push(crate::RequestMessage::Assistant {                            content: None,                            tool_calls: vec![tool_call],                            reasoning_content: current_reasoning.take(),                        });                    }                }                MessageContent::ToolResult(tool_result) => {                    let content: Vec<MessagePart> = tool_result                        .content                        .iter()                        .map(|part| match part {                            LanguageModelToolResultContent::Text(text) => MessagePart::Text {                                text: text.to_string(),                            },                            LanguageModelToolResultContent::Image(image) => MessagePart::Image {                                image_url: ImageUrl {                                    url: image.to_base64_url(),                                    detail: None,                                },                            },                        })                        .collect();                     messages.push(crate::RequestMessage::Tool {                        content: content.into(),                        tool_call_id: tool_result.tool_use_id.to_string(),                    });                }            }        }    }     crate::Request {        model: model_id.into(),        messages,        stream,        stream_options: if stream {            Some(crate::StreamOptions::default())        } else {            None        },        stop: request.stop,        temperature: request.temperature.or(Some(1.0)),        max_completion_tokens: max_output_tokens,        parallel_tool_calls: if supports_parallel_tool_calls && !request.tools.is_empty() {            Some(supports_parallel_tool_calls)        } else {            None        },        prompt_cache_key: if supports_prompt_cache_key {            request.thread_id        } else {            None        },        tools: request            .tools            .into_iter()            .map(|tool| crate::ToolDefinition::Function {                function: FunctionDefinition {                    name: tool.name,                    description: Some(tool.description),                    parameters: Some(tool.input_schema),                },            })            .collect(),        tool_choice: request.tool_choice.map(|choice| match choice {            LanguageModelToolChoice::Auto => crate::ToolChoice::Auto,            LanguageModelToolChoice::Any => crate::ToolChoice::Required,            LanguageModelToolChoice::None => crate::ToolChoice::None,        }),        reasoning_effort,    }} pub fn into_open_ai_response(    request: LanguageModelRequest,    model_id: &str,    supports_parallel_tool_calls: bool,    supports_prompt_cache_key: bool,    max_output_tokens: Option<u64>,    reasoning_effort: Option<ReasoningEffort>,) -> ResponseRequest {    let stream = !model_id.starts_with("o1-");     let LanguageModelRequest {        thread_id,        prompt_id: _,        intent: _,        messages,        tools,        tool_choice,        stop: _,        temperature,        thinking_allowed: _,        thinking_effort: _,        speed: _,    } = request;     let mut input_items = Vec::new();    for (index, message) in messages.into_iter().enumerate() {        append_message_to_response_items(message, index, &mut input_items);    }     let tools: Vec<_> = tools        .into_iter()        .map(|tool| crate::responses::ToolDefinition::Function {            name: tool.name,            description: Some(tool.description),            parameters: Some(tool.input_schema),            strict: None,        })        .collect();     ResponseRequest {        model: model_id.into(),        input: input_items,        stream,        temperature,        top_p: None,        max_output_tokens,        parallel_tool_calls: if tools.is_empty() {            None        } else {            Some(supports_parallel_tool_calls)        },        tool_choice: tool_choice.map(|choice| match choice {            LanguageModelToolChoice::Auto => crate::ToolChoice::Auto,            LanguageModelToolChoice::Any => crate::ToolChoice::Required,            LanguageModelToolChoice::None => crate::ToolChoice::None,        }),        tools,        prompt_cache_key: if supports_prompt_cache_key {            thread_id        } else {            None        },        reasoning: reasoning_effort.map(|effort| crate::responses::ReasoningConfig {            effort,            summary: Some(crate::responses::ReasoningSummaryMode::Auto),        }),    }} fn append_message_to_response_items(    message: LanguageModelRequestMessage,    index: usize,    input_items: &mut Vec<ResponseInputItem>,) {    let mut content_parts: Vec<ResponseInputContent> = Vec::new();     for content in message.content {        match content {            MessageContent::Text(text) => {                push_response_text_part(&message.role, text, &mut content_parts);            }            MessageContent::Thinking { text, .. } => {                push_response_text_part(&message.role, text, &mut content_parts);            }            MessageContent::RedactedThinking(_) => {}            MessageContent::Image(image) => {                push_response_image_part(&message.role, image, &mut content_parts);            }            MessageContent::ToolUse(tool_use) => {                flush_response_parts(&message.role, index, &mut content_parts, input_items);                let call_id = tool_use.id.to_string();                input_items.push(ResponseInputItem::FunctionCall(ResponseFunctionCallItem {                    call_id,                    name: tool_use.name.to_string(),                    arguments: tool_use.raw_input,                }));            }            MessageContent::ToolResult(tool_result) => {                flush_response_parts(&message.role, index, &mut content_parts, input_items);                let output = match tool_result.content.as_slice() {                    [LanguageModelToolResultContent::Text(text)] => {                        ResponseFunctionCallOutputContent::Text(text.to_string())                    }                    _ => {                        let parts = tool_result                            .content                            .into_iter()                            .map(|part| match part {                                LanguageModelToolResultContent::Text(text) => {                                    ResponseInputContent::Text {                                        text: text.to_string(),                                    }                                }                                LanguageModelToolResultContent::Image(image) => {                                    ResponseInputContent::Image {                                        image_url: image.to_base64_url(),                                    }                                }                            })                            .collect();                        ResponseFunctionCallOutputContent::List(parts)                    }                };                input_items.push(ResponseInputItem::FunctionCallOutput(                    ResponseFunctionCallOutputItem {                        call_id: tool_result.tool_use_id.to_string(),                        output,                    },                ));            }        }    }     flush_response_parts(&message.role, index, &mut content_parts, input_items);} fn push_response_text_part(    role: &Role,    text: impl Into<String>,    parts: &mut Vec<ResponseInputContent>,) {    let text = text.into();    if text.trim().is_empty() {        return;    }     match role {        Role::Assistant => parts.push(ResponseInputContent::OutputText {            text,            annotations: Vec::new(),        }),        _ => parts.push(ResponseInputContent::Text { text }),    }} fn push_response_image_part(    role: &Role,    image: LanguageModelImage,    parts: &mut Vec<ResponseInputContent>,) {    match role {        Role::Assistant => parts.push(ResponseInputContent::OutputText {            text: "[image omitted]".to_string(),            annotations: Vec::new(),        }),        _ => parts.push(ResponseInputContent::Image {            image_url: image.to_base64_url(),        }),    }} fn flush_response_parts(    role: &Role,    _index: usize,    parts: &mut Vec<ResponseInputContent>,    input_items: &mut Vec<ResponseInputItem>,) {    if parts.is_empty() {        return;    }     let item = ResponseInputItem::Message(ResponseMessageItem {        role: match role {            Role::User => crate::Role::User,            Role::Assistant => crate::Role::Assistant,            Role::System => crate::Role::System,        },        content: parts.clone(),    });     input_items.push(item);    parts.clear();} fn add_message_content_part(    new_part: MessagePart,    role: Role,    messages: &mut Vec<crate::RequestMessage>,) {    match (role, messages.last_mut()) {        (Role::User, Some(crate::RequestMessage::User { content }))        | (            Role::Assistant,            Some(crate::RequestMessage::Assistant {                content: Some(content),                ..            }),        )        | (Role::System, Some(crate::RequestMessage::System { content, .. })) => {            content.push_part(new_part);        }        _ => {            messages.push(match role {                Role::User => crate::RequestMessage::User {                    content: crate::MessageContent::from(vec![new_part]),                },                Role::Assistant => crate::RequestMessage::Assistant {                    content: Some(crate::MessageContent::from(vec![new_part])),                    tool_calls: Vec::new(),                    reasoning_content: None,                },                Role::System => crate::RequestMessage::System {                    content: crate::MessageContent::from(vec![new_part]),                },            });        }    }} pub struct OpenAiEventMapper {    tool_calls_by_index: HashMap<usize, RawToolCall>,} impl OpenAiEventMapper {    pub fn new() -> Self {        Self {            tool_calls_by_index: HashMap::default(),        }    }     pub fn map_stream(        mut self,        events: Pin<Box<dyn Send + Stream<Item = Result<ResponseStreamEvent>>>>,    ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>    {        events.flat_map(move |event| {            futures::stream::iter(match event {                Ok(event) => self.map_event(event),                Err(error) => vec![Err(LanguageModelCompletionError::from(anyhow!(error)))],            })        })    }     pub fn map_event(        &mut self,        event: ResponseStreamEvent,    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {        let mut events = Vec::new();        if let Some(usage) = event.usage {            events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {                input_tokens: usage.prompt_tokens,                output_tokens: usage.completion_tokens,                cache_creation_input_tokens: 0,                cache_read_input_tokens: 0,            })));        }         let Some(choice) = event.choices.first() else {            return events;        };         if let Some(delta) = choice.delta.as_ref() {            if let Some(reasoning_content) = delta.reasoning_content.clone() {                if !reasoning_content.is_empty() {                    events.push(Ok(LanguageModelCompletionEvent::Thinking {                        text: reasoning_content,                        signature: None,                    }));                }            }            if let Some(content) = delta.content.clone() {                if !content.is_empty() {                    events.push(Ok(LanguageModelCompletionEvent::Text(content)));                }            }             if let Some(tool_calls) = delta.tool_calls.as_ref() {                for tool_call in tool_calls {                    let entry = self.tool_calls_by_index.entry(tool_call.index).or_default();                     if let Some(tool_id) = tool_call.id.clone() {                        entry.id = tool_id;                    }                     if let Some(function) = tool_call.function.as_ref() {                        if let Some(name) = function.name.clone() {                            entry.name = name;                        }                         if let Some(arguments) = function.arguments.clone() {                            entry.arguments.push_str(&arguments);                        }                    }                     if !entry.id.is_empty() && !entry.name.is_empty() {                        if let Ok(input) = serde_json::from_str::<serde_json::Value>(                            &fix_streamed_json(&entry.arguments),                        ) {                            events.push(Ok(LanguageModelCompletionEvent::ToolUse(                                LanguageModelToolUse {                                    id: entry.id.clone().into(),                                    name: entry.name.as_str().into(),                                    is_input_complete: false,                                    input,                                    raw_input: entry.arguments.clone(),                                    thought_signature: None,                                },                            )));                        }                    }                }            }        }         match choice.finish_reason.as_deref() {            Some("stop") => {                events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));

Lines 1–500 of 1725.