From adb44ab399ea4e1d68a4255310a9cc5791e5b67c Mon Sep 17 00:00:00 2001 From: imnotdev25 <85677268+imnotdev25@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:39:38 +0530 Subject: [PATCH 1/2] feat: support OpenAI-compatible providers for local AI chat and embeddings Local AI was hard-wired to Ollama. This introduces an AIClient abstraction with two backends, so chat, completion and embeddings can also run against any endpoint speaking the OpenAI REST API (OpenAI, LM Studio, vLLM, llama.cpp, OpenRouter, Groq, ...). LocalAISetting gains a `provider` discriminator and an `api_key`; the existing server_url / chat_model / embedding_model fields now apply to both backends. Both new fields are `#[serde(default)]` so settings written by earlier versions still load. The backend swap happens inside LLMOllama rather than at its call sites, so the ten chain, completion and database modules that consume it are untouched. Embeddings request `dimensions: 768` and every returned vector is verified, because the vector store column is declared `float[768]`. A model that returns a different size now fails with an actionable message instead of corrupting the store. Also fixes the embedding model name being ignored during indexing: it was read from settings only for the resource pre-flight check, while indexing used a hardcoded "nomic-embed-text". async-openai is used directly rather than langchain's OpenAiEmbedder, which cannot request `dimensions`. It was already compiled as a non-optional dependency of the langchain-rust fork. Co-Authored-By: Claude Opus 4.8 --- frontend/rust-lib/Cargo.lock | 1 + frontend/rust-lib/flowy-ai/Cargo.toml | 3 + .../flowy-ai/src/embeddings/context.rs | 49 +- .../src/embeddings/document_indexer.rs | 15 +- .../flowy-ai/src/embeddings/embedder.rs | 45 +- .../flowy-ai/src/embeddings/scheduler.rs | 37 +- .../rust-lib/flowy-ai/src/embeddings/store.rs | 33 +- frontend/rust-lib/flowy-ai/src/entities.rs | 42 ++ .../flowy-ai/src/local_ai/chat/llm.rs | 88 ++-- .../flowy-ai/src/local_ai/chat/llm_chat.rs | 6 +- .../flowy-ai/src/local_ai/chat/mod.rs | 17 +- .../rust-lib/flowy-ai/src/local_ai/client.rs | 449 ++++++++++++++++++ .../flowy-ai/src/local_ai/controller.rs | 81 ++-- .../rust-lib/flowy-ai/src/local_ai/mod.rs | 1 + .../rust-lib/flowy-ai/src/model_select.rs | 8 +- .../rust-lib/flowy-ai/src/search/summary.rs | 53 +-- 16 files changed, 703 insertions(+), 225 deletions(-) create mode 100644 frontend/rust-lib/flowy-ai/src/local_ai/client.rs diff --git a/frontend/rust-lib/Cargo.lock b/frontend/rust-lib/Cargo.lock index 6729ed9687234..755e335fda9c7 100644 --- a/frontend/rust-lib/Cargo.lock +++ b/frontend/rust-lib/Cargo.lock @@ -2204,6 +2204,7 @@ dependencies = [ "allo-isolate", "anyhow", "arc-swap", + "async-openai", "async-stream", "async-trait", "base64 0.21.5", diff --git a/frontend/rust-lib/flowy-ai/Cargo.toml b/frontend/rust-lib/flowy-ai/Cargo.toml index ad859bd16f8d4..913a79774b567 100644 --- a/frontend/rust-lib/flowy-ai/Cargo.toml +++ b/frontend/rust-lib/flowy-ai/Cargo.toml @@ -41,6 +41,9 @@ base64 = "0.21.5" futures-util = "0.3.30" flowy-storage-pub = { workspace = true } ollama-rs.workspace = true +# Already compiled as a non-optional dependency of the langchain-rust fork; used directly here for +# the OpenAI-compatible backend because langchain's OpenAiEmbedder cannot request `dimensions`. +async-openai = "0.28.1" schemars = "0.8.22" twox-hash = { version = "2.1.0", features = ["xxhash64"] } async-trait.workspace = true diff --git a/frontend/rust-lib/flowy-ai/src/embeddings/context.rs b/frontend/rust-lib/flowy-ai/src/embeddings/context.rs index 5605a1ea0b40c..fc75f55f91bc3 100644 --- a/frontend/rust-lib/flowy-ai/src/embeddings/context.rs +++ b/frontend/rust-lib/flowy-ai/src/embeddings/context.rs @@ -1,15 +1,17 @@ use crate::embeddings::scheduler::EmbeddingScheduler; +use crate::local_ai::client::AIClient; use arc_swap::ArcSwapOption; use flowy_error::{ErrorCode, FlowyError, FlowyResult}; use flowy_sqlite_vec::db::VectorSqliteDB; use lib_infra::util::get_operating_system; -use ollama_rs::Ollama; use std::path::PathBuf; use std::sync::{Arc, OnceLock}; use tracing::{error, info, warn}; pub struct EmbedContext { - ollama: ArcSwapOption, + client: ArcSwapOption, + /// Embedding model name for the active client. Provider-specific, so it travels with the client. + embedding_model: ArcSwapOption, vector_db: ArcSwapOption, scheduler: ArcSwapOption, } @@ -19,7 +21,8 @@ impl EmbedContext { static INSTANCE: OnceLock> = OnceLock::new(); INSTANCE.get_or_init(|| { Arc::new(EmbedContext { - ollama: ArcSwapOption::empty(), + client: ArcSwapOption::empty(), + embedding_model: ArcSwapOption::empty(), vector_db: ArcSwapOption::empty(), scheduler: Default::default(), }) @@ -50,19 +53,31 @@ impl EmbedContext { } } - pub fn set_ollama(&self, ollama: Option>) { - if let Some(ollama) = ollama { - if let Some(o) = self.ollama.load().as_ref() { - if o.uri() == ollama.uri() { - info!("[Embedding] Ollama does not change"); - return; - } + pub fn set_client(&self, client: Option>, embedding_model: &str) { + if let Some(client) = client { + let unchanged = self + .client + .load() + .as_ref() + .is_some_and(|c| c.same_as(&client)) + && self + .embedding_model + .load() + .as_ref() + .is_some_and(|m| m.as_str() == embedding_model); + if unchanged { + info!("[Embedding] ai client does not change"); + return; } - self.ollama.store(Some(ollama)); + self.client.store(Some(client)); + self + .embedding_model + .store(Some(Arc::new(embedding_model.to_string()))); self.try_create_scheduler(); } else { - self.ollama.store(None); + self.client.store(None); + self.embedding_model.store(None); if let Some(s) = self.scheduler.swap(None) { info!("[Embedding] Stopping scheduler"); let _ = s.stop_tx.send(()); @@ -80,9 +95,13 @@ impl EmbedContext { } fn try_create_scheduler(&self) { - if let (Some(ollama), Some(vector_db)) = (self.ollama.load_full(), self.vector_db.load_full()) { + if let (Some(client), Some(model), Some(vector_db)) = ( + self.client.load_full(), + self.embedding_model.load_full(), + self.vector_db.load_full(), + ) { info!("[Embedding] Creating scheduler"); - match EmbeddingScheduler::new(ollama, vector_db) { + match EmbeddingScheduler::new(client, model.to_string(), vector_db) { Ok(s) => { info!("[Embedding] create scheduler successfully"); self.scheduler.store(Some(s)); @@ -90,7 +109,7 @@ impl EmbedContext { Err(err) => error!("[Embedding] Failed to create scheduler: {}", err), } } else { - info!("[Embedding] Ollama or vector db is not initialized, remove embedding scheduler"); + info!("[Embedding] ai client or vector db is not initialized, remove embedding scheduler"); self.scheduler.store(None); } } diff --git a/frontend/rust-lib/flowy-ai/src/embeddings/document_indexer.rs b/frontend/rust-lib/flowy-ai/src/embeddings/document_indexer.rs index 10b8446f6abdf..0f66ae4213484 100644 --- a/frontend/rust-lib/flowy-ai/src/embeddings/document_indexer.rs +++ b/frontend/rust-lib/flowy-ai/src/embeddings/document_indexer.rs @@ -3,7 +3,6 @@ use crate::embeddings::indexer::{EmbeddingModel, Indexer}; use flowy_ai_pub::entities::{EmbeddedChunk, SOURCE, SOURCE_ID, SOURCE_NAME}; use flowy_error::FlowyError; use lib_infra::async_trait::async_trait; -use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest}; use serde_json::json; use text_splitter::{ChunkConfig, TextSplitter}; use tracing::{debug, error, trace, warn}; @@ -54,25 +53,21 @@ impl Indexer for DocumentIndexer { contents.push(chunks[i].content.as_ref().unwrap().to_owned()); } - let request = GenerateEmbeddingsRequest::new( - embedder.model().name().to_string(), - EmbeddingsInput::Multiple(contents), - ); - let resp = embedder.embed(request).await?; - if resp.embeddings.len() != valid_indices.len() { + let embeddings = embedder.embed(contents).await?; + if embeddings.len() != valid_indices.len() { error!( "[Embedding] requested {} embeddings, received {} embeddings", valid_indices.len(), - resp.embeddings.len() + embeddings.len() ); return Err(FlowyError::internal().with_context(format!( "Mismatch in number of embeddings requested and received: {} vs {}", valid_indices.len(), - resp.embeddings.len() + embeddings.len() ))); } - for (index, embedding) in resp.embeddings.into_iter().enumerate() { + for (index, embedding) in embeddings.into_iter().enumerate() { let chunk_idx = valid_indices[index]; chunks[chunk_idx].embeddings = Some(embedding); } diff --git a/frontend/rust-lib/flowy-ai/src/embeddings/embedder.rs b/frontend/rust-lib/flowy-ai/src/embeddings/embedder.rs index c8c2b629b0bc1..f1f84f6b7ffd6 100644 --- a/frontend/rust-lib/flowy-ai/src/embeddings/embedder.rs +++ b/frontend/rust-lib/flowy-ai/src/embeddings/embedder.rs @@ -1,41 +1,30 @@ -use crate::embeddings::indexer::EmbeddingModel; +use crate::local_ai::client::AIClient; use flowy_error::FlowyResult; -use ollama_rs::Ollama; -use ollama_rs::generation::embeddings::GenerateEmbeddingsResponse; -use ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest; use std::sync::Arc; +/// Pairs the configured backend with the embedding model to use. +/// +/// The model name lives here because it is provider-specific (`nomic-embed-text:latest` for Ollama +/// versus `text-embedding-3-small` for OpenAI-compatible endpoints) and callers should not need to +/// know which backend is active. Regardless of provider the vectors are always +/// [`crate::local_ai::client::EMBEDDING_DIMENSION`] long. #[derive(Debug, Clone)] -pub enum Embedder { - Ollama(OllamaEmbedder), +pub struct Embedder { + client: Arc, + model: String, } impl Embedder { - pub async fn embed( - &self, - request: GenerateEmbeddingsRequest, - ) -> FlowyResult { - match self { - Embedder::Ollama(ollama) => ollama.embed(request).await, - } + pub fn new(client: Arc, model: String) -> Self { + Self { client, model } } - pub fn model(&self) -> EmbeddingModel { - EmbeddingModel::NomicEmbedText + /// Returns one vector per input string, in the same order. + pub async fn embed(&self, input: Vec) -> FlowyResult>> { + self.client.embed(&self.model, input).await } -} - -#[derive(Debug, Clone)] -pub struct OllamaEmbedder { - pub ollama: Arc, -} -impl OllamaEmbedder { - pub async fn embed( - &self, - request: GenerateEmbeddingsRequest, - ) -> FlowyResult { - let resp = self.ollama.generate_embeddings(request).await?; - Ok(resp) + pub fn model_name(&self) -> &str { + &self.model } } diff --git a/frontend/rust-lib/flowy-ai/src/embeddings/scheduler.rs b/frontend/rust-lib/flowy-ai/src/embeddings/scheduler.rs index 08693f2633cf1..aa41be6e5ec9f 100644 --- a/frontend/rust-lib/flowy-ai/src/embeddings/scheduler.rs +++ b/frontend/rust-lib/flowy-ai/src/embeddings/scheduler.rs @@ -1,5 +1,6 @@ -use crate::embeddings::embedder::{Embedder, OllamaEmbedder}; -use crate::embeddings::indexer::IndexerProvider; +use crate::embeddings::embedder::Embedder; +use crate::embeddings::indexer::{EmbeddingModel, IndexerProvider}; +use crate::local_ai::client::AIClient; use crate::search::summary::{LLMDocument, summarize_documents}; use flowy_ai_pub::cloud::search_dto::{ SearchContentType, SearchDocumentResponseItem, SearchResult, SearchSummaryResult, Summary, @@ -8,8 +9,6 @@ use flowy_ai_pub::entities::{EmbeddingRecord, UnindexedCollab, UnindexedData}; use flowy_error::{ErrorCode, FlowyError, FlowyResult}; use flowy_sqlite::internal::derives::multiconnection::chrono::Utc; use flowy_sqlite_vec::db::VectorSqliteDB; -use ollama_rs::Ollama; -use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest}; use std::sync::{Arc, Weak}; use tokio::select; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; @@ -23,14 +22,16 @@ pub struct EmbeddingScheduler { indexer_provider: Arc, write_embedding_tx: UnboundedSender, generate_embedding_tx: mpsc::Sender, - ollama: Arc, + client: Arc, + embedding_model: String, vector_db: Arc, pub(crate) stop_tx: tokio::sync::broadcast::Sender<()>, } impl EmbeddingScheduler { pub fn new( - ollama: Arc, + client: Arc, + embedding_model: String, vector_db: Arc, ) -> FlowyResult> { let indexer_provider = IndexerProvider::new(); @@ -42,7 +43,8 @@ impl EmbeddingScheduler { indexer_provider, write_embedding_tx, generate_embedding_tx, - ollama, + client, + embedding_model, vector_db, stop_tx, }); @@ -67,10 +69,10 @@ impl EmbeddingScheduler { } pub(crate) fn create_embedder(&self) -> Result { - let embedder = Embedder::Ollama(OllamaEmbedder { - ollama: self.ollama.clone(), - }); - Ok(embedder) + Ok(Embedder::new( + self.client.clone(), + self.embedding_model.clone(), + )) } pub async fn index_collab(&self, data: UnindexedCollab) -> FlowyResult<()> { @@ -99,13 +101,8 @@ impl EmbeddingScheduler { query: &str, ) -> FlowyResult> { let embedder = self.create_embedder()?; - let request = GenerateEmbeddingsRequest::new( - embedder.model().name().to_string(), - EmbeddingsInput::Single(query.to_string()), - ); - - let resp = embedder.embed(request).await?; - match resp.embeddings.first() { + let embeddings = embedder.embed(vec![query.to_string()]).await?; + match embeddings.first() { None => Ok(vec![]), Some(query_embed) => { let result = self @@ -155,7 +152,7 @@ impl EmbeddingScheduler { }) .collect::>(); - let resp = summarize_documents(&self.ollama, question, model_name, docs) + let resp = summarize_documents(&self.client, question, model_name, docs) .await .map_err(|err| { error!("[Embedding] Failed to generate summary: {}", err); @@ -295,7 +292,7 @@ async fn spawn_generate_embeddings( match indexer.create_embedded_chunks_from_text( record.object_id, paragraphs, - embedder.model(), + EmbeddingModel::NomicEmbedText, ) { Ok(mut chunks) => { if let Some(fragment_ids) = existing_embeddings.get(&record.object_id) { diff --git a/frontend/rust-lib/flowy-ai/src/embeddings/store.rs b/frontend/rust-lib/flowy-ai/src/embeddings/store.rs index 031d0fb7b802c..2624d78bc64b6 100644 --- a/frontend/rust-lib/flowy-ai/src/embeddings/store.rs +++ b/frontend/rust-lib/flowy-ai/src/embeddings/store.rs @@ -1,7 +1,8 @@ use crate::embeddings::document_indexer::split_text_into_chunks; -use crate::embeddings::embedder::{Embedder, OllamaEmbedder}; +use crate::embeddings::embedder::Embedder; use crate::embeddings::indexer::{EmbeddingModel, IndexerProvider}; use crate::local_ai::chat::retriever::MultipleSourceRetrieverStore; +use crate::local_ai::client::AIClient; use async_trait::async_trait; use flowy_ai_pub::cloud::CollabType; use flowy_ai_pub::entities::{RAG_IDS, SOURCE_ID}; @@ -9,10 +10,8 @@ use flowy_error::{FlowyError, FlowyResult}; use flowy_sqlite_vec::db::VectorSqliteDB; use flowy_sqlite_vec::entities::{EmbeddedContent, SqliteEmbeddedDocument}; use futures::stream::{self, StreamExt}; -use langchain_rust::llm::client::OllamaClient; use langchain_rust::schemas::Document; use langchain_rust::vectorstore::{VecStoreOptions, VectorStore}; -use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest}; use serde_json::Value; use std::collections::HashMap; use std::error::Error; @@ -22,28 +21,33 @@ use uuid::Uuid; #[derive(Clone)] pub struct SqliteVectorStore { - ollama: Weak, + client: Weak, + embedding_model: String, vector_db: Weak, indexer_provider: Arc, } impl SqliteVectorStore { - pub fn new(ollama: Weak, vector_db: Weak) -> Self { + pub fn new( + client: Weak, + embedding_model: String, + vector_db: Weak, + ) -> Self { Self { - ollama, + client, + embedding_model, vector_db, indexer_provider: IndexerProvider::new(), } } pub(crate) fn create_embedder(&self) -> Result { - let ollama = self - .ollama + let client = self + .client .upgrade() - .ok_or_else(|| FlowyError::internal().with_context("Ollama reference was dropped"))?; + .ok_or_else(|| FlowyError::internal().with_context("AI client reference was dropped"))?; - let embedder = Embedder::Ollama(OllamaEmbedder { ollama }); - Ok(embedder) + Ok(Embedder::new(client, self.embedding_model.clone())) } pub(crate) async fn select_all_embedded_documents( @@ -107,12 +111,7 @@ impl MultipleSourceRetrieverStore for SqliteVectorStore { // Create embedder and generate embedding for query let embedder = self.create_embedder()?; - let request = GenerateEmbeddingsRequest::new( - embedder.model().name().to_string(), - EmbeddingsInput::Single(query.to_string()), - ); - - let embedding = embedder.embed(request).await?.embeddings; + let embedding = embedder.embed(vec![query.to_string()]).await?; if embedding.is_empty() { return Ok(Vec::new()); } diff --git a/frontend/rust-lib/flowy-ai/src/entities.rs b/frontend/rust-lib/flowy-ai/src/entities.rs index f63f4bc22d48c..2aaa80fc5e840 100644 --- a/frontend/rust-lib/flowy-ai/src/entities.rs +++ b/frontend/rust-lib/flowy-ai/src/entities.rs @@ -1,3 +1,4 @@ +use crate::local_ai::client::AIProvider; use crate::local_ai::controller::LocalAISetting; use crate::local_ai::resource::PendingResource; use flowy_ai_pub::cloud::{ @@ -671,8 +672,37 @@ impl From for ResponseFormat { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ProtoBuf_Enum)] +pub enum AIProviderPB { + /// Talks to an Ollama server; `server_url` is the Ollama host and `api_key` is unused. + #[default] + Ollama = 0, + /// Talks to any endpoint implementing the OpenAI REST API. `server_url` is the API base + /// (for example `https://api.openai.com/v1`) and `api_key` is sent as the bearer token. + OpenAI = 1, +} + +impl From for AIProviderPB { + fn from(value: AIProvider) -> Self { + match value { + AIProvider::Ollama => AIProviderPB::Ollama, + AIProvider::OpenAI => AIProviderPB::OpenAI, + } + } +} + +impl From for AIProvider { + fn from(value: AIProviderPB) -> Self { + match value { + AIProviderPB::Ollama => AIProvider::Ollama, + AIProviderPB::OpenAI => AIProvider::OpenAI, + } + } +} + #[derive(Default, ProtoBuf, Validate, Clone, Debug)] pub struct LocalAISettingPB { + /// Ollama host, or the OpenAI-compatible API base, depending on `provider`. #[pb(index = 1)] #[validate(custom(function = "required_not_empty_str"))] pub server_url: String, @@ -684,6 +714,14 @@ pub struct LocalAISettingPB { #[pb(index = 3)] #[validate(custom(function = "required_not_empty_str"))] pub embedding_model_name: String, + + #[pb(index = 4)] + pub provider: AIProviderPB, + + /// Bearer token for OpenAI-compatible providers. Not validated: local servers such as LM Studio + /// and llama.cpp accept requests without one. + #[pb(index = 5)] + pub api_key: String, } impl From for LocalAISettingPB { @@ -692,6 +730,8 @@ impl From for LocalAISettingPB { server_url: value.ollama_server_url, global_chat_model: value.chat_model_name, embedding_model_name: value.embedding_model_name, + provider: value.provider.into(), + api_key: value.api_key, } } } @@ -702,6 +742,8 @@ impl From for LocalAISetting { ollama_server_url: value.server_url, chat_model_name: value.global_chat_model, embedding_model_name: value.embedding_model_name, + provider: value.provider.into(), + api_key: value.api_key, } } } diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm.rs b/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm.rs index 0fbb639e0e8c8..e78fbe8201f02 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm.rs @@ -1,47 +1,40 @@ use async_trait::async_trait; use futures::Stream; use langchain_rust::language_models::llm::LLM; -use langchain_rust::language_models::{GenerateResult, LLMError, TokenUsage}; +use langchain_rust::language_models::{GenerateResult, LLMError}; use langchain_rust::schemas::{Message, StreamData}; -use ollama_rs::error::OllamaError; use std::pin::Pin; use std::sync::Arc; -use tokio_stream::StreamExt; -use ollama_rs::Ollama; -use ollama_rs::generation::chat::request::ChatMessageRequest; +use crate::local_ai::client::AIClient; use ollama_rs::generation::parameters::FormatType; use ollama_rs::models::ModelOptions; +/// The [`LLM`] implementation backing every chain in this crate. +/// +/// The name is historical: it now dispatches to whichever backend the user configured — Ollama or +/// an OpenAI-compatible endpoint — through [`AIClient`]. Keeping this type concrete rather than +/// making the chains generic over the backend is what keeps the provider change contained to this +/// file. #[derive(Debug, Clone)] pub struct LLMOllama { pub model_name: String, - ollama: Arc, + client: Arc, format: Option, + /// Ollama-specific sampling parameters. The OpenAI-compatible backend ignores these. options: Option, } -impl Default for LLMOllama { - fn default() -> Self { - LLMOllama { - model_name: "llama3.1".to_string(), - ollama: Arc::new(Ollama::default()), - format: None, - options: None, - } - } -} - impl LLMOllama { pub fn new( model: &str, - ollama: Arc, + client: Arc, format: Option, options: Option, ) -> Self { LLMOllama { model_name: model.to_string(), - ollama, + client, format, options, } @@ -65,55 +58,34 @@ impl LLMOllama { pub fn set_model(&mut self, model: &str) { self.model_name = model.to_string(); } - - fn generate_request(&self, messages: &[Message]) -> ChatMessageRequest { - let mapped_messages = messages.iter().map(|message| message.into()).collect(); - let mut request = ChatMessageRequest::new(self.model_name.clone(), mapped_messages); - if let Some(option) = &self.options { - request = request.options(option.clone()) - } - if let Some(format) = &self.format { - request = request.format(format.clone()); - } - request - } } #[async_trait] impl LLM for LLMOllama { async fn generate(&self, messages: &[Message]) -> Result { - let request = self.generate_request(messages); - let result = self.ollama.send_chat_messages(request).await?; - let generation = result.message.content; - let tokens = result.final_data.map(|final_data| { - let prompt_tokens = final_data.prompt_eval_count as u32; - let completion_tokens = final_data.eval_count as u32; - TokenUsage { - prompt_tokens, - completion_tokens, - total_tokens: prompt_tokens + completion_tokens, - } - }); - - Ok(GenerateResult { tokens, generation }) + self + .client + .chat( + &self.model_name, + messages, + self.format.as_ref(), + self.options.as_ref(), + ) + .await } async fn stream( &self, messages: &[Message], ) -> Result> + Send>>, LLMError> { - let request = self.generate_request(messages); - let result = self.ollama.send_chat_messages_stream(request).await?; - - let stream = result.map(|data| match data { - Ok(data) => Ok(StreamData::new( - serde_json::to_value(&data).unwrap_or_default(), - None, - data.message.content, - )), - Err(_) => Err(OllamaError::Other("Stream error".to_string()).into()), - }); - - Ok(Box::pin(stream)) + self + .client + .chat_stream( + &self.model_name, + messages, + self.format.as_ref(), + self.options.as_ref(), + ) + .await } } diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm_chat.rs b/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm_chat.rs index 5ea86c0a018e2..0c225bf6e91fb 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm_chat.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/chat/llm_chat.rs @@ -19,7 +19,7 @@ use langchain_rust::memory::SimpleMemory; use langchain_rust::prompt_args; use langchain_rust::schemas::{Document, Message}; use langchain_rust::vectorstore::{VecStoreOptions, VectorStore}; -use ollama_rs::Ollama; +use crate::local_ai::client::AIClient; use serde_json::json; use std::collections::HashMap; use std::sync::{Arc, Weak}; @@ -30,7 +30,7 @@ pub struct LLMChat { store: Option, chain: ConversationalRetrieverChain, #[allow(dead_code)] - client: Arc, + client: Arc, prompt: AFContextPrompt, info: LLMChatInfo, } @@ -38,7 +38,7 @@ pub struct LLMChat { impl LLMChat { pub fn new( info: LLMChatInfo, - client: Arc, + client: Arc, store: Option, user_service: Option>, retriever_sources: Vec>, diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/chat/mod.rs b/frontend/rust-lib/flowy-ai/src/local_ai/chat/mod.rs index e89910f39875d..f3b22ac63ba34 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/chat/mod.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/chat/mod.rs @@ -23,7 +23,7 @@ use flowy_ai_pub::user_service::AIUserService; use flowy_database_pub::cloud::{SummaryRowContent, TranslateRowContent}; use flowy_error::{FlowyError, FlowyResult}; use futures_util::StreamExt; -use ollama_rs::Ollama; +use crate::local_ai::client::AIClient; use serde_json::Value; use std::collections::HashMap; use std::path::PathBuf; @@ -32,7 +32,7 @@ use tokio::sync::RwLock; use tracing::warn; use uuid::Uuid; -type OllamaClientRef = Arc>>>; +type AIClientRef = Arc>>>; pub struct LLMChatInfo { pub chat_id: Uuid, @@ -47,7 +47,7 @@ pub type RetrieversSources = RwLock>>; pub struct LLMChatController { chat_by_id: DashMap>>, store: RwLock>, - client: OllamaClientRef, + client: AIClientRef, user_service: Weak, retriever_sources: RetrieversSources, } @@ -73,11 +73,12 @@ impl LLMChatController { #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))] pub async fn initialize( &self, - ollama: Weak, + client: Weak, + embedding_model: String, vector_db: Weak, ) { - let store = SqliteVectorStore::new(ollama.clone(), vector_db); - *self.client.write().await = Some(ollama); + let store = SqliteVectorStore::new(client.clone(), embedding_model, vector_db); + *self.client.write().await = Some(client); *self.store.write().await = Some(store); } @@ -94,9 +95,9 @@ impl LLMChatController { .read() .await .as_ref() - .ok_or_else(|| FlowyError::local_ai().with_context("Ollama client not initialized"))? + .ok_or_else(|| FlowyError::local_ai().with_context("AI client not initialized"))? .upgrade() - .ok_or_else(|| FlowyError::local_ai().with_context("Ollama client has been dropped"))? + .ok_or_else(|| FlowyError::local_ai().with_context("AI client has been dropped"))? .clone(); let entry = self.chat_by_id.entry(info.chat_id); let retriever_sources = self diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/client.rs b/frontend/rust-lib/flowy-ai/src/local_ai/client.rs new file mode 100644 index 0000000000000..aab503617dba0 --- /dev/null +++ b/frontend/rust-lib/flowy-ai/src/local_ai/client.rs @@ -0,0 +1,449 @@ +use async_openai::Client as OpenAIHttpClient; +use async_openai::config::OpenAIConfig; +use async_openai::types::{ + ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestMessage, + ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestUserMessageArgs, + CreateChatCompletionRequestArgs, CreateEmbeddingRequestArgs, EmbeddingInput, ResponseFormat, +}; +use flowy_error::{FlowyError, FlowyResult}; +use futures::Stream; +use langchain_rust::language_models::{GenerateResult, LLMError, TokenUsage}; +use langchain_rust::schemas::{Message, MessageType, StreamData}; +use ollama_rs::Ollama; +use ollama_rs::generation::chat::request::ChatMessageRequest; +use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest}; +use ollama_rs::generation::parameters::FormatType; +use ollama_rs::models::ModelOptions; +use serde::{Deserialize, Serialize}; +use std::pin::Pin; +use std::sync::Arc; +use tokio_stream::StreamExt; + +/// The vector store column is declared as `float[768]` (see +/// flowy-sqlite-vec/migrations/001-init/up.sql), so every provider must yield vectors of exactly +/// this size or the embeddings cannot be stored. +pub const EMBEDDING_DIMENSION: usize = 768; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum AIProvider { + #[default] + Ollama, + /// Any endpoint speaking the OpenAI REST API: OpenAI itself, LM Studio, vLLM, llama.cpp, + /// OpenRouter, Groq, Together, ... + OpenAI, +} + +impl AIProvider { + pub fn from_index(index: i64) -> Self { + match index { + 1 => AIProvider::OpenAI, + _ => AIProvider::Ollama, + } + } + + pub fn index(&self) -> i64 { + match self { + AIProvider::Ollama => 0, + AIProvider::OpenAI => 1, + } + } +} + +/// A chat/embedding backend. Both variants expose the same operations so the rest of the crate +/// never branches on the provider. +pub enum AIClient { + Ollama(Arc), + OpenAI { + client: Box>, + base_url: String, + api_key: String, + }, +} + +impl std::fmt::Debug for AIClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AIClient::Ollama(o) => write!(f, "AIClient::Ollama({})", o.url_str()), + AIClient::OpenAI { base_url, .. } => write!(f, "AIClient::OpenAI({})", base_url), + } + } +} + +impl AIClient { + pub fn new(provider: AIProvider, base_url: &str, api_key: &str) -> FlowyResult { + match provider { + AIProvider::Ollama => { + let ollama = Ollama::try_new(base_url).map_err(|err| { + FlowyError::local_ai().with_context(format!("invalid ollama url {}: {}", base_url, err)) + })?; + Ok(AIClient::Ollama(Arc::new(ollama))) + }, + AIProvider::OpenAI => { + // Trailing slashes produce `//chat/completions` on some gateways. + let base_url = base_url.trim_end_matches('/').to_string(); + let mut config = OpenAIConfig::new().with_api_base(&base_url); + // Local servers (LM Studio, llama.cpp) usually accept any key, but async-openai always + // sends the header, so only override when the user actually supplied one. + if !api_key.is_empty() { + config = config.with_api_key(api_key); + } + Ok(AIClient::OpenAI { + client: Box::new(OpenAIHttpClient::with_config(config)), + base_url, + api_key: api_key.to_string(), + }) + }, + } + } + + pub fn base_url(&self) -> String { + match self { + AIClient::Ollama(o) => o.url_str().to_string(), + AIClient::OpenAI { base_url, .. } => base_url.clone(), + } + } + + pub fn provider(&self) -> AIProvider { + match self { + AIClient::Ollama(_) => AIProvider::Ollama, + AIClient::OpenAI { .. } => AIProvider::OpenAI, + } + } + + /// Whether this client is already configured exactly as requested, so callers can skip a reload. + /// The url crate normalises `http://host:1234` to a trailing slash, hence the trim on both sides. + pub fn matches(&self, provider: AIProvider, base_url: &str, api_key: &str) -> bool { + if self.provider() != provider { + return false; + } + if self.base_url().trim_end_matches('/') != base_url.trim_end_matches('/') { + return false; + } + match self { + AIClient::Ollama(_) => true, + AIClient::OpenAI { api_key: key, .. } => key == api_key, + } + } + + /// Whether two clients are configured identically. + pub fn same_as(&self, other: &AIClient) -> bool { + match other { + AIClient::Ollama(_) => self.matches(other.provider(), &other.base_url(), ""), + AIClient::OpenAI { api_key, .. } => self.matches(other.provider(), &other.base_url(), api_key), + } + } + + pub fn as_ollama(&self) -> Option<&Arc> { + match self { + AIClient::Ollama(o) => Some(o), + AIClient::OpenAI { .. } => None, + } + } + + /// Model names available on the backend. + pub async fn list_models(&self) -> FlowyResult> { + match self { + AIClient::Ollama(ollama) => { + let models = ollama + .list_local_models() + .await + .map_err(|err| FlowyError::local_ai().with_context(err.to_string()))?; + Ok(models.into_iter().map(|m| m.name).collect()) + }, + AIClient::OpenAI { client, .. } => { + let models = client + .models() + .list() + .await + .map_err(|err| FlowyError::local_ai().with_context(err.to_string()))?; + Ok(models.data.into_iter().map(|m| m.id).collect()) + }, + } + } + + /// Embed `input`, returning one vector per input string. Every vector is verified to be + /// [`EMBEDDING_DIMENSION`] long so a mis-sized model fails here rather than at insert time. + pub async fn embed(&self, model: &str, input: Vec) -> FlowyResult>> { + let embeddings = match self { + AIClient::Ollama(ollama) => { + let request = + GenerateEmbeddingsRequest::new(model.to_string(), EmbeddingsInput::Multiple(input)); + ollama + .generate_embeddings(request) + .await + .map_err(|err| FlowyError::local_ai().with_context(err.to_string()))? + .embeddings + }, + AIClient::OpenAI { client, .. } => { + let request = CreateEmbeddingRequestArgs::default() + .model(model) + .input(EmbeddingInput::StringArray(input)) + // text-embedding-3-* are Matryoshka models and truncate natively, which lets us reuse + // the existing float[768] column instead of migrating the vector store. + .dimensions(EMBEDDING_DIMENSION as u32) + .build() + .map_err(|err| FlowyError::local_ai().with_context(err.to_string()))?; + + let response = client + .embeddings() + .create(request) + .await + .map_err(|err| FlowyError::local_ai().with_context(err.to_string()))?; + + let mut data = response.data; + // The API does not guarantee ordering; `index` is authoritative. + data.sort_by_key(|e| e.index); + data.into_iter().map(|e| e.embedding).collect() + }, + }; + + if let Some(actual) = embeddings.iter().map(|e| e.len()).find(|n| *n != EMBEDDING_DIMENSION) { + return Err(FlowyError::local_ai().with_context(format!( + "embedding model '{}' returned {}-dimensional vectors, but {} are required. \ + Pick a model that supports {0} dimensions, or one that honours the `dimensions` parameter \ + (for example text-embedding-3-small).", + model, actual, EMBEDDING_DIMENSION + ))); + } + + Ok(embeddings) + } + + pub async fn chat( + &self, + model: &str, + messages: &[Message], + format: Option<&FormatType>, + options: Option<&ModelOptions>, + ) -> Result { + match self { + AIClient::Ollama(ollama) => { + let request = ollama_request(model, messages, format, options); + let result = ollama.send_chat_messages(request).await?; + let tokens = result.final_data.map(|final_data| { + let prompt_tokens = final_data.prompt_eval_count as u32; + let completion_tokens = final_data.eval_count as u32; + TokenUsage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + } + }); + Ok(GenerateResult { + tokens, + generation: result.message.content, + }) + }, + AIClient::OpenAI { client, .. } => { + let request = openai_request(model, messages, format, false)?; + let response = client.chat().create(request).await?; + let generation = response + .choices + .into_iter() + .next() + .and_then(|c| c.message.content) + .unwrap_or_default(); + let tokens = response.usage.map(|usage| TokenUsage { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + total_tokens: usage.total_tokens, + }); + Ok(GenerateResult { tokens, generation }) + }, + } + } + + pub async fn chat_stream( + &self, + model: &str, + messages: &[Message], + format: Option<&FormatType>, + options: Option<&ModelOptions>, + ) -> Result> + Send>>, LLMError> { + match self { + AIClient::Ollama(ollama) => { + let request = ollama_request(model, messages, format, options); + let result = ollama.send_chat_messages_stream(request).await?; + let stream = result.map(|data| match data { + Ok(data) => Ok(StreamData::new( + serde_json::to_value(&data).unwrap_or_default(), + None, + data.message.content, + )), + Err(_) => Err(ollama_rs::error::OllamaError::Other("Stream error".to_string()).into()), + }); + Ok(Box::pin(stream)) + }, + AIClient::OpenAI { client, .. } => { + let request = openai_request(model, messages, format, true)?; + let stream = client.chat().create_stream(request).await?; + let stream = stream.map(|chunk| match chunk { + Ok(chunk) => { + let content = chunk + .choices + .first() + .and_then(|c| c.delta.content.clone()) + .unwrap_or_default(); + Ok(StreamData::new( + serde_json::to_value(&chunk).unwrap_or_default(), + None, + content, + )) + }, + Err(err) => Err(LLMError::OpenAIError(err)), + }); + Ok(Box::pin(stream)) + }, + } + } +} + +fn ollama_request( + model: &str, + messages: &[Message], + format: Option<&FormatType>, + options: Option<&ModelOptions>, +) -> ChatMessageRequest { + let mapped = messages.iter().map(|message| message.into()).collect(); + let mut request = ChatMessageRequest::new(model.to_string(), mapped); + if let Some(options) = options { + request = request.options(options.clone()); + } + if let Some(format) = format { + request = request.format(format.clone()); + } + request +} + +fn openai_request( + model: &str, + messages: &[Message], + format: Option<&FormatType>, + stream: bool, +) -> Result { + let mut mapped = messages + .iter() + .map(openai_message) + .collect::, LLMError>>()?; + + let mut builder = CreateChatCompletionRequestArgs::default(); + builder.model(model).stream(stream); + + if let Some(format) = format { + // Ollama takes a JSON schema directly; the OpenAI-compatible equivalent that works across + // gateways is `json_object` plus the schema in the prompt. Strict `json_schema` mode is + // rejected by many compatible servers. + builder.response_format(ResponseFormat::JsonObject); + if let Ok(schema) = serde_json::to_value(format) { + mapped.push( + ChatCompletionRequestSystemMessageArgs::default() + .content(format!( + "Respond with a single JSON object matching this JSON Schema. \ + Output JSON only, with no markdown fences or commentary.\n{}", + schema + )) + .build()? + .into(), + ); + } + } + + builder.messages(mapped); + Ok(builder.build()?) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::local_ai::controller::LocalAISetting; + + /// Settings written before OpenAI-compatible support existed must still load, otherwise every + /// existing install silently resets to defaults on upgrade. + #[test] + fn settings_stored_before_provider_support_still_load() { + let stored = r#"{"ollama_server_url":"http://localhost:11434", + "chat_model_name":"llama3.1:latest", + "embedding_model_name":"nomic-embed-text:latest"}"#; + + let setting: LocalAISetting = serde_json::from_str(stored).unwrap(); + assert_eq!(setting.provider, AIProvider::Ollama); + assert!(setting.api_key.is_empty()); + assert_eq!(setting.ollama_server_url, "http://localhost:11434"); + assert_eq!(setting.chat_model_name, "llama3.1:latest"); + } + + #[test] + fn client_identity_tracks_provider_url_and_key() { + let client = AIClient::new(AIProvider::OpenAI, "https://api.openai.com/v1", "k1").unwrap(); + + // A trailing slash is the same endpoint and must not force a reload. + assert!(client.matches(AIProvider::OpenAI, "https://api.openai.com/v1/", "k1")); + // A rotated key must force a reload, or the old key keeps being used. + assert!(!client.matches(AIProvider::OpenAI, "https://api.openai.com/v1", "k2")); + assert!(!client.matches(AIProvider::Ollama, "https://api.openai.com/v1", "k1")); + } + + /// Same URL, different protocol: these are not interchangeable clients. + #[test] + fn ollama_and_openai_at_one_url_are_distinct() { + let ollama = AIClient::new(AIProvider::Ollama, "http://localhost:11434", "").unwrap(); + let openai = AIClient::new(AIProvider::OpenAI, "http://localhost:11434", "").unwrap(); + assert!(!ollama.same_as(&openai)); + assert!(!openai.same_as(&ollama)); + } + + /// Ollama takes a JSON schema natively; the OpenAI path has to carry it in the prompt instead. + #[test] + fn structured_output_becomes_json_mode_plus_schema_message() { + let messages = vec![Message::new_human_message("hi")]; + let format = FormatType::Json; + + let plain = openai_request("gpt-4o", &messages, None, false).unwrap(); + assert_eq!(plain.messages.len(), 1); + assert!(plain.response_format.is_none()); + + let structured = openai_request("gpt-4o", &messages, Some(&format), false).unwrap(); + assert_eq!(structured.messages.len(), 2); + assert!(matches!( + structured.response_format, + Some(ResponseFormat::JsonObject) + )); + } + + #[test] + fn message_roles_map_to_openai_roles() { + let mapped = openai_message(&Message::new_system_message("s")).unwrap(); + assert!(matches!( + mapped, + ChatCompletionRequestMessage::System(_) + )); + let mapped = openai_message(&Message::new_ai_message("a")).unwrap(); + assert!(matches!( + mapped, + ChatCompletionRequestMessage::Assistant(_) + )); + let mapped = openai_message(&Message::new_human_message("h")).unwrap(); + assert!(matches!(mapped, ChatCompletionRequestMessage::User(_))); + } +} + +fn openai_message(message: &Message) -> Result { + let content = message.content.clone(); + Ok(match message.message_type { + MessageType::SystemMessage => ChatCompletionRequestSystemMessageArgs::default() + .content(content) + .build()? + .into(), + MessageType::AIMessage => ChatCompletionRequestAssistantMessageArgs::default() + .content(content) + .build()? + .into(), + // Tool results carry no tool_call_id here, so surface them as user turns rather than dropping. + MessageType::HumanMessage | MessageType::ToolMessage => { + ChatCompletionRequestUserMessageArgs::default() + .content(content) + .build()? + .into() + }, + }) +} diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/controller.rs b/frontend/rust-lib/flowy-ai/src/local_ai/controller.rs index 64943d1b7eb63..f13b43b21f1f4 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/controller.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/controller.rs @@ -11,6 +11,7 @@ use lib_infra::async_trait::async_trait; use std::collections::HashMap; use crate::local_ai::chat::{LLMChatController, LLMChatInfo}; +use crate::local_ai::client::{AIClient, AIProvider}; use crate::stream_message::StreamMessage; use arc_swap::ArcSwapOption; use flowy_ai_pub::cloud::AIModel; @@ -20,8 +21,6 @@ use flowy_ai_pub::persistence::{ use flowy_ai_pub::user_service::AIUserService; use futures_util::SinkExt; use lib_infra::util::get_operating_system; -use ollama_rs::Ollama; -use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest}; use serde::{Deserialize, Serialize}; use std::ops::Deref; use std::path::PathBuf; @@ -34,6 +33,12 @@ pub struct LocalAISetting { pub ollama_server_url: String, pub chat_model_name: String, pub embedding_model_name: String, + /// `#[serde(default)]` on the new fields keeps settings stored before OpenAI-compatible support + /// existed deserializable; without it every existing user silently falls back to defaults. + #[serde(default)] + pub provider: AIProvider, + #[serde(default)] + pub api_key: String, } impl Default for LocalAISetting { @@ -42,6 +47,8 @@ impl Default for LocalAISetting { ollama_server_url: "http://localhost:11434".to_string(), chat_model_name: "llama3.1:latest".to_string(), embedding_model_name: "nomic-embed-text:latest".to_string(), + provider: AIProvider::Ollama, + api_key: String::new(), } } } @@ -54,7 +61,7 @@ pub struct LocalAIController { current_chat_id: ArcSwapOption, store_preferences: Weak, user_service: Arc, - pub(crate) ollama: ArcSwapOption, + pub(crate) ollama: ArcSwapOption, } impl Deref for LocalAIController { @@ -100,34 +107,48 @@ impl LocalAIController { if !self.is_enabled_on_workspace(workspace_id) { #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))] { - trace!("[Local AI] local ai is disabled, clear ollama client",); + trace!("[Local AI] local ai is disabled, clear ai client",); let shared = crate::embeddings::context::EmbedContext::shared(); - shared.set_ollama(None); + shared.set_client(None, ""); self.ollama.store(None); } return; } let setting = self.resource.get_llm_setting(); - if let Some(ollama) = self.ollama.load_full() { - if ollama.url_str() == setting.ollama_server_url { - info!("[Local AI] ollama client is already initialized"); + if let Some(client) = self.ollama.load_full() { + if client.matches( + setting.provider, + &setting.ollama_server_url, + &setting.api_key, + ) { + info!("[Local AI] ai client is already initialized"); return; } } - info!("[Local AI] reloading ollama client"); - match Ollama::try_new(&setting.ollama_server_url).map(Arc::new) { + info!("[Local AI] reloading {:?} client", setting.provider); + match AIClient::new( + setting.provider, + &setting.ollama_server_url, + &setting.api_key, + ) + .map(Arc::new) + { Ok(new_ollama) => { #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))] { - info!("[Local AI] reload ollama client successfully"); + info!("[Local AI] reload ai client successfully"); let shared = crate::embeddings::context::EmbedContext::shared(); - shared.set_ollama(Some(new_ollama.clone())); + shared.set_client(Some(new_ollama.clone()), &setting.embedding_model_name); if let Some(vc) = shared.get_vector_db() { self .llm_controller - .initialize(Arc::downgrade(&new_ollama), Arc::downgrade(&vc)) + .initialize( + Arc::downgrade(&new_ollama), + setting.embedding_model_name.clone(), + Arc::downgrade(&vc), + ) .await; } else { error!("[Local AI] vector db is not initialized"); @@ -255,14 +276,14 @@ impl LocalAIController { { match self.ollama.load_full() { None => vec![], - Some(ollama) => ollama - .list_local_models() + Some(client) => client + .list_models() .await .map(|models| { models .into_iter() - .filter(|m| filter_fn(&m.name.to_lowercase())) - .map(|m| AIModel::local(m.name, String::new())) + .filter(|name| filter_fn(&name.to_lowercase())) + .map(|name| AIModel::local(name, String::new())) .collect() }) .unwrap_or_default(), @@ -274,25 +295,17 @@ impl LocalAIController { let mut conn = self.user_service.sqlite_connection(uid)?; match select_local_ai_model(&mut conn, model_name) { None => { - let ollama = self + let client = self .ollama .load_full() - .ok_or_else(|| FlowyError::local_ai().with_context("ollama is not initialized"))?; - - let request = GenerateEmbeddingsRequest::new( - model_name.to_string(), - EmbeddingsInput::Single("Hello".to_string()), - ); - - let model_type = match ollama.generate_embeddings(request).await { - Ok(value) => { - if value.embeddings.is_empty() { - ModelType::Chat - } else { - ModelType::Embedding - } - }, - Err(_) => ModelType::Chat, + .ok_or_else(|| FlowyError::local_ai().with_context("ai client is not initialized"))?; + + // Probe by embedding: a chat model rejects the embeddings endpoint. `embed` also enforces + // the 768-dimension requirement, so a wrongly-sized embedding model is classified as Chat + // rather than being silently accepted and failing later at insert time. + let model_type = match client.embed(model_name, vec!["Hello".to_string()]).await { + Ok(embeddings) if !embeddings.is_empty() => ModelType::Embedding, + _ => ModelType::Chat, }; upsert_local_ai_model( diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/mod.rs b/frontend/rust-lib/flowy-ai/src/local_ai/mod.rs index a040ec5d234c0..9e9f833759fec 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/mod.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/mod.rs @@ -1,3 +1,4 @@ +pub mod client; pub mod controller; mod request; pub mod resource; diff --git a/frontend/rust-lib/flowy-ai/src/model_select.rs b/frontend/rust-lib/flowy-ai/src/model_select.rs index 43cda1a3dfeb1..1acf2e1050381 100644 --- a/frontend/rust-lib/flowy-ai/src/model_select.rs +++ b/frontend/rust-lib/flowy-ai/src/model_select.rs @@ -305,14 +305,14 @@ impl ModelSource for LocalAiSource { async fn list_chat_models(&self, _workspace_id: &Uuid) -> Vec { match self.controller.ollama.load_full() { None => vec![], - Some(ollama) => ollama - .list_local_models() + Some(client) => client + .list_models() .await .map(|models| { models .into_iter() - .filter(|m| !m.name.contains("embed")) - .map(|m| AIModel::local(m.name, String::new())) + .filter(|name| !name.contains("embed")) + .map(|name| AIModel::local(name, String::new())) .collect() }) .unwrap_or_default(), diff --git a/frontend/rust-lib/flowy-ai/src/search/summary.rs b/frontend/rust-lib/flowy-ai/src/search/summary.rs index ee04d8655a84a..7d0090fb0e562 100644 --- a/frontend/rust-lib/flowy-ai/src/search/summary.rs +++ b/frontend/rust-lib/flowy-ai/src/search/summary.rs @@ -1,7 +1,6 @@ +use crate::local_ai::client::AIClient; use flowy_error::FlowyError; -use ollama_rs::Ollama; -use ollama_rs::generation::chat::request::ChatMessageRequest; -use ollama_rs::generation::chat::{ChatMessage, MessageRole}; +use langchain_rust::schemas::Message; use ollama_rs::generation::parameters::{FormatType, JsonStructure}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -60,7 +59,7 @@ fn convert_documents_to_text(documents: Vec) -> String { } pub async fn summarize_documents( - client: &Ollama, + client: &AIClient, question: &str, model_name: &str, documents: Vec, @@ -68,32 +67,30 @@ pub async fn summarize_documents( let documents_text = convert_documents_to_text(documents); let context = format!("{}\n\n##Context##\n{}", SYSTEM_PROMPT, documents_text); let messages = vec![ - ChatMessage::new(MessageRole::System, context), - ChatMessage::new(MessageRole::User, question.to_string()), + Message::new_system_message(context), + Message::new_human_message(question), ]; let format = FormatType::StructuredJson(JsonStructure::new::()); - let request = ChatMessageRequest::new(model_name.to_string(), messages).format(format); - match client.send_chat_messages(request).await { - Ok(resp) => { - if resp.final_data.is_some() { - let resp: SummarySearchSchema = serde_json::from_str(&resp.message.content)?; - let resp = SummarySearchResponse { - summaries: vec![SearchSummary { - content: resp.answer, - highlights: resp.highlights, - sources: resp - .sources - .into_iter() - .flat_map(|s| Uuid::parse_str(&s).ok()) - .collect(), - }], - }; - Ok(resp) - } else { - Ok(SummarySearchResponse { summaries: vec![] }) - } + match client + .chat(model_name, &messages, Some(&format), None) + .await + { + Ok(result) if !result.generation.trim().is_empty() => { + let resp: SummarySearchSchema = serde_json::from_str(&result.generation)?; + Ok(SummarySearchResponse { + summaries: vec![SearchSummary { + content: resp.answer, + highlights: resp.highlights, + sources: resp + .sources + .into_iter() + .flat_map(|s| Uuid::parse_str(&s).ok()) + .collect(), + }], + }) }, + Ok(_) => Ok(SummarySearchResponse { summaries: vec![] }), Err(err) => { error!("Error generating summary: {}", err); Ok(SummarySearchResponse { summaries: vec![] }) @@ -103,12 +100,12 @@ pub async fn summarize_documents( #[cfg(test)] mod tests { + use crate::local_ai::client::{AIClient, AIProvider}; use crate::search::summary::{LLMDocument, summarize_documents}; - use ollama_rs::Ollama; #[tokio::test] async fn summarize_documents_test() { - let ollama = Ollama::try_new("http://localhost:11434").unwrap(); + let ollama = AIClient::new(AIProvider::Ollama, "http://localhost:11434", "").unwrap(); let docs = vec![ ("Rust is a multiplayer survival game developed by Facepunch Studios, first released in early access in December 2013 and fully launched in February 2018. It has since become one of the most popular games in the survival genre, known for its harsh environment, intricate crafting system, and player-driven dynamics. The game is available on Windows, macOS, and PlayStation, with a community-driven approach to updates and content additions.", uuid::Uuid::new_v4()), ("Rust is a modern, system-level programming language designed with a focus on performance, safety, and concurrency. It was created by Mozilla and first released in 2010, with its 1.0 version launched in 2015. Rust is known for providing the control and performance of languages like C and C++, but with built-in safety features that prevent common programming errors, such as memory leaks, data races, and buffer overflows.", uuid::Uuid::new_v4()), From e9ccf2d62d95ac1b5de52ca674c7493adeb2a143 Mon Sep 17 00:00:00 2001 From: imnotdev25 <85677268+imnotdev25@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:49:50 +0530 Subject: [PATCH 2/2] fix: report required dimension in embedding size error, drop dead provider mapping The mismatch message mixed implicit `{}` with an explicit `{0}`, so the final clause resolved to the first argument and told users to "pick a model that supports text-embedding-3-large dimensions" instead of naming 768. Valid Rust that compiles and never panics, but the guidance it printed was wrong. The check moves into `ensure_dimensions` so the guard and its message can be tested without a live endpoint, covered by two new tests. Verified the new test fails against the previous format string before fixing it. Also removes AIProvider::from_index/index. They had no call sites, and deleting them leaves AIProviderPB as the single provider index mapping rather than a second one that could drift. Co-Authored-By: Claude Opus 4.8 --- .../rust-lib/flowy-ai/src/local_ai/client.rs | 71 ++++++++++++------- 1 file changed, 46 insertions(+), 25 deletions(-) diff --git a/frontend/rust-lib/flowy-ai/src/local_ai/client.rs b/frontend/rust-lib/flowy-ai/src/local_ai/client.rs index aab503617dba0..d8f672018baf7 100644 --- a/frontend/rust-lib/flowy-ai/src/local_ai/client.rs +++ b/frontend/rust-lib/flowy-ai/src/local_ai/client.rs @@ -33,22 +33,6 @@ pub enum AIProvider { OpenAI, } -impl AIProvider { - pub fn from_index(index: i64) -> Self { - match index { - 1 => AIProvider::OpenAI, - _ => AIProvider::Ollama, - } - } - - pub fn index(&self) -> i64 { - match self { - AIProvider::Ollama => 0, - AIProvider::OpenAI => 1, - } - } -} - /// A chat/embedding backend. Both variants expose the same operations so the rest of the crate /// never branches on the provider. pub enum AIClient { @@ -197,15 +181,7 @@ impl AIClient { }, }; - if let Some(actual) = embeddings.iter().map(|e| e.len()).find(|n| *n != EMBEDDING_DIMENSION) { - return Err(FlowyError::local_ai().with_context(format!( - "embedding model '{}' returned {}-dimensional vectors, but {} are required. \ - Pick a model that supports {0} dimensions, or one that honours the `dimensions` parameter \ - (for example text-embedding-3-small).", - model, actual, EMBEDDING_DIMENSION - ))); - } - + ensure_dimensions(model, &embeddings)?; Ok(embeddings) } @@ -298,6 +274,26 @@ impl AIClient { } } +/// Rejects vectors the store cannot hold. Extracted from `embed` so the guard and its message are +/// testable without a live endpoint. +fn ensure_dimensions(model: &str, embeddings: &[Vec]) -> FlowyResult<()> { + match embeddings + .iter() + .map(|e| e.len()) + .find(|n| *n != EMBEDDING_DIMENSION) + { + None => Ok(()), + Some(actual) => Err(FlowyError::local_ai().with_context(format!( + "embedding model '{model}' returned {actual}-dimensional vectors, but {required} are \ + required. Pick a model that supports {required} dimensions, or one that honours the \ + `dimensions` parameter (for example text-embedding-3-small).", + model = model, + actual = actual, + required = EMBEDDING_DIMENSION, + ))), + } +} + fn ollama_request( model: &str, messages: &[Message], @@ -410,6 +406,31 @@ mod tests { )); } + #[test] + fn dimension_guard_accepts_correct_size_and_rejects_others() { + assert!(ensure_dimensions("m", &[vec![0.0; EMBEDDING_DIMENSION]]).is_ok()); + assert!(ensure_dimensions("m", &[]).is_ok()); + // A single bad vector in an otherwise valid batch must still be caught. + let mixed = vec![vec![0.0; EMBEDDING_DIMENSION], vec![0.0; 1536]]; + assert!(ensure_dimensions("m", &mixed).is_err()); + } + + /// The message must report the *required* dimension, not echo the model name back. A previous + /// revision used `{0}`, which silently resolved to the model argument. + #[test] + fn dimension_error_reports_required_size_not_model_name() { + let err = ensure_dimensions("text-embedding-3-large", &[vec![0.0; 3072]]).unwrap_err(); + let msg = err.to_string(); + + assert!(msg.contains("returned 3072-dimensional"), "{msg}"); + assert!(msg.contains("but 768 are required"), "{msg}"); + assert!(msg.contains("supports 768 dimensions"), "{msg}"); + assert!( + !msg.contains("supports text-embedding-3-large dimensions"), + "model name leaked into the dimension slot: {msg}" + ); + } + #[test] fn message_roles_map_to_openai_roles() { let mapped = openai_message(&Message::new_system_message("s")).unwrap();