feat: 完成图片生成功能

2026-01-02 19:26:09 +08:00
parent 3f53eb14ab
commit 560a76558a
14 changed files with 251 additions and 134 deletions
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Jobs/ImageGenerationJob.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Jobs/ImageGenerationJob.cs
@@ -30,32 +30,93 @@ public class ImageGenerationJob : AsyncBackgroundJob<ImageGenerationJobArgs>, IT

    public override async Task ExecuteAsync(ImageGenerationJobArgs args)
    {
-        _logger.LogInformation("开始执行图片生成任务，TaskId: {TaskId}, ModelId: {ModelId}, UserId: {UserId}",
-            args.TaskId, args.ModelId, args.UserId);
+        var task = await _imageStoreTaskRepository.GetFirstAsync(x => x.Id == args.TaskId);
+        if (task is null)
+        {
+            throw new UserFriendlyException($"{args.TaskId} 图片生成任务不存在");
+        }

+        _logger.LogInformation("开始执行图片生成任务，TaskId: {TaskId}, ModelId: {ModelId}, UserId: {UserId}",
+            task.Id, task.ModelId, task.UserId);
        try
        {
-            var request = JsonSerializer.Deserialize<JsonElement>(args.RequestJson);
+            // 构建 Gemini API 请求对象
+            var parts = new List<object>
+            {
+                new { text = task.Prompt }
+            };

+            // 添加参考图（如果有）
+            foreach (var prefixBase64 in task.ReferenceImagesPrefixBase64)
+            {
+                var (mimeType, base64Data) = ParsePrefixBase64(prefixBase64);
+                parts.Add(new
+                {
+                    inline_data = new
+                    {
+                        mime_type = mimeType,
+                        data = base64Data
+                    }
+                });
+            }
+
+            var requestObj = new
+            {
+                contents = new[]
+                {
+                    new { parts }
+                }
+            };
+
+            var request = JsonSerializer.Deserialize<JsonElement>(
+                JsonSerializer.Serialize(requestObj));
+
+            //里面生成成功已经包含扣款了
            await _aiGateWayManager.GeminiGenerateContentImageForStatisticsAsync(
-                args.TaskId,
-                args.ModelId,
+                task.Id,
+                task.ModelId,
                request,
-                args.UserId);
+                task.UserId,
+                tokenId:task.TokenId);

+            
            _logger.LogInformation("图片生成任务完成，TaskId: {TaskId}", args.TaskId);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "图片生成任务失败，TaskId: {TaskId}, Error: {Error}", args.TaskId, ex.Message);

-            // 更新任务状态为失败
-            var task = await _imageStoreTaskRepository.GetFirstAsync(x => x.Id == args.TaskId);
-            if (task != null)
-            {
-                task.TaskStatus = TaskStatusEnum.Fail;
-                await _imageStoreTaskRepository.UpdateAsync(task);
-            }
+            task.TaskStatus = TaskStatusEnum.Fail;
+            task.ErrorInfo = ex.Message;
+
+            await _imageStoreTaskRepository.UpdateAsync(task);
        }
    }
-}
+
+    /// <summary>
+    /// 解析带前缀的 Base64 字符串，提取 mimeType 和纯 base64 数据
+    /// </summary>
+    private static (string mimeType, string base64Data) ParsePrefixBase64(string prefixBase64)
+    {
+        // 默认值
+        var mimeType = "image/png";
+        var base64Data = prefixBase64;
+
+        if (prefixBase64.Contains(","))
+        {
+            var parts = prefixBase64.Split(',');
+            if (parts.Length == 2)
+            {
+                var header = parts[0];
+                if (header.Contains(":") && header.Contains(";"))
+                {
+                    mimeType = header.Split(':')[1].Split(';')[0];
+                }
+
+                base64Data = parts[1];
+            }
+        }
+
+        return (mimeType, base64Data);
+    }
+}
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Jobs/ImageGenerationJobArgs.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Jobs/ImageGenerationJobArgs.cs
@@ -9,19 +9,4 @@ public class ImageGenerationJobArgs
    /// 图片任务ID
    /// </summary>
    public Guid TaskId { get; set; }
-
-    /// <summary>
-    /// 模型ID
-    /// </summary>
-    public string ModelId { get; set; } = string.Empty;
-
-    /// <summary>
-    /// 请求JSON字符串
-    /// </summary>
-    public string RequestJson { get; set; } = string.Empty;
-
-    /// <summary>
-    /// 用户ID
-    /// </summary>
-    public Guid UserId { get; set; }
 }
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Services/Chat/AiChatService.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Services/Chat/AiChatService.cs
@@ -86,7 +86,7 @@ public class AiChatService : ApplicationService
    }

    /// <summary>
-    /// 获取模型列表
+    /// 获取对话模型列表
    /// </summary>
    /// <returns></returns>
    public async Task<List<ModelGetListOutput>> GetModelAsync()
@@ -98,16 +98,9 @@ public class AiChatService : ApplicationService
            .Select(x => new ModelGetListOutput
            {
                Id = x.Id,
-                Category = "chat",
                ModelId = x.ModelId,
                ModelName = x.Name,
                ModelDescribe = x.Description,
-                ModelPrice = 0,
-                ModelType = "1",
-                ModelShow = "0",
-                SystemPrompt = null,
-                ApiHost = null,
-                ApiKey = null,
                Remark = x.Description,
                IsPremiumPackage = PremiumPackageConst.ModeIds.Contains(x.ModelId)
            }).ToListAsync();
