Designing Idempotent Consumers in Event-Driven AI Systems

Tuhin Paul
5d
130
0
1

Article

In event-driven architectures, message brokers like Amazon SQS or Kafka guarantee at-least-once delivery. This means that if a network blip occurs, a timeout happens, or a worker crashes, the broker will redeliver the message.

In traditional systems, a duplicate message might just result in a harmless database unique constraint violation. But in Enterprise Multi-Agent LLM/RAG systems, a duplicate message is a disaster. It means you pay for expensive LLM inferences twice, you might trigger duplicate external API calls, and you risk corrupting your application state.

To survive the "at-least-once" reality, you must design Idempotent Consumers.

In this article, we will explore how to design idempotent consumers, apply it to a real-world enterprise use case, and implement it end-to-end using LangGraph, RAG, and AWS DynamoDB.

Part 1: The Core Concepts of Idempotency

Idempotency means that applying the same operation multiple times yields the exact same result as applying it once.

To achieve this in an event-driven consumer, we rely on three pillars:

Idempotency Keys: A unique identifier for the business operation (e.g., client_id + event_type), usually passed in the message headers.
Idempotency Store: A fast, durable database (like DynamoDB or Redis) used to track the state of the key (IN_PROGRESS, COMPLETED, FAILED).
Conditional Writes / Locking: Ensuring that only one worker can "claim" and process the idempotency key at a time, preventing race conditions.

Why is this critical for LLM/RAG Agents?

Cost Control: LLM API calls cost money. If a message is processed twice, you pay twice.
Side-Effect Prevention: If an agent sends an email or creates a bank account, doing it twice breaks the business logic.
State Corruption: RAG systems often update vector stores or relational databases. Duplicate writes can create duplicate embeddings or orphaned records.

Part 2: The Use Case - Enterprise Client Onboarding & KYC

The Scenario: A digital bank receives a new client onboarding request. The payload includes the client's details and a reference to uploaded KYC documents (PDFs).

This requires a Multi-Agent LangGraph RAG pipeline:

Document Extraction Agent (RAG): Retrieves information from the uploaded PDFs using a Vector DB.
KYC Compliance Agent (RAG): Queries the internal compliance policy Vector DB to check for sanctions or PEP (Politically Exposed Person) flags.
Account Provisioning Agent: Calls the core banking API to create the actual account.

The Failure Scenario:
The Account Provisioning Agent successfully creates the account in the core banking system. However, right before it can delete the message from SQS, the container crashes. SQS redelivers the message. Without idempotency, the system creates a duplicate bank account.

Part 3: End-to-End Code Implementation

Prerequisites: boto3, langgraph, langchain, langchain-openai

Step 1: Provisioning the Idempotency Store (DynamoDB)

We use DynamoDB because its conditional writes are perfect for distributed locking and idempotency tracking.

import boto3
import json
import time
import uuid
from typing import TypedDict, Annotated, Sequence
from langgraph.graph import StateGraph, END
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
import operator

dynamodb = boto3.resource('dynamodb')
sqs = boto3.client('sqs')

# Create Idempotency Table
table = dynamodb.create_table(
    TableName='AgentIdempotencyStore',
    KeySchema=[{'AttributeName': 'IdempotencyKey', 'KeyType': 'HASH'}],
    AttributeDefinitions=[{'AttributeName': 'IdempotencyKey', 'AttributeType': 'S'}],
    BillingMode='PAY_PER_REQUEST'
)
table.wait_until_exists()

# Create SQS Queue
queue = sqs.create_queue(QueueName='client-onboarding-queue')
queue_url = queue['QueueUrl']

Step 2: Defining the Multi-Agent LangGraph RAG

We define the state and the agents. Notice how we include the idempotency_key in the state.

# 1. Define the Shared State
class OnboardingState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    client_id: str
    idempotency_key: str
    extracted_data: dict
    compliance_status: str
    account_created: bool

# 2. Define the Agent Nodes (Mocked RAG & API calls for brevity)
def document_extraction_agent(state: OnboardingState):
    """RAG Agent: Queries Vector DB for document data"""
    print(f"[Doc Agent] Retrieving data for client {state['client_id']} from Vector DB...")
    # Mocking RAG retrieval
    extracted = {"name": "John Doe", "passport": "A123456", "risk_level": "low"}
    return {"extracted_data": extracted, "messages": [AIMessage(content="Docs extracted.")]}

