diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs index 6d63f022..53136d08 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs +++ b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs @@ -8,12 +8,22 @@ public sealed partial class RustService public async Task ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false) { var streamId = Guid.NewGuid().ToString(); - var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}"; + var extractImagesQueryValue = extractImages ? "true" : "false"; + var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImagesQueryValue}"; var request = new HttpRequestMessage(HttpMethod.Get, requestUri); var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead); if (!response.IsSuccessStatusCode) + { + var responseBody = await response.Content.ReadAsStringAsync(); + this.logger?.LogError( + "Failed to read arbitrary file data from Rust runtime. Status: {StatusCode}, reason: '{ReasonPhrase}', path: '{Path}', body: '{Body}'", + response.StatusCode, + response.ReasonPhrase, + path, + responseBody); return string.Empty; + } var resultBuilder = new StringBuilder(); diff --git a/runtime/src/file_data.rs b/runtime/src/file_data.rs index ee48af13..d41a9daa 100644 --- a/runtime/src/file_data.rs +++ b/runtime/src/file_data.rs @@ -5,6 +5,7 @@ use crate::pandoc::PandocProcessBuilder; use crate::pdfium::PdfiumInit; use async_stream::stream; use axum::extract::Query; +use axum::extract::rejection::QueryRejection; use axum::response::sse::{Event, Sse}; use base64::{engine::general_purpose, Engine as _}; use calamine::{open_workbook_auto, Reader}; @@ -15,7 +16,7 @@ use pptx_to_md::{ImageHandlingMode, ParserConfig, PptxContainer}; use serde::{Deserialize, Serialize}; use std::path::Path; use std::pin::Pin; -use log::{debug, error}; +use log::{debug, error, warn}; use tokio::io::AsyncBufReadExt; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -90,32 +91,49 @@ pub struct ExtractDataQuery { pub async fn extract_data( _token: APIToken, - Query(query): Query, + query: std::result::Result, QueryRejection>, ) -> Sse>> { + let query = match query { + Ok(Query(query)) => Ok(query), + Err(e) => { + let message = format!("Invalid query for '/retrieval/fs/extract': {e}"); + warn!("{message}"); + Err(message) + }, + }; + let stream = stream! { - let stream_result = stream_data(&query.path, query.extract_images).await; - let id_ref = &query.stream_id; + match query { + Ok(query) => { + let stream_result = stream_data(&query.path, query.extract_images).await; + let id_ref = &query.stream_id; - match stream_result { - Ok(mut stream) => { - while let Some(chunk) = stream.next().await { - match chunk { - Ok(mut chunk) => { - chunk.set_stream_id(id_ref); - yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}")))); - }, + match stream_result { + Ok(mut stream) => { + while let Some(chunk) = stream.next().await { + match chunk { + Ok(mut chunk) => { + chunk.set_stream_id(id_ref); + yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}")))); + }, - Err(e) => { - yield Ok(Event::default().json_data(format!("Error: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error: {e}")))); - break; - }, + Err(e) => { + yield Ok(Event::default().json_data(format!("Error: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error: {e}")))); + break; + }, + } + } + }, + + Err(e) => { + yield Ok(Event::default().json_data(format!("Error starting stream: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error starting stream: {e}")))); } - } + }; }, Err(e) => { yield Ok(Event::default().json_data(format!("Error starting stream: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error starting stream: {e}")))); - } + }, } };