codekingpro/portable-devtools

1import logging2from typing import Any, Callable, Dict, List, Optional3 4from langchain_core.callbacks import CallbackManagerForLLMRun5from langchain_core.language_models.llms import LLM6from pydantic import model_validator7 8logger = logging.getLogger(__name__)9 10 11def default_guardrail_violation_handler(violation: dict) -> str:12    """Default guardrail violation handler.13 14    Args:15        violation (dict): The violation dictionary.16 17    Returns:18        str: The canned response.19    """20    if violation.get("canned_response"):21        return violation["canned_response"]22    guardrail_name = (23        f"Guardrail {violation.get('offending_guardrail')}"24        if violation.get("offending_guardrail")25        else "A guardrail"26    )27    raise ValueError(28        f"{guardrail_name} was violated without a proper guardrail violation handler."29    )30 31 32class LayerupSecurity(LLM):33    """Layerup Security LLM service."""34 35    llm: LLM36    layerup_api_key: str37    layerup_api_base_url: str = "https://api.uselayerup.com/v1"38    prompt_guardrails: Optional[List[str]] = []39    response_guardrails: Optional[List[str]] = []40    mask: bool = False41    metadata: Optional[Dict[str, Any]] = {}42    handle_prompt_guardrail_violation: Callable[[dict], str] = (43        default_guardrail_violation_handler44    )45    handle_response_guardrail_violation: Callable[[dict], str] = (46        default_guardrail_violation_handler47    )48    client: Any  #: :meta private:49 50    @model_validator(mode="before")51    @classmethod52    def validate_layerup_sdk(cls, values: Dict[str, Any]) -> Any:53        try:54            from layerup_security import LayerupSecurity as LayerupSecuritySDK55 56            values["client"] = LayerupSecuritySDK(57                api_key=values["layerup_api_key"],58                base_url=values["layerup_api_base_url"],59            )60        except ImportError:61            raise ImportError(62                "Could not import LayerupSecurity SDK. "63                "Please install it with `pip install LayerupSecurity`."64            )65        return values66 67    @property68    def _llm_type(self) -> str:69        return "layerup_security"70 71    def _call(72        self,73        prompt: str,74        stop: Optional[List[str]] = None,75        run_manager: Optional[CallbackManagerForLLMRun] = None,76        **kwargs: Any,77    ) -> str:78        messages = [{"role": "user", "content": prompt}]79        unmask_response = None80 81        if self.mask:82            messages, unmask_response = self.client.mask_prompt(messages, self.metadata)83 84        if self.prompt_guardrails:85            security_response = self.client.execute_guardrails(86                self.prompt_guardrails, messages, prompt, self.metadata87            )88            if not security_response["all_safe"]:89                return self.handle_prompt_guardrail_violation(security_response)90 91        result = self.llm._call(92            messages[0]["content"], run_manager=run_manager, **kwargs93        )94 95        if self.mask and unmask_response:96            result = unmask_response(result)97 98        messages.append({"role": "assistant", "content": result})99 100        if self.response_guardrails:101            security_response = self.client.execute_guardrails(102                self.response_guardrails, messages, result, self.metadata103            )104            if not security_response["all_safe"]:105                return self.handle_response_guardrail_violation(security_response)106 107        return result108