LiteLLM

from __future__ import annotations

from typing import Any, List, Optional

import litellm
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth


class TrustGuard(CustomGuardrail):
    def __init__(
        self,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
        timeout: float = 10.0,
        fail_open: bool = False,
        scope: str = "current_turn",
        **kwargs: Any,
    ) -> None:
        self.api_base = (api_base or "").strip()
        self.api_key = (api_key or "").strip()
        self.timeout = float(timeout)
        self.fail_open = bool(fail_open)
        self.scope = scope
        self.http = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )
        # Sets guardrail_name and event_hook. Without it the proxy cannot match
        # this instance to the modes declared in config.yaml.
        super().__init__(**kwargs)

    async def _evaluate(
        self,
        payload: dict,
        direction: str,
        data: dict,
        user_api_key_dict: Optional[UserAPIKeyAuth],
    ) -> dict:
        body: dict = {"payload": payload, "direction": direction, "protocol": "llm"}

        # /v1/evaluate rejects unknown fields and empty identifiers, so only add
        # these when there is a real value.
        session_id = self._get_session_id_from_request_data(data)
        if session_id:
            body["session_id"] = str(session_id)
        consumer_id = self._consumer_id(data, user_api_key_dict)
        if consumer_id:
            body["consumer_id"] = str(consumer_id)

        try:
            response = await self.http.post(
                url=self.api_base,
                json=body,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
                timeout=self.timeout,
            )
        except Exception as exc:
            return self._unavailable(f"{type(exc).__name__}: {exc}")

        if response.status_code != 200:
            return self._unavailable(f"HTTP {response.status_code}: {response.text[:200]}")

        return response.json()

    def _unavailable(self, reason: str) -> dict:
        if self.fail_open:
            verbose_proxy_logger.warning(
                "TrustGuard unreachable, failing open (traffic NOT inspected): %s", reason
            )
            return {"status": "allow", "findings": []}
        raise HTTPException(
            status_code=503,
            detail={"error": "TrustGuard unavailable", "guardrail": self.guardrail_name},
        )

    def _block(self, result: dict) -> None:
        # 400 so the caller sees a client error. A bare exception would surface
        # as a 500 and look like an outage.
        raise HTTPException(
            status_code=400,
            detail={
                "error": "Blocked by TrustGuard",
                "guardrail": self.guardrail_name,
                "findings": result.get("findings"),
                "trace_id": result.get("trace_id"),
            },
        )

    @staticmethod
    def _consumer_id(
        data: dict, user_api_key_dict: Optional[UserAPIKeyAuth]
    ) -> Optional[str]:
        if user_api_key_dict is not None:
            for attr in ("key_alias", "user_email", "user_id", "team_alias"):
                value = getattr(user_api_key_dict, attr, None)
                if value:
                    return str(value)
        metadata = data.get("metadata") or data.get("litellm_metadata") or {}
        return metadata.get("user_api_key_alias") or metadata.get("user_api_key_user_id")

    def _in_scope(self, messages: Optional[list]) -> list:
        """Which messages to send for inspection."""
        if not messages:
            return []
        if self.scope == "transcript":
            return [m for m in messages if isinstance(m, dict)]

        last_user = -1
        for index, message in enumerate(messages):
            if isinstance(message, dict) and message.get("role") == "user":
                last_user = index
        if last_user < 0:
            return []
        # Assistant turns are dropped: they are model output and the output hook
        # already covers them. Tool results keep role="tool", which is what the
        # indirect prompt injection detector scopes itself to.
        return [
            message
            for message in messages[last_user:]
            if isinstance(message, dict) and message.get("role") in ("user", "tool")
        ]

    @staticmethod
    def _inspection_messages(messages: list) -> List[dict]:
        """Text-bearing messages to send for inspection."""
        inspected: List[dict] = []
        for message in messages:
            content = message.get("content")
            role = message.get("role") or "user"
            if isinstance(content, str) and content:
                inspected.append({"role": role, "content": content})
            elif isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and isinstance(part.get("text"), str):
                        inspected.append({"role": role, "content": part["text"]})
        return inspected

    def _log(self, direction: str, result: dict) -> None:
        findings = result.get("findings") or []
        detail = [
            "{}:{}/{}".format(
                (f.get("source") or {}).get("plugin")
                or (f.get("source") or {}).get("gate_name")
                or "?",
                (f.get("signal") or {}).get("type") or "-",
                (f.get("outcome") or {}).get("action") or "-",
            )
            for f in findings
        ]
        verbose_proxy_logger.info(
            "TrustGuard %s -> status=%s findings=[%s] trace_id=%s",
            direction,
            result.get("status"),
            ", ".join(detail),
            result.get("trace_id"),
        )

    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: DualCache,
        data: dict,
        call_type: str,
    ) -> Optional[dict]:
        messages = self._inspection_messages(self._in_scope(data.get("messages")))
        if not messages:
            return None

        payload = {"messages": messages}
        result = await self._evaluate(payload, "input", data, user_api_key_dict)
        self._log("input", result)

        status = result.get("status")
        if status == "block":
            self._block(result)
        if status == "transform":
            verbose_proxy_logger.warning(
                "TrustGuard returned transform for input, but this guardrail does not "
                "apply transformed_payload"
            )

        return data

    async def async_post_call_success_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        response: Any,
    ) -> Any:
        if not isinstance(response, litellm.ModelResponse):
            return None

        choices = [
            choice
            for choice in response.choices
            if isinstance(choice, litellm.Choices)
            and isinstance(getattr(choice.message, "content", None), str)
            and choice.message.content
        ]
        if not choices:
            return None

        text = "\n\n".join(choice.message.content for choice in choices)
        result = await self._evaluate({"input": text}, "output", data, user_api_key_dict)
        self._log("output", result)

        status = result.get("status")
        if status == "block":
            self._block(result)
        if status == "transform":
            verbose_proxy_logger.warning(
                "TrustGuard returned transform for output, but this guardrail does not "
                "apply transformed_payload"
            )

        return response