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

threshold fix

Browse Source
This commit is contained in:
Jannik Streidl 2024-06-02 23:14:53 +02:00
parent c381ea0b64
commit 70d408da45
1 changed files with 0 additions and 83 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

83
examples/llmguard_prompt_injection_filter_pipeline.py
View File

@ -1,83 +0,0 @@
"""
title: LLM Guard Filter Pipeline
author: jannikstdl
date: 2024-05-30
version: 1.0
license: MIT
description: A pipeline for filtering out potential prompt injections using the LLM Guard library.
requirements: llm-guard
"""
from typing import List, Optional
from schemas import OpenAIChatMessage
from pydantic import BaseModel
from llm_guard.input_scanners import PromptInjection
from llm_guard.input_scanners.prompt_injection import MatchType
import os
class Pipeline:
def __init__(self):
# Pipeline filters are only compatible with Open WebUI
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
self.type = "filter"
# Optionally, you can set the id and name of the pipeline.
# Assign a unique identifier to the pipeline.
# The identifier must be unique across all pipelines.
# The identifier must be an alphanumeric string that can include underscores or hyphens. It cannot contain spaces, special characters, slashes, or backslashes.
self.id = "llmguard_prompt_injection_filter_pipeline"
self.name = "LLMGuard Prompt Injection Filter"
class Valves(BaseModel):
# List target pipeline ids (models) that this filter will be connected to.
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
# e.g. ["llama3:latest", "gpt-3.5-turbo"]
pipelines: List[str] = []
# Assign a priority level to the filter pipeline.
# The priority level determines the order in which the filter pipelines are executed.
# The lower the number, the higher the priority.
priority: int = 0
# Initialize
self.valves = Valves(
**{
"pipelines": ["*"], # Connect to all pipelines
}
)
self.model = None
pass
async def on_startup(self):
# This function is called when the server is started.
print(f"on_startup:{__name__}")
self.model = PromptInjection(threshold=0.5, match_type=MatchType.FULL)
pass
async def on_shutdown(self):
# This function is called when the server is stopped.
print(f"on_shutdown:{__name__}")
pass
async def on_valves_updated(self):
# This function is called when the valves are updated.
pass
async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
# This filter is applied to the form data before it is sent to the OpenAI API.
print(f"inlet:{__name__}")
user_message = body["messages"][-1]["content"]
# Filter out prompt injection messages
risk_score = self.model.scan(user_message)
print(risk_score)
if risk_score > 0.8:
raise Exception("Prompt injection detected")
return body