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

using System.Runtime.CompilerServices;
using System.Text;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using OpenAI.Chat;
using Yi.Framework.AiHub.Domain.Extensions;
using Yi.Framework.AiHub.Domain.Shared.Dtos;

namespace Yi.Framework.AiHub.Domain.AiChat.Impl;

public class AzureRestChatService : IChatService
{
    public AzureRestChatService()
    {
    }

    public async IAsyncEnumerable<CompleteChatResponse> CompleteChatAsync(AiModelDescribe aiModelDescribe,
        List<ChatMessage> messages,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        // 设置API URL
        var apiUrl = $"{aiModelDescribe.Endpoint}";

        // 准备请求内容
        var requestBody = new
        {
            messages = messages.Select(x => new
            {
                role = x.GetRoleAsString(),
                content = x.Content.FirstOrDefault()?.Text
            }).ToList(),
            stream = true,
            // max_tokens = 2048,
            // temperature = 0.8,
            // top_p = 0.1,
            // presence_penalty = 0,
            // frequency_penalty = 0,
            model = aiModelDescribe.ModelId
        };

        // 序列化请求内容为JSON
        string jsonBody = JsonConvert.SerializeObject(requestBody);

        using var httpClient = new HttpClient()
        {
            //10分钟超时
            Timeout = TimeSpan.FromSeconds(600)
        };
        // 设置请求头
        httpClient.DefaultRequestHeaders.Add("Authorization", $"Bearer {aiModelDescribe.ApiKey}");
        // 其他头信息如Content-Type在StringContent中设置

        // 构造 POST 请求
        var request = new HttpRequestMessage(HttpMethod.Post, apiUrl);

        // 设置请求内容（示例）
        request.Content = new StringContent(jsonBody, Encoding.UTF8, "application/json");

        // 发送POST请求
        HttpResponseMessage response =
            await httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);

        if (!response.IsSuccessStatusCode)
        {
            throw new UserFriendlyException($"当前模型不可用：{aiModelDescribe.ModelId}，状态码:{response.StatusCode},原因：{response.ReasonPhrase}");
        }
        // 确认响应成功
        // response.EnsureSuccessStatusCode();

        // 读取响应内容
        var responseStream = await response.Content.ReadAsStreamAsync(cancellationToken);
        // 从流中读取数据并输出到控制台
        using var streamReader = new StreamReader(responseStream);
        while (await streamReader.ReadLineAsync(cancellationToken) is { } line)
        {
            var result = new CompleteChatResponse();
            try
            {
                var jsonObj = MapToJObject(line);
                if (jsonObj is not null)
                {
                    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 JObject? MapToJObject(string line)
    {
        if (line == "data: [DONE]"||string.IsNullOrWhiteSpace(line) )
        {
            return null;
        }

        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();
        }

        return null;
    }

    private TokenUsage? GetTokenUsage(JObject? jsonObj)
    {
        var usage = jsonObj.SelectToken("usage");
        if (usage is not null && usage.Type != JTokenType.Null)
        {
            var result = new TokenUsage();
            var completionTokens = usage["completion_tokens"];
            if (completionTokens is not null && completionTokens.Type != JTokenType.Null)
            {
                result.OutputTokenCount = completionTokens.ToObject<int>();
            }

            var promptTokens = usage["prompt_tokens"];
            if (promptTokens is not null && promptTokens.Type != JTokenType.Null)
            {
                result.InputTokenCount = promptTokens.ToObject<int>();
            }

            var totalTokens = usage["total_tokens"];
            if (totalTokens is not null && totalTokens.Type != JTokenType.Null)
            {
                result.TotalTokenCount = totalTokens.ToObject<int>();
            }

            return result;
        }

        return null;
    }
}