添加Gemini生图功能到article_auto_image_matching.py - 未匹配成功时自动生成图片

article_auto_image_matching.py

@@ -364,9 +364,305 @@ class ArticleImageMatcher:
             self.log_to_database('ERROR', error_msg, traceback.format_exc())
             return False
     
+    def generate_image_with_gemini(self, prompt: str, article_tags: List[str], article_id: int) -> Optional[str]:
+        """
+        使用Gemini生成图片并上传到服务器
+        
+        Args:
+            prompt: 图片生成提示词
+            article_tags: 文章标签列表，用于查询department和keywords
+            article_id: 文章ID，用于关联图片
+        
+        Returns:
+            上传后的图片URL，失败返回None
+        """
+        try:
+            # 导入必要的库
+            from google import genai
+            from google.genai.client import HttpOptions
+
+            client = genai.Client(
+                http_options=HttpOptions(base_url="https://work.poloapi.com"),
+                api_key="sk-V4tPnDgzFPa7nxWrvKnNJsW8ZcBXXPuGmjfgvPVRnwpHoeob"
+            )
+            
+            logger.info(f"正在调用Gemini API生成图片，提示词: {prompt[:50]}...")
+            
+            # 生成内容
+            response = client.models.generate_content(
+                model="gemini-3-pro-image-preview",
+                contents=[prompt],
+            )
+            
+            # 检查是否有候选答案
+            if not response.candidates:
+                raise Exception("Gemini API未返回任何候选答案")
+            
+            # 处理响应
+            candidate = response.candidates[0]
+            if not candidate.content or not candidate.content.parts:
+                raise Exception("Gemini API返回的候选答案中没有内容部分")
+            
+            for part in candidate.content.parts:
+                if hasattr(part, 'inline_data') and part.inline_data is not None:
+                    image_data = part.inline_data
+                    if image_data.data is not None:
+                        # 生成唯一的文件名（基于时间戳）
+                        timestamp_ms = int(time.time() * 1000)
+                        image_filename = f"{timestamp_ms}.png"
+                        today_date = datetime.now().strftime("%Y%m%d")
+                        image_url_path = f"{today_date}/{image_filename}"
+                        
+                        temp_filename = f"temp_generated_image_{timestamp_ms}.png"
+                        # 保存图片数据到临时文件
+                        with open(temp_filename, 'wb') as f:
+                            f.write(image_data.data)
+                        logger.info(f"Gemini生成图片成功: {temp_filename}")
+                        
+                        # 先将图片信息插入数据库
+                        image_info = self.insert_generated_image_to_db(image_filename, image_url_path, article_tags)
+                        
+                        if not image_info:
+                            raise Exception("插入图片信息到数据库失败")
+                        
+                        logger.info(f"图片信息已插入数据库，tag_image_id: {image_info['tag_image_id']}, image_id: {image_info['image_id']}")
+                        
+                        # 使用tag_image_id上传图片到服务器
+                        uploaded_url = self.upload_image_to_server(temp_filename, image_info['tag_image_id'])
+                        
+                        # 将文章与图片的关联信息插入ai_article_images表
+                        article_image_id = self.insert_article_image_relation_for_generated(
+                            article_id=article_id,
+                            image_id=image_info['image_id'],
+                            image_url=image_info['image_url'],
+                            image_thumb_url=image_info['image_thumb_url'],
+                            tag_image_id=image_info['tag_image_id'],
+                            keywords_id=image_info['keywords_id'],
+                            keywords_name=image_info['keywords_name'],
+                            department_id=image_info['department_id'],
+                            department_name=image_info['department_name'],
+                            image_source=0  # 默认值
+                        )
+                        
+                        if article_image_id:
+                            logger.info(f"文章图片关联信息已创建，ai_article_images.id: {article_image_id}")
+                        
+                        # 删除临时文件
+                        os.remove(temp_filename)
+                        
+                        logger.info(f"图片已上传到服务器: {uploaded_url}")
+                        return uploaded_url
+            
+            raise Exception("Gemini API未返回有效的图片数据")
+            
+        except ImportError:
+            logger.error("错误：未安装google-genai库，请运行 'pip install google-genai' 进行安装")
+            return None
+        except Exception as e:
+            error_msg = f"Gemini生成图片异常: {e}"
+            logger.error(error_msg)
+            self.log_to_database('ERROR', error_msg, traceback.format_exc())
+            return None
+    
+    def insert_generated_image_to_db(self, image_name: str, image_url: str, article_tags: List[str]) -> Optional[Dict]:
+        """
+        将Gemini生成的图片信息插入数据库
+        
+        Args:
+            image_name: 图片文件名
+            image_url: 图片URL路径
+            article_tags: 文章标签列表
+        
+        Returns:
+            包含插入信息的字典
+        """
+        try:
+            connection = self.db_manager.get_connection()
+            try:
+                with connection.cursor(pymysql.cursors.DictCursor) as cursor:
+                    # 根据文章标签查询ai_image_tags表
+                    if article_tags:
+                        query = """
+                        SELECT department_name, keywords_name, department_id, keywords_id, tag_id
+                        FROM ai_image_tags 
+                        WHERE tag_name = %s 
+                        LIMIT 1
+                        """
+                        cursor.execute(query, (article_tags[0],))
+                        tag_info = cursor.fetchone()
+                        
+                        if tag_info:
+                            department = tag_info['department_name']
+                            keywords = tag_info['keywords_name']
+                            department_id = tag_info['department_id']
+                            keywords_id = tag_info['keywords_id']
+                            tag_id = tag_info['tag_id']
+                            tag_name = article_tags[0]
+                        else:
+                            department = "AI生成"
+                            keywords = "AI图片"
+                            department_id = 1
+                            keywords_id = 1
+                            tag_id = 1
+                            tag_name = article_tags[0] if article_tags else "AI生成"
+                    else:
+                        department = "AI生成"
+                        keywords = "AI图片"
+                        department_id = 1
+                        keywords_id = 1
+                        tag_id = 1
+                        tag_name = "AI生成"
+                    
+                    # 插入ai_images表
+                    insert_image_query = """
+                    INSERT INTO ai_images 
+                    (image_name, image_url, image_thumb_url, department, keywords, image_type, upload_user_id, status)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+                    """
+                    cursor.execute(insert_image_query, (
+                        image_name, image_url, '', department, keywords, 'medical', 1, 'active'
+                    ))
+                    image_id = cursor.lastrowid
+                    logger.info(f"图片信息已插入ai_images表，image_id: {image_id}")
+                    
+                    # 插入ai_image_tags表
+                    insert_tag_query = """
+                    INSERT INTO ai_image_tags 
+                    (image_id, image_name, image_url, image_thumb_url, tag_id, tag_name, 
+                     keywords_id, keywords_name, department_id, department_name, 
+                     image_source, created_user_id, image_attached_article_count)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    """
+                    cursor.execute(insert_tag_query, (
+                        image_id, image_name, image_url, '', tag_id, tag_name,
+                        keywords_id, keywords, department_id, department,
+                        3, 1, 0  # image_source: 3表示AI生成
+                    ))
+                    tag_image_id = cursor.lastrowid
+                    logger.info(f"图片标签信息已插入ai_image_tags表，tag_image_id: {tag_image_id}")
+                    
+                    connection.commit()
+                    
+                    return {
+                        'tag_image_id': tag_image_id,
+                        'image_id': image_id,
+                        'image_url': image_url,
+                        'image_thumb_url': '',
+                        'keywords_id': keywords_id,
+                        'keywords_name': keywords,
+                        'department_id': department_id,
+                        'department_name': department
+                    }
+            finally:
+                connection.close()
+        except Exception as e:
+            logger.error(f"插入图片信息到数据库失败: {e}")
+            return None
+    
+    def upload_image_to_server(self, image_path: str, tag_image_id: int) -> str:
+        """
+        上传图片到服务器
+        
+        Args:
+            image_path: 本地图片路径
+            tag_image_id: 图片标签ID
+        
+        Returns:
+            服务器上的图片URL
+        """
+        base_url = "http://47.99.184.230:8324"
+        jwt_token = self.login_and_get_jwt_token(base_url)
+        
+        if not jwt_token:
+            raise Exception("获取JWT token失败，无法上传图片")
+        
+        upload_url = f"{base_url}/api/images/upload"
+        headers = {'Authorization': f'Bearer {jwt_token}'}
+        
+        with open(image_path, 'rb') as image_file:
+            files = {'file': image_file}
+            data = {'tag_image_id': tag_image_id}
+            
+            response = requests.post(upload_url, headers=headers, files=files, data=data)
+            
+            logger.info(f"图片上传响应状态码: {response.status_code}")
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result.get('code') == 200:
+                    return result['data']['http_image_url']
+                else:
+                    raise Exception(f"图片上传失败: {result.get('message', '未知错误')}")
+            else:
+                raise Exception(f"图片上传请求失败，状态码: {response.status_code}")
+    
+    def login_and_get_jwt_token(self, base_url: str) -> Optional[str]:
+        """登录获取JWT token"""
+        login_url = f"{base_url}/api/auth/login"
+        login_data = {"username": "user010", "password": "@5^2W6R7"}
+        
+        try:
+            response = requests.post(login_url, json=login_data, headers={'Content-Type': 'application/json'})
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result.get('code') == 200:
+                    logger.info("JWT token获取成功")
+                    return result['data']['token']
+            
+            logger.error(f"登录失败: {response.status_code}")
+            return None
+        except Exception as e:
+            logger.error(f"登录异常: {e}")
+            return None
+    
+    def insert_article_image_relation_for_generated(self, article_id: int, image_id: int, image_url: str, 
+                                                     image_thumb_url: str, tag_image_id: int, keywords_id: int, 
+                                                     keywords_name: str, department_id: int, department_name: str, 
+                                                     image_source: int = 0) -> Optional[int]:
+        """
+        将文章与生成图片的关联信息插入ai_article_images表
+        """
+        try:
+            connection = self.db_manager.get_connection()
+            try:
+                with connection.cursor(pymysql.cursors.DictCursor) as cursor:
+                    # 查询当前文章下已有图片的最大sort_order
+                    query_max_sort = """
+                    SELECT COALESCE(MAX(sort_order), 0) as max_sort_order
+                    FROM ai_article_images
+                    WHERE article_id = %s
+                    """
+                    cursor.execute(query_max_sort, (article_id,))
+                    result = cursor.fetchone()
+                    max_sort_order = result['max_sort_order'] if result else 0
+                    new_sort_order = max_sort_order + 1
+                    
+                    # 插入ai_article_images表
+                    insert_query = """
+                    INSERT INTO ai_article_images 
+                    (article_id, image_id, image_url, image_thumb_url, image_tag_id, sort_order,
+                     keywords_id, keywords_name, department_id, department_name, image_source)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    """
+                    cursor.execute(insert_query, (
+                        article_id, image_id, image_url, image_thumb_url, tag_image_id, new_sort_order,
+                        keywords_id, keywords_name, department_id, department_name, image_source
+                    ))
+                    article_image_id = cursor.lastrowid
+                    logger.info(f"文章图片关联信息已插入ai_article_images表，id: {article_image_id}")
+                    
+                    connection.commit()
+                    return article_image_id
+            finally:
+                connection.close()
+        except Exception as e:
+            logger.error(f"插入文章图片关联信息失败: {e}")
+            return None
+    
     def match_article_with_images(self, article_data: Dict, available_images: List[Dict]) -> bool:
         """
-        为单篇文章匹配图片
+        为单篇文章匹配图片，未成功匹配时调用Gemini生图
         
         Args:
             article_data: 文章数据
@@ -377,6 +673,7 @@ class ArticleImageMatcher:
         """
         article_id = article_data['article_id']
         article_title = article_data.get('title', '')
