Cryptocurrency  

AWS Cognito User Pools vs. Identity Pools in LangGraph

Building an enterprise-grade Retrieval-Augmented Generation (RAG) application requires more than just connecting a Large Language Model (LLM) to a vector database. In a multi-tenant or multi-department enterprise environment, security, identity management, and fine-grained access control are paramount. This is where Amazon Cognito becomes the backbone of your AI architecture. However, developers often confuse its two main components: User Pools and Identity Pools. In this end-to-end guide, we will demystify the differences between the two, explore a real-world enterprise use case, and implement a secure, multi-agent LangGraph RAG system that leverages both.

Part 1: The Core Differences (The "Who" vs. The "What")

To secure an application, you need to answer two questions: Who is this user? and What are they allowed to do? Cognito splits these responsibilities.

1. Cognito User Pools (Authentication / "Who are you?")

  • Purpose: Acts as a managed user directory. It handles sign-up, sign-in, Multi-Factor Authentication (MFA), and account recovery.

  • Output: Issues JSON Web Tokens (JWTs) — specifically ID Tokens, Access Tokens, and Refresh Tokens.

  • Enterprise Use: Integrates with corporate Identity Providers (IdP) like Okta, Microsoft Entra ID (Azure AD), or PingIdentity via SAML/OIDC for Single Sign-On (SSO).

  • Key Feature: You can add Custom Attributes (e.g., custom:department, custom:clearance_level) to the ID Token, which is crucial for application-level logic.

2. Cognito Identity Pools (Authorization / "What can you access in AWS?")

  • Purpose: Exchanges your User Pool tokens (or external IdP tokens) for temporary, scoped AWS credentials (Access Key, Secret Key, Session Token).

  • Output: IAM Roles mapped to authenticated (or unauthenticated) users.

  • Enterprise Use: Allows the frontend client or a specific backend microservice to directly access AWS resources (like S3, OpenSearch, or Bedrock) without hardcoding long-term IAM credentials.

  • Key Feature: Enables fine-grained AWS resource access. For example, an authenticated user can be given an IAM role that only allows reading from s3://company-data/finance/ and not s3://company-data/hr/.

The Golden Rule: User Pools authenticate the user to your application. Identity Pools authorize the user to access AWS resources.

10

Part 2: Real-Time Enterprise Use Case

The Scenario: GlobalCorp is deploying an internal "Knowledge Assistant" built with a Multi-Agent LangGraph RAG pipeline.

  • The Problem: The company has highly sensitive documents. HR documents must only be accessible to the HR department; Finance documents only to Finance. Furthermore, the LLM inference (Amazon Bedrock) and Vector Store (Amazon OpenSearch) are hosted on AWS.

  • The Solution:

    1. User Pool: Employees log in via corporate SSO. The ID Token includes a custom claim: custom:department.

    2. Identity Pool: The frontend exchanges the User Pool token for temporary AWS credentials, allowing it to securely invoke the LangGraph API Gateway endpoint.

    3. LangGraph RAG: The backend extracts the department from the User Pool ID token and passes it to the Retriever Agent, which applies a metadata filter to the Vector DB, ensuring the LLM only generates answers based on authorized documents.

Part 3: Architecture Overview

  1. Frontend (React/Next.js): User logs in via Cognito User Pool. Frontend exchanges the token via Cognito Identity Pool to get temporary AWS credentials.

  2. API Gateway: Receives the request. The temporary credentials validate the request.

  3. Backend (FastAPI/Lambda): Validates the User Pool ID Token, extracts the custom:department claim.

  4. LangGraph Orchestrator:

    • Supervisor Agent: Routes the query.

    • Retriever Agent: Queries OpenSearch using the user's department as a metadata filter.

    • Generator Agent: Synthesizes the final answer using Amazon Bedrock.

Part 4: Code Implementation

Step 1: Frontend Authentication (React + AWS Amplify)

Here, we configure both the User Pool (for login) and the Identity Pool (for AWS resource access).

// src/auth/config.ts
import { Amplify } from 'aws-amplify';
import { signIn, fetchAuthSession } from 'aws-amplify/auth';

Amplify.configure({
  Auth: {
    Cognito: {
      // USER POOL: Handles App Authentication
      userPoolId: 'us-east-1_XXXXXXXXX',
      userPoolClientId: 'XXXXXXXXXXXXXXXXXXXXXXXXXX',
      
      // IDENTITY POOL: Handles AWS Resource Authorization
      identityPoolId: 'us-east-1:xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx',
    }
  }
});

export async function authenticateUser(username: string, password: string) {
  // 1. Authenticate via User Pool
  const user = await signIn({ username, password });
  
  // 2. Get Identity Pool credentials (Temporary AWS Creds)
  const session = await fetchAuthSession();
  const awsCredentials = session.credentials;
  
  // 3. Extract User Pool ID Token (Contains custom claims like department)
  const idToken = session.tokens?.idToken?.toString();
  
  return { awsCredentials, idToken, userDepartment: session.tokens?.idToken?.payload['custom:department'] };
}

Step 2: Backend LangGraph Multi-Agent RAG (Python)

Now, let's build the LangGraph pipeline. We will define a state that carries the user's context, and a Retriever Agent that enforces document-level security based on the Cognito User Pool claims.

# rag_graph.py
import os
import jwt
from typing import TypedDict, Annotated, Literal
from langgraph.graph import StateGraph, END
from langchain_aws import ChatBedrock
from langchain_aws import AmazonOpenSearchVectorStore
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.documents import Document

