feat: 完成用量统计功能模块

Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/AiChat/Impl/AzureRestChatService.cs

@@ -14,7 +14,8 @@ public class AzureRestChatService : IChatService
     {
     }
 
-    public async IAsyncEnumerable<string> CompleteChatAsync(AiModelDescribe aiModelDescribe, List<ChatMessage> messages,
+    public async IAsyncEnumerable<CompleteChatResponse> CompleteChatAsync(AiModelDescribe aiModelDescribe,
+        List<ChatMessage> messages,
         [EnumeratorCancellation] CancellationToken cancellationToken)
     {
         // 设置API URL
@@ -61,51 +62,65 @@ public class AzureRestChatService : IChatService
         var responseStream = await response.Content.ReadAsStreamAsync(cancellationToken);
         // 从流中读取数据并输出到控制台
         using var streamReader = new StreamReader(responseStream);
-        string line;
-        while ((line = await streamReader.ReadLineAsync(cancellationToken)) != null)
+        while (await streamReader.ReadLineAsync(cancellationToken) is { } line)
         {
-            var result = GetContent(line);
-            if (result is not null)
+            var result = new CompleteChatResponse();
+            try
             {
-                yield return result;
+                var jsonObj = MapToJObject(line);
+                var content = GetContent(jsonObj);
+                var tokenUsage = GetTokenUsage(jsonObj);
+                result= new CompleteChatResponse
+                {
+                    TokenUsage = tokenUsage,
+                    IsFinish = tokenUsage is not null,
+                    Content = content
+                };
             }
+            catch (Exception e)
+            {
+                Console.WriteLine("解析失败");
+            }
+
+            yield return result;
         }
     }
 
-    private string? GetContent(string line)
+    private JObject? MapToJObject(string line)
     {
         if (string.IsNullOrWhiteSpace(line))
             return null;
         string prefix = "data: ";
         line = line.Substring(prefix.Length);
+        return JObject.Parse(line);
+    }
 
+    private string? GetContent(JObject? jsonObj)
+    {
+        var contentToken = jsonObj.SelectToken("choices[0].delta.content");
+        if (contentToken != null && contentToken.Type != JTokenType.Null)
+        {
+            return contentToken.ToString();
+        }
 
-        try
+        return null;
+    }
+
+    private TokenUsage? GetTokenUsage(JObject? jsonObj)
+    {
+        var usage = jsonObj.SelectToken("usage");
+        if (usage is not null && usage.Type != JTokenType.Null)
         {
-            // 解析为JObject
-            var jsonObj = JObject.Parse(line);
-            var content = jsonObj["choices"][0]["delta"]["content"].ToString();
-            return content;
-            // // 判断choices是否存在且是数组，并且有元素
-            // if (jsonObj.TryGetValue("choices", out var choicesToken) && choicesToken is JArray choicesArray &&
-            //     choicesArray.Count > 0)
-            // {
-            //     var firstChoice = choicesArray[0] as JObject;
-            //     // 判断delta字段是否存在
-            //     if (firstChoice.TryGetValue("delta", out var deltaToken))
-            //     {
-            //         // 获取content字段
-            //         if (deltaToken.Type == JTokenType.Object && ((JObject)deltaToken).TryGetValue("content", out var contentToken))
-            //         {
-            //             return contentToken.ToString();
-            //         }
-            //     }
-            // }
-        }
-        catch (Exception)
-        {
-            // 解析失败
-            return null;
+            var result = new TokenUsage()
+            {
+                OutputTokenCount = usage["completion_tokens"].ToObject<int>(),
+                InputTokenCount = usage["prompt_tokens"].ToObject<int>(),
+                TotalTokenCount = usage["total_tokens"].ToObject<int>()
+            };
+
+            return result;
         }
+
+        return null;
     }
 }