diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index da4c706..b7d90f5 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -29,7 +29,7 @@ chrono = { version = "0.4.42", features = ["serde"] } anyhow = "1.0.100" ignore = "0.4.25" dashmap = "6.1.0" -async-openai = { version = "0.31.1", features = ["responses", "embedding", "chat-completion", "chat-completion-types"] } +async-openai = { version = "0.31.1", features = ["responses", "embedding", "chat-completion", "chat-completion-types", "byot"] } specta = { version = "=2.0.0-rc.22", features = ["derive", "chrono", "serde"] } specta-typescript = "0.0.9" tauri-specta = { version = "=2.0.0-rc.21", features = ["derive", "typescript"] } diff --git a/src-tauri/migrations/20251205113555_init_schema.sql b/src-tauri/migrations/20251205113555_init_schema.sql index 0f20a47..eabaa51 100644 --- a/src-tauri/migrations/20251205113555_init_schema.sql +++ b/src-tauri/migrations/20251205113555_init_schema.sql @@ -59,7 +59,7 @@ CREATE TABLE IF NOT EXISTS file_chunks ( file_id INTEGER NOT NULL, chunk_index INTEGER NOT NULL, -- index of chunk in file - content TEXT NOT NULL, -- text part + content TEXT, -- text part -- optional for syntax highlighting start_char_idx INTEGER, diff --git a/src-tauri/src/ai.rs b/src-tauri/src/ai.rs index 3caaa82..d763467 100644 --- a/src-tauri/src/ai.rs +++ b/src-tauri/src/ai.rs @@ -17,6 +17,8 @@ use async_openai::{ }; use base64::Engine; +use crate::ai::{self}; + pub mod embedding; pub mod llm; @@ -78,6 +80,29 @@ impl AI { Ok(response) } + pub async fn create_embedding_vecbox( + &self, + input: ai::embedding::vecbox::VecBoxEmbeddingInput, + model: String, + ) -> Result { + let content_parts = input.to_content_parts(); + + let request = ai::embedding::vecbox::VecBoxEmbeddingRequest { + model: Some(model), + input: content_parts, + instruction: input.instruction, + }; + + let response: CreateEmbeddingResponse = self + .client + .embeddings() + .create_byot(&request) + .await + .context("failed to generate vecBox embedding")?; + + Ok(response) + } + pub async fn request_llm(&self, input: InputParam, model: String) -> Result { let args = CreateResponseArgs::default() .input(input) diff --git a/src-tauri/src/ai/embedding.rs b/src-tauri/src/ai/embedding.rs index e69de29..c5dbdd0 100644 --- a/src-tauri/src/ai/embedding.rs +++ b/src-tauri/src/ai/embedding.rs @@ -0,0 +1 @@ +pub mod vecbox; diff --git a/src-tauri/src/ai/embedding/vecbox.rs b/src-tauri/src/ai/embedding/vecbox.rs new file mode 100644 index 0000000..fd7a9fa --- /dev/null +++ b/src-tauri/src/ai/embedding/vecbox.rs @@ -0,0 +1,46 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum EmbeddingContentPart { + Text { text: String }, + ImageUrl { image_url: EmbeddingImageUrl }, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct EmbeddingImageUrl { + pub url: String, +} + +#[derive(Debug, Serialize)] +pub struct VecBoxEmbeddingRequest { + pub model: Option, + pub input: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub instruction: Option, +} + +#[derive(Debug, Clone)] +pub struct VecBoxEmbeddingInput { + pub text: Option, + pub instruction: Option, + pub image_url: Option, +} + +impl VecBoxEmbeddingInput { + pub fn to_content_parts(&self) -> Vec { + let mut parts = Vec::new(); + + if let Some(ref text) = self.text { + parts.push(EmbeddingContentPart::Text { text: text.clone() }); + } + + if let Some(ref url) = self.image_url { + parts.push(EmbeddingContentPart::ImageUrl { + image_url: EmbeddingImageUrl { url: url.clone() }, + }); + } + + parts + } +} diff --git a/src-tauri/src/database/chunks.rs b/src-tauri/src/database/chunks.rs index d0a56e4..fcecff4 100644 --- a/src-tauri/src/database/chunks.rs +++ b/src-tauri/src/database/chunks.rs @@ -7,7 +7,7 @@ pub struct AddFileChunk { pub file_id: i32, pub chunk_index: i32, - pub content: String, + pub content: Option, pub start_char_idx: Option, pub end_char_idx: Option, diff --git a/src-tauri/src/database/models.rs b/src-tauri/src/database/models.rs index da0a2ab..d95cda4 100644 --- a/src-tauri/src/database/models.rs +++ b/src-tauri/src/database/models.rs @@ -1,9 +1,20 @@ use specta::Type; -use serde::{Deserialize, Serialize}; -use sqlx::FromRow; use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; use sqlx::types::Json; +use sqlx::FromRow; + +// EMBEDDING BACKEND TYPE + +#[derive(Debug, Serialize, Deserialize, Clone, Type, PartialEq, Default)] +pub enum EmbeddingBackendType { + #[default] + #[serde(rename = "openai_compat")] + OpenAICompat, + #[serde(rename = "vecbox")] + VecBox, +} // APP CONFIG #[derive(Debug, Serialize, Deserialize, Clone, Type)] @@ -16,21 +27,25 @@ pub struct AppConfig { #[serde(default = "default_parallelism")] #[specta(type = i32)] pub indexer_parallelism: usize, // How many files vecDir needs to index in parallel (2-4) - + // AI Settings (Global defaults) pub default_openai_url: Option, } // Default values -fn default_theme() -> String { "system".to_string() } -fn default_parallelism() -> usize { 2 } +fn default_theme() -> String { + "system".to_string() +} +fn default_parallelism() -> usize { + 2 +} impl Default for AppConfig { fn default() -> Self { Self { theme: default_theme(), indexer_parallelism: default_parallelism(), - default_openai_url: None + default_openai_url: None, } } } @@ -49,7 +64,7 @@ pub struct LLMConfig { pub api_key: String, pub model: String, - + pub text_processing_prompt: AIPrompt, pub image_processing_prompt: AIPrompt, pub default_processing_prompt: AIPrompt, @@ -57,12 +72,23 @@ pub struct LLMConfig { #[derive(Debug, Serialize, Deserialize, Clone, Type)] pub struct EmbeddingConfig { + #[serde(default)] + pub backend_type: EmbeddingBackendType, + pub api_base_url: String, pub api_key: String, pub model: String, pub dimensions: i32, + + pub text_processing_prompt: AIPrompt, + pub image_processing_prompt: AIPrompt, + pub default_processing_prompt: AIPrompt, + + pub search_prompt: AIPrompt, + + pub multimodal: bool, } // SPACES @@ -77,7 +103,7 @@ pub struct Space { #[specta(type = EmbeddingConfig)] pub embedding_config: Json, - #[specta(type = LLMConfig)] + #[specta(type = LLMConfig)] pub llm_config: Json, pub created_at: DateTime, @@ -112,7 +138,7 @@ pub struct FileMetadata { pub modified_at_fs: DateTime, pub last_indexed_at: Option>, pub content_hash: Option, - + pub indexing_status: String, pub indexing_error_message: Option, } @@ -124,17 +150,17 @@ pub struct FileChunk { pub file_id: i32, pub chunk_index: i32, - pub content: String, + pub content: Option, pub start_char_idx: Option, - pub end_char_idx: Option + pub end_char_idx: Option, } // VECTOR SEARCH RESULT #[derive(Debug, FromRow, Serialize, Type)] pub struct VectorSearchResult { pub chunk_id: i32, - pub content: String, + pub content: Option, pub file_id: i32, pub absolute_path: String, diff --git a/src-tauri/src/indexer/processor.rs b/src-tauri/src/indexer/processor.rs index 9154786..967b4c2 100644 --- a/src-tauri/src/indexer/processor.rs +++ b/src-tauri/src/indexer/processor.rs @@ -6,22 +6,18 @@ use tauri::{AppHandle, Emitter}; use tauri_specta::Event; use crate::{ - ai::AI, + ai::{self, AI}, database::{ - self, - chunks::{add_chunk, add_chunks_batch, AddFileChunk}, + chunks::{add_chunks_batch, AddFileChunk}, files::{ get_pending_files_for_space, mark_file_as_indexed, mark_file_as_indexed_batch, MarkFileAsIndexed, }, - models::LLMConfig, + models::{EmbeddingBackendType, EmbeddingConfig, LLMConfig, Space}, spaces::get_space_by_id, DbPool, }, - status::{ - self, - events::{StatusEvent, StatusType}, - }, + status::events::{StatusEvent, StatusType}, }; async fn process_image( @@ -59,12 +55,37 @@ pub async fn process_space( space_id: i32, limit: i32, ) -> Result<()> { - let mut descriptions: Vec = Vec::new(); - let mut processed_files: Vec = Vec::new(); - let space = get_space_by_id(pool, space_id) .await .context("failed to get space in process_space")?; + let embedding_config = &space.embedding_config.0; + + match embedding_config.backend_type { + EmbeddingBackendType::VecBox if embedding_config.multimodal => { + process_space_vecbox_multimodal(app_handle, pool, space, limit) + .await + .context("failed to process space using vecBox multimodal embedding")?; + } + EmbeddingBackendType::OpenAICompat | EmbeddingBackendType::VecBox => { + process_space_llm(app_handle, pool, space, limit) + .await + .context("failed to process space using LLM description embedding")?; + } + } + + Ok(()) +} + +pub async fn process_space_llm( + app_handle: AppHandle, + pool: &DbPool, + space: Space, + limit: i32, +) -> Result<()> { + let mut descriptions: Vec = Vec::new(); + let mut processed_files: Vec = Vec::new(); + + let space_id = space.id; let llm_config = space.llm_config.0; let embedding_config = space.embedding_config.0; @@ -166,7 +187,7 @@ pub async fn process_space( file_id: file_id, chunk_index: 0, - content: content.clone(), + content: Some(content.clone()), start_char_idx: None, end_char_idx: None, @@ -220,3 +241,140 @@ pub async fn process_space( Ok(()) } + +pub async fn process_space_vecbox_multimodal( + app_handle: AppHandle, + pool: &DbPool, + space: Space, + limit: i32, +) -> Result<()> { + let space_id = space.id; + + let embedding_config = space.embedding_config.0; + + let ai_client_embedding = AI::new(&embedding_config.api_base_url, &embedding_config.api_key) + .context("failed to create openai client")?; + + let pending_files = get_pending_files_for_space(pool, space_id, limit) + .await + .context("failed to get pending indexed files")?; + + if pending_files.is_empty() { + return Ok(()); + } + + let total_files = pending_files.len() as i32; + for (i, file) in pending_files.iter().enumerate() { + let file_path = file.absolute_path.clone(); + let guess = mime_guess::from_path(&file_path); + let mime_type = guess.first_or_octet_stream(); + + let input = match mime_type.type_() { + mime::IMAGE => ai::embedding::vecbox::VecBoxEmbeddingInput { + instruction: Some( + embedding_config + .image_processing_prompt + .system_prompt + .clone(), + ), + text: Some( + embedding_config + .image_processing_prompt + .user_prompt + .clone(), + ), + image_url: Some( + ai_client_embedding + .image_to_base64(&file_path) + .await + .context("failed to get image base64 url")?, + ), + }, + _ => { + continue; + } + }; + + let embeddings_response = ai_client_embedding + .create_embedding_vecbox(input, embedding_config.model.clone()) + .await + .context("failed to create vecBox embedding during processing space")?; + + if embeddings_response.data.is_empty() { + continue; + } + + let mut chunks_to_add: Vec = Vec::new(); + let mut updates_to_add: Vec = Vec::new(); + + for embedding_item in embeddings_response.data { + let file_id = file.id; + + let embedding_result = ai_client_embedding + .prepare_matroshka(embedding_item.embedding.clone(), 768) + .context("failed to prepare matroshka to 768 dim"); + + if embedding_result.is_err() { + println!("failed to create batch embedding {:?}", embedding_result); + continue; + } + + let embedding = embedding_result.unwrap(); + + let file_chunk = AddFileChunk { + file_id: file_id, + + chunk_index: 0, + content: None, + + start_char_idx: None, + end_char_idx: None, + + embedding: embedding, + }; + + chunks_to_add.push(file_chunk); + + let update: MarkFileAsIndexed = MarkFileAsIndexed { + file_id: file_id, + description: None, + }; + + updates_to_add.push(update); + } + + add_chunks_batch(pool, chunks_to_add) + .await + .context("failed to add chunks in batch in process_space")?; + + mark_file_as_indexed_batch(pool, updates_to_add) + .await + .context("failed to update file indexing status in batch in process_space")?; + + StatusEvent { + status: StatusType::Processing, + message: Some("Processing Space".to_string()), + total: Some(total_files), + processed: Some(i as i32), + } + .emit(&app_handle)?; + } + + StatusEvent { + status: StatusType::Idle, + message: None, + total: None, + processed: None, + } + .emit(&app_handle)?; + + StatusEvent { + status: StatusType::Notification, + message: Some("Space Processed".to_string()), + total: None, + processed: None, + } + .emit(&app_handle)?; + + Ok(()) +} diff --git a/src/components/settings/createSpace.tsx b/src/components/settings/createSpace.tsx index 9207f07..79d08e4 100644 --- a/src/components/settings/createSpace.tsx +++ b/src/components/settings/createSpace.tsx @@ -2,10 +2,11 @@ import type { EmbeddingConfig, LLMConfig } from "@/lib/vecdir/bindings"; import { useForm } from "@tanstack/react-form"; import { useNavigate } from "@tanstack/react-router"; -import { Brain, FileImage, FileType, FolderPen, SquaresIntersect } from "lucide-react"; +import { Brain, FileImage, FileType, FolderPen, Search, SquaresIntersect } from "lucide-react"; import { toast } from "sonner"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { createSpace } from "@/lib/vecdir/spaces/createSpace"; @@ -53,6 +54,26 @@ export function CreateSpace() { api_base_url: "http://127.0.0.1:1234/v1", // LM Studio api_key: "lmstudio", + + text_processing_prompt: { + system_prompt: "you are a text processing LLM for RAG purposes", + user_prompt: "describe this text file", + }, + image_processing_prompt: { + system_prompt: "you are an image descriptor. you describe images very precisely to use these descriptions in RAG", + user_prompt: "describe this image", + }, + default_processing_prompt: { + system_prompt: "you are a file processing LLM for RAG purposes", + user_prompt: "describe this file based on its metadata: ", + }, + + search_prompt: { + system_prompt: "you are a file processing LLM for RAG purposes", + user_prompt: "describe this file based on its metadata: ", + }, + + multimodal: true, }, }; const form = useForm({ @@ -345,6 +366,173 @@ export function CreateSpace() { )} + !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(!field.state.value)} + className="border-border" + /> +
+ )} +
+ +
+ +

Text Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Image Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Default Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Search Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
[state.canSubmit, state.isSubmitting]} diff --git a/src/components/settings/settings.tsx b/src/components/settings/settings.tsx index 3244438..c2e4a94 100644 --- a/src/components/settings/settings.tsx +++ b/src/components/settings/settings.tsx @@ -1,11 +1,12 @@ import type { EmbeddingConfig, IndexedRoot, LLMConfig } from "@/lib/vecdir/bindings"; import { useForm } from "@tanstack/react-form"; import { open } from "@tauri-apps/plugin-dialog"; -import { Brain, FileImage, FileType, FolderPen, SquareEqual, X } from "lucide-react"; +import { Brain, FileImage, FileType, FolderPen, Search, SquareEqual, X } from "lucide-react"; import { useEffect, useState } from "react"; import { toast } from "sonner"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardFooter } from "@/components/ui/card"; +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { addRoot } from "@/lib/vecdir/roots/createRoot"; @@ -408,6 +409,173 @@ export function Settings() { )} + !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(!field.state.value)} + className="border-border" + /> +
+ )} +
+ +
+ +

