Dual-Layer Rate Limiting Middleware for Enterprise LangGraph RAG

Tuhin Paul
18h
143
0
0

Article

In the enterprise AI landscape of 2026, Large Language Models (LLMs) are no longer just a novelty; they are critical infrastructure. However, LLM inference is inherently expensive and computationally heavy. When you expose a multi-agent Retrieval-Augmented Generation (RAG) system via an API, you open the door to abuse, budget-draining prompt injection loops, and the "noisy neighbor" problem. If a single client spams your RAG endpoint, they can exhaust your LLM token quotas, spike your cloud compute costs, and degrade latency for all other users. This article provides an end-to-end guide to implementing a Dual-Layer Rate Limiting Middleware in FastAPI. We will enforce limits based on both Client IP (for anonymous/internal traffic) and API Key (for authenticated partners), returning standard 429 Too Many Requests responses. Finally, we will integrate this middleware with a production-grade Multi-Agent LangGraph RAG system featuring persistent memory and state.

Part 1: The Enterprise Use Case

Scenario: "Acme Corp" has deployed an internal/external AI Knowledge Assistant. The system uses a multi-agent LangGraph architecture to retrieve data from Confluence, Jira, and internal wikis, and synthesize answers.

The Access Tiers:

Authenticated Partners (API Key): Have a strict limit (e.g., 60 requests per minute) to prevent automated script abuse.
Internal Employees (Corporate IP): Have a higher limit (e.g., 120 requests per minute) as they are trusted internal users.
Anonymous Web Chat (Public IP): Have the lowest limit (e.g., 20 requests per minute) to prevent bot attacks.

The Architecture:

FastAPI Gateway: Handles HTTP routing, authentication extraction, and rate limiting.
Redis: Acts as the distributed, high-performance counter for rate limits.
LangGraph Backend: Executes the multi-agent RAG workflow, maintaining conversation state via a Checkpointer.

Part 2: Designing the Dual-Layer Rate Limiter

To build an enterprise-grade rate limiter, we must address three challenges:

Distributed State: If your FastAPI app runs on multiple replicas (Kubernetes pods), in-memory counters will fail. We must use Redis.
Race Conditions: A simple GET then SET in Redis is not atomic. We will use Redis INCR with an expiration to ensure thread safety.
Header Standards: When rejecting a request with a 429, we must return standard headers (X-RateLimit-Limit, X-RateLimit-Remaining, Retry-After) so client applications can handle the backoff gracefully.

Part 3: The Code Implementation

Below is the complete, end-to-end implementation.

1. Dependencies

pip install fastapi uvicorn redis langchain-openai langgraph langchain-core pydantic

2. The Redis Rate Limiter & FastAPI Middleware

import time
import logging
from typing import Optional, Tuple
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import redis.asyncio as redis

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize Redis connection (Enterprise: Use a connection pool)
redis_client = redis.from_url("redis://localhost:6379", decode_responses=True)

