yym68686 commited on
Commit
aeec83c
·
1 Parent(s): f3fce0a

Add feature: support load balancing for multiple keys in a single channel, enabled by default.

Browse files
Files changed (4) hide show
  1. README.md +4 -3
  2. request.py +6 -6
  3. test/test_nostream.py +121 -0
  4. utils.py +8 -0
README.md CHANGED
@@ -21,10 +21,9 @@
21
  - 同时支持 Anthropic、Gemini、Vertex API。Vertex 同时支持 Claude 和 Gemini API。
22
  - 支持 OpenAI、 Anthropic、Gemini、Vertex 原生 tool use 函数调用。
23
  - 支持 OpenAI、Anthropic、Gemini、Vertex 原生识图 API。
24
- - 支持负载均衡,支持 Vertex 区域负载均衡,支持 Vertex 高并发,最高可将 Gemini,Claude 并发提高 (API数量 * 区域数量) 倍。除了 Vertex 区域负载均衡,所有 API 均支持渠道级负载均衡,提高沉浸式翻译体验。
25
  - 支持自动重试,当一个 API 渠道响应失败时,自动重试下一个 API 渠道。
26
  - 支持细粒度的权限控制。支持使用通配符设置 API key 可用渠道的特定模型。
27
- - 支持多个 API Key。
28
 
29
  ## Configuration
30
 
@@ -42,7 +41,9 @@ providers:
42
 
43
  - provider: anthropic
44
  base_url: https://api.anthropic.com/v1/messages
45
- api: sk-ant-api03-bNnAOJyA-xQw_twAA
 
 
46
  model:
47
  - claude-3-5-sonnet-20240620: claude-3-5-sonnet # 重命名模型,claude-3-5-sonnet-20240620 是服务商的模型名称,claude-3-5-sonnet 是重命名后的名字,可以使用简洁的名字代替原来复杂的名称,选填
48
  tools: true # 是否支持工具,如生成代码、生成文档等,默认是 true,选填
 
21
  - 同时支持 Anthropic、Gemini、Vertex API。Vertex 同时支持 Claude 和 Gemini API。
22
  - 支持 OpenAI、 Anthropic、Gemini、Vertex 原生 tool use 函数调用。
23
  - 支持 OpenAI、Anthropic、Gemini、Vertex 原生识图 API。
24
+ - 支持三种负载均衡,默认同时开启。1. 支持单个渠道多个 API Key 自动开启 API key 级别的轮训负载均衡。2. 支持 Vertex 区域级负载均衡,支持 Vertex 高并发,最高可将 Gemini,Claude 并发提高 (API数量 * 区域数量) 倍。3. 除了 Vertex 区域级负载均衡,所有 API 均支持渠道级负载均衡,提高沉浸式翻译体验。
25
  - 支持自动重试,当一个 API 渠道响应失败时,自动重试下一个 API 渠道。
26
  - 支持细粒度的权限控制。支持使用通配符设置 API key 可用渠道的特定模型。
 
27
 
28
  ## Configuration
29
 
 
41
 
42
  - provider: anthropic
43
  base_url: https://api.anthropic.com/v1/messages
44
+ api: # 支持多个 API Key,多个 key 自动开启轮训负载均衡,至少一个 key,必填
45
+ - sk-ant-api03-bNnAOJyA-xQw_twAA
46
+ - sk-ant-api02-bNnxxxx
47
  model:
48
  - claude-3-5-sonnet-20240620: claude-3-5-sonnet # 重命名模型,claude-3-5-sonnet-20240620 是服务商的模型名称,claude-3-5-sonnet 是重命名后的名字,可以使用简洁的名字代替原来复杂的名称,选填
49
  tools: true # 是否支持工具,如生成代码、生成文档等,默认是 true,选填
request.py CHANGED
@@ -43,9 +43,9 @@ async def get_gemini_payload(request, engine, provider):
43
  gemini_stream = "streamGenerateContent"
44
  url = provider['base_url']
45
  if url.endswith("v1beta"):
46
- url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'])
47
  if url.endswith("v1"):
48
- url = "https://generativelanguage.googleapis.com/v1/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'])
49
 
50
  messages = []
51
  systemInstruction = None
