✨ Feature: Add feature: Support vertex ai API and vertex tool use invocation.

.gitignore

@@ -4,4 +4,5 @@ api.yaml
 __pycache__
 .vscode
 node_modules
-.wrangler
+.wrangler
+.pytest_cache

README.md

@@ -12,14 +12,15 @@
 
 ## Introduction
 
-这是一个统一管理大模型API的项目，可以通过一个统一的API接口调用多个后端服务，统一转换为 OpenAI 格式，支持负载均衡。目前支持的后端服务有：OpenAI、Anthropic、DeepBricks、OpenRouter、Gemini等。
+这是一个统一管理大模型API的项目，可以通过一个统一的API接口调用多个后端服务，统一转换为 OpenAI 格式，支持负载均衡。目前支持的后端服务有：OpenAI、Anthropic、DeepBricks、OpenRouter、Gemini、Vertex 等。
 
 ## Features
 
 - 统一管理多个后端服务
 - 支持负载均衡
+- 支持 OpenAI, Anthropic, Gemini, Vertex 函数调用
 - 支持多个模型
-- 支持多个API Key
+- 支持多个 API Key
 
 ## Configuration
 
@@ -48,9 +49,18 @@ providers:
     api: AIzaSyAN2k6IRdgw
     model:
       - gemini-1.5-pro
-      - gemini-1.5-flash
+      - gemini-1.5-flash-exp-0827: gemini-1.5-flash
     tools: false
 
+  - 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: [email protected] # 描述： Google Cloud Vertex AI服务账号的电子邮件地址。格式： 通常是形如 "[email protected]" 的字符串。获取方式： 在创建服务账号时生成，也可以在Google Cloud Console的"IAM与管理"部分查看服务账号详情获得。
+    model:
+      - gemini-1.5-pro
+      - gemini-1.5-flash
+    tools: true
+
   - provider: other-provider
     base_url: https://api.xxx.com/v1/messages
     api: sk-bNnAOJyA-xQw_twAA

main.py

@@ -21,7 +21,18 @@ from urllib.parse import urlparse
 async def lifespan(app: FastAPI):
     # 启动时的代码
     timeout = httpx.Timeout(connect=15.0, read=20.0, write=30.0, pool=30.0)
-    app.state.client = httpx.AsyncClient(timeout=timeout)
+    default_headers = {
+        "User-Agent": "curl/7.68.0",  # 模拟 curl 的 User-Agent
+        "Accept": "*/*",  # curl 的默认 Accept 头
+    }
+    app.state.client = httpx.AsyncClient(
+        timeout=timeout,
+        headers=default_headers,
+        http2=True,  # 禁用 HTTP/2
+        verify=True,  # 保持 SSL 验证（如需禁用，设为 False，但不建议）
+        follow_redirects=True,  # 自动跟随重定向
+    )
+    # app.state.client = httpx.AsyncClient(timeout=timeout)
     app.state.config, app.state.api_keys_db, app.state.api_list = await load_config(app)
     yield
     # 关闭时的代码
@@ -45,6 +56,8 @@ async def process_request(request: RequestModel, provider: Dict):
     engine = None
     if parsed_url.netloc == 'generativelanguage.googleapis.com':
         engine = "gemini"
+    elif parsed_url.netloc == 'aiplatform.googleapis.com':
+        engine = "vertex"
     elif parsed_url.netloc == 'api.anthropic.com' or parsed_url.path.endswith("v1/messages"):
         engine = "claude"
     elif parsed_url.netloc == 'openrouter.ai':
@@ -59,6 +72,7 @@ async def process_request(request: RequestModel, provider: Dict):
 
     if provider.get("engine"):
         engine = provider["engine"]
+
     logger.info(f"provider: {provider['provider']:<10} model: {request.model:<10} engine: {engine}")
 
     url, headers, payload = await get_payload(request, engine, provider)

request.py

@@ -165,10 +165,204 @@ async def get_gemini_payload(request, engine, provider):
 
     return url, headers, payload
 
