assistant: Fix issues when configuring different providers (#15072)

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>

crates/language_model/src/provider/ollama.rs

@@ -2,21 +2,24 @@ use anyhow::{anyhow, Result};
 use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
 use gpui::{AnyView, AppContext, AsyncAppContext, ModelContext, Subscription, Task};
 use http_client::HttpClient;
-use ollama::{get_models, stream_chat_completion, ChatMessage, ChatOptions, ChatRequest};
+use ollama::{
+    get_models, preload_model, stream_chat_completion, ChatMessage, ChatOptions, ChatRequest,
+};
 use settings::{Settings, SettingsStore};
 use std::{sync::Arc, time::Duration};
 use ui::{prelude::*, ButtonLike, ElevationIndex};
 
 use crate::{
     settings::AllLanguageModelSettings, LanguageModel, LanguageModelId, LanguageModelName,
-    LanguageModelProvider, LanguageModelProviderName, LanguageModelProviderState,
-    LanguageModelRequest, Role,
+    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
+    LanguageModelProviderState, LanguageModelRequest, Role,
 };
 
 const OLLAMA_DOWNLOAD_URL: &str = "https://ollama.com/download";
 const OLLAMA_LIBRARY_URL: &str = "https://ollama.com/library";
 
-const PROVIDER_NAME: &str = "ollama";
+const PROVIDER_ID: &str = "ollama";
+const PROVIDER_NAME: &str = "Ollama";
 
 #[derive(Default, Debug, Clone, PartialEq)]
 pub struct OllamaSettings {
@@ -32,14 +35,14 @@ pub struct OllamaLanguageModelProvider {
 struct State {
     http_client: Arc<dyn HttpClient>,
     available_models: Vec<ollama::Model>,
-    settings: OllamaSettings,
     _subscription: Subscription,
 }
 
 impl State {
-    fn fetch_models(&self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
+    fn fetch_models(&self, cx: &ModelContext<Self>) -> Task<Result<()>> {
+        let settings = &AllLanguageModelSettings::get_global(cx).ollama;
         let http_client = self.http_client.clone();
-        let api_url = self.settings.api_url.clone();
+        let api_url = settings.api_url.clone();
 
         // As a proxy for the server being "authenticated", we'll check if its up by fetching the models
         cx.spawn(|this, mut cx| async move {
@@ -66,23 +69,25 @@ impl State {
 
 impl OllamaLanguageModelProvider {
     pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut AppContext) -> Self {
-        Self {
+        let this = Self {
             http_client: http_client.clone(),
             state: cx.new_model(|cx| State {
                 http_client,
                 available_models: Default::default(),
-                settings: OllamaSettings::default(),
                 _subscription: cx.observe_global::<SettingsStore>(|this: &mut State, cx| {
-                    this.settings = AllLanguageModelSettings::get_global(cx).ollama.clone();
+                    this.fetch_models(cx).detach_and_log_err(cx);
                     cx.notify();
                 }),
             }),
-        }
+        };
+        this.fetch_models(cx).detach_and_log_err(cx);
+        this
     }
 
     fn fetch_models(&self, cx: &AppContext) -> Task<Result<()>> {
+        let settings = &AllLanguageModelSettings::get_global(cx).ollama;
         let http_client = self.http_client.clone();
-        let api_url = self.state.read(cx).settings.api_url.clone();
+        let api_url = settings.api_url.clone();
 
         let state = self.state.clone();
         // As a proxy for the server being "authenticated", we'll check if its up by fetching the models
@@ -117,6 +122,10 @@ impl LanguageModelProviderState for OllamaLanguageModelProvider {
 }
 
 impl LanguageModelProvider for OllamaLanguageModelProvider {
+    fn id(&self) -> LanguageModelProviderId {
+        LanguageModelProviderId(PROVIDER_ID.into())
+    }
+
     fn name(&self) -> LanguageModelProviderName {
         LanguageModelProviderName(PROVIDER_NAME.into())
     }
@@ -131,12 +140,20 @@ impl LanguageModelProvider for OllamaLanguageModelProvider {
                     id: LanguageModelId::from(model.name.clone()),
                     model: model.clone(),
                     http_client: self.http_client.clone(),
-                    state: self.state.clone(),
                 }) as Arc<dyn LanguageModel>
             })
             .collect()
     }
 
+    fn load_model(&self, model: Arc<dyn LanguageModel>, cx: &AppContext) {
+        let settings = &AllLanguageModelSettings::get_global(cx).ollama;
+        let http_client = self.http_client.clone();
+        let api_url = settings.api_url.clone();
+        let id = model.id().0.to_string();
+        cx.spawn(|_| async move { preload_model(http_client, &api_url, &id).await })
+            .detach_and_log_err(cx);
+    }
+
     fn is_authenticated(&self, cx: &AppContext) -> bool {
         !self.state.read(cx).available_models.is_empty()
     }
@@ -167,7 +184,6 @@ impl LanguageModelProvider for OllamaLanguageModelProvider {
 pub struct OllamaLanguageModel {
     id: LanguageModelId,
     model: ollama::Model,
-    state: gpui::Model<State>,
     http_client: Arc<dyn HttpClient>,
 }
 
@@ -211,6 +227,14 @@ impl LanguageModel for OllamaLanguageModel {
         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 max_token_count(&self) -> usize {
         self.model.max_token_count()
     }
@@ -219,10 +243,6 @@ impl LanguageModel for OllamaLanguageModel {
         format!("ollama/{}", self.model.id())
     }
 
-    fn provider_name(&self) -> LanguageModelProviderName {
-        LanguageModelProviderName(PROVIDER_NAME.into())
-    }
-
     fn count_tokens(
         &self,
         request: LanguageModelRequest,
@@ -248,11 +268,9 @@ impl LanguageModel for OllamaLanguageModel {
         let request = self.to_ollama_request(request);
 
         let http_client = self.http_client.clone();
-        let Ok((api_url, low_speed_timeout)) = cx.read_model(&self.state, |state, _| {
-            (
-                state.settings.api_url.clone(),
-                state.settings.low_speed_timeout,
-            )
+        let Ok((api_url, low_speed_timeout)) = cx.update(|cx| {
+            let settings = &AllLanguageModelSettings::get_global(cx).ollama;
+            (settings.api_url.clone(), settings.low_speed_timeout)
         }) else {
             return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
         };