Estratégias serverless para streaming de respostas LLM

Por KyungYong Shim, Delivery Consultant-AppMod.

Aplicações modernas de IA generativa frequentemente precisam transmitir saídas de modelos de linguagem grande (LLM) para usuários em tempo real. Em vez de esperar por uma resposta completa, o streaming entrega resultados parciais à medida que ficam disponíveis, o que melhora significativamente a experiência do usuário para interfaces de chat e tarefas de IA de longa duração. Esta publicação compara três abordagens Serverless para lidar com streaming de LLM do Amazon Bedrock na Amazon Web Services (AWS), o que ajuda você a escolher a melhor opção para sua aplicação.

URLs de função do AWS Lambda com streaming de resposta
APIs WebSocket do Amazon API Gateway
Assinaturas GraphQL do AWS AppSync

Cobrimos como cada opção funciona, os detalhes de implementação, autenticação com o Amazon Cognito e quando escolher uma em vez das outras.

URLs de função do Lambda com streaming de resposta

URLs de função do AWS Lambda fornecem um endpoint HTTP(S) direto para invocar sua função Lambda. O streaming de resposta permite que sua função envie fragmentos incrementais de dados de volta ao chamador sem fazer buffer da resposta inteira. Esta abordagem é ideal para encaminhar a saída transmitida do Amazon Bedrock, proporcionando uma experiência de usuário mais rápida. O streaming é suportado no Node.js 18+. No Node.js, você envolve seu handler com awslambda.streamifyResponse(), que fornece um stream para gravar dados e que os envia imediatamente para a resposta HTTP.

Arquitetura

A figura a seguir mostra a arquitetura.

O cliente faz uma chamada fetch() para a URL da função Lambda.
O Lambda invoca InvokeModelWithResponseStream usando o AWS SDK for JavaScript.
À medida que os tokens chegam do Amazon Bedrock, eles são gravados no stream de resposta.

Etapas de implementação

Criar uma função Lambda com streaming: Use um runtime Node.js 18+ ou posterior (necessário para streaming nativo). Instale o AWS SDK para chamar o Amazon Bedrock. No código do handler, envolva a função com awslambda.streamifyResponse e faça streaming da saída do modelo. Por exemplo, no Node.js você pode fazer o seguinte:

const bedrock = new BedrockRuntimeClient({region: "us-east-1"});

// Please consider adding more details when you use it for your application.
exports.handler = awslambda.streamifyResponse(async (event, responseStream) => 
{
    // 1. Parse input (e.g., prompt from event)
    const prompt = event.body?.prompt;
    // 2. Call Amazon Bedrock with streaming (using AWS SDK for Amazon Bedrock)
    const command = new InvokeModelWithResponseStreamCommand({ modelId: "YOUR_MODEL_ID", body: { prompt }});
    const response = await bedrock.send(command);
    // 3. Stream Bedrock tokens to client
    for await (const event of response.body) {
        if (event.content) {
            responseStream.write(event.content); // write partial output
        }
    }
    // 4. End stream when done
    responseStream.end();
});

Este trecho de código usa o async iterable do SDK do Amazon Bedrock para ler o stream de eventos de tokens e grava cada um no responseStream.
Configurar a função do Lambda: a função de execução deve permitir a invocação do Amazon Bedrock (como bedrock:InvokeModelWithResponseStream no Amazon Resource Name (ARN) do modelo LLM).

Autenticação com o Amazon Cognito

URLs de função do Lambda podem ser definidas como “None” (público) ou “AWS_IAM”. A autenticação nativa de token do Cognito User Pool não é suportada, portanto você precisa implementar uma solução.

Verificação JWT no Lambda: Permita acesso público e verifique um JWT válido do Amazon Cognito no cabeçalho da solicitação dentro do seu código Lambda. Isso requer esforço de desenvolvimento.

// Initialize Cognito JWT Verifier
const { CognitoJwtVerifier } = require('aws-jwt-verify');

const jwtVerifier = CognitoJwtVerifier.create({
  userPoolId: USER_POOL_ID,
  tokenUse: 'id',
  clientId: USER_POOL_CLIENT_ID
});

