from typing import Any from collections import OrderedDict from strands import Agent, tool import asyncio from strands.agent.conversation_manager.null_conversation_manager import NullConversationManager from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model from mcp_client.client import get_awsinfrastructuretools_mcp_client from memory.session import get_memory_session_manager app = BedrockAgentCoreApp() log = app.logger # Define MCP clients for all configured MCP servers (gateways and/or remote MCP) mcp_clients = [] aws_infra_client = get_awsinfrastructuretools_mcp_client() if aws_infra_client: mcp_clients.append(aws_infra_client) DEFAULT_SYSTEM_PROMPT = """ Eres un asistente especializado en infraestructura AWS. Tu función es consultar y explicar el estado de la infraestructura AWS utilizando exclusivamente las herramientas disponibles. Antes de responder sobre EC2 o Lambda, consulta las herramientas correspondientes: - Para EC2, llama a listar_instancias_ec2 sin parámetros ({}). - Para Lambda, llama a listar_lambdas con {"operation": "listar_lambdas"}. Usa los nombres completos publicados por el gateway, incluidos sus prefijos. Respeta el esquema de cada herramienta y no inventes valores de operation. Si el usuario solicita EC2 y Lambda juntos, realiza dos llamadas independientes: primero consulta EC2 y después Lambda. No combines ambas consultas en una sola operación. Intenta ambas consultas aunque una falle y resume cada resultado por separado. Solo declara una consulta fallida cuando la herramienta devuelva un error o no esté disponible; no deduzcas que Lambda no está soportado a partir del resultado de EC2. No inventes recursos, estados, direcciones IP ni configuraciones. Si no puedes obtener una información mediante las herramientas disponibles, indícalo claramente. Actualmente tu función es únicamente de consulta. No debes iniciar, detener, modificar ni eliminar recursos AWS. """ # Define a collection of tools used by the model tools = [] _INLINE_FUNCTION_NAMES = set() # Define a simple function tool @tool def add_numbers(a: int, b: int) -> int: """Return the sum of two numbers""" return a+b tools.append(add_numbers) # Add MCP clients to tools for mcp_client in mcp_clients: if mcp_client: tools.append(mcp_client) def _make_conversation_manager(): return NullConversationManager() def agent_factory(): cache = {} def get_or_create_agent(session_id, user_id): _actor_id = user_id key = f"{session_id}/{_actor_id}" if key not in cache: cache[key] = Agent( model=load_model(), session_manager=get_memory_session_manager(session_id, _actor_id), conversation_manager=_make_conversation_manager(), system_prompt=DEFAULT_SYSTEM_PROMPT, tools=tools, hooks=[ ], ) return cache[key] return get_or_create_agent get_or_create_agent = agent_factory() def strip_trailing_tool_use(messages: Any) -> list[dict]: """Strip toolUse blocks from the tail until the last message has none.""" if not isinstance(messages, list): raise ValueError("messages must be a list") messages = list(messages) while messages: last = messages[-1] if not isinstance(last, dict): raise ValueError("each message must be an object") original_content = last.get("content", []) if not isinstance(original_content, list) or not all(isinstance(block, dict) for block in original_content): raise ValueError("each message content value must be a list of content blocks") content = [block for block in original_content if "toolUse" not in block] if len(content) == len(original_content): break if content: messages[-1] = {**last, "content": content} break messages.pop() return messages def _extract_prompt(payload: dict): """Accept validated harness messages, tool results, or a plain prompt string.""" if not isinstance(payload, dict): raise ValueError("payload must be a JSON object") if "messages" in payload: return strip_trailing_tool_use(payload["messages"]) if "tool_results" in payload: tool_results = payload["tool_results"] if not isinstance(tool_results, list) or not all( isinstance(tool_result, dict) and isinstance(tool_result.get("toolUseId"), str) for tool_result in tool_results ): raise ValueError("tool_results must contain objects with a toolUseId string") return [{"role": "user", "content": [{"toolResult": { "toolUseId": tr["toolUseId"], "status": tr.get("status", "success"), "content": tr.get("content", []), }} for tr in tool_results]}] prompt = payload.get("prompt", "") if not isinstance(prompt, str): raise ValueError("prompt must be a string") return prompt def _has_inline_function_call(messages) -> bool: """Return True if messages contains an assistant toolUse for an inline function tool.""" if not _INLINE_FUNCTION_NAMES or not isinstance(messages, list): return False for msg in messages: if msg.get("role") == "assistant": for block in msg.get("content", []): if isinstance(block, dict) and block.get("toolUse", {}).get("name") in _INLINE_FUNCTION_NAMES: return True return False def _is_inline_function_call(event: dict) -> bool: """Check if a contentBlockStart event is for an inline function tool.""" if not _INLINE_FUNCTION_NAMES: return False cbs = event.get("contentBlockStart", {}) start = cbs.get("start", {}) tool_use = start.get("toolUse") if isinstance(start, dict) else None return tool_use is not None and tool_use.get("name") in _INLINE_FUNCTION_NAMES @app.entrypoint async def invoke(payload, context): log.info("Invoking Agent.....") session_id = getattr(context, 'session_id', 'default-session') user_id = getattr(context, 'user_id', 'default-user') agent = get_or_create_agent(session_id, user_id) prompt = _extract_prompt(payload) async for event in agent.stream_async( prompt, ): if not isinstance(event, dict) or "event" not in event: continue cbs = event["event"].get("contentBlockStart") if cbs is not None and not cbs.get("start"): continue yield event if __name__ == "__main__": app.run()