diff --git a/plugins/AIChat/main.py b/plugins/AIChat/main.py index 313daea..4acb797 100644 --- a/plugins/AIChat/main.py +++ b/plugins/AIChat/main.py @@ -488,7 +488,7 @@ class AIChat(PluginBase): async def _generate_image_description(self, image_base64: str, prompt: str, config: dict) -> str: """ - 使用 AI 生成图片描述 + 使用 Gemini API 生成图片描述 Args: image_base64: 图片的 base64 数据 @@ -498,26 +498,39 @@ class AIChat(PluginBase): Returns: 图片描述文本,失败返回空字符串 """ + import json try: api_config = self.config["api"] description_model = config.get("model", api_config["model"]) + api_url = api_config.get("gemini_url", "https://api.functen.cn/v1beta/models") - # 构建消息 - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": prompt}, - {"type": "image_url", "image_url": {"url": image_base64}} - ] - } - ] + # 处理 base64 数据 + image_data = image_base64 + mime_type = "image/jpeg" + if image_data.startswith("data:"): + mime_type = image_data.split(";")[0].split(":")[1] + image_data = image_data.split(",", 1)[1] + + # 构建 Gemini 格式请求 + full_url = f"{api_url}/{description_model}:streamGenerateContent?alt=sse" payload = { - "model": description_model, - "messages": messages, - "max_tokens": config.get("max_tokens", 1000), - "stream": True + "contents": [ + { + "parts": [ + {"text": prompt}, + { + "inline_data": { + "mime_type": mime_type, + "data": image_data + } + } + ] + } + ], + "generationConfig": { + "maxOutputTokens": config.get("max_tokens", 1000) + } } headers = { @@ -525,12 +538,12 @@ class AIChat(PluginBase): "Authorization": f"Bearer {api_config['api_key']}" } - timeout = aiohttp.ClientTimeout(total=api_config["timeout"]) + timeout = aiohttp.ClientTimeout(total=api_config.get("timeout", 120)) # 配置代理 connector = None proxy_config = self.config.get("proxy", {}) - if proxy_config.get("enabled", False): + if proxy_config.get("enabled", False) and PROXY_SUPPORT: proxy_type = proxy_config.get("type", "socks5").upper() proxy_host = proxy_config.get("host", "127.0.0.1") proxy_port = proxy_config.get("port", 7890) @@ -542,43 +555,37 @@ class AIChat(PluginBase): else: proxy_url = f"{proxy_type}://{proxy_host}:{proxy_port}" - if PROXY_SUPPORT: - try: - connector = ProxyConnector.from_url(proxy_url) - except Exception as e: - logger.warning(f"代理配置失败,将直连: {e}") - connector = None + try: + connector = ProxyConnector.from_url(proxy_url) + except Exception as e: + logger.warning(f"代理配置失败,将直连: {e}") async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: - async with session.post( - api_config["url"], - json=payload, - headers=headers - ) as resp: + async with session.post(full_url, json=payload, headers=headers) as resp: if resp.status != 200: error_text = await resp.text() logger.error(f"图片描述 API 返回错误: {resp.status}, {error_text[:200]}") return "" - # 流式接收响应 - import json + # 流式接收 Gemini 响应 description = "" async for line in resp.content: line = line.decode('utf-8').strip() - if not line or line == "data: [DONE]": + if not line or not line.startswith("data: "): continue - if line.startswith("data: "): - try: - data = json.loads(line[6:]) - delta = data.get("choices", [{}])[0].get("delta", {}) - content = delta.get("content", "") - if content: - description += content - except: - pass + try: + data = json.loads(line[6:]) + candidates = data.get("candidates", []) + if candidates: + parts = candidates[0].get("content", {}).get("parts", []) + for part in parts: + if "text" in part: + description += part["text"] + except: + pass - logger.debug(f"图片描述生成成功: {description}") + logger.debug(f"图片描述生成成功: {description[:100]}...") return description.strip() except Exception as e: @@ -788,6 +795,7 @@ class AIChat(PluginBase): Gemini 格式: {"functionCall": {"name": "...", "args": {...}}} 转换为内部格式: {"id": "...", "function": {"name": "...", "arguments": "..."}} """ + import json tool_calls = [] for i, part in enumerate(response_parts): if "functionCall" in part: @@ -812,6 +820,7 @@ class AIChat(PluginBase): tool_calls: 工具调用列表 tool_results: 工具执行结果列表 """ + import json new_contents = contents.copy() # 添加 model 的工具调用响应