mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-17 22:57:59 +00:00
Kept the runtime API responsive during blocking work (#1043)
This commit is contained in:
1 parent
4616d1a4b2
commit
fb756dd767
11 files changed
+389
-77
No files matched your search
@@ -130,6 +130,25 @@ Key modules:
|
|||||||
- `pandoc.rs` - Integration with Pandoc for document conversion
|
- `pandoc.rs` - Integration with Pandoc for document conversion
|
||||||
- `log.rs` - Logging infrastructure using `flexi_logger`
|
- `log.rs` - Logging infrastructure using `flexi_logger`
|
||||||
|
|
||||||
|
**Runtime API handlers never block.** All calls of the .NET app share one HTTP/2 connection, and the
|
||||||
|
task driving it may wait on exactly the Tokio worker which a blocking handler occupies, so a single
|
||||||
|
blocking handler can hold up the whole app. Therefore:
|
||||||
|
|
||||||
|
- File and disk access, the OS keyring, waiting for a `std::sync::Mutex`, and CPU-bound work such as
|
||||||
|
loading a tokenizer or scanning text run in `tokio::task::spawn_blocking`. Follow
|
||||||
|
`run_qdrant_edge_request` in `qdrant_edge_database.rs`, `prepare_image` and `prepare_image_sync` in
|
||||||
|
`image.rs`, or `sanitize_batch` in `prompt_injection/api.rs`.
|
||||||
|
- When an API offers a callback, as the file dialogs do, await the callback instead of calling the
|
||||||
|
blocking variant; see `await_dialog` in `file_actions.rs`.
|
||||||
|
- Short, bounded work, such as a single metadata lookup or writing a log line, may stay on the worker.
|
||||||
|
- When the blocking task fails, answer with an error, never with an empty value the app could take
|
||||||
|
for a valid answer.
|
||||||
|
- Never keep a `std::sync::MutexGuard` alive across an `.await`, not even with an explicit `drop`
|
||||||
|
before it: the future of the handler is then no longer `Send`, and Axum refuses it. Scope the guard
|
||||||
|
in a block instead.
|
||||||
|
- `runtime_api::test_support::assert_runtime_stays_free` tests a handler which waits for a lock. Run
|
||||||
|
such a test with `#[tokio::test(flavor = "multi_thread", worker_threads = 1)]`.
|
||||||
|
|
||||||
### .NET App (`app/MindWork AI Studio/`)
|
### .NET App (`app/MindWork AI Studio/`)
|
||||||
**Entry point:** `app/MindWork AI Studio/Program.cs`
|
**Entry point:** `app/MindWork AI Studio/Program.cs`
|
||||||
|
|
||||||
|
|||||||
@@ -26,6 +26,7 @@
|
|||||||
- Fixed AI Studio sometimes rebuilding the index of a data source from scratch when the configuration of your organization was applied while the data source was being updated.
|
- Fixed AI Studio sometimes rebuilding the index of a data source from scratch when the configuration of your organization was applied while the data source was being updated.
|
||||||
- Fixed justified texts being hyphenated by the rules of English even when AI Studio shows another language, such as German. Screen readers now also know which language AI Studio uses.
|
- Fixed justified texts being hyphenated by the rules of English even when AI Studio shows another language, such as German. Screen readers now also know which language AI Studio uses.
|
||||||
- Fixed the log file of AI Studio showing the wrong time for some entries when AI Studio was busy. Each entry now shows when it actually happened.
|
- Fixed the log file of AI Studio showing the wrong time for some entries when AI Studio was busy. Each entry now shows when it actually happened.
|
||||||
|
- Fixed AI Studio pausing for a moment while it opened the local index of a data source, such as a mailbox. Other work, for example applying the configuration of your organization, had to wait until then.
|
||||||
- Updated the code contributions on the supporters page, which now thank everyone who has contributed code to AI Studio so far.
|
- Updated the code contributions on the supporters page, which now thank everyone who has contributed code to AI Studio so far.
|
||||||
- Upgraded several libraries to improve security.
|
- Upgraded several libraries to improve security.
|
||||||
- Upgraded to Rust v1.99.0
|
- Upgraded to Rust v1.99.0
|
||||||
@@ -77,6 +77,26 @@ fn release_clipboard<B>(clipboard: &mut Option<B>) -> bool {
|
|||||||
|
|
||||||
/// Sets the clipboard text to the provided encrypted text.
|
/// Sets the clipboard text to the provided encrypted text.
|
||||||
pub async fn set_clipboard(_token: APIToken, encrypted_text: String) -> Json<SetClipboardResponse> {
|
pub async fn set_clipboard(_token: APIToken, encrypted_text: String) -> Json<SetClipboardResponse> {
|
||||||
|
//
|
||||||
|
// The clipboard backend of the operating system may take a moment to answer, a failed write is
|
||||||
|
// retried with a fresh backend, and the lock waits for a write still running. None of this
|
||||||
|
// belongs on a runtime worker, because all calls of the app share one HTTP/2 connection.
|
||||||
|
//
|
||||||
|
match tokio::task::spawn_blocking(move || set_clipboard_sync(encrypted_text)).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(error) => {
|
||||||
|
let issue = format!("The clipboard worker failed: {error}");
|
||||||
|
error!(Source = "Clipboard"; "{issue}");
|
||||||
|
Json(SetClipboardResponse {
|
||||||
|
success: false,
|
||||||
|
issue,
|
||||||
|
})
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Performs the blocking part of [`set_clipboard`].
|
||||||
|
fn set_clipboard_sync(encrypted_text: String) -> Json<SetClipboardResponse> {
|
||||||
let encrypted_text = EncryptedText::new(encrypted_text);
|
let encrypted_text = EncryptedText::new(encrypted_text);
|
||||||
|
|
||||||
// Decrypt this text first:
|
// Decrypt this text first:
|
||||||
|
|||||||
+25
-24
@@ -1,4 +1,5 @@
|
|||||||
use crate::api_token::APIToken;
|
use crate::api_token::APIToken;
|
||||||
|
use axum::http::StatusCode;
|
||||||
use axum::Json;
|
use axum::Json;
|
||||||
use log::{debug, error, info, warn};
|
use log::{debug, error, info, warn};
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
@@ -641,35 +642,35 @@ struct EnterpriseSourceData {
|
|||||||
encryption_secret: String,
|
encryption_secret: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn read_enterprise_env_config_id(_token: APIToken) -> String {
|
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> Result<String, (StatusCode, String)> {
|
||||||
debug!("Trying to read the effective enterprise configuration ID.");
|
|
||||||
resolve_effective_enterprise_config_source()
|
|
||||||
.configs
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
.map(|config| config.id)
|
|
||||||
.unwrap_or_default()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn read_enterprise_env_config_server_url(_token: APIToken) -> String {
|
|
||||||
debug!("Trying to read the effective enterprise configuration server URL.");
|
|
||||||
resolve_effective_enterprise_config_source()
|
|
||||||
.configs
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
.map(|config| config.server_url)
|
|
||||||
.unwrap_or_default()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> String {
|
|
||||||
debug!("Trying to read the effective enterprise configuration encryption secret.");
|
debug!("Trying to read the effective enterprise configuration encryption secret.");
|
||||||
resolve_effective_enterprise_secret_source().encryption_secret
|
read_enterprise_sources(|| resolve_effective_enterprise_secret_source().encryption_secret).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns all enterprise configurations from the effective source.
|
/// Returns all enterprise configurations from the effective source.
|
||||||
pub async fn read_enterprise_configs(_token: APIToken) -> Json<Vec<EnterpriseConfig>> {
|
pub async fn read_enterprise_configs(_token: APIToken) -> Result<Json<Vec<EnterpriseConfig>>, (StatusCode, String)> {
|
||||||
info!("Trying to read the effective enterprise configurations.");
|
info!("Trying to read the effective enterprise configurations.");
|
||||||
Json(resolve_effective_enterprise_config_source().configs)
|
read_enterprise_sources(|| Json(resolve_effective_enterprise_config_source().configs)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reads the enterprise sources on the blocking pool.
|
||||||
|
///
|
||||||
|
/// Every call reads the policy files and, on Windows, the registry again. That is quick most of
|
||||||
|
/// the time, but a slow disk must not hold up every other call of the app, because they all share
|
||||||
|
/// one HTTP/2 connection.
|
||||||
|
///
|
||||||
|
/// A failure is answered as an error, never as an empty value: the app would read an empty value
|
||||||
|
/// as "no configuration of your organization" and act on it.
|
||||||
|
async fn read_enterprise_sources<T, F>(read: F) -> Result<T, (StatusCode, String)>
|
||||||
|
where
|
||||||
|
T: Send + 'static,
|
||||||
|
F: FnOnce() -> T + Send + 'static,
|
||||||
|
{
|
||||||
|
tokio::task::spawn_blocking(read).await.map_err(|error| {
|
||||||
|
let issue = format!("The enterprise configuration worker failed: {error}");
|
||||||
|
error!("{issue}");
|
||||||
|
(StatusCode::INTERNAL_SERVER_ERROR, issue)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn resolve_external_http_custom_root_certificate_policy() -> ExternalHttpCustomRootCertificatePolicy {
|
pub fn resolve_external_http_custom_root_certificate_policy() -> ExternalHttpCustomRootCertificatePolicy {
|
||||||
|
|||||||
+51
-17
@@ -123,26 +123,25 @@ pub async fn select_directory(
|
|||||||
Query(query): Query<SelectDirectoryQuery>,
|
Query(query): Query<SelectDirectoryQuery>,
|
||||||
previous_directory: Option<Json<PreviousDirectory>>,
|
previous_directory: Option<Json<PreviousDirectory>>,
|
||||||
) -> Json<DirectorySelectionResponse> {
|
) -> Json<DirectorySelectionResponse> {
|
||||||
let main_window_lock = MAIN_WINDOW.lock().unwrap();
|
let dialog = MAIN_WINDOW
|
||||||
let main_window = match main_window_lock.as_ref() {
|
.lock()
|
||||||
Some(window) => window,
|
.unwrap()
|
||||||
None => {
|
.as_ref()
|
||||||
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
|
.map(|w| w.dialog().file().set_parent(w).set_title(&query.title));
|
||||||
return Json(DirectorySelectionResponse {
|
|
||||||
user_cancelled: true,
|
let Some(mut dialog) = dialog else {
|
||||||
selected_directory: String::from(""),
|
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
|
||||||
});
|
return Json(DirectorySelectionResponse {
|
||||||
}
|
user_cancelled: true,
|
||||||
|
selected_directory: String::from(""),
|
||||||
|
});
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut dialog = main_window.dialog().file().set_parent(main_window).set_title(&query.title);
|
|
||||||
if let Some(previous) = previous_directory {
|
if let Some(previous) = previous_directory {
|
||||||
dialog = dialog.set_directory(previous.path.clone());
|
dialog = dialog.set_directory(previous.path.clone());
|
||||||
}
|
}
|
||||||
|
|
||||||
drop(main_window_lock);
|
let folder_path = await_dialog(|answer| dialog.pick_folder(answer)).await;
|
||||||
|
|
||||||
let folder_path = dialog.blocking_pick_folder();
|
|
||||||
match folder_path {
|
match folder_path {
|
||||||
Some(path) => {
|
Some(path) => {
|
||||||
match path.into_path() {
|
match path.into_path() {
|
||||||
@@ -203,7 +202,7 @@ pub async fn select_file(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Show the file dialog and get the selected file path:
|
// Show the file dialog and get the selected file path:
|
||||||
let file_path = file_dialog.blocking_pick_file();
|
let file_path = await_dialog(|answer| file_dialog.pick_file(answer)).await;
|
||||||
match file_path {
|
match file_path {
|
||||||
Some(path) => match path.into_path() {
|
Some(path) => match path.into_path() {
|
||||||
Ok(pb) => {
|
Ok(pb) => {
|
||||||
@@ -262,7 +261,7 @@ pub async fn select_files(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Show the file dialog and get the selected file path:
|
// Show the file dialog and get the selected file path:
|
||||||
let file_paths = file_dialog.blocking_pick_files();
|
let file_paths = await_dialog(|answer| file_dialog.pick_files(answer)).await;
|
||||||
match file_paths {
|
match file_paths {
|
||||||
Some(paths) => {
|
Some(paths) => {
|
||||||
let converted: Vec<String> = paths.into_iter().filter_map(|p| p.into_path().ok()).map(|pb| pb.to_string_lossy().to_string()).collect();
|
let converted: Vec<String> = paths.into_iter().filter_map(|p| p.into_path().ok()).map(|pb| pb.to_string_lossy().to_string()).collect();
|
||||||
@@ -314,7 +313,7 @@ pub async fn save_file(_token: APIToken, payload: Json<SaveFileOptions>) -> Json
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Displays the file dialogue box and select the file:
|
// Displays the file dialogue box and select the file:
|
||||||
let file_path = file_dialog.blocking_save_file();
|
let file_path = await_dialog(|answer| file_dialog.save_file(answer)).await;
|
||||||
match file_path {
|
match file_path {
|
||||||
Some(path) => match path.into_path() {
|
Some(path) => match path.into_path() {
|
||||||
Ok(pb) => {
|
Ok(pb) => {
|
||||||
@@ -882,6 +881,26 @@ fn apply_filter<R: tauri::Runtime>(file_dialog: FileDialogBuilder<R>, filter: &O
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Shows a file dialog and waits for the user's answer without holding up the runtime.
|
||||||
|
///
|
||||||
|
/// A dialog stays open until the user answers it, which can take minutes. The blocking variants of
|
||||||
|
/// the dialog API would hold a runtime worker that long, and with it every other call of the app,
|
||||||
|
/// because they all share one HTTP/2 connection. The plugin shows the dialog on the main thread and
|
||||||
|
/// hands the answer to a callback, so this function awaits that callback instead.
|
||||||
|
async fn await_dialog<T: Send + 'static>(show: impl FnOnce(Box<dyn FnOnce(Option<T>) + Send>)) -> Option<T> {
|
||||||
|
let (sender, receiver) = tokio::sync::oneshot::channel();
|
||||||
|
show(Box::new(move |answer| {
|
||||||
|
// The receiver is only gone when the request was dropped, and then nobody waits for the answer:
|
||||||
|
let _ = sender.send(answer);
|
||||||
|
}));
|
||||||
|
|
||||||
|
receiver.await.unwrap_or_else(|_| {
|
||||||
|
// The plugin drops the callback without calling it when it cannot reach the main thread:
|
||||||
|
error!(Source = "Tauri"; "The file dialog ended without an answer.");
|
||||||
|
None
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
fn split_save_file_path(file_path: &str) -> (Option<PathBuf>, Option<String>) {
|
fn split_save_file_path(file_path: &str) -> (Option<PathBuf>, Option<String>) {
|
||||||
let path = Path::new(file_path);
|
let path = Path::new(file_path);
|
||||||
let directory = path
|
let directory = path
|
||||||
@@ -1047,6 +1066,21 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_dialog_answer_arrives_through_its_callback() {
|
||||||
|
let answer = await_dialog(|answer| {
|
||||||
|
std::thread::spawn(move || answer(Some("Quarterly briefing.html")));
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
assert_eq!(answer, Some("Quarterly briefing.html"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_dialog_without_an_answer_counts_as_cancelled() {
|
||||||
|
let answer = await_dialog::<String>(drop).await;
|
||||||
|
assert_eq!(answer, None);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn save_file_options_accept_the_previous_file_contract() {
|
fn save_file_options_accept_the_previous_file_contract() {
|
||||||
let options: SaveFileOptions = serde_json::from_str(
|
let options: SaveFileOptions = serde_json::from_str(
|
||||||
|
|||||||
@@ -524,7 +524,7 @@ pub async fn extract_data(
|
|||||||
// whole file, and a failure here means we cannot answer the request at all.
|
// whole file, and a failure here means we cannot answer the request at all.
|
||||||
//
|
//
|
||||||
let tokenizer = if query.include_token_count {
|
let tokenizer = if query.include_token_count {
|
||||||
match crate::tokenizer::get_tokenizer(&query.tokenizer_path) {
|
match crate::tokenizer::get_tokenizer_off_worker(query.tokenizer_path.clone()).await {
|
||||||
Ok(tokenizer) => Some(tokenizer),
|
Ok(tokenizer) => Some(tokenizer),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
let error = ExtractionError::new(ExtractionErrorCode::InvalidRequest, format!("The tokenizer could not be loaded: {e}"));
|
let error = ExtractionError::new(ExtractionErrorCode::InvalidRequest, format!("The tokenizer could not be loaded: {e}"));
|
||||||
|
|||||||
@@ -45,13 +45,30 @@ pub struct SanitizeBatchResponse {
|
|||||||
pub results: Vec<SanitizeResponse>,
|
pub results: Vec<SanitizeResponse>,
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn sanitize(_token: APIToken, Json(request): Json<SanitizeRequest>) -> Json<SanitizeResponse> {
|
pub async fn sanitize(
|
||||||
let (sanitized_text, report) = sanitize_text(&request.text);
|
_token: APIToken,
|
||||||
|
Json(request): Json<SanitizeRequest>,
|
||||||
|
) -> Result<Json<SanitizeResponse>, (StatusCode, String)> {
|
||||||
|
//
|
||||||
|
// A single web page or retrieval context can still be large, and scanning it is CPU-bound. Like
|
||||||
|
// a batch, it goes to the blocking pool, so it does not stall every other call the app makes
|
||||||
|
// meanwhile.
|
||||||
|
//
|
||||||
|
tokio::task::spawn_blocking(move || {
|
||||||
|
let (sanitized_text, report) = sanitize_text(&request.text);
|
||||||
|
|
||||||
Json(SanitizeResponse {
|
Json(SanitizeResponse {
|
||||||
sanitized_text,
|
sanitized_text,
|
||||||
findings: report.findings,
|
findings: report.findings,
|
||||||
redacted_count: report.redacted_count,
|
redacted_count: report.redacted_count,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|error| {
|
||||||
|
(
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("The prompt injection filter failed: {error}"),
|
||||||
|
)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -473,15 +473,25 @@ fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
||||||
let status = QDRANT_EDGE_STATUS.lock().unwrap();
|
let (current_status, mut unavailable_reason) = {
|
||||||
let current_status = status.status;
|
let status = QDRANT_EDGE_STATUS.lock().unwrap();
|
||||||
let unavailable_reason = status.unavailable_reason.clone();
|
(status.status, status.unavailable_reason.clone())
|
||||||
drop(status);
|
};
|
||||||
|
|
||||||
let database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
|
//
|
||||||
let database_info = database_guard
|
// The info needs the lock of the database, which another request may hold for seconds while
|
||||||
.as_ref()
|
// it loads a store, and it reads the stores directory. Both belong on the blocking pool, see
|
||||||
.and_then(|database| database.info().ok());
|
// run_qdrant_edge_request.
|
||||||
|
//
|
||||||
|
let database_info = match tokio::task::spawn_blocking(read_qdrant_edge_info).await {
|
||||||
|
Ok(database_info) => database_info,
|
||||||
|
Err(error) => {
|
||||||
|
let reason = format!("The Qdrant Edge worker failed: {error}");
|
||||||
|
error!(Source = "Qdrant Edge"; "Reading the Qdrant Edge info failed: {reason}");
|
||||||
|
unavailable_reason = Some(reason);
|
||||||
|
None
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
let is_available = current_status == QdrantEdgeStatus::Available && database_info.is_some();
|
let is_available = current_status == QdrantEdgeStatus::Available && database_info.is_some();
|
||||||
Json(QdrantEdgeServiceInfo {
|
Json(QdrantEdgeServiceInfo {
|
||||||
@@ -495,40 +505,48 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn read_qdrant_edge_info() -> Option<QdrantEdgeInfo> {
|
||||||
|
QDRANT_EDGE_DATABASE
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|database| database.info().ok())
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<QdrantEdgeEnsureStoreResult>> {
|
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<QdrantEdgeEnsureStoreResult>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
|
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<()>> {
|
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.insert_embedding(&request.store_name, request.points)
|
database.insert_embedding(&request.store_name, request.points)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn search_qdrant_edge_embeddings(_token: APIToken, Json(request): Json<SearchQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<Vec<QdrantEdgeSearchResult>>> {
|
pub async fn search_qdrant_edge_embeddings(_token: APIToken, Json(request): Json<SearchQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<Vec<QdrantEdgeSearchResult>>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.search_embedding(&request.store_name, request.vector, request.max_matches, request.point_ids)
|
database.search_embedding(&request.store_name, request.vector, request.max_matches, request.point_ids)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeResponse<()>> {
|
pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.delete_embedding_by_file(&request.store_name, &request.file_path)
|
database.delete_embedding_by_file(&request.store_name, &request.file_path)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.optimize_store(&request.store_name)
|
database.optimize_store(&request.store_name)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||||
execute_qdrant_edge_request(|database| {
|
run_qdrant_edge_request(move |database| {
|
||||||
database.delete_store(&request.store_name)
|
database.delete_store(&request.store_name)
|
||||||
})
|
}).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn start_qdrant_edge_database<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {
|
pub fn start_qdrant_edge_database<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {
|
||||||
@@ -569,6 +587,32 @@ pub fn stop_qdrant_edge_database() {
|
|||||||
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
|
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Runs one request against the database on the blocking pool.
|
||||||
|
///
|
||||||
|
/// A request waits for the lock of the database and may load a store from disk, which together
|
||||||
|
/// can take seconds. On a runtime worker, that would hold up every other call of the app as well:
|
||||||
|
/// they all share one HTTP/2 connection, and the task driving it may be waiting in the LIFO slot of
|
||||||
|
/// exactly this worker, where no other worker can take it over.
|
||||||
|
async fn run_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
||||||
|
where
|
||||||
|
T: Serialize + Send + 'static,
|
||||||
|
F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<T> + Send + 'static,
|
||||||
|
{
|
||||||
|
match tokio::task::spawn_blocking(move || execute_qdrant_edge_request(operation)).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(error) => {
|
||||||
|
let issue = format!("The Qdrant Edge worker failed: {error}");
|
||||||
|
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
|
||||||
|
Json(QdrantEdgeResponse {
|
||||||
|
success: false,
|
||||||
|
issue,
|
||||||
|
issue_code: "",
|
||||||
|
data: None,
|
||||||
|
})
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
||||||
where
|
where
|
||||||
T: Serialize,
|
T: Serialize,
|
||||||
@@ -1033,6 +1077,7 @@ fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::runtime_api::test_support::assert_runtime_stays_free;
|
||||||
use qdrant_edge::WalOptions;
|
use qdrant_edge::WalOptions;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1254,6 +1299,27 @@ mod tests {
|
|||||||
fs::remove_dir_all(test_directory).unwrap();
|
fs::remove_dir_all(test_directory).unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||||
|
async fn a_request_waiting_for_the_database_leaves_the_runtime_free() {
|
||||||
|
let request = optimize_qdrant_edge_store(
|
||||||
|
APIToken::from_hex_text("test"),
|
||||||
|
Json(OptimizeQdrantEdgeStoreRequest { store_name: SEARCH_STORE.to_string() }),
|
||||||
|
);
|
||||||
|
|
||||||
|
let response = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
|
||||||
|
|
||||||
|
// No database runs in the tests, so the request ends as soon as it gets the lock.
|
||||||
|
assert!(!response.success);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||||
|
async fn the_info_waiting_for_the_database_leaves_the_runtime_free() {
|
||||||
|
let request = qdrant_edge_info(APIToken::from_hex_text("test"));
|
||||||
|
|
||||||
|
let info = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
|
||||||
|
assert!(!info.is_available);
|
||||||
|
}
|
||||||
|
|
||||||
const SEARCH_STORE: &str = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
const SEARCH_STORE: &str = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
||||||
const POINT_X: &str = "0b5f1e8a-3c2d-4e6f-9a1b-7c8d9e0f1a2b";
|
const POINT_X: &str = "0b5f1e8a-3c2d-4e6f-9a1b-7c8d9e0f1a2b";
|
||||||
const POINT_Y: &str = "1c6a2f9b-4d3e-4f70-8b2c-8d9e0f1a2b3c";
|
const POINT_Y: &str = "1c6a2f9b-4d3e-4f70-8b2c-8d9e0f1a2b3c";
|
||||||
|
|||||||
@@ -64,8 +64,6 @@ pub fn start_runtime_api() {
|
|||||||
.route("/system/runtime/info", get(crate::environment::get_runtime_info))
|
.route("/system/runtime/info", get(crate::environment::get_runtime_info))
|
||||||
.route("/system/language", get(crate::environment::read_user_language))
|
.route("/system/language", get(crate::environment::read_user_language))
|
||||||
.route("/system/username", get(crate::environment::read_user_name))
|
.route("/system/username", get(crate::environment::read_user_name))
|
||||||
.route("/system/enterprise/config/id", get(crate::environment::read_enterprise_env_config_id))
|
|
||||||
.route("/system/enterprise/config/server", get(crate::environment::read_enterprise_env_config_server_url))
|
|
||||||
.route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
|
.route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
|
||||||
.route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
|
.route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
|
||||||
.route("/retrieval/fs/extract", get(crate::file_data::extract_data))
|
.route("/retrieval/fs/extract", get(crate::file_data::extract_data))
|
||||||
@@ -113,3 +111,51 @@ fn install_rustls_crypto_provider() {
|
|||||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) mod test_support {
|
||||||
|
use std::sync::Mutex;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
/// How long the lock stays held while the request waits for it.
|
||||||
|
const LOCK_DURATION: Duration = Duration::from_secs(1);
|
||||||
|
|
||||||
|
/// How long the probe may wait for the runtime. Far below the lock duration, so a request which
|
||||||
|
/// blocks the worker while it waits fails the test every time.
|
||||||
|
const PROBE_LIMIT: Duration = Duration::from_millis(200);
|
||||||
|
|
||||||
|
/// Starts a request while another thread holds the given lock, as if a long operation were
|
||||||
|
/// running, and checks that a probe still gets through the runtime meanwhile. Returns the
|
||||||
|
/// response once the lock is free again.
|
||||||
|
///
|
||||||
|
/// Call it from a test with a single worker, `#[tokio::test(flavor = "multi_thread",
|
||||||
|
/// worker_threads = 1)]`. A request which blocks that worker while it waits then holds up the
|
||||||
|
/// probe, just as in the app it holds up the one HTTP/2 connection which all calls share.
|
||||||
|
pub(crate) async fn assert_runtime_stays_free<T, F>(lock: &'static Mutex<T>, request: F) -> F::Output
|
||||||
|
where
|
||||||
|
T: Send + 'static,
|
||||||
|
F: Future + Send + 'static,
|
||||||
|
F::Output: Send + 'static,
|
||||||
|
{
|
||||||
|
let (locked_sender, locked_receiver) = std::sync::mpsc::channel();
|
||||||
|
let lock_holder = std::thread::spawn(move || {
|
||||||
|
let guard = lock.lock().unwrap();
|
||||||
|
locked_sender.send(()).unwrap();
|
||||||
|
std::thread::sleep(LOCK_DURATION);
|
||||||
|
drop(guard);
|
||||||
|
});
|
||||||
|
|
||||||
|
locked_receiver.recv().unwrap();
|
||||||
|
let request = tokio::spawn(request);
|
||||||
|
|
||||||
|
let probe_started = Instant::now();
|
||||||
|
tokio::spawn(async {}).await.unwrap();
|
||||||
|
let probe_duration = probe_started.elapsed();
|
||||||
|
|
||||||
|
lock_holder.join().unwrap();
|
||||||
|
let response = request.await.unwrap();
|
||||||
|
|
||||||
|
assert!(probe_duration < PROBE_LIMIT, "The probe waited {probe_duration:?} for the runtime.");
|
||||||
|
response
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -78,8 +78,37 @@ pub fn init_secret_store() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Runs one request against the credential store on the blocking pool.
|
||||||
|
///
|
||||||
|
/// The credential store of the operating system may take a moment to answer, and on macOS and Linux
|
||||||
|
/// it may wait until the user unlocks it, which can take minutes. On a runtime worker, that would
|
||||||
|
/// hold up every other call of the app, because they all share one HTTP/2 connection.
|
||||||
|
async fn run_secret_store_request<T, F>(request: F) -> Result<T, String>
|
||||||
|
where
|
||||||
|
T: Send + 'static,
|
||||||
|
F: FnOnce() -> T + Send + 'static,
|
||||||
|
{
|
||||||
|
tokio::task::spawn_blocking(request).await.map_err(|error| {
|
||||||
|
let issue = format!("The secret store worker failed: {error}");
|
||||||
|
error!(Source = "Secret Store"; "{issue}");
|
||||||
|
issue
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// Stores a secret in the secret store using the operating system's keyring.
|
/// Stores a secret in the secret store using the operating system's keyring.
|
||||||
pub async fn store_secret(_token: APIToken, request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
|
pub async fn store_secret(_token: APIToken, request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
|
||||||
|
match run_secret_store_request(move || store_secret_sync(request)).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(issue) => Json(StoreSecretResponse {
|
||||||
|
success: false,
|
||||||
|
issue,
|
||||||
|
issue_code: SecretStoreIssueCode::Unknown,
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Performs the blocking part of [`store_secret`].
|
||||||
|
fn store_secret_sync(request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
|
||||||
let user_name = request.user_name.as_str();
|
let user_name = request.user_name.as_str();
|
||||||
let decrypted_text = match ENCRYPTION.decrypt(&request.secret) {
|
let decrypted_text = match ENCRYPTION.decrypt(&request.secret) {
|
||||||
Ok(text) => text,
|
Ok(text) => text,
|
||||||
@@ -145,6 +174,19 @@ pub struct StoreSecretResponse {
|
|||||||
|
|
||||||
/// Retrieves a secret from the secret store using the operating system's keyring.
|
/// Retrieves a secret from the secret store using the operating system's keyring.
|
||||||
pub async fn get_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<RequestedSecret> {
|
pub async fn get_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<RequestedSecret> {
|
||||||
|
match run_secret_store_request(move || get_secret_sync(request)).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(issue) => Json(RequestedSecret {
|
||||||
|
success: false,
|
||||||
|
secret: EncryptedText::new(String::from("")),
|
||||||
|
issue,
|
||||||
|
issue_code: SecretStoreIssueCode::Unknown,
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Performs the blocking part of [`get_secret`].
|
||||||
|
fn get_secret_sync(request: Json<RequestSecret>) -> Json<RequestedSecret> {
|
||||||
let user_name = request.user_name.as_str();
|
let user_name = request.user_name.as_str();
|
||||||
let service = format!("mindwork-ai-studio::{}", request.destination);
|
let service = format!("mindwork-ai-studio::{}", request.destination);
|
||||||
let entry = match Entry::new(service.as_str(), user_name) {
|
let entry = match Entry::new(service.as_str(), user_name) {
|
||||||
@@ -227,6 +269,19 @@ pub struct RequestedSecret {
|
|||||||
|
|
||||||
/// Deletes a secret from the secret store using the operating system's keyring.
|
/// Deletes a secret from the secret store using the operating system's keyring.
|
||||||
pub async fn delete_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
|
pub async fn delete_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
|
||||||
|
match run_secret_store_request(move || delete_secret_sync(request)).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(issue) => Json(DeleteSecretResponse {
|
||||||
|
success: false,
|
||||||
|
was_entry_found: false,
|
||||||
|
issue,
|
||||||
|
issue_code: SecretStoreIssueCode::Unknown,
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Performs the blocking part of [`delete_secret`].
|
||||||
|
fn delete_secret_sync(request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
|
||||||
let user_name = request.user_name.as_str();
|
let user_name = request.user_name.as_str();
|
||||||
let service = format!("mindwork-ai-studio::{}", request.destination);
|
let service = format!("mindwork-ai-studio::{}", request.destination);
|
||||||
let entry = match Entry::new(service.as_str(), user_name) {
|
let entry = match Entry::new(service.as_str(), user_name) {
|
||||||
|
|||||||
@@ -104,33 +104,58 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
|
pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
|
||||||
match get_token_count(&req.tokenizer_path, &req.text) {
|
match run_tokenizer_work(move || get_token_count(&req.tokenizer_path, &req.text)).await {
|
||||||
Ok(count) => Json(TokenizerResponse::available(count)),
|
Ok(count) => Json(TokenizerResponse::available(count)),
|
||||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn validate_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
|
pub async fn validate_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
|
||||||
match handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())) {
|
match run_tokenizer_work(move || handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone()))).await {
|
||||||
Ok(count) => Json(TokenizerResponse::available(count)),
|
Ok(count) => Json(TokenizerResponse::available(count)),
|
||||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn store_tokenizer(_token: APIToken, payload: Json<TokenizerStorage>) -> Json<TokenizerResponse> {
|
pub async fn store_tokenizer(_token: APIToken, payload: Json<TokenizerStorage>) -> Json<TokenizerResponse> {
|
||||||
match handle_tokenizer_store(&payload) {
|
match run_tokenizer_work(move || handle_tokenizer_store(&payload).map_err(|e| e.to_string())).await {
|
||||||
Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
|
Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
|
||||||
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
|
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>) -> Json<TokenizerResponse> {
|
pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>) -> Json<TokenizerResponse> {
|
||||||
match handle_tokenizer_delete(&payload) {
|
match run_tokenizer_work(move || handle_tokenizer_delete(&payload).map_err(|e| e.to_string())).await {
|
||||||
Ok(_) => Json(TokenizerResponse::stored(String::new())),
|
Ok(_) => Json(TokenizerResponse::stored(String::new())),
|
||||||
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
|
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Gets the tokenizer for the given path, loading it on the blocking pool when it is not cached yet.
|
||||||
|
pub async fn get_tokenizer_off_worker(path: String) -> Result<Arc<Tokenizer>, String> {
|
||||||
|
run_tokenizer_work(move || get_tokenizer(&path)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Runs tokenizer work on the blocking pool.
|
||||||
|
///
|
||||||
|
/// Loading a tokenizer reads and parses its file, which takes most of a second for a large one, and
|
||||||
|
/// it waits for the storage lock while another request copies or deletes a tokenizer. Encoding a
|
||||||
|
/// long text is CPU-bound as well. On a runtime worker, any of this would hold up every other call
|
||||||
|
/// of the app, because they all share one HTTP/2 connection.
|
||||||
|
async fn run_tokenizer_work<T, F>(work: F) -> Result<T, String>
|
||||||
|
where
|
||||||
|
T: Send + 'static,
|
||||||
|
F: FnOnce() -> Result<T, String> + Send + 'static,
|
||||||
|
{
|
||||||
|
tokio::task::spawn_blocking(work)
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|error| {
|
||||||
|
let issue = format!("The tokenizer worker failed: {error}");
|
||||||
|
error!(Source = "Tokenizer"; "{issue}");
|
||||||
|
Err(issue)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
|
fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
|
||||||
validate_tokenizer_file(path)
|
validate_tokenizer_file(path)
|
||||||
}
|
}
|
||||||
@@ -240,7 +265,7 @@ fn tokenizer_cache() -> &'static RwLock<HashMap<PathBuf, Arc<Tokenizer>>> {
|
|||||||
TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
|
TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
|
fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
|
||||||
let resolved_path = resolve_tokenizer_path(path)?;
|
let resolved_path = resolve_tokenizer_path(path)?;
|
||||||
let tokenizer_path = fs::canonicalize(&resolved_path)
|
let tokenizer_path = fs::canonicalize(&resolved_path)
|
||||||
.map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
|
.map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
|
||||||
@@ -303,3 +328,31 @@ fn load_tokenizer_from_file(path: &PathBuf) -> Result<Tokenizer, String> {
|
|||||||
Tokenizer::from_file(path)
|
Tokenizer::from_file(path)
|
||||||
.map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
|
.map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::runtime_api::test_support::assert_runtime_stays_free;
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||||
|
async fn a_token_count_waiting_for_the_storage_leaves_the_runtime_free() {
|
||||||
|
//
|
||||||
|
// Any file will do: the request waits for the storage lock before it reads the tokenizer.
|
||||||
|
// Since this file is no tokenizer, the request fails once it gets the lock.
|
||||||
|
//
|
||||||
|
let tokenizer_path = std::env::temp_dir().join(format!("ai-studio-tokenizer-{}.json", std::process::id()));
|
||||||
|
fs::write(&tokenizer_path, "This is not a tokenizer.").unwrap();
|
||||||
|
|
||||||
|
let request = token_count(
|
||||||
|
APIToken::from_hex_text("test"),
|
||||||
|
Json(SetTokenText {
|
||||||
|
text: "Hello, world!".to_string(),
|
||||||
|
tokenizer_path: tokenizer_path.to_string_lossy().to_string(),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let response = assert_runtime_stays_free(&TOKENIZER_STORAGE_LOCK, request).await;
|
||||||
|
fs::remove_file(&tokenizer_path).unwrap();
|
||||||
|
assert!(!response.success);
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user