From aeec83c903a464ee3762278a151ede2c7fd782f0 Mon Sep 17 00:00:00 2001 From: yym68686 Date: Wed, 4 Sep 2024 23:45:01 +0800 Subject: [PATCH] Add feature: support load balancing for multiple keys in a single channel, enabled by default. --- README.md | 7 +-- request.py | 12 ++--- test/test_nostream.py | 121 ++++++++++++++++++++++++++++++++++++++++++ utils.py | 8 +++ 4 files changed, 139 insertions(+), 9 deletions(-) create mode 100644 test/test_nostream.py diff --git a/README.md b/README.md index 8241d32..63d8b11 100644 --- a/README.md +++ b/README.md @@ -21,10 +21,9 @@ - 同时支持 Anthropic、Gemini、Vertex API。Vertex 同时支持 Claude 和 Gemini API。 - 支持 OpenAI、 Anthropic、Gemini、Vertex 原生 tool use 函数调用。 - 支持 OpenAI、Anthropic、Gemini、Vertex 原生识图 API。 -- 支持负载均衡，支持 Vertex 区域负载均衡，支持 Vertex 高并发，最高可将 Gemini，Claude 并发提高 （API数量 * 区域数量） 倍。除了 Vertex 区域负载均衡，所有 API 均支持渠道级负载均衡，提高沉浸式翻译体验。 +- 支持三种负载均衡，默认同时开启。1. 支持单个渠道多个 API Key 自动开启 API key 级别的轮训负载均衡。2. 支持 Vertex 区域级负载均衡，支持 Vertex 高并发，最高可将 Gemini，Claude 并发提高 （API数量 * 区域数量） 倍。3. 除了 Vertex 区域级负载均衡，所有 API 均支持渠道级负载均衡，提高沉浸式翻译体验。 - 支持自动重试，当一个 API 渠道响应失败时，自动重试下一个 API 渠道。 - 支持细粒度的权限控制。支持使用通配符设置 API key 可用渠道的特定模型。 -- 支持多个 API Key。 ## Configuration @@ -42,7 +41,9 @@ providers: - provider: anthropic base_url: https://api.anthropic.com/v1/messages - api: sk-ant-api03-bNnAOJyA-xQw_twAA + api: # 支持多个 API Key，多个 key 自动开启轮训负载均衡，至少一个 key，必填 + - sk-ant-api03-bNnAOJyA-xQw_twAA + - sk-ant-api02-bNnxxxx model: - claude-3-5-sonnet-20240620: claude-3-5-sonnet # 重命名模型，claude-3-5-sonnet-20240620 是服务商的模型名称，claude-3-5-sonnet 是重命名后的名字，可以使用简洁的名字代替原来复杂的名称，选填 tools: true # 是否支持工具，如生成代码、生成文档等，默认是 true，选填 diff --git a/request.py b/request.py index 8df44e8..8960231 100644 --- a/request.py +++ b/request.py @@ -43,9 +43,9 @@ async def get_gemini_payload(request, engine, provider): gemini_stream = "streamGenerateContent" url = provider['base_url'] if url.endswith("v1beta"): - url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api']) + url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'].next()) if url.endswith("v1"): - url = "https://generativelanguage.googleapis.com/v1/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api']) + url = "https://generativelanguage.googleapis.com/v1/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'].next()) messages = [] systemInstruction = None @@ -492,7 +492,7 @@ async def get_gpt_payload(request, engine, provider): 'Content-Type': 'application/json', } if provider.get("api"): - headers['Authorization'] = f"Bearer {provider['api']}" + headers['Authorization'] = f"Bearer {provider['api'].next()}" url = provider['base_url'] messages = [] @@ -556,7 +556,7 @@ async def get_openrouter_payload(request, engine, provider): 'Content-Type': 'application/json' } if provider.get("api"): - headers['Authorization'] = f"Bearer {provider['api']}" + headers['Authorization'] = f"Bearer {provider['api'].next()}" url = provider['base_url'] @@ -640,7 +640,7 @@ async def get_claude_payload(request, engine, provider): model = provider['model'][request.model] headers = { "content-type": "application/json", - "x-api-key": f"{provider['api']}", + "x-api-key": f"{provider['api'].next()}", "anthropic-version": "2023-06-01", "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" if "claude-3-5-sonnet" in model else "tools-2024-05-16", } @@ -753,7 +753,7 @@ async def get_dalle_payload(request, engine, provider): "Content-Type": "application/json", } if provider.get("api"): - headers['Authorization'] = f"Bearer {provider['api']}" + headers['Authorization'] = f"Bearer {provider['api'].next()}" url = provider['base_url'] url = BaseAPI(url).image_url diff --git a/test/test_nostream.py b/test/test_nostream.py new file mode 100644 index 0000000..7378d7c --- /dev/null +++ b/test/test_nostream.py @@ -0,0 +1,121 @@ +import requests +import base64 +import json +import os +from datetime import datetime + +# 設置API密鑰和自定義base URL +API_KEY = '' +BASE_URL = 'http://localhost:8000/v1' +SAVE_DIR = 'safe_output' # 保存 JSON 輸出的目錄 + +def ensure_save_directory(): + if not os.path.exists(SAVE_DIR): + os.makedirs(SAVE_DIR) + +def image_to_base64(image_path): + with open(image_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + +def get_model_response(image_base64): + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {API_KEY}" + } + + tools = [ + { + "type": "function", + "function": { + "name": "extract_underlined_text", + "description": "從圖片中提取紅色下劃線的文字", + "parameters": { + "type": "object", + "properties": { + "underlined_text": { + "type": "array", + "items": {"type": "string"}, + "description": "紅色下劃線的文字列表" + } + }, + "required": ["underlined_text"] + } + } + } + ] + + payload = { + + "model": "claude-3-5-sonnet", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "請仔細分析圖片，並提取所有使用紅色筆在單字、單詞或句子下方畫有橫線的文字。只提取有紅色下劃線的文字，忽略其他未標記的文字。將結果以 JSON 格式輸出，格式為 {\"underlined_text\": [\"文字1\", \"文字2\", ...]}。" + }, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{image_base64}" + } + } + ] + } + ], + "stream": True, + "tools": tools, + "tool_choice": {"type": "function", "function": {"name": "extract_underlined_text"}}, + "max_tokens": 300 + } + + try: + response = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload, timeout=30) + response.raise_for_status() + return response.json() + except requests.exceptions.RequestException as e: + return f"Error: {e}" + +def save_json_output(data): + ensure_save_directory() + timestamp = datetime.now().strftime("%Y%m%d%H%M%S") + filename = f"{SAVE_DIR}/output_{timestamp}.json" + with open(filename, 'w', encoding='utf-8') as f: + json.dump(data, f, ensure_ascii=False, indent=2) + return filename + +def main(image_path): + image_base64 = image_to_base64(image_path) + + response = get_model_response(image_base64) + + print("模型回應:") + print(json.dumps(response, indent=2, ensure_ascii=False)) + + if isinstance(response, str) and response.startswith("Error"): + print(response) + return + + if 'choices' in response and response['choices']: + message = response['choices'][0]['message'] + if 'tool_calls' in message: + tool_call = message['tool_calls'][0] + if tool_call['function']['name'] == 'extract_underlined_text': + function_args = json.loads(tool_call['function']['arguments']) + print("\n提取的紅色下劃線文字:") + print(json.dumps(function_args, indent=2, ensure_ascii=False)) + + # 保存 JSON 輸出 + saved_file = save_json_output(function_args) + print(f"\nJSON 輸出已保存至: {saved_file}") + else: + print("\n模型調用了未預期的函數。") + else: + print("\n模型沒有調用工具。") + else: + print("\n無法解析回應。") + +if __name__ == "__main__": + image_path = "00001 (8).jpg" # 替換為您的圖像路徑 + main(image_path) \ No newline at end of file diff --git a/utils.py b/utils.py index 231d61e..f27154e 100644 --- a/utils.py +++ b/utils.py @@ -15,7 +15,15 @@ def update_config(config_data): provider['model'] = model_dict if provider.get('project_id'): provider['base_url'] = 'https://aiplatform.googleapis.com/' + + if provider.get('api'): + if isinstance(provider.get('api'), str): + provider['api'] = CircularList([provider.get('api')]) + if isinstance(provider.get('api'), list): + provider['api'] = CircularList(provider.get('api')) + config_data['providers'][index] = provider + api_keys_db = config_data['api_keys'] api_list = [item["api"] for item in api_keys_db] # logger.info(json.dumps(config_data, indent=4, ensure_ascii=False))