mirror of
https://github.com/Darkatse/TauriTavern.git
synced 2026-10-19 15:45:12 +00:00
feat(vectors): implement native compatibility layer
This commit is contained in:
parent
8925afe5bf
commit
abd51f095e
37 changed files with 4411 additions and 101 deletions
|
|
@ -56,6 +56,7 @@ src-tauri/
|
|||
├── tt-adapter-http
|
||||
├── tt-adapter-mcp
|
||||
├── tt-adapter-provider-http
|
||||
├── tt-adapter-vector
|
||||
├── tt-adapter-tokenization
|
||||
├── tt-adapter-storage-core
|
||||
├── tt-adapter-storage-userdata
|
||||
|
|
@ -76,7 +77,8 @@ src-tauri/
|
|||
| `tt-application` | 用例服务、业务编排、任务协调、policy 执行 |
|
||||
| `tt-adapter-http` | 共享 HTTP client pool/profile/helper |
|
||||
| `tt-adapter-mcp` | RMCP client、Streamable HTTP lifecycle、bounded response、tools/list pagination 与 tool validation |
|
||||
| `tt-adapter-provider-http` | LLM、SD、Translate、TTS、provider metadata 的 HTTP repository |
|
||||
| `tt-adapter-provider-http` | LLM、SD、Translate、TTS、embedding、provider metadata 的 HTTP repository |
|
||||
| `tt-adapter-vector` | Vector 的 ACID 本地索引与本地 embedding runtime;不承载 provider HTTP |
|
||||
| `tt-adapter-tokenization` | tokenizer concrete repository |
|
||||
| `tt-adapter-storage-core` | `DataDirectory`、基础文件系统 helper、chat/settings/user/theme/secret/quick reply/prompt cache/asset/llm connection/extension-store |
|
||||
| `tt-adapter-storage-userdata` | character、world info、agent workspace、agent profile、skill local package store、PNG card metadata |
|
||||
|
|
@ -260,7 +262,8 @@ adapter 是外层细节,但不是可以任意堆放的 common bucket。一个
|
|||
| 角色卡、世界书、Agent workspace/profile、Skill package | `tt-adapter-storage-userdata` | skill 是 local package store,不是普通 JSON repo |
|
||||
| 第三方扩展安装、更新、发现 | `tt-adapter-extension` | 不归入 storage-userdata |
|
||||
| avatar/background/user media/host resource 文件读取 | `tt-adapter-media` | 保持浏览器资源契约 |
|
||||
| LLM/SD/Translate/TTS/provider metadata HTTP | `tt-adapter-provider-http` | 复用 `tt-adapter-http` |
|
||||
| LLM/SD/Translate/TTS/provider metadata/embedding HTTP | `tt-adapter-provider-http` | 复用 `tt-adapter-http` |
|
||||
| Vector 索引与本地 embedding runtime | `tt-adapter-vector` | 独立持久化/推理依赖,不下沉到 storage-core |
|
||||
| 通用 HTTP pool/profile/helper | `tt-adapter-http` | 不放 provider 业务规则 |
|
||||
| LAN/TT Sync runtime | `tt-adapter-sync` | Tauri event/UI adapter 留 host composition |
|
||||
| Data Archive import/export executor | `tt-adapter-archive` | Tauri picker/share 留 host infrastructure |
|
||||
|
|
@ -318,6 +321,7 @@ TauriTavern 私有状态放在 `_tauritavern` 下,例如 agent workspace、age
|
|||
| 媒体 Range / browser resource contract | `docs/CurrentState/MediaAssetContract.md` |
|
||||
| Logging / Dev observability | `docs/CurrentState/LoggingObservability.md` |
|
||||
| Native provider API formats | `docs/CurrentState/NativeApiFormats.md` |
|
||||
| Vector 兼容层 | `docs/CurrentState/VectorApi.md` |
|
||||
| Sync | `docs/CurrentState/Sync.md` |
|
||||
| Agent 总览 | `docs/AgentArchitecture.md` |
|
||||
| Agent 细节 | `docs/Agent/README.md` |
|
||||
|
|
|
|||
|
|
@ -98,3 +98,7 @@
|
|||
21. `docs/CurrentState/MCP.md`
|
||||
- MCP registration、persistent catalog、user test call 与模型工具调用的当前实现快照
|
||||
- 包含 crate/Host ABI、identity、outcome、cache、Agent/Legacy 消费、故障域与明确未支持边界
|
||||
|
||||
22. `docs/CurrentState/VectorApi.md`
|
||||
- SillyTavern 1.18.0 Vector 七路由兼容层的当前实现快照
|
||||
- 包含 typed command 链路、redb 事务索引、本地/远端 embedding 边界、精确 cosine 基线与明确未支持范围
|
||||
|
|
|
|||
57
docs/CurrentState/VectorApi.md
Normal file
57
docs/CurrentState/VectorApi.md
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
# Vector 兼容层当前契约
|
||||
|
||||
本文记录 TauriTavern 对 SillyTavern 1.18.0 Vector 扩展的原生兼容边界。
|
||||
|
||||
## 1. 可观察契约
|
||||
|
||||
同源兼容层实现以下 `POST` 路由:
|
||||
|
||||
- `/api/vector/list`
|
||||
- `/api/vector/insert`
|
||||
- `/api/vector/delete`
|
||||
- `/api/vector/query`
|
||||
- `/api/vector/query-multi`
|
||||
- `/api/vector/purge`
|
||||
- `/api/vector/purge-all`
|
||||
|
||||
前端请求经 `src/tauri/main/routes/vector-routes.js` 汇入一个 typed Tauri command,presentation 只转交 DTO,输入校验、provider 选择、secret 与 iOS policy 均由 `VectorService` 编排。
|
||||
|
||||
## 2. 边界与存储
|
||||
|
||||
```text
|
||||
same-origin route
|
||||
-> vector_handle
|
||||
-> VectorService
|
||||
-> VectorRepository -> tt-adapter-vector / redb
|
||||
-> LocalEmbeddingRepository -> tt-adapter-vector / FastEmbed
|
||||
-> RemoteEmbeddingRepository -> tt-adapter-provider-http
|
||||
```
|
||||
|
||||
索引位于 `default-user/vectors/tauritavern-v1.redb`。模型缓存位于 `_cache/embedding-models/`,purge 只删除索引记录,不重复下载模型。数据库与模型均懒加载;初始化失败只令当前 Vector 请求失败,不影响应用启动。
|
||||
|
||||
每个 scope 由 collection、source、model 以及会改变 embedding 空间的 endpoint/account 设置共同确定。插入先完成整个 batch 的 embedding、数量/维度/有限值/非零校验与 L2 归一化,再以单个 redb write transaction 提交 metadata 和 `f32` bytes。相同 hash/index/text 的重试是幂等 upsert;delete 仍按上游 hash 语义删除所有分块。
|
||||
|
||||
文件索引以整份文件作为一个 `/insert` 工作单元;`VectorService` 可在内部对 provider 分批请求,但只有全部 embedding 成功后才提交索引,避免可恢复的中途失败留下“已有 hash、实际不完整”的状态。
|
||||
|
||||
## 3. 检索基线
|
||||
|
||||
当前使用归一化向量的精确 cosine scan,并按分数降序返回。`query-multi` 只计算一次 query embedding,再跨 collection 做全局 threshold 与 top-K。`hashes` 和 `metadata` 始终来自同一组 threshold 后的结果,避免上游实现的数组错位。
|
||||
|
||||
这是有意选择的 60/95 基线:聊天与文件集合通常远小于需要 ANN 生命周期的规模。只有真实 profile 显示 collection scan 延迟不可接受时,才在 `VectorRepository` 后替换为 ANN;GraphRAG、知识图谱抽取、reranker 和后台索引任务不属于 SillyTavern Vector 兼容层。
|
||||
|
||||
## 4. Embedding source
|
||||
|
||||
- `transformers`:Vector 设置可选择以下本地模型,首次使用时下载到 `_cache/embedding-models/` 并缓存;未携带 `model` 的旧调用继续使用 Jina v2,避免升级时静默切换已有索引。
|
||||
- `jinaai/jina-embeddings-v2-base-en`:兼容旧默认。
|
||||
- `BAAI/bge-m3`:当前只消费 dense 输出。
|
||||
- `Qwen/Qwen3-Embedding-0.6B`:query 使用检索 instruction,document 保持原文。
|
||||
- `google/embeddinggemma-300m`:使用 FastEmbed 的 Q4 ONNX 版本,query/document 分别使用官方 retrieval prompt。
|
||||
- `webllm`:沿用前端预计算映射。
|
||||
- `koboldcpp`:兼容 `/api/backends/kobold/embed`,并把服务端报告的实际 embedding model 带入 Vector scope,避免同一 endpoint 换模型后混合向量空间。
|
||||
- OpenAI-compatible、Cohere、Nomic、Google AI Studio、Vertex AI、Extras、Ollama、llama.cpp、vLLM:由 provider HTTP adapter 按各自协议发送。
|
||||
|
||||
远端 source 在调用前读取 active secret,并遵守 iOS source allowlist 与 endpoint override policy。未知 source、缺失凭据、非法 URL、provider 数量或维度不一致都会显式失败,不创建或部分写入索引。
|
||||
|
||||
与上游一致,Vector 兼容层不区分交互式与批量请求,也不额外施加生成期或 embedding 专属超时;调用会等待本地推理或远端请求完成并正常传播错误。
|
||||
|
||||
本地模型 id、量化/runtime 版本、最大 token 长度和 prompt 版本共同组成 Vector profile。模型切换不会混合向量空间;repository 同时只保留一个已加载 runtime,切换时释放旧模型以限制常驻内存。
|
||||
|
|
@ -121,6 +121,7 @@ const APPLICATION_FORBIDDEN_PACKAGES = new Set([
|
|||
'tt-adapter-storage-userdata',
|
||||
'tt-adapter-sync',
|
||||
'tt-adapter-tokenization',
|
||||
'tt-adapter-vector',
|
||||
'yup-oauth2',
|
||||
'zip',
|
||||
]);
|
||||
|
|
@ -149,6 +150,7 @@ const APPLICATION_FORBIDDEN_SOURCE_PATTERNS = [
|
|||
['tt-adapter-storage-userdata', /\btt_adapter_storage_userdata::/],
|
||||
['tt-adapter-sync', /\btt_adapter_sync::/],
|
||||
['tt-adapter-tokenization', /\btt_adapter_tokenization::/],
|
||||
['tt-adapter-vector', /\btt_adapter_vector::/],
|
||||
['main crate', /\btauritavern(_lib)?::/],
|
||||
['provider oauth client', /\byup_oauth2::/],
|
||||
['concrete hyper client', /\bhyper_util::/],
|
||||
|
|
@ -297,6 +299,50 @@ const ADAPTER_PROVIDER_HTTP_FORBIDDEN_SOURCE_PATTERNS = [
|
|||
['zip', /\bzip::/],
|
||||
];
|
||||
|
||||
const ADAPTER_VECTOR_FORBIDDEN_PACKAGES = new Set([
|
||||
'axum',
|
||||
'miktik',
|
||||
'qrcode',
|
||||
'reqwest',
|
||||
'rmcp',
|
||||
'tar',
|
||||
'tauri',
|
||||
'tauritavern',
|
||||
'tt-adapter-archive',
|
||||
'tt-adapter-extension',
|
||||
'tt-adapter-http',
|
||||
'tt-adapter-media',
|
||||
'tt-adapter-mcp',
|
||||
'tt-adapter-provider-http',
|
||||
'tt-adapter-storage-core',
|
||||
'tt-adapter-storage-userdata',
|
||||
'tt-adapter-sync',
|
||||
'tt-adapter-tokenization',
|
||||
'tt-application',
|
||||
'ttsync-core',
|
||||
]);
|
||||
|
||||
const ADAPTER_VECTOR_FORBIDDEN_SOURCE_PATTERNS = [
|
||||
...ADAPTER_FORBIDDEN_SOURCE_PATTERNS,
|
||||
['axum', /\baxum::/],
|
||||
['image', /\bimage::/],
|
||||
['miktik', /\bmiktik::/],
|
||||
['qrcode', /\bqrcode::/],
|
||||
['reqwest', /\breqwest::/],
|
||||
['rmcp', /\brmcp::/],
|
||||
['tar', /\btar::/],
|
||||
['tt-adapter-http', /\btt_adapter_http::/],
|
||||
['tt-adapter-media', /\btt_adapter_media::/],
|
||||
['tt-adapter-mcp', /\btt_adapter_mcp::/],
|
||||
['tt-adapter-provider-http', /\btt_adapter_provider_http::/],
|
||||
['tt-adapter-storage-core', /\btt_adapter_storage_core::/],
|
||||
['tt-adapter-storage-userdata', /\btt_adapter_storage_userdata::/],
|
||||
['tt-adapter-sync', /\btt_adapter_sync::/],
|
||||
['tt-adapter-tokenization', /\btt_adapter_tokenization::/],
|
||||
['ttsync-core', /\bttsync_core::/],
|
||||
['zip', /\bzip::/],
|
||||
];
|
||||
|
||||
const ADAPTER_MCP_FORBIDDEN_PACKAGES = new Set([
|
||||
'async-compression',
|
||||
'axum',
|
||||
|
|
@ -495,6 +541,7 @@ const CRATES = [
|
|||
crateConfig('tt-adapter-sync', ADAPTER_SYNC_FORBIDDEN_PACKAGES, ADAPTER_SYNC_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-archive', ADAPTER_ARCHIVE_FORBIDDEN_PACKAGES, ADAPTER_ARCHIVE_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-provider-http', ADAPTER_PROVIDER_HTTP_FORBIDDEN_PACKAGES, ADAPTER_PROVIDER_HTTP_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-vector', ADAPTER_VECTOR_FORBIDDEN_PACKAGES, ADAPTER_VECTOR_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-extension', ADAPTER_EXTENSION_FORBIDDEN_PACKAGES, ADAPTER_EXTENSION_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-storage-core', ADAPTER_STORAGE_CORE_FORBIDDEN_PACKAGES, ADAPTER_STORAGE_CORE_FORBIDDEN_SOURCE_PATTERNS),
|
||||
crateConfig('tt-adapter-storage-userdata', ADAPTER_STORAGE_USERDATA_FORBIDDEN_PACKAGES, ADAPTER_STORAGE_USERDATA_FORBIDDEN_SOURCE_PATTERNS),
|
||||
|
|
|
|||
1063
src-tauri/Cargo.lock
generated
1063
src-tauri/Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -11,6 +11,7 @@ members = [
|
|||
"crates/tt-adapter-sync",
|
||||
"crates/tt-adapter-archive",
|
||||
"crates/tt-adapter-provider-http",
|
||||
"crates/tt-adapter-vector",
|
||||
"crates/tt-adapter-extension",
|
||||
"crates/tt-adapter-storage-core",
|
||||
"crates/tt-adapter-storage-userdata",
|
||||
|
|
|
|||
|
|
@ -57,6 +57,7 @@ tt-adapter-sync = { path = "../tt-adapter-sync" }
|
|||
tt-adapter-archive = { path = "../tt-adapter-archive" }
|
||||
tt-adapter-extension = { path = "../tt-adapter-extension" }
|
||||
tt-adapter-provider-http = { path = "../tt-adapter-provider-http" }
|
||||
tt-adapter-vector = { path = "../tt-adapter-vector" }
|
||||
tt-adapter-storage-core = { path = "../tt-adapter-storage-core" }
|
||||
tt-adapter-storage-userdata = { path = "../tt-adapter-storage-userdata" }
|
||||
tt-adapter-media = { path = "../tt-adapter-media" }
|
||||
|
|
|
|||
|
|
@ -14,8 +14,8 @@ use tt_adapter_media::{
|
|||
FileAvatarRepository, FileBackgroundRepository, FileImageMetadataRepository,
|
||||
};
|
||||
use tt_adapter_provider_http::{
|
||||
HttpChatCompletionRepository, HttpProviderMetadataRepository, HttpStableDiffusionRepository,
|
||||
HttpTranslateRepository, HttpTtsRepository,
|
||||
HttpChatCompletionRepository, HttpEmbeddingRepository, HttpProviderMetadataRepository,
|
||||
HttpStableDiffusionRepository, HttpTranslateRepository, HttpTtsRepository,
|
||||
};
|
||||
use tt_adapter_storage_core::{
|
||||
DataDirectory, FileAssetRepository, FileChatRepository, FileExtensionStoreRepository,
|
||||
|
|
@ -30,6 +30,7 @@ use tt_adapter_storage_userdata::FileCharacterRepository;
|
|||
use tt_adapter_storage_userdata::FileWorldInfoRepository;
|
||||
use tt_adapter_storage_userdata::{FileSkillRepository, FileSpriteRepository};
|
||||
use tt_adapter_tokenization::MiktikTokenizerRepository;
|
||||
use tt_adapter_vector::{FastEmbedLocalEmbeddingRepository, RedbVectorRepository};
|
||||
use tt_domain::errors::DomainError;
|
||||
use tt_domain::models::settings::ChatBackupSettings;
|
||||
use tt_ports::repositories::agent_invocation_repository::AgentInvocationRepository;
|
||||
|
|
@ -69,6 +70,9 @@ use tt_ports::repositories::tts_repository::TtsRepository;
|
|||
use tt_ports::repositories::update_repository::UpdateRepository;
|
||||
use tt_ports::repositories::user_directory_repository::UserDirectoryRepository;
|
||||
use tt_ports::repositories::user_repository::UserRepository;
|
||||
use tt_ports::repositories::vector_repository::{
|
||||
LocalEmbeddingRepository, RemoteEmbeddingRepository, VectorRepository,
|
||||
};
|
||||
use tt_ports::repositories::workspace_repository::WorkspaceRepository;
|
||||
use tt_ports::repositories::world_info_repository::WorldInfoRepository;
|
||||
use tt_ports::settings::ChatBackupRuntime;
|
||||
|
|
@ -118,6 +122,9 @@ pub(in crate::app::composition) struct AppRepositories {
|
|||
pub(in crate::app::composition) tts_repository: Arc<dyn TtsRepository>,
|
||||
pub(in crate::app::composition) world_info_repository: Arc<dyn WorldInfoRepository>,
|
||||
pub(in crate::app::composition) update_repository: Arc<dyn UpdateRepository>,
|
||||
pub(in crate::app::composition) vector_repository: Arc<dyn VectorRepository>,
|
||||
pub(in crate::app::composition) remote_embedding_repository: Arc<dyn RemoteEmbeddingRepository>,
|
||||
pub(in crate::app::composition) local_embedding_repository: Arc<dyn LocalEmbeddingRepository>,
|
||||
}
|
||||
|
||||
pub(super) async fn build(
|
||||
|
|
@ -296,6 +303,16 @@ pub(super) async fn build(
|
|||
let update_repository: Arc<dyn UpdateRepository> =
|
||||
Arc::new(GitHubUpdateRepository::new(http_client_pool.clone()));
|
||||
|
||||
let vector_root = default_user_dir.join("vectors");
|
||||
let vector_repository: Arc<dyn VectorRepository> = Arc::new(RedbVectorRepository::new(
|
||||
vector_root.join("tauritavern-v1.redb"),
|
||||
));
|
||||
let remote_embedding_repository: Arc<dyn RemoteEmbeddingRepository> =
|
||||
Arc::new(HttpEmbeddingRepository::new(http_client_pool));
|
||||
let local_embedding_repository: Arc<dyn LocalEmbeddingRepository> = Arc::new(
|
||||
FastEmbedLocalEmbeddingRepository::new(data_root.join("_cache").join("embedding-models")),
|
||||
);
|
||||
|
||||
Ok(AppRepositories {
|
||||
character_repository,
|
||||
chat_repository,
|
||||
|
|
@ -337,5 +354,8 @@ pub(super) async fn build(
|
|||
tts_repository,
|
||||
world_info_repository,
|
||||
update_repository,
|
||||
vector_repository,
|
||||
remote_embedding_repository,
|
||||
local_embedding_repository,
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ use tt_application::services::tts_service::TtsService;
|
|||
use tt_application::services::update_service::UpdateService;
|
||||
use tt_application::services::user_directory_service::UserDirectoryService;
|
||||
use tt_application::services::user_service::UserService;
|
||||
use tt_application::services::vector_service::VectorService;
|
||||
use tt_application::services::world_info_service::WorldInfoService;
|
||||
use tt_domain::errors::DomainError;
|
||||
|
||||
|
|
@ -125,6 +126,13 @@ pub(super) async fn build(
|
|||
repositories.secret_repository.clone(),
|
||||
ios_policy.clone(),
|
||||
));
|
||||
let vector_service = Arc::new(VectorService::new(
|
||||
repositories.vector_repository.clone(),
|
||||
repositories.remote_embedding_repository.clone(),
|
||||
repositories.local_embedding_repository.clone(),
|
||||
repositories.secret_repository.clone(),
|
||||
ios_policy.clone(),
|
||||
));
|
||||
let agent_services = agent::build(
|
||||
&repositories,
|
||||
skill_service.clone(),
|
||||
|
|
@ -255,6 +263,7 @@ pub(super) async fn build(
|
|||
llm_connection_service,
|
||||
mcp_service,
|
||||
provider_metadata_service,
|
||||
vector_service,
|
||||
tokenization_service,
|
||||
stable_diffusion_service,
|
||||
translate_service,
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ use tt_application::services::tts_service::TtsService;
|
|||
use tt_application::services::update_service::UpdateService;
|
||||
use tt_application::services::user_directory_service::UserDirectoryService;
|
||||
use tt_application::services::user_service::UserService;
|
||||
use tt_application::services::vector_service::VectorService;
|
||||
use tt_application::services::world_info_service::WorldInfoService;
|
||||
use tt_domain::errors::DomainError;
|
||||
use tt_domain::ios_policy::IosPolicyActivationReport;
|
||||
|
|
@ -107,6 +108,7 @@ pub(crate) struct AppServices {
|
|||
pub(crate) llm_connection_service: Arc<LlmConnectionService>,
|
||||
pub(crate) mcp_service: Arc<McpService>,
|
||||
pub(crate) provider_metadata_service: Arc<ProviderMetadataService>,
|
||||
pub(crate) vector_service: Arc<VectorService>,
|
||||
pub(crate) tokenization_service: Arc<TokenizationService>,
|
||||
pub(crate) stable_diffusion_service: Arc<StableDiffusionService>,
|
||||
pub(crate) translate_service: Arc<TranslateService>,
|
||||
|
|
|
|||
|
|
@ -52,4 +52,5 @@ pub mod update_commands;
|
|||
pub mod upload_staging_commands;
|
||||
pub mod user_commands;
|
||||
pub mod user_directory_commands;
|
||||
pub mod vector_commands;
|
||||
pub mod world_info_commands;
|
||||
|
|
|
|||
|
|
@ -365,6 +365,7 @@ pub fn invoke_handler() -> impl Fn(tauri::ipc::Invoke<tauri::Wry>) -> bool + Sen
|
|||
// Stable diffusion (local chain) commands
|
||||
super::stable_diffusion_commands::sd_handle,
|
||||
super::stable_diffusion_commands::cancel_sd_request,
|
||||
super::vector_commands::vector_handle,
|
||||
// Translate commands
|
||||
super::translate_commands::translate_text,
|
||||
// TTS commands
|
||||
|
|
|
|||
|
|
@ -0,0 +1,23 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use tauri::State;
|
||||
|
||||
use crate::app::AppState;
|
||||
use crate::presentation::commands::helpers::{log_command, map_command_error};
|
||||
use crate::presentation::errors::CommandError;
|
||||
use tt_application::dto::vector_dto::{VectorRouteRequestDto, VectorRouteResponseDto};
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn vector_handle(
|
||||
path: String,
|
||||
request: VectorRouteRequestDto,
|
||||
app_state: State<'_, Arc<AppState>>,
|
||||
) -> Result<VectorRouteResponseDto, CommandError> {
|
||||
log_command(format!("vector_handle {}", path.trim()));
|
||||
app_state
|
||||
.services
|
||||
.vector_service
|
||||
.handle_request(&path, request)
|
||||
.await
|
||||
.map_err(map_command_error("Vector request failed"))
|
||||
}
|
||||
|
|
@ -6,15 +6,14 @@ publish = false
|
|||
|
||||
[dependencies]
|
||||
reqwest = { workspace = true, features = ["blocking", "http2", "rustls", "socks"] }
|
||||
rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs", "std", "tls12"] }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tt-domain = { path = "../tt-domain" }
|
||||
tt-ports = { path = "../tt-ports" }
|
||||
|
||||
[target.'cfg(target_os = "android")'.dependencies]
|
||||
rustls = "0.23"
|
||||
webpki-roots = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
rcgen = "0.14"
|
||||
rustls = "0.23"
|
||||
|
|
|
|||
|
|
@ -57,6 +57,7 @@ pub struct HttpClientPool {
|
|||
|
||||
impl HttpClientPool {
|
||||
pub fn new(product_user_agent: impl Into<String>) -> Self {
|
||||
install_rustls_crypto_provider();
|
||||
let product_user_agent = product_user_agent.into();
|
||||
assert!(
|
||||
!product_user_agent.trim().is_empty(),
|
||||
|
|
@ -140,6 +141,15 @@ impl HttpClientPool {
|
|||
}
|
||||
}
|
||||
|
||||
fn install_rustls_crypto_provider() {
|
||||
static INSTALL: std::sync::Once = std::sync::Once::new();
|
||||
INSTALL.call_once(|| {
|
||||
// Workspace dependencies may compile both rustls providers. Choosing the provider already
|
||||
// used by reqwest keeps every TLS consumer deterministic instead of relying on features.
|
||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||
});
|
||||
}
|
||||
|
||||
impl RequestProxyRuntime for HttpClientPool {
|
||||
fn validate_request_proxy_settings(
|
||||
&self,
|
||||
|
|
@ -574,10 +584,11 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn git_blocking_builder_validates_server_certificates() {
|
||||
let pool = pool();
|
||||
let (config, root) = test_tls_config();
|
||||
|
||||
let (untrusted_url, untrusted_request, untrusted_handle) = tls_server(Arc::clone(&config));
|
||||
let client = pool().git_blocking_client_builder().build().unwrap();
|
||||
let client = pool.git_blocking_client_builder().build().unwrap();
|
||||
assert!(client.get(untrusted_url).send().is_err());
|
||||
assert!(
|
||||
!untrusted_request
|
||||
|
|
@ -587,7 +598,7 @@ mod tests {
|
|||
untrusted_handle.join().expect("untrusted TLS server");
|
||||
|
||||
let (trusted_url, trusted_request, trusted_handle) = tls_server(config);
|
||||
let client = pool()
|
||||
let client = pool
|
||||
.git_blocking_client_builder()
|
||||
.tls_certs_only([root])
|
||||
.build()
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ mod openai;
|
|||
mod openai_responses;
|
||||
mod response_body;
|
||||
mod vertexai;
|
||||
mod vertexai_auth;
|
||||
pub(crate) mod vertexai_auth;
|
||||
mod workers_ai;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ struct CachedServiceAccount {
|
|||
static SERVICE_ACCOUNT_CACHE: OnceLock<RwLock<HashMap<String, CachedServiceAccount>>> =
|
||||
OnceLock::new();
|
||||
|
||||
pub(super) async fn get_service_account_access_token(
|
||||
pub(crate) async fn get_service_account_access_token(
|
||||
service_account_json: &str,
|
||||
) -> Result<String, DomainError> {
|
||||
let cache_key = sha256_hex(service_account_json);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,659 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
|
||||
use reqwest::{Client, RequestBuilder, StatusCode, Url};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use tt_adapter_http::{HttpClientPool, HttpClientProfile};
|
||||
use tt_domain::errors::DomainError;
|
||||
use tt_ports::repositories::vector_repository::{
|
||||
RemoteEmbeddingBatch, RemoteEmbeddingProtocol, RemoteEmbeddingRepository,
|
||||
RemoteEmbeddingRequest, VertexEmbeddingAuth,
|
||||
};
|
||||
|
||||
use crate::http_chat_completion_repository::vertexai_auth;
|
||||
|
||||
const NOMIC_EMBEDDING_URL: &str = "https://api-atlas.nomic.ai/v1/embedding/text";
|
||||
const GOOGLE_AI_BASE: &str = "https://generativelanguage.googleapis.com/v1beta";
|
||||
|
||||
pub struct HttpEmbeddingRepository {
|
||||
http_clients: Arc<HttpClientPool>,
|
||||
}
|
||||
|
||||
impl HttpEmbeddingRepository {
|
||||
pub fn new(http_clients: Arc<HttpClientPool>) -> Self {
|
||||
Self { http_clients }
|
||||
}
|
||||
|
||||
fn client(&self) -> Result<Client, DomainError> {
|
||||
self.http_clients.client(HttpClientProfile::Default)
|
||||
}
|
||||
|
||||
async fn send_json(
|
||||
&self,
|
||||
request: RequestBuilder,
|
||||
provider: &str,
|
||||
) -> Result<Value, DomainError> {
|
||||
let response = request.send().await.map_err(|error| {
|
||||
DomainError::upstream_failure(crate::http_error::reqwest_transport_failure(&error))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(map_error_response(provider, response).await);
|
||||
}
|
||||
|
||||
let endpoint = response.url().clone();
|
||||
response.json().await.map_err(|error| {
|
||||
if error.is_decode() {
|
||||
DomainError::InternalError(format!(
|
||||
"{provider} embedding response is not valid JSON: {error}"
|
||||
))
|
||||
} else {
|
||||
DomainError::upstream_failure(crate::http_error::reqwest_body_failure(
|
||||
&error,
|
||||
Some(&endpoint),
|
||||
))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn embed_vertex(
|
||||
&self,
|
||||
model: String,
|
||||
region: String,
|
||||
auth: VertexEmbeddingAuth,
|
||||
texts: Vec<String>,
|
||||
) -> Result<Vec<Vec<f32>>, DomainError> {
|
||||
let region = normalized_region(®ion);
|
||||
let host = vertex_host(®ion);
|
||||
let instances = texts
|
||||
.into_iter()
|
||||
.map(|content| json!({ "content": content }))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let request = match auth {
|
||||
VertexEmbeddingAuth::Express {
|
||||
api_key,
|
||||
project_id,
|
||||
} => {
|
||||
let path = match project_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|id| !id.is_empty())
|
||||
{
|
||||
Some(project_id) => format!(
|
||||
"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:predict"
|
||||
),
|
||||
None => format!("/v1/publishers/google/models/{model}:predict"),
|
||||
};
|
||||
let url = Url::parse(&format!("{host}{path}")).map_err(invalid_url("Vertex AI"))?;
|
||||
self.client()?
|
||||
.post(url)
|
||||
.header("x-goog-api-key", api_key)
|
||||
.json(&json!({ "instances": instances }))
|
||||
}
|
||||
VertexEmbeddingAuth::ServiceAccount { json: credentials } => {
|
||||
let project_id = service_account_project_id(&credentials)?;
|
||||
let token = vertexai_auth::get_service_account_access_token(&credentials).await?;
|
||||
let url = Url::parse(&format!(
|
||||
"{host}/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:predict"
|
||||
))
|
||||
.map_err(invalid_url("Vertex AI"))?;
|
||||
self.client()?
|
||||
.post(url)
|
||||
.header(AUTHORIZATION, format!("Bearer {token}"))
|
||||
.json(&json!({ "instances": instances }))
|
||||
}
|
||||
};
|
||||
|
||||
let body = self.send_json(request, "Google Vertex AI").await?;
|
||||
let predictions = required_array(&body, "predictions", "Google Vertex AI")?;
|
||||
predictions
|
||||
.iter()
|
||||
.map(|prediction| {
|
||||
parse_embedding(
|
||||
prediction
|
||||
.get("embeddings")
|
||||
.and_then(|value| value.get("values"))
|
||||
.ok_or_else(|| {
|
||||
invalid_response("Google Vertex AI", "missing embeddings.values")
|
||||
})?,
|
||||
"Google Vertex AI",
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RemoteEmbeddingRepository for HttpEmbeddingRepository {
|
||||
async fn embed(
|
||||
&self,
|
||||
request: RemoteEmbeddingRequest,
|
||||
) -> Result<RemoteEmbeddingBatch, DomainError> {
|
||||
let RemoteEmbeddingRequest {
|
||||
protocol,
|
||||
texts,
|
||||
is_query,
|
||||
} = request;
|
||||
if texts.is_empty() && !matches!(&protocol, RemoteEmbeddingProtocol::KoboldCpp { .. }) {
|
||||
return Ok(RemoteEmbeddingBatch {
|
||||
embeddings: Vec::new(),
|
||||
reported_model: None,
|
||||
});
|
||||
}
|
||||
let expected_count = texts.len();
|
||||
let mut reported_model = None;
|
||||
|
||||
let embeddings = match protocol {
|
||||
RemoteEmbeddingProtocol::OpenAi {
|
||||
provider,
|
||||
base_url,
|
||||
api_key,
|
||||
model,
|
||||
omit_model,
|
||||
headers,
|
||||
} => {
|
||||
let url = joined_url(&base_url, "embeddings")?;
|
||||
let mut body = json!({ "input": texts, "model": model });
|
||||
if omit_model {
|
||||
body["model"] = Value::Null;
|
||||
}
|
||||
|
||||
let mut request = self
|
||||
.client()?
|
||||
.post(url)
|
||||
.header(AUTHORIZATION, format!("Bearer {api_key}"))
|
||||
.json(&body);
|
||||
for (name, value) in headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
let body = self.send_json(request, &provider).await?;
|
||||
parse_indexed_embeddings(&body, "data", &provider)
|
||||
}
|
||||
RemoteEmbeddingProtocol::Cohere { api_key, model } => {
|
||||
let request = self
|
||||
.client()?
|
||||
.post("https://api.cohere.ai/v2/embed")
|
||||
.header(AUTHORIZATION, format!("Bearer {api_key}"))
|
||||
.json(&json!({
|
||||
"texts": texts,
|
||||
"model": model,
|
||||
"embedding_types": ["float"],
|
||||
"input_type": if is_query { "search_query" } else { "search_document" },
|
||||
"truncate": "END",
|
||||
}));
|
||||
let body = self.send_json(request, "Cohere").await?;
|
||||
parse_embedding_array(
|
||||
body.get("embeddings")
|
||||
.and_then(|value| value.get("float"))
|
||||
.ok_or_else(|| invalid_response("Cohere", "missing embeddings.float"))?,
|
||||
"Cohere",
|
||||
)
|
||||
}
|
||||
RemoteEmbeddingProtocol::Nomic { api_key } => {
|
||||
let request = self
|
||||
.client()?
|
||||
.post(NOMIC_EMBEDDING_URL)
|
||||
.header(AUTHORIZATION, format!("Bearer {api_key}"))
|
||||
.json(&json!({
|
||||
"texts": texts,
|
||||
"model": "nomic-embed-text-v1.5",
|
||||
}));
|
||||
let body = self.send_json(request, "Nomic AI").await?;
|
||||
parse_embedding_array(
|
||||
body.get("embeddings")
|
||||
.ok_or_else(|| invalid_response("Nomic AI", "missing embeddings"))?,
|
||||
"Nomic AI",
|
||||
)
|
||||
}
|
||||
RemoteEmbeddingProtocol::Extras { base_url, api_key } => {
|
||||
let url = root_path_url(&base_url, "/api/embeddings/compute", "Extras")?;
|
||||
let mut request = self.client()?.post(url).json(&json!({ "text": texts }));
|
||||
if let Some(api_key) = api_key.filter(|value| !value.trim().is_empty()) {
|
||||
request = request.header(AUTHORIZATION, format!("Bearer {api_key}"));
|
||||
}
|
||||
let body = self.send_json(request, "Extras").await?;
|
||||
parse_embedding_array(
|
||||
body.get("embedding")
|
||||
.ok_or_else(|| invalid_response("Extras", "missing embedding"))?,
|
||||
"Extras",
|
||||
)
|
||||
}
|
||||
RemoteEmbeddingProtocol::GoogleAiStudio { api_key, model } => {
|
||||
let url = Url::parse(&format!(
|
||||
"{GOOGLE_AI_BASE}/models/{model}:batchEmbedContents"
|
||||
))
|
||||
.map_err(invalid_url("Google AI Studio"))?;
|
||||
let requests = texts
|
||||
.into_iter()
|
||||
.map(|text| {
|
||||
json!({
|
||||
"model": format!("models/{model}"),
|
||||
"content": { "parts": [{ "text": text }] },
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let request = self
|
||||
.client()?
|
||||
.post(url)
|
||||
.header("x-goog-api-key", api_key)
|
||||
.json(&json!({ "requests": requests }));
|
||||
let body = self.send_json(request, "Google AI Studio").await?;
|
||||
required_array(&body, "embeddings", "Google AI Studio")?
|
||||
.iter()
|
||||
.map(|entry| {
|
||||
parse_embedding(
|
||||
entry.get("values").ok_or_else(|| {
|
||||
invalid_response("Google AI Studio", "missing embeddings.values")
|
||||
})?,
|
||||
"Google AI Studio",
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
RemoteEmbeddingProtocol::VertexAi {
|
||||
model,
|
||||
region,
|
||||
auth,
|
||||
} => self.embed_vertex(model, region, auth, texts).await,
|
||||
RemoteEmbeddingProtocol::Ollama {
|
||||
base_url,
|
||||
model,
|
||||
keep,
|
||||
} => {
|
||||
let url = root_path_url(&base_url, "/api/embed", "Ollama")?;
|
||||
let mut body = Map::from_iter([
|
||||
("input".to_string(), json!(texts)),
|
||||
("model".to_string(), json!(model)),
|
||||
("truncate".to_string(), Value::Bool(true)),
|
||||
]);
|
||||
if keep {
|
||||
body.insert("keep_alive".to_string(), json!(-1));
|
||||
}
|
||||
let request = self.client()?.post(url).json(&body);
|
||||
let body = self.send_json(request, "Ollama").await?;
|
||||
parse_embedding_array(
|
||||
body.get("embeddings")
|
||||
.ok_or_else(|| invalid_response("Ollama", "missing embeddings"))?,
|
||||
"Ollama",
|
||||
)
|
||||
}
|
||||
RemoteEmbeddingProtocol::LlamaCpp { base_url, api_key } => {
|
||||
let url = v1_embeddings_url(&base_url, "Llama.cpp")?;
|
||||
let request = optional_bearer(
|
||||
self.client()?.post(url).json(&json!({ "input": texts })),
|
||||
api_key,
|
||||
);
|
||||
let body = self.send_json(request, "Llama.cpp").await?;
|
||||
parse_indexed_embeddings(&body, "data", "Llama.cpp")
|
||||
}
|
||||
RemoteEmbeddingProtocol::Vllm {
|
||||
base_url,
|
||||
api_key,
|
||||
model,
|
||||
} => {
|
||||
let url = v1_embeddings_url(&base_url, "vLLM")?;
|
||||
let request = optional_bearer(
|
||||
self.client()?
|
||||
.post(url)
|
||||
.json(&json!({ "input": texts, "model": model })),
|
||||
api_key,
|
||||
);
|
||||
let body = self.send_json(request, "vLLM").await?;
|
||||
parse_indexed_embeddings(&body, "data", "vLLM")
|
||||
}
|
||||
RemoteEmbeddingProtocol::KoboldCpp { base_url, api_key } => {
|
||||
let url = root_path_url(&base_url, "/api/extra/embeddings", "KoboldCpp")?;
|
||||
let request = optional_bearer(
|
||||
self.client()?
|
||||
.post(url)
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.json(&json!({ "input": texts })),
|
||||
api_key,
|
||||
);
|
||||
let body = self.send_json(request, "KoboldCpp").await?;
|
||||
reported_model = Some(parse_reported_model(&body, "KoboldCpp")?);
|
||||
parse_kobold_embeddings(&body)
|
||||
}
|
||||
}?;
|
||||
validate_embedding_count(&embeddings, expected_count)?;
|
||||
Ok(RemoteEmbeddingBatch {
|
||||
embeddings,
|
||||
reported_model,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn joined_url(base_url: &str, path: &str) -> Result<Url, DomainError> {
|
||||
Url::parse(&format!(
|
||||
"{}/{}",
|
||||
base_url.trim().trim_end_matches('/'),
|
||||
path.trim_start_matches('/')
|
||||
))
|
||||
.map_err(invalid_url("embedding provider"))
|
||||
}
|
||||
|
||||
fn root_path_url(base_url: &str, path: &str, provider: &'static str) -> Result<Url, DomainError> {
|
||||
let mut url = Url::parse(base_url.trim()).map_err(invalid_url(provider))?;
|
||||
url.set_path(path);
|
||||
url.set_query(None);
|
||||
url.set_fragment(None);
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn v1_embeddings_url(base_url: &str, provider: &'static str) -> Result<Url, DomainError> {
|
||||
let base = base_url.trim().trim_end_matches('/');
|
||||
let base = base.strip_suffix("/v1").unwrap_or(base);
|
||||
Url::parse(&format!("{base}/v1/embeddings")).map_err(invalid_url(provider))
|
||||
}
|
||||
|
||||
fn optional_bearer(request: RequestBuilder, api_key: Option<String>) -> RequestBuilder {
|
||||
match api_key
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
Some(api_key) => request.header(AUTHORIZATION, format!("Bearer {api_key}")),
|
||||
None => request,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_indexed_embeddings(
|
||||
body: &Value,
|
||||
field: &str,
|
||||
provider: &str,
|
||||
) -> Result<Vec<Vec<f32>>, DomainError> {
|
||||
let entries = required_array(body, field, provider)?;
|
||||
let mut embeddings = entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(position, entry)| {
|
||||
let index = parse_embedding_index(entry, position, provider)?;
|
||||
let embedding = parse_embedding(
|
||||
entry
|
||||
.get("embedding")
|
||||
.ok_or_else(|| invalid_response(provider, "missing embedding"))?,
|
||||
provider,
|
||||
)?;
|
||||
Ok((index, embedding))
|
||||
})
|
||||
.collect::<Result<Vec<_>, DomainError>>()?;
|
||||
embeddings.sort_by_key(|(index, _)| *index);
|
||||
validate_embedding_indices(&embeddings, provider)?;
|
||||
Ok(embeddings
|
||||
.into_iter()
|
||||
.map(|(_, embedding)| embedding)
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn parse_kobold_embeddings(body: &Value) -> Result<Vec<Vec<f32>>, DomainError> {
|
||||
let entries = required_array(body, "data", "KoboldCpp")?;
|
||||
let mut normalized = Vec::with_capacity(entries.len());
|
||||
for (position, entry) in entries.iter().enumerate() {
|
||||
let entry = entry
|
||||
.as_array()
|
||||
.and_then(|entries| entries.first())
|
||||
.unwrap_or(entry);
|
||||
let index = parse_embedding_index(entry, position, "KoboldCpp")?;
|
||||
let embedding = parse_embedding(
|
||||
entry
|
||||
.get("embedding")
|
||||
.ok_or_else(|| invalid_response("KoboldCpp", "missing embedding"))?,
|
||||
"KoboldCpp",
|
||||
)?;
|
||||
normalized.push((index, embedding));
|
||||
}
|
||||
normalized.sort_by_key(|(index, _)| *index);
|
||||
validate_embedding_indices(&normalized, "KoboldCpp")?;
|
||||
Ok(normalized
|
||||
.into_iter()
|
||||
.map(|(_, embedding)| embedding)
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn parse_reported_model(body: &Value, provider: &str) -> Result<String, DomainError> {
|
||||
body.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|model| !model.is_empty())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| invalid_response(provider, "missing model"))
|
||||
}
|
||||
|
||||
fn parse_embedding_index(
|
||||
entry: &Value,
|
||||
position: usize,
|
||||
provider: &str,
|
||||
) -> Result<u64, DomainError> {
|
||||
match entry.get("index") {
|
||||
Some(index) => index.as_u64().ok_or_else(|| {
|
||||
invalid_response(provider, "embedding index is not an unsigned integer")
|
||||
}),
|
||||
None => Ok(position as u64),
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_embedding_indices(
|
||||
embeddings: &[(u64, Vec<f32>)],
|
||||
provider: &str,
|
||||
) -> Result<(), DomainError> {
|
||||
if embeddings
|
||||
.iter()
|
||||
.enumerate()
|
||||
.any(|(expected, (actual, _))| *actual != expected as u64)
|
||||
{
|
||||
return Err(invalid_response(
|
||||
provider,
|
||||
"embedding indices are missing or duplicated",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_embedding_count(
|
||||
embeddings: &[Vec<f32>],
|
||||
expected_count: usize,
|
||||
) -> Result<(), DomainError> {
|
||||
if embeddings.len() != expected_count {
|
||||
return Err(invalid_response(
|
||||
"Embedding provider",
|
||||
&format!(
|
||||
"returned {} vectors for {expected_count} inputs",
|
||||
embeddings.len()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_embedding_array(value: &Value, provider: &str) -> Result<Vec<Vec<f32>>, DomainError> {
|
||||
value
|
||||
.as_array()
|
||||
.ok_or_else(|| invalid_response(provider, "embeddings is not an array"))?
|
||||
.iter()
|
||||
.map(|embedding| parse_embedding(embedding, provider))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_embedding(value: &Value, provider: &str) -> Result<Vec<f32>, DomainError> {
|
||||
let values = value
|
||||
.as_array()
|
||||
.ok_or_else(|| invalid_response(provider, "embedding is not an array"))?;
|
||||
if values.is_empty() {
|
||||
return Err(invalid_response(provider, "embedding is empty"));
|
||||
}
|
||||
|
||||
values
|
||||
.iter()
|
||||
.map(|value| {
|
||||
let value = value
|
||||
.as_f64()
|
||||
.ok_or_else(|| invalid_response(provider, "embedding contains a non-number"))?;
|
||||
let value = value as f32;
|
||||
if !value.is_finite() {
|
||||
return Err(invalid_response(
|
||||
provider,
|
||||
"embedding contains a non-finite number",
|
||||
));
|
||||
}
|
||||
Ok(value)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn required_array<'a>(
|
||||
body: &'a Value,
|
||||
field: &str,
|
||||
provider: &str,
|
||||
) -> Result<&'a Vec<Value>, DomainError> {
|
||||
body.get(field)
|
||||
.and_then(Value::as_array)
|
||||
.ok_or_else(|| invalid_response(provider, &format!("missing {field} array")))
|
||||
}
|
||||
|
||||
fn service_account_project_id(credentials: &str) -> Result<String, DomainError> {
|
||||
serde_json::from_str::<Value>(credentials)
|
||||
.map_err(|error| {
|
||||
DomainError::InvalidData(format!(
|
||||
"Vertex AI service account JSON parse failed: {error}"
|
||||
))
|
||||
})?
|
||||
.get("project_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
DomainError::InvalidData(
|
||||
"Vertex AI service account JSON is missing project_id".to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn normalized_region(region: &str) -> String {
|
||||
let region = region.trim().to_ascii_lowercase();
|
||||
if region.is_empty() {
|
||||
"us-central1".to_string()
|
||||
} else {
|
||||
region
|
||||
}
|
||||
}
|
||||
|
||||
fn vertex_host(region: &str) -> String {
|
||||
match region {
|
||||
"global" => "https://aiplatform.googleapis.com".to_string(),
|
||||
"us" | "eu" => format!("https://aiplatform.{region}.rep.googleapis.com"),
|
||||
_ => format!("https://{region}-aiplatform.googleapis.com"),
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_url(provider: &'static str) -> impl FnOnce(url::ParseError) -> DomainError {
|
||||
move |error| DomainError::InvalidData(format!("Invalid {provider} embedding URL: {error}"))
|
||||
}
|
||||
|
||||
fn invalid_response(provider: &str, detail: &str) -> DomainError {
|
||||
DomainError::InternalError(format!(
|
||||
"{provider} embedding response is invalid: {detail}"
|
||||
))
|
||||
}
|
||||
|
||||
async fn map_error_response(provider: &str, response: reqwest::Response) -> DomainError {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let message = response_error_message(&body);
|
||||
match status {
|
||||
StatusCode::BAD_REQUEST => DomainError::InvalidData(message),
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
|
||||
DomainError::AuthenticationError(message)
|
||||
}
|
||||
StatusCode::TOO_MANY_REQUESTS => DomainError::rate_limited(format!(
|
||||
"{provider} embedding request was rate limited: {message}"
|
||||
)),
|
||||
status if status.is_server_error() => DomainError::transient(format!(
|
||||
"{provider} embedding request failed with status {}: {message}",
|
||||
status.as_u16()
|
||||
)),
|
||||
_ => DomainError::InternalError(format!(
|
||||
"{provider} embedding request failed with status {}: {message}",
|
||||
status.as_u16()
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn response_error_message(body: &str) -> String {
|
||||
let parsed = serde_json::from_str::<Value>(body).ok();
|
||||
parsed
|
||||
.as_ref()
|
||||
.and_then(|value| {
|
||||
value
|
||||
.pointer("/error/message")
|
||||
.or_else(|| value.get("message"))
|
||||
.or_else(|| value.get("detail"))
|
||||
})
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|message| !message.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
let body = body.trim();
|
||||
(!body.is_empty()).then(|| body.chars().take(1_000).collect())
|
||||
})
|
||||
.unwrap_or_else(|| "Upstream embedding request failed".to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn openai_embeddings_are_sorted_by_index() {
|
||||
let body = json!({
|
||||
"data": [
|
||||
{ "index": 1, "embedding": [0.0, 1.0] },
|
||||
{ "index": 0, "embedding": [1.0, 0.0] }
|
||||
]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
parse_indexed_embeddings(&body, "data", "test").unwrap(),
|
||||
vec![vec![1.0, 0.0], vec![0.0, 1.0]]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_explicit_embedding_index_is_rejected() {
|
||||
let body = json!({ "data": [{ "index": "0", "embedding": [1.0] }] });
|
||||
assert!(parse_indexed_embeddings(&body, "data", "test").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_embeddings_fail_instead_of_becoming_zero_vectors() {
|
||||
let body = json!({ "embeddings": [[0.0, "bad"]] });
|
||||
assert!(parse_embedding_array(&body["embeddings"], "test").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_v1_urls_keep_exactly_one_v1_segment() {
|
||||
assert_eq!(
|
||||
v1_embeddings_url("http://localhost:8080/v1", "test")
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"http://localhost:8080/v1/embeddings"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_count_must_match_the_requested_batch() {
|
||||
assert!(validate_embedding_count(&[vec![1.0]], 2).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reported_model_must_be_present_and_non_empty() {
|
||||
assert_eq!(
|
||||
parse_reported_model(&json!({ "model": " embed.gguf " }), "test").unwrap(),
|
||||
"embed.gguf"
|
||||
);
|
||||
assert!(parse_reported_model(&json!({}), "test").is_err());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
mod endpoint_url;
|
||||
mod file_replace;
|
||||
mod http_chat_completion_repository;
|
||||
mod http_embedding_repository;
|
||||
mod http_error;
|
||||
mod http_provider_metadata_repository;
|
||||
mod http_stable_diffusion_repository;
|
||||
|
|
@ -10,6 +11,7 @@ mod workers_ai_endpoint;
|
|||
mod workers_ai_models;
|
||||
|
||||
pub use http_chat_completion_repository::HttpChatCompletionRepository;
|
||||
pub use http_embedding_repository::HttpEmbeddingRepository;
|
||||
pub use http_provider_metadata_repository::HttpProviderMetadataRepository;
|
||||
pub use http_stable_diffusion_repository::HttpStableDiffusionRepository;
|
||||
pub use http_translate_repository::HttpTranslateRepository;
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ filetime = "0.2"
|
|||
local-ip-address = "0.6"
|
||||
rand = { workspace = true }
|
||||
reqwest = { workspace = true, features = ["json", "rustls", "stream"] }
|
||||
rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs", "std", "tls12"] }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tokio = { version = "1", default-features = false, features = ["fs", "io-util", "macros", "net", "rt-multi-thread", "sync", "time"] }
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ pub async fn spawn_lan_sync_server(
|
|||
inbound: Arc<dyn LanInboundRequestHandler>,
|
||||
errors: Arc<dyn LanServerErrorReporter>,
|
||||
) -> Result<LanSyncServerHandle, DomainError> {
|
||||
install_rustls_crypto_provider();
|
||||
let identity = store.load_or_create_identity().await?;
|
||||
let tls = SelfManagedTls::load_or_create(&store.state_dir()).map_err(sync_error_to_domain)?;
|
||||
let spki_sha256 = tls.spki_sha256().to_string();
|
||||
|
|
@ -112,6 +113,15 @@ pub async fn spawn_lan_sync_server(
|
|||
spawn_router(addr, Arc::new(tls), spki_sha256, app, errors).await
|
||||
}
|
||||
|
||||
fn install_rustls_crypto_provider() {
|
||||
static INSTALL: std::sync::Once = std::sync::Once::new();
|
||||
INSTALL.call_once(|| {
|
||||
// LAN TLS can run without the HTTP adapter, so this concrete TLS owner selects the same
|
||||
// provider explicitly when workspace dependencies compile more than one rustls backend.
|
||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||
});
|
||||
}
|
||||
|
||||
async fn spawn_router(
|
||||
addr: SocketAddr,
|
||||
tls: Arc<dyn TlsProvider>,
|
||||
|
|
|
|||
22
src-tauri/crates/tt-adapter-vector/Cargo.toml
Normal file
22
src-tauri/crates/tt-adapter-vector/Cargo.toml
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
[package]
|
||||
name = "tt-adapter-vector"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1"
|
||||
candle-core = "0.11.0"
|
||||
candle-nn = "0.11.0"
|
||||
fastembed = { version = "5.17.4", default-features = false, features = ["hf-hub-rustls-tls", "ort-download-binaries-rustls-tls", "qwen3"] }
|
||||
hf-hub = { version = "0.5.0", default-features = false }
|
||||
redb = "4.1.0"
|
||||
serde_json = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
tokio = { version = "1", default-features = false, features = ["macros", "rt", "sync"] }
|
||||
tokenizers = { version = "0.22.2", default-features = false, features = ["onig"] }
|
||||
tt-domain = { path = "../tt-domain" }
|
||||
tt-ports = { path = "../tt-ports" }
|
||||
|
||||
[dev-dependencies]
|
||||
uuid = { workspace = true }
|
||||
28
src-tauri/crates/tt-adapter-vector/build.rs
Normal file
28
src-tauri/crates/tt-adapter-vector/build.rs
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
let target = std::env::var("TARGET").unwrap_or_default();
|
||||
if !target.ends_with("apple-darwin") {
|
||||
return;
|
||||
}
|
||||
|
||||
// ORT resolves this through generic CC, which may point at the Android NDK
|
||||
// in a multi-target Tauri shell. xcrun is the authoritative macOS toolchain.
|
||||
let Ok(output) = Command::new("xcrun")
|
||||
.args(["clang", "--print-resource-dir"])
|
||||
.output()
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if !output.status.success() {
|
||||
return;
|
||||
}
|
||||
|
||||
let runtime_dir = PathBuf::from(String::from_utf8_lossy(&output.stdout).trim())
|
||||
.join("lib")
|
||||
.join("darwin");
|
||||
if runtime_dir.is_dir() {
|
||||
println!("cargo:rustc-link-search=native={}", runtime_dir.display());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,218 @@
|
|||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use candle_core::{DType, Device};
|
||||
use candle_nn::VarBuilder;
|
||||
use fastembed::{
|
||||
EmbeddingModel, Qwen3Config, Qwen3Model, Qwen3TextEmbedding, TextEmbedding, TextInitOptions,
|
||||
};
|
||||
use hf_hub::api::sync::ApiBuilder;
|
||||
use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
|
||||
|
||||
use tt_domain::errors::DomainError;
|
||||
use tt_ports::repositories::vector_repository::{
|
||||
LocalEmbeddingModel, LocalEmbeddingRepository, LocalEmbeddingRequest,
|
||||
};
|
||||
|
||||
const QWEN_MAX_LENGTH: usize = 8_192;
|
||||
const QWEN_RETRIEVAL_INSTRUCTION: &str =
|
||||
"Given a conversation context, retrieve relevant passages and memories";
|
||||
|
||||
pub struct FastEmbedLocalEmbeddingRepository {
|
||||
cache_dir: PathBuf,
|
||||
// ponytail: one model lock bounds memory and inference concurrency; add a worker pool only
|
||||
// if measured queue latency justifies loading or scheduling more runtime state.
|
||||
model: Arc<Mutex<Option<LoadedEmbeddingModel>>>,
|
||||
}
|
||||
|
||||
enum LoadedEmbeddingModel {
|
||||
Onnx {
|
||||
model: LocalEmbeddingModel,
|
||||
runtime: TextEmbedding,
|
||||
},
|
||||
Qwen3(Qwen3TextEmbedding),
|
||||
}
|
||||
|
||||
impl LoadedEmbeddingModel {
|
||||
fn load(model: LocalEmbeddingModel, cache_dir: &Path) -> Result<Self, fastembed::Error> {
|
||||
let (runtime_model, max_length) = match model {
|
||||
LocalEmbeddingModel::JinaV2BaseEn => (EmbeddingModel::JinaEmbeddingsV2BaseEN, 512),
|
||||
LocalEmbeddingModel::BgeM3 => (EmbeddingModel::BGEM3, 8_192),
|
||||
LocalEmbeddingModel::EmbeddingGemma300M => {
|
||||
(EmbeddingModel::EmbeddingGemma300MQ4, 2_048)
|
||||
}
|
||||
LocalEmbeddingModel::Qwen3Embedding06B => {
|
||||
return Ok(Self::Qwen3(load_qwen3(cache_dir)?));
|
||||
}
|
||||
};
|
||||
let options = TextInitOptions::new(runtime_model)
|
||||
.with_cache_dir(cache_dir.to_path_buf())
|
||||
.with_max_length(max_length)
|
||||
.with_intra_threads(2)
|
||||
.with_show_download_progress(false);
|
||||
Ok(Self::Onnx {
|
||||
model,
|
||||
runtime: TextEmbedding::try_new(options)?,
|
||||
})
|
||||
}
|
||||
|
||||
const fn model(&self) -> LocalEmbeddingModel {
|
||||
match self {
|
||||
Self::Onnx { model, .. } => *model,
|
||||
Self::Qwen3(_) => LocalEmbeddingModel::Qwen3Embedding06B,
|
||||
}
|
||||
}
|
||||
|
||||
fn embed(&mut self, texts: &[String]) -> Result<Vec<Vec<f32>>, fastembed::Error> {
|
||||
match self {
|
||||
Self::Onnx { runtime, .. } => runtime.embed(texts, Some(10)),
|
||||
Self::Qwen3(runtime) => runtime.embed(texts).map_err(fastembed::Error::new),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FastEmbedLocalEmbeddingRepository {
|
||||
pub fn new(cache_dir: PathBuf) -> Self {
|
||||
Self {
|
||||
cache_dir,
|
||||
model: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LocalEmbeddingRepository for FastEmbedLocalEmbeddingRepository {
|
||||
async fn embed(&self, request: LocalEmbeddingRequest) -> Result<Vec<Vec<f32>>, DomainError> {
|
||||
let cache_dir = self.cache_dir.clone();
|
||||
let model = self.model.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let texts = prepare_texts(request.model, request.texts, request.is_query);
|
||||
let mut loaded = model.lock().map_err(|_| {
|
||||
DomainError::InternalError("Local embedding model lock was poisoned".to_string())
|
||||
})?;
|
||||
|
||||
if loaded.as_ref().map(LoadedEmbeddingModel::model) != Some(request.model) {
|
||||
*loaded = None;
|
||||
*loaded = Some(
|
||||
LoadedEmbeddingModel::load(request.model, &cache_dir).map_err(|error| {
|
||||
DomainError::transient(format!(
|
||||
"Failed to initialize local embedding model {}: {error}",
|
||||
request.model.id()
|
||||
))
|
||||
})?,
|
||||
);
|
||||
}
|
||||
|
||||
loaded
|
||||
.as_mut()
|
||||
.expect("local embedding model was initialized")
|
||||
.embed(&texts)
|
||||
.map_err(|error| {
|
||||
DomainError::InternalError(format!(
|
||||
"Local embedding inference failed for {}: {error}",
|
||||
request.model.id()
|
||||
))
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|error| {
|
||||
DomainError::InternalError(format!("Local embedding task failed: {error}"))
|
||||
})?
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_texts(model: LocalEmbeddingModel, texts: Vec<String>, is_query: bool) -> Vec<String> {
|
||||
texts
|
||||
.into_iter()
|
||||
.map(|text| match (model, is_query) {
|
||||
(LocalEmbeddingModel::Qwen3Embedding06B, true) => {
|
||||
format!("Instruct: {QWEN_RETRIEVAL_INSTRUCTION}\nQuery:{text}")
|
||||
}
|
||||
(LocalEmbeddingModel::EmbeddingGemma300M, true) => {
|
||||
format!("task: search result | query: {text}")
|
||||
}
|
||||
(LocalEmbeddingModel::EmbeddingGemma300M, false) => {
|
||||
format!("title: none | text: {text}")
|
||||
}
|
||||
_ => text,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn load_qwen3(cache_dir: &Path) -> Result<Qwen3TextEmbedding, fastembed::Error> {
|
||||
// FastEmbed's convenience loader does not accept a cache path. Assemble its public Qwen
|
||||
// runtime here so every downloaded model remains inside the selected TauriTavern data root.
|
||||
let repo = ApiBuilder::from_env()
|
||||
.with_cache_dir(cache_dir.to_path_buf())
|
||||
.with_progress(false)
|
||||
.build()?
|
||||
.model(LocalEmbeddingModel::Qwen3Embedding06B.id().to_string());
|
||||
|
||||
let config: Qwen3Config = serde_json::from_slice(&std::fs::read(repo.get("config.json")?)?)?;
|
||||
let weights = [repo.get("model.safetensors")?];
|
||||
let device = Device::Cpu;
|
||||
// SAFETY: the immutable safetensors path comes from hf-hub's content-addressed cache and
|
||||
// remains owned by that cache for the lifetime of the mapped Qwen runtime.
|
||||
let variables = unsafe { VarBuilder::from_mmaped_safetensors(&weights, DType::F32, &device)? };
|
||||
let model = Qwen3Model::new(config, variables)?;
|
||||
let tokenizer_path = repo.get("tokenizer.json")?;
|
||||
let mut tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
|
||||
.map_err(|error| fastembed::Error::msg(error.to_string()))?;
|
||||
tokenizer.with_padding(Some(PaddingParams {
|
||||
strategy: PaddingStrategy::BatchLongest,
|
||||
direction: tokenizers::PaddingDirection::Left,
|
||||
..Default::default()
|
||||
}));
|
||||
tokenizer
|
||||
.with_truncation(Some(TruncationParams {
|
||||
max_length: QWEN_MAX_LENGTH,
|
||||
..Default::default()
|
||||
}))
|
||||
.map_err(fastembed::Error::msg)?;
|
||||
|
||||
Ok(Qwen3TextEmbedding::new(model, tokenizer))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn asymmetric_models_receive_query_and_document_prompts() {
|
||||
let qwen_query = prepare_texts(
|
||||
LocalEmbeddingModel::Qwen3Embedding06B,
|
||||
vec!["remember this".to_string()],
|
||||
true,
|
||||
);
|
||||
let qwen_document = prepare_texts(
|
||||
LocalEmbeddingModel::Qwen3Embedding06B,
|
||||
vec!["remember this".to_string()],
|
||||
false,
|
||||
);
|
||||
assert_eq!(
|
||||
qwen_query,
|
||||
[format!(
|
||||
"Instruct: {QWEN_RETRIEVAL_INSTRUCTION}\nQuery:remember this"
|
||||
)]
|
||||
);
|
||||
assert_eq!(qwen_document, ["remember this"]);
|
||||
|
||||
assert_eq!(
|
||||
prepare_texts(
|
||||
LocalEmbeddingModel::EmbeddingGemma300M,
|
||||
vec!["memory".to_string()],
|
||||
true,
|
||||
),
|
||||
["task: search result | query: memory"]
|
||||
);
|
||||
assert_eq!(
|
||||
prepare_texts(
|
||||
LocalEmbeddingModel::EmbeddingGemma300M,
|
||||
vec!["memory".to_string()],
|
||||
false,
|
||||
),
|
||||
["title: none | text: memory"]
|
||||
);
|
||||
}
|
||||
}
|
||||
5
src-tauri/crates/tt-adapter-vector/src/lib.rs
Normal file
5
src-tauri/crates/tt-adapter-vector/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
mod fastembed_local_embedding_repository;
|
||||
mod redb_vector_repository;
|
||||
|
||||
pub use fastembed_local_embedding_repository::FastEmbedLocalEmbeddingRepository;
|
||||
pub use redb_vector_repository::RedbVectorRepository;
|
||||
608
src-tauri/crates/tt-adapter-vector/src/redb_vector_repository.rs
Normal file
608
src-tauri/crates/tt-adapter-vector/src/redb_vector_repository.rs
Normal file
|
|
@ -0,0 +1,608 @@
|
|||
use std::collections::HashSet;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
use tt_domain::errors::DomainError;
|
||||
use tt_ports::repositories::vector_repository::{
|
||||
VectorMatch, VectorRecord, VectorRepository, VectorScope,
|
||||
};
|
||||
|
||||
const METADATA: TableDefinition<&str, &[u8]> = TableDefinition::new("metadata_v1");
|
||||
const VECTORS: TableDefinition<&str, &[u8]> = TableDefinition::new("vectors_v1");
|
||||
const SCOPES: TableDefinition<&str, u32> = TableDefinition::new("scopes_v1");
|
||||
|
||||
pub struct RedbVectorRepository {
|
||||
database_path: PathBuf,
|
||||
database: OnceCell<Arc<Database>>,
|
||||
}
|
||||
|
||||
impl RedbVectorRepository {
|
||||
pub fn new(database_path: PathBuf) -> Self {
|
||||
Self {
|
||||
database_path,
|
||||
database: OnceCell::new(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn database(&self) -> Result<Arc<Database>, DomainError> {
|
||||
let database_path = self.database_path.clone();
|
||||
self.database
|
||||
.get_or_try_init(|| async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
if let Some(parent) = database_path.parent() {
|
||||
std::fs::create_dir_all(parent).map_err(|error| {
|
||||
storage_error("create vector storage directory", error)
|
||||
})?;
|
||||
}
|
||||
|
||||
let database = Database::create(&database_path)
|
||||
.map_err(|error| storage_error("open vector database", error))?;
|
||||
let transaction = database
|
||||
.begin_write()
|
||||
.map_err(|error| storage_error("initialize vector database", error))?;
|
||||
{
|
||||
transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("initialize metadata table", error))?;
|
||||
transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("initialize vector table", error))?;
|
||||
transaction
|
||||
.open_table(SCOPES)
|
||||
.map_err(|error| storage_error("initialize scope table", error))?;
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| storage_error("commit vector database schema", error))?;
|
||||
Ok(Arc::new(database))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| {
|
||||
DomainError::InternalError(format!(
|
||||
"Vector database initialization task failed: {error}"
|
||||
))
|
||||
})?
|
||||
})
|
||||
.await
|
||||
.cloned()
|
||||
}
|
||||
|
||||
async fn run_blocking<T, F>(&self, operation: F) -> Result<T, DomainError>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(Arc<Database>) -> Result<T, DomainError> + Send + 'static,
|
||||
{
|
||||
let database = self.database().await?;
|
||||
tokio::task::spawn_blocking(move || operation(database))
|
||||
.await
|
||||
.map_err(|error| {
|
||||
DomainError::InternalError(format!("Vector storage task failed: {error}"))
|
||||
})?
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VectorRepository for RedbVectorRepository {
|
||||
async fn list_hashes(&self, scope: &VectorScope) -> Result<Vec<i64>, DomainError> {
|
||||
let prefix = scope_prefix(scope);
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_read()
|
||||
.map_err(|error| storage_error("read vector hashes", error))?;
|
||||
let table = transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?;
|
||||
|
||||
let mut hashes = Vec::new();
|
||||
let end = prefix_end(&prefix);
|
||||
for entry in table
|
||||
.range(prefix.as_str()..end.as_str())
|
||||
.map_err(|error| storage_error("scan vector metadata", error))?
|
||||
{
|
||||
let (_, value) =
|
||||
entry.map_err(|error| storage_error("read vector metadata", error))?;
|
||||
let metadata = serde_json::from_slice::<
|
||||
tt_ports::repositories::vector_repository::VectorMetadata,
|
||||
>(value.value())
|
||||
.map_err(|error| storage_error("decode vector metadata", error))?;
|
||||
hashes.push(metadata.hash);
|
||||
}
|
||||
Ok(hashes)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn upsert(
|
||||
&self,
|
||||
scope: &VectorScope,
|
||||
records: Vec<VectorRecord>,
|
||||
) -> Result<(), DomainError> {
|
||||
if records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let dimension = validate_records(&records)?;
|
||||
let prefix = scope_prefix(scope);
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_write()
|
||||
.map_err(|error| storage_error("begin vector upsert", error))?;
|
||||
{
|
||||
let mut scopes = transaction
|
||||
.open_table(SCOPES)
|
||||
.map_err(|error| storage_error("open vector scopes", error))?;
|
||||
if let Some(stored) = scopes
|
||||
.get(prefix.as_str())
|
||||
.map_err(|error| storage_error("read vector dimension", error))?
|
||||
{
|
||||
if stored.value() != dimension {
|
||||
return Err(DomainError::Conflict(format!(
|
||||
"Embedding dimension changed for the selected vector source: stored {}, received {dimension}",
|
||||
stored.value()
|
||||
)));
|
||||
}
|
||||
} else {
|
||||
scopes
|
||||
.insert(prefix.as_str(), dimension)
|
||||
.map_err(|error| storage_error("store vector dimension", error))?;
|
||||
}
|
||||
|
||||
let mut metadata_table = transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?;
|
||||
let mut vector_table = transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("open vector values", error))?;
|
||||
|
||||
for record in records {
|
||||
let key = format!("{}{}", prefix, item_id(&record));
|
||||
let metadata = serde_json::to_vec(&record.metadata)
|
||||
.map_err(|error| storage_error("encode vector metadata", error))?;
|
||||
let embedding = encode_embedding(&record.embedding);
|
||||
metadata_table
|
||||
.insert(key.as_str(), metadata.as_slice())
|
||||
.map_err(|error| storage_error("store vector metadata", error))?;
|
||||
vector_table
|
||||
.insert(key.as_str(), embedding.as_slice())
|
||||
.map_err(|error| storage_error("store vector value", error))?;
|
||||
}
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| storage_error("commit vector upsert", error))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_hashes(&self, scope: &VectorScope, hashes: &[i64]) -> Result<(), DomainError> {
|
||||
if hashes.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let prefix = scope_prefix(scope);
|
||||
let hashes = hashes.iter().copied().collect::<HashSet<_>>();
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_write()
|
||||
.map_err(|error| storage_error("begin vector delete", error))?;
|
||||
{
|
||||
let mut metadata_table = transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?;
|
||||
let end = prefix_end(&prefix);
|
||||
let keys = metadata_table
|
||||
.range(prefix.as_str()..end.as_str())
|
||||
.map_err(|error| storage_error("scan vector metadata", error))?
|
||||
.map(|entry| {
|
||||
let (key, value) =
|
||||
entry.map_err(|error| storage_error("read vector metadata", error))?;
|
||||
let metadata = serde_json::from_slice::<
|
||||
tt_ports::repositories::vector_repository::VectorMetadata,
|
||||
>(value.value())
|
||||
.map_err(|error| storage_error("decode vector metadata", error))?;
|
||||
Ok((key.value().to_string(), metadata.hash))
|
||||
})
|
||||
.collect::<Result<Vec<_>, DomainError>>()?;
|
||||
|
||||
let mut vector_table = transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("open vector values", error))?;
|
||||
for (key, hash) in keys {
|
||||
if hashes.contains(&hash) {
|
||||
metadata_table
|
||||
.remove(key.as_str())
|
||||
.map_err(|error| storage_error("delete vector metadata", error))?;
|
||||
vector_table
|
||||
.remove(key.as_str())
|
||||
.map_err(|error| storage_error("delete vector value", error))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| storage_error("commit vector delete", error))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn query(
|
||||
&self,
|
||||
scope: &VectorScope,
|
||||
embedding: Vec<f32>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<VectorMatch>, DomainError> {
|
||||
validate_embedding(&embedding)?;
|
||||
let prefix = scope_prefix(scope);
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_read()
|
||||
.map_err(|error| storage_error("begin vector query", error))?;
|
||||
let scopes = transaction
|
||||
.open_table(SCOPES)
|
||||
.map_err(|error| storage_error("open vector scopes", error))?;
|
||||
let Some(stored_dimension) = scopes
|
||||
.get(prefix.as_str())
|
||||
.map_err(|error| storage_error("read vector dimension", error))?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if stored_dimension.value() as usize != embedding.len() {
|
||||
return Err(DomainError::Conflict(format!(
|
||||
"Embedding dimension changed for the selected vector source: stored {}, received {}",
|
||||
stored_dimension.value(),
|
||||
embedding.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let vectors = transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("open vector values", error))?;
|
||||
let mut scored = Vec::new();
|
||||
let end = prefix_end(&prefix);
|
||||
for entry in vectors
|
||||
.range(prefix.as_str()..end.as_str())
|
||||
.map_err(|error| storage_error("scan vector values", error))?
|
||||
{
|
||||
let (key, value) =
|
||||
entry.map_err(|error| storage_error("read vector value", error))?;
|
||||
let stored = decode_embedding(value.value())?;
|
||||
if stored.len() != embedding.len() {
|
||||
return Err(DomainError::InternalError(format!(
|
||||
"Vector index contains an invalid dimension for key {}",
|
||||
key.value()
|
||||
)));
|
||||
}
|
||||
let score = stored
|
||||
.iter()
|
||||
.zip(&embedding)
|
||||
.map(|(left, right)| left * right)
|
||||
.sum::<f32>();
|
||||
scored.push((key.value().to_string(), score));
|
||||
}
|
||||
|
||||
// ponytail: exact scan/sort is the deliberate baseline; add ANN only after
|
||||
// measured collection-scale latency warrants its index lifecycle complexity.
|
||||
scored.sort_by(|left, right| right.1.total_cmp(&left.1));
|
||||
scored.truncate(limit);
|
||||
|
||||
let metadata = transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?;
|
||||
scored
|
||||
.into_iter()
|
||||
.map(|(key, score)| {
|
||||
let value = metadata
|
||||
.get(key.as_str())
|
||||
.map_err(|error| storage_error("read vector metadata", error))?
|
||||
.ok_or_else(|| {
|
||||
DomainError::InternalError(format!(
|
||||
"Vector index is missing metadata for key {key}"
|
||||
))
|
||||
})?;
|
||||
let metadata = serde_json::from_slice(value.value())
|
||||
.map_err(|error| storage_error("decode vector metadata", error))?;
|
||||
Ok(VectorMatch { metadata, score })
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn purge_collection(&self, collection_id: &str) -> Result<(), DomainError> {
|
||||
let prefix = collection_prefix(collection_id);
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_write()
|
||||
.map_err(|error| storage_error("begin vector collection purge", error))?;
|
||||
{
|
||||
let end = prefix_end(&prefix);
|
||||
let mut metadata = transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?;
|
||||
metadata
|
||||
.retain_in(prefix.as_str()..end.as_str(), |_, _| false)
|
||||
.map_err(|error| storage_error("purge vector metadata", error))?;
|
||||
|
||||
let mut vectors = transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("open vector values", error))?;
|
||||
vectors
|
||||
.retain_in(prefix.as_str()..end.as_str(), |_, _| false)
|
||||
.map_err(|error| storage_error("purge vector values", error))?;
|
||||
|
||||
let mut scopes = transaction
|
||||
.open_table(SCOPES)
|
||||
.map_err(|error| storage_error("open vector scopes", error))?;
|
||||
scopes
|
||||
.retain_in(prefix.as_str()..end.as_str(), |_, _| false)
|
||||
.map_err(|error| storage_error("purge vector scopes", error))?;
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| storage_error("commit vector collection purge", error))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn purge_all(&self) -> Result<(), DomainError> {
|
||||
self.run_blocking(move |database| {
|
||||
let transaction = database
|
||||
.begin_write()
|
||||
.map_err(|error| storage_error("begin vector purge", error))?;
|
||||
{
|
||||
transaction
|
||||
.open_table(METADATA)
|
||||
.map_err(|error| storage_error("open vector metadata", error))?
|
||||
.retain(|_, _| false)
|
||||
.map_err(|error| storage_error("purge vector metadata", error))?;
|
||||
transaction
|
||||
.open_table(VECTORS)
|
||||
.map_err(|error| storage_error("open vector values", error))?
|
||||
.retain(|_, _| false)
|
||||
.map_err(|error| storage_error("purge vector values", error))?;
|
||||
transaction
|
||||
.open_table(SCOPES)
|
||||
.map_err(|error| storage_error("open vector scopes", error))?
|
||||
.retain(|_, _| false)
|
||||
.map_err(|error| storage_error("purge vector scopes", error))?;
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| storage_error("commit vector purge", error))
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_records(records: &[VectorRecord]) -> Result<u32, DomainError> {
|
||||
let dimension = records[0].embedding.len();
|
||||
if dimension > u32::MAX as usize {
|
||||
return Err(DomainError::InvalidData(
|
||||
"Embedding dimension is too large".to_string(),
|
||||
));
|
||||
}
|
||||
for record in records {
|
||||
validate_embedding(&record.embedding)?;
|
||||
if record.embedding.len() != dimension {
|
||||
return Err(DomainError::InvalidData(
|
||||
"Embedding batch contains mixed dimensions".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(dimension as u32)
|
||||
}
|
||||
|
||||
fn validate_embedding(embedding: &[f32]) -> Result<(), DomainError> {
|
||||
if embedding.is_empty() {
|
||||
return Err(DomainError::InvalidData(
|
||||
"Embedding must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if embedding.iter().any(|value| !value.is_finite()) {
|
||||
return Err(DomainError::InvalidData(
|
||||
"Embedding contains a non-finite value".to_string(),
|
||||
));
|
||||
}
|
||||
if embedding
|
||||
.iter()
|
||||
.map(|value| f64::from(*value).powi(2))
|
||||
.sum::<f64>()
|
||||
<= f64::EPSILON
|
||||
{
|
||||
return Err(DomainError::InvalidData(
|
||||
"Embedding must not be a zero vector".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn collection_prefix(collection_id: &str) -> String {
|
||||
format!("c/{}/", sha256_hex(collection_id.as_bytes()))
|
||||
}
|
||||
|
||||
fn scope_prefix(scope: &VectorScope) -> String {
|
||||
let mut identity = Sha256::new();
|
||||
identity.update((scope.source.len() as u64).to_le_bytes());
|
||||
identity.update(scope.source.as_bytes());
|
||||
identity.update((scope.profile.len() as u64).to_le_bytes());
|
||||
identity.update(scope.profile.as_bytes());
|
||||
format!(
|
||||
"{}s/{}/",
|
||||
collection_prefix(&scope.collection_id),
|
||||
hex_digest(identity.finalize())
|
||||
)
|
||||
}
|
||||
|
||||
fn item_id(record: &VectorRecord) -> String {
|
||||
let mut identity = Sha256::new();
|
||||
identity.update(record.metadata.hash.to_le_bytes());
|
||||
identity.update(record.metadata.index.to_le_bytes());
|
||||
identity.update(record.metadata.text.as_bytes());
|
||||
hex_digest(identity.finalize())
|
||||
}
|
||||
|
||||
fn prefix_end(prefix: &str) -> String {
|
||||
format!("{prefix}~")
|
||||
}
|
||||
|
||||
fn sha256_hex(value: &[u8]) -> String {
|
||||
hex_digest(Sha256::digest(value))
|
||||
}
|
||||
|
||||
fn hex_digest(digest: impl AsRef<[u8]>) -> String {
|
||||
const HEX: &[u8; 16] = b"0123456789abcdef";
|
||||
let digest = digest.as_ref();
|
||||
let mut output = String::with_capacity(digest.len() * 2);
|
||||
for byte in digest {
|
||||
output.push(HEX[(byte >> 4) as usize] as char);
|
||||
output.push(HEX[(byte & 0x0f) as usize] as char);
|
||||
}
|
||||
output
|
||||
}
|
||||
|
||||
fn encode_embedding(embedding: &[f32]) -> Vec<u8> {
|
||||
embedding
|
||||
.iter()
|
||||
.flat_map(|value| value.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn decode_embedding(bytes: &[u8]) -> Result<Vec<f32>, DomainError> {
|
||||
let chunks = bytes.chunks_exact(4);
|
||||
if !chunks.remainder().is_empty() {
|
||||
return Err(DomainError::InternalError(
|
||||
"Vector index contains malformed embedding bytes".to_string(),
|
||||
));
|
||||
}
|
||||
let embedding = chunks
|
||||
.map(|chunk| f32::from_le_bytes(chunk.try_into().expect("four-byte chunk")))
|
||||
.collect::<Vec<_>>();
|
||||
validate_embedding(&embedding)?;
|
||||
Ok(embedding)
|
||||
}
|
||||
|
||||
fn storage_error(action: &str, error: impl std::fmt::Display) -> DomainError {
|
||||
DomainError::InternalError(format!("Failed to {action}: {error}"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tt_ports::repositories::vector_repository::VectorMetadata;
|
||||
use uuid::Uuid;
|
||||
|
||||
struct TempDatabase(PathBuf);
|
||||
|
||||
impl TempDatabase {
|
||||
fn new() -> Self {
|
||||
let root = std::env::temp_dir().join(format!("tauritavern-vector-{}", Uuid::new_v4()));
|
||||
Self(root.join("index.redb"))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TempDatabase {
|
||||
fn drop(&mut self) {
|
||||
if let Some(root) = self.0.parent() {
|
||||
let _ = std::fs::remove_dir_all(root);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn record(hash: i64, text: &str, embedding: Vec<f32>) -> VectorRecord {
|
||||
VectorRecord {
|
||||
metadata: VectorMetadata {
|
||||
hash,
|
||||
text: text.to_string(),
|
||||
index: 0,
|
||||
},
|
||||
embedding,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exact_query_is_idempotent_and_purge_is_collection_scoped() {
|
||||
let temp = TempDatabase::new();
|
||||
let repository = RedbVectorRepository::new(temp.0.clone());
|
||||
let scope = VectorScope {
|
||||
collection_id: "chat-a".to_string(),
|
||||
source: "test".to_string(),
|
||||
profile: "model-a".to_string(),
|
||||
};
|
||||
let other_scope = VectorScope {
|
||||
collection_id: "chat-b".to_string(),
|
||||
..scope.clone()
|
||||
};
|
||||
|
||||
repository
|
||||
.upsert(
|
||||
&scope,
|
||||
vec![
|
||||
record(1, "east", vec![1.0, 0.0]),
|
||||
record(2, "north", vec![0.0, 1.0]),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
repository
|
||||
.upsert(&scope, vec![record(1, "east", vec![1.0, 0.0])])
|
||||
.await
|
||||
.unwrap();
|
||||
repository
|
||||
.upsert(&other_scope, vec![record(3, "west", vec![-1.0, 0.0])])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(repository.list_hashes(&scope).await.unwrap().len(), 2);
|
||||
let matches = repository.query(&scope, vec![1.0, 0.0], 2).await.unwrap();
|
||||
assert_eq!(matches[0].metadata.hash, 1);
|
||||
assert_eq!(matches[0].score, 1.0);
|
||||
|
||||
assert!(matches!(
|
||||
repository
|
||||
.upsert(&scope, vec![record(4, "wrong", vec![1.0, 0.0, 0.0])])
|
||||
.await,
|
||||
Err(DomainError::Conflict(_))
|
||||
));
|
||||
assert_eq!(repository.list_hashes(&scope).await.unwrap().len(), 2);
|
||||
|
||||
repository.delete_hashes(&scope, &[1]).await.unwrap();
|
||||
assert_eq!(repository.list_hashes(&scope).await.unwrap(), vec![2]);
|
||||
|
||||
repository.purge_collection("chat-a").await.unwrap();
|
||||
assert!(repository.list_hashes(&scope).await.unwrap().is_empty());
|
||||
assert_eq!(repository.list_hashes(&other_scope).await.unwrap(), vec![3]);
|
||||
|
||||
drop(repository);
|
||||
let reopened = RedbVectorRepository::new(temp.0.clone());
|
||||
assert_eq!(reopened.list_hashes(&other_scope).await.unwrap(), vec![3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_writers_commit_complete_transactions() {
|
||||
let temp = TempDatabase::new();
|
||||
let repository = Arc::new(RedbVectorRepository::new(temp.0.clone()));
|
||||
let first = VectorScope {
|
||||
collection_id: "first".to_string(),
|
||||
source: "test".to_string(),
|
||||
profile: "model".to_string(),
|
||||
};
|
||||
let second = VectorScope {
|
||||
collection_id: "second".to_string(),
|
||||
..first.clone()
|
||||
};
|
||||
|
||||
tokio::try_join!(
|
||||
repository.upsert(&first, vec![record(1, "one", vec![1.0, 0.0])]),
|
||||
repository.upsert(&second, vec![record(2, "two", vec![0.0, 1.0])]),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(repository.list_hashes(&first).await.unwrap(), vec![1]);
|
||||
assert_eq!(repository.list_hashes(&second).await.unwrap(), vec![2]);
|
||||
}
|
||||
}
|
||||
|
|
@ -24,4 +24,5 @@ pub mod tokenization_dto;
|
|||
pub mod tts_dto;
|
||||
pub mod user_directory_dto;
|
||||
pub mod user_dto;
|
||||
pub mod vector_dto;
|
||||
pub mod world_info_dto;
|
||||
|
|
|
|||
67
src-tauri/crates/tt-application/src/dto/vector_dto.rs
Normal file
67
src-tauri/crates/tt-application/src/dto/vector_dto.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct VectorRouteRequestDto {
|
||||
#[serde(default)]
|
||||
pub collection_id: String,
|
||||
pub collection_ids: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub source: String,
|
||||
pub items: Option<Vec<VectorItemDto>>,
|
||||
pub hashes: Option<Vec<i64>>,
|
||||
#[serde(default)]
|
||||
pub search_text: String,
|
||||
pub top_k: Option<i64>,
|
||||
pub threshold: Option<f32>,
|
||||
#[serde(default)]
|
||||
pub model: String,
|
||||
#[serde(default)]
|
||||
pub api_url: String,
|
||||
#[serde(default)]
|
||||
pub keep: bool,
|
||||
#[serde(default)]
|
||||
pub is_query: bool,
|
||||
#[serde(default)]
|
||||
pub extras_url: String,
|
||||
#[serde(default)]
|
||||
pub extras_key: String,
|
||||
#[serde(default)]
|
||||
pub embeddings: HashMap<String, Vec<f32>>,
|
||||
pub texts: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub siliconflow_endpoint: String,
|
||||
#[serde(default)]
|
||||
pub workers_ai_account_id: String,
|
||||
#[serde(default)]
|
||||
pub vertexai_auth_mode: String,
|
||||
#[serde(default)]
|
||||
pub vertexai_region: String,
|
||||
#[serde(default)]
|
||||
pub vertexai_express_project_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct VectorItemDto {
|
||||
pub hash: i64,
|
||||
pub text: String,
|
||||
#[serde(default)]
|
||||
pub index: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum VectorRouteResponseKindDto {
|
||||
Json,
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct VectorRouteResponseDto {
|
||||
pub status: u16,
|
||||
pub kind: VectorRouteResponseKindDto,
|
||||
pub body: Value,
|
||||
}
|
||||
|
|
@ -62,4 +62,5 @@ pub mod update_service;
|
|||
pub mod user_directory_service;
|
||||
pub mod user_media_service;
|
||||
pub mod user_service;
|
||||
pub mod vector_service;
|
||||
pub mod world_info_service;
|
||||
|
|
|
|||
1023
src-tauri/crates/tt-application/src/services/vector_service.rs
Normal file
1023
src-tauri/crates/tt-application/src/services/vector_service.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -37,5 +37,6 @@ pub mod tts_repository;
|
|||
pub mod update_repository;
|
||||
pub mod user_directory_repository;
|
||||
pub mod user_repository;
|
||||
pub mod vector_repository;
|
||||
pub mod workspace_repository;
|
||||
pub mod world_info_repository;
|
||||
|
|
|
|||
210
src-tauri/crates/tt-ports/src/repositories/vector_repository.rs
Normal file
210
src-tauri/crates/tt-ports/src/repositories/vector_repository.rs
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use tt_domain::errors::DomainError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub enum LocalEmbeddingModel {
|
||||
#[default]
|
||||
JinaV2BaseEn,
|
||||
BgeM3,
|
||||
Qwen3Embedding06B,
|
||||
EmbeddingGemma300M,
|
||||
}
|
||||
|
||||
impl LocalEmbeddingModel {
|
||||
pub const SUPPORTED: [Self; 4] = [
|
||||
Self::JinaV2BaseEn,
|
||||
Self::BgeM3,
|
||||
Self::Qwen3Embedding06B,
|
||||
Self::EmbeddingGemma300M,
|
||||
];
|
||||
|
||||
pub const fn id(self) -> &'static str {
|
||||
match self {
|
||||
Self::JinaV2BaseEn => "jinaai/jina-embeddings-v2-base-en",
|
||||
Self::BgeM3 => "BAAI/bge-m3",
|
||||
Self::Qwen3Embedding06B => "Qwen/Qwen3-Embedding-0.6B",
|
||||
Self::EmbeddingGemma300M => "google/embeddinggemma-300m",
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn profile(self) -> &'static str {
|
||||
match self {
|
||||
// Preserve the original profile so existing Jina indexes remain reusable.
|
||||
Self::JinaV2BaseEn => "jinaai/jina-embeddings-v2-base-en",
|
||||
Self::BgeM3 => "BAAI/bge-m3|dense-v1|max=8192",
|
||||
Self::Qwen3Embedding06B => "Qwen/Qwen3-Embedding-0.6B|f32-v1|max=8192|query-prompt-v1",
|
||||
Self::EmbeddingGemma300M => {
|
||||
"google/embeddinggemma-300m|q4-v1|max=2048|retrieval-prompt-v1"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_id(value: &str) -> Option<Self> {
|
||||
Self::SUPPORTED
|
||||
.into_iter()
|
||||
.find(|model| model.id() == value)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct VectorScope {
|
||||
pub collection_id: String,
|
||||
pub source: String,
|
||||
pub profile: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct VectorMetadata {
|
||||
pub hash: i64,
|
||||
pub text: String,
|
||||
pub index: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct VectorRecord {
|
||||
pub metadata: VectorMetadata,
|
||||
pub embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct VectorMatch {
|
||||
pub metadata: VectorMetadata,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VectorRepository: Send + Sync {
|
||||
async fn list_hashes(&self, scope: &VectorScope) -> Result<Vec<i64>, DomainError>;
|
||||
|
||||
async fn upsert(
|
||||
&self,
|
||||
scope: &VectorScope,
|
||||
records: Vec<VectorRecord>,
|
||||
) -> Result<(), DomainError>;
|
||||
|
||||
async fn delete_hashes(&self, scope: &VectorScope, hashes: &[i64]) -> Result<(), DomainError>;
|
||||
|
||||
async fn query(
|
||||
&self,
|
||||
scope: &VectorScope,
|
||||
embedding: Vec<f32>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<VectorMatch>, DomainError>;
|
||||
|
||||
async fn purge_collection(&self, collection_id: &str) -> Result<(), DomainError>;
|
||||
|
||||
async fn purge_all(&self) -> Result<(), DomainError>;
|
||||
}
|
||||
|
||||
pub struct RemoteEmbeddingRequest {
|
||||
pub protocol: RemoteEmbeddingProtocol,
|
||||
pub texts: Vec<String>,
|
||||
pub is_query: bool,
|
||||
}
|
||||
|
||||
pub struct RemoteEmbeddingBatch {
|
||||
pub embeddings: Vec<Vec<f32>>,
|
||||
pub reported_model: Option<String>,
|
||||
}
|
||||
|
||||
pub struct LocalEmbeddingRequest {
|
||||
pub model: LocalEmbeddingModel,
|
||||
pub texts: Vec<String>,
|
||||
pub is_query: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum RemoteEmbeddingProtocol {
|
||||
OpenAi {
|
||||
provider: String,
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
omit_model: bool,
|
||||
headers: HashMap<String, String>,
|
||||
},
|
||||
Cohere {
|
||||
api_key: String,
|
||||
model: String,
|
||||
},
|
||||
Nomic {
|
||||
api_key: String,
|
||||
},
|
||||
Extras {
|
||||
base_url: String,
|
||||
api_key: Option<String>,
|
||||
},
|
||||
GoogleAiStudio {
|
||||
api_key: String,
|
||||
model: String,
|
||||
},
|
||||
VertexAi {
|
||||
model: String,
|
||||
region: String,
|
||||
auth: VertexEmbeddingAuth,
|
||||
},
|
||||
Ollama {
|
||||
base_url: String,
|
||||
model: String,
|
||||
keep: bool,
|
||||
},
|
||||
LlamaCpp {
|
||||
base_url: String,
|
||||
api_key: Option<String>,
|
||||
},
|
||||
Vllm {
|
||||
base_url: String,
|
||||
api_key: Option<String>,
|
||||
model: String,
|
||||
},
|
||||
KoboldCpp {
|
||||
base_url: String,
|
||||
api_key: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum VertexEmbeddingAuth {
|
||||
Express {
|
||||
api_key: String,
|
||||
project_id: Option<String>,
|
||||
},
|
||||
ServiceAccount {
|
||||
json: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RemoteEmbeddingRepository: Send + Sync {
|
||||
async fn embed(
|
||||
&self,
|
||||
request: RemoteEmbeddingRequest,
|
||||
) -> Result<RemoteEmbeddingBatch, DomainError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait LocalEmbeddingRepository: Send + Sync {
|
||||
async fn embed(&self, request: LocalEmbeddingRequest) -> Result<Vec<Vec<f32>>, DomainError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn local_embedding_ids_round_trip_and_profiles_are_distinct() {
|
||||
let profiles = LocalEmbeddingModel::SUPPORTED.map(|model| {
|
||||
assert_eq!(LocalEmbeddingModel::from_id(model.id()), Some(model));
|
||||
model.profile()
|
||||
});
|
||||
|
||||
for (index, profile) in profiles.iter().enumerate() {
|
||||
assert!(!profiles[..index].contains(profile));
|
||||
}
|
||||
assert!(LocalEmbeddingModel::from_id("unknown").is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -57,6 +57,7 @@ const getBatchSize = () => ['transformers', 'ollama'].includes(settings.source)
|
|||
const settings = {
|
||||
// For both
|
||||
source: 'transformers',
|
||||
transformers_model: 'jinaai/jina-embeddings-v2-base-en',
|
||||
alt_endpoint_url: '',
|
||||
use_alt_endpoint: false,
|
||||
include_wi: false,
|
||||
|
|
@ -131,11 +132,10 @@ const cachedSummaries = new Map();
|
|||
* @type {Set<number>}
|
||||
*/
|
||||
const skippedHashes = new Set();
|
||||
/**
|
||||
* Error causes treated as fatal; abort Vectorize All rather than skip.
|
||||
* @type {Set<string>}
|
||||
*/
|
||||
const FATAL_CAUSES = new Set(['account_id_missing', 'api_key_missing', 'api_url_missing', 'api_model_missing', 'extras_module_missing', 'webllm_not_supported', 'summary_endpoint_invalid', 'vector_endpoint_unavailable']);
|
||||
/** @type {Set<string>} */
|
||||
const VECTOR_SOURCE_ABORT_CAUSES = new Set(['account_id_missing', 'api_key_missing', 'api_url_missing', 'api_model_missing', 'embedding_auth_failed', 'embedding_rate_limited', 'embedding_unavailable', 'extras_module_missing', 'webllm_not_supported', 'vector_endpoint_unavailable']);
|
||||
/** @type {Set<string>} */
|
||||
const VECTOR_JOB_ABORT_CAUSES = new Set([...VECTOR_SOURCE_ABORT_CAUSES, 'summary_endpoint_invalid', 'vector_index_conflict']);
|
||||
const vectorApiRequiresUrl = ['llamacpp', 'vllm', 'ollama', 'koboldcpp'];
|
||||
|
||||
/**
|
||||
|
|
@ -494,7 +494,7 @@ async function synchronizeChat(batchSize = 5) {
|
|||
try {
|
||||
await insertVectorItems(chatId, chunkedBatch);
|
||||
} catch (insertError) {
|
||||
if (FATAL_CAUSES.has(insertError?.cause)) {
|
||||
if (VECTOR_JOB_ABORT_CAUSES.has(insertError?.cause)) {
|
||||
throw insertError;
|
||||
}
|
||||
console.warn('Vectors: insert failed for batch; marking for skip', insertError);
|
||||
|
|
@ -511,17 +511,25 @@ async function synchronizeChat(batchSize = 5) {
|
|||
} catch (error) {
|
||||
/**
|
||||
* Gets the error message for a given cause
|
||||
* @param {string} cause Error cause key
|
||||
* @param {Error} error Vector request error
|
||||
* @returns {string} Error message
|
||||
*/
|
||||
function getErrorMessage(cause) {
|
||||
switch (cause) {
|
||||
function getErrorMessage(error) {
|
||||
switch (error?.cause) {
|
||||
case 'api_key_missing':
|
||||
return 'API key missing. Save it in the "API Connections" panel.';
|
||||
case 'api_url_missing':
|
||||
return 'API URL missing. Save it in the "API Connections" panel.';
|
||||
case 'api_model_missing':
|
||||
return 'Vectorization Source Model is required, but not set.';
|
||||
case 'embedding_auth_failed':
|
||||
return 'Embedding provider authentication failed. Check the active API key.';
|
||||
case 'embedding_rate_limited':
|
||||
return 'Embedding provider rate limit reached. Retry later.';
|
||||
case 'embedding_unavailable':
|
||||
return 'Embedding service is temporarily unavailable. Retry later.';
|
||||
case 'vector_index_conflict':
|
||||
return 'The vector index is incompatible with the current embedding model. Purge it and retry.';
|
||||
case 'extras_module_missing':
|
||||
return 'Extras API must provide an "embeddings" module.';
|
||||
case 'webllm_not_supported':
|
||||
|
|
@ -533,15 +541,15 @@ async function synchronizeChat(batchSize = 5) {
|
|||
case 'summary_failed':
|
||||
return 'Summarization failed after the configured number of retries.';
|
||||
case 'vector_endpoint_unavailable':
|
||||
return 'Vector Storage backend is not implemented in the native TauriTavern backend.';
|
||||
return 'Vector Storage is unavailable in this TauriTavern build.';
|
||||
default:
|
||||
return 'Check server console for more details';
|
||||
return error?.message || 'Check server console for more details';
|
||||
}
|
||||
}
|
||||
|
||||
console.error('Vectors: Failed to synchronize chat', error);
|
||||
|
||||
const message = getErrorMessage(error.cause);
|
||||
const message = getErrorMessage(error);
|
||||
toastr.error(message, 'Vectorization failed', { preventDuplicates: true });
|
||||
return null;
|
||||
} finally {
|
||||
|
|
@ -639,6 +647,9 @@ async function processFiles(chat) {
|
|||
}
|
||||
} catch (error) {
|
||||
console.error('Vectors: Failed to retrieve files', error);
|
||||
if (VECTOR_SOURCE_ABORT_CAUSES.has(error?.cause)) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -702,6 +713,9 @@ async function injectDataBankChunks(queryText, collectionIds) {
|
|||
setExtensionPrompt(EXTENSION_PROMPT_TAG_DB, insertedText, settings.file_position_db, settings.file_depth_db, settings.include_wi, settings.file_depth_role_db);
|
||||
} catch (error) {
|
||||
console.error('Vectors: Failed to insert Data Bank chunks', error);
|
||||
if (VECTOR_SOURCE_ABORT_CAUSES.has(error?.cause)) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -740,7 +754,6 @@ async function vectorizeFile(fileText, fileName, collectionId, chunkSize, overla
|
|||
fileText = translatedText;
|
||||
}
|
||||
|
||||
const batchSize = getBatchSize();
|
||||
const toastBody = $('<span>').text('This may take a while. Please wait...');
|
||||
toast = toastr.info(toastBody, `Ingesting file ${escapeHtml(fileName)}`, { closeButton: false, escapeHtml: false, timeOut: 0, extendedTimeOut: 0 });
|
||||
const overlapSize = Math.round(chunkSize * overlapPercent / 100);
|
||||
|
|
@ -755,19 +768,19 @@ async function vectorizeFile(fileText, fileName, collectionId, chunkSize, overla
|
|||
|
||||
const items = chunks.map((chunk, index) => ({ hash: getStringHash(chunk), text: chunk, index: index }));
|
||||
|
||||
for (let i = 0; i < items.length; i += batchSize) {
|
||||
toastBody.text(`${i}/${items.length} (${Math.round((i / items.length) * 100)}%) chunks processed`);
|
||||
const chunkedBatch = items.slice(i, i + batchSize);
|
||||
await insertVectorItems(collectionId, chunkedBatch);
|
||||
}
|
||||
toastBody.text(`Embedding ${items.length} chunks...`);
|
||||
await insertVectorItems(collectionId, items);
|
||||
|
||||
toastr.clear(toast);
|
||||
console.log(`Vectors: Inserted ${chunks.length} vector items for file ${fileName} into ${collectionId}`);
|
||||
return true;
|
||||
} catch (error) {
|
||||
toastr.clear(toast);
|
||||
toastr.error(String(error), 'Failed to vectorize file', { preventDuplicates: true });
|
||||
console.error('Vectors: Failed to vectorize file', error);
|
||||
if (VECTOR_SOURCE_ABORT_CAUSES.has(error?.cause)) {
|
||||
throw error;
|
||||
}
|
||||
toastr.error(String(error), 'Failed to vectorize file', { preventDuplicates: true });
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
|
@ -795,7 +808,15 @@ async function rearrangeChat(chat, _contextSize, _abort, type) {
|
|||
}
|
||||
|
||||
if (settings.enabled_world_info) {
|
||||
await activateWorldInfo(chat);
|
||||
try {
|
||||
await activateWorldInfo(chat);
|
||||
} catch (error) {
|
||||
if (VECTOR_SOURCE_ABORT_CAUSES.has(error?.cause)) {
|
||||
throw error;
|
||||
}
|
||||
toastr.warning('World Info vector retrieval skipped. Chat retrieval will continue.', 'Vector Storage', { preventDuplicates: true });
|
||||
console.error('Vectors: Failed to activate World Info', error);
|
||||
}
|
||||
}
|
||||
|
||||
if (!settings.enabled_chats) {
|
||||
|
|
@ -859,7 +880,7 @@ async function rearrangeChat(chat, _contextSize, _abort, type) {
|
|||
const insertedText = getPromptText(queriedMessages);
|
||||
setExtensionPrompt(EXTENSION_PROMPT_TAG, insertedText, settings.position, settings.depth, settings.include_wi);
|
||||
} catch (error) {
|
||||
toastr.error('Generation interceptor aborted. Check browser console for more details.', 'Vector Storage');
|
||||
toastr.warning('Vector retrieval skipped. Generation will continue.', 'Vector Storage', { preventDuplicates: true });
|
||||
console.error('Vectors: Failed to rearrange chat', error);
|
||||
}
|
||||
}
|
||||
|
|
@ -937,6 +958,9 @@ async function getQueryText(chat, initiator) {
|
|||
function getVectorsRequestBody(args = {}) {
|
||||
const body = Object.assign({}, args);
|
||||
switch (settings.source) {
|
||||
case 'transformers':
|
||||
body.model = settings.transformers_model;
|
||||
break;
|
||||
case 'extras':
|
||||
body.extrasUrl = extension_settings.apiUrl;
|
||||
body.extrasKey = extension_settings.apiKey;
|
||||
|
|
@ -968,6 +992,9 @@ function getVectorsRequestBody(args = {}) {
|
|||
body.apiUrl = settings.use_alt_endpoint ? settings.alt_endpoint_url : textgenerationwebui_settings.server_urls[textgen_types.VLLM];
|
||||
body.model = extension_settings.vectors.vllm_model;
|
||||
break;
|
||||
case 'koboldcpp':
|
||||
body.apiUrl = settings.use_alt_endpoint ? settings.alt_endpoint_url : textgenerationwebui_settings.server_urls[textgen_types.KOBOLDCPP];
|
||||
break;
|
||||
case 'webllm':
|
||||
body.model = extension_settings.vectors.webllm_model;
|
||||
break;
|
||||
|
|
@ -1005,16 +1032,17 @@ function getVectorsRequestBody(args = {}) {
|
|||
/**
|
||||
* Gets additional arguments for vector requests.
|
||||
* @param {string[]} items Items to embed
|
||||
* @param {boolean} [isQuery=false] Whether embeddings use query semantics
|
||||
* @returns {Promise<object>} Additional arguments
|
||||
*/
|
||||
async function getAdditionalArgs(items) {
|
||||
async function getAdditionalArgs(items, isQuery = false) {
|
||||
const args = {};
|
||||
switch (settings.source) {
|
||||
case 'webllm':
|
||||
args.embeddings = await createWebLlmEmbeddings(items);
|
||||
break;
|
||||
case 'koboldcpp': {
|
||||
const { embeddings, model } = await createKoboldCppEmbeddings(items);
|
||||
const { embeddings, model } = await createKoboldCppEmbeddings(items, isQuery);
|
||||
args.embeddings = embeddings;
|
||||
args.model = model;
|
||||
break;
|
||||
|
|
@ -1029,12 +1057,14 @@ async function getAdditionalArgs(items) {
|
|||
* @param {string} action Action being performed
|
||||
* @param {string} [collectionId] Optional collection ID
|
||||
*/
|
||||
function throwVectorResponseError(response, action, collectionId = '') {
|
||||
async function throwVectorResponseError(response, action, collectionId = '') {
|
||||
const status = response?.status;
|
||||
const statusText = status ? ` (HTTP ${status})` : '';
|
||||
const collectionText = collectionId ? ` for collection ${collectionId}` : '';
|
||||
const cause = [404, 405, 501].includes(status) ? 'vector_endpoint_unavailable' : undefined;
|
||||
throw new Error(`Failed to ${action}${collectionText}${statusText}`, cause ? { cause } : undefined);
|
||||
const payload = await response.clone().json().catch(() => null);
|
||||
const cause = payload?.cause || ([404, 405, 501].includes(status) ? 'vector_endpoint_unavailable' : undefined);
|
||||
const message = payload?.message || `Failed to ${action}${collectionText}${statusText}`;
|
||||
throw new Error(message, cause ? { cause } : undefined);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -1055,7 +1085,7 @@ async function getSavedHashes(collectionId) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'get saved hashes', collectionId);
|
||||
await throwVectorResponseError(response, 'get saved hashes', collectionId);
|
||||
}
|
||||
|
||||
const hashes = await response.json();
|
||||
|
|
@ -1084,7 +1114,7 @@ async function insertVectorItems(collectionId, items) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'insert vector items', collectionId);
|
||||
await throwVectorResponseError(response, 'insert vector items', collectionId);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1159,7 +1189,7 @@ async function deleteVectorItems(collectionId, hashes) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'delete vector items', collectionId);
|
||||
await throwVectorResponseError(response, 'delete vector items', collectionId);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1170,7 +1200,8 @@ async function deleteVectorItems(collectionId, hashes) {
|
|||
* @returns {Promise<{ hashes: number[], metadata: object[]}>} - Hashes of the results
|
||||
*/
|
||||
async function queryCollection(collectionId, searchText, topK) {
|
||||
const args = await getAdditionalArgs([searchText]);
|
||||
throwIfSourceInvalid();
|
||||
const args = await getAdditionalArgs([searchText], true);
|
||||
const response = await fetch('/api/vector/query', {
|
||||
method: 'POST',
|
||||
headers: getRequestHeaders(),
|
||||
|
|
@ -1185,7 +1216,7 @@ async function queryCollection(collectionId, searchText, topK) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'query collection', collectionId);
|
||||
await throwVectorResponseError(response, 'query collection', collectionId);
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
|
|
@ -1200,7 +1231,8 @@ async function queryCollection(collectionId, searchText, topK) {
|
|||
* @returns {Promise<Record<string, { hashes: number[], metadata: object[] }>>} - Results mapped to collection IDs
|
||||
*/
|
||||
async function queryMultipleCollections(collectionIds, searchText, topK, threshold) {
|
||||
const args = await getAdditionalArgs([searchText]);
|
||||
throwIfSourceInvalid();
|
||||
const args = await getAdditionalArgs([searchText], true);
|
||||
const response = await fetch('/api/vector/query-multi', {
|
||||
method: 'POST',
|
||||
headers: getRequestHeaders(),
|
||||
|
|
@ -1215,7 +1247,7 @@ async function queryMultipleCollections(collectionIds, searchText, topK, thresho
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'query multiple collections');
|
||||
await throwVectorResponseError(response, 'query multiple collections');
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
|
|
@ -1227,10 +1259,6 @@ async function queryMultipleCollections(collectionIds, searchText, topK, thresho
|
|||
*/
|
||||
async function purgeFileVectorIndex(fileUrl) {
|
||||
try {
|
||||
if (!settings.enabled_files) {
|
||||
return;
|
||||
}
|
||||
|
||||
console.log(`Vectors: Purging file vector index for ${fileUrl}`);
|
||||
const collectionId = getFileCollectionId(fileUrl);
|
||||
|
||||
|
|
@ -1244,15 +1272,32 @@ async function purgeFileVectorIndex(fileUrl) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'purge vector index', collectionId);
|
||||
await throwVectorResponseError(response, 'purge vector index', collectionId);
|
||||
}
|
||||
|
||||
console.log(`Vectors: Purged vector index for collection ${collectionId}`);
|
||||
return true;
|
||||
} catch (error) {
|
||||
console.error('Vectors: Failed to purge file', error);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Purges file indexes without letting one failure block the rest.
|
||||
* @param {{url: string}[]} files File attachments
|
||||
* @returns {Promise<boolean>} Whether every index was purged
|
||||
*/
|
||||
async function purgeFileVectorIndexes(files) {
|
||||
let allSuccess = true;
|
||||
for (const file of files) {
|
||||
if (!await purgeFileVectorIndex(file.url)) {
|
||||
allSuccess = false;
|
||||
}
|
||||
}
|
||||
return allSuccess;
|
||||
}
|
||||
|
||||
/**
|
||||
* Purges the vector index for a collection.
|
||||
* @param {string} collectionId Collection ID to purge
|
||||
|
|
@ -1274,7 +1319,7 @@ async function purgeVectorIndex(collectionId) {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'purge vector index', collectionId);
|
||||
await throwVectorResponseError(response, 'purge vector index', collectionId);
|
||||
}
|
||||
|
||||
console.log(`Vectors: Purged vector index for collection ${collectionId}`);
|
||||
|
|
@ -1299,7 +1344,7 @@ async function purgeAllVectorIndexes() {
|
|||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throwVectorResponseError(response, 'purge all vector indexes');
|
||||
await throwVectorResponseError(response, 'purge all vector indexes');
|
||||
}
|
||||
|
||||
console.log('Vectors: Purged all vector indexes');
|
||||
|
|
@ -1315,6 +1360,7 @@ function toggleSettings() {
|
|||
$('#vectors_files_settings').toggle(!!settings.enabled_files);
|
||||
$('#vectors_chats_settings').toggle(!!settings.enabled_chats);
|
||||
$('#vectors_world_info_settings').toggle(!!settings.enabled_world_info);
|
||||
$('#transformers_vectorsModel').toggle(source === 'transformers');
|
||||
$('#together_vectorsModel').toggle(source === 'togetherai');
|
||||
$('#openai_vectorsModel').toggle(source === 'openai');
|
||||
$('#electronhub_vectorsModel').toggle(source === 'electronhub');
|
||||
|
|
@ -1457,22 +1503,25 @@ async function createWebLlmEmbeddings(items) {
|
|||
}
|
||||
|
||||
/**
|
||||
* Creates KoboldCpp embeddings for a list of items.
|
||||
* Creates KoboldCpp embeddings and keeps the server-reported model as part of
|
||||
* the Vector scope so changing the loaded embedding model cannot mix spaces.
|
||||
* @param {string[]} items Items to embed
|
||||
* @param {boolean} [isQuery=false] Whether this is a query embedding
|
||||
* @returns {Promise<{embeddings: Record<string, number[]>, model: string}>} Calculated embeddings
|
||||
*/
|
||||
async function createKoboldCppEmbeddings(items) {
|
||||
async function createKoboldCppEmbeddings(items, isQuery = false) {
|
||||
const response = await fetch('/api/backends/kobold/embed', {
|
||||
method: 'POST',
|
||||
headers: getRequestHeaders(),
|
||||
body: JSON.stringify({
|
||||
items: items,
|
||||
items,
|
||||
isQuery,
|
||||
server: settings.use_alt_endpoint ? settings.alt_endpoint_url : textgenerationwebui_settings.server_urls[textgen_types.KOBOLDCPP],
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error('Failed to get KoboldCpp embeddings');
|
||||
await throwVectorResponseError(response, 'create KoboldCpp embeddings');
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
|
|
@ -1480,19 +1529,15 @@ async function createKoboldCppEmbeddings(items) {
|
|||
throw new Error('Invalid response from KoboldCpp embeddings');
|
||||
}
|
||||
|
||||
const embeddings = /** @type {Record<string, number[]>} */ ({});
|
||||
const embeddings = /** @type {Record<string, number[]>} */ (Object.create(null));
|
||||
for (let i = 0; i < data.embeddings.length; i++) {
|
||||
if (!Array.isArray(data.embeddings[i]) || data.embeddings[i].length === 0) {
|
||||
throw new Error('KoboldCpp returned an empty embedding. Reduce the chunk size and/or size threshold and try again.');
|
||||
}
|
||||
|
||||
embeddings[items[i]] = data.embeddings[i];
|
||||
}
|
||||
|
||||
return {
|
||||
embeddings: embeddings,
|
||||
model: data.model,
|
||||
};
|
||||
return { embeddings, model: data.model };
|
||||
}
|
||||
|
||||
async function onPurgeClick() {
|
||||
|
|
@ -1609,7 +1654,7 @@ async function onVectorizeAllFilesClick() {
|
|||
}
|
||||
} catch (error) {
|
||||
console.error('Vectors: Failed to vectorize all files', error);
|
||||
toastr.error('Failed to vectorize all files', 'Vectorization failed');
|
||||
toastr.error(error?.message || 'Failed to vectorize all files', 'Vectorization failed');
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1619,11 +1664,12 @@ async function onPurgeFilesClick() {
|
|||
const chatAttachments = getContext().chat.filter(x => Array.isArray(x.extra?.files)).map(x => x.extra.files).flat();
|
||||
const allFiles = [...dataBank, ...chatAttachments];
|
||||
|
||||
for (const file of allFiles) {
|
||||
await purgeFileVectorIndex(file.url);
|
||||
const allSuccess = await purgeFileVectorIndexes(allFiles);
|
||||
if (allSuccess) {
|
||||
toastr.success('All files purged', 'Purge successful');
|
||||
} else {
|
||||
toastr.warning('Some files failed to purge. Check browser console for details.', 'Purge incomplete');
|
||||
}
|
||||
|
||||
toastr.success('All files purged', 'Purge successful');
|
||||
} catch (error) {
|
||||
console.error('Vectors: Failed to purge all files', error);
|
||||
toastr.error('Failed to purge all files', 'Purge failed');
|
||||
|
|
@ -1774,6 +1820,11 @@ export async function init() {
|
|||
saveSettingsDebounced();
|
||||
toggleSettings();
|
||||
});
|
||||
$('#vectors_transformers_model').val(settings.transformers_model).on('change', () => {
|
||||
settings.transformers_model = String($('#vectors_transformers_model').val());
|
||||
Object.assign(extension_settings.vectors, settings);
|
||||
saveSettingsDebounced();
|
||||
});
|
||||
$('#vector_altEndpointUrl_enabled').prop('checked', settings.use_alt_endpoint).on('input', () => {
|
||||
settings.use_alt_endpoint = $('#vector_altEndpointUrl_enabled').prop('checked');
|
||||
Object.assign(extension_settings.vectors, settings);
|
||||
|
|
@ -2121,11 +2172,9 @@ export async function init() {
|
|||
name: 'db-purge',
|
||||
callback: async () => {
|
||||
const dataBank = getDataBankAttachments();
|
||||
|
||||
for (const file of dataBank) {
|
||||
await purgeFileVectorIndex(file.url);
|
||||
if (!await purgeFileVectorIndexes(dataBank)) {
|
||||
throw new Error('Some Data Bank vector indexes failed to purge');
|
||||
}
|
||||
|
||||
return '';
|
||||
},
|
||||
aliases: ['databank-purge', 'data-bank-purge'],
|
||||
|
|
|
|||
|
|
@ -32,6 +32,20 @@
|
|||
<option value="webllm" data-i18n="WebLLM Extension">WebLLM Extension</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="flex-container flexFlowColumn" id="transformers_vectorsModel">
|
||||
<label for="vectors_transformers_model" data-i18n="Vectorization Model">
|
||||
Vectorization Model
|
||||
</label>
|
||||
<select id="vectors_transformers_model" class="text_pole">
|
||||
<option value="jinaai/jina-embeddings-v2-base-en">Jina Embeddings v2 Base English (legacy)</option>
|
||||
<option value="BAAI/bge-m3">BGE-M3</option>
|
||||
<option value="Qwen/Qwen3-Embedding-0.6B">Qwen3-Embedding-0.6B</option>
|
||||
<option value="google/embeddinggemma-300m">EmbeddingGemma 300M (Q4)</option>
|
||||
</select>
|
||||
<small>
|
||||
The selected model is downloaded on first use. Each model uses an isolated vector index.
|
||||
</small>
|
||||
</div>
|
||||
<div class="flex-container flexFlowColumn" id="chutes_vectorsModel">
|
||||
<label for="vectors_chutes_model" data-i18n="Vectorization Model">
|
||||
Vectorization Model
|
||||
|
|
|
|||
|
|
@ -257,6 +257,7 @@ export {};
|
|||
* | 'upload_user_image'
|
||||
* | 'upload_user_file'
|
||||
* | 'verify_user_files'
|
||||
* | 'vector_handle'
|
||||
* | 'view_secrets'
|
||||
* | 'wait_for_backend_ready'
|
||||
* | 'write_skill_file'
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import { extractErrorText, resolveHostErrorResponse } from '../kernel/host-error-response.js';
|
||||
|
||||
const SUPPORTED_VECTOR_ENDPOINTS = new Set([
|
||||
'list',
|
||||
'insert',
|
||||
|
|
@ -12,18 +14,59 @@ function normalizeEndpoint(wildcard) {
|
|||
return String(wildcard || '').replace(/^\/+/, '');
|
||||
}
|
||||
|
||||
export function registerVectorRoutes(router, _context, { jsonResponse }) {
|
||||
router.post('/api/vector/*', async ({ wildcard }) => {
|
||||
function vectorErrorCause(status) {
|
||||
if (status === 401) return 'embedding_auth_failed';
|
||||
if (status === 409) return 'vector_index_conflict';
|
||||
if (status === 429) return 'embedding_rate_limited';
|
||||
if (status >= 500) return 'embedding_unavailable';
|
||||
return undefined;
|
||||
}
|
||||
|
||||
async function invokeVector(context, jsonResponse, path, request) {
|
||||
try {
|
||||
const response = await context.safeInvoke('vector_handle', { path, request });
|
||||
const status = Number(response?.status);
|
||||
const safeStatus = Number.isInteger(status) && status >= 100 && status <= 599 ? status : 500;
|
||||
|
||||
if (response?.kind === 'empty') {
|
||||
return new Response(null, { status: safeStatus });
|
||||
}
|
||||
if (response?.kind === 'json') {
|
||||
return jsonResponse(response.body ?? null, safeStatus);
|
||||
}
|
||||
return jsonResponse({ error: true, cause: 'embedding_unavailable', message: 'Invalid native vector response' }, 500);
|
||||
} catch (error) {
|
||||
const resolved = resolveHostErrorResponse(extractErrorText(error));
|
||||
return jsonResponse({
|
||||
error: true,
|
||||
cause: vectorErrorCause(resolved.status),
|
||||
message: resolved.body,
|
||||
}, resolved.status);
|
||||
}
|
||||
}
|
||||
|
||||
export function registerVectorRoutes(router, context, { jsonResponse }) {
|
||||
router.post('/api/vector/*', async ({ body, wildcard }) => {
|
||||
const endpoint = normalizeEndpoint(wildcard);
|
||||
|
||||
if (!SUPPORTED_VECTOR_ENDPOINTS.has(endpoint)) {
|
||||
return jsonResponse({ error: `Unsupported vector endpoint: ${endpoint}` }, 404);
|
||||
}
|
||||
|
||||
return jsonResponse({
|
||||
error: true,
|
||||
cause: 'vector_endpoint_unavailable',
|
||||
message: 'Vector Storage backend is not implemented in the native TauriTavern backend yet.',
|
||||
}, 501);
|
||||
return invokeVector(
|
||||
context,
|
||||
jsonResponse,
|
||||
endpoint,
|
||||
body && typeof body === 'object' && !Array.isArray(body) ? body : {},
|
||||
);
|
||||
});
|
||||
|
||||
router.post('/api/backends/kobold/embed', async ({ body }) => {
|
||||
return invokeVector(context, jsonResponse, 'koboldcpp-embed', {
|
||||
source: 'koboldcpp',
|
||||
apiUrl: body?.server,
|
||||
...(body?.isQuery === true ? { isQuery: true } : {}),
|
||||
texts: body?.items,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,7 +28,8 @@ function extractDeclaration(source, marker) {
|
|||
const start = source.indexOf(marker);
|
||||
assert.notEqual(start, -1, `Missing declaration marker: ${marker}`);
|
||||
|
||||
const braceStart = source.indexOf('{', start);
|
||||
const functionBody = source.startsWith('function ', start) ? source.indexOf(') {', start) : -1;
|
||||
const braceStart = functionBody === -1 ? source.indexOf('{', start) : functionBody + 2;
|
||||
assert.notEqual(braceStart, -1, `Missing declaration body: ${marker}`);
|
||||
|
||||
let depth = 0;
|
||||
|
|
@ -86,6 +87,7 @@ async function getVectorsRuntimeSource() {
|
|||
vectorsRuntimeSourcePromise = readFile(new URL('../src/scripts/extensions/vectors/index.js', import.meta.url), 'utf8')
|
||||
.then(source => [
|
||||
extractDeclaration(source, 'const remoteEmbeddingEndpoints = {'),
|
||||
extractDeclaration(source, 'function getVectorsRequestBody(args = {})'),
|
||||
extractDeclaration(source, 'function toggleSettings()'),
|
||||
extractDeclaration(source, 'async function loadRemoteEmbeddingModels(source)'),
|
||||
].join('\n\n'));
|
||||
|
|
@ -150,6 +152,7 @@ async function createVectorsRuntimeHarness({
|
|||
workers_ai_model: '',
|
||||
siliconflow_model: '',
|
||||
electronhub_model: '',
|
||||
transformers_model: 'jinaai/jina-embeddings-v2-base-en',
|
||||
...settingsPatch,
|
||||
};
|
||||
const extension_settings = { vectors: {} };
|
||||
|
|
@ -196,7 +199,7 @@ async function createVectorsRuntimeHarness({
|
|||
'console',
|
||||
'vectorApiRequiresUrl',
|
||||
'loadWebLlmModels',
|
||||
`${source}\nreturn { remoteEmbeddingEndpoints, toggleSettings, loadRemoteEmbeddingModels };`,
|
||||
`${source}\nreturn { remoteEmbeddingEndpoints, getVectorsRequestBody, toggleSettings, loadRemoteEmbeddingModels };`,
|
||||
);
|
||||
|
||||
const runtime = factory(
|
||||
|
|
@ -453,21 +456,102 @@ test('vectors source toggle exposes remote model failures without clearing curre
|
|||
assert.equal(harness.savedSettings.length, 0);
|
||||
});
|
||||
|
||||
test('native vector endpoints fail fast while backend is not implemented', async () => {
|
||||
test('vectors local model selection is forwarded and shown only for transformers', async () => {
|
||||
const harness = await createVectorsRuntimeHarness({
|
||||
settings: {
|
||||
source: 'transformers',
|
||||
transformers_model: 'Qwen/Qwen3-Embedding-0.6B',
|
||||
},
|
||||
});
|
||||
|
||||
assert.deepEqual(harness.runtime.getVectorsRequestBody({ collectionId: 'chat-1' }), {
|
||||
collectionId: 'chat-1',
|
||||
model: 'Qwen/Qwen3-Embedding-0.6B',
|
||||
});
|
||||
harness.runtime.toggleSettings();
|
||||
assert.equal(harness.elements.get('#transformers_vectorsModel').visible, true);
|
||||
|
||||
harness.settings.source = 'openai';
|
||||
harness.runtime.toggleSettings();
|
||||
assert.equal(harness.elements.get('#transformers_vectorsModel').visible, false);
|
||||
});
|
||||
|
||||
test('native vector endpoints bridge the seven upstream routes to one typed command', async () => {
|
||||
const router = createRouteRegistry();
|
||||
registerVectorRoutes(router, {}, { jsonResponse });
|
||||
const calls = [];
|
||||
registerVectorRoutes(router, {
|
||||
safeInvoke: async (command, payload) => {
|
||||
calls.push({ command, payload });
|
||||
if (payload.path === 'list') {
|
||||
return { status: 200, kind: 'json', body: [17, 23] };
|
||||
}
|
||||
if (payload.path === 'query') {
|
||||
throw new Error('Bad request: searchText is required');
|
||||
}
|
||||
if (payload.path === 'koboldcpp-embed') {
|
||||
return { status: 200, kind: 'json', body: { embeddings: [[1, 0]], model: 'embed.gguf' } };
|
||||
}
|
||||
return { status: 200, kind: 'empty', body: null };
|
||||
},
|
||||
}, { jsonResponse });
|
||||
|
||||
const response = await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/vector/insert',
|
||||
body: { collectionId: 'chat-1', items: [] },
|
||||
body: { collectionId: 'chat-1', items: [], source: 'transformers' },
|
||||
});
|
||||
|
||||
assert.equal(response.status, 501);
|
||||
assert.deepEqual(await response.json(), {
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal(await response.text(), '');
|
||||
assert.deepEqual(calls, [{
|
||||
command: 'vector_handle',
|
||||
payload: {
|
||||
path: 'insert',
|
||||
request: { collectionId: 'chat-1', items: [], source: 'transformers' },
|
||||
},
|
||||
}]);
|
||||
|
||||
const listResponse = await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/vector/list',
|
||||
body: { collectionId: 'chat-1' },
|
||||
});
|
||||
assert.deepEqual(await listResponse.json(), [17, 23]);
|
||||
|
||||
const koboldResponse = await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/backends/kobold/embed',
|
||||
body: { server: 'http://localhost:5001', items: ['hello'] },
|
||||
});
|
||||
assert.deepEqual(await koboldResponse.json(), { embeddings: [[1, 0]], model: 'embed.gguf' });
|
||||
assert.deepEqual(calls.at(-1), {
|
||||
command: 'vector_handle',
|
||||
payload: {
|
||||
path: 'koboldcpp-embed',
|
||||
request: {
|
||||
source: 'koboldcpp',
|
||||
apiUrl: 'http://localhost:5001',
|
||||
texts: ['hello'],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/backends/kobold/embed',
|
||||
body: { server: 'http://localhost:5001', items: ['query'], isQuery: true },
|
||||
});
|
||||
assert.equal(calls.at(-1).payload.request.isQuery, true);
|
||||
|
||||
const invalidQuery = await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/vector/query',
|
||||
body: { collectionId: 'chat-1' },
|
||||
});
|
||||
assert.equal(invalidQuery.status, 400);
|
||||
assert.deepEqual(await invalidQuery.json(), {
|
||||
error: true,
|
||||
cause: 'vector_endpoint_unavailable',
|
||||
message: 'Vector Storage backend is not implemented in the native TauriTavern backend yet.',
|
||||
message: 'Bad request: searchText is required',
|
||||
});
|
||||
assert.equal(router.canHandle('POST', '/api/vector/query'), true);
|
||||
|
||||
|
|
@ -483,7 +567,31 @@ test('native vector endpoints fail fast while backend is not implemented', async
|
|||
});
|
||||
});
|
||||
|
||||
test('top-level native routes register vector fail-fast before provider routes', async () => {
|
||||
test('native vector routes classify operation-wide recoverable failures', async () => {
|
||||
const failures = [
|
||||
['Unauthorized: bad key', 401, 'embedding_auth_failed'],
|
||||
['Conflict: stale dimensions', 409, 'vector_index_conflict'],
|
||||
['Too many requests: retry later', 429, 'embedding_rate_limited'],
|
||||
['Internal server error: provider offline', 500, 'embedding_unavailable'],
|
||||
];
|
||||
|
||||
for (const [message, status, cause] of failures) {
|
||||
const router = createRouteRegistry();
|
||||
registerVectorRoutes(router, {
|
||||
safeInvoke: async () => { throw new Error(message); },
|
||||
}, { jsonResponse });
|
||||
|
||||
const response = await router.handle({
|
||||
method: 'POST',
|
||||
path: '/api/vector/insert',
|
||||
body: {},
|
||||
});
|
||||
assert.equal(response.status, status);
|
||||
assert.equal((await response.json()).cause, cause);
|
||||
}
|
||||
});
|
||||
|
||||
test('top-level native routes register vector compatibility before provider routes', async () => {
|
||||
const source = await readFile(new URL('../src/tauri/main/routes/index.js', import.meta.url), 'utf8');
|
||||
const aiIndex = source.indexOf('registerAiRoutes(router, context, responses);');
|
||||
const vectorIndex = source.indexOf('registerVectorRoutes(router, context, responses);');
|
||||
|
|
@ -497,20 +605,35 @@ test('top-level native routes register vector fail-fast before provider routes',
|
|||
assert.ok(vectorIndex < providerIndex);
|
||||
});
|
||||
|
||||
test('vectors runtime keeps upstream 1.18 summary skip settings and native fatal route semantics', async () => {
|
||||
test('vectors runtime keeps upstream 1.18 summary skip settings and native route semantics', async () => {
|
||||
const source = await readFile(new URL('../src/scripts/extensions/vectors/index.js', import.meta.url), 'utf8');
|
||||
const settingsHtml = await readFile(new URL('../src/scripts/extensions/vectors/settings.html', import.meta.url), 'utf8');
|
||||
const processFiles = extractDeclaration(source, 'async function processFiles(');
|
||||
const vectorizeFile = extractDeclaration(source, 'async function vectorizeFile(');
|
||||
const rearrangeChat = extractDeclaration(source, 'async function rearrangeChat(');
|
||||
const insertVectorItems = extractDeclaration(source, 'async function insertVectorItems(');
|
||||
const purgeFiles = extractDeclaration(source, 'async function purgeFileVectorIndexes(');
|
||||
|
||||
assert.match(source, /summary_retries: 2/);
|
||||
assert.match(source, /summary_threshold: 200/);
|
||||
assert.match(source, /keep_hidden: false/);
|
||||
assert.match(source, /const skippedHashes = new Set\(\);/);
|
||||
assert.match(source, /vector_endpoint_unavailable/);
|
||||
assert.match(source, /function throwVectorResponseError\(response, action, collectionId = ''\)/);
|
||||
assert.match(source, /async function throwVectorResponseError\(response, action, collectionId = ''\)/);
|
||||
assert.match(source, /\[404, 405, 501\]\.includes\(status\)/);
|
||||
assert.match(source, /settings\.keep_hidden \|\| !x\.is_system/);
|
||||
assert.match(source, /summarize\(toSummarize, settings\.summary_source, \{ skipOnFailure: true \}\)/);
|
||||
assert.match(source, /skippedHashes\.add\(item\.hash\)/);
|
||||
assert.match(source, /case 'koboldcpp':\s+body\.apiUrl =/);
|
||||
assert.match(source, /case 'transformers':\s+body\.model = settings\.transformers_model/);
|
||||
assert.match(source, /createKoboldCppEmbeddings/);
|
||||
assert.match(vectorizeFile, /await insertVectorItems\(collectionId, items\)/);
|
||||
assert.doesNotMatch(vectorizeFile, /items\.slice\(/);
|
||||
assert.doesNotMatch(insertVectorItems, /interactive/);
|
||||
assert.match(processFiles, /VECTOR_SOURCE_ABORT_CAUSES\.has\(error\?\.cause\)/);
|
||||
assert.match(rearrangeChat, /World Info vector retrieval skipped\. Chat retrieval will continue\./);
|
||||
assert.match(purgeFiles, /if \(!await purgeFileVectorIndex\(file\.url\)\)/);
|
||||
assert.match(purgeFiles, /return allSuccess/);
|
||||
|
||||
for (const id of ['vectors_keep_hidden', 'vectors_summary_retries', 'vectors_summary_threshold']) {
|
||||
assert.match(settingsHtml, new RegExp(`id="${id}"`));
|
||||
|
|
@ -518,6 +641,13 @@ test('vectors runtime keeps upstream 1.18 summary skip settings and native fatal
|
|||
}
|
||||
|
||||
assert.match(settingsHtml, /gemini-embedding-2-preview/);
|
||||
for (const model of [
|
||||
'BAAI/bge-m3',
|
||||
'Qwen/Qwen3-Embedding-0.6B',
|
||||
'google/embeddinggemma-300m',
|
||||
]) {
|
||||
assert.match(settingsHtml, new RegExp(model.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')));
|
||||
}
|
||||
});
|
||||
|
||||
test('unrelated secret mutations leave provider metadata caches intact', async () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue