python 使用 Google Gemini API

python 使用 Google Gemini API

• 注册APIKEY : Google AI Studio [免费]

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岁了。
这张图上面是蓝莓松饼、咖啡和花。
'''