✨ Feature: Add feature: support Gemini API tool use

README.md

@@ -44,7 +44,7 @@ providers:
     tools: true
 
   - provider: gemini
-    base_url: https://generativelanguage.googleapis.com/v1/models/{model}:{stream}?key={api_key} # base_url 支持变量替换，{model} 会被替换为模型名称，{stream} 会被替换为 stream 参数，{api_key} 会被替换为 api_key 参数, 仅供 Gemini 模型使用，必填
+    base_url: https://generativelanguage.googleapis.com/v1beta # base_url 支持 v1beta/v1, 仅供 Gemini 模型使用，必填
     api: AIzaSyAN2k6IRdgw
     model:
       - gemini-1.5-pro

request.py

@@ -39,32 +39,70 @@ async def get_gemini_payload(request, engine, provider):
     headers = {
         'Content-Type': 'application/json'
     }
-    url = provider['base_url']
     model = provider['model'][request.model]
     if request.stream:
         gemini_stream = "streamGenerateContent"
-    url = url.format(model=model, stream=gemini_stream, api_key=provider['api'])
+    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'])
+    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'])
 
     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)
-                    # print("text_message", text_message)
                     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}]
-        if msg.role != "system":
+            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})
-        if msg.role == "system":
+        elif msg.role == "system":
             systemInstruction = {"parts": content}
 
 
@@ -96,7 +134,6 @@ async def get_gemini_payload(request, engine, provider):
         'model',
         'messages',
         'stream',
-        'tools',
         'tool_choice',
         'temperature',
         'top_p',
@@ -112,7 +149,19 @@ async def get_gemini_payload(request, engine, provider):
 
     for field, value in request.model_dump(exclude_unset=True).items():
         if field not in miss_fields and value is not None:
-            payload[field] = value
+            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
 

response.py

@@ -25,7 +25,7 @@ async def generate_sse_response(timestamp, model, content=None, tools_id=None, f
     if function_call_content:
         sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"function":{"arguments": function_call_content}}]}
     if tools_id and function_call_name:
-        sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"id":tools_id,"type":"function","function":{"name":function_call_name,"arguments":""}}]}
+        sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"id": tools_id,"type":"function","function":{"name": function_call_name, "arguments":""}}]}
         # sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"function":{"id": tools_id, "name": function_call_name}}]}
     if role:
         sample_data["choices"][0]["delta"] = {"role": role, "content": ""}
@@ -48,6 +48,9 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
                 error_json = error_str
             yield {"error": f"fetch_gpt_response_stream HTTP Error {response.status_code}", "details": error_json}
         buffer = ""
+        revicing_function_call = False
+        function_full_response = "{"
+        need_function_call = False
         async for chunk in response.aiter_text():
             buffer += chunk
             while "\n" in buffer:
@@ -63,18 +66,23 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
                     except json.JSONDecodeError:
                         logger.error(f"无法解析JSON: {line}")
 
-        # # 处理缓冲区中剩余的内容
-        # if buffer:
-        #     # print(buffer)
-        #     if '\"text\": \"' in buffer:
-        #         try:
-        #             json_data = json.loads(buffer)
-        #             content = json_data.get('text', '')
-        #             content = "\n".join(content.split("\\n"))
-        #             sse_string = await generate_sse_response(timestamp, model, content)
-        #             yield sse_string
-        #         except json.JSONDecodeError:
-        #             print(f"无法解析JSON: {buffer}")
+                if line and ('\"functionCall\": {' in line or revicing_function_call):
+                    revicing_function_call = True
+                    need_function_call = True
+                    if ']' in line:
+                        revicing_function_call = False
+                        continue
+
+                    function_full_response += line
+
+        if need_function_call:
+            function_call = json.loads(function_full_response)
+            function_call_name = function_call["functionCall"]["name"]
+            sse_string = await generate_sse_response(timestamp, model, content=None, tools_id="chatcmpl-9inWv0yEtgn873CxMBzHeCeiHctTV", function_call_name=function_call_name)
+            yield sse_string
+            function_full_response = json.dumps(function_call["functionCall"]["args"])
+            sse_string = await generate_sse_response(timestamp, model, content=None, tools_id="chatcmpl-9inWv0yEtgn873CxMBzHeCeiHctTV", function_call_name=None, function_call_content=function_full_response)
+            yield sse_string
 
 async def fetch_gpt_response_stream(client, url, headers, payload, max_redirects=5):
     redirect_count = 0

test/test_vertex copy.py

@@ -0,0 +1,190 @@
+import json
+import base64
+import time
+import httpx
+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"]
+
+def ask_stream(prompt, client_email, private_key, project_id, engine):
+    payload = {
+    "contents": [
+        {
+            "role": "user",
+            "parts": [
+                {
+                    "text": prompt
+                }
+            ]
+        }
+    ],
+    "system_instruction": {
+        "parts": [
+            {
+                "text": "You are Gemini, a large language model trained by Google. Respond conversationally"
+            }
+        ]
+    },
+    # "safety_settings": [
+    #     {
+    #         "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": 256,
+        "top_k": 40,
+        "top_p": 0.95
+    },
+    "tools": [
+        {
+            "function_declarations": [
+                {
+                    "name": "get_search_results",
+                    "description": "Search Google to enhance knowledge.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "prompt": {
+                                "type": "string",
+                                "description": "The prompt to search."
+                            }
+                        },
+                        "required": [
+                            "prompt"
+                        ]
+                    }
+                },
+                {
+                    "name": "get_url_content",
+                    "description": "Get the webpage content of a URL",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "url": {
+                                "type": "string",
+                                "description": "the URL to request"
+                            }
+                        },
+                        "required": [
+                            "url"
+                        ]
+                    }
+                }
+            ]
+        }
+    ],
+    "tool_config": {
+        "function_calling_config": {
+            "mode": "AUTO"
+        }
+    }
+}
+    # payload = {
+    #     "contents": [
+    #         {
+    #             "role": "user",
+    #             "parts": [
+    #                 {
+    #                     "text": prompt
+    #                 }
+    #             ]
+    #         },
+    #     ],
+    #     "generationConfig": {
+    #         "temperature": 0.2,
+    #         "maxOutputTokens": 256,
+    #         "topK": 40,
+    #         "topP": 0.95
+    #     }
+    # }
+
+    access_token = get_access_token(client_email, private_key)
+    headers = {
+        'Authorization': f"Bearer {access_token}",
+        'Content-Type': "application/json"
+    }
+
+    MODEL_ID = engine
+    PROJECT_ID = project_id
+    stream = "generateContent"
+    with httpx.Client() as client:
+        response = client.post(
+            f"https://us-central1-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/us-central1/publishers/google/models/{MODEL_ID}:{stream}",
+            json=payload,
+            headers=headers,
+            timeout=600,
+        )
+        response.raise_for_status()
+        return response.json()
+
+# 使用示例
+client_email, private_key, project_id = SERVICE_ACCOUNT_KEY["client_email"], SERVICE_ACCOUNT_KEY["private_key"], SERVICE_ACCOUNT_KEY["project_id"]
+engine = "gemini-1.5-pro"
+user_input = input("请输入您的问题： ")
+result = ask_stream(user_input, client_email, private_key, project_id, engine)
+print(json.dumps(result, ensure_ascii=False, indent=2))