# --- 1. Define the LangGraph State ---
class AgentState(TypedDict):
    query: str
    user_department: str  # Injected from Cognito User Pool ID Token
    context: list[Document]
    answer: str
    error: str

# --- 2. Initialize AWS Resources (Using Identity Pool Creds in Prod) ---
# In production, the backend assumes an IAM role or uses the temp creds passed from frontend
llm = ChatBedrock(model_id="anthropic.claude-3-sonnet-20240229-v1:0", region_name="us-east-1")
vector_store = AmazonOpenSearchVectorStore(
    opensearch_url=os.getenv("OPENSEARCH_URL"),
    region_name="us-east-1"
)

# --- 3. Define the Agents (Nodes) ---

def retriever_agent(state: AgentState):
    """Retrieves documents, strictly filtered by the user's department."""
    query = state["query"]
    department = state["user_department"]
    
    print(f"[Retriever Agent] Searching for '{query}' restricted to department: {department}")
    
    # Apply metadata filter based on Cognito User Pool custom claim
    search_kwargs = {
        "k": 5,
        "filter": {"department": department} # Document-level security!
    }
    
    docs = vector_store.similarity_search(query, **search_kwargs)
    return {"context": docs}

def generator_agent(state: AgentState):
    """Generates the final answer using the retrieved context."""
    query = state["query"]
    context = state["context"]
    
    context_text = "\n\n".join([doc.page_content for doc in context])
    
    prompt = ChatPromptTemplate.from_template(
        "You are an enterprise assistant. Answer the question based ONLY on the context.\n"
        "Context: {context}\n"
        "Question: {query}"
    )
    
    chain = prompt | llm
    response = chain.invoke({"context": context_text, "query": query})
    
    return {"answer": response.content}

def supervisor_agent(state: AgentState) -> Literal["retriever", "end"]:
    """Supervisor decides the next step."""
    if not state.get("context"):
        return "retriever"
    return "end"

# --- 4. Build the LangGraph ---
def build_rag_graph():
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("supervisor", supervisor_agent)
    workflow.add_node("retriever", retriever_agent)
    workflow.add_node("generator", generator_agent)
    
    # Set entry point and edges
    workflow.set_entry_point("supervisor")
    workflow.add_conditional_edges("supervisor", supervisor_agent, {
        "retriever": "retriever",
        "end": END
    })
    workflow.add_edge("retriever", "generator")
    workflow.add_edge("generator", END)
    
    return workflow.compile()

rag_app = build_rag_graph()

# --- 5. API Endpoint (FastAPI) ---
from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import requests

app = FastAPI()
security = HTTPBearer()

# JWKS client to verify Cognito User Pool tokens
jwks_url = f"https://cognito-idp.us-east-1.amazonaws.com/us-east-1_XXXXXXXXX/.well-known/jwks.json"

def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """Validates the Cognito User Pool ID Token and extracts claims."""
    token = credentials.credentials
    
    # In production, use a proper JWT library with JWKS verification (e.g., python-jose)
    # This is a simplified representation
    try:
        # Decode without verification for demonstration; ALWAYS verify in prod!
        payload = jwt.decode(token, options={"verify_signature": False})
        
        # Verify issuer and audience match your Cognito User Pool
        if payload.get("iss") != "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_XXXXXXXXX":
            raise HTTPException(status_code=401, detail="Invalid token issuer")
            
        return payload
    except Exception:
        raise HTTPException(status_code=401, detail="Invalid token")

@app.post("/ask")
def ask_question(query: str, user: dict = Depends(get_current_user)):
    # Extract the custom department claim from the User Pool ID Token
    user_department = user.get("custom:department", "public")
    
    # Invoke the LangGraph Multi-Agent RAG pipeline
    initial_state = {
        "query": query,
        "user_department": user_department,
        "context": [],
        "answer": "",
        "error": ""
    }
    
    final_state = rag_app.invoke(initial_state)
    
    return {
        "answer": final_state["answer"],
        "department_scoped": user_department
    }

Part 5: Why this Architecture is "Enterprise-Ready"

  1. Zero Trust Document Retrieval: By passing the custom:department claim from the User Pool directly into the LangGraph Retriever Agent's metadata filter, we ensure that the LLM never even "sees" unauthorized documents. This prevents prompt injection attacks that try to trick the LLM into revealing hidden context.

  2. No Hardcoded Secrets: The frontend uses the Identity Pool to get temporary AWS credentials. The backend uses IAM roles. There are no long-term access keys sitting in environment variables.

  3. Scalable Multi-Agent Orchestration: LangGraph allows us to easily add more agents later (e.g., a Guardrail Agent to check for PII, or an Action Agent to execute API calls) without rewriting the core authentication logic. The user context flows seamlessly through the graph state.

  4. Corporate SSO Integration: Because Cognito User Pools support SAML and OIDC, employees can use their existing corporate credentials (Okta/Entra ID) to access the AI tool, satisfying strict enterprise compliance requirements.

Conclusion

Understanding the boundary between Cognito User Pools (Application Identity) and Cognito Identity Pools (AWS Infrastructure Authorization) is critical for cloud-native AI applications. By combining Cognito's robust identity management with a Multi-Agent LangGraph RAG pipeline, you can build enterprise AI systems that are not only intelligent but also strictly secure, context-aware, and compliant with organizational data governance policies.

Membership not found