def kyc_compliance_agent(state: OnboardingState):
    """RAG Agent: Queries Vector DB for compliance policies"""
    print(f"[KYC Agent] Checking compliance policies in Vector DB...")
    # Mocking RAG retrieval
    return {"compliance_status": "CLEARED", "messages": [AIMessage(content="KYC cleared.")]}

def account_provisioning_agent(state: OnboardingState):
    """Action Agent: Creates account in Core Banking System"""
    print(f"[Provisioning Agent] Creating account for {state['client_id']} in Core Banking...")
    # CRITICAL: This is the side-effect that MUST NOT run twice!
    time.sleep(1) # Simulate API latency
    return {"account_created": True, "messages": [AIMessage(content="Account created.")]}

# 3. Build the LangGraph
workflow = StateGraph(OnboardingState)
workflow.add_node("extract", document_extraction_agent)
workflow.add_node("compliance", kyc_compliance_agent)
workflow.add_node("provision", account_provisioning_agent)

workflow.set_entry_point("extract")
workflow.add_edge("extract", "compliance")
workflow.add_edge("compliance", "provision")
workflow.add_edge("provision", END)

app = workflow.compile()

Step 3: The Idempotent Consumer Wrapper

This is the most critical part. We write a decorator that intercepts the SQS message, checks the idempotency store, and either processes the LangGraph or safely ignores the duplicate.

from functools import wraps
from botocore.exceptions import ClientError

def idempotent_consumer(func):
    """
    Decorator to ensure the wrapped function is executed exactly once 
    per Idempotency Key, even if the SQS message is redelivered.
    """
    @wraps(func)
    def wrapper(event, context):
        for record in event.get('Records', []):
            body = json.loads(record['body'])
            idempotency_key = body.get('idempotency_key')
            
            if not idempotency_key:
                raise ValueError("Missing idempotency_key in payload!")

            # 1. Check Idempotency Store
            try:
                # Attempt to create the record with status IN_PROGRESS.
                # The ConditionExpression ensures this only succeeds if the key DOES NOT exist.
                table.put_item(
                    Item={
                        'IdempotencyKey': idempotency_key,
                        'Status': 'IN_PROGRESS',
                        'Timestamp': int(time.time())
                    },
                    ConditionExpression='attribute_not_exists(IdempotencyKey)'
                )
                print(f"[Idempotency] Claimed key: {idempotency_key}")
                
            except ClientError as e:
                if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
                    # The key already exists. Let's check its status.
                    existing_item = table.get_item(Key={'IdempotencyKey': idempotency_key}).get('Item')
                    status = existing_item.get('Status')
                    
                    if status == 'COMPLETED':
                        print(f"[Idempotency] Key {idempotency_key} already completed. Skipping.")
                        # We must delete the message from SQS so it doesn't loop forever
                        sqs.delete_message(
                            QueueUrl=queue_url,
                            ReceiptHandle=record['receiptHandle']
                        )
                        continue # Move to next record in batch
                        
                    elif status == 'IN_PROGRESS':
                        print(f"[Idempotency] Key {idempotency_key} is currently being processed by another worker. Failing back to queue.")
                        # Don't delete from SQS. Let the visibility timeout expire and it will retry.
                        continue 
                else:
                    raise e

            # 2. Execute the actual business logic (LangGraph)
            try:
                print(f"[Worker] Executing LangGraph for {idempotency_key}...")
                func(body)
                
                # 3. Mark as COMPLETED
                table.update_item(
                    Key={'IdempotencyKey': idempotency_key},
                    UpdateExpression='SET #s = :s',
                    ExpressionAttributeNames={'#s': 'Status'},
                    ExpressionAttributeValues={':s': 'COMPLETED'}
                )
                print(f"[Idempotency] Marked key {idempotency_key} as COMPLETED.")
                
                # 4. Delete message from SQS
                sqs.delete_message(
                    QueueUrl=queue_url,
                    ReceiptHandle=record['receiptHandle']
                )
                
            except Exception as e:
                # Mark as FAILED so it can be retried or sent to DLQ
                table.update_item(
                    Key={'IdempotencyKey': idempotency_key},
                    UpdateExpression='SET #s = :s, #err = :err',
                    ExpressionAttributeNames={'#s': 'Status', '#err': 'Error'},
                    ExpressionAttributeValues={':s': 'FAILED', ':err': str(e)}
                )
                print(f"[Idempotency] Marked key {idempotency_key} as FAILED: {e}")
                raise e

    return wrapper