+        article_content = article_data.get('content', '')
         coze_tag = article_data.get('coze_tag', '')
         
         try:
@@ -419,8 +716,22 @@ class ArticleImageMatcher:
                 else:
                     return False
             else:
-                logger.info(f"文章 {article_id} 未找到合适的匹配图片")
+                # 未找到合适的匹配图片，使用Gemini生成图片
-                self.log_to_database('WARNING', f"文章未找到匹配图片", f"文章ID: {article_id}")
+                logger.info(f"文章 {article_id} 未找到合适的匹配图片，调用Gemini生成图片")
+                self.log_to_database('WARNING', f"文章未找到匹配图片，尝试生成图片", f"文章ID: {article_id}")
+                
+                # 构建生成提示词
+                prompt = f"与'{article_title}'相关的插图，标签: {', '.join(article_tags)}"
+                generated_image_url = self.generate_image_with_gemini(prompt, article_tags, article_id)
+                
+                if generated_image_url:
+                    logger.info(f"文章 {article_id} 成功生成图片: {generated_image_url}")
+                    self.log_to_database('INFO', f"文章生成图片成功", 
+                                       f"文章ID: {article_id}, 图片URL: {generated_image_url}")
+                    return True
+                else:
+                    logger.error(f"文章 {article_id} 生成图片失败")
+                    self.log_to_database('ERROR', f"文章生成图片失败", f"文章ID: {article_id}")
                     return False
                 
         except Exception as e: