Add Gemini region load balancing.

README.md

@@ -55,10 +55,14 @@ providers:
   - provider: vertex
     project_id: gen-lang-client-xxxxxxxxxxxxxx #    描述： 您的Google Cloud项目ID。格式： 字符串，通常由小写字母、数字和连字符组成。获取方式： 在Google Cloud Console的项目选择器中可以找到您的项目ID。
     private_key: "-----BEGIN PRIVATE KEY-----\nxxxxx\n-----END PRIVATE" # 描述： Google Cloud Vertex AI服务账号的私钥。格式： 一个JSON格式的字符串，包含服务账号的私钥信息。获取方式： 在Google Cloud Console中创建服务账号，生成JSON格式的密钥文件，然后将其内容设置为此环境变量的值。
-    client_email: xxxxxxxxxx@developer.gserviceaccount.com # 描述： Google Cloud Vertex AI服务账号的电子邮件地址。格式： 通常是形如 "[email protected]" 的字符串。获取方式： 在创建服务账号时生成，也可以在Google Cloud Console的"IAM与管理"部分查看服务账号详情获得。
+    client_email: xxxxxxxxxx@xxxxxxx.gserviceaccount.com # 描述： Google Cloud Vertex AI服务账号的电子邮件地址。格式： 通常是形如 "[email protected]" 的字符串。获取方式： 在创建服务账号时生成，也可以在Google Cloud Console的"IAM与管理"部分查看服务账号详情获得。
     model:
       - gemini-1.5-pro
       - gemini-1.5-flash
+      - claude-3-5-sonnet@20240620: claude-3-5-sonnet
+      - claude-3-opus@20240229: claude-3-opus
+      - claude-3-sonnet@20240229: claude-3-sonnet
+      - claude-3-haiku@20240307: claude-3-haiku
     tools: true
 
   - provider: other-provider

request.py

@@ -1,6 +1,6 @@
 import json
 from models import RequestModel
-from log_config import logger
+from utils import c35s, c3s, c3o, c3h, CircularList
 
 async def get_image_message(base64_image, engine = None):
     if "gpt" == engine:
@@ -222,19 +222,168 @@ def get_access_token(client_email, private_key):
         response.raise_for_status()
         return response.json()["access_token"]
 
-async def get_vertex_payload(request, engine, provider):
+async def get_vertex_gemini_payload(request, engine, provider):
     headers = {
         'Content-Type': 'application/json'
     }
     if provider.get("client_email") and provider.get("private_key"):
         access_token = get_access_token(provider['client_email'], provider['private_key'])
         headers['Authorization'] = f"Bearer {access_token}"
-    model = provider['model'][request.model]
+    if provider.get("project_id"):
+        project_id = provider.get("project_id")
+
     if request.stream:
         gemini_stream = "streamGenerateContent"