// Verify JWT token from Cognito
async function verifyToken(token) {
  try {
    if (!token) throw new Error('No authorization token provided');
    
    // Remove 'Bearer ' prefix if present
    if (token.startsWith('Bearer ')) {
      token = token.slice(7);
    }

    // Verify the token using Cognito JWT Verifier
    const payload = await jwtVerifier.verify(token);
    logger.info(`Verified token for user: ${payload.sub}`);
    
    return payload;
  } catch (error) {
    logger.error(`Token verification failed: ${error.message}`);
    throw new Error(`Invalid token: ${error.message}`);
  }
}

//...

    // Verify authentication
    let userId;
    try {
      const authHeader = event.headers?.Authorization;
      const payload = await verifyToken(authHeader);
      userId = payload.sub;
      logger.info(`Authenticated user: ${userId}`);
    } catch (error) {
      responseStream.write(`data: ${JSON.stringify({ type: 'error', error: 'Unauthorized', message: error.message })}\n\n`);
      return;
    }

Autorização IAM com identidade do Amazon Cognito: Use credenciais AWS obtidas do Amazon Cognito. Uma configuração mais complexa, especialmente para aplicações web, é potencialmente excessiva para uma única função.

Prós e contras das URLs de função do Lambda

Prós:

Simplicidade: Não são necessários API Gateway ou outros serviços, o que minimiza a sobrecarga operacional.
Baixa latência, alta taxa de transferência: A resposta é entregue diretamente do Lambda para o cliente. Isso produz excelente desempenho de Time To First Byte (TTFB), sem buffer intermediário.
Implementação direta: Para desenvolvedores Node.js, ativar streaming é tão direto quanto um wrapper e gravar em um stream. Isso é ideal para protótipos rápidos ou microsserviços de função única.
Menor custo para baixa concorrência: Você paga apenas pelo tempo de execução do Lambda. Não há custo de conexão persistente, o que é o mesmo que com WebSocket ou AWS AppSync. Se as invocações forem pouco frequentes ou curtas, então isso pode ser econômico.

Contras:

Suporte limitado de runtime: O streaming nativo é suportado apenas no Node.js.
Sem autenticação de user pool integrada: Ao contrário do API Gateway ou AWS AppSync, URLs do Lambda não suportam diretamente autorizadores de user pool do Amazon Cognito. Você deve lidar com autenticação através do AWS Identity and Access Management (IAM) ou validação manual de token, adicionando algum esforço de desenvolvimento e possíveis armadilhas de segurança se feito incorretamente.
Complexidade no tratamento de erros: O streaming torna a propagação de erros mais complicada. Se ocorrer um erro no meio do stream, então você precisa decidir como informar o cliente.

WebSocket do API Gateway para streaming

APIs WebSocket do API Gateway estabelecem conexões persistentes e com estado entre clientes e seu backend. Isso é ideal para aplicações em tempo real que precisam de mensagens iniciadas pelo servidor. O cliente se conecta uma vez, envia um prompt para o Amazon Bedrock através do WebSocket, e o servidor envia a resposta transmitida de volta pela mesma conexão.

Arquitetura

A figura a seguir mostra a arquitetura.

O cliente se conecta através da URL WebSocket e armazena o connectionId.
O cliente envia um prompt através de uma rota personalizada para o LLMHandler.
O Lambda como LLMHandler invoca o Amazon Bedrock e transmite de volta através do WebSocket.
O cliente se desconecta através do DisconnectHandler e remove o connectionId.

Etapas de implementação

Criar uma API WebSocket no API Gateway com rotas
1. $connect: Conectado ao Lambda ConnectHandler.
2. $disconnect: Conectado ao Lambda DisconnectHandler.
3. $stream: Todas as mensagens vão para o Lambda StreamHandler.

Criar Lambda Authorizer

Recebe a solicitação de conexão com token na query string.
Valida o token JWT contra o Amazon Cognito.

Retorna política Allow/Deny para a conexão.

def lambda_handler(event, context):
    # Extract token from querystring
    token = event.get("queryStringParameters", {}).get("token", "")
    
    # Validate JWT token against Cognito
    if validate_token(token):
        return {
            "isAuthorized": True,
            # Optionally include context that other handlers can access
            "context": {
                "userId": extracted_user_id
            }
        }
    else:
        return {"isAuthorized": False}

Criar Connection Handler

