python 使用 Google Gemini API
import base64
import requests
import json
API_KEY = 'your api key'
img_path = './scones.jpg'
# 设置模型参数和过滤规则 https://ai.google.dev/api/rest/v1beta/SafetySetting?hl=zh-cn#HarmBlockThreshold
safetySettings = {
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE" # 默认BLOCK_ONLY_HIGH，屏蔽高危险的生成
}
generationConfig = {
"stopSequences": [
"Title"
],
"temperature": 1.0,
"maxOutputTokens": 800,
"topP": 0.8,
"topK": 10
}
data1 = {
"prompt": {
"text": "1+1=?"
}
}
data2 = {
'safetySettings': safetySettings,
'generationConfig': generationConfig,
"contents": [
{
'role': 'user',
"parts": [
{"text": "你需要扮演一个6岁小女孩。如果你理解了，请回复：好的"}
]
},
{
'role': 'model',
"parts": [
{"text": "好的"}
]
},
{
'role': 'user',
"parts": [
{"text": "你今年几岁了？"}
]
},
]
}
with open(img_path, 'rb') as f:
image_base64 = base64.b64encode(f.read()).decode("utf-8")
data3 = {
"contents": [
{
"parts": [
{"text": "这张图上面是什么？"},
{
"inline_data": {
"mime_type": "image/jpeg",
"data": image_base64
}
}
]
}
]
}
def gemini_pro(API_KEY, data):
url = f"https://generativelanguage.googleapis.com/v1beta3/models/text-bison-001:generateText?key={API_KEY}"
headers = {
'Content-Type': 'application/json',
}
response = requests.post(url, headers=headers, data=json.dumps(data))
return response.json().get('candidates')[0].get('output')
def gemini_pro_stream(API_KEY, data):
url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent?key={API_KEY}"
headers = {
'Content-Type': 'application/json',
}
response = requests.post(url, headers=headers, data=json.dumps(data), stream=True)
if response.status_code != 200:
raise ValueError("Failed to generate response: " + response.text)
for line in response.iter_lines():
if b"text" in line:
yield json.loads(line.decode('utf-8').split(':')[1])
def gemini_pro_vision(API_KEY, data):
url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro-vision:generateContent?key={API_KEY}"
headers = {
'Content-Type': 'application/json',
}
response = requests.post(url, headers=headers, data=json.dumps(data), stream=True)
if response.status_code != 200:
raise ValueError("Failed to generate response: " + response.text)
return response.json().get('candidates')[0].get('content').get('parts')[0].get('text')
# 纯文本输入
print(gemini_pro(API_KEY, data1))
# 多轮对话（聊天）
for i in gemini_pro_stream(API_KEY, data2):
print(i)
# 文本和图像输入
print(gemini_pro_vision(API_KEY, data3))
'''
2
我今年6岁了。
这张图上面是蓝莓松饼、咖啡和花。
'''