evandro/OllaHub
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(chat): implement speculative web search and embeddings commands

Browse Source 
- 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:
Evandro Fonseca Junior committed 2025-11-27 21:19:16 -03:00
1 parent ba03131a78
commit e53e49c8f8
7 files changed
+1407 -40
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

app/chat/page.tsx
+76 -3
View File
@@ -255,6 +255,34 @@ export default function ChatPage() {
}
}, [settings.queryPreprocessing, webSearch.isEnabled, preprocess]);
// ========== EXECUÇÃO ESPECULATIVA DE BUSCA ==========
// Heurística rápida para detectar se a query provavelmente precisa de busca web
const mightNeedWebSearch = useCallback((text: string): boolean => {
if (!webSearch.isEnabled) return false;
const lowerText = text.toLowerCase();
// Padrões que indicam alta probabilidade de precisar de busca
const questionPatterns = /\b(o que|como|quando|onde|quem|porque|por que|qual|quais|quanto|quantos|me explique|me diga|existe|há)\b/i;
const recentPatterns = /\b(hoje|ontem|esta semana|este mês|este ano|atualmente|recentemente|últimas notícias|novo|novos|última|últimas|2024|2025)\b/i;
const searchPatterns = /\?$|buscar|pesquisar|encontrar|procurar|descobrir/i;
// Se contém padrões de pergunta ou busca, provavelmente precisa de pesquisa
if (questionPatterns.test(lowerText) || recentPatterns.test(lowerText) || searchPatterns.test(lowerText)) {
chatLog.debug('[SpeculativeSearch] Query might need web search (pattern match)');
return true;
}
// Se tem mais de 5 palavras e termina com ?, provavelmente é uma pergunta
const wordCount = text.trim().split(/\s+/).length;
if (wordCount > 5 && text.trim().endsWith('?')) {
chatLog.debug('[SpeculativeSearch] Query might need web search (question detected)');
return true;
}
return false;
}, [webSearch.isEnabled]);
const handleSend = async (content: string) => {
if (!selectedModel) return;
@@ -271,6 +299,22 @@ export default function ChatPage() {
// Limpar erros anteriores
setError(null);
// ========== EXECUÇÃO ESPECULATIVA ==========
// Iniciar busca web especulativamente se a query parecer precisar
// Isso reduz latência sobrepondo busca com pré-processamento
let speculativeSearchPromise: Promise<ScrapedContent[]> | null = null;
const shouldSpeculate = mightNeedWebSearch(content);
if (shouldSpeculate) {
chatLog.info('[SpeculativeSearch] Starting speculative web search in parallel with preprocessing');
// Iniciar busca sem await - será usada depois se confirmada
speculativeSearchPromise = webSearch.smartSearchRag(content, settings.webSearch?.maxResults || 5)
.catch(err => {
chatLog.debug(`[SpeculativeSearch] Speculative search failed (will use normal flow): ${err}`);
return [] as ScrapedContent[];
});
}
// ========== PRÉ-PROCESSAMENTO ==========
let preprocessed: PreprocessedQuery;
let finalContent: string;
@@ -339,7 +383,8 @@ export default function ChatPage() {
// Função auxiliar para executar pesquisa web
const executeWebResearch = async (
content: string,
preprocessed: PreprocessedQuery
preprocessed: PreprocessedQuery,
speculativeResults?: ScrapedContent[] | null
): Promise<{
knowledgeBaseContext: string;
scrapedSources: ScrapedContent[];
@@ -353,6 +398,12 @@ export default function ChatPage() {
if (webSearch.isEnabled && preprocessed.shouldSearch) {
chatLog.info('\n--- STARTING DEEP RESEARCH PIPELINE ---');
// Se temos resultados especulativos, usá-los como seed inicial
if (speculativeResults && speculativeResults.length > 0) {
chatLog.info(`[SpeculativeSearch] Using ${speculativeResults.length} speculative results as seed`);
deepResearch.addToKnowledgeBase('speculative_seed', speculativeResults);
}
// Adicionar mensagem de processo: Pesquisa Web
addThinkingMessage('web-research', 'Pesquisando na Web...', 'running');
const webResearchStart = Date.now();
@@ -614,10 +665,25 @@ export default function ChatPage() {
// ========== LÓGICA DE DEEP RESEARCH (Knowledge Base Aggregation) ==========
// Removido: setThinkingStep e setProcessSteps - agora usando mensagens thinking agrupadas
// Executar pesquisa web
// Aguardar resultado da busca especulativa se ela foi iniciada e shouldSearch é true
let speculativeResults: ScrapedContent[] | null = null;
if (speculativeSearchPromise && preprocessed.shouldSearch) {
chatLog.info('[SpeculativeSearch] Waiting for speculative search results...');
try {
speculativeResults = await speculativeSearchPromise;
chatLog.info(`[SpeculativeSearch] Got ${speculativeResults.length} speculative results`);
} catch (err) {
chatLog.debug(`[SpeculativeSearch] Failed to get speculative results: ${err}`);
}
} else if (speculativeSearchPromise && !preprocessed.shouldSearch) {
chatLog.info('[SpeculativeSearch] Discarding speculative search (preprocessing determined search not needed)');
}
// Executar pesquisa web (usando resultados especulativos como seed se disponíveis)
const { knowledgeBaseContext, scrapedSources, validationReport } = await executeWebResearch(
content,
preprocessed
preprocessed,
speculativeResults
);
// Passo 5: Formular resposta
@@ -1151,6 +1217,13 @@ Ao responder sobre fatos atuais ou notícias, inicie mencionando explicitamente
'type' in msg.metadata &&
msg.metadata.type === 'thinking';
// Pular mensagens de sistema que não são thinking (system prompt)
// Essas mensagens nunca devem aparecer no chat
const isSystemPrompt = msg.role === 'system' && !isThinkingMessage;
if (isSystemPrompt) {
continue;
}
if (isThinkingMessage) {
const thinkingMeta = msg.metadata as ThinkingMessageMetadata;
currentThinkingGroup.push(thinkingMeta);
src-tauri/Cargo.lock
Generated
+504 -14
View File
@@ -39,6 +39,7 @@ dependencies = [
"cfg-if",
"getrandom 0.3.4",
"once_cell",
"serde",
"version_check",
"zerocopy",
]
@@ -110,6 +111,8 @@ dependencies = [
"headless_chrome",
"html2text",
"log",
"ndarray",
"ort",
"rand 0.8.5",
"readability",
"regex",
@@ -126,6 +129,7 @@ dependencies = [
"tauri-plugin-log",
"tauri-plugin-notification",
"tauri-plugin-shell",
"tokenizers",
"tokio",
"tokio-cron-scheduler",
"url",
@@ -363,7 +367,7 @@ dependencies = [
"quote",
"serde",
"serde_json",
"ureq",
"ureq 2.12.1",
]
[[package]]
@@ -372,6 +376,12 @@ version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
[[package]]
name = "base64"
version = "0.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
[[package]]
name = "base64"
version = "0.21.7"
@@ -654,6 +664,15 @@ dependencies = [
"toml 0.9.8",
]
[[package]]
name = "castaway"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
dependencies = [
"rustversion",
]
[[package]]
name = "cc"
version = "1.2.46"
@@ -689,7 +708,7 @@ version = "0.15.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02"
dependencies = [
"smallvec",
"smallvec 1.15.1",
"target-lexicon",
]
@@ -739,6 +758,21 @@ dependencies = [
"memchr",
]
[[package]]
name = "compact_str"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a"
dependencies = [
"castaway",
"cfg-if",
"itoa",
"rustversion",
"ryu",
"serde",
"static_assertions",
]
[[package]]
name = "concurrent-queue"
version = "2.5.0"
@@ -748,6 +782,19 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "console"
version = "0.15.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8"
dependencies = [
"encode_unicode",
"libc",
"once_cell",
"unicode-width 0.2.2",
"windows-sys 0.59.0",
]
[[package]]
name = "constant_time_eq"
version = "0.1.5"
@@ -858,6 +905,25 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.21"
@@ -887,7 +953,7 @@ dependencies = [
"phf 0.10.1",
"proc-macro2",
"quote",
"smallvec",
"smallvec 1.15.1",
"syn 1.0.109",
]
@@ -901,7 +967,7 @@ dependencies = [
"dtoa-short",
"itoa",
"phf 0.11.3",
"smallvec",
"smallvec 1.15.1",
]
[[package]]
@@ -994,12 +1060,31 @@ dependencies = [
"syn 2.0.110",
]
[[package]]
name = "dary_heap"
version = "0.3.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06d2e3287df1c007e74221c49ca10a95d557349e54b3a75dc2fb14712c751f04"
dependencies = [
"serde",
]
[[package]]
name = "data-encoding"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a2330da5de22e8a3cb63252ce2abb30116bf5265e89c0e01bc17015ce30a476"
[[package]]
name = "der"
version = "0.7.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
dependencies = [
"pem-rfc7468",
"zeroize",
]
[[package]]
name = "deranged"
version = "0.5.5"
@@ -1219,6 +1304,12 @@ version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "12a0bb14ac04a9fcf170d0bbbef949b44cc492f4452bd20c095636956f653642"
[[package]]
name = "either"
version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
[[package]]
name = "embed-resource"
version = "3.0.6"
@@ -1239,6 +1330,12 @@ version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ef6b89e5b37196644d8796de5268852ff179b44e96276cf4290264843743bb7"
[[package]]
name = "encode_unicode"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0"
[[package]]
name = "encoding_rs"
version = "0.8.35"
@@ -1318,6 +1415,15 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "esaxx-rs"
version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
dependencies = [
"cc",
]
[[package]]
name = "event-listener"
version = "5.4.1"
@@ -1385,6 +1491,18 @@ dependencies = [
"rustc_version",
]
[[package]]
name = "filetime"
version = "0.2.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bc0505cd1b6fa6580283f6bdf70a73fcf4aba1184038c90902b92b3dd0df63ed"
dependencies = [
"cfg-if",
"libc",
"libredox",
"windows-sys 0.60.2",
]
[[package]]
name = "find-msvc-tools"
version = "0.1.5"
@@ -1735,7 +1853,7 @@ dependencies = [
"libc",
"once_cell",
"pin-project-lite",
"smallvec",
"smallvec 1.15.1",
"thiserror 1.0.69",
]
@@ -1771,7 +1889,7 @@ dependencies = [
"libc",
"memchr",
"once_cell",
"smallvec",
"smallvec 1.15.1",
"thiserror 1.0.69",
]
@@ -2143,7 +2261,7 @@ dependencies = [
"itoa",
"pin-project-lite",
"pin-utils",
"smallvec",
"smallvec 1.15.1",
"tokio",
"want",
]
@@ -2289,7 +2407,7 @@ dependencies = [
"icu_normalizer_data",
"icu_properties",
"icu_provider",
"smallvec",
"smallvec 1.15.1",
"zerovec",
]
@@ -2347,7 +2465,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
dependencies = [
"idna_adapter",
"smallvec",
"smallvec 1.15.1",
"utf8_iter",
]
@@ -2384,6 +2502,19 @@ dependencies = [
"serde_core",
]
[[package]]
name = "indicatif"
version = "0.17.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
dependencies = [
"console",
"number_prefix",
"portable-atomic",
"unicode-width 0.2.2",
"web-time",
]
[[package]]
name = "infer"
version = "0.19.0"
@@ -2437,6 +2568,15 @@ dependencies = [
"once_cell",
]
[[package]]
name = "itertools"
version = "0.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.15"
@@ -2579,7 +2719,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e9ec52138abedcc58dc17a7c6c0c00a2bdb4f3427c7f63fa97fd0d859155caf"
dependencies = [
"gtk-sys",
"libloading",
"libloading 0.7.4",
"once_cell",
]
@@ -2599,6 +2739,16 @@ dependencies = [
"winapi",
]
[[package]]
name = "libloading"
version = "0.8.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
dependencies = [
"cfg-if",
"windows-link 0.2.1",
]
[[package]]
name = "libredox"
version = "0.1.10"
@@ -2607,6 +2757,7 @@ checksum = "416f7e718bdb06000964960ffa43b4335ad4012ae8b99060261aa4a8088d5ccb"
dependencies = [
"bitflags 2.10.0",
"libc",
"redox_syscall",
]
[[package]]
@@ -2668,6 +2819,22 @@ dependencies = [
"time",
]
[[package]]
name = "macro_rules_attribute"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520"
dependencies = [
"macro_rules_attribute-proc_macro",
"paste",
]
[[package]]
name = "macro_rules_attribute-proc_macro"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30"
[[package]]
name = "malloc_buf"
version = "0.0.6"
@@ -2734,6 +2901,16 @@ version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5"
[[package]]
name = "matrixmultiply"
version = "0.3.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
dependencies = [
"autocfg",
"rawpointer",
]
[[package]]
name = "memchr"
version = "2.7.6"
@@ -2782,6 +2959,28 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "monostate"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
dependencies = [
"monostate-impl",
"serde",
"serde_core",
]
[[package]]
name = "monostate-impl"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.110",
]
[[package]]
name = "muda"
version = "0.17.1"
@@ -2820,6 +3019,21 @@ dependencies = [
"tempfile",
]
[[package]]
name = "ndarray"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
dependencies = [
"matrixmultiply",
"num-complex",
"num-integer",
"num-traits",
"portable-atomic",
"portable-atomic-util",
"rawpointer",
]
[[package]]
name = "ndk"
version = "0.9.0"
@@ -2921,6 +3135,15 @@ dependencies = [
"winapi",
]
[[package]]
name = "num-complex"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
dependencies = [
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.1.0"
@@ -2938,6 +3161,15 @@ dependencies = [
"syn 1.0.109",
]
[[package]]
name = "num-integer"
version = "0.1.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
dependencies = [
"num-traits",
]
[[package]]
name = "num-traits"
version = "0.2.19"
@@ -2978,6 +3210,12 @@ dependencies = [
"libc",
]
[[package]]
name = "number_prefix"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
[[package]]
name = "objc"
version = "0.2.7"
@@ -3287,6 +3525,28 @@ version = "1.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
[[package]]
name = "onig"
version = "6.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0"
dependencies = [
"bitflags 2.10.0",
"libc",
"once_cell",
"onig_sys",
]
[[package]]
name = "onig_sys"
version = "69.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc"
dependencies = [
"cc",
"pkg-config",
]
[[package]]
name = "open"
version = "5.3.3"
@@ -3359,6 +3619,32 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "ort"
version = "2.0.0-rc.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fa7e49bd669d32d7bc2a15ec540a527e7764aec722a45467814005725bcd721"
dependencies = [
"libloading 0.8.9",
"ndarray",
"ort-sys",
"smallvec 2.0.0-alpha.10",
"tracing",
]
[[package]]
name = "ort-sys"
version = "2.0.0-rc.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e2aba9f5c7c479925205799216e7e5d07cc1d4fa76ea8058c60a9a30f6a4e890"
dependencies = [
"flate2",
"pkg-config",
"sha2",
"tar",
"ureq 3.1.4",
]
[[package]]
name = "os_pipe"
version = "1.2.3"
@@ -3419,7 +3705,7 @@ dependencies = [
"cfg-if",
"libc",
"redox_syscall",
"smallvec",
"smallvec 1.15.1",
"windows-link 0.2.1",
]
@@ -3434,6 +3720,12 @@ dependencies = [
"subtle",
]
[[package]]
name = "paste"
version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
[[package]]
name = "pathdiff"
version = "0.2.3"
@@ -3452,6 +3744,15 @@ dependencies = [
"sha2",
]
[[package]]
name = "pem-rfc7468"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412"
dependencies = [
"base64ct",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@@ -3677,6 +3978,21 @@ version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22686f4785f02a4fcc856d3b3bb19bf6c8160d103f7a99cc258bddd0251dc7f2"
[[package]]
name = "portable-atomic"
version = "1.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483"
[[package]]
name = "portable-atomic-util"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507"
dependencies = [
"portable-atomic",
]
[[package]]
name = "potential_utf"
version = "0.1.4"
@@ -3950,6 +4266,43 @@ version = "0.6.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
[[package]]
name = "rawpointer"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
[[package]]
name = "rayon"
version = "1.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f"
dependencies = [
"either",
"rayon-core",
]
[[package]]
name = "rayon-cond"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
dependencies = [
"either",
"itertools",
"rayon",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
dependencies = [
"crossbeam-deque",
"crossbeam-utils",
]
[[package]]
name = "readability"
version = "0.3.0"
@@ -4213,7 +4566,7 @@ dependencies = [
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
"smallvec 1.15.1",
]
[[package]]
@@ -4445,7 +4798,7 @@ dependencies = [
"phf_codegen 0.8.0",
"precomputed-hash",
"servo_arc 0.2.0",
"smallvec",
"smallvec 1.15.1",
]
[[package]]
@@ -4464,7 +4817,7 @@ dependencies = [
"phf_codegen 0.10.0",
"precomputed-hash",
"servo_arc 0.3.0",
"smallvec",
"smallvec 1.15.1",
]
[[package]]
@@ -4761,6 +5114,12 @@ version = "1.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
[[package]]
name = "smallvec"
version = "2.0.0-alpha.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51d44cfb396c3caf6fbfd0ab422af02631b69ddd96d2eff0b0f0724f9024051b"
[[package]]
name = "socket2"
version = "0.5.10"
@@ -4840,6 +5199,18 @@ dependencies = [
"system-deps",
]
[[package]]
name = "spm_precompiled"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
dependencies = [
"base64 0.13.1",
"nom",
"serde",
"unicode-segmentation",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
@@ -5074,6 +5445,17 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
[[package]]
name = "tar"
version = "0.4.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d863878d212c87a19c1a610eb53bb01fe12951c0501cf5a0d65f724914a667a"
dependencies = [
"filetime",
"libc",
"xattr",
]
[[package]]
name = "target-lexicon"
version = "0.12.16"
@@ -5522,6 +5904,40 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokenizers"
version = "0.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a620b996116a59e184c2fa2dfd8251ea34a36d0a514758c6f966386bd2e03476"
dependencies = [
"ahash 0.8.12",
"aho-corasick",
"compact_str",
"dary_heap",
"derive_builder",
"esaxx-rs",
"getrandom 0.3.4",
"indicatif",
"itertools",
"log",
"macro_rules_attribute",
"monostate",
"onig",
"paste",
"rand 0.9.2",
"rayon",
"rayon-cond",
"regex",
"regex-syntax",
"serde",
"serde_json",
"spm_precompiled",
"thiserror 2.0.17",
"unicode-normalization-alignments",
"unicode-segmentation",
"unicode_categories",
]
[[package]]
name = "tokio"
version = "1.48.0"
@@ -5885,6 +6301,15 @@ version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
[[package]]
name = "unicode-normalization-alignments"
version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
dependencies = [
"smallvec 1.15.1",
]
[[package]]
name = "unicode-segmentation"
version = "1.12.0"
@@ -5903,6 +6328,12 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
[[package]]
name = "unicode_categories"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
[[package]]
name = "untrusted"
version = "0.9.0"
@@ -5926,6 +6357,36 @@ dependencies = [
"webpki-roots 0.26.11",
]
[[package]]
name = "ureq"
version = "3.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d39cb1dbab692d82a977c0392ffac19e188bd9186a9f32806f0aaa859d75585a"
dependencies = [
"base64 0.22.1",
"der",
"log",
"native-tls",
"percent-encoding",
"rustls-pki-types",
"socks",
"ureq-proto",
"utf-8",
"webpki-root-certs",
]
[[package]]
name = "ureq-proto"
version = "0.5.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d81f9efa9df032be5934a46a068815a10a042b494b6a58cb0a1a97bb5467ed6f"
dependencies = [
"base64 0.22.1",
"http 1.3.1",
"httparse",
"log",
]
[[package]]
name = "url"
version = "2.5.7"
@@ -6151,6 +6612,16 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "web-time"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb"
dependencies = [
"js-sys",
"wasm-bindgen",
]
[[package]]
name = "webkit2gtk"
version = "2.0.1"
@@ -6195,6 +6666,15 @@ dependencies = [
"system-deps",
]
[[package]]
name = "webpki-root-certs"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ee3e3b5f5e80bc89f30ce8d0343bf4e5f12341c51f3e26cbeecbc7c85443e85b"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "webpki-roots"
version = "0.26.11"
@@ -6911,6 +7391,16 @@ dependencies = [
"pkg-config",
]
[[package]]
name = "xattr"
version = "1.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156"
dependencies = [
"libc",
"rustix",
]
[[package]]
name = "xdg-home"
version = "1.3.0"
src-tauri/Cargo.toml
+7
View File
@@ -47,3 +47,10 @@ tauri-plugin-autostart = "2"
dirs = "5.0"
rfd = "0.14"
rusqlite = { version = "0.31", features = ["bundled"] }
# ONNX Runtime para embeddings locais (ranking de relevância)
# Usando load-dynamic para evitar conflito de RuntimeLibrary (MD vs MT) no Windows
# entre ort (dynamic) e tokenizers/esaxx-rs (static)
ort = { version = "2.0.0-rc.10", features = ["load-dynamic"] }
ndarray = "0.16"
tokenizers = "0.21"
src-tauri/src/embeddings.rs
+566
View File
@@ -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);
}
}
src-tauri/src/lib.rs
+147 -14
View File
@@ -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!())
src-tauri/src/scheduler_loop.rs
+5 -6
View File
@@ -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;
}
}
src-tauri/src/web_scraper.rs
+102 -3
View File
@@ -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 || {