LiteLLM integration

[robomotic]

Hello sir!
This work in progress: https://github.com/robomotic/litellm/tree/nova
Will let you know in due progress.

[robomotic]

Added you @fr0gger as collaborator in the repo too.

[fr0gger]

Excellent, I have a prototype for NOVA Litellm but it is not working yet. The best way to add it would be to use the code from the NOVA MCP and translate it to the guardrail features of Litellm :) https://github.com/fr0gger/nova_mcp

The code below is a prototype I was working on for LiteLLM not working correctly yet, and I don't have a lot of time to focus on it for now, maybe it can be useful.

"""
+-------------------------------------------------------------+
|         NovaGuard: Nova Security Guardrail for LiteLLM      |
+-------------------------------------------------------------+
|  Checks prompts against Nova rules and blocks on match.     |
|  LLM evaluator is selected from LiteLLM config via          |
|  the 'llm_model' parameter, which must match a model_name   |
|  in your model_list.                                        |
+-------------------------------------------------------------+
"""

import os
import sys
import logging
import json
from typing import Dict, List, Literal, Optional, Union

from fastapi import HTTPException

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
    CustomGuardrail,
    log_guardrail_information,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata

# Try importing Nova components
try:
    from nova.core.parser import NovaParser
    from nova.core.matcher import NovaMatcher
    from nova.evaluators.llm import (
        OpenAIEvaluator,
        AnthropicEvaluator,
        AzureOpenAIEvaluator,
        OllamaEvaluator,
        GroqEvaluator,
    )
except ImportError:
    raise ImportError(
        "Nova package not found. Please install Nova and ensure it is in your PYTHONPATH."
    )

GUARDRAIL_NAME = "nova_guard"

