""" 一个使用 google-genai SDK 向 Gemini API 发送请求的客户端。 该脚本提供了一个与Gemini模型交互的函数，支持文本提示和多种文件输入 (例如本地文件路径、文件字节流、网络URL)。 **依赖库:** - google-genai: Google官方最新的GenAI Python SDK。 安装命令: pip install google-genai - filetype: 用于从文件内容识别MIME类型。 安装命令: pip install filetype """ import os import sys import logging import urllib.request from typing import List, Union, Optional, Dict, Any from concurrent.futures import ThreadPoolExecutor, as_completed from google import genai from google.genai import types from google.genai.errors import APIError import filetype current_dir = os.path.dirname(os.path.abspath(__file__)) root_dir = os.path.dirname(current_dir) sys.path.insert(0, root_dir) from utils import llm_account_helper logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' ) def generate_text( prompt: str, model: str = 'gemini-2.5-flash', files_input: Optional[List[Union[str, bytes]]] = None ) -> str: """ 向Gemini API发送一个包含提示和可选文件(本地路径、URL或字节流)的请求。 Args: prompt (str): 必需的，给模型的文本提示。 model (str): 可选的，要使用的模型名称。默认为 'gemini-2.5-flash'。 files_input (Optional[List[Union[str, bytes]]]): 可选的，请求中要包含的文件列表。 列表中的每一项可以是： - 字符串(str): 本地文件路径，或者以 http://, https:// 开头的网络URL。 - 字节(bytes): 内存中的文件内容。 Returns: str: 从Gemini模型返回的文本响应。 """ api_key = llm_account_helper.get_api_key('gemini') client = genai.Client(api_key=api_key) uploaded_files_to_delete = [] try: contents = [] if files_input: logging.info(f"正在处理 {len(files_input)} 个文件输入...") for file_item in files_input: if isinstance(file_item, str): # --- 情况 A: 处理网络 URL --- if file_item.lower().startswith(('http://', 'https://')): logging.info(f"检测到 URL，正在下载: {file_item}") try: # 使用标准库 urllib 下载。添加 User-Agent 以避免部分服务器拒绝请求。 req = urllib.request.Request( file_item, headers={'User-Agent': 'Mozilla/5.0 (Compatible; GeminiClient/1.0)'} ) with urllib.request.urlopen(req) as response: file_data = response.read() # 下载成功后，复用下方的字节流处理逻辑 mime_type = filetype.guess_mime(file_data) logging.info(f"URL 下载完成 (MIME: {mime_type})，已添加到请求。") contents.append(types.Part.from_bytes(data=file_data, mime_type=mime_type)) except Exception as e: logging.error(f"下载 URL 失败: {e}") raise ValueError(f"无法处理 URL '{file_item}': {e}") continue # 处理完 URL 后跳过当前循环，避免进入本地文件判断 # --- 情况 B: 处理本地文件路径 --- if not os.path.exists(file_item): raise FileNotFoundError(f"本地文件 '{file_item}' 不存在。") logging.info(f"正在上传本地文件: {file_item}") # 使用 File API 上传本地文件 (适合大文件) uploaded_file = client.files.upload(file=file_item) contents.append(uploaded_file) uploaded_files_to_delete.append(uploaded_file) elif isinstance(file_item, bytes): # --- 情况 C: 处理内存字节流 --- mime_type = filetype.guess_mime(file_item) logging.info(f"正在处理内存字节流 (MIME: {mime_type})") # 直接将小文件数据内嵌到请求中 contents.append(types.Part.from_bytes(data=file_item, mime_type=mime_type)) else: raise ValueError( f"不支持的输入类型: {type(file_item)}。仅支持本地路径(str)、URL(str)或字节流(bytes)。" ) contents.append(prompt) logging.info(f"正在向模型 '{model}' 发送请求...") response = client.models.generate_content( model=model, contents=contents ) return response.text except APIError as e: logging.error(f"Gemini API 调用错误: {e}") raise except Exception as e: logging.error(f"执行过程中发生未知错误: {e}") raise finally: # 清理通过 File API 上传的文件 if uploaded_files_to_delete: logging.info(f"正在清理 {len(uploaded_files_to_delete)} 个已上传的临时文件...") for f in uploaded_files_to_delete: try: client.files.delete(name=f.name) except Exception as e: logging.warning(f"清理文件 {f.name} 失败: {e}") def concurrent_generate_text( prompt_file_pairs: List[Dict[str, Union[str, Optional[List[Union[str, bytes]]]]]], model: str = 'gemini-2.5-flash', max_workers: int = 10 ) -> List[Dict[str, Any]]: """ 并发执行多个Gemini请求，每个请求对应一个prompt和文件数组。 Args: prompt_file_pairs (List[Dict[str, Union[str, Optional[List[Union[str, bytes]]]]]]): 包含多个prompt和对应文件列表的字典对象列表，每个字典应包含'prompt'和'files'键 model (str): 要使用的模型名称。默认为 'gemini-2.5-flash' max_workers (int): 最大并发线程数。默认为 10 Returns: List[Dict[str, Any]]: 每个元素是包含成功失败状态、失败原因、返回数据的字典对象 结果数组中元素的位置与输入的prompt_file_pairs位置一一对应 """ results = [None] * len(prompt_file_pairs) # 预先分配结果列表，保持与输入顺序一致 def process_single_request(pair_idx: int, prompt: str, files: Optional[List[Union[str, bytes]]]) -> Dict[str, Any]: try: response_text = generate_text(prompt=prompt, model=model, files_input=files) return { "success": True, "data": response_text, "error_message": None } except Exception as e: return { "success": False, "data": None, "error_message": str(e) } with ThreadPoolExecutor(max_workers=max_workers) as executor: # 提交所有任务 future_to_index = { executor.submit(process_single_request, i, prompt_file_pairs[i]['prompt'], prompt_file_pairs[i].get('files')): i for i in range(len(prompt_file_pairs)) } # 收集结果并按原始顺序放置 for future in as_completed(future_to_index): idx = future_to_index[future] try: result = future.result() results[idx] = result except Exception as e: # 即使future执行出现异常，也确保在正确位置设置结果 results[idx] = { "success": False, "data": None, "error_message": str(e) } return results if __name__ == '__main__': print("--- Gemini 请求客户端演示 (已支持 URL) ---") # try: # # 示例: 使用网络图片 URL # print("\n--- 运行示例: 多模态提示 (使用网络 URL) ---") # image_url = "https://www.google.com/images/branding/googlelogo/2x/googlelogo_color_272x92dp.png" # prompt = "这张图片里的 logo 是什么公司的？" # # print(f"[提示]: {prompt}") # print(f"[输入 URL]: {image_url}") # # response = generate_text( # prompt=prompt, # files_input=[image_url] # ) # print(f"[Gemini 回复]:\n{response}") # # except Exception as e: # print(f"\n演示运行失败: {e}") # 示例: 使用并发请求 print("\n--- 运行示例: 并发请求 ---") prompt_file_pairs = [ { "prompt": "图片上是什么？", "files": ["https://www.google.com/images/branding/googlelogo/2x/googlelogo_color_272x92dp.png"] # 第一个请求不带文件 }, { "prompt": "图片中是谁？", "files": ["https://inews.gtimg.com/om_bt/O6k8mse8MT9ki8fba5c7RK1j1xLFqT-FFZ9RirryqjENkAA/641"] # 第二个请求不带文件 } ] try: concurrent_results = concurrent_generate_text( prompt_file_pairs=prompt_file_pairs, model='gemini-2.5-flash' ) for i, result in enumerate(concurrent_results): print(f"请求 {i+1}:") if result['success']: print(f" 成功: {result['data'][:50]}...") # 只显示前50个字符 else: print(f" 失败: {result['error_message']}") except Exception as e: print(f"\n并发请求演示失败: {e}")