O Lambda de conexão é executado após autorização bem-sucedida.
Recebe o connectionId único da nova conexão.
Armazena informações de conexão no Amazon DynamoDB (opcional).

Retorna status 200 para completar a conexão.

def lambda_handler(event, context):
    # Extract connectionId
    connection_id = event.get("requestContext", {}).get("connectionId")
    
    # Optionally store in DynamoDB
    # dynamodb.put_item(...)
    
    # Connection established successfully
    return {"statusCode": 200}

Criar Disconnect Handler

O Lambda de desconexão é acionado automaticamente quando os clientes se desconectam.
Recebe o connectionId da conexão terminada.
Limpa quaisquer dados de conexão armazenados.

Retorna status 200

def lambda_handler(event, context):
    # Extract connectionId
    connection_id = event.get("requestContext", {}).get("connectionId")
    
    # Optionally remove from DynamoDB
    # dynamodb.delete_item(...)
    
    # Disconnection handled successfully
    return {"statusCode": 200}

Criar LLM Handler

Recebe mensagens enviadas para a rota stream.
Extrai o prompt do corpo da mensagem.
Chama o modelo do Amazon Bedrock com resposta de streaming.

Transmite tokens de volta para o cliente usando o ID de conexão.

def lambda_handler(event, context):
    # Extract connectionId and domain details for sending responses
    connection_id = event["requestContext"]["connectionId"]
    domain = event["requestContext"]["domainName"]
    stage = event["requestContext"]["stage"]
    
    # Parse message body to get the prompt
    body = json.loads(event.get("body", "{}"))
    prompt = body.get("prompt", "")
    
    # Create API Gateway management client for sending responses
    api_client = boto3.client(
        'apigatewaymanagementapi',
        endpoint_url=f'https://{domain}/{stage}'
    )
    
    # Call Amazon Bedrock with streaming response
    response = bedrock_client.invoke_model_with_response_stream(...)
    
    # Stream tokens back to client
    for chunk in response["body"]:
        # Extract token from chunk
        token = process_chunk(chunk)
        
        # Send token directly back through the WebSocket
        api_client.post_to_connection(
            ConnectionId=connection_id,
            Data=json.dumps({"token": token, "isComplete": False})
        )
    
    # Send completion message
    api_client.post_to_connection(
        ConnectionId=connection_id,
        Data=json.dumps({"token": "", "isComplete": True})
    )
    
    return {"statusCode": 200}

Autenticação com o Amazon Cognito

Proteger uma API WebSocket com o Amazon Cognito precisa de um pouco mais de trabalho. O WebSocket do API Gateway não tem um autorizador de User Pool do Amazon Cognito integrado:

Lambda authorizer com autenticação JWT: O API Gateway invoca seu Lambda authorizer na conexão, validando o JWT do Amazon Cognito (passado como parâmetro de consulta). O Lambda gera uma política IAM concedendo acesso e a retorna.
Autenticação IAM para WebSockets: Os clientes assinam solicitações com SigV4 usando credenciais AWS de um Identity Pool do Amazon Cognito. O API Gateway avalia a solicitação contra políticas IAM.

Prós e contras das APIs WebSocket do API Gateway

Prós:

Comunicação bidirecional em tempo real: WebSockets são ideais para aplicações onde o servidor precisa enviar dados como a resposta do LLM sem solicitações explícitas.
Conexão persistente para conversas de múltiplas rodadas: Após o handshake inicial, a mesma conexão pode ser reutilizada para prompts e respostas subsequentes, evitando latência de configuração repetida. Isso é ótimo para uma UI de chat onde o usuário faz várias perguntas em uma sessão.
Escalabilidade: O API Gateway é um serviço gerenciado que pode lidar com 500 conexões/segundo e 10.000 solicitações/segundo entre APIs, o que pode ser aumentado mediante solicitação.

Contras:

Maior complexidade de desenvolvimento: Quando comparado à simplicidade de uma URL Lambda direta, uma API WebSocket envolve múltiplos Lambdas e coordenação para gerenciar o estado da conexão.
Implementação de autenticação personalizada: Não há integração integrada de user pool do Amazon Cognito, portanto você deve implementar um Lambda authorizer.
Gerenciamento de timeout: O timeout de integração do API Gateway é de 29 s, portanto sua função Lambda deve retornar a resposta prontamente.

