Refactor Server Sent Events parsing

src/Nyanlabs.Umogen.Core/Models/ApiChoice.cs

@@ -1,8 +1,3 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
-
 namespace Nyanlabs.Umogen.Core.Models;
 
 public class ApiChoice

src/Nyanlabs.Umogen.Core/Models/ApiResponse.cs

@@ -1,8 +1,3 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
-
 namespace Nyanlabs.Umogen.Core.Models;
 
 public class ApiResponse

src/Nyanlabs.Umogen.Core/Models/UsageStats.cs

@@ -1,8 +1,4 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
 using System.Text.Json.Serialization;
-using System.Threading.Tasks;
 
 namespace Nyanlabs.Umogen.Core.Models;
 

src/Nyanlabs.Umogen.Core/SseConsumer.cs

@@ -0,0 +1,45 @@
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Linq;
+using System.Text.Json;
+using System.Threading.Tasks;
+
+namespace Nyanlabs.Umogen.Core;
+
+public class SseConsumer<TEvent> : IDisposable, IAsyncDisposable
+{
+    public SseConsumer(Stream stream)
+    {
+        _stream = stream;
+    }
+
+    private readonly Stream _stream;
+
+    public async IAsyncEnumerable<TEvent> ReadEvents()
+    {
+        var reader = new StreamReader(_stream);
+
+        while (!reader.EndOfStream)
+        {
+            string line = (await reader.ReadLineAsync())!;
+            //Second \n
+            await reader.ReadLineAsync();
+
+            //Remove 'data: '
+            line = line[6..].Trim();
+
+            yield return JsonSerializer.Deserialize<TEvent>(line, Umogen.JSON_OPTS) ?? throw new DataException("Invalid data...");
+        }
+    }
+
+    public void Dispose()
+    {
+        _stream.Dispose();
+    }
+
+    public ValueTask DisposeAsync()
+    {
+        return _stream.DisposeAsync();
+    }
+}

src/Nyanlabs.Umogen.Core/UmoProcess.cs

@@ -1,5 +1,4 @@
 using System.Data;
-using System.Net.Http.Json;
 using System.Text;
 using System.Text.Json;
 
@@ -38,34 +37,18 @@ public class UmoProcess
             throw new DataException($"Cannot process request: {resp.StatusCode}.");
         }
 
-        var apiResp = await resp.Content.ReadAsStreamAsync();
+        var apiStream = await resp.Content.ReadAsStreamAsync();
 
-        var reader = new StreamReader(apiResp);
+        using var sseConsumer = new SseConsumer<ApiResponse>(apiStream);
-
+        await foreach (var chunk in sseConsumer.ReadEvents())
-        string? line;
-        while ((line = await reader.ReadLineAsync()) != null)
         {
-            if (line.StartsWith("data:"))
+            ApiChoice choice = chunk.Choices.First();
-            {
-                line = line[5..].Trim();
-            }
-
-            if (string.IsNullOrWhiteSpace(line))
-            {
-                continue;
-            }
-
-            File.AppendAllText("/tmp/raw.md", line + "\n");
-
-            ApiResponse respChunk = JsonSerializer.Deserialize<ApiResponse>(line, options: Umogen.JSON_OPTS)!;
-            ApiChoice choice = respChunk.Choices.First();
-
             if (choice.FinishReason == "stop")
             {
                 break;
             }
 
-            yield return choice.Delta?.Content ?? ":<";
+            yield return choice.Delta?.Content ?? throw new DataException("Message delta is null");
         }
     }
 }

src/Nyanlabs.Umogen.Core/Umogen.cs

@@ -5,7 +5,7 @@ namespace Nyanlabs.Umogen.Core;
 public class Umogen
 {
     public const string DEFAULT_API_KEY_FILE = "umogenkey.secret";
-    public const string PROMPT = "You write Polish legal contracts in Polish language in markdown format.";
+    public const string PROMPT = "You write Polish legal contracts in Polish language in markdown format. Fill unknown fields with signature placeholder (underscore) with enough space to sign it with pen.";
     public static readonly JsonSerializerOptions JSON_OPTS = new(JsonSerializerDefaults.Web)
     {
         PropertyNamingPolicy = JsonNamingPolicy.SnakeCaseLower