-async def get_gpt_payload(request, engine, provider):
+import time
+import httpx
+import base64
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.serialization import load_pem_private_key
+
+def create_jwt(client_email, private_key):
+    # JWT Header
+    header = json.dumps({
+        "alg": "RS256",
+        "typ": "JWT"
+    }).encode()
+
+    # JWT Payload
+    now = int(time.time())
+    payload = json.dumps({
+        "iss": client_email,
+        "scope": "https://www.googleapis.com/auth/cloud-platform",
+        "aud": "https://oauth2.googleapis.com/token",
+        "exp": now + 3600,
+        "iat": now
+    }).encode()
+
+    # Encode header and payload
+    segments = [
+        base64.urlsafe_b64encode(header).rstrip(b'='),
+        base64.urlsafe_b64encode(payload).rstrip(b'=')
+    ]
+
+    # Create signature
+    signing_input = b'.'.join(segments)
+    private_key = load_pem_private_key(private_key.encode(), password=None)
+    signature = private_key.sign(
+        signing_input,
+        padding.PKCS1v15(),
+        hashes.SHA256()
+    )
+
+    segments.append(base64.urlsafe_b64encode(signature).rstrip(b'='))
+    return b'.'.join(segments).decode()
+
+def get_access_token(client_email, private_key):
+    jwt = create_jwt(client_email, private_key)
+
+    with httpx.Client() as client:
+        response = client.post(
+            "https://oauth2.googleapis.com/token",
+            data={
+                "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+                "assertion": jwt
+            },
+            headers={'Content-Type': "application/x-www-form-urlencoded"}
+        )
+        response.raise_for_status()
+        return response.json()["access_token"]
+
+async def get_vertex_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 request.stream:
+        gemini_stream = "streamGenerateContent"
+    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)
+
+    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_gpt_payload(request, engine, provider):
+    headers = {
+        'Content-Type': 'application/json',
+    }
     if provider.get("api"):
         headers['Authorization'] = f"Bearer {provider['api']}"
     url = provider['base_url']
@@ -426,6 +620,8 @@ 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 == "claude":
         return await get_claude_payload(request, engine, provider)
     elif engine == "gpt":

requirements.txt

@@ -1,4 +1,4 @@
-httpx
+httpx[http2]
 pyyaml
 pytest
 uvicorn

response.py

@@ -202,7 +202,7 @@ async def fetch_response(client, url, headers, payload):
 
 async def fetch_response_stream(client, url, headers, payload, engine, model):
     try:
-        if engine == "gemini":
+        if engine == "gemini" or engine == "vertex":
             async for chunk in fetch_gemini_response_stream(client, url, headers, payload, model):
                 yield chunk
         elif engine == "claude":

test/provider_test.py

@@ -74,6 +74,7 @@ def test_request_model(test_client, api_key, get_model):
     }
 
     headers = {
+        'Content-Type': 'application/json',
         "Authorization": f"Bearer {api_key}"
     }
 

test/test_httpx.py

@@ -0,0 +1,64 @@
+import httpx
+import asyncio
+import ssl
+import logging
+
+# 设置日志
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+async def make_request():
+    # SSL 上下文设置
+    # ssl_context = ssl.create_default_context()
+    # ssl_context.set_alpn_protocols(["h2", "http/1.1"])
+
+    # 创建自定义传输
+    transport = httpx.AsyncHTTPTransport(
+        http2=True,
+        # verify=ssl_context,
+        verify=False,
+        retries=1
+    )
+
+    # 设置头部
+    headers = {
+        "User-Agent": "curl/8.7.1",
+        "Accept": "*/*",
+        "Content-Type": "application/json",
+        "Authorization": "Bearer sk-xxxxxxx"
+    }
+
+    # 请求数据
+    data = {
+        "model": "gpt-4o",
+        "messages": [
+            {
+                "role": "user",
+                "content": "say test"
+            }
+        ],
+        "stream": True
+    }
+
+    async with httpx.AsyncClient(transport=transport) as client:
+        try:
+            response = await client.post(
+                "https://api.xxxxxxxxxx.me/v1/chat/completions",
+                headers=headers,
+                json=data,
+                timeout=30.0
+            )
+
+            logger.info(f"Status Code: {response.status_code}")
+            logger.info(f"Headers: {response.headers}")
+
+            # 处理流式响应
+            async for line in response.aiter_lines():
+                if line:
+                    print(line)
+
+        except httpx.RequestError as e:
+            logger.error(f"An error occurred while requesting {e.request.url!r}.")
+
+# 运行异步函数
+asyncio.run(make_request())

utils.py

@@ -13,6 +13,8 @@ def update_config(config_data):
             if type(model) == dict:
                 model_dict.update({new: old for old, new in model.items()})
         provider['model'] = model_dict
+        if provider.get('project_id'):
+            provider['base_url'] = 'https://aiplatform.googleapis.com/'
         config_data['providers'][index] = provider
     api_keys_db = config_data['api_keys']
     api_list = [item["api"] for item in api_keys_db]