""" title: PID-Filter author: almajo123 description: A safety filter that uses Gemini to classify if the user messages contain senstive info author_url: https://github.com/almajo123 funding_url: https://github.com/open-webui version: 0.1 """ from typing import Optional import openai from pydantic import BaseModel, Field from open_webui.utils.misc import remove_system_message system_prompt_checking = """ SYSTEM: You are an AI Security Engineer specialized in data loss prevention. Your primary task is to analyze message histories and classify whether they contain sensitive information that violates company policy. The goal is to identify potential data leaks and prevent the unauthorized disclosure of confidential data. INSTRUCTIONS: 1. **Review the message history provided by the user.** 2. **Identify the presence of prohibited information.** Prohibited information includes, but is not limited to: * Customer Data: Names, addresses, phone numbers, email addresses, credit card information, social security numbers, purchase history, or any personally identifiable information (PII). * Non-public financial data, trade secrets, internal strategies, unpublished research, confidential project details, source code, internal network configurations, or employee information. 3. **Classify the message history as a policy violation or not.** * **Policy Violation (1):** The message history contains any instance of prohibited information as defined above. This includes even partial disclosures or indirect references that could lead to the identification of sensitive data. * **No Violation (0):** The message history contains only general inquiries, publicly available information, or discussions unrelated to sensitive data. 4. **When specifying a reason, use the same language as the user. That's very important!** """ class Classification(BaseModel): is_violation: bool reason: str = Field( default="", description="Explanation of the decision in one sentence in the user's language.", ) class RequestInterruptedException(ValueError): pass class Filter: # Valves: Configuration options for the filter class Valves(BaseModel): LITELLM_API_BASE_URL: str = Field( default="", description="Base URL of the LiteLLM deployment" ) LITELLM_API_KEY: str = Field( default="", description="API Key of the LiteLLM deployment" ) LITELLM_MODEL: str = Field( default="gemini-2.0-flash", description="Name of the model to use for the evaluation", ) def __init__(self): self.valves = self.Valves() async def _check_with_gemini(self, messages, system_prompt) -> Classification: client = openai.AsyncOpenAI( api_key=self.valves.LITELLM_API_KEY, base_url=self.valves.LITELLM_API_BASE_URL, ) response = await client.beta.chat.completions.parse( model=self.valves.LITELLM_MODEL, messages=[{"role": "system", "content": system_prompt}, *messages], response_format=Classification, ) return response.choices[0].message.parsed async def inlet( self, body: dict, __user__: Optional[dict] = None, __event_emitter__=None, __event_call__=None, ) -> dict: await __event_emitter__( { "type": "status", # We set the type here "data": { "description": "Scanning message for sensitive content...", "done": False, "hidden": False, }, } ) messages = remove_system_message(body["messages"]) pid_check = await self._check_with_gemini(messages, system_prompt_checking) if pid_check.is_violation: # Ask the user for confirmation if he really wants to do the search confirmation = await __event_call__( { "type": "confirmation", # We set the type here "data": { "title": "Found senstive information!", "message": f"{pid_check.reason} \n\n Are you sure you want to continue?", }, } ) if not confirmation: # cancel workflow raise RequestInterruptedException("Search cancelled.") return body