Ollama support

README.md

@@ -133,7 +133,16 @@ layerrule = ignore_alpha 0.35,match:class zlaunch
 
 ## AI Mode
 
-To enable AI mode, run the daemon with the `GEMINI_API_KEY` `OPENAI_API_KEY` or `OPENROUTER_API_KEY` env var set to an
+To enable AI mode with Ollama, set the `OLLAMA_URL` and `OLLAMA_MODEL` environment variables.
+- `OLLAMA_URL`: The url of your Ollama server (`http://127.0.0.1:11434`).
+- `OLLAMA_MODEL`: The model to use (`llama3.2:latest`).
+
+Example usage:
+```bash
+OLLAMA_URL="http://127.0.0.1:11434" OLLAMA_MODEL="llama3.2:latest" zlaunch
+```
+
+To enable AI mode with cloud models, run the daemon with the `GEMINI_API_KEY` `OPENAI_API_KEY` or `OPENROUTER_API_KEY` env var set to an
 appropriate key. Model can be chosen for OpenRouter with the `OPENROUTER_MODEL` env var.
 
 ## License

src/ai/client.rs

@@ -1,36 +1,53 @@
 //! LLM API client for streaming AI responses.
 
-use anyhow::{Context, Result};
+use anyhow::{Result, anyhow};
 use futures::Stream;
-use futures::stream::StreamExt;
+use futures::stream::{StreamExt, once};
 use llm::LLMProvider;
 use llm::builder::{LLMBackend, LLMBuilder};
 use llm::chat::ChatMessage;
 use std::env;
 use std::pin::Pin;
 
