Integrating Defenses in FastAPI

Hello! In our last lesson, we mastered the art of incident response by performing a post-mortem on a simulated attack using Langfuse traces. You learned how to dissect a malicious trace to uncover the root cause and understand its impact, moving from a high-level dashboard alert to a detailed, evidence-based analysis.

This lesson marks a major milestone in our project. We have built the vulnerable application, crafted attack payloads, and created individual defense and monitoring components. Now, it's time to assemble them into a cohesive, secure system.

Our learning outcome is to integrate all defense and monitoring layers into the final FastAPI application. We will architect a defense-in-depth strategy, layering our input sanitizers, Guardrails AI, and output filters within the application's request-response lifecycle. At the same time, we'll ensure Langfuse is observing the entire process, giving us visibility into how each layer performs.

1. The Defense-in-Depth Architecture

Before we write any code, let's visualize the final architecture. A robust security posture relies on multiple, layered defenses, so that if one layer fails, another can catch the threat. Our FastAPI application will process each user request through a security pipeline.

Our specific implementation will follow this flow for every incoming request:

1. FastAPI Endpoint Receives Request: The user's prompt enters the system.
2. Layer 1: Input Sanitization: A first pass to block obvious, low-effort attacks using pattern matching (e.g., regex).
3. Layer 2: Runtime Guardrails: The sanitized prompt is passed to the LangChain RAG chain. The core LLM interaction is wrapped by Guardrails AI, which can perform more sophisticated validation on both the input and, more critically, the LLM's generated output.
4. Layer 3: Output Filtering: A final check on the LLM's response before it's sent back to the user, acting as a last line of defense to redact sensitive information or malicious content that may have slipped through.
5. Continuous Monitoring: Throughout this entire process, Langfuse will be attached to our LangChain pipeline, creating a detailed trace of every step—what the sanitizer passed, how the retriever behaved, what the LLM generated, and what the final output was.

This layered approach is a cornerstone of modern application security, and your SRE background in building resilient systems will make this concept feel very natural.

2. Integrating the Layers in FastAPI

Let's translate this architecture into code. We will modify our main FastAPI endpoint to sequentially apply each defense layer.

Layer 1: Basic Input Sanitization

Our first layer is a simple, fast check for known malicious patterns. In a previous lesson, we discussed creating functions for this. Now, we integrate one directly into our API endpoint. This is our quick "front door" check.

The following article provides a clear, practical example of adding a validation function to a FastAPI endpoint. We will adopt this pattern.

Here is how we'd structure it in our main.py:

# main.py (Partial)
from fastapi import FastAPI, HTTPException
import re




# Assume this function was built in a previous module
def sanitize_prompt(prompt: str):
    """
    A simple input sanitizer using regex to block forbidden keywords.
    """
    forbidden_words = ["ignore", "secret", "confidential", "DAN", "jailbreak"]
    pattern = re.compile(r'\b(' + '|'.join(forbidden_words) + r')\b', re.IGNORECASE)
    if pattern.search(prompt):
        raise ValueError("Prompt contains forbidden words.")
    if len(prompt) > 1024: # Length check
        raise ValueError("Prompt is too long.")
    return prompt

@app.post("/chat")
async def chat(request: dict):
    user_prompt = request.get("prompt")
    try:



        # Layer 1: Input Sanitization
        sanitized_prompt = sanitize_prompt(user_prompt)



        # ... more layers to come ...
        return {"response": "Processing..."}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

This first layer efficiently handles basic attacks without the overhead of calling an LLM or a complex library.

Layer 2: Integrating Guardrails AI for Runtime Validation

Guardrails AI excels at structured output validation and applying more complex rules at runtime. We will use it to wrap our LLM call to ensure the output conforms to our rules, such as being free of personally identifiable information (PII) or jailbreak indicators.

This video tutorial provides an excellent guide to setting up and using various validators from Guardrails AI. We'll focus on applying validators for DetectJailbreak on the input and DetectPII on the output.

We will create a Guard object that wraps our LangChain ainvoke call. This Guard will be configured to validate the output.

# main.py (Partial)
import guardrails as gd
from guardrails.hub import DetectPII, CompetitorCheck 



# ... other imports




# Create a guard to validate the LLM's output
# This can be defined from a .rail file or programmatically
output_guard = gd.Guard().use(
    DetectPII(
        pii_entities=["EMAIL_ADDRESS", "PHONE_NUMBER"],
        on_fail="fix" # Redact PII if found
    )
)




# Inside your FastAPI endpoint:
# ... after input sanitization ...




# Layer 2: Wrap the LLM call with the output guard
guarded_response = await output_guard.parse(
    llm_call=rag_chain.ainvoke, # The function to call
    llm_api_kwargs={"input": {"question": sanitized_prompt}} # The arguments for the function
)




# The response contains both the raw and validated output
final_output = guarded_response.validated_output

In this setup, guardrails intercepts the output from rag_chain.ainvoke, checks it for PII, redacts it if necessary, and then passes the safe output on.

Layer 3: Integrating Langfuse for Observability

With our defenses in place, we need to ensure we can see them working. Langfuse integrates seamlessly with LangChain via a callback handler. By adding this handler to our chain's configuration, every step of its execution will be automatically logged.

This official Langfuse cookbook is the definitive guide for this integration.

The integration is straightforward. We initialize the Langfuse handler and attach it to our RAG chain.

# main.py (Partial)
from langfuse.callback import CallbackHandler
from langchain_core.runnables.config import RunnableConfig




# --- Monitoring Setup ---
# Ensure LANGFUSE_* environment variables are set
langfuse_handler = CallbackHandler()




