From 067a5365d28b4733375943c5d4f2a9d6ff2e0835 Mon Sep 17 00:00:00 2001
From: femsci <nya@femboy.science>
Date: Fri, 17 Nov 2023 23:38:35 +0100
Subject: [PATCH] Refactor Server Sent Events parsing

---
 src/Nyanlabs.Umogen.Core/Models/ApiChoice.cs  |  5 ---
 .../Models/ApiResponse.cs                     |  5 ---
 src/Nyanlabs.Umogen.Core/Models/UsageStats.cs |  4 --
 src/Nyanlabs.Umogen.Core/SseConsumer.cs       | 45 +++++++++++++++++++
 src/Nyanlabs.Umogen.Core/UmoProcess.cs        | 27 +++--------
 src/Nyanlabs.Umogen.Core/Umogen.cs            |  2 +-
 6 files changed, 51 insertions(+), 37 deletions(-)
 create mode 100644 src/Nyanlabs.Umogen.Core/SseConsumer.cs

diff --git a/src/Nyanlabs.Umogen.Core/Models/ApiChoice.cs b/src/Nyanlabs.Umogen.Core/Models/ApiChoice.cs
index 68462c9..6fe57f6 100644
--- a/src/Nyanlabs.Umogen.Core/Models/ApiChoice.cs
+++ b/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
diff --git a/src/Nyanlabs.Umogen.Core/Models/ApiResponse.cs b/src/Nyanlabs.Umogen.Core/Models/ApiResponse.cs
index 8421f3b..fca01a0 100644
--- a/src/Nyanlabs.Umogen.Core/Models/ApiResponse.cs
+++ b/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
diff --git a/src/Nyanlabs.Umogen.Core/Models/UsageStats.cs b/src/Nyanlabs.Umogen.Core/Models/UsageStats.cs
index 45077e6..2ab1b60 100644
--- a/src/Nyanlabs.Umogen.Core/Models/UsageStats.cs
+++ b/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;
 
diff --git a/src/Nyanlabs.Umogen.Core/SseConsumer.cs b/src/Nyanlabs.Umogen.Core/SseConsumer.cs
new file mode 100644
index 0000000..80c7558
--- /dev/null
+++ b/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();
+    }
+}
diff --git a/src/Nyanlabs.Umogen.Core/UmoProcess.cs b/src/Nyanlabs.Umogen.Core/UmoProcess.cs
index 9257e31..875a61a 100644
--- a/src/Nyanlabs.Umogen.Core/UmoProcess.cs
+++ b/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);
-
-        string? line;
-        while ((line = await reader.ReadLineAsync()) != null)
+        using var sseConsumer = new SseConsumer<ApiResponse>(apiStream);
+        await foreach (var chunk in sseConsumer.ReadEvents())
         {
-            if (line.StartsWith("data:"))
-            {
-                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();
-
+            ApiChoice choice = chunk.Choices.First();
             if (choice.FinishReason == "stop")
             {
                 break;
             }
 
-            yield return choice.Delta?.Content ?? ":<";
+            yield return choice.Delta?.Content ?? throw new DataException("Message delta is null");
         }
     }
 }
diff --git a/src/Nyanlabs.Umogen.Core/Umogen.cs b/src/Nyanlabs.Umogen.Core/Umogen.cs
index b9781e6..6331e4f 100644
--- a/src/Nyanlabs.Umogen.Core/Umogen.cs
+++ b/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