""" Title: Gemini Pipe / Choose Models Original Author: Coker (https://linux.do/u/coker/summary) Modified by: Barry Orlando (https://SolutionsHub.ai) Version: 1.2 License: MIT Description: Modified to return all Gemini models, not just Imagen models. """ import json import random import httpx from typing import List, AsyncGenerator, Callable, Awaitable from pydantic import BaseModel, Field import re class Pipe: class Valves(BaseModel): GOOGLE_API_KEYS: str = Field( default="", description="API Keys for Google, use , to split" ) BASE_URL: str = Field( default="https://generativelanguage.googleapis.com/v1beta", description="API Base Url", ) OPEN_SEARCH_INFO: bool = Field( default=True, description="Open search info show " ) IMAGE_NUM: int = Field(default=2, description="1-4") IMAGE_RATIO: str = Field( default="16:9", description="1:1, 3:4, 4:3, 16:9, 9:16" ) # --- NEW VALVE ADDED HERE --- ALLOWED_MODELS: str = Field( default="", description="Comma-separated list of Gemini model IDs to expose (e.g., gemini-pro,gemini-1.5-flash-latest,gemini-pro-vision). Leave blank to show all." ) # ---------------------------- def __init__(self): self.type = "manifold" self.name = "Google: " self.valves = self.Valves() self.OPEN_SEARCH_MODELS = ["gemini-2.5-pro-exp-03-25"] self.emitter = None self.open_search = False self.open_image = False self.open_think = False self.think_first = True def get_google_models(self) -> List[dict]: """ Return ALL Gemini models (not just Imagen models). • Pulls the full model catalog from Gemini. • Returns a list of all available models, with ID and display name. • All error text is in English. """ try: # Pick a key at random (you may have comma-separated keys) self.GOOGLE_API_KEY = random.choice( self.valves.GOOGLE_API_KEYS.split(",") ).strip() if not self.GOOGLE_API_KEY: return [{"id": "error", "name": "Error: API Key not found"}] url = f"{self.valves.BASE_URL}/models?key={self.GOOGLE_API_KEY}" response = httpx.get(url, timeout=10) response.raise_for_status() # raise on 4xx/5xx data = response.json() # Return all models models = [ { "id": mdl["name"].split("/")[-1], "name": mdl.get("displayName", mdl["name"].split("/")[-1]), } for mdl in data.get("models", []) ] if not models: return [{"id": "error", "name": "No models found"}] return models except Exception as exc: return [{"id": "error", "name": f"Could not fetch models: {exc}"}] async def emit_status( self, message: str = "", done: bool = False, ): if self.emitter: await self.emitter( { "type": "status", "data": { "description": message, "done": done, }, } ) def pipes(self) -> List[dict]: # --- MODIFIED PIPES METHOD --- all_models = self.get_google_models() allowed_models_str = self.valves.ALLOWED_MODELS.strip() if allowed_models_str: # Parse the comma-separated string into a list of IDs allowed_model_ids = [m.strip() for m in allowed_models_str.split(",") if m.strip()] # Filter the models based on the allowed_model_ids filtered_models = [] for model in all_models: if model["id"] in allowed_model_ids: filtered_models.append(model) return filtered_models else: # If the valve is empty, return all models (current default behavior) return all_models # ----------------------------- def create_search_link(self, idx, web): return f'\n{idx:02d}: [**{web["title"]}**]({web["uri"]})' def create_think_info(self, think_info): pass def _get_safety_settings(self, model: str): if ( model == "gemini-2.0-flash-exp" or model == "gemini-2.0-flash-exp-image-generation" ): return [ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"}, ] return [ {"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"}, {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}, ] def split_image(self, content): pattern = r"!\[image\]\(data:([^;]+);base64,([^)]+)\)" matches = re.findall(pattern, content) image_data_list = [] if matches: for mime_type, base64_data in matches: image_data_list.append({"mimeType": mime_type, "data": base64_data}) content = re.sub(r"!\[image\]\(data:[^;]+;base64,[^)]+\)", "", content) if not content: content = "请参考图片内容" return content, image_data_list def convert_message(self, message) -> dict: new_message = { "role": "model" if message["role"] == "assistant" else "user", "parts": [], } if isinstance(message.get("content"), str): if not message["role"] == "assistant": new_message["parts"].append({"text": message["content"]}) return new_message content, image_data_list = self.split_image(message["content"]) new_message["parts"].append({"text": content}) if image_data_list: for image_data in image_data_list: new_message["parts"].append( { "inline_data": { "mime_type": image_data["mimeType"], "data": image_data["data"], } } ) return new_message if isinstance(message.get("content"), list): for content in message["content"]: if content["type"] == "text": new_message["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] new_message["parts"].append( { "inline_data": { "mime_type": "image/jpeg", "data": image_data, } } ) return new_message async def do_parts(self, parts): res = "" if not parts or not isinstance(parts, list): return "Error: No parts found" for part in parts: if "text" in part: res += part["text"] if "inlineData" in part and part["inlineData"]: try: res += f'\n ![image](data:{part["inlineData"]["mimeType"]};base64,{part["inlineData"]["data"]}) \n' except: pass return res async def pipe( self, body: dict, __event_emitter__: Callable[[dict], Awaitable[None]] = None, ) -> AsyncGenerator[str, None]: self.emitter = __event_emitter__ self.GOOGLE_API_KEY = random.choice( self.valves.GOOGLE_API_KEYS.split(",") ).strip() self.base_url = self.valves.BASE_URL if not self.GOOGLE_API_KEY: yield "Error: GOOGLE_API_KEY is not set" return try: model_id = body["model"] if "." in model_id: model_id = model_id.split(".", 1)[1] if "imagen" in model_id: await self.emit_status(message="🐎 图像生成中……") async for res in self.gen_image( body["messages"][-1]["content"], model_id ): yield res return 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"]}] } continue contents.append(self.convert_message(message)) # if len(str(contents)) >1000: # yield "contents :"+str(contents)[:1000]+"......" request_data["contents"] = contents if model_id.endswith("-search"): model_id = model_id[:-7] request_data["tools"] = [{"googleSearch": {}}] self.open_search = True await self.emit_status(message="🔍 我好像在搜索……") elif "thinking" in model_id: await self.emit_status(message="🧐 我好像在思考……") self.open_think = True self.think_first = True elif model_id.endswith("-image-generation"): request_data["generationConfig"]["response_modalities"] = [ "Text", "Image", ] self.open_image = True # elif model_id.endswith("-image"): # model_id = model_id[:-6] # request_data["generationConfig"]["response_modalities"]=['Text', 'Image'] # self.open_image=True else: await self.emit_status(message="🚀 飞速生成中……") request_data["safetySettings"] = self._get_safety_settings(model_id) params = {"key": self.GOOGLE_API_KEY} if stream: url = f"{self.valves.BASE_URL}/models/{model_id}:streamGenerateContent" params["alt"] = "sse" else: url = f"{self.valves.BASE_URL}/models/{model_id}:generateContent" headers = {"Content-Type": "application/json"} async with httpx.AsyncClient() as client: if stream: async with client.stream( "POST", url, json=request_data, headers=headers, params=params, timeout=120, ) as response: if response.status_code != 200: error_content = await response.aread() yield f"Error: HTTP {response.status_code}: {error_content.decode('utf-8')}" await self.emit_status(message="❌ 生成失败", done=True) 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"]: try: parts = data["candidates"][0]["content"][ "parts" ] except: if ( "finishReason" in data["candidates"][0] and data["candidates"][0][ "finishReason" ] != "STOP" ): yield "\n---\n" + "异常结束: " + data[ "candidates" ][0]["finishReason"] return else: continue text = await self.do_parts(parts) yield text try: if ( self.open_search and self.valves.OPEN_SEARCH_INFO and data["candidates"][0][ "groundingMetadata" ]["groundingChunks"] ): yield "\n---------------------------------\n" groundingChunks = data["candidates"][0][ "groundingMetadata" ]["groundingChunks"] for idx, groundingChunk in enumerate( groundingChunks, 1 ): if "web" in groundingChunk: yield self.create_search_link( idx, groundingChunk["web"] ) except Exception as e: pass except Exception as e: # yield f"Error parsing stream: {str(e)}" pass await self.emit_status(message="🎉 生成成功", done=True) else: response = await client.post( url, json=request_data, headers=headers, params=params, timeout=120, ) if response.status_code != 200: yield f"Error: HTTP {response.status_code}: {response.text}" return data = response.json() res = "" if "candidates" in data and data["candidates"]: parts = data["candidates"][0]["content"]["parts"] res = await self.do_parts(parts) try: if ( self.open_search and self.valves.OPEN_SEARCH_INFO and data["candidates"][0]["groundingMetadata"][ "groundingChunks" ] ): res += "\n---------------------------------\n" groundingChunks = data["candidates"][0][ "groundingMetadata" ]["groundingChunks"] for idx, groundingChunk in enumerate( groundingChunks, 1 ): if "web" in groundingChunk: res += self.create_search_link( idx, groundingChunk["web"] ) except Exception as e: pass await self.emit_status( message="🎉 Generation Succeeded", done=True ) yield res else: yield "No response data" except Exception as e: yield f"Error: {str(e)}" await self.emit_status(message="❌ Generation Failed", done=True) async def gen_image(self, prompt: str, model: str) -> AsyncGenerator[str, None]: """ Generate an image using the specified model via the Google API. Yields status messages and, if successful, yields a markdown-formatted image with its size. """ url = f"{self.base_url}/models/{model}:predict" params = {"key": self.GOOGLE_API_KEY} headers = {"Content-Type": "application/json"} request_data = { "instances": [{"prompt": prompt}], "parameters": { "sampleCount": self.valves.IMAGE_NUM, "personGeneration": "allow_adult", "aspectRatio": self.valves.IMAGE_RATIO, }, } async with httpx.AsyncClient() as client: response = await client.post( url, json=request_data, headers=headers, params=params, timeout=120 ) if response.status_code != 200: yield f"Error: HTTP {response.status_code}: {response.text}" await self.emit_status(message="❌ Generation failed", done=True) return data = response.json() await self.emit_status(message="🎉 Generation successful", done=True) if "predictions" in data and isinstance(data["predictions"], list): yield f"Images generated: {len(data['predictions'])}\n\n" for index, prediction in enumerate(data["predictions"]): base64_str = prediction.get("bytesBase64Encoded") if base64_str: size_bytes = len(base64_str) * 3 / 4 if size_bytes >= 1024 * 1024: size = round(size_bytes / (1024 * 1024), 1) unit = "MB" else: size = round(size_bytes / 1024, 1) unit = "KB" yield f"Image {index+1} size: {size} {unit}\n" yield f'![image](data:{prediction["mimeType"]};base64,{base64_str}) \n\n' else: yield "No image data found"