Assinatura GraphQL do AWS AppSync

O AWS AppSync é um serviço GraphQL totalmente gerenciado que simplifica a construção de APIs em tempo real. Ele lida com conexões WebSocket e fan-out de cliente automaticamente. Os clientes assinam uma assinatura GraphQL, e um resolvedor Lambda envia os tokens transmitidos do Amazon Bedrock de volta.

Arquitetura

A figura a seguir mostra a arquitetura.

O cliente chama uma mutation startStream. O AppSync invoca o Lambda de Solicitação.
O Lambda de Solicitação retorna imediatamente um sessionId único e envia a tarefa de processamento para uma fila do Amazon Simple Queue Service (Amazon SQS).
O cliente usa o sessionId para assinar uma subscription GraphQL onTokenReceived.
O Lambda de Processamento (acionado pelo Amazon SQS) invoca o Amazon Bedrock e, para cada token, chama uma mutation publishToken no AWS AppSync.
O AWS AppSync envia automaticamente o token para todos os clientes inscritos com o sessionId correspondente.

Etapas de implementação

Projetar o Schema GraphQL: definir tipos e operações.

type StreamResponse {
  sessionId: String!
  status: String!
  message: String
  timestamp: String!
  error: String
}

type TokenEvent {
  sessionId: String!
  token: String!
  isComplete: Boolean!
  timestamp: String!
}

type Mutation {
  startStream(prompt: String!): StreamResponse!
  publishToken(sessionId: String!, token: String!, isComplete: Boolean!): TokenEvent!
}

