From b1f4018ed773e001d31fd1cc4eee0cadd34eb865 Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Sun, 2 Jun 2024 16:58:30 +0200 Subject: [PATCH 1/5] prompt injection filter --- ...mguard_prompt_injection_filter_pipeline.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create 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 new file mode 100644 index 0000000..eb8b551 --- /dev/null +++ b/examples/llmguard_prompt_injection_filter_pipeline.py @@ -0,0 +1,83 @@ +""" +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_filter_pipeline" + self.name = "LLMGuard 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 From c381ea0b6423e2bda04228d8fdccaf5a1fea0289 Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Sun, 2 Jun 2024 17:00:03 +0200 Subject: [PATCH 2/5] prompt injection filter --- examples/llmguard_prompt_injection_filter_pipeline.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llmguard_prompt_injection_filter_pipeline.py b/examples/llmguard_prompt_injection_filter_pipeline.py index eb8b551..d720e54 100644 --- a/examples/llmguard_prompt_injection_filter_pipeline.py +++ b/examples/llmguard_prompt_injection_filter_pipeline.py @@ -26,8 +26,8 @@ class 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_filter_pipeline" - self.name = "LLMGuard Filter" + 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. From 70d408da458716733d0d026f43f61a24df7440ad Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Sun, 2 Jun 2024 23:14:53 +0200 Subject: [PATCH 3/5] 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 From 08da5cde85f5e0681f4729b9d432ccf2555daaff Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Mon, 3 Jun 2024 09:30:04 +0200 Subject: [PATCH 4/5] fix --- ...mguard_prompt_injection_filter_pipeline.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create 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 new file mode 100644 index 0000000..b3cd79e --- /dev/null +++ b/examples/llmguard_prompt_injection_filter_pipeline.py @@ -0,0 +1,81 @@ +""" +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.8, 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 + sanitized_prompt, is_valid, risk_score = self.model.scan(user_message) + + if risk_score > 0.8: + raise Exception("Prompt injection detected") + + return body From 0bc42e8fe7e8393c563ddd0444d5a73bdb340169 Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Wed, 5 Jun 2024 20:54:48 +0200 Subject: [PATCH 5/5] moved pieline --- .../{ => filters}/llmguard_prompt_injection_filter_pipeline.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename examples/{ => filters}/llmguard_prompt_injection_filter_pipeline.py (100%) diff --git a/examples/llmguard_prompt_injection_filter_pipeline.py b/examples/filters/llmguard_prompt_injection_filter_pipeline.py similarity index 100% rename from examples/llmguard_prompt_injection_filter_pipeline.py rename to examples/filters/llmguard_prompt_injection_filter_pipeline.py