1use anyhow::{Result, anyhow};
2use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream};
3use futures::{Stream, TryFutureExt, stream};
4use gpui::{AnyView, App, AsyncApp, Context, Subscription, Task};
5use http_client::HttpClient;
6use language_model::{
7 AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
8 LanguageModelRequestTool, LanguageModelToolChoice, LanguageModelToolUse,
9 LanguageModelToolUseId, StopReason,
10};
11use language_model::{
12 LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
13 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
14 LanguageModelRequest, RateLimiter, Role,
15};
16use ollama::{
17 ChatMessage, ChatOptions, ChatRequest, ChatResponseDelta, KeepAlive, OllamaFunctionTool,
18 OllamaToolCall, get_models, preload_model, show_model, stream_chat_completion,
19};
20use schemars::JsonSchema;
21use serde::{Deserialize, Serialize};
22use settings::{Settings, SettingsStore};
23use std::pin::Pin;
24use std::sync::atomic::{AtomicU64, Ordering};
25use std::{collections::BTreeMap, sync::Arc};
26use ui::{ButtonLike, Indicator, List, prelude::*};
27use util::ResultExt;
28
29use crate::AllLanguageModelSettings;
30use crate::ui::InstructionListItem;
31
32const OLLAMA_DOWNLOAD_URL: &str = "https://ollama.com/download";
33const OLLAMA_LIBRARY_URL: &str = "https://ollama.com/library";
34const OLLAMA_SITE: &str = "https://ollama.com/";
35
36const PROVIDER_ID: &str = "ollama";
37const PROVIDER_NAME: &str = "Ollama";
38
39#[derive(Default, Debug, Clone, PartialEq)]
40pub struct OllamaSettings {
41 pub api_url: String,
42 pub available_models: Vec<AvailableModel>,
43}
44
45#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
46pub struct AvailableModel {
47 /// The model name in the Ollama API (e.g. "llama3.2:latest")
48 pub name: String,
49 /// The model's name in Zed's UI, such as in the model selector dropdown menu in the assistant panel.
50 pub display_name: Option<String>,
51 /// The Context Length parameter to the model (aka num_ctx or n_ctx)
52 pub max_tokens: usize,
53 /// The number of seconds to keep the connection open after the last request
54 pub keep_alive: Option<KeepAlive>,
55 /// Whether the model supports tools
56 pub supports_tools: Option<bool>,
57}
58
59pub struct OllamaLanguageModelProvider {
60 http_client: Arc<dyn HttpClient>,
61 state: gpui::Entity<State>,
62}
63
64pub struct State {
65 http_client: Arc<dyn HttpClient>,
66 available_models: Vec<ollama::Model>,
67 fetch_model_task: Option<Task<Result<()>>>,
68 _subscription: Subscription,
69}
70
71impl State {
72 fn is_authenticated(&self) -> bool {
73 !self.available_models.is_empty()
74 }
75
76 fn fetch_models(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
77 let settings = &AllLanguageModelSettings::get_global(cx).ollama;
78 let http_client = Arc::clone(&self.http_client);
79 let api_url = settings.api_url.clone();
80
81 // As a proxy for the server being "authenticated", we'll check if its up by fetching the models
82 cx.spawn(async move |this, cx| {
83 let models = get_models(http_client.as_ref(), &api_url, None).await?;
84
85 let tasks = models
86 .into_iter()
87 // Since there is no metadata from the Ollama API
88 // indicating which models are embedding models,
89 // simply filter out models with "-embed" in their name
90 .filter(|model| !model.name.contains("-embed"))
91 .map(|model| {
92 let http_client = Arc::clone(&http_client);
93 let api_url = api_url.clone();
94 async move {
95 let name = model.name.as_str();
96 let capabilities = show_model(http_client.as_ref(), &api_url, name).await?;
97 let ollama_model = ollama::Model::new(
98 name,
99 None,
100 None,
101 Some(capabilities.supports_tools()),
102 );
103 Ok(ollama_model)
104 }
105 });
106
107 // Rate-limit capability fetches
108 // since there is an arbitrary number of models available
109 let mut ollama_models: Vec<_> = futures::stream::iter(tasks)
110 .buffer_unordered(5)
111 .collect::<Vec<Result<_>>>()
112 .await
113 .into_iter()
114 .collect::<Result<Vec<_>>>()?;
115
116 ollama_models.sort_by(|a, b| a.name.cmp(&b.name));
117
118 this.update(cx, |this, cx| {
119 this.available_models = ollama_models;
120 cx.notify();
121 })
122 })
123 }
124
125 fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
126 let task = self.fetch_models(cx);
127 self.fetch_model_task.replace(task);
128 }
129
130 fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
131 if self.is_authenticated() {
132 return Task::ready(Ok(()));
133 }
134
135 let fetch_models_task = self.fetch_models(cx);
136 cx.spawn(async move |_this, _cx| Ok(fetch_models_task.await?))
137 }
138}
139
140impl OllamaLanguageModelProvider {
141 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
142 let this = Self {
143 http_client: http_client.clone(),
144 state: cx.new(|cx| {
145 let subscription = cx.observe_global::<SettingsStore>({
146 let mut settings = AllLanguageModelSettings::get_global(cx).ollama.clone();
147 move |this: &mut State, cx| {
148 let new_settings = &AllLanguageModelSettings::get_global(cx).ollama;
149 if &settings != new_settings {
150 settings = new_settings.clone();
151 this.restart_fetch_models_task(cx);
152 cx.notify();
153 }
154 }
155 });
156
157 State {
158 http_client,
159 available_models: Default::default(),
160 fetch_model_task: None,
161 _subscription: subscription,
162 }
163 }),
164 };
165 this.state
166 .update(cx, |state, cx| state.restart_fetch_models_task(cx));
167 this
168 }
169}
170
171impl LanguageModelProviderState for OllamaLanguageModelProvider {
172 type ObservableEntity = State;
173
174 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
175 Some(self.state.clone())
176 }
177}
178
179impl LanguageModelProvider for OllamaLanguageModelProvider {
180 fn id(&self) -> LanguageModelProviderId {
181 LanguageModelProviderId(PROVIDER_ID.into())
182 }
183
184 fn name(&self) -> LanguageModelProviderName {
185 LanguageModelProviderName(PROVIDER_NAME.into())
186 }
187
188 fn icon(&self) -> IconName {
189 IconName::AiOllama
190 }
191
192 fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
193 self.provided_models(cx).into_iter().next()
194 }
195
196 fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
197 self.default_model(cx)
198 }
199
200 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
201 let mut models: BTreeMap<String, ollama::Model> = BTreeMap::default();
202
203 // Add models from the Ollama API
204 for model in self.state.read(cx).available_models.iter() {
205 models.insert(model.name.clone(), model.clone());
206 }
207
208 // Override with available models from settings
209 for model in AllLanguageModelSettings::get_global(cx)
210 .ollama
211 .available_models
212 .iter()
213 {
214 models.insert(
215 model.name.clone(),
216 ollama::Model {
217 name: model.name.clone(),
218 display_name: model.display_name.clone(),
219 max_tokens: model.max_tokens,
220 keep_alive: model.keep_alive.clone(),
221 supports_tools: model.supports_tools,
222 },
223 );
224 }
225
226 models
227 .into_values()
228 .map(|model| {
229 Arc::new(OllamaLanguageModel {
230 id: LanguageModelId::from(model.name.clone()),
231 model: model.clone(),
232 http_client: self.http_client.clone(),
233 request_limiter: RateLimiter::new(4),
234 }) as Arc<dyn LanguageModel>
235 })
236 .collect()
237 }
238
239 fn load_model(&self, model: Arc<dyn LanguageModel>, cx: &App) {
240 let settings = &AllLanguageModelSettings::get_global(cx).ollama;
241 let http_client = self.http_client.clone();
242 let api_url = settings.api_url.clone();
243 let id = model.id().0.to_string();
244 cx.spawn(async move |_| preload_model(http_client, &api_url, &id).await)
245 .detach_and_log_err(cx);
246 }
247
248 fn is_authenticated(&self, cx: &App) -> bool {
249 self.state.read(cx).is_authenticated()
250 }
251
252 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
253 self.state.update(cx, |state, cx| state.authenticate(cx))
254 }
255
256 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
257 let state = self.state.clone();
258 cx.new(|cx| ConfigurationView::new(state, window, cx))
259 .into()
260 }
261
262 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
263 self.state.update(cx, |state, cx| state.fetch_models(cx))
264 }
265}
266
267pub struct OllamaLanguageModel {
268 id: LanguageModelId,
269 model: ollama::Model,
270 http_client: Arc<dyn HttpClient>,
271 request_limiter: RateLimiter,
272}
273
274impl OllamaLanguageModel {
275 fn to_ollama_request(&self, request: LanguageModelRequest) -> ChatRequest {
276 ChatRequest {
277 model: self.model.name.clone(),
278 messages: request
279 .messages
280 .into_iter()
281 .map(|msg| match msg.role {
282 Role::User => ChatMessage::User {
283 content: msg.string_contents(),
284 },
285 Role::Assistant => ChatMessage::Assistant {
286 content: msg.string_contents(),
287 tool_calls: None,
288 },
289 Role::System => ChatMessage::System {
290 content: msg.string_contents(),
291 },
292 })
293 .collect(),
294 keep_alive: self.model.keep_alive.clone().unwrap_or_default(),
295 stream: true,
296 options: Some(ChatOptions {
297 num_ctx: Some(self.model.max_tokens),
298 stop: Some(request.stop),
299 temperature: request.temperature.or(Some(1.0)),
300 ..Default::default()
301 }),
302 tools: request.tools.into_iter().map(tool_into_ollama).collect(),
303 }
304 }
305}
306
307impl LanguageModel for OllamaLanguageModel {
308 fn id(&self) -> LanguageModelId {
309 self.id.clone()
310 }
311
312 fn name(&self) -> LanguageModelName {
313 LanguageModelName::from(self.model.display_name().to_string())
314 }
315
316 fn provider_id(&self) -> LanguageModelProviderId {
317 LanguageModelProviderId(PROVIDER_ID.into())
318 }
319
320 fn provider_name(&self) -> LanguageModelProviderName {
321 LanguageModelProviderName(PROVIDER_NAME.into())
322 }
323
324 fn supports_tools(&self) -> bool {
325 self.model.supports_tools.unwrap_or(false)
326 }
327
328 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
329 match choice {
330 LanguageModelToolChoice::Auto => false,
331 LanguageModelToolChoice::Any => false,
332 LanguageModelToolChoice::None => false,
333 }
334 }
335
336 fn telemetry_id(&self) -> String {
337 format!("ollama/{}", self.model.id())
338 }
339
340 fn max_token_count(&self) -> usize {
341 self.model.max_token_count()
342 }
343
344 fn count_tokens(
345 &self,
346 request: LanguageModelRequest,
347 _cx: &App,
348 ) -> BoxFuture<'static, Result<usize>> {
349 // There is no endpoint for this _yet_ in Ollama
350 // see: https://github.com/ollama/ollama/issues/1716 and https://github.com/ollama/ollama/issues/3582
351 let token_count = request
352 .messages
353 .iter()
354 .map(|msg| msg.string_contents().chars().count())
355 .sum::<usize>()
356 / 4;
357
358 async move { Ok(token_count) }.boxed()
359 }
360
361 fn stream_completion(
362 &self,
363 request: LanguageModelRequest,
364 cx: &AsyncApp,
365 ) -> BoxFuture<
366 'static,
367 Result<
368 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
369 >,
370 > {
371 let request = self.to_ollama_request(request);
372
373 let http_client = self.http_client.clone();
374 let Ok(api_url) = cx.update(|cx| {
375 let settings = &AllLanguageModelSettings::get_global(cx).ollama;
376 settings.api_url.clone()
377 }) else {
378 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
379 };
380
381 let future = self.request_limiter.stream(async move {
382 let stream = stream_chat_completion(http_client.as_ref(), &api_url, request).await?;
383 let stream = map_to_language_model_completion_events(stream);
384 Ok(stream)
385 });
386
387 future.map_ok(|f| f.boxed()).boxed()
388 }
389}
390
391fn map_to_language_model_completion_events(
392 stream: Pin<Box<dyn Stream<Item = anyhow::Result<ChatResponseDelta>> + Send>>,
393) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
394 // Used for creating unique tool use ids
395 static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
396
397 struct State {
398 stream: Pin<Box<dyn Stream<Item = anyhow::Result<ChatResponseDelta>> + Send>>,
399 used_tools: bool,
400 }
401
402 // We need to create a ToolUse and Stop event from a single
403 // response from the original stream
404 let stream = stream::unfold(
405 State {
406 stream,
407 used_tools: false,
408 },
409 async move |mut state| {
410 let response = state.stream.next().await?;
411
412 let delta = match response {
413 Ok(delta) => delta,
414 Err(e) => {
415 let event = Err(LanguageModelCompletionError::Other(anyhow!(e)));
416 return Some((vec![event], state));
417 }
418 };
419
420 let mut events = Vec::new();
421
422 match delta.message {
423 ChatMessage::User { content } => {
424 events.push(Ok(LanguageModelCompletionEvent::Text(content)));
425 }
426 ChatMessage::System { content } => {
427 events.push(Ok(LanguageModelCompletionEvent::Text(content)));
428 }
429 ChatMessage::Assistant {
430 content,
431 tool_calls,
432 } => {
433 // Check for tool calls
434 if let Some(tool_call) = tool_calls.and_then(|v| v.into_iter().next()) {
435 match tool_call {
436 OllamaToolCall::Function(function) => {
437 let tool_id = format!(
438 "{}-{}",
439 &function.name,
440 TOOL_CALL_COUNTER.fetch_add(1, Ordering::Relaxed)
441 );
442 let event =
443 LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse {
444 id: LanguageModelToolUseId::from(tool_id),
445 name: Arc::from(function.name),
446 raw_input: function.arguments.to_string(),
447 input: function.arguments,
448 is_input_complete: true,
449 });
450 events.push(Ok(event));
451 state.used_tools = true;
452 }
453 }
454 } else {
455 events.push(Ok(LanguageModelCompletionEvent::Text(content)));
456 }
457 }
458 };
459
460 if delta.done {
461 if state.used_tools {
462 state.used_tools = false;
463 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
464 } else {
465 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
466 }
467 }
468
469 Some((events, state))
470 },
471 );
472
473 stream.flat_map(futures::stream::iter)
474}
475
476struct ConfigurationView {
477 state: gpui::Entity<State>,
478 loading_models_task: Option<Task<()>>,
479}
480
481impl ConfigurationView {
482 pub fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
483 let loading_models_task = Some(cx.spawn_in(window, {
484 let state = state.clone();
485 async move |this, cx| {
486 if let Some(task) = state
487 .update(cx, |state, cx| state.authenticate(cx))
488 .log_err()
489 {
490 task.await.log_err();
491 }
492 this.update(cx, |this, cx| {
493 this.loading_models_task = None;
494 cx.notify();
495 })
496 .log_err();
497 }
498 }));
499
500 Self {
501 state,
502 loading_models_task,
503 }
504 }
505
506 fn retry_connection(&self, cx: &mut App) {
507 self.state
508 .update(cx, |state, cx| state.fetch_models(cx))
509 .detach_and_log_err(cx);
510 }
511}
512
513impl Render for ConfigurationView {
514 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
515 let is_authenticated = self.state.read(cx).is_authenticated();
516
517 let ollama_intro =
518 "Get up & running with Llama 3.3, Mistral, Gemma 2, and other LLMs with Ollama.";
519
520 if self.loading_models_task.is_some() {
521 div().child(Label::new("Loading models...")).into_any()
522 } else {
523 v_flex()
524 .gap_2()
525 .child(
526 v_flex().gap_1().child(Label::new(ollama_intro)).child(
527 List::new()
528 .child(InstructionListItem::text_only("Ollama must be running with at least one model installed to use it in the assistant."))
529 .child(InstructionListItem::text_only(
530 "Once installed, try `ollama run llama3.2`",
531 )),
532 ),
533 )
534 .child(
535 h_flex()
536 .w_full()
537 .justify_between()
538 .gap_2()
539 .child(
540 h_flex()
541 .w_full()
542 .gap_2()
543 .map(|this| {
544 if is_authenticated {
545 this.child(
546 Button::new("ollama-site", "Ollama")
547 .style(ButtonStyle::Subtle)
548 .icon(IconName::ArrowUpRight)
549 .icon_size(IconSize::XSmall)
550 .icon_color(Color::Muted)
551 .on_click(move |_, _, cx| cx.open_url(OLLAMA_SITE))
552 .into_any_element(),
553 )
554 } else {
555 this.child(
556 Button::new(
557 "download_ollama_button",
558 "Download Ollama",
559 )
560 .style(ButtonStyle::Subtle)
561 .icon(IconName::ArrowUpRight)
562 .icon_size(IconSize::XSmall)
563 .icon_color(Color::Muted)
564 .on_click(move |_, _, cx| {
565 cx.open_url(OLLAMA_DOWNLOAD_URL)
566 })
567 .into_any_element(),
568 )
569 }
570 })
571 .child(
572 Button::new("view-models", "All Models")
573 .style(ButtonStyle::Subtle)
574 .icon(IconName::ArrowUpRight)
575 .icon_size(IconSize::XSmall)
576 .icon_color(Color::Muted)
577 .on_click(move |_, _, cx| cx.open_url(OLLAMA_LIBRARY_URL)),
578 ),
579 )
580 .map(|this| {
581 if is_authenticated {
582 this.child(
583 ButtonLike::new("connected")
584 .disabled(true)
585 .cursor_style(gpui::CursorStyle::Arrow)
586 .child(
587 h_flex()
588 .gap_2()
589 .child(Indicator::dot().color(Color::Success))
590 .child(Label::new("Connected"))
591 .into_any_element(),
592 ),
593 )
594 } else {
595 this.child(
596 Button::new("retry_ollama_models", "Connect")
597 .icon_position(IconPosition::Start)
598 .icon_size(IconSize::XSmall)
599 .icon(IconName::Play)
600 .on_click(cx.listener(move |this, _, _, cx| {
601 this.retry_connection(cx)
602 })),
603 )
604 }
605 })
606 )
607 .into_any()
608 }
609 }
610}
611
612fn tool_into_ollama(tool: LanguageModelRequestTool) -> ollama::OllamaTool {
613 ollama::OllamaTool::Function {
614 function: OllamaFunctionTool {
615 name: tool.name,
616 description: Some(tool.description),
617 parameters: Some(tool.input_schema),
618 },
619 }
620}