From 70d408da458716733d0d026f43f61a24df7440ad Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Sun, 2 Jun 2024 23:14:53 +0200 Subject: [PATCH] threshold fix --- ...mguard_prompt_injection_filter_pipeline.py | 83 ------------------- 1 file changed, 83 deletions(-) delete mode 100644 examples/llmguard_prompt_injection_filter_pipeline.py diff --git a/examples/llmguard_prompt_injection_filter_pipeline.py b/examples/llmguard_prompt_injection_filter_pipeline.py deleted file mode 100644 index d720e54..0000000 --- a/examples/llmguard_prompt_injection_filter_pipeline.py +++ /dev/null @@ -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