diff --git a/examples/pipelines/providers/google_manifold_pipeline.py b/examples/pipelines/providers/google_manifold_pipeline.py index 9b15e4d..66778ee 100644 --- a/examples/pipelines/providers/google_manifold_pipeline.py +++ b/examples/pipelines/providers/google_manifold_pipeline.py @@ -1,4 +1,13 @@ -"""A manifold to integrate Google's GenAI models into Open-WebUI""" +""" +title: Google GenAI Manifold Pipeline +author: Marc Lopez (refactor by justinh-rahb) +date: 2024-06-06 +version: 1.1 +license: MIT +description: A pipeline for generating text using Google's GenAI models in Open-WebUI. +requirements: google-generativeai +environment_variables: GOOGLE_API_KEY +""" from typing import List, Union, Iterator import os @@ -6,6 +15,7 @@ import os from pydantic import BaseModel import google.generativeai as genai +from google.generativeai.types import GenerationConfig class Pipeline: @@ -72,56 +82,78 @@ class Pipeline: def pipe( self, user_message: str, model_id: str, messages: List[dict], body: dict ) -> Union[str, Iterator]: - """The pipe function (connects open-webui to google-genai) + if not self.valves.GOOGLE_API_KEY: + return "Error: GOOGLE_API_KEY is not set" - Args: - user_message (str): The last message input by the user - model_id (str): The model to use - messages (List[dict]): The chat history - body (dict): The raw request body in OpenAI's "chat/completions" style + try: + genai.configure(api_key=self.valves.GOOGLE_API_KEY) - Returns: - str: The complete response + if model_id.startswith("google_genai."): + model_id = model_id[12:] + model_id = model_id.lstrip(".") - Yields: - Iterator[str]: Yields a new message part every time it is received - """ + if not model_id.startswith("gemini-"): + return f"Error: Invalid model name format: {model_id}" - print(f"pipe:{__name__}") + print(f"Pipe function called for model: {model_id}") + print(f"Stream mode: {body.get('stream', False)}") - system_prompt = None - google_messages = [] - for message in messages: - google_role = "" - if message["role"] == "user": - google_role = "user" - elif message["role"] == "assistant": - google_role = "model" - elif message["role"] == "system": - system_prompt = message["content"] - continue # System promt is not inyected as a message - google_messages.append( - genai.protos.Content( - role=google_role, - parts=[ - genai.protos.Part( - text=message["content"], - ), - ], - ) + system_message = next((msg["content"] for msg in messages if msg["role"] == "system"), None) + + contents = [] + for message in messages: + if message["role"] != "system": + 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"]}] + }) + + if system_message: + contents.insert(0, {"role": "user", "parts": [{"text": f"System: {system_message}"}]}) + + model = genai.GenerativeModel(model_name=model_id) + + generation_config = GenerationConfig( + temperature=body.get("temperature", 0.7), + top_p=body.get("top_p", 0.9), + top_k=body.get("top_k", 40), + max_output_tokens=body.get("max_tokens", 8192), + stop_sequences=body.get("stop", []), ) - response = genai.GenerativeModel( - f"models/{model_id}", # we have to add the "models/" part again - system_instruction=system_prompt, - ).generate_content( - google_messages, - stream=body["stream"], - ) + safety_settings = body.get("safety_settings") - if body["stream"]: - for chunk in response: + response = model.generate_content( + contents, + generation_config=generation_config, + safety_settings=safety_settings, + stream=body.get("stream", False), + ) + + if body.get("stream", False): + return self.stream_response(response) + else: + return response.text + + except Exception as e: + print(f"Error generating content: {e}") + return f"An error occurred: {str(e)}" + + def stream_response(self, response): + for chunk in response: + if chunk.text: yield chunk.text - return "" - - return response.text