+fn get_keys() -> Option<(String, LLMBackend)> {
+    [
+        ("OLLAMA_URL", LLMBackend::Ollama),
+        ("GEMINI_API_KEY", LLMBackend::Google),
+        ("OPENAI_API_KEY", LLMBackend::OpenAI),
+        ("OPENROUTER_API_KEY", LLMBackend::OpenRouter),
+    ]
+    .iter()
+    .find_map(|(var_name, backend)| {
+        env::var(var_name)
+            .ok()
+            .map(|value| (value, backend.clone()))
+    })
+}
+
 /// LLM client for AI queries.
 pub struct LLMClient {
     llm: Box<dyn LLMProvider>,
+    backend: LLMBackend,
 }
 
 impl LLMClient {
     /// Create a new LLM client.
     /// Returns None if no valid API_KEY environment variable is set.
     pub fn new() -> Option<Self> {
-        // Find the first available environment variable
-        let (api_key, backend) = [
-            ("GEMINI_API_KEY", LLMBackend::Google),
-            ("OPENAI_API_KEY", LLMBackend::OpenAI),
-            ("OPENROUTER_API_KEY", LLMBackend::OpenRouter),
-        ]
-        .iter()
-        .find_map(|(var_name, backend)| env::var(var_name).ok().map(|key| (key, backend)))?;
-
-        LLMBuilder::new()
-            .backend(backend.clone())
-            .api_key(&api_key)
+        let (api_key, backend) = get_keys()?;
+
+        let mut builder = LLMBuilder::new().backend(backend.clone());
+
+        builder = match backend {
+            LLMBackend::Ollama => builder.base_url(&api_key),
+            _ => builder.api_key(&api_key),
+        };
+
+        let llm = builder
             .model(match backend {
+                LLMBackend::Ollama => {
+                    env::var("OLLAMA_MODEL").unwrap_or_else(|_| "llama3.2:latest".to_string())
+                }
                 LLMBackend::Google => "gemini-flash-latest".to_string(),
                 LLMBackend::OpenAI => "gpt-5-mini".to_string(),
                 LLMBackend::OpenRouter => env::var("OPENROUTER_MODEL")
@@ -40,34 +57,71 @@ impl LLMClient {
             .max_tokens(2000)
             .temperature(0.7)
             .build()
-            .ok()
-            .map(|llm| Self { llm })
+            .ok()?;
+
+        Some(Self {
+            llm,
+            backend: backend.clone(),
+        })
+    }
+
+    /// Return true if any LLM is configured.
+    pub fn is_configured() -> bool {
+        get_keys().is_some()
     }
 
     /// Stream a response for the given query.
     /// Returns a stream of tokens (strings).
+    /// Ollama doesnt support streaming for some reason so it uses normal chat().
     pub async fn stream_query(
         &self,
         messages: &[ChatMessage],
     ) -> Result<Pin<Box<dyn Stream<Item = Result<String>> + Send>>> {
-        let stream = self
-            .llm
-            .chat_stream(messages)
-            .await
-            .context("Failed to initiate streaming chat")?;
+        if self.backend == LLMBackend::Ollama {
+            let response = self.llm.chat(messages).await?;
+
+            let text = response
+                .text()
+                .ok_or_else(|| anyhow!("LLM response missing text"))?;
+
+            let result = once(async move { Ok(text) });
+
+            Ok(Box::pin(result))
+        } else {
+            let stream = self.llm.chat_stream(messages).await?;
 
-        // Convert LLMError to anyhow::Error
-        let result_stream =
-            stream.map(|result| result.map_err(|e| anyhow::Error::msg(e.to_string())));
+            // Convert LLMError to anyhow::Error
+            let result_stream =
+                stream.map(|result| result.map_err(|e| anyhow::Error::msg(e.to_string())));
 
-        Ok(Box::pin(result_stream))
+            Ok(Box::pin(result_stream))
+        }
     }
+
+    // Keeping this here bc if llm crate adds streaming for ollama we can probably just uncomment this and remove above code.
+
+    // pub async fn stream_query(
+    //     &self,
+    //     messages: &[ChatMessage],
+    // ) -> Result<Pin<Box<dyn Stream<Item = Result<String>> + Send>>> {
+    //     let stream = self
+    //         .llm
+    //         .chat_stream(messages)
+    //         .await
+    //         .context("Failed to initiate streaming chat")?;
+
+    //     // Convert LLMError to anyhow::Error
+    //     let result_stream =
+    //         stream.map(|result| result.map_err(|e| anyhow::Error::msg(e.to_string())));
+
+    //     Ok(Box::pin(result_stream))
+    // }
 }
 
 impl Default for LLMClient {
     fn default() -> Self {
         Self::new().expect(
-            "GEMINI_API_KEY, OPENAI_API_KEY or OPENROUTER_API_KEY environment variable not set",
+            "GEMINI_API_KEY, OPENAI_API_KEY, OPENROUTER_API_KEY or OLLAMA_URL environment variable not set",
         )
     }
 }

src/ai/streaming.rs

@@ -31,6 +31,10 @@ static CLIENT: OnceCell<LLMClient> = OnceCell::const_new();
 /// - `Ok("")` when streaming completes successfully
 /// - `Err(error)` if an error occurs
 pub fn spawn_stream(messages: Vec<ChatMessage>) -> Option<Receiver<Result<String, String>>> {
+    if !LLMClient::is_configured() {
+        return None;
+    }
+
     // Create channel for communication between Tokio thread and caller
     let (tx, rx) = flume::unbounded::<Result<String, String>>();
 

src/ui/delegates/item_delegate.rs

@@ -1,3 +1,4 @@
+use crate::ai::LLMClient;
 use crate::calculator::evaluate_expression;
 use crate::config::{ConfigModule, config};
 use crate::items::{ActionItem, AiItem, CalculatorItem, ListItem, SearchItem, SubmenuItem};
@@ -155,6 +156,8 @@ impl ItemListDelegate {
     fn process_query(&mut self, query: &str) {
         // Get the config disabled modules
         let disabled_modules = config().disabled_modules.unwrap_or_default();
+        let ai_enabled =
+            !disabled_modules.contains(&ConfigModule::Ai) && LLMClient::is_configured();
 
         // Check for calculator expression
         if !disabled_modules.contains(&ConfigModule::Calculator)
@@ -185,7 +188,7 @@ impl ItemListDelegate {
         // 2. Else if search trigger (!g, !ddg, etc.) → only show that search provider
         // 3. Else if query not empty → always show AI item + all search providers at bottom
 
-        if !disabled_modules.contains(&ConfigModule::Ai) && has_ai_trigger {
+        if ai_enabled && has_ai_trigger {
             // Only show AI item when !ai trigger is used
             let ai_query = trimmed.strip_prefix("!ai").unwrap().trim();
             if !ai_query.is_empty() {
@@ -199,7 +202,7 @@ impl ItemListDelegate {
         } else if !trimmed.is_empty() {
             // Always show AI item and all search providers when query is not empty
             // These appear at the bottom in "Search and AI" section
-            if !disabled_modules.contains(&ConfigModule::Ai) {
+            if ai_enabled {
                 self.ai_item = Some(AiItem::new(trimmed.to_string()));
             }
             if !disabled_modules.contains(&ConfigModule::Search)