@@ -492,7 +492,7 @@ async def get_gpt_payload(request, engine, provider):
492
  'Content-Type': 'application/json',
493
  }
494
  if provider.get("api"):
495
- headers['Authorization'] = f"Bearer {provider['api']}"
496
  url = provider['base_url']
497
 
498
  messages = []
@@ -556,7 +556,7 @@ async def get_openrouter_payload(request, engine, provider):
556
  'Content-Type': 'application/json'
557
  }
558
  if provider.get("api"):
559
- headers['Authorization'] = f"Bearer {provider['api']}"
560
 
561
  url = provider['base_url']
562
 
@@ -640,7 +640,7 @@ async def get_claude_payload(request, engine, provider):
640
  model = provider['model'][request.model]
641
  headers = {
642
  "content-type": "application/json",
643
- "x-api-key": f"{provider['api']}",
644
  "anthropic-version": "2023-06-01",
645
  "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" if "claude-3-5-sonnet" in model else "tools-2024-05-16",
646
  }
@@ -753,7 +753,7 @@ async def get_dalle_payload(request, engine, provider):
753
  "Content-Type": "application/json",
754
  }
755
  if provider.get("api"):
756
- headers['Authorization'] = f"Bearer {provider['api']}"
757
  url = provider['base_url']
758
  url = BaseAPI(url).image_url
759
 
 
43
  gemini_stream = "streamGenerateContent"
44
  url = provider['base_url']
45
  if url.endswith("v1beta"):
46
+ url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'].next())
47
  if url.endswith("v1"):
48
+ url = "https://generativelanguage.googleapis.com/v1/models/{model}:{stream}?key={api_key}".format(model=model, stream=gemini_stream, api_key=provider['api'].next())
49
 
50
  messages = []
51
  systemInstruction = None
 
492
  'Content-Type': 'application/json',
493
  }
494
  if provider.get("api"):
495
+ headers['Authorization'] = f"Bearer {provider['api'].next()}"
496
  url = provider['base_url']
497
 
498
  messages = []
 
556
  'Content-Type': 'application/json'
557
  }
558
  if provider.get("api"):
559
+ headers['Authorization'] = f"Bearer {provider['api'].next()}"
560
 
561
  url = provider['base_url']
562
 
 
640
  model = provider['model'][request.model]
641
  headers = {
642
  "content-type": "application/json",
643
+ "x-api-key": f"{provider['api'].next()}",
644
  "anthropic-version": "2023-06-01",
645
  "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" if "claude-3-5-sonnet" in model else "tools-2024-05-16",
646
  }
 
753
  "Content-Type": "application/json",
754
  }
755
  if provider.get("api"):
756
+ headers['Authorization'] = f"Bearer {provider['api'].next()}"
757
  url = provider['base_url']
758
  url = BaseAPI(url).image_url
759
 