class DualLayerRateLimiter:
    """Handles the logic for checking and incrementing rate limits in Redis."""
    
    def __init__(self, redis_conn: redis.Redis):
        self.redis = redis_conn

    async def check_limit(self, identifier: str, limit: int, window_seconds: int) -> Tuple[bool, int, int]:
        """
        Checks the rate limit. Uses a sliding/fixed window approach.
        Returns: (is_allowed, remaining_requests, limit)
        """
        # Create a time-bounded key (Fixed window per minute)
        current_window = int(time.time() // window_seconds)
        redis_key = f"ratelimit:{identifier}:{current_window}"
        
        # Atomic increment. If key doesn't exist, it's created with value 1.
        current_count = await self.redis.incr(redis_key)
        
        # Set expiration only on the first request of the window to prevent race conditions
        if current_count == 1:
            await self.redis.expire(redis_key, window_seconds + 1) # +1s buffer
            
        is_allowed = current_count <= limit
        remaining = max(0, limit - current_count)
        
        return is_allowed, remaining, limit

class RateLimitMiddleware(BaseHTTPMiddleware):
    """
    FastAPI Middleware that intercepts requests, extracts IP/API Key,
    and enforces dual-layer rate limiting.
    """
    def __init__(self, app, limiter: DualLayerRateLimiter):
        super().__init__(app)
        self.limiter = limiter
        # Define tiers: (limit per minute)
        self.api_key_limit = 60
        self.internal_ip_limit = 120
        self.public_ip_limit = 20
        self.window = 60 # 60 seconds

    async def dispatch(self, request: Request, call_next):
        # Skip rate limiting for health checks and docs
        if request.url.path in ["/health", "/docs", "/openapi.json"]:
            return await call_next(request)

        # 1. Extract Client IP (Handling reverse proxies like AWS ALB / Nginx)
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            client_ip = forwarded_for.split(",")[0].strip()
        else:
            client_ip = request.client.host if request.client else "unknown"

        # 2. Extract API Key
        api_key = request.headers.get("X-API-Key") or request.headers.get("Authorization", "").replace("Bearer ", "")

        # 3. Determine Identifier and Limit Tier
        if api_key and api_key != "":
            identifier = f"key:{api_key}"
            limit = self.api_key_limit
            tier = "API_KEY"
        elif client_ip.startswith("10.") or client_ip.startswith("192.168."): # Mock internal IP check
            identifier = f"ip:{client_ip}"
            limit = self.internal_ip_limit
            tier = "INTERNAL_IP"
        else:
            identifier = f"ip:{client_ip}"
            limit = self.public_ip_limit
            tier = "PUBLIC_IP"

        # 4. Check Limit
        is_allowed, remaining, max_limit = await self.limiter.check_limit(identifier, limit, self.window)

        if not is_allowed:
            logger.warning(f"Rate limit exceeded for {tier} {identifier}. Limit: {max_limit}")
            return JSONResponse(
                status_code=429,
                content={"error": "Too Many Requests", "message": "Rate limit exceeded. Please slow down."},
                headers={
                    "X-RateLimit-Limit": str(max_limit),
                    "X-RateLimit-Remaining": "0",
                    "Retry-After": str(self.window),
                    "X-RateLimit-Tier": tier
                }
            )

        # 5. Proceed and inject rate limit headers into the successful response
        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(max_limit)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Tier"] = tier
        return response

3. The Multi-Agent LangGraph RAG Backend

Now, we build the AI backend. We will use a multi-agent approach: a Router Agent directs the query to a RAG Retriever Agent, which passes context to a Synthesizer Agent. We use LangGraph's Checkpointer for memory.

from typing import TypedDict, Annotated, Sequence, Literal
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.tools import tool

# --- 1. Define the State ---
class GraphState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], "The conversation history"]
    user_id: str
    retrieved_context: str
    next_agent: str

# --- 2. Define the Tools (RAG) ---
@tool
def search_internal_knowledge_base(query: str) -> str:
    """Searches the enterprise vector database for internal policies and documentation."""
    # In production, this queries Pinecone/Milvus/Weaviate
    return "Acme Corp Policy: Employees are entitled to 25 days of PTO. Rollover is permitted up to 5 days."

# --- 3. Define the Agent Nodes ---
llm = ChatOpenAI(model="gpt-4o", temperature=0)

async def router_node(state: GraphState):
    """Determines if the query requires RAG or can be answered directly."""
    # Simplified routing logic for demonstration
    messages = state["messages"]
    last_msg = messages[-1].content.lower()
    
    if "policy" in last_msg or "pto" in last_msg or "hr" in last_msg:
        return {"next_agent": "rag_retriever"}
    return {"next_agent": "synthesizer"}

async def rag_retriever_node(state: GraphState):
    """Executes the RAG tool and stores context in state."""
    last_msg = state["messages"][-1].content
    # In a real agent, the LLM would decide the tool arguments
    context = search_internal_knowledge_base.invoke({"query": last_msg})
    return {"retrieved_context": context, "next_agent": "synthesizer"}

async def synthesizer_node(state: GraphState):
    """Synthesizes the final answer using the LLM and retrieved context."""
    context = state.get("retrieved_context", "No specific context retrieved.")
    messages = state["messages"]
    
    prompt = [
        SystemMessage(content=f"You are an enterprise AI assistant. Use the following context to answer the user:\n\nContext: {context}\n\nIf the context doesn't contain the answer, state that you don't know."),
        *messages
    ]
    
    response = await llm.ainvoke(prompt)
    return {"messages": [response], "next_agent": "end"}

# --- 4. Build and Compile the Graph ---
def route_next(state: GraphState) -> Literal["rag_retriever", "synthesizer", "end"]:
    return state["next_agent"]

workflow = StateGraph(GraphState)
workflow.add_node("router", router_node)
workflow.add_node("rag_retriever", rag_retriever_node)
workflow.add_node("synthesizer", synthesizer_node)

workflow.set_entry_point("router")
workflow.add_conditional_edges("router", route_next)
workflow.add_edge("rag_retriever", "synthesizer")
workflow.add_edge("synthesizer", END)

# Initialize Memory (Enterprise: Use AsyncPostgresSaver)
memory = MemorySaver()
rag_graph = workflow.compile(checkpointer=memory)

4. Tying it Together: The FastAPI App

from pydantic import BaseModel

app = FastAPI(title="Enterprise RAG Gateway")

