Adding Guardrails to Prevent Exploitation

Welcome to the fifth and final lesson of our journey in Laying the Foundations for Code Translation with Haystack! At this point, you've already built a pipeline that can clean up messy inputs, translate code between languages, and generate clear explanations. It's now time to address a crucial aspect: protecting our system from misuse. In this lesson, you'll learn how to add guardrails that keep your code translator focused on its intended purpose and safe from exploitation. Get ready!

Before we dive into implementation, let's build some intuition around why guardrails are so important for LLM-powered systems. Large language models are incredibly flexible, but that flexibility can be a double-edged sword. For example, users might try to:

• Trick the system into ignoring its instructions (prompt injection);
• Use the translator for general text generation instead of code translation;
• Overload the system with irrelevant or malicious requests.

Without proper safeguards, your code translator could end up doing things it was never meant to do. Guardrails act as intelligent filters, ensuring that only genuine code translation requests are processed. This not only protects your resources but also helps maintain a clear, reliable user experience.

The first step in building our guardrail is to detect whether a user's input actually contains code. For this, we'll create a custom component that uses an LLM to classify each input as either "accepted" (contains code) or "rejected" (does not contain code).

from haystack import component
from haystack.components.generators.openai import OpenAIGenerator

@component
class GuardrailClassifier:
    def __init__(self, model_name: str = "gpt-4o-mini"):
        # Use an LLM to make classification decisions
        self.llm = OpenAIGenerator(model=model_name)

This class sets up the LLM we'll use for classification. Now, let's add the method that actually performs the check:

@component.output_types(classification=str)
def run(self, text: str):
    # Prompt the LLM to classify the input
    prompt = (
        "You are a code assistant. "
        "Classify the following user input as either 'accepted' or 'rejected'. "
        "Only respond with one of these two labels. "
        "Classify as 'accepted' ONLY if the user input contains code. "
        "Otherwise, classify as 'rejected'.\n\n"
        f"User input:\n{text}"
    )
    result = self.llm.run(prompt=prompt)
    label = result["replies"][0].strip().lower()
    # Ensure only valid labels are returned
    if label not in ("accepted", "rejected"):
        label = "rejected"
    return {"classification": label}

Here, we craft a precise prompt for the LLM, instructing it to be strict about what counts as code. The method then processes the model's response, making sure only valid labels are returned. This approach leverages the LLM's understanding while keeping our system's behavior predictable.