title: Gemini description: "https://ai.google.dev/" author: DavidBueso author_url: "" # Version: 0.1.1 # Adjust version as needed license: "" # Gemini API license import os import json import requests from typing import List, Union, Generator, Iterator from pydantic import BaseModel, Field class Pipe: class Valves(BaseModel): API_KEY: str = Field( default=os.getenv("GEMINI_API_KEY", ""), description="Get your key here: https://ai.google.dev/", ) def __init__(self): # Validate API Key if not self.valves.API_KEY: raise ValueError("Gemini API key is required. Set GEMINI_API_KEY environment variable.") self.type = "manifold" self.id = "gemini" self.name = "gemini/" self.valves = self.Valves() self.base_url = "https://generativelanguage.googleapis.com/v1beta" self.model_map = { "gemini-pro": "models/gemini-pro", "gemini-pro-vision": "models/gemini-pro-vision", } def get_model_ids(self): return list(self.model_map.keys()) def pipes(self) -> List[dict]: model_ids = self.get_model_ids() return [{"id": model_id, "name": model_id} for model_id in model_ids] def _prepare_gemini_request(self, body: dict): """ Prepare the request body for Gemini API """ # Prepare contents contents = [] system_context = "" for message in body.get("messages", []): role = message.get("role") content = message.get("content") if role == "system": # Prepend system message as context system_context += f"System instruction: {content}\n" elif role == "user": # If system context exists, prepend it to the first user message if system_context and not contents: contents.append({ "role": "user", "parts": [{"text": system_context + content}] }) else: contents.append({ "role": "user", "parts": [{"text": content}] }) elif role == "assistant": contents.append({ "role": "model", "parts": [{"text": content}] }) # Prepare request body request_body = { "contents": contents, "generation_config": { "maxOutputTokens": body.get("max_tokens", 4096), "temperature": body.get("temperature", 0.7), "topP": body.get("top_p", 1.0), } } # Add tools if present if body.get("tools"): request_body["tools"] = [{"function_declarations": body["tools"]}] return request_body def pipe( self, body: dict, ) -> Union[str, Generator, Iterator]: # Validate model model_id = body["model"].removeprefix(self.id + ".") api_model_path = self.model_map.get(model_id) if not api_model_path: return f"Error: Model ID '{model_id}' is not supported." # Construct full URL url = f"{self.base_url}/{api_model_path}:streamGenerateContent" # Prepare request headers and parameters headers = { "Content-Type": "application/json", } params = { "key": self.valves.API_KEY } # Prepare Gemini-specific request body gemini_body = self._prepare_gemini_request(body) try: # Make the API request r = requests.post( url=url, params=params, headers=headers, json=gemini_body, stream=True, ) r.raise_for_status() # Streaming response if body.get("stream", False): def stream_response(): for line in r.iter_lines(decode_unicode=True): if line: try: chunk = json.loads(line) if "candidates" in chunk and chunk["candidates"]: for candidate in chunk["candidates"]: if "content" in candidate and "parts" in candidate["content"]: text = "".join([part.get("text", "") for part in candidate["content"]["parts"]]) if text: yield f"data: {json.dumps({'choices': [{'delta': {'content': text}}]})}\n\n" elif "toolCalls" in candidate: yield f"data: {json.dumps({'choices': [{'delta': {'tool_calls': candidate['toolCalls']}}]})}\n\n" except (json.JSONDecodeError, KeyError): continue yield "data: [DONE]\n\n" return stream_response() # Non-streaming response else: response_json = r.json() if "candidates" in response_json and response_json["candidates"]: first_candidate = response_json["candidates"][0] if "content" in first_candidate and "parts" in first_candidate["content"]: text = "".join([part.get("text", "") for part in first_candidate["content"]["parts"]]) return { "choices": [ { "message": { "content": text, "role": "assistant" } } ] } elif "toolCalls" in first_candidate: return { "choices": [ { "message": { "tool_calls": first_candidate["toolCalls"], "role": "assistant" } } ] } return {"error": "No valid response from Gemini"} except requests.exceptions.RequestException as e: return f"Error: {e}" except json.JSONDecodeError as e: return f"Error decoding response: {e}" # Version: 0.1.2