@@ -202,7 +195,7 @@ public class AiChatService : ApplicationService
    [HttpPost("ai-chat/agent/send")]
    public async Task PostAgentSendAsync([FromBody] AgentSendInput input, CancellationToken cancellationToken)
    {
-        var tokenValidation = await _tokenManager.ValidateTokenAsync(input.Token, input.ModelId);
+        var tokenValidation = await _tokenManager.ValidateTokenAsync(input.TokenId, input.ModelId);

        await _aiBlacklistManager.VerifiyAiBlacklist(tokenValidation.UserId);
        // 验证用户是否为VIP
@@ -232,7 +225,7 @@ public class AiChatService : ApplicationService
        await _chatManager.AgentCompleteChatStreamAsync(_httpContextAccessor.HttpContext,
            input.SessionId,
            input.Content,
-            input.Token,
+            tokenValidation.Token,
            tokenValidation.TokenId,
            input.ModelId,
            tokenValidation.UserId,
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Services/Chat/AiImageService.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Application/Services/Chat/AiImageService.cs
@@ -31,14 +31,14 @@ public class AiImageService : ApplicationService
    private readonly PremiumPackageManager _premiumPackageManager;
    private readonly IGuidGenerator _guidGenerator;
    private readonly IWebHostEnvironment _webHostEnvironment;
-
+    private readonly TokenManager _tokenManager;
    public AiImageService(
        ISqlSugarRepository<ImageStoreTaskAggregateRoot> imageTaskRepository,
        IBackgroundJobManager backgroundJobManager,
        AiBlacklistManager aiBlacklistManager,
        PremiumPackageManager premiumPackageManager,
        IGuidGenerator guidGenerator,
-        IWebHostEnvironment webHostEnvironment)
+        IWebHostEnvironment webHostEnvironment, TokenManager tokenManager)
    {
        _imageTaskRepository = imageTaskRepository;
        _backgroundJobManager = backgroundJobManager;
@@ -46,6 +46,7 @@ public class AiImageService : ApplicationService
        _premiumPackageManager = premiumPackageManager;
        _guidGenerator = guidGenerator;
        _webHostEnvironment = webHostEnvironment;
+        _tokenManager = tokenManager;
    }

    /// <summary>
@@ -58,10 +59,17 @@ public class AiImageService : ApplicationService
    public async Task<Guid> GenerateAsync([FromBody] ImageGenerationInput input)
    {
        var userId = CurrentUser.GetId();
-
+        
        // 黑名单校验
        await _aiBlacklistManager.VerifiyAiBlacklist(userId);

+        //校验token
+        if (input.TokenId is not null)
+        {
+            await _tokenManager.ValidateTokenAsync(input.TokenId, input.ModelId);
+        }
+
+        
        // VIP校验
        if (!CurrentUser.IsAiVip())
        {
@@ -82,32 +90,21 @@ public class AiImageService : ApplicationService
        var task = new ImageStoreTaskAggregateRoot
        {
            Prompt = input.Prompt,
-            ReferenceImagesBase64 = input.ReferenceImagesBase64 ?? new List<string>(),
+            ReferenceImagesPrefixBase64 = input.ReferenceImagesPrefixBase64 ?? new List<string>(),
            ReferenceImagesUrl = new List<string>(),
            TaskStatus = TaskStatusEnum.Processing,
-            UserId = userId
+            UserId = userId,
+            TokenId = input.TokenId
        };

        await _imageTaskRepository.InsertAsync(task);
-        var taskId = task.Id;
-
-        // 构建请求JSON
-        var requestJson = JsonSerializer.Serialize(new
-        {
-            prompt = input.Prompt,
-            referenceImages = input.ReferenceImagesBase64
-        });
-
        // 入队后台任务
        await _backgroundJobManager.EnqueueAsync(new ImageGenerationJobArgs
        {
-            TaskId = taskId,
-            ModelId = input.ModelId,
-            RequestJson = requestJson,
-            UserId = userId
+            TaskId = task.Id,
        });

-        return taskId;
+        return task.Id;
    }

    /// <summary>
@@ -130,9 +127,9 @@ public class AiImageService : ApplicationService
        {
            Id = task.Id,
            Prompt = task.Prompt,
-            ReferenceImagesBase64 = task.ReferenceImagesBase64,
-            ReferenceImagesUrl = task.ReferenceImagesUrl,
-            StoreBase64 = task.StoreBase64,
+            // ReferenceImagesBase64 = task.ReferenceImagesBase64,
+            // ReferenceImagesUrl = task.ReferenceImagesUrl,
+            // StoreBase64 = task.StoreBase64,
            StoreUrl = task.StoreUrl,
            TaskStatus = task.TaskStatus,
            CreationTime = task.CreationTime
@@ -234,9 +231,9 @@ public class AiImageService : ApplicationService
            {
                Id = x.Id,
                Prompt = x.Prompt,
-                ReferenceImagesBase64 = x.ReferenceImagesBase64,
-                ReferenceImagesUrl = x.ReferenceImagesUrl,
-                StoreBase64 = x.StoreBase64,
+                // ReferenceImagesBase64 = x.ReferenceImagesBase64,
+                // ReferenceImagesUrl = x.ReferenceImagesUrl,
+                // StoreBase64 = x.StoreBase64,
                StoreUrl = x.StoreUrl,
                TaskStatus = x.TaskStatus,
                CreationTime = x.CreationTime