Text Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Image Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Default Processing Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
+ +
+ +

Search Prompt

+
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="System Prompt" + className="border-border" + /> +
+ )} +
+ !value ? "This field is required!" : undefined }}> + {field => ( +
+ + field.handleChange(e.target.value)} + placeholder="User Prompt" + className="border-border" + /> +
+ )} +
+
[state.canSubmit, state.isSubmitting]} diff --git a/src/lib/vecdir/bindings.ts b/src/lib/vecdir/bindings.ts index 8213c61..f441a0c 100644 --- a/src/lib/vecdir/bindings.ts +++ b/src/lib/vecdir/bindings.ts @@ -136,7 +136,8 @@ statusEvent: "status-event" export type AIPrompt = { system_prompt: string; user_prompt: string } export type AppConfig = { theme?: string; indexer_parallelism?: number; default_openai_url: string | null } export type BackendReadyEvent = null -export type EmbeddingConfig = { api_base_url: string; api_key: string; model: string; dimensions: number } +export type EmbeddingBackendType = "openai_compat" | "vecbox" +export type EmbeddingConfig = { backend_type?: EmbeddingBackendType; api_base_url: string; api_key: string; model: string; dimensions: number; text_processing_prompt: AIPrompt; image_processing_prompt: AIPrompt; default_processing_prompt: AIPrompt; search_prompt: AIPrompt; multimodal: boolean } export type ErrorEvent = { message: string; context: string | null } export type FileMetadata = { id: number; root_id: number; absolute_path: string; filename: string; file_extension: string; file_size: number; description: string | null; modified_at_fs: string; last_indexed_at: string | null; content_hash: string | null; indexing_status: string; indexing_error_message: string | null } export type IndexedRoot = { id: number; space_id: number; path: string; status: string } @@ -144,7 +145,7 @@ export type LLMConfig = { api_base_url: string; api_key: string; model: string; export type Space = { id: number; name: string; description: string | null; embedding_config: EmbeddingConfig; llm_config: LLMConfig; created_at: string } export type StatusEvent = { status: StatusType; message: string | null; total: number | null; processed: number | null } export type StatusType = "Idle" | "Indexing" | "Processing" | "Notification" | "Error" -export type VectorSearchResult = { chunk_id: number; content: string; file_id: number; absolute_path: string; filename: string; distance: number } +export type VectorSearchResult = { chunk_id: number; content: string | null; file_id: number; absolute_path: string; filename: string; distance: number } /** tauri-specta globals **/