class NovaGuard(CustomGuardrail):
    def __init__(
        self,
        rules_dir: Optional[str] = None,
        log_file: Optional[str] = None,
        llm_model: Optional[str] = None,
        model_list: Optional[List[Dict]] = None,
        **kwargs,
    ):
        """
        rules_dir: Directory containing .nov Nova rule files
        log_file: Path to log file for security events
        llm_model: Name of the LLM to use for LLM-based rules (must match model_name in model_list)
        model_list: List of model configs from LiteLLM (injected by proxy)
        """
        self.rules_dir = rules_dir or os.path.join(os.getcwd(), "nova_rules")
        self.log_file = log_file or os.path.join(os.getcwd(), "logs", "nova_guard.log")
        self.llm_model = llm_model
        self.model_list = model_list or []
        self.rules = []
        self.llm_evaluator = None

        # Setup logging
        os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
        self.logger = logging.getLogger("nova_guard")
        self.logger.setLevel(logging.INFO)
        file_handler = logging.FileHandler(self.log_file)
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        if not self.logger.hasHandlers():
            self.logger.addHandler(file_handler)

        # Load rules
        self._load_rules()

        # Setup LLM evaluator
        self._setup_llm_evaluator()

        super().__init__(**kwargs)

    def _load_rules(self):
        self.rules = []
        if not os.path.isdir(self.rules_dir):
            self.logger.warning(f"Rules directory not found: {self.rules_dir}")
            return
        for root, _, files in os.walk(self.rules_dir):
            for file in files:
                if file.endswith('.nov'):
                    rule_path = os.path.join(root, file)
                    try:
                        with open(rule_path, 'r') as f:
                            rule_text = f.read()
                        parser = NovaParser()
                        rule = parser.parse(rule_text)
                        if rule:
                            self.rules.append((rule, rule_path))
                            self.logger.info(f"Loaded rule from {rule_path}")
                    except Exception as e:
                        self.logger.error(f"Error loading rule {rule_path}: {e}")

    def _setup_llm_evaluator(self):
        if not self.llm_model or not self.model_list:
            self.logger.warning("No llm_model or model_list provided; LLM-based rules will be skipped.")
            return
        # Find the model config in model_list
        model_cfg = next((m for m in self.model_list if m.get("model_name") == self.llm_model), None)
        if not model_cfg:
            self.logger.warning(f"llm_model '{self.llm_model}' not found in model_list; LLM-based rules will be skipped.")
            return
        params = model_cfg.get("litellm_params", {})
        # Instantiate the correct evaluator
        if "openai" in params.get("model", ""):
            api_key = os.environ.get(params.get("api_key", "").replace("os.environ/", ""), "")
            model = params.get("model")
            self.llm_evaluator = OpenAIEvaluator(api_key=api_key, model=model)
        elif "anthropic" in params.get("model", ""):
            api_key = os.environ.get(params.get("api_key", "").replace("os.environ/", ""), "")
            model = params.get("model")
            self.llm_evaluator = AnthropicEvaluator(api_key=api_key, model=model)
        elif "azure" in params.get("model", ""):
            api_key = os.environ.get(params.get("api_key", "").replace("os.environ/", ""), "")
            endpoint = params.get("api_base") or params.get("endpoint")
            deployment = params.get("deployment_name") or params.get("model")
            self.llm_evaluator = AzureOpenAIEvaluator(api_key=api_key, endpoint=endpoint, deployment_name=deployment)
        elif "ollama" in params.get("model", ""):
            host = os.environ.get(params.get("api_base", "").replace("os.environ/", ""), "")
            model = params.get("model")
            self.llm_evaluator = OllamaEvaluator(host=host, model=model)
        elif "groq" in params.get("model", ""):
            api_key = os.environ.get(params.get("api_key", "").replace("os.environ/", ""), "")
            model = params.get("model")
            self.llm_evaluator = GroqEvaluator(api_key=api_key, model=model)
        else:
            self.logger.warning(f"Unknown LLM provider for model: {params.get('model', '')}; LLM-based rules will be skipped.")

    def _find_matching_rule(self, prompt: str) -> Optional[Dict]:
        for rule, rule_path in self.rules:
            try:
                matcher = NovaMatcher(rule, llm_evaluator=self.llm_evaluator, create_llm_evaluator=False)
                result = matcher.check_prompt(prompt)
                if result.get("matched", False):
                    result["rule_file"] = rule_path
                    return result
            except Exception as e:
                self.logger.error(f"Error checking prompt against rule {getattr(rule, 'name', 'unknown')} in {rule_path}: {e}")
        return None

    def _log_event(self, event: Dict):
        try:
            self.logger.warning(json.dumps(event))
        except Exception as e:
            self.logger.error(f"Error logging event: {e}")

    async def _check(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        call_type: Literal[
            "completion",
            "text_completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "pass_through_endpoint",
            "rerank",
            "responses",
        ],
    ):
        if (
            await should_proceed_based_on_metadata(
                data=data,
                guardrail_name=GUARDRAIL_NAME,
            )
            is False
        ):
            return

        # Extract prompt text
        prompt = ""
        if "messages" in data and isinstance(data["messages"], list):
            prompt = " ".join([m.get("content", "") for m in data["messages"] if m.get("content")])
        elif "input" in data and isinstance(data["input"], str):
            prompt = data["input"]
        elif "input" in data and isinstance(data["input"], list):
            prompt = "\n".join(data["input"])

        if not prompt.strip():
            return  # Empty prompt, allow

        match = self._find_matching_rule(prompt)
        if match:
            # Log the match
            event = {
                "event": "SECURITY_ALERT",
                "rule_name": match.get("rule_name", "unknown"),
                "description": match.get("meta", {}).get("description", ""),
                "severity": match.get("meta", {}).get("severity", ""),
                "prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt,
                "user_id": getattr(user_api_key_dict, "user_id", "unknown"),
            }
            self._log_event(event)
            # Block the prompt
            raise HTTPException(
                status_code=400,
                detail={
                    "error": "Prompt blocked by nova_guard security policy.",
                    "rule_name": match.get("rule_name", "unknown"),
                    "description": match.get("meta", {}).get("description", ""),
                    "severity": match.get("meta", {}).get("severity", ""),
                },
            )
        # If no match, allow
        return None

    @log_guardrail_information
    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: litellm.DualCache,
        data: Dict,
        call_type: Literal[
            "completion",
            "text_completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "pass_through_endpoint",
            "rerank",
        ],
    ) -> Optional[Union[Exception, str, Dict]]:
        from litellm.types.guardrails import GuardrailEventHooks

        if self.event_hook is None:
            return None
        else:
            if (
                self.should_run_guardrail(
                    data=data, event_type=GuardrailEventHooks.pre_call
                )
                is not True
            ):
                return None

        return await self._check(
            data=data, user_api_key_dict=user_api_key_dict, call_type=call_type
        )

    @log_guardrail_information
    async def async_moderation_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        call_type: Literal[
            "completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "responses",
        ],
    ):
        from litellm.types.guardrails import GuardrailEventHooks

        if self.event_hook is None:
            return
        else:
            event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
            if self.should_run_guardrail(data=data, event_type=event_type) is not True:
                return

        return await self._check(
            data=data, user_api_key_dict=user_api_key_dict, call_type=call_type
        )

[robomotic]

Hello,
thanks for sharing, I could certainly re-use part of that code and the mcp.
I think the may issue for now is that because some rule could use an LLM call, we end up in inception mode where we then potentially have to use the registered LLM APIs within LiteLLM to perform the query.
For now what I will do is to keep it within the standard separate API then we can make it more sophisticated.
I am also listening to the LiteLLM developers on Discord for feedback.
Great.

[robomotic]

Getting there @fr0gger I worked out all the logic!

[robomotic]

@fr0gger it works!!!!!

[fr0gger]

Fantastic @robomotic you are a legend 🤩