Merge pull request #165 from endsock/addpic

添加OPENAI图片协议，支持data:image/${base64}

Author: CJackHwang

.gitignore

@@ -7,6 +7,7 @@ yarn-debug.log*
 yarn-error.log*
 pnpm-debug.log*
 lerna-debug.log*
+/upload_images
 
 # Diagnostic reports (https://nodejs.org/api/report.html)
 report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json

api_utils/request_processor.py

@@ -833,7 +833,7 @@ async def _process_request_refactored(
         await _handle_model_switching(req_id, context, check_client_disconnected)
         await _handle_parameter_cache(req_id, context)
 
-        prepared_prompt = await _prepare_and_validate_request(req_id, request, check_client_disconnected)
+        prepared_prompt,image_list = await _prepare_and_validate_request(req_id, request, check_client_disconnected)
 
         # 使用PageController处理页面交互
         # 注意：聊天历史清空已移至队列处理锁释放后执行
@@ -850,7 +850,7 @@ async def _process_request_refactored(
         # 优化：在提交提示前再次检查客户端连接，避免不必要的后台请求
         check_client_disconnected("提交提示前最终检查")
 
-        await page_controller.submit_prompt(prepared_prompt, check_client_disconnected)
+        await page_controller.submit_prompt(prepared_prompt,image_list, check_client_disconnected)
 
         # 响应处理仍然需要在这里，因为它决定了是流式还是非流式，并设置future
         response_result = await _handle_response_processing(

api_utils/request_processor_backup.py

@@ -198,7 +198,7 @@ def check_client_disconnected(*args):
             raise HTTPException(status_code=400, detail=f"[{req_id}] 无效请求: {e}")
 
         # 准备提示
-        prepared_prompt = prepare_combined_prompt(request.messages, req_id)
+        prepared_prompt,image_list = prepare_combined_prompt(request.messages, req_id)
         check_client_disconnected("After Prompt Prep: ")
 
         # 这里需要添加完整的处理逻辑 - 由于函数太长，暂时返回简化响应

api_utils/utils.py

@@ -10,7 +10,11 @@
 from typing import Any, Dict, List, Optional, AsyncGenerator
 from asyncio import Queue
 from models import Message
-
+import re
+import base64
+import requests
+import os
+import hashlib
 
 
 # --- SSE生成函数 ---
@@ -194,6 +198,48 @@ def validate_chat_request(messages: List[Message], req_id: str) -> Dict[str, Opt
     }
 
 
+def extract_base64_to_local(base64_data: str) -> str:
+    output_dir = os.path.join(os.path.dirname(__file__), '..', 'upload_images')
+    match = re.match(r"data:image/(\w+);base64,(.*)", base64_data)
+    if not match:
+        print("错误: Base64 数据格式不正确。")
+        return None
+
+    image_type = match.group(1)  # 例如 "png", "jpeg"
+    encoded_image_data = match.group(2)
+
+    try:
+        # 解码 Base64 字符串
+        decoded_image_data = base64.b64decode(encoded_image_data)
+    except base64.binascii.Error as e:
+        print(f"错误: Base64 解码失败 - {e}")
+        return None
+
+    # 计算图片数据的 MD5 值
+    md5_hash = hashlib.md5(decoded_image_data).hexdigest()
+
+    # 确定文件扩展名和完整文件路径
+    file_extension = f".{image_type}"
+    output_filepath = os.path.join(output_dir, f"{md5_hash}{file_extension}")
+
+    # 确保输出目录存在
+    os.makedirs(output_dir, exist_ok=True)
+
+    if os.path.exists(output_filepath):
+        print(f"文件已存在，跳过保存: {output_filepath}")
+        return output_filepath
+
+    # 保存图片到文件
+    try:
+        with open(output_filepath, "wb") as f:
+            f.write(decoded_image_data)
+        print(f"图片已成功保存到: {output_filepath}")
+        return output_filepath
+    except IOError as e:
+        print(f"错误: 保存文件失败 - {e}")
+        return None
+
+
 # --- 提示准备函数 ---
 def prepare_combined_prompt(messages: List[Message], req_id: str) -> str:
     """准备组合提示"""
@@ -238,6 +284,7 @@ def prepare_combined_prompt(messages: List[Message], req_id: str) -> str:
         role = msg.role or 'unknown'
         role_prefix_ui = f"{role_map_ui.get(role, role.capitalize())}:\n"
         current_turn_parts = [role_prefix_ui]
+        images_list = []
 
         content = msg.content or ''
         content_str = ""
@@ -252,6 +299,15 @@ def prepare_combined_prompt(messages: List[Message], req_id: str) -> str:
                     text_parts.append(item.text or '')
                 elif isinstance(item, dict) and item.get('type') == 'text':
                     text_parts.append(item.get('text', ''))
+                elif hasattr(item, 'type') and item.type == 'image_url':
+                    image_url_value = item.image_url.url
+                    if image_url_value.startswith("data:image/"):
+                        try:
+                            # 提取 Base64 字符串
+                            image_full_path = extract_base64_to_local(image_url_value)
+                            images_list.append(image_full_path)
+                        except (ValueError, requests.exceptions.RequestException, Exception) as e:
+                            print(f"处理 Base64 图片并上传到 Imgur 失败: {e}")
                 else:
                     logger.warning(f"[{req_id}] (准备提示) 警告: 在索引 {i} 的消息中忽略非文本或未知类型的 content item")
             content_str = "\n".join(text_parts).strip()
@@ -302,7 +358,7 @@ def prepare_combined_prompt(messages: List[Message], req_id: str) -> str:
     preview_text = final_prompt[:300].replace('\n', '\\n')
     logger.info(f"[{req_id}] (准备提示) 组合提示长度: {len(final_prompt)}。预览: '{preview_text}...'")
 
-    return final_prompt 
+    return final_prompt,images_list
 
 
 def estimate_tokens(text: str) -> int:

browser_utils/page_controller.py

@@ -12,7 +12,7 @@
     MAT_CHIP_REMOVE_BUTTON_SELECTOR, TOP_P_INPUT_SELECTOR, SUBMIT_BUTTON_SELECTOR,
     CLEAR_CHAT_BUTTON_SELECTOR, CLEAR_CHAT_CONFIRM_BUTTON_SELECTOR, OVERLAY_SELECTOR,
     PROMPT_TEXTAREA_SELECTOR, RESPONSE_CONTAINER_SELECTOR, RESPONSE_TEXT_SELECTOR,
-    EDIT_MESSAGE_BUTTON_SELECTOR
+    EDIT_MESSAGE_BUTTON_SELECTOR,USE_URL_CONTEXT_SELECTOR,UPLOAD_BUTTON_SELECTOR
 )
 from config import (
     CLICK_TIMEOUT_MS, WAIT_FOR_ELEMENT_TIMEOUT_MS, CLEAR_CHAT_VERIFY_TIMEOUT_MS,
@@ -59,6 +59,30 @@ async def adjust_parameters(self, request_params: Dict[str, Any], page_params_ca
         await self._adjust_top_p(top_p_to_set, check_client_disconnected)
         await self._check_disconnect(check_client_disconnected, "End Parameter Adjustment")
 
+        # 调整URL CONTEXT
+        await self._open_url_content(check_client_disconnected)
+
+    async def _open_url_content(self,check_client_disconnected: Callable):
+        try:
+            collapse_tools_locator = self.page.locator('button[aria-label="Expand or collapse tools"]')
+            grandparent_locator = collapse_tools_locator.locator("xpath=../..")
+
+            # 3. 获取祖父级元素的 class 属性值
+            # get_attribute 返回一个包含所有 class 的字符串，例如 "menu dropdown active"
+            class_string = await grandparent_locator.get_attribute("class")
+
+            # 4. 在 Python 中进行判断
+            # 确保 class_string 不是 None，并且 'expanded' 是一个独立的 class
+            if class_string and "expanded" not in class_string.split():
+                await collapse_tools_locator.click(timeout=CLICK_TIMEOUT_MS)
+                await asyncio.sleep(0.5)
+            use_url_content_selector = self.page.locator(USE_URL_CONTEXT_SELECTOR)
+            is_checked = await use_url_content_selector.get_attribute("aria-checked")
+            if "false" == is_checked:
+                await use_url_content_selector.click(timeout=CLICK_TIMEOUT_MS)
+                await self._check_disconnect(check_client_disconnected, "点击URLCONTEXT")
+        except Exception as e:
+            self.logger.error(f"[{self.req_id}] ❌ 操作USE_URL_CONTEXT_SELECTOR时发生错误:{e}。")
 
     async def _adjust_temperature(self, temperature: float, page_params_cache: dict, params_cache_lock: asyncio.Lock, check_client_disconnected: Callable):
         """调整温度参数。"""
@@ -76,6 +100,7 @@ async def _adjust_temperature(self, temperature: float, page_params_cache: dict,
             self.logger.info(f"[{self.req_id}] 请求温度 ({clamped_temp}) 与缓存值 ({cached_temp}) 不一致或缓存中无值。需要与页面交互。")
             temp_input_locator = self.page.locator(TEMPERATURE_INPUT_SELECTOR)
 
+
             try:
                 await expect_async(temp_input_locator).to_be_visible(timeout=5000)
                 await self._check_disconnect(check_client_disconnected, "温度调整 - 输入框可见后")
@@ -421,7 +446,7 @@ async def _verify_chat_cleared(self, check_client_disconnected: Callable):
         except Exception as verify_err:
             self.logger.warning(f"[{self.req_id}] ⚠️ 警告: 清空聊天验证失败 (最后响应容器未隐藏): {verify_err}")
 
-    async def submit_prompt(self, prompt: str, check_client_disconnected: Callable):
+    async def submit_prompt(self, prompt: str,image_list: List, check_client_disconnected: Callable):
         """提交提示到页面。"""
         self.logger.info(f"[{self.req_id}] 填充并提交提示 ({len(prompt)} chars)...")
         prompt_textarea_locator = self.page.locator(PROMPT_TEXTAREA_SELECTOR)
@@ -446,8 +471,34 @@ async def submit_prompt(self, prompt: str, check_client_disconnected: Callable):
             await autosize_wrapper_locator.evaluate('(element, text) => { element.setAttribute("data-value", text); }', prompt)
             await self._check_disconnect(check_client_disconnected, "After Input Fill")
 
+            # 上传
+            if len(image_list) > 0:
+                try:
+                    # 1. 监听文件选择器
+                    #    page.expect_file_chooser() 会返回一个上下文管理器
+                    #    当文件选择器出现时，它会得到 FileChooser 对象
+                    function_btn_localtor = self.page.locator('button[aria-label="Insert assets such as images, videos, files, or audio"]')
+                    await function_btn_localtor.click()
+                    asyncio.sleep(0.5)
+                    async with self.page.expect_file_chooser() as fc_info:
+                        # 2. 点击那个会触发文件选择的普通按钮
+                        upload_btn_localtor = self.page.locator(UPLOAD_BUTTON_SELECTOR)
+                        await upload_btn_localtor.click()
+                        print("点击了 JS 上传按钮，等待文件选择器...")
+
+                    # 3. 获取文件选择器对象
+                    file_chooser = await fc_info.value
+                    print("文件选择器已出现。")
+
+                    # 4. 设置要上传的文件
+                    await file_chooser.set_files(image_list)
+                    print(f"已将 '{image_list}' 设置到文件选择器。")
+
+                except Exception as e:
+                    print(f"在上传文件时发生错误: {e}")
+
             # 等待发送按钮启用
-            wait_timeout_ms_submit_enabled = 40000
+            wait_timeout_ms_submit_enabled = 100000
             try:
                 await self._check_disconnect(check_client_disconnected, "填充提示后等待发送按钮启用 - 前置检查")
                 await expect_async(submit_button_locator).to_be_enabled(timeout=wait_timeout_ms_submit_enabled)

config/__init__.py

@@ -65,6 +65,8 @@
     'MAT_CHIP_REMOVE_BUTTON_SELECTOR',
     'TOP_P_INPUT_SELECTOR',
     'TEMPERATURE_INPUT_SELECTOR',
+    'USE_URL_CONTEXT_SELECTOR',
+    'UPLOAD_BUTTON_SELECTOR',
 
     # 设置配置
     'DEBUG_LOGS_ENABLED',

config/selectors.py

@@ -12,6 +12,7 @@
 SUBMIT_BUTTON_SELECTOR = 'button[aria-label="Run"].run-button'
 CLEAR_CHAT_BUTTON_SELECTOR = 'button[data-test-clear="outside"][aria-label="Clear chat"]'
 CLEAR_CHAT_CONFIRM_BUTTON_SELECTOR = 'button.mdc-button:has-text("Continue")'
+UPLOAD_BUTTON_SELECTOR = 'button[aria-label="Upload File"]'
 
 # --- 响应相关选择器 ---
 RESPONSE_CONTAINER_SELECTOR = 'ms-chat-turn .chat-turn-container.model'
@@ -39,4 +40,5 @@
 STOP_SEQUENCE_INPUT_SELECTOR = 'input[aria-label="Add stop token"]'
 MAT_CHIP_REMOVE_BUTTON_SELECTOR = 'mat-chip-set mat-chip-row button[aria-label*="Remove"]'
 TOP_P_INPUT_SELECTOR = 'div.settings-item-column:has(h3:text-is("Top P")) input[type="number"].slider-input'
-TEMPERATURE_INPUT_SELECTOR = 'div[data-test-id="temperatureSliderContainer"] input[type="number"].slider-input' 
+TEMPERATURE_INPUT_SELECTOR = 'div[data-test-id="temperatureSliderContainer"] input[type="number"].slider-input'
+USE_URL_CONTEXT_SELECTOR = 'button[aria-label="Browse the url context"]'

models/chat.py

@@ -13,11 +13,13 @@ class ToolCall(BaseModel):
     type: str = "function"
     function: FunctionCall
 
+class ImageURL(BaseModel):
+    url: str
 
 class MessageContentItem(BaseModel):
     type: str
     text: Optional[str] = None
-
+    image_url: Optional[ImageURL] = None
 
 class Message(BaseModel):
     role: str