+    model = provider['model'][request.model]
+    location = CircularList(["us-central1", "us-east4", "us-west1", "us-west4", "europe-west1", "europe-west2"])
+    url = "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}:{stream}".format(LOCATION=location.next(), PROJECT_ID=project_id, MODEL_ID=model, stream=gemini_stream)
+
+    messages = []
+    systemInstruction = None
+    function_arguments = None
+    for msg in request.messages:
+        if msg.role == "assistant":
+            msg.role = "model"
+        tool_calls = None
+        if isinstance(msg.content, list):
+            content = []
+            for item in msg.content:
+                if item.type == "text":
+                    text_message = await get_text_message(msg.role, item.text, engine)
+                    content.append(text_message)
+                elif item.type == "image_url":
+                    image_message = await get_image_message(item.image_url.url, engine)
+                    content.append(image_message)
+        else:
+            content = [{"text": msg.content}]
+            tool_calls = msg.tool_calls
+
+        if tool_calls:
+            tool_call = tool_calls[0]
+            function_arguments = {
+                "functionCall": {
+                    "name": tool_call.function.name,
+                    "args": json.loads(tool_call.function.arguments)
+                }
+            }
+            messages.append(
+                {
+                    "role": "model",
+                    "parts": [function_arguments]
+                }
+            )
+        elif msg.role == "tool":
+            function_call_name = function_arguments["functionCall"]["name"]
+            messages.append(
+                {
+                    "role": "function",
+                    "parts": [{
+                    "functionResponse": {
+                        "name": function_call_name,
+                        "response": {
+                            "name": function_call_name,
+                            "content": {
+                                "result": msg.content,
+                            }
+                        }
+                    }
+                    }]
+                }
+            )
+        elif msg.role != "system":
+            messages.append({"role": msg.role, "parts": content})
+        elif msg.role == "system":
+            systemInstruction = {"parts": content}
+
+
+    payload = {
+        "contents": messages,
+        # "safetySettings": [
+        #     {
+        #         "category": "HARM_CATEGORY_HARASSMENT",
+        #         "threshold": "BLOCK_NONE"
+        #     },
+        #     {
+        #         "category": "HARM_CATEGORY_HATE_SPEECH",
+        #         "threshold": "BLOCK_NONE"
+        #     },
+        #     {
+        #         "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        #         "threshold": "BLOCK_NONE"
+        #     },
+        #     {
+        #         "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        #         "threshold": "BLOCK_NONE"
+        #     }
+        # ]
+        "generationConfig": {
+            "temperature": 0.5,
+            "max_output_tokens": 8192,
+            "top_k": 40,
+            "top_p": 0.95
+        },
+    }
+    if systemInstruction:
+        payload["system_instruction"] = systemInstruction
+
+    miss_fields = [
+        'model',
+        'messages',
+        'stream',
+        'tool_choice',
+        'temperature',
+        'top_p',
+        'max_tokens',
+        'presence_penalty',
+        'frequency_penalty',
+        'n',
+        'user',
+        'include_usage',
+        'logprobs',
+        'top_logprobs'
+    ]
+
+    for field, value in request.model_dump(exclude_unset=True).items():
+        if field not in miss_fields and value is not None:
+            if field == "tools":
+                payload.update({
+                    "tools": [{
+                        "function_declarations": [tool["function"] for tool in value]
+                    }],
+                    "tool_config": {
+                        "function_calling_config": {
+                            "mode": "AUTO"
+                        }
+                    }
+                })
+            else:
+                payload[field] = value
+
+    return url, headers, payload
+
+async def get_vertex_claude_payload(request, engine, provider):
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    if provider.get("client_email") and provider.get("private_key"):
+        access_token = get_access_token(provider['client_email'], provider['private_key'])
+        headers['Authorization'] = f"Bearer {access_token}"
     if provider.get("project_id"):
         project_id = provider.get("project_id")
-        url = "https://us-central1-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/us-central1/publishers/google/models/{MODEL_ID}:{stream}".format(PROJECT_ID=project_id, MODEL_ID=model, stream=gemini_stream)
+
+    model = provider['model'][request.model]
+    if "claude-3-5-sonnet" in model:
+        location = c35s
+    elif "claude-3-opus" in model:
+        location = c3o
+    elif "claude-3-sonnet" in model:
+        location = c3s
+    elif "claude-3-haiku" in model:
+        location = c3h
+
+    if request.stream:
+        claude_stream = "streamRawPredict"
+    url = "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:{stream}".format(LOCATION=location.next(), PROJECT_ID=project_id, MODEL=model, stream=claude_stream)
 
     messages = []
     systemInstruction = None
@@ -620,8 +769,10 @@ async def get_claude_payload(request, engine, provider):
 async def get_payload(request: RequestModel, engine, provider):
     if engine == "gemini":
         return await get_gemini_payload(request, engine, provider)
-    elif engine == "vertex":
-        return await get_vertex_payload(request, engine, provider)
+    elif engine == "vertex" and "gemini" in provider['model'][request.model]:
+        return await get_vertex_gemini_payload(request, engine, provider)
+    elif engine == "vertex" and "claude" in provider['model'][request.model]:
+        return await get_vertex_claude_payload(request, engine, provider)
     elif engine == "claude":
         return await get_claude_payload(request, engine, provider)
     elif engine == "gpt":

requirements.txt

@@ -1,5 +1,6 @@
-httpx[http2]
 pyyaml
 pytest
 uvicorn
-fastapi
+fastapi
+httpx[http2]
+cryptography

utils.py

@@ -185,4 +185,39 @@ def get_all_models(config):
                 }
                 all_models.append(model_info)
 
-    return all_models
+    return all_models
+
+# 【GCP-Vertex AI 目前有這些區域可用】 https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude?hl=zh_cn
+# c3.5s
+# us-east5
+# europe-west1
+
+# c3s
+# us-east5
+# us-central1
+# asia-southeast1
+
+# c3o
+# us-east5
+
+# c3h
+# us-east5
+# us-central1
+# europe-west1
+# europe-west4
+from collections import deque
+class CircularList:
+    def __init__(self, items):
+        self.queue = deque(items)
+
+    def next(self):
+        if not self.queue:
+            return None
+        item = self.queue.popleft()
+        self.queue.append(item)
+        return item
+
+c35s = CircularList(["us-east5", "europe-west1"])
+c3s = CircularList(["us-east5", "us-central1", "asia-southeast1"])
+c3o = CircularList(["us-east5"])
+c3h = CircularList(["us-east5", "us-central1", "europe-west1", "europe-west4"])