use anyhow::{Context as _, Result, anyhow};use collections::{BTreeMap, HashMap};use credentials_provider::CredentialsProvider;use editor::{Editor, EditorElement, EditorStyle};use futures::Stream;use futures::{FutureExt, StreamExt, future::BoxFuture};use gpui::{ AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,};use http_client::HttpClient;use language_model::{ AuthenticateError, LanguageModel, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason,};use open_ai::{Model, ResponseStreamEvent, stream_completion};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use settings::{Settings, SettingsStore};use std::pin::Pin;use std::str::FromStr as _;use std::sync::Arc;use strum::IntoEnumIterator;use theme::ThemeSettings;use ui::{Icon, IconName, List, Tooltip, prelude::*};use util::{ResultExt, maybe}; use crate::{AllLanguageModelSettings, ui::InstructionListItem}; const PROVIDER_ID: &str = "openai";const PROVIDER_NAME: &str = "OpenAI"; #[derive(Default, Clone, Debug, PartialEq)]pub struct OpenAiSettings { pub api_url: String, pub available_models: Vec<AvailableModel>, pub needs_setting_migration: bool,} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]pub struct AvailableModel { pub name: String, pub display_name: Option<String>, pub max_tokens: usize, pub max_output_tokens: Option<u32>, pub max_completion_tokens: Option<u32>,} pub struct OpenAiLanguageModelProvider { http_client: Arc<dyn HttpClient>, state: gpui::Entity<State>,} pub struct State { api_key: Option<String>, api_key_from_env: bool, _subscription: Subscription,} const OPENAI_API_KEY_VAR: &str = "OPENAI_API_KEY"; impl State { fn is_authenticated(&self) -> bool { self.api_key.is_some() } fn reset_api_key(&self, cx: &mut Context<Self>) -> Task<Result<()>> { let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { credentials_provider .delete_credentials(&api_url, &cx) .await .log_err(); this.update(cx, |this, cx| { this.api_key = None; this.api_key_from_env = false; cx.notify(); }) }) } fn set_api_key(&mut self, api_key: String, cx: &mut Context<Self>) -> Task<Result<()>> { let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { credentials_provider .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx) .await .log_err(); this.update(cx, |this, cx| { this.api_key = Some(api_key); cx.notify(); }) }) } fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> { if self.is_authenticated() { return Task::ready(Ok(())); } let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { let (api_key, from_env) = if let Ok(api_key) = std::env::var(OPENAI_API_KEY_VAR) { (api_key, true) } else { let (_, api_key) = credentials_provider .read_credentials(&api_url, &cx) .await? .ok_or(AuthenticateError::CredentialsNotFound)?; ( String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?, false, ) }; this.update(cx, |this, cx| { this.api_key = Some(api_key); this.api_key_from_env = from_env; cx.notify(); })?; Ok(()) }) }} impl OpenAiLanguageModelProvider { pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self { let state = cx.new(|cx| State { api_key: None, api_key_from_env: false, _subscription: cx.observe_global::<SettingsStore>(|_this: &mut State, cx| { cx.notify(); }), }); Self { http_client, state } }} impl LanguageModelProviderState for OpenAiLanguageModelProvider { type ObservableEntity = State; fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> { Some(self.state.clone()) }} impl LanguageModelProvider for OpenAiLanguageModelProvider { fn id(&self) -> LanguageModelProviderId { LanguageModelProviderId(PROVIDER_ID.into()) } fn name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn icon(&self) -> IconName { IconName::AiOpenAi } fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> { let model = open_ai::Model::default(); Some(Arc::new(OpenAiLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, state: self.state.clone(), http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), })) } fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> { let mut models = BTreeMap::default(); // Add base models from open_ai::Model::iter() for model in open_ai::Model::iter() { if !matches!(model, open_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), model); } } // Override with available models from settings for model in &AllLanguageModelSettings::get_global(cx) .openai .available_models { models.insert( model.name.clone(), open_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, max_output_tokens: model.max_output_tokens, max_completion_tokens: model.max_completion_tokens, }, ); } models .into_values() .map(|model| { Arc::new(OpenAiLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, state: self.state.clone(), http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), }) as Arc<dyn LanguageModel> }) .collect() } fn is_authenticated(&self, cx: &App) -> bool { self.state.read(cx).is_authenticated() } fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> { self.state.update(cx, |state, cx| state.authenticate(cx)) } fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView { cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx)) .into() } fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> { self.state.update(cx, |state, cx| state.reset_api_key(cx)) }} pub struct OpenAiLanguageModel { id: LanguageModelId, model: open_ai::Model, state: gpui::Entity<State>, http_client: Arc<dyn HttpClient>, request_limiter: RateLimiter,} impl OpenAiLanguageModel { fn stream_completion( &self, request: open_ai::Request, cx: &AsyncApp, ) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<ResponseStreamEvent>>>> { let http_client = self.http_client.clone(); let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| { let settings = &AllLanguageModelSettings::get_global(cx).openai; (state.api_key.clone(), settings.api_url.clone()) }) else { return futures::future::ready(Err(anyhow!("App state dropped"))).boxed(); }; let future = self.request_limiter.stream(async move { let api_key = api_key.ok_or_else(|| anyhow!("Missing OpenAI API Key"))?; let request = stream_completion(http_client.as_ref(), &api_url, &api_key, request); let response = request.await?; Ok(response) }); async move { Ok(future.await?.boxed()) }.boxed() }} impl LanguageModel for OpenAiLanguageModel { fn id(&self) -> LanguageModelId { self.id.clone() } fn name(&self) -> LanguageModelName { LanguageModelName::from(self.model.display_name().to_string()) } fn provider_id(&self) -> LanguageModelProviderId { LanguageModelProviderId(PROVIDER_ID.into()) } fn provider_name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn supports_tools(&self) -> bool { true } fn telemetry_id(&self) -> String { format!("openai/{}", self.model.id()) } fn max_token_count(&self) -> usize { self.model.max_token_count() } fn max_output_tokens(&self) -> Option<u32> { self.model.max_output_tokens() } fn count_tokens( &self, request: LanguageModelRequest, cx: &App, ) -> BoxFuture<'static, Result<usize>> { count_open_ai_tokens(request, self.model.clone(), cx) } fn stream_completion( &self, request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, Result<futures::stream::BoxStream<'static, Result<LanguageModelCompletionEvent>>>, > { let request = into_open_ai(request, &self.model, self.max_output_tokens()); let completions = self.stream_completion(request, cx); async move { Ok(map_to_language_model_completion_events(completions.await?).boxed()) } .boxed() }} pub fn into_open_ai( request: LanguageModelRequest, model: &Model, max_output_tokens: Option<u32>,) -> open_ai::Request { let stream = !model.id().starts_with("o1-"); let mut messages = Vec::new(); for message in request.messages { for content in message.content { match content { MessageContent::Text(text) => messages.push(match message.role { Role::User => open_ai::RequestMessage::User { content: text }, Role::Assistant => open_ai::RequestMessage::Assistant { content: Some(text), tool_calls: Vec::new(), }, Role::System => open_ai::RequestMessage::System { content: text }, }), MessageContent::Image(_) => {} MessageContent::ToolUse(tool_use) => { let tool_call = open_ai::ToolCall { id: tool_use.id.to_string(), content: open_ai::ToolCallContent::Function { function: open_ai::FunctionContent { name: tool_use.name.to_string(), arguments: serde_json::to_string(&tool_use.input) .unwrap_or_default(), }, }, }; if let Some(last_assistant_message) = messages.iter_mut().rfind(|message| { matches!(message, open_ai::RequestMessage::Assistant { .. }) }) { if let open_ai::RequestMessage::Assistant { tool_calls, .. } = last_assistant_message { tool_calls.push(tool_call); } } else { messages.push(open_ai::RequestMessage::Assistant { content: None, tool_calls: vec![tool_call], }); } } MessageContent::ToolResult(tool_result) => { messages.push(open_ai::RequestMessage::Tool { content: tool_result.content.to_string(), tool_call_id: tool_result.tool_use_id.to_string(), }); } } } } open_ai::Request { model: model.id().into(), messages, stream, stop: request.stop, temperature: request.temperature.unwrap_or(1.0), max_tokens: max_output_tokens, parallel_tool_calls: if model.supports_parallel_tool_calls() && !request.tools.is_empty() { // Disable parallel tool calls, as the Agent currently expects a maximum of one per turn. Some(false) } else { None }, tools: request .tools .into_iter() .map(|tool| open_ai::ToolDefinition::Function { function: open_ai::FunctionDefinition { name: tool.name, description: Some(tool.description), parameters: Some(tool.input_schema), }, }) .collect(), tool_choice: None, }} pub fn map_to_language_model_completion_events( events: Pin<Box<dyn Send + Stream<Item = Result<ResponseStreamEvent>>>>,) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> { #[derive(Default)] struct RawToolCall { id: String, name: String, arguments: String, } struct State { events: Pin<Box<dyn Send + Stream<Item = Result<ResponseStreamEvent>>>>, tool_calls_by_index: HashMap<usize, RawToolCall>, } futures::stream::unfold( State { events, tool_calls_by_index: HashMap::default(), }, |mut state| async move { if let Some(event) = state.events.next().await { match event { Ok(event) => { let Some(choice) = event.choices.first() else { return Some(( vec![Err(anyhow!("Response contained no choices"))], state, )); }; let mut events = Vec::new(); if let Some(content) = choice.delta.content.clone() { events.push(Ok(LanguageModelCompletionEvent::Text(content))); } if let Some(tool_calls) = choice.delta.tool_calls.as_ref() { for tool_call in tool_calls { let entry = state .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); } } } } match choice.finish_reason.as_deref() { Some("stop") => { events.push(Ok(LanguageModelCompletionEvent::Stop( StopReason::EndTurn, ))); } Some("tool_calls") => { events.extend(state.tool_calls_by_index.drain().map( |(_, tool_call)| { maybe!({ Ok(LanguageModelCompletionEvent::ToolUse( LanguageModelToolUse { id: tool_call.id.into(), name: tool_call.name.as_str().into(), input: serde_json::Value::from_str( &tool_call.arguments, )?, }, )) }) }, )); Lines 1–500 of 741. Load more
use anyhow::{Context as _, Result, anyhow}; use collections::{BTreeMap, HashMap}; use credentials_provider::CredentialsProvider; use editor::{Editor, EditorElement, EditorStyle}; use futures::Stream; use futures::{FutureExt, StreamExt, future::BoxFuture}; use gpui::{ AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace, }; use http_client::HttpClient; use language_model::{ AuthenticateError, LanguageModel, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason, }; use open_ai::{Model, ResponseStreamEvent, stream_completion}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use settings::{Settings, SettingsStore}; use std::pin::Pin; use std::str::FromStr as _; use std::sync::Arc; use strum::IntoEnumIterator; use theme::ThemeSettings; use ui::{Icon, IconName, List, Tooltip, prelude::*}; use util::{ResultExt, maybe}; use crate::{AllLanguageModelSettings, ui::InstructionListItem}; const PROVIDER_ID: &str = "openai"; const PROVIDER_NAME: &str = "OpenAI"; #[derive(Default, Clone, Debug, PartialEq)] pub struct OpenAiSettings { pub api_url: String, pub available_models: Vec<AvailableModel>, pub needs_setting_migration: bool, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct AvailableModel { pub name: String, pub display_name: Option<String>, pub max_tokens: usize, pub max_output_tokens: Option<u32>, pub max_completion_tokens: Option<u32>, } pub struct OpenAiLanguageModelProvider { http_client: Arc<dyn HttpClient>, state: gpui::Entity<State>, } pub struct State { api_key: Option<String>, api_key_from_env: bool, _subscription: Subscription, } const OPENAI_API_KEY_VAR: &str = "OPENAI_API_KEY"; impl State { fn is_authenticated(&self) -> bool { self.api_key.is_some() } fn reset_api_key(&self, cx: &mut Context<Self>) -> Task<Result<()>> { let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { credentials_provider .delete_credentials(&api_url, &cx) .await .log_err(); this.update(cx, |this, cx| { this.api_key = None; this.api_key_from_env = false; cx.notify(); }) }) } fn set_api_key(&mut self, api_key: String, cx: &mut Context<Self>) -> Task<Result<()>> { let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { credentials_provider .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx) .await .log_err(); this.update(cx, |this, cx| { this.api_key = Some(api_key); cx.notify(); }) }) } fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> { if self.is_authenticated() { return Task::ready(Ok(())); } let credentials_provider = <dyn CredentialsProvider>::global(cx); let api_url = AllLanguageModelSettings::get_global(cx) .openai .api_url .clone(); cx.spawn(async move |this, cx| { let (api_key, from_env) = if let Ok(api_key) = std::env::var(OPENAI_API_KEY_VAR) { (api_key, true) } else { let (_, api_key) = credentials_provider .read_credentials(&api_url, &cx) .await? .ok_or(AuthenticateError::CredentialsNotFound)?; ( String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?, false, ) }; this.update(cx, |this, cx| { this.api_key = Some(api_key); this.api_key_from_env = from_env; cx.notify(); })?; Ok(()) }) } } impl OpenAiLanguageModelProvider { pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self { let state = cx.new(|cx| State { api_key: None, api_key_from_env: false, _subscription: cx.observe_global::<SettingsStore>(|_this: &mut State, cx| { cx.notify(); }), }); Self { http_client, state } } } impl LanguageModelProviderState for OpenAiLanguageModelProvider { type ObservableEntity = State; fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> { Some(self.state.clone()) } } impl LanguageModelProvider for OpenAiLanguageModelProvider { fn id(&self) -> LanguageModelProviderId { LanguageModelProviderId(PROVIDER_ID.into()) } fn name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn icon(&self) -> IconName { IconName::AiOpenAi } fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> { let model = open_ai::Model::default(); Some(Arc::new(OpenAiLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, state: self.state.clone(), http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), })) } fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> { let mut models = BTreeMap::default(); // Add base models from open_ai::Model::iter() for model in open_ai::Model::iter() { if !matches!(model, open_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), model); } } // Override with available models from settings for model in &AllLanguageModelSettings::get_global(cx) .openai .available_models { models.insert( model.name.clone(), open_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, max_output_tokens: model.max_output_tokens, max_completion_tokens: model.max_completion_tokens, }, ); } models .into_values() .map(|model| { Arc::new(OpenAiLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, state: self.state.clone(), http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), }) as Arc<dyn LanguageModel> }) .collect() } fn is_authenticated(&self, cx: &App) -> bool { self.state.read(cx).is_authenticated() } fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> { self.state.update(cx, |state, cx| state.authenticate(cx)) } fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView { cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx)) .into() } fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> { self.state.update(cx, |state, cx| state.reset_api_key(cx)) } } pub struct OpenAiLanguageModel { id: LanguageModelId, model: open_ai::Model, state: gpui::Entity<State>, http_client: Arc<dyn HttpClient>, request_limiter: RateLimiter, } impl OpenAiLanguageModel { fn stream_completion( &self, request: open_ai::Request, cx: &AsyncApp, ) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<ResponseStreamEvent>>>> { let http_client = self.http_client.clone(); let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| { let settings = &AllLanguageModelSettings::get_global(cx).openai; (state.api_key.clone(), settings.api_url.clone()) }) else { return futures::future::ready(Err(anyhow!("App state dropped"))).boxed(); }; let future = self.request_limiter.stream(async move { let api_key = api_key.ok_or_else(|| anyhow!("Missing OpenAI API Key"))?; let request = stream_completion(http_client.as_ref(), &api_url, &api_key, request); let response = request.await?; Ok(response) }); async move { Ok(future.await?.boxed()) }.boxed() } } impl LanguageModel for OpenAiLanguageModel { fn id(&self) -> LanguageModelId { self.id.clone() } fn name(&self) -> LanguageModelName { LanguageModelName::from(self.model.display_name().to_string()) } fn provider_id(&self) -> LanguageModelProviderId { LanguageModelProviderId(PROVIDER_ID.into()) } fn provider_name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn supports_tools(&self) -> bool { true } fn telemetry_id(&self) -> String { format!("openai/{}", self.model.id()) } fn max_token_count(&self) -> usize { self.model.max_token_count() } fn max_output_tokens(&self) -> Option<u32> { self.model.max_output_tokens() } fn count_tokens( &self, request: LanguageModelRequest, cx: &App, ) -> BoxFuture<'static, Result<usize>> { count_open_ai_tokens(request, self.model.clone(), cx) } fn stream_completion( &self, request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, Result<futures::stream::BoxStream<'static, Result<LanguageModelCompletionEvent>>>, > { let request = into_open_ai(request, &self.model, self.max_output_tokens()); let completions = self.stream_completion(request, cx); async move { Ok(map_to_language_model_completion_events(completions.await?).boxed()) } .boxed() } } pub fn into_open_ai( request: LanguageModelRequest, model: &Model, max_output_tokens: Option<u32>, ) -> open_ai::Request { let stream = !model.id().starts_with("o1-"); let mut messages = Vec::new(); for message in request.messages { for content in message.content { match content { MessageContent::Text(text) => messages.push(match message.role { Role::User => open_ai::RequestMessage::User { content: text }, Role::Assistant => open_ai::RequestMessage::Assistant { content: Some(text), tool_calls: Vec::new(), }, Role::System => open_ai::RequestMessage::System { content: text }, }), MessageContent::Image(_) => {} MessageContent::ToolUse(tool_use) => { let tool_call = open_ai::ToolCall { id: tool_use.id.to_string(), content: open_ai::ToolCallContent::Function { function: open_ai::FunctionContent { name: tool_use.name.to_string(), arguments: serde_json::to_string(&tool_use.input) .unwrap_or_default(), }, }, }; if let Some(last_assistant_message) = messages.iter_mut().rfind(|message| { matches!(message, open_ai::RequestMessage::Assistant { .. }) }) { if let open_ai::RequestMessage::Assistant { tool_calls, .. } = last_assistant_message { tool_calls.push(tool_call); } } else { messages.push(open_ai::RequestMessage::Assistant { content: None, tool_calls: vec![tool_call], }); } } MessageContent::ToolResult(tool_result) => { messages.push(open_ai::RequestMessage::Tool { content: tool_result.content.to_string(), tool_call_id: tool_result.tool_use_id.to_string(), }); } } } } open_ai::Request { model: model.id().into(), messages, stream, stop: request.stop, temperature: request.temperature.unwrap_or(1.0), max_tokens: max_output_tokens, parallel_tool_calls: if model.supports_parallel_tool_calls() && !request.tools.is_empty() { // Disable parallel tool calls, as the Agent currently expects a maximum of one per turn. Some(false) } else { None }, tools: request .tools .into_iter() .map(|tool| open_ai::ToolDefinition::Function { function: open_ai::FunctionDefinition { name: tool.name, description: Some(tool.description), parameters: Some(tool.input_schema), }, }) .collect(), tool_choice: None, } } pub fn map_to_language_model_completion_events( events: Pin<Box<dyn Send + Stream<Item = Result<ResponseStreamEvent>>>>, ) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> { #[derive(Default)] struct RawToolCall { id: String, name: String, arguments: String, } struct State { events: Pin<Box<dyn Send + Stream<Item = Result<ResponseStreamEvent>>>>, tool_calls_by_index: HashMap<usize, RawToolCall>, } futures::stream::unfold( State { events, tool_calls_by_index: HashMap::default(), }, |mut state| async move { if let Some(event) = state.events.next().await { match event { Ok(event) => { let Some(choice) = event.choices.first() else { return Some(( vec![Err(anyhow!("Response contained no choices"))], state, )); }; let mut events = Vec::new(); if let Some(content) = choice.delta.content.clone() { events.push(Ok(LanguageModelCompletionEvent::Text(content))); } if let Some(tool_calls) = choice.delta.tool_calls.as_ref() { for tool_call in tool_calls { let entry = state .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); } } } } match choice.finish_reason.as_deref() { Some("stop") => { events.push(Ok(LanguageModelCompletionEvent::Stop( StopReason::EndTurn, ))); } Some("tool_calls") => { events.extend(state.tool_calls_by_index.drain().map( |(_, tool_call)| { maybe!({ Ok(LanguageModelCompletionEvent::ToolUse( LanguageModelToolUse { id: tool_call.id.into(), name: tool_call.name.as_str().into(), input: serde_json::Value::from_str( &tool_call.arguments, )?, }, )) }) }, ));