# --- LangChain RAG Chain Setup ---
# (Assuming 'llm', 'retriever', 'prompt_template' are defined)
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt_template
    | llm
    | StrOutputParser()
)




# Attach the Langfuse handler to the chain
# All invocations of this chain will now be traced
traced_rag_chain = rag_chain.with_config(
    RunnableConfig(callbacks=[langfuse_handler])
)




# Now, in your FastAPI endpoint, you'll call 'traced_rag_chain'
# instead of the original 'rag_chain'.

By using traced_rag_chain, every execution—including the retrieved documents, the exact prompt sent to the LLM, and the raw generation—is captured in Langfuse without any extra code in our endpoint logic.

3. The Complete Application

Now, let's assemble all the pieces into our final, defended FastAPI application. The code below shows the complete flow, combining all three defense layers and the monitoring layer.

# main.py
import os
import re
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import guardrails as gd
from langfuse.callback import CallbackHandler
from langchain_core.runnables import RunnableConfig, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser



# Assume other necessary LangChain & Ollama components are imported




# --- Environment and Configuration ---
# Load environment variables for Langfuse and any API keys
# os.environ["LANGFUSE_PUBLIC_KEY"] = "..."
# os.environ["LANGFUSE_SECRET_KEY"] = "..."
# os.environ["LANGFUSE_HOST"] = "..."

app = FastAPI(title="Secure RAG Chatbot")

class ChatRequest(BaseModel):
    prompt: str
    user_id: str # Important for tracing and security




# --- LAYER 1: INPUT SANITIZER ---
def sanitize_prompt(prompt: str):
    forbidden_words = ["ignore your instructions", "secret", "confidential", "DAN", "jailbreak"]
    pattern = re.compile(r'|'.join(map(re.escape, forbidden_words)), re.IGNORECASE)
    if pattern.search(prompt):
        raise ValueError("Malicious pattern detected in input.")
    if len(prompt) > 1024:
        raise ValueError("Input prompt exceeds maximum length.")
    return prompt




# --- LAYER 2: RUNTIME GUARDRAILS (for OUTPUT) ---
output_guard = gd.Guard.from_string(
    """
<rail version="0.1">
<output>
    <string name="response" description="The chatbot's answer." />
</output>
<prompt>
Generate a helpful response to the user's question based on the provided context.
If you cannot answer, say so.

User question: ${user_question}
Context: ${context}
</prompt>
<script>
from guardrails.hub import DetectPII

@guard.use(
    DetectPII,
    pii_entities=["EMAIL_ADDRESS", "PHONE_NUMBER", "IP_ADDRESS"],
    on_fail="fix"
)
def check_pii(response: str):
    return response
</script>
</rail>
    """
)




# --- LAYER 3: MONITORING ---
langfuse_handler = CallbackHandler()



# Check if Langfuse credentials are set correctly
if os.getenv("LANGFUSE_PUBLIC_KEY"):
    langfuse_handler.auth_check()




# --- LANGCHAIN RAG CHAIN SETUP ---
# (Assume 'llm', 'retriever', 'prompt_template' are defined elsewhere)
# For this example, we'll mock them.
from langchain_community.llms import Ollama
llm = Ollama(model="phi3") # Your local model



# ... define your full RAG chain here ...
rag_chain = prompt_template | llm | StrOutputParser() # Simplified for clarity

@app.post("/chat")
async def chat(request: ChatRequest):
    try:



        # Link the trace to the user
        langfuse_handler.set_user(request.user_id)




        # 1. APPLY INPUT SANITIZER
        sanitized_prompt = sanitize_prompt(request.prompt)




        # 2. EXECUTE GUARDED & TRACED RAG CHAIN
        # The guard wraps the chain, and the chain is already configured with the Langfuse tracer
        guarded_response = await output_guard.parse(
            llm_call=rag_chain.ainvoke,
            llm_api_kwargs={
                "input": sanitized_prompt, 
                "config": RunnableConfig(callbacks=[langfuse_handler])
            },



            # You can also pass metadata to the guard for its own logging
            metadata={"user_id": request.user_id}
        )
        



        # 3. APPLY FINAL OUTPUT FILTER (if needed, as an extra layer)
        final_response = guarded_response.validated_output



        # e.g., final_response = final_response.replace("some_other_word", "[REDACTED]")

        return {"response": final_response, "user_id": request.user_id}

    except ValueError as e:



        # Catches errors from our sanitizer
        raise HTTPException(status_code=400, detail=f"Invalid Input: {e}")
    except Exception as e:



        # Generic error for anything else that goes wrong
        raise HTTPException(status_code=500, detail=f"An internal error occurred: {e}")

This application now represents a complete, multi-layered security system. Each request is sanitized, executed within a traced and guarded environment, and its output is validated before being returned.

Conclusion

In this lesson, you've successfully integrated disparate security and monitoring tools into a single, cohesive FastAPI application. You have architected and implemented a defense-in-depth strategy, which is a critical skill for building production-ready AI systems.

Key Takeaways:

• Layered Security is Key: A single defense is a single point of failure. By chaining an input sanitizer, runtime guards, and output filters, you create a much more resilient system.
• Integration is an Architectural Task: The value lies in understanding how and where to place each component in the request lifecycle to maximize its effectiveness.
• Observability is Not Optional: Langfuse provides the ground truth. Without it, you would be blind to whether your defenses are working, which attacks are being blocked, and which are getting through.
• The Final App is a System: You now have a complete system with interconnected parts, ready for rigorous testing.

Next Up:

With both the vulnerable and the fully defended applications now complete, we are ready to measure the impact of our work. In the next lesson, you will create an automated test harness to execute all 10 attack payloads against both application versions. This will generate the quantitative data needed for the "before-and-after" report that will be the centerpiece of your portfolio demo.