# Initialize Limiter and Middleware
limiter = DualLayerRateLimiter(redis_client)
app.add_middleware(RateLimitMiddleware, limiter=limiter)

class ChatRequest(BaseModel):
    message: str
    user_id: str
    thread_id: str # Used for LangGraph memory checkpointing

@app.post("/v1/chat")
async def chat_endpoint(payload: ChatRequest, request: Request):
    """
    The main endpoint. Protected by the RateLimitMiddleware.
    Invokes the LangGraph asynchronously.
    """
    # Extract user identity for LangGraph state
    user_id = payload.user_id
    thread_id = payload.thread_id
    
    # Prepare LangGraph config for memory persistence
    config = {"configurable": {"thread_id": thread_id}}
    
    # Initial state input
    inputs = {
        "messages": [HumanMessage(content=payload.message)],
        "user_id": user_id,
        "retrieved_context": "",
        "next_agent": ""
    }
    
    try:
        # Invoke the multi-agent graph asynchronously
        final_state = await rag_graph.ainvoke(inputs, config)
        
        # Extract the final AI response
        ai_message = final_state["messages"][-1].content
        
        return {
            "response": ai_message,
            "thread_id": thread_id,
            "metadata": {
                "agents_used": ["router", "rag_retriever", "synthesizer"],
                "user_id": user_id
            }
        }
    except Exception as e:
        logger.error(f"LangGraph execution failed: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal AI processing error")

@app.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": time.time()}

Part 4: Testing the Implementation

To test this, run the FastAPI app (uvicorn main:app --reload) and use curl to simulate traffic.

Test 1: Normal Request (API Key)

curl -X POST "http://localhost:8000/v1/chat" \
-H "X-API-Key: partner_key_123" \
-H "Content-Type: application/json" \
-d '{"message": "What is the PTO policy?", "user_id": "user_1", "thread_id": "thread_99"}' \
-i

Expected Response: 200 OK with headers X-RateLimit-Limit: 60 and X-RateLimit-Remaining: 59. The AI will correctly retrieve and synthesize the PTO policy using the RAG agent.

Test 2: Triggering the 429 Rate Limit

Run a quick loop to exhaust the public IP limit (20 requests/minute).

for i in {1..25}; do
  curl -s -o /dev/null -w "%{http_code}\n" -X POST "http://localhost:8000/v1/chat" \
  -H "Content-Type: application/json" \
  -d '{"message": "Hello", "user_id": "user_2", "thread_id": "thread_100"}'
done

Expected Output: The first 20 requests will return 200. The remaining 5 will return 429.

If you inspect the headers of the 429 response:

HTTP/1.1 429 Too Many Requests
X-RateLimit-Limit: 20
X-RateLimit-Remaining: 0
Retry-After: 60
X-RateLimit-Tier: PUBLIC_IP
Content-Type: application/json

{"error": "Too Many Requests", "message": "Rate limit exceeded. Please slow down."}

Test 3: Verifying Memory and State

Send a follow-up message using the same thread_id to prove the LangGraph Checkpointer is maintaining state across requests.

curl -X POST "http://localhost:8000/v1/chat" \
-H "X-API-Key: partner_key_123" \
-H "Content-Type: application/json" \
-d '{"message": "Can I rollover more than 5 days?", "user_id": "user_1", "thread_id": "thread_99"}'

Expected Response: The AI will remember the previous context about the 5-day rollover limit without needing the RAG tool to fetch it again, demonstrating successful state persistence.

Part 5: Enterprise Best Practices & Next Steps

Implementing this architecture gives you a robust, production-ready AI gateway. However, to fully harden this for a Fortune 500 deployment, consider the following enhancements:

Atomic Lua Scripts in Redis: While INCR and EXPIRE work well, a strict enterprise environment should use a Redis Lua script to perform the check-and-set atomically, preventing edge-case race conditions during high-concurrency spikes.
Token-Based Rate Limiting: Instead of limiting by requests, limit by LLM tokens. A request that generates 10,000 tokens should cost more against the rate limit than a request generating 100 tokens. You can intercept the LangGraph output, calculate token usage, and decrement a Redis token bucket.
Distributed Checkpointing: In the code above, we used MemorySaver for brevity. In production, you must swap this for AsyncPostgresSaver or AsyncRedisSaver so that conversation memory survives FastAPI pod restarts and scales horizontally.
Circuit Breakers: Combine this rate limiter with a circuit breaker (like pybreaker). If the underlying vector database or LLM provider goes down, the circuit breaker should immediately return a 503 without wasting time executing the LangGraph nodes.

By combining FastAPI's dual-layer rate limiting with LangGraph's stateful multi-agent execution, you ensure that your enterprise AI is not only intelligent and context-aware but also financially predictable and highly resilient to abuse.