Merge pull request #131 from justinh-rahb/gemini-fix

Fix: Google GenAI pipeline

examples/pipelines/providers/google_manifold_pipeline.py

@@ -1,11 +1,21 @@
-"""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
 
 from pydantic import BaseModel
 
 import google.generativeai as genai
+from google.generativeai.types import GenerationConfig
 
 
 class Pipeline:
@@ -72,56 +82,78 @@ def update_pipelines(self) -> None:
     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)
-
-        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
-
-        Returns:
-            str: The complete response
-
-        Yields:
-            Iterator[str]: Yields a new message part every time it is received
-        """
-
-        print(f"pipe:{__name__}")
-
-        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"],
-                        ),
-                    ],
-                )
+        if not self.valves.GOOGLE_API_KEY:
+            return "Error: GOOGLE_API_KEY is not set"
+
+        try:
+            genai.configure(api_key=self.valves.GOOGLE_API_KEY)
+
+            if model_id.startswith("google_genai."):
+                model_id = model_id[12:]
+            model_id = model_id.lstrip(".")
+
+            if not model_id.startswith("gemini-"):
+                return f"Error: Invalid model name format: {model_id}"
+
+            print(f"Pipe function called for model: {model_id}")
+            print(f"Stream mode: {body.get('stream', False)}")
+
+            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:
-                yield chunk.text
-            return ""
+            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
 
-        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