type Subscription {
  onTokenReceived(sessionId: String!): TokenEvent

Criar o Request Handler (Lambda de Solicitação)

Recebe a mutation GraphQL com o prompt.
Gera um ID de sessão único.
Envia o prompt e o ID de sessão para a fila SQS.

Retorna o ID de sessão para o cliente imediatamente.

def lambda_handler(event, context):
    # Extract prompt from GraphQL event
    prompt = event["arguments"]["prompt"]
    
    # Generate unique session ID
    session_id = str(uuid.uuid4())
    
    # Send message to SQS queue
    sqs_client.send_message(
        QueueUrl="your-sqs-queue-url",
        MessageBody=json.dumps({
            "prompt": prompt,
            "sessionId": session_id
        })
    )
    
    # Return session ID to client
    return {
        "sessionId": session_id,
        "status": "streaming_started",
        "timestamp": datetime.datetime.utcnow().isoformat()
    }

Criar o Processing Handler (Lambda de Processamento)

É acionado por mensagens do Amazon SQS.
Chama o Amazon Bedrock com streaming ativado.

Para cada token gerado, chama a mutation publishToken do AppSync.

def lambda_handler(event, context):
    # Process SQS event records
    for record in event["Records"]:
        body = json.loads(record["body"])
        prompt = body["prompt"]
        session_id = body["sessionId"]
        
        # Call Amazon Bedrock with streaming
        response = bedrock_client.invoke_model_with_response_stream(...)
        
        # Process streaming response
        for chunk in response["body"]:
            # Extract token from chunk
            token = process_chunk(chunk)
            
            # Publish token to AppSync
            publish_token_to_appsync(
                session_id=session_id,
                token=token,
                is_complete=False
            )
        
        # Send completion token
        publish_token_to_appsync(
            session_id=session_id,
            token="",
            is_complete=True
        )

Configurar Resolvedores GraphQL
1. Resolvedor StartStream: Conectar ao Lambda de Solicitação.
2. Resolvedor PublishToken: Acionar assinatura com uma fonte de dados NONE.

Configuração de assinatura do cliente

Fazer uma mutation startStream.

const { sessionId } = await client.mutate({
  mutation: START_STREAM,
  variables: { prompt }
});

Assinar para receber tokens.

client.subscribe({
  query: ON_TOKEN_RECEIVED,
  variables: { sessionId }
}).subscribe({
  next: ({ data }) => {
    if (data.onTokenReceived.isComplete) {
      // Handle completion
    } else {
      // Append token to UI
      appendToken(data.onTokenReceived.token);
    }
  }
});

Autenticação com o Amazon Cognito

O AWS AppSync se integra perfeitamente com User Pools do Amazon Cognito. Definir o modo de autenticação da API para User Pool do Amazon Cognito requer um JWT válido para cada operação GraphQL. Esta é a opção mais amigável ao desenvolvedor para autenticação. O AWS AppSync lida com o handshake e a atualização de token.

Prós e contras das assinaturas do AWS AppSync

Prós:

Protocolo em tempo real totalmente gerenciado: Você não lida com WebSockets brutos ou IDs de conexão. O AWS AppSync estabelece e mantém automaticamente um WebSocket seguro para assinaturas (não é necessário um Lambda de conexão ou desconexão).
Autenticação simplificada: Suporte integrado para tokens de User Pool do Amazon Cognito significa que você pode proteger a API sem escrever autorizadores personalizados.

Contras:

Potencial sobrecarga e complexidade: Para um caso direto (um prompt—um stream), introduzir GraphQL e AWS AppSync pode ser visto como excesso de engenharia se sua aplicação não usa GraphQL para outros casos de uso.
Limite de 30 segundos do resolvedor: O AWS AppSync tem um limite de 30 segundos para resolvedores de mutation, portanto você precisa projetar a solicitação inicial para iniciar o processo e retornar imediatamente, contando com uma assinatura para transmitir os resultados progressivamente para evitar bloquear o usuário.

Conclusão

A interface de streaming do Amazon Bedrock desbloqueia experiências LLM fluidas e de baixa latência. Você pode usar a arquitetura Serverless AWS correta para entregar respostas transmitidas de forma segura, escalável e econômica.

URLs de função do Lambda com streaming: Aplicações diretas de usuário único e protótipos.
WebSocket do API Gateway: Conversas de múltiplas rodadas, aplicações colaborativas.
AppSync: Aplicações complexas já usando GraphQL.

Cada método é Serverless, pronto para produção e totalmente integrado com o Amazon Cognito para controle de acesso seguro. A AWS fornece a flexibilidade para projetar experiências de usuário de IA de alta qualidade em escala.

Consulte o código-fonte de exemplo no GitHub para mais detalhes.

Tabela comparativa

Recurso	URLS DE FUNÇÃO DO LAMBDA	APIS WEBSOCKET DO API GATEWAY	ASSINATURAS GRAPHQL DO APPSYNC
Complexidade	Mais baixa	Média	Alta
Foco em tempo real	Limitado	Forte	Forte
Autenticação	Requer lógica personalizada	Requer lógica personalizada	Suporte integrado ao Amazon Cognito
Escalabilidade	Boa	Boa	Excelente
Suporte GraphQL	Nenhum	Nenhum	Nativo
Casos de uso	Perguntas e respostas	Chatbots, aplicações em tempo real	Aplicações complexas, cenários multiusuário
Custo	Pagamento por invocação	Tempo de conexão e execução do Lambda	Preços baseados em solicitação/conexão

Este conteúdo foi traduzido do post original do blog, que pode ser encontrado aqui.

Autores

KyungYong Shim, Delivery Consultant-AppMod

Tradutores

Rodrigo Peres é Arquiteto de Soluções na AWS, com mais de 20 anos de experiência trabalhando com arquitetura de soluções, desenvolvimento de sistemas e modernização de sistemas legados.

Daniel Abib é arquiteto de soluções sênior na AWS, com mais de 25 anos trabalhando com gerenciamento de projetos, arquiteturas de soluções escaláveis, desenvolvimento de sistemas e CI/CD, microsserviços, arquitetura Serverless & Containers e segurança. Ele trabalha apoiando clientes corporativos, ajudando-os em sua jornada para a nuvem.

https://www.linkedin.com/in/danielabib/

O blog da AWS

Estratégias serverless para streaming de respostas LLM

URLs de função do Lambda com streaming de resposta

Arquitetura

Etapas de implementação

Autenticação com o Amazon Cognito

Prós e contras das URLs de função do Lambda

WebSocket do API Gateway para streaming

Arquitetura

Etapas de implementação

Autenticação com o Amazon Cognito

Prós e contras das APIs WebSocket do API Gateway

Assinatura GraphQL do AWS AppSync

Arquitetura

Etapas de implementação

Autenticação com o Amazon Cognito

Prós e contras das assinaturas do AWS AppSync

Conclusão

Tabela comparativa

Autores

Tradutores

Aprenda

Recursos

Desenvolvedores

Ajuda