# The actual business logic function
@idempotent_consumer
def process_onboarding_event(payload):
    initial_state = {
        "messages": [HumanMessage(content=f"Onboard client {payload['client_id']}")],
        "client_id": payload['client_id'],
        "idempotency_key": payload['idempotency_key'],
        "extracted_data": {},
        "compliance_status": "",
        "account_created": False
    }
    
    # Invoke the Multi-Agent LangGraph
    final_state = app.invoke(initial_state)
    print(f"Final State: Account Created = {final_state['account_created']}")

Step 4: Simulating the Event-Driven Flow (and the Crash)

Let's simulate the exact failure scenario to prove the idempotency works.

def simulate_sqs_event(payload, receipt_handle):
    """Simulates the JSON payload AWS Lambda receives from SQS"""
    return {
        "Records": [
            {
                "body": json.dumps(payload),
                "receiptHandle": receipt_handle
            }
        ]
    }

if __name__ == "__main__":
    client_id = "CLIENT_999"
    idempotency_key = f"onboarding_{client_id}_v1"
    
    payload = {
        "client_id": client_id,
        "idempotency_key": idempotency_key,
        "document_s3_uri": "s3://kyc-docs/client_999.pdf"
    }

    # --- SCENARIO 1: First attempt (Success) ---
    print("\n--- ATTEMPT 1: Normal Execution ---")
    event1 = simulate_sqs_event(payload, "handle_1")
    process_onboarding_event(event1, None)

    # --- SCENARIO 2: SQS Redelivery (The Duplicate) ---
    print("\n--- ATTEMPT 2: Simulating SQS Redelivery (Duplicate Message) ---")
    # In reality, SQS redelivers the exact same message with the same body.
    # We use a new receipt handle, but the idempotency_key in the body is identical.
    event2 = simulate_sqs_event(payload, "handle_2") 
    process_onboarding_event(event2, None)

Expected Output:

--- ATTEMPT 1: Normal Execution ---
[Idempotency] Claimed key: onboarding_CLIENT_999_v1
[Worker] Executing LangGraph for onboarding_CLIENT_999_v1...
[Doc Agent] Retrieving data for client CLIENT_999 from Vector DB...
[KYC Agent] Checking compliance policies in Vector DB...
[Provisioning Agent] Creating account for CLIENT_999 in Core Banking...
Final State: Account Created = True
[Idempotency] Marked key onboarding_CLIENT_999_v1 as COMPLETED.

--- ATTEMPT 2: Simulating SQS Redelivery (Duplicate Message) ---
[Idempotency] Key onboarding_CLIENT_999_v1 already completed. Skipping.

Notice how in Attempt 2, the LangGraph never executes. The Account Provisioning Agent is not called a second time. The system safely ignores the duplicate, deletes the SQS message, and saves the cost of the LLM inference.

Part 4: Advanced Enterprise Considerations

When moving this to a production enterprise environment, keep these advanced patterns in mind:

Time-To-Live (TTL) on Idempotency Records:
Don't let your DynamoDB table grow infinitely. Set a TTL on the IdempotencyStore table (e.g., 24 hours). If a message is redelivered after 24 hours, it's treated as a new event.
Handling "IN_PROGRESS" Race Conditions:
If Worker A crashes while IN_PROGRESS, Worker B will see the IN_PROGRESS status and fail the message back to the queue. To prevent infinite loops if Worker A is truly dead, include a Timestamp in the idempotency record. If Worker B sees an IN_PROGRESS record older than the SQS visibility timeout, it can forcefully overwrite it and take over.
Idempotency at the API Gateway:
Generate the idempotency_key at the API Gateway (using a combination of User-ID and a client-generated Request-ID) and pass it down to EventBridge/SQS. This ensures the key is consistent across the entire distributed trace.
Dead Letter Queues (DLQ):
If an agent fails 3 times, the message should go to a DLQ. The idempotency record should be marked as FAILED so that when an engineer manually replays the message from the DLQ, the system knows it's a retry, not a brand new request.

Conclusion

In event-driven AI systems, you cannot trust the network. Messages will be duplicated.

By wrapping your LangGraph multi-agent pipelines in an Idempotent Consumer pattern, you transform your system from a fragile, state-corrupting prototype into a resilient, enterprise-grade platform. You ensure that your RAG retrievals are efficient, your LLM costs are strictly controlled, and your downstream side-effects happen exactly once.