background embed

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

@@ -1,5 +1,6 @@
 using System.Text;
 using System.Text.Json;
+using System.Runtime.CompilerServices;
 
 namespace AIStudio.Tools.Services;
 
@@ -48,6 +49,9 @@ public sealed partial class RustService
                 }
                 catch (JsonException)
                 {
+                    if (this.TryLogSseErrorMessage(jsonContent, path))
+                        continue;
+
                     this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
                 }
             }
@@ -65,4 +69,77 @@ public sealed partial class RustService
         
         return resultBuilder.ToString();
     }
-}
+
+    public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
+    {
+        var streamId = Guid.NewGuid().ToString();
+        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}";
+        using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
+        using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
+
+        if (!response.IsSuccessStatusCode)
+            yield break;
+
+        string? finalContentChunk = null;
+        try
+        {
+            await using var stream = await response.Content.ReadAsStreamAsync(token);
+            using var reader = new StreamReader(stream);
+
+            while (!reader.EndOfStream && !token.IsCancellationRequested)
+            {
+                var line = await reader.ReadLineAsync(token);
+                if (string.IsNullOrWhiteSpace(line))
+                    continue;
+
+                if (!line.StartsWith("data:", StringComparison.InvariantCulture))
+                    continue;
+
+                var jsonContent = line[5..];
+                ContentStreamSseEvent? sseEvent = null;
+                try
+                {
+                    sseEvent = JsonSerializer.Deserialize<ContentStreamSseEvent>(jsonContent);
+                }
+                catch (JsonException)
+                {
+                    if (this.TryLogSseErrorMessage(jsonContent, path))
+                        continue;
+
+                    this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
+                }
+
+                if (sseEvent is null)
+                    continue;
+
+                var content = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages);
+                if (!string.IsNullOrWhiteSpace(content))
+                    yield return content;
+            }
+        }
+        finally
+        {
+            finalContentChunk = ContentStreamSseHandler.Clear(streamId);
+        }
+
+        if (!string.IsNullOrWhiteSpace(finalContentChunk))
+            yield return finalContentChunk;
+    }
+
+    private bool TryLogSseErrorMessage(string jsonContent, string path)
+    {
+        try
+        {
+            var errorMessage = JsonSerializer.Deserialize<string>(jsonContent);
+            if (string.IsNullOrWhiteSpace(errorMessage))
+                return false;
+
+            this.logger?.LogError("Rust retrieval stream error for '{Path}': {ErrorMessage}", path, errorMessage);
+            return true;
+        }
+        catch (JsonException)
+        {
+            return false;
+        }
+    }
+}