feat(chat): implement speculative web search and embeddings commands
- Added speculative search functionality to enhance query processing by detecting when a web search may be needed based on user input patterns. - Integrated commands for managing embeddings, including downloading models, checking availability, calculating relevance scores, generating embeddings, and pruning context. - Optimized chat streaming by buffering tokens to reduce overhead and improve responsiveness. - Implemented lazy initialization for the headless browser to conserve memory until needed. - Enhanced web scraping capabilities with a static scraping method for faster content retrieval.
This commit is contained in:
1 parent
6cc3f71f38
commit
5fb1489eff
7 files changed
+1407
-40
No files matched your search
@@ -0,0 +1,566 @@
|
||||
//! Módulo de Embeddings ONNX para ranking de relevância
|
||||
//!
|
||||
//! Usa o modelo all-MiniLM-L6-v2 para gerar embeddings de texto
|
||||
//! e calcular similaridade de cosseno para ranking de resultados de busca.
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use ndarray::Array2;
|
||||
use ort::session::{Session, builder::GraphOptimizationLevel};
|
||||
use ort::value::Value;
|
||||
use std::io::Read;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
/// Flag para controlar se o ort já foi inicializado
|
||||
static ORT_INITIALIZED: OnceLock<Result<(), String>> = OnceLock::new();
|
||||
|
||||
/// Dimensão dos embeddings do modelo all-MiniLM-L6-v2
|
||||
pub const EMBEDDING_DIM: usize = 384;
|
||||
|
||||
/// Tamanho máximo de tokens para o modelo
|
||||
const MAX_SEQ_LENGTH: usize = 256;
|
||||
|
||||
/// Global lazy-initialized embedding model (com Mutex para permitir mutabilidade)
|
||||
static EMBEDDING_MODEL: OnceLock<Result<Arc<Mutex<EmbeddingModel>>, String>> = OnceLock::new();
|
||||
|
||||
/// Modelo de embeddings para cálculo de similaridade semântica
|
||||
pub struct EmbeddingModel {
|
||||
session: Session,
|
||||
tokenizer: Tokenizer,
|
||||
}
|
||||
|
||||
impl EmbeddingModel {
|
||||
/// Carrega o modelo ONNX e tokenizer
|
||||
pub fn new(model_path: &str, tokenizer_path: &str) -> Result<Self> {
|
||||
// Verificar se arquivos existem
|
||||
if !Path::new(model_path).exists() {
|
||||
return Err(anyhow!("Model file not found: {}", model_path));
|
||||
}
|
||||
if !Path::new(tokenizer_path).exists() {
|
||||
return Err(anyhow!("Tokenizer file not found: {}", tokenizer_path));
|
||||
}
|
||||
|
||||
log::info!("[Embeddings] Loading ONNX model from: {}", model_path);
|
||||
|
||||
// Criar sessão ONNX
|
||||
let session = Session::builder()?
|
||||
.with_optimization_level(GraphOptimizationLevel::Level3)?
|
||||
.commit_from_file(model_path)?;
|
||||
|
||||
// Carregar tokenizer
|
||||
let tokenizer = Tokenizer::from_file(tokenizer_path)
|
||||
.map_err(|e| anyhow!("Failed to load tokenizer: {}", e))?;
|
||||
|
||||
log::info!("[Embeddings] Model loaded successfully");
|
||||
|
||||
Ok(Self { session, tokenizer })
|
||||
}
|
||||
|
||||
/// Gera embedding para um texto
|
||||
pub fn embed(&mut self, text: &str) -> Result<Vec<f32>> {
|
||||
// Tokenizar texto
|
||||
let encoding = self.tokenizer
|
||||
.encode(text, true)
|
||||
.map_err(|e| anyhow!("Tokenization failed: {}", e))?;
|
||||
|
||||
let mut input_ids: Vec<i64> = encoding.get_ids()
|
||||
.iter()
|
||||
.map(|&id| id as i64)
|
||||
.collect();
|
||||
let mut attention_mask: Vec<i64> = encoding.get_attention_mask()
|
||||
.iter()
|
||||
.map(|&m| m as i64)
|
||||
.collect();
|
||||
let mut token_type_ids: Vec<i64> = encoding.get_type_ids()
|
||||
.iter()
|
||||
.map(|&t| t as i64)
|
||||
.collect();
|
||||
|
||||
// Truncar/pad para MAX_SEQ_LENGTH
|
||||
input_ids.truncate(MAX_SEQ_LENGTH);
|
||||
attention_mask.truncate(MAX_SEQ_LENGTH);
|
||||
token_type_ids.truncate(MAX_SEQ_LENGTH);
|
||||
|
||||
while input_ids.len() < MAX_SEQ_LENGTH {
|
||||
input_ids.push(0);
|
||||
attention_mask.push(0);
|
||||
token_type_ids.push(0);
|
||||
}
|
||||
|
||||
// Criar arrays para inferência (batch size = 1)
|
||||
let input_ids_array = Array2::from_shape_vec((1, MAX_SEQ_LENGTH), input_ids)?;
|
||||
let attention_mask_array = Array2::from_shape_vec((1, MAX_SEQ_LENGTH), attention_mask)?;
|
||||
let token_type_ids_array = Array2::from_shape_vec((1, MAX_SEQ_LENGTH), token_type_ids)?;
|
||||
|
||||
// Criar inputs ONNX
|
||||
let input_ids_value = Value::from_array(input_ids_array)?;
|
||||
let attention_mask_value = Value::from_array(attention_mask_array)?;
|
||||
let token_type_ids_value = Value::from_array(token_type_ids_array)?;
|
||||
|
||||
// Executar inferência usando vetor de inputs
|
||||
let inputs: Vec<(std::borrow::Cow<str>, ort::session::SessionInputValue)> = vec![
|
||||
("input_ids".into(), input_ids_value.into()),
|
||||
("attention_mask".into(), attention_mask_value.into()),
|
||||
("token_type_ids".into(), token_type_ids_value.into()),
|
||||
];
|
||||
|
||||
let outputs = self.session.run(inputs)?;
|
||||
|
||||
// Extrair output (last_hidden_state ou sentence_embedding dependendo do modelo)
|
||||
// Para all-MiniLM-L6-v2, fazemos mean pooling do last_hidden_state
|
||||
let output = outputs.get("last_hidden_state")
|
||||
.or_else(|| outputs.get("sentence_embedding"))
|
||||
.ok_or_else(|| anyhow!("Output tensor not found"))?;
|
||||
|
||||
let (shape, raw_data) = output.try_extract_tensor::<f32>()?;
|
||||
// Shape implementa Deref para [i64], então podemos usar diretamente
|
||||
let dims: &[i64] = &*shape;
|
||||
let data: &[f32] = raw_data;
|
||||
|
||||
// Mean pooling: média ao longo da dimensão de sequência
|
||||
let embedding = if dims.len() == 3 {
|
||||
// Shape: (batch, seq_len, hidden_dim)
|
||||
let seq_len = dims[1] as usize;
|
||||
let hidden_dim = dims[2] as usize;
|
||||
|
||||
let mut pooled = vec![0.0f32; hidden_dim];
|
||||
for i in 0..seq_len {
|
||||
for j in 0..hidden_dim {
|
||||
let idx = i * hidden_dim + j;
|
||||
pooled[j] += data[idx];
|
||||
}
|
||||
}
|
||||
for v in &mut pooled {
|
||||
*v /= seq_len as f32;
|
||||
}
|
||||
|
||||
// Normalizar L2
|
||||
let norm: f32 = pooled.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for v in &mut pooled {
|
||||
*v /= norm;
|
||||
}
|
||||
}
|
||||
|
||||
pooled
|
||||
} else {
|
||||
// Shape: (batch, hidden_dim) - já pooled
|
||||
let hidden_dim = dims[1] as usize;
|
||||
let mut embedding = Vec::with_capacity(hidden_dim);
|
||||
for j in 0..hidden_dim {
|
||||
embedding.push(data[j]);
|
||||
}
|
||||
|
||||
// Normalizar L2
|
||||
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for v in &mut embedding {
|
||||
*v /= norm;
|
||||
}
|
||||
}
|
||||
|
||||
embedding
|
||||
};
|
||||
|
||||
Ok(embedding)
|
||||
}
|
||||
|
||||
/// Calcula embeddings em batch (mais eficiente para múltiplos textos)
|
||||
pub fn embed_batch(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
|
||||
// Para simplificar, processa um por um (pode ser otimizado para batch real)
|
||||
let mut results = Vec::with_capacity(texts.len());
|
||||
for text in texts {
|
||||
results.push(self.embed(text)?);
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
/// Calcula similaridade de cosseno entre dois vetores
|
||||
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
if a.len() != b.len() || a.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
|
||||
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
|
||||
if norm_a == 0.0 || norm_b == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
dot / (norm_a * norm_b)
|
||||
}
|
||||
|
||||
/// URLs para download do modelo (Hugging Face)
|
||||
const MODEL_URL: &str = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx";
|
||||
const TOKENIZER_URL: &str = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json";
|
||||
|
||||
/// URL para download da biblioteca ONNX Runtime (Windows x64)
|
||||
#[cfg(target_os = "windows")]
|
||||
const ORT_DLL_URL: &str = "https://github.com/microsoft/onnxruntime/releases/download/v1.20.1/onnxruntime-win-x64-1.20.1.zip";
|
||||
|
||||
/// Inicializa o ONNX Runtime com a biblioteca dinâmica
|
||||
fn init_ort_runtime(app_data_dir: &Path) -> Result<()> {
|
||||
let result = ORT_INITIALIZED.get_or_init(|| {
|
||||
let ort_dir = app_data_dir.join("ort");
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
let dll_path = ort_dir.join("onnxruntime.dll");
|
||||
#[cfg(target_os = "macos")]
|
||||
let dll_path = ort_dir.join("libonnxruntime.dylib");
|
||||
#[cfg(target_os = "linux")]
|
||||
let dll_path = ort_dir.join("libonnxruntime.so");
|
||||
|
||||
if dll_path.exists() {
|
||||
// Definir variável de ambiente para o ort encontrar a DLL
|
||||
std::env::set_var("ORT_DYLIB_PATH", &dll_path);
|
||||
log::info!("[Embeddings] ORT_DYLIB_PATH set to: {:?}", dll_path);
|
||||
Ok(())
|
||||
} else {
|
||||
// Tentar usar a biblioteca do sistema
|
||||
log::warn!("[Embeddings] ONNX Runtime library not found at {:?}, will try system path", dll_path);
|
||||
Ok(())
|
||||
}
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(()) => Ok(()),
|
||||
Err(e) => Err(anyhow!("{}", e))
|
||||
}
|
||||
}
|
||||
|
||||
/// Baixa um arquivo de uma URL para o caminho especificado
|
||||
async fn download_file(url: &str, path: &Path) -> Result<()> {
|
||||
log::info!("[Embeddings] Downloading: {} -> {:?}", url, path);
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(300)) // 5 min timeout
|
||||
.build()?;
|
||||
|
||||
let response = client.get(url).send().await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(anyhow!("Download failed with status: {}", response.status()));
|
||||
}
|
||||
|
||||
let bytes = response.bytes().await?;
|
||||
|
||||
// Criar diretório pai se não existir
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
std::fs::write(path, bytes)?;
|
||||
log::info!("[Embeddings] Downloaded successfully: {:?}", path);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Baixa e extrai a biblioteca ONNX Runtime
|
||||
#[cfg(target_os = "windows")]
|
||||
async fn ensure_ort_library(app_data_dir: &Path) -> Result<()> {
|
||||
let ort_dir = app_data_dir.join("ort");
|
||||
let dll_path = ort_dir.join("onnxruntime.dll");
|
||||
|
||||
if dll_path.exists() {
|
||||
log::info!("[Embeddings] ONNX Runtime library already exists");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
log::info!("[Embeddings] Downloading ONNX Runtime library...");
|
||||
|
||||
// Criar diretório
|
||||
std::fs::create_dir_all(&ort_dir)?;
|
||||
|
||||
// Baixar arquivo zip
|
||||
let zip_path = ort_dir.join("onnxruntime.zip");
|
||||
download_file(ORT_DLL_URL, &zip_path).await?;
|
||||
|
||||
// Extrair DLL do zip
|
||||
log::info!("[Embeddings] Extracting ONNX Runtime library...");
|
||||
|
||||
let file = std::fs::File::open(&zip_path)?;
|
||||
let mut archive = zip::ZipArchive::new(file)?;
|
||||
|
||||
// Procurar pela DLL dentro do zip
|
||||
for i in 0..archive.len() {
|
||||
let mut entry = archive.by_index(i)?;
|
||||
let name = entry.name().to_string();
|
||||
|
||||
if name.ends_with("onnxruntime.dll") {
|
||||
let mut outfile = std::fs::File::create(&dll_path)?;
|
||||
std::io::copy(&mut entry, &mut outfile)?;
|
||||
log::info!("[Embeddings] Extracted onnxruntime.dll");
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Remover arquivo zip
|
||||
let _ = std::fs::remove_file(&zip_path);
|
||||
|
||||
if !dll_path.exists() {
|
||||
return Err(anyhow!("Failed to extract onnxruntime.dll from archive"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
async fn ensure_ort_library(_app_data_dir: &Path) -> Result<()> {
|
||||
// Em outros sistemas, assumimos que a biblioteca está no sistema
|
||||
log::info!("[Embeddings] Using system ONNX Runtime library");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Garante que os arquivos do modelo existem, baixando se necessário
|
||||
pub async fn ensure_model_files(app_data_dir: &Path) -> Result<(std::path::PathBuf, std::path::PathBuf)> {
|
||||
let models_dir = app_data_dir.join("models");
|
||||
let model_path = models_dir.join("all-MiniLM-L6-v2.onnx");
|
||||
let tokenizer_path = models_dir.join("tokenizer.json");
|
||||
|
||||
// Criar diretório de modelos
|
||||
std::fs::create_dir_all(&models_dir)?;
|
||||
|
||||
// Baixar biblioteca ONNX Runtime se necessário (apenas Windows)
|
||||
ensure_ort_library(app_data_dir).await?;
|
||||
|
||||
// Inicializar ort com o caminho da DLL
|
||||
init_ort_runtime(app_data_dir)?;
|
||||
|
||||
// Baixar modelo se não existir
|
||||
if !model_path.exists() {
|
||||
log::info!("[Embeddings] Model not found, downloading...");
|
||||
download_file(MODEL_URL, &model_path).await?;
|
||||
}
|
||||
|
||||
// Baixar tokenizer se não existir
|
||||
if !tokenizer_path.exists() {
|
||||
log::info!("[Embeddings] Tokenizer not found, downloading...");
|
||||
download_file(TOKENIZER_URL, &tokenizer_path).await?;
|
||||
}
|
||||
|
||||
Ok((model_path, tokenizer_path))
|
||||
}
|
||||
|
||||
/// Verifica se o modelo está disponível
|
||||
pub fn is_model_available(app_data_dir: &Path) -> bool {
|
||||
let models_dir = app_data_dir.join("models");
|
||||
let model_path = models_dir.join("all-MiniLM-L6-v2.onnx");
|
||||
let tokenizer_path = models_dir.join("tokenizer.json");
|
||||
|
||||
model_path.exists() && tokenizer_path.exists()
|
||||
}
|
||||
|
||||
/// Obtém ou inicializa o modelo global de embeddings
|
||||
pub fn get_or_init_model(app_data_dir: &Path) -> Result<Arc<Mutex<EmbeddingModel>>> {
|
||||
// Inicializar ort com o caminho da DLL antes de criar o modelo
|
||||
init_ort_runtime(app_data_dir)?;
|
||||
|
||||
let result = EMBEDDING_MODEL.get_or_init(|| {
|
||||
let model_path = app_data_dir.join("models").join("all-MiniLM-L6-v2.onnx");
|
||||
let tokenizer_path = app_data_dir.join("models").join("tokenizer.json");
|
||||
|
||||
match EmbeddingModel::new(
|
||||
model_path.to_str().unwrap_or(""),
|
||||
tokenizer_path.to_str().unwrap_or("")
|
||||
) {
|
||||
Ok(model) => Ok(Arc::new(Mutex::new(model))),
|
||||
Err(e) => Err(format!("Failed to load embedding model: {}", e))
|
||||
}
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(model) => Ok(model.clone()),
|
||||
Err(e) => Err(anyhow!("{}", e))
|
||||
}
|
||||
}
|
||||
|
||||
/// Calcula scores de relevância para múltiplos textos em relação a uma query
|
||||
pub fn rank_by_relevance(
|
||||
model: &mut EmbeddingModel,
|
||||
query: &str,
|
||||
texts: &[&str],
|
||||
) -> Result<Vec<(usize, f32)>> {
|
||||
let query_embedding = model.embed(query)?;
|
||||
|
||||
let mut scores: Vec<(usize, f32)> = Vec::with_capacity(texts.len());
|
||||
for (idx, text) in texts.iter().enumerate() {
|
||||
let text_embedding = model.embed(text).unwrap_or_else(|_| vec![0.0; EMBEDDING_DIM]);
|
||||
let score = cosine_similarity(&query_embedding, &text_embedding);
|
||||
scores.push((idx, score));
|
||||
}
|
||||
|
||||
// Ordenar por score decrescente
|
||||
scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
Ok(scores)
|
||||
}
|
||||
|
||||
/// Poda o contexto mantendo apenas os parágrafos mais relevantes
|
||||
///
|
||||
/// Argumentos:
|
||||
/// - model: Modelo de embeddings
|
||||
/// - query: Query do usuário
|
||||
/// - context: Contexto completo (texto separado por parágrafos)
|
||||
/// - max_tokens: Número máximo de tokens (aproximado por palavras)
|
||||
/// - min_score: Score mínimo de relevância (0.0 a 1.0)
|
||||
///
|
||||
/// Retorna: Contexto podado com os parágrafos mais relevantes
|
||||
pub fn prune_context(
|
||||
model: &mut EmbeddingModel,
|
||||
query: &str,
|
||||
context: &str,
|
||||
max_tokens: usize,
|
||||
min_score: f32,
|
||||
) -> Result<String> {
|
||||
// Dividir contexto em parágrafos
|
||||
let paragraphs: Vec<&str> = context
|
||||
.split("\n\n")
|
||||
.map(|p| p.trim())
|
||||
.filter(|p| !p.is_empty() && p.len() > 20) // Filtrar parágrafos muito curtos
|
||||
.collect();
|
||||
|
||||
if paragraphs.is_empty() {
|
||||
return Ok(context.to_string());
|
||||
}
|
||||
|
||||
// Calcular embedding da query
|
||||
let query_embedding = model.embed(query)?;
|
||||
|
||||
// Calcular scores para cada parágrafo
|
||||
let mut scored_paragraphs: Vec<(f32, &str, usize)> = Vec::with_capacity(paragraphs.len());
|
||||
for (idx, &p) in paragraphs.iter().enumerate() {
|
||||
let embedding = model.embed(p).unwrap_or_else(|_| vec![0.0; EMBEDDING_DIM]);
|
||||
let score = cosine_similarity(&query_embedding, &embedding);
|
||||
scored_paragraphs.push((score, p, idx));
|
||||
}
|
||||
|
||||
// Ordenar por score decrescente
|
||||
scored_paragraphs.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
// Filtrar por score mínimo
|
||||
scored_paragraphs.retain(|(score, _, _)| *score >= min_score);
|
||||
|
||||
// Selecionar parágrafos até atingir max_tokens
|
||||
let mut result = Vec::new();
|
||||
let mut total_tokens = 0;
|
||||
|
||||
for (score, paragraph, original_idx) in scored_paragraphs {
|
||||
let paragraph_tokens = paragraph.split_whitespace().count();
|
||||
|
||||
if total_tokens + paragraph_tokens > max_tokens {
|
||||
// Se ainda não temos nenhum parágrafo, incluir o primeiro mesmo que exceda
|
||||
if result.is_empty() {
|
||||
result.push((original_idx, paragraph));
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
result.push((original_idx, paragraph));
|
||||
total_tokens += paragraph_tokens;
|
||||
}
|
||||
|
||||
// Ordenar pelo índice original para manter a ordem do contexto
|
||||
result.sort_by_key(|(idx, _)| *idx);
|
||||
|
||||
// Reconstruir contexto
|
||||
let pruned: Vec<&str> = result.into_iter().map(|(_, p)| p).collect();
|
||||
|
||||
log::info!(
|
||||
"[Embeddings] Context pruned: {} paragraphs -> {} paragraphs, {} tokens",
|
||||
paragraphs.len(),
|
||||
pruned.len(),
|
||||
total_tokens
|
||||
);
|
||||
|
||||
Ok(pruned.join("\n\n"))
|
||||
}
|
||||
|
||||
/// Versão simplificada de poda usando apenas BM25-like (sem embeddings)
|
||||
/// Útil quando o modelo não está disponível
|
||||
pub fn prune_context_bm25(
|
||||
query: &str,
|
||||
context: &str,
|
||||
max_tokens: usize,
|
||||
) -> String {
|
||||
let paragraphs: Vec<&str> = context
|
||||
.split("\n\n")
|
||||
.map(|p| p.trim())
|
||||
.filter(|p| !p.is_empty() && p.len() > 20)
|
||||
.collect();
|
||||
|
||||
if paragraphs.is_empty() {
|
||||
return context.to_string();
|
||||
}
|
||||
|
||||
let query_lower = query.to_lowercase();
|
||||
let query_terms: Vec<&str> = query_lower
|
||||
.split_whitespace()
|
||||
.filter(|w| w.len() > 2)
|
||||
.collect();
|
||||
|
||||
// Calcular score BM25-like simples
|
||||
let mut scored: Vec<(f32, &str, usize)> = paragraphs.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, &p)| {
|
||||
let p_lower = p.to_lowercase();
|
||||
let mut score = 0.0f32;
|
||||
|
||||
for term in &query_terms {
|
||||
let count = p_lower.matches(term).count();
|
||||
if count > 0 {
|
||||
// TF com saturação logarítmica
|
||||
score += (1.0 + (count as f32).ln()) * (1.0 / (1.0 + query_terms.len() as f32));
|
||||
}
|
||||
}
|
||||
|
||||
(score, p, idx)
|
||||
})
|
||||
.collect();
|
||||
|
||||
scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
let mut result = Vec::new();
|
||||
let mut total_tokens = 0;
|
||||
|
||||
for (_, paragraph, original_idx) in scored {
|
||||
let paragraph_tokens = paragraph.split_whitespace().count();
|
||||
|
||||
if total_tokens + paragraph_tokens > max_tokens {
|
||||
if result.is_empty() {
|
||||
result.push((original_idx, paragraph));
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
result.push((original_idx, paragraph));
|
||||
total_tokens += paragraph_tokens;
|
||||
}
|
||||
|
||||
result.sort_by_key(|(idx, _)| *idx);
|
||||
|
||||
let pruned: Vec<&str> = result.into_iter().map(|(_, p)| p).collect();
|
||||
pruned.join("\n\n")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_cosine_similarity() {
|
||||
let a = vec![1.0, 0.0, 0.0];
|
||||
let b = vec![1.0, 0.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
|
||||
|
||||
let c = vec![0.0, 1.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &c) - 0.0).abs() < 0.001);
|
||||
|
||||
let d = vec![-1.0, 0.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &d) - (-1.0)).abs() < 0.001);
|
||||
}
|
||||
}
|
||||
|
||||
+147
-14
@@ -21,6 +21,7 @@ mod sources_config;
|
||||
mod system_monitor;
|
||||
mod intent_classifier;
|
||||
mod db;
|
||||
mod embeddings;
|
||||
|
||||
use web_scraper::{
|
||||
ScrapedContent,
|
||||
@@ -2984,12 +2985,18 @@ async fn chat_stream(
|
||||
return Err(error_msg);
|
||||
}
|
||||
|
||||
// 5. Processar stream e emitir tokens
|
||||
// IMPORTANTE: O Ollama envia tokens INCREMENTAIS (cada chunk contém apenas o novo conteúdo)
|
||||
// 5. Processar stream e emitir tokens COM BUFFERING
|
||||
// OTIMIZAÇÃO: Acumular tokens e emitir em batches para reduzir overhead da bridge
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
let mut full_content = String::new();
|
||||
|
||||
// Buffer de tokens para reduzir eventos na bridge
|
||||
let mut token_buffer = String::new();
|
||||
let mut last_emit = std::time::Instant::now();
|
||||
const EMIT_INTERVAL_MS: u64 = 16; // ~60fps para sincronizar com RAF do frontend
|
||||
const MAX_BUFFER_CHARS: usize = 50; // Emitir quando buffer tiver ~50 chars
|
||||
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
let chunk = chunk_result.map_err(|e| format!("Stream error: {}", e))?;
|
||||
let chunk_str = String::from_utf8_lossy(&chunk);
|
||||
@@ -3013,20 +3020,23 @@ async fn chat_stream(
|
||||
// Extrair conteúdo do chunk (Ollama envia tokens incrementais)
|
||||
if let Some(message) = json.get("message") {
|
||||
if let Some(content) = message.get("content").and_then(|c| c.as_str()) {
|
||||
// O Ollama envia apenas o NOVO conteúdo em cada chunk, não o acumulado
|
||||
// Então podemos emitir diretamente
|
||||
if !content.is_empty() {
|
||||
full_content.push_str(content);
|
||||
token_buffer.push_str(content);
|
||||
|
||||
// Emitir token para frontend
|
||||
let token_event = ChatTokenEvent {
|
||||
session_id: session_id.clone(),
|
||||
content: content.to_string(),
|
||||
done: false,
|
||||
};
|
||||
|
||||
if let Err(e) = window.emit("chat-token", &token_event) {
|
||||
log::warn!("Erro ao emitir token: {}", e);
|
||||
// Emitir buffer quando: tempo >= 16ms OU buffer >= 50 chars
|
||||
let elapsed = last_emit.elapsed().as_millis() as u64;
|
||||
if elapsed >= EMIT_INTERVAL_MS || token_buffer.len() >= MAX_BUFFER_CHARS {
|
||||
let token_event = ChatTokenEvent {
|
||||
session_id: session_id.clone(),
|
||||
content: std::mem::take(&mut token_buffer),
|
||||
done: false,
|
||||
};
|
||||
|
||||
if let Err(e) = window.emit("chat-token", &token_event) {
|
||||
log::warn!("Erro ao emitir token: {}", e);
|
||||
}
|
||||
last_emit = std::time::Instant::now();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3034,6 +3044,16 @@ async fn chat_stream(
|
||||
|
||||
// Verificar se stream terminou
|
||||
if is_done {
|
||||
// Flush do buffer residual antes de finalizar
|
||||
if !token_buffer.is_empty() {
|
||||
let flush_event = ChatTokenEvent {
|
||||
session_id: session_id.clone(),
|
||||
content: std::mem::take(&mut token_buffer),
|
||||
done: false,
|
||||
};
|
||||
let _ = window.emit("chat-token", &flush_event);
|
||||
}
|
||||
|
||||
// Emitir evento final
|
||||
let final_event = ChatTokenEvent {
|
||||
session_id: session_id.clone(),
|
||||
@@ -3127,6 +3147,113 @@ async fn chat_stream(
|
||||
Ok(session_id)
|
||||
}
|
||||
|
||||
// ============== COMANDOS DE EMBEDDINGS ==============
|
||||
|
||||
/// Baixa o modelo de embeddings se não existir
|
||||
#[command]
|
||||
async fn download_embedding_model(app_handle: AppHandle) -> Result<bool, String> {
|
||||
let app_data_dir = app_handle.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| format!("Failed to get app data dir: {}", e))?;
|
||||
|
||||
match embeddings::ensure_model_files(&app_data_dir).await {
|
||||
Ok(_) => {
|
||||
log::info!("[Embeddings] Model files ready");
|
||||
Ok(true)
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!("[Embeddings] Failed to ensure model files: {}", e);
|
||||
Err(format!("Failed to download model: {}", e))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Verifica se o modelo de embeddings está disponível
|
||||
#[command]
|
||||
fn is_embedding_model_available(app_handle: AppHandle) -> Result<bool, String> {
|
||||
let app_data_dir = app_handle.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| format!("Failed to get app data dir: {}", e))?;
|
||||
|
||||
Ok(embeddings::is_model_available(&app_data_dir))
|
||||
}
|
||||
|
||||
/// Calcula scores de relevância para textos em relação a uma query
|
||||
#[command]
|
||||
fn calculate_relevance_scores(
|
||||
app_handle: AppHandle,
|
||||
query: String,
|
||||
texts: Vec<String>,
|
||||
) -> Result<Vec<(usize, f32)>, String> {
|
||||
let app_data_dir = app_handle.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| format!("Failed to get app data dir: {}", e))?;
|
||||
|
||||
let model_arc = embeddings::get_or_init_model(&app_data_dir)
|
||||
.map_err(|e| format!("Failed to load model: {}", e))?;
|
||||
|
||||
let mut model = model_arc.lock()
|
||||
.map_err(|e| format!("Failed to lock model: {}", e))?;
|
||||
|
||||
let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
|
||||
|
||||
embeddings::rank_by_relevance(&mut model, &query, &text_refs)
|
||||
.map_err(|e| format!("Failed to calculate relevance: {}", e))
|
||||
}
|
||||
|
||||
/// Gera embedding para um texto
|
||||
#[command]
|
||||
fn generate_embedding(
|
||||
app_handle: AppHandle,
|
||||
text: String,
|
||||
) -> Result<Vec<f32>, String> {
|
||||
let app_data_dir = app_handle.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| format!("Failed to get app data dir: {}", e))?;
|
||||
|
||||
let model_arc = embeddings::get_or_init_model(&app_data_dir)
|
||||
.map_err(|e| format!("Failed to load model: {}", e))?;
|
||||
|
||||
let mut model = model_arc.lock()
|
||||
.map_err(|e| format!("Failed to lock model: {}", e))?;
|
||||
|
||||
model.embed(&text)
|
||||
.map_err(|e| format!("Failed to generate embedding: {}", e))
|
||||
}
|
||||
|
||||
/// Poda o contexto mantendo apenas os parágrafos mais relevantes
|
||||
#[command]
|
||||
fn prune_context(
|
||||
app_handle: AppHandle,
|
||||
query: String,
|
||||
context: String,
|
||||
max_tokens: Option<usize>,
|
||||
min_score: Option<f32>,
|
||||
) -> Result<String, String> {
|
||||
let max_tokens = max_tokens.unwrap_or(2000);
|
||||
let min_score = min_score.unwrap_or(0.3);
|
||||
|
||||
let app_data_dir = app_handle.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| format!("Failed to get app data dir: {}", e))?;
|
||||
|
||||
// Tentar usar embeddings se modelo disponível
|
||||
if embeddings::is_model_available(&app_data_dir) {
|
||||
let model_arc = embeddings::get_or_init_model(&app_data_dir)
|
||||
.map_err(|e| format!("Failed to load model: {}", e))?;
|
||||
|
||||
let mut model = model_arc.lock()
|
||||
.map_err(|e| format!("Failed to lock model: {}", e))?;
|
||||
|
||||
embeddings::prune_context(&mut model, &query, &context, max_tokens, min_score)
|
||||
.map_err(|e| format!("Failed to prune context: {}", e))
|
||||
} else {
|
||||
// Fallback para BM25-like
|
||||
log::info!("[PruneContext] Using BM25 fallback (embedding model not available)");
|
||||
Ok(embeddings::prune_context_bm25(&query, &context, max_tokens))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
tauri::Builder::default()
|
||||
@@ -3267,7 +3394,13 @@ pub fn run() {
|
||||
get_downloaded_installer_path,
|
||||
check_ollama_full,
|
||||
auto_start_ollama,
|
||||
classify_intent
|
||||
classify_intent,
|
||||
// Embeddings commands
|
||||
download_embedding_model,
|
||||
is_embedding_model_available,
|
||||
calculate_relevance_scores,
|
||||
generate_embedding,
|
||||
prune_context
|
||||
])
|
||||
.manage(Arc::new(Mutex::new(HashMap::<String, McpProcessHandle>::new())) as McpProcessMap)
|
||||
.run(tauri::generate_context!())
|
||||
|
||||
@@ -85,14 +85,13 @@ pub async fn reload_scheduled_tasks(
|
||||
return;
|
||||
}
|
||||
|
||||
// Obter browser - precisa acessar via app_handle
|
||||
// Obter browser - usando lazy initialization global
|
||||
let browser_arc = {
|
||||
// Criar browser diretamente se necessário
|
||||
use crate::web_scraper::create_browser;
|
||||
match create_browser() {
|
||||
Ok(b) => Arc::new(b),
|
||||
use crate::web_scraper::get_or_create_browser;
|
||||
match get_or_create_browser() {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
log::error!("Erro ao criar browser para task {}: {}", task_id, e);
|
||||
log::error!("Erro ao obter browser para task {}: {}", task_id, e);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ use anyhow::Result;
|
||||
use headless_chrome::{Browser, LaunchOptions, Tab};
|
||||
use reqwest::header::USER_AGENT;
|
||||
use scraper::{Html, Selector};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
use url::Url;
|
||||
use rand::Rng;
|
||||
@@ -10,6 +10,37 @@ use tokio::sync::Semaphore;
|
||||
use regex::Regex;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Lazy-initialized global browser instance
|
||||
/// Evita criar o browser no startup, economizando ~500MB de RAM até ser necessário
|
||||
static LAZY_BROWSER: OnceLock<Mutex<Option<Arc<Browser>>>> = OnceLock::new();
|
||||
|
||||
/// Obtém ou cria a instância global do browser (lazy initialization)
|
||||
pub fn get_or_create_browser() -> Result<Arc<Browser>> {
|
||||
let mutex = LAZY_BROWSER.get_or_init(|| Mutex::new(None));
|
||||
let mut guard = mutex.lock().map_err(|e| anyhow::anyhow!("Browser mutex poisoned: {}", e))?;
|
||||
|
||||
if guard.is_none() {
|
||||
log::info!("[LazyBrowser] Initializing headless browser on first use...");
|
||||
let browser = create_browser()?;
|
||||
*guard = Some(Arc::new(browser));
|
||||
log::info!("[LazyBrowser] Browser initialized successfully");
|
||||
}
|
||||
|
||||
Ok(guard.as_ref().unwrap().clone())
|
||||
}
|
||||
|
||||
/// Limpa a instância do browser (para liberar memória quando não em uso)
|
||||
pub fn clear_browser() {
|
||||
if let Some(mutex) = LAZY_BROWSER.get() {
|
||||
if let Ok(mut guard) = mutex.lock() {
|
||||
if guard.is_some() {
|
||||
log::info!("[LazyBrowser] Clearing browser instance to free memory");
|
||||
*guard = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resultado da extração de conteúdo de uma URL
|
||||
#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
|
||||
pub struct ScrapedContent {
|
||||
@@ -1302,7 +1333,7 @@ pub async fn search_and_scrape_with_config(
|
||||
if connection_closed && !failed_urls.is_empty() {
|
||||
let retry_concurrency = std::cmp::min(3, config.max_concurrent_tabs.max(1));
|
||||
let semaphore = Arc::new(Semaphore::new(retry_concurrency));
|
||||
let browser_new = Arc::new(create_browser()?);
|
||||
let browser_new = get_or_create_browser()?;
|
||||
let mut retry_handles = Vec::new();
|
||||
for url in failed_urls.clone() {
|
||||
let browser_clone = browser_new.clone();
|
||||
@@ -1341,11 +1372,79 @@ pub async fn search_and_scrape_with_config(
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Busca e extrai conteúdo de uma única URL
|
||||
/// Scraping estático usando apenas reqwest (sem headless browser)
|
||||
/// Muito mais rápido (~100ms vs ~3s) e consome menos RAM
|
||||
/// Retorna None se o conteúdo for insuficiente (SPA/JavaScript-heavy)
|
||||
pub async fn scrape_url_static(url: &str) -> Result<Option<ScrapedContent>> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(8))
|
||||
.redirect(reqwest::redirect::Policy::limited(5))
|
||||
.build()?;
|
||||
|
||||
let user_agent = get_random_user_agent();
|
||||
let start_time = Instant::now();
|
||||
|
||||
log::debug!("[StaticScrape] Fetching: {}", url);
|
||||
|
||||
let response = match client
|
||||
.get(url)
|
||||
.header(USER_AGENT, user_agent)
|
||||
.header("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8")
|
||||
.header("Accept-Language", "pt-BR,pt;q=0.9,en-US;q=0.8,en;q=0.7")
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
log::debug!("[StaticScrape] Request failed for {}: {}", url, e);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
if !response.status().is_success() {
|
||||
log::debug!("[StaticScrape] HTTP {} for {}", response.status(), url);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let html = match response.text().await {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
log::debug!("[StaticScrape] Failed to read body for {}: {}", url, e);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let duration = start_time.elapsed().as_millis();
|
||||
|
||||
// Usar extract_paragraph_fallback existente
|
||||
let result = extract_paragraph_fallback(url, &html);
|
||||
|
||||
if let Some(ref content) = result {
|
||||
log::info!("[StaticScrape] Success for {} ({} chars, {}ms)",
|
||||
url, content.content.len(), duration);
|
||||
} else {
|
||||
log::debug!("[StaticScrape] Insufficient content for {} ({}ms)", url, duration);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Busca e extrai conteúdo de uma única URL (híbrido: tenta estático primeiro)
|
||||
pub async fn scrape_url(
|
||||
url: &str,
|
||||
browser: Arc<Browser>,
|
||||
) -> Result<ScrapedContent> {
|
||||
// OTIMIZAÇÃO: Tentar scraping estático primeiro (muito mais rápido)
|
||||
if let Ok(Some(content)) = scrape_url_static(url).await {
|
||||
// Se conseguiu conteúdo suficiente (>500 chars), usar resultado estático
|
||||
if content.content.len() > 500 {
|
||||
log::info!("[ScrapeHybrid] Using static result for {} ({} chars)", url, content.content.len());
|
||||
return Ok(content);
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: usar headless browser para SPAs/JS-heavy pages
|
||||
log::info!("[ScrapeHybrid] Falling back to headless for {}", url);
|
||||
let browser_clone = browser.clone();
|
||||
let url_str = url.to_string();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
|
||||
Reference in new issue
Block a user