test/test_nostream.py ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+ import base64
3
+ import json
4
+ import os
5
+ from datetime import datetime
6
+
7
+ # 設置API密鑰和自定義base URL
8
+ API_KEY = ''
9
+ BASE_URL = 'http://localhost:8000/v1'
10
+ SAVE_DIR = 'safe_output' # 保存 JSON 輸出的目錄
11
+
12
+ def ensure_save_directory():
13
+ if not os.path.exists(SAVE_DIR):
14
+ os.makedirs(SAVE_DIR)
15
+
16
+ def image_to_base64(image_path):
17
+ with open(image_path, "rb") as image_file:
18
+ return base64.b64encode(image_file.read()).decode('utf-8')
19
+
20
+ def get_model_response(image_base64):
21
+ headers = {
22
+ "Content-Type": "application/json",
23
+ "Authorization": f"Bearer {API_KEY}"
24
+ }
25
+
26
+ tools = [
27
+ {
28
+ "type": "function",
29
+ "function": {
30
+ "name": "extract_underlined_text",
31
+ "description": "從圖片中提取紅色下劃線的文字",
32
+ "parameters": {
33
+ "type": "object",
34
+ "properties": {
35
+ "underlined_text": {
36
+ "type": "array",
37
+ "items": {"type": "string"},
38
+ "description": "紅色下劃線的文字列表"
39
+ }
40
+ },
41
+ "required": ["underlined_text"]
42
+ }
43
+ }
44
+ }
45
+ ]
46
+
47
+ payload = {
48
+
49
+ "model": "claude-3-5-sonnet",
50
+ "messages": [
51
+ {
52
+ "role": "user",
53
+ "content": [
54
+ {
55
+ "type": "text",
56
+ "text": "請仔細分析圖片,並提取所有使用紅色筆在單字、單詞或句子下方畫有橫線的文字。只提取有紅色下劃線的文字,忽略其他未標記的文字。將結果以 JSON 格式輸出,格式為 {\"underlined_text\": [\"文字1\", \"文字2\", ...]}。"
57
+ },
58
+ {
59
+ "type": "image_url",
60
+ "image_url": {
61
+ "url": f"data:image/jpeg;base64,{image_base64}"
62
+ }
63
+ }
64
+ ]
65
+ }
66
+ ],
67
+ "stream": True,
68
+ "tools": tools,
69
+ "tool_choice": {"type": "function", "function": {"name": "extract_underlined_text"}},
70
+ "max_tokens": 300
71
+ }
72
+
73
+ try:
74
+ response = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload, timeout=30)
75
+ response.raise_for_status()
76
+ return response.json()
77
+ except requests.exceptions.RequestException as e:
78
+ return f"Error: {e}"
79
+
80
+ def save_json_output(data):
81
+ ensure_save_directory()
82
+ timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
83
+ filename = f"{SAVE_DIR}/output_{timestamp}.json"
84
+ with open(filename, 'w', encoding='utf-8') as f:
85
+ json.dump(data, f, ensure_ascii=False, indent=2)
86
+ return filename
87
+
88
+ def main(image_path):
89
+ image_base64 = image_to_base64(image_path)
90
+
91
+ response = get_model_response(image_base64)
92
+
93
+ print("模型回應:")
94
+ print(json.dumps(response, indent=2, ensure_ascii=False))
95
+
96
+ if isinstance(response, str) and response.startswith("Error"):
97
+ print(response)
98
+ return
99
+
100
+ if 'choices' in response and response['choices']:
101
+ message = response['choices'][0]['message']
102
+ if 'tool_calls' in message:
103
+ tool_call = message['tool_calls'][0]
104
+ if tool_call['function']['name'] == 'extract_underlined_text':
105
+ function_args = json.loads(tool_call['function']['arguments'])
106
+ print("\n提取的紅色下劃線文字:")
107
+ print(json.dumps(function_args, indent=2, ensure_ascii=False))
108
+
109
+ # 保存 JSON 輸出
110
+ saved_file = save_json_output(function_args)
111
+ print(f"\nJSON 輸出已保存至: {saved_file}")
112
+ else:
113
+ print("\n模型調用了未預期的函數。")
114
+ else:
115
+ print("\n模型沒有調用工具。")
116
+ else:
117
+ print("\n無法解析回應。")
118
+
119
+ if __name__ == "__main__":
120
+ image_path = "00001 (8).jpg" # 替換為您的圖像路徑
121
+ main(image_path)
utils.py CHANGED
@@ -15,7 +15,15 @@ def update_config(config_data):
15
  provider['model'] = model_dict
16
  if provider.get('project_id'):
17
  provider['base_url'] = 'https://aiplatform.googleapis.com/'
 
 
 
 
 
 
 
18
  config_data['providers'][index] = provider
 
19
  api_keys_db = config_data['api_keys']
20
  api_list = [item["api"] for item in api_keys_db]
21
  # logger.info(json.dumps(config_data, indent=4, ensure_ascii=False))
 
15
  provider['model'] = model_dict
16
  if provider.get('project_id'):
17
  provider['base_url'] = 'https://aiplatform.googleapis.com/'
18
+
19
+ if provider.get('api'):
20
+ if isinstance(provider.get('api'), str):
21
+ provider['api'] = CircularList([provider.get('api')])
22
+ if isinstance(provider.get('api'), list):
23
+ provider['api'] = CircularList(provider.get('api'))
24
+
25
  config_data['providers'][index] = provider
26
+
27
  api_keys_db = config_data['api_keys']
28
  api_list = [item["api"] for item in api_keys_db]
29
  # logger.info(json.dumps(config_data, indent=4, ensure_ascii=False))