OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-11 08:01:08 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #191 from open-webui/gemini-permissive-safety

Browse Source 
Add Gemini Manifold Pipeline Safety Settings
This commit is contained in:
Justin Hayes 2024-08-01 14:29:52 -04:00 committed by GitHub
commit 342566765f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
examples/pipelines/providers/google_manifold_pipeline.py
View File

@ -2,7 +2,7 @@
title: Google GenAI Manifold Pipeline
author: Marc Lopez (refactor by justinh-rahb)
date: 2024-06-06
version: 1.1
version: 1.2
license: MIT
description: A pipeline for generating text using Google's GenAI models in Open-WebUI.
requirements: google-generativeai
@ -12,7 +12,7 @@ environment_variables: GOOGLE_API_KEY
from typing import List, Union, Iterator
import os
from pydantic import BaseModel
from pydantic import BaseModel, Field
import google.generativeai as genai
from google.generativeai.types import GenerationConfig
@ -25,13 +25,17 @@ class Pipeline:
"""Options to change from the WebUI"""
GOOGLE_API_KEY: str = ""
USE_PERMISSIVE_SAFETY: bool = Field(default=False)
def __init__(self):
self.type = "manifold"
self.id = "google_genai"
self.name = "Google: "
self.valves = self.Valves(**{"GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY", "")})
self.valves = self.Valves(**{
"GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY", ""),
"USE_PERMISSIVE_SAFETY": False
})
self.pipelines = []
genai.configure(api_key=self.valves.GOOGLE_API_KEY)
@ -135,7 +139,15 @@ class Pipeline:
stop_sequences=body.get("stop", []),
)
safety_settings = body.get("safety_settings")
if self.valves.USE_PERMISSIVE_SAFETY:
safety_settings = {
genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_NONE,
genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
}
else:
safety_settings = body.get("safety_settings")
response = model.generate_content(
contents,