lambda_handler.py

import importlib.resources
import json
import os
import sys

import boto3
from mcp import StdioServerParameters, stdio_client
from strands import Agent
from strands.models.openai import OpenAIModel
from strands.tools.mcp import MCPClient

ssm = boto3.client("ssm")
GROQ_PARAM_NAME = os.environ.get("GROQ_PARAM_ARN")
SYSTEM_PROMPT = "Você é um assistente da AWS, especializado em responder perguntas sobre os serviços da AWS. Use os recursos disponíveis para fornecer respostas precisas e úteis."


def get_groq_api_key():
    """Busca a chave de API do AWS Parameter Store. Levanta exceção em caso de erro."""
    if not GROQ_PARAM_NAME:
        raise ValueError("Variável de ambiente GROQ_PARAM_ARN não definida.")

    try:
        param = ssm.get_parameter(Name=GROQ_PARAM_NAME, WithDecryption=True)

        return param["Parameter"]["Value"]
    except Exception as e:
        raise RuntimeError(
            "Não foi possível obter a chave de API do Parameter Store."
        ) from e


def get_groq_model():
    """Inicializa e retorna o modelo Groq. Levanta exceção em caso de erro."""
    api_key = get_groq_api_key()

    return OpenAIModel(
        client_args={
            "api_key": api_key,
            "base_url": "https://api.groq.com/openai/v1",
        },
        model_id="openai/gpt-oss-20b",
        params={"max_tokens": 2048, "temperature": 0.5},
    )


def get_mcp_client():
    """Configura e retorna o cliente MCP."""
    pkg_path = importlib.resources.files("awslabs.aws_documentation_mcp_server")

    server_script_path = str(pkg_path / "server.py")

    process_env = os.environ.copy()

    if os.environ.get("AWS_LAMBDA_FUNCTION_NAME"):
        layer_path = "/opt/python"
        current_pythonpath = os.environ.get("PYTHONPATH", "")
        new_pythonpath = f"{layer_path}:{current_pythonpath}"
        process_env["PYTHONPATH"] = new_pythonpath

    return MCPClient(
        lambda: stdio_client(
            StdioServerParameters(
                command=sys.executable, args=[server_script_path], env=process_env
            )
        )
    )


try:
    groq_model = get_groq_model()
    mcp_client = get_mcp_client()
except Exception as e:
    groq_model = None
    mcp_client = None
    initialization_error = e


def lambda_handler(event, context):
    if groq_model is None or mcp_client is None:
        return {
            "statusCode": 500,
            "body": json.dumps(
                {"error": "Falha na inicialização da configuração do agente."}
            ),
        }

    try:
        prompt = event.get("prompt")

        if not prompt:
            return {"statusCode": 400, "body": json.dumps({"error": "Missing prompt"})}

        with mcp_client as mpc:
            tools = mpc.list_tools_sync()
            agent = Agent(model=groq_model, tools=tools, system_prompt=SYSTEM_PROMPT)
            result = agent(prompt)

            return {"statusCode": 200, "body": json.dumps({"result": str(result)})}

    except Exception:
        return {
            "statusCode": 500,
            "body": json.dumps(
                {"error": "Ocorreu um erro interno no processamento do agente."}
            ),
        }