feat: add groq chat stream (#88)

tryAGI · Oct 25, 2024 · 847afee · 847afee
1 parent e6c1162
commit 847afee
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/src/Google.VertexAI/src/VertexAIChatModel.cs b/src/Google.VertexAI/src/VertexAIChatModel.cs
@@ -19,8 +19,7 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
         ChatSettings? settings = null,
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
-            if (request == null)
-                throw new ArgumentNullException(nameof(request));
+            request = request ?? throw new ArgumentNullException(nameof(request));
 
             var prompt = ToPrompt(request.Messages);
             var watch = Stopwatch.StartNew();

diff --git a/src/Groq/src/GroqChatModel.cs b/src/Groq/src/GroqChatModel.cs
@@ -1,6 +1,6 @@
-using System.Diagnostics;
+using GroqSharp.Models;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
-using GroqSharp.Models;
 using MessageGroq = GroqSharp.Models.Message;
 
 namespace LangChain.Providers.Groq;
@@ -21,23 +21,39 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
         var prompt = ToPrompt(request.Messages);
         Provider.Api.SetModel(Id);
         var watch = Stopwatch.StartNew();
-        var response = await Provider.Api.CreateChatCompletionAsync(prompt).ConfigureAwait(false);
 
+        if (settings?.UseStreaming == true)
+        {
+            var responseStream = Provider.Api.CreateChatCompletionStreamAsync(prompt);
+            await foreach (var response in responseStream.WithCancellation(cancellationToken))
+            {
+                yield return BuildChatResponse(response, request, settings, watch);
+            }
+        }
+        else
+        {
+            var response = await Provider.Api.CreateChatCompletionAsync(prompt).ConfigureAwait(false);
+            yield return BuildChatResponse(response, request, settings, watch);
+        }
+
+    }
+
+    private ChatResponse BuildChatResponse(string response, ChatRequest request, ChatSettings? settings, Stopwatch watch)
+    {
         var usage = Usage.Empty with
         {
             Time = watch.Elapsed,
         };
-        AddUsage(usage);
-        provider.AddUsage(usage);
 
         var result = request.Messages.ToList();
+        provider.AddUsage(usage);
         result.Add(response.AsAiMessage());
 
-        yield return new ChatResponse
+        return new ChatResponse
         {
             Messages = result,
             Usage = usage,
-            UsedSettings = ChatSettings.Default,
+            UsedSettings = settings ?? ChatSettings.Default,
         };
     }