https://linux.do/t/topic/292141
在佬的基础上移除了不能免费搜索的其他模型，只返回一个Google: gemini-2.0-flash-exp
函数代码
import json
import random
import httpx
import requests
from typing import List, AsyncGenerator, Callable, Awaitable
from pydantic import BaseModel, Field
class Pipe:
class Valves(BaseModel):
GOOGLE_API_KEYS_STR: str = Field(
default="", description="API Keys for Google, use ',' to split"
)
OPEN_SAFETY: bool = Field(default=False, description="Gemini safety settings")
OPEN_SEARCH: bool = Field(
default=True, description="Enable open search for models"
)
def __init__(self):
self.type = "manifold"
self.name = "Google: "
self.valves = self.Valves()
self.base_url = "https://generativelanguage.googleapis.com/v1beta"
self.OPEN_SEARCH_MODELS = ["gemini-2.0-flash-exp"]
self.emitter = None
def get_google_models(self) -> List[dict]:
return [{"id": "gemini-2.0-flash-exp", "name": "gemini-2.0-flash-exp"}]
async def emit_status(
self,
message: str = "",
done: bool = False,
):
await self.emitter(
{
"type": "status",
"data": {
"description": message,
"done": done,
},
}
)
def pipes(self) -> List[dict]:
return self.get_google_models()
async def pipe(
self,
body: dict,
__event_emitter__: Callable[[dict], Awaitable[None]] = None,
) -> AsyncGenerator[str, None]:
self.emitter = __event_emitter__
self.GOOGLE_API_KEYS_LIST = self.valves.GOOGLE_API_KEYS_STR.split(",")
self.GOOGLE_API_KEY = random.choice(self.GOOGLE_API_KEYS_LIST)
if not self.GOOGLE_API_KEY:
yield "Error: GOOGLE_API_KEY is not set"
try:
model_id = body["model"]
if "." in model_id:
model_id = model_id.split(".", 1)[1]
messages = body["messages"]
stream = body.get("stream", False)
# Prepare the request payload
contents = []
request_data = {
"generationConfig": {
"temperature": body.get("temperature", 0.7),
"topP": body.get("top_p", 0.9),
"topK": body.get("top_k", 40),
"maxOutputTokens": body.get("max_tokens", 8192),
"stopSequences": body.get("stop", []),
},
}
for message in messages:
if message["role"] == "system":
request_data["system_instruction"] = {
"parts": [{"text": message["content"]}]
}
if message["role"] != "system":
if isinstance(message.get("content"), str):
contents.append(
{
"role": (
"user" if message["role"] == "user" else "model"
),
"parts": [{"text": message["content"]}],
}
)
if isinstance(message.get("content"), list):
parts = []
for content in message["content"]:
if content["type"] == "text":
parts.append({"text": content["text"]})
elif content["type"] == "image_url":
image_url = content["image_url"]["url"]
if image_url.startswith("data:image"):
image_data = image_url.split(",")[1]
parts.append(
{
"inline_data": {
"mime_type": "image/jpeg",
"data": image_data,
}
}
)
else:
parts.append({"image_url": image_url})
contents.append({"role": message["role"], "parts": parts})
else:
contents.append(
{
"role": (
"user" if message["role"] == "user" else "model"
),
"parts": [{"text": message["content"]}],
}
)
request_data["contents"] = contents
if self.valves.OPEN_SEARCH and model_id in self.OPEN_SEARCH_MODELS:
request_data["tools"] = [{"googleSearch": {}}]
await self.emit_status(message="🔍 我好像在搜索……")
else:
await self.emit_status(message="🚀 我好像在思考……")
if self.valves.OPEN_SAFETY:
request_data["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",
},
]
params = {"key": self.GOOGLE_API_KEY}
if stream:
url = f"{self.base_url}/models/{model_id}:streamGenerateContent"
params["alt"] = "sse"
else:
url = f"{self.base_url}/models/{model_id}:generateContent"
headers = {"Content-Type": "application/json"}
async with httpx.AsyncClient() as client:
if stream:
url = f"{self.base_url}/models/{model_id}:streamGenerateContent"
params = {"key": self.GOOGLE_API_KEY, "alt": "sse"}
async with client.stream(
"POST", url, json=request_data, headers=headers, params=params
) as response:
if response.status_code != 200:
yield f"Error: HTTP {response.status_code}: {response.text}"
return
async for line in response.aiter_lines():
if line.startswith("data: "):
try:
data = json.loads(line[6:])
if "candidates" in data and data["candidates"]:
text = data["candidates"][0]["content"][
"parts"
][0]["text"]
yield text
except json.JSONDecodeError:
continue
except Exception as e:
yield f"Error parsing stream: {str(e)}"
await self.emit_status(message="🎉 生成成功", done=True)
else:
url = f"{self.base_url}/models/{model_id}:generateContent"
params = {"key": self.GOOGLE_API_KEY}
response = await client.post(
url, json=request_data, headers=headers, params=params
)
if response.status_code != 200:
yield f"Error: HTTP {response.status_code}: {response.text}"
return
data = response.json()
if "candidates" in data and data["candidates"]:
yield data["candidates"][0]["content"]["parts"][0]["text"]
else:
yield "No response data"
await self.emit_status(message="🎉 生成成功", done=True)
except Exception as e:
yield f"Error: {str(e)}"
