bedrock-agentcore-starter-toolkit 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

bedrock_agentcore_starter_toolkit/services/import_agent/scripts/bedrock_to_langchain.py

@@ -0,0 +1,382 @@
+# pylint: disable=consider-using-f-string, line-too-long
+# ruff: noqa: E501
+"""Bedrock Agent to LangChain Translator.
+
+This script translates AWS Bedrock Agent configurations into equivalent LangChain code.
+"""
+
+import os
+import textwrap
+
+from .base_bedrock_translate import BaseBedrockTranslator
+
+
+class BedrockLangchainTranslation(BaseBedrockTranslator):
+    """Class to translate Bedrock Agent configurations to LangChain code."""
+
+    def __init__(self, agent_config, debug: bool, output_dir: str, enabled_primitives: dict):
+        """Initialize the BedrockLangchainTranslation class."""
+        super().__init__(agent_config, debug, output_dir, enabled_primitives)
+
+        self.imports_code += self.generate_imports()
+        self.tools_code = self.generate_action_groups_code(platform="langchain")
+        self.memory_code = self.generate_memory_configuration(memory_saver="InMemorySaver")
+        self.collaboration_code = self.generate_collaboration_code()
+        self.kb_code = self.generate_knowledge_base_code()
+        self.models_code = self.generate_model_configurations()
+        self.agent_setup_code = self.generate_agent_setup()
+        self.usage_code = self.generate_example_usage()
+
+        # Observability
+        if self.observability_enabled:
+            self.imports_code += """
+    from opentelemetry.instrumentation.langchain import LangchainInstrumentor
+    LangchainInstrumentor().instrument()
+    """
+
+        # Format prompts code
+        self.prompts_code = textwrap.fill(
+            self.prompts_code, width=150, break_long_words=False, replace_whitespace=False
+        )
+
+        self.code_sections = [
+            self.imports_code,
+            self.models_code,
+            self.prompts_code,
+            self.collaboration_code,
+            self.tools_code,
+            self.memory_code,
+            self.kb_code,
+            self.agent_setup_code,
+            self.usage_code,
+        ]
+
+    def generate_imports(self) -> str:
+        """Generate import statements for LangChain components."""
+        return """
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+    from langchain_aws import ChatBedrock
+    from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever
+
+    from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
+    from langchain_core.globals import set_verbose, set_debug
+
+    from langchain.tools import tool
+
+    from langgraph.prebuilt import create_react_agent, InjectedState
+    from langgraph.checkpoint.memory import InMemorySaver
+    """
+
+    def generate_model_configurations(self) -> str:
+        """Generate LangChain model configurations from Bedrock agent config."""
+        model_configs = []
+
+        for i, config in enumerate(self.prompt_configs):
+            prompt_type = config.get("promptType", f"CUSTOM_{i}")
+            inference_config = config.get("inferenceConfiguration", {})
+
+            # Skip KB Generation if no knowledge bases are defined
+            if prompt_type == "KNOWLEDGE_BASE_RESPONSE_GENERATION" and not self.knowledge_bases:
+                continue
+
+            # Build model configuration string
+            model_config = f"""
+    # {prompt_type} LLM configuration
+    llm_{prompt_type} = ChatBedrock(
+        model_id="{self.model_id}",
+        region_name="{self.agent_region}",
+        provider="{self.agent_info["model"]["providerName"].lower()}",
+        model_kwargs={{
+            {f'"top_k": {inference_config.get("topK", 250)},' if self.agent_info["model"]["providerName"].lower() in ["anthropic", "amazon"] else ""}
+            "top_p":{inference_config.get("topP", 1.0)},
+            "temperature": {inference_config.get("temperature", 0)},
+            "max_tokens": {inference_config.get("maximumLength", 2048)},
+            {f'"stop_sequences": {repr(inference_config.get("stopSequences", []))},'.strip() if self.agent_info["model"]["providerName"].lower() in ["anthropic", "amazon"] else ""}
+        }}"""
+
+            # Add guardrails if available
+            if self.guardrail_config:
+                model_config += f""",
+        guardrails={self.guardrail_config}"""
+
+            model_config += "\n)"
+            model_configs.append(model_config)
+
+            self.generate_prompt(config)
+
+        return "\n".join(model_configs)
+
+    def generate_knowledge_base_code(self) -> str:
+        """Generate code for knowledge base retrievers."""
+        if not self.knowledge_bases:
+            return ""
+
+        kb_code = ""
+
+        for kb in self.knowledge_bases:
+            kb_name = kb.get("name", "")
+            kb_description = kb.get("description", "")
+            kb_id = kb.get("knowledgeBaseId", "")
+            kb_region_name = kb.get("knowledgeBaseArn", "").split(":")[3]
+
+            kb_code += f"""retriever_{kb_name} = AmazonKnowledgeBasesRetriever(
+        knowledge_base_id="{kb_id}",
+        retrieval_config={{"vectorSearchConfiguration": {{"numberOfResults": 5}}}},
+        region_name="{kb_region_name}"
+    )
+
+    retriever_tool_{kb_name} = retriever_{kb_name}.as_tool(name="kb_{kb_name}", description="{kb_description}")
+
+    """
+            self.tools.append(f"retriever_tool_{kb_name}")
+
+        return kb_code
+
+    def generate_collaboration_code(self) -> str:
+        """Generate code for multi-agent collaboration."""
+        if not self.multi_agent_enabled or not self.collaborators:
+            return ""
+
+        collaborator_code = ""
+
+        # Create the collaborators
+        for i, collaborator in enumerate(self.collaborators):
+            collaborator_name = collaborator.get("collaboratorName", "")
+            collaborator_file_name = f"langchain_collaborator_{collaborator_name}"
+            collaborator_path = os.path.join(self.output_dir, f"{collaborator_file_name}.py")
+
+            # Recursively translate the collaborator agent to LangChain
+            BedrockLangchainTranslation(
+                collaborator, debug=self.debug, output_dir=self.output_dir, enabled_primitives=self.enabled_primitives
+            ).translate_bedrock_to_langchain(collaborator_path)
+
+            self.imports_code += (
+                f"\nfrom {collaborator_file_name} import invoke_agent as invoke_{collaborator_name}_collaborator"
+            )
+
+            # conversation relay
+            relay_conversation_history = collaborator.get("relayConversationHistory", "DISABLED") == "TO_COLLABORATOR"
+
+            # Create tool to invoke the collaborator
+            collaborator_code += """
+    @tool
+    def invoke_{0}(query: str, state: Annotated[dict, InjectedState]) -> str:
+        \"\"\"Invoke the collaborator agent/specialist with the following description: {1}\"\"\"
+        {2}
+        invoke_agent_response = invoke_{0}_collaborator(query{3})
+        tools_used.update([msg.name for msg in invoke_agent_response if isinstance(msg, ToolMessage)])
+        return invoke_agent_response
+        """.format(
+                collaborator_name,
+                self.collaborator_descriptions[i],
+                "relay_history = state.get('messages', [])[:-1]" if relay_conversation_history else "",
+                ", relay_history" if relay_conversation_history else "",
+            )
+
+            # Add the tool to the list of tools
+            self.tools.append(f"invoke_{collaborator_name}")
+
+        return collaborator_code
+
+    def generate_agent_setup(self) -> str:
+        """Generate agent setup code."""
+        agent_code = f"tools = [{','.join(self.tools)}]\ntools_used = set()"
+
+        if self.action_groups and self.tools_code:
+            agent_code += """\ntools += action_group_tools"""
+
+        if self.gateway_enabled:
+            agent_code += """\ntools += mcp_tools"""
+
+        memory_retrieve_code = (
+            ""
+            if not self.memory_enabled
+            else (
+                "memory_synopsis = memory_manager.get_memory_synopsis()"
+                if not self.agentcore_memory_enabled
+                else """
+            memories = memory_client.retrieve_memories(memory_id=memory_id, namespace=f'/summaries/{user_id}', query="Retrieve the most recent session sumamries.", actor_id=user_id, top_k=20)
+            memory_synopsis = "\\n".join([m.get("content", {}).get("text", "") for m in memories])
+"""
+            )
+        )
+
+        # Create agent based on available components
+        agent_code += """
+    config = {{"configurable": {{"thread_id": "1"}}}}
+    set_verbose({})
+    set_debug({})
+
+    _agent = None
+    first_turn = True
+    last_input = ""
+    user_id = ""
+    {}
+
+    # agent update loop
+    def get_agent():
+
+        global _agent, user_id, memory_id
+
+        {}
+            {}
+            system_prompt = ORCHESTRATION_TEMPLATE
+            {}
+            _agent = create_react_agent(
+                model=llm_ORCHESTRATION,
+                prompt=system_prompt,
+                tools=tools,
+                checkpointer=checkpointer_STM,
+                debug={}
+            )
+
+        return _agent
+""".format(
+            self.debug,
+            self.debug,
+            'last_agent = ""' if self.multi_agent_enabled and self.supervision_type == "SUPERVISOR_ROUTER" else "",
+            (
+                "if _agent is None or memory_manager.has_memory_changed():"
+                if self.memory_enabled and not self.agentcore_memory_enabled
+                else "if _agent is None:"
+            ),
+            memory_retrieve_code,
+            (
+                "system_prompt = system_prompt.replace('$memory_synopsis$', memory_synopsis)"
+                if self.memory_enabled
+                else ""
+            ),
+            self.debug,
+        )
+
+        # Generate routing code if needed
+        routing_code = self.generate_routing_code()
+
+        # Set up relay parameter definition based on whether we're accepting relays
+        relay_param_def = ", relayed_messages = []" if self.is_accepting_relays else ""
+
+        # Add relay handling code if needed
+        relay_code = (
+            """if relayed_messages:
+            agent.update_state(config, {"messages": relayed_messages})"""
+            if self.is_accepting_relays
+            else ""
+        )
+
+        # Set up preprocessing code if enabled
+        preprocess_code = ""
+        if "PRE_PROCESSING" in self.enabled_prompts:
+            preprocess_code = """
+        pre_process_output = llm_PRE_PROCESSING.invoke([SystemMessage(PRE_PROCESSING_TEMPLATE), HumanMessage(question)])
+        question += "\\n<PRE_PROCESSING>{}</PRE_PROCESSING>".format(pre_process_output.content)
+"""
+            if self.debug:
+                preprocess_code += '        print("PREPROCESSING_OUTPUT: {pre_process_output}")'
+
+        # Memory recording code
+        memory_add_user = (
+            """
+        memory_manager.add_message({'role': 'user', 'content': question})"""
+            if self.memory_enabled and not self.agentcore_memory_enabled
+            else ""
+        )
+
+        memory_add_assistant = (
+            """
+        memory_manager.add_message({'role': 'assistant', 'content': str(response)})"""
+            if self.memory_enabled and not self.agentcore_memory_enabled
+            else ""
+        )
+
+        # KB optimization code if enabled
+        kb_code = ""
+        if self.single_kb_optimization_enabled:
+            kb_name = self.knowledge_bases[0]["name"]
+            kb_code = f"""
+        if first_turn:
+            search_results = retriever_{kb_name}.invoke(question)
+            response = llm_KNOWLEDGE_BASE_RESPONSE_GENERATION.invoke([SystemMessage(KB_GENERATION_TEMPLATE.replace("$search_results$, search_results)), HumanMessage(question))])
+            first_turn = False
+"""
+
+        # Post-processing code
+        post_process_code = (
+            """
+        post_process_prompt = POST_PROCESSING_TEMPLATE.replace("$question$", question).replace("$latest_response$", response["messages"][-1].content).replace("$responses$", str(response["messages"]))
+        post_process_output = llm_POST_PROCESSING.invoke([HumanMessage(post_process_prompt)])
+        return [AIMessage(post_process_output.content)]"""
+            if "POST_PROCESSING" in self.enabled_prompts
+            else "return response['messages']"
+        )
+
+        # Combine it all into the invoke_agent function
+        agent_code += f"""
+    def invoke_agent(question: str{relay_param_def}):
+        {"global last_agent" if self.supervision_type == "SUPERVISOR_ROUTER" else ""}
+        {"global first_turn" if self.single_kb_optimization_enabled else ""}
+        global last_input, memory_id
+        last_input = question
+        agent = get_agent()
+        {relay_code}
+        {routing_code}
+        {preprocess_code}
+        {memory_add_user}
+
+        response = asyncio.run(agent.ainvoke({{"messages": [{{"role": "user", "content": question}}]}}, config))
+        {memory_add_assistant}
+        {kb_code}
+        {post_process_code}
+        """
+
+        agent_code += self.generate_entrypoint_code("langchain")
+
+        return agent_code
+
+    def generate_routing_code(self):
+        """Generate routing code for supervisor router."""
+        if not self.multi_agent_enabled or self.supervision_type != "SUPERVISOR_ROUTER":
+            return ""
+
+        code = """
+        conversation = agent.checkpointer.get(config)
+        if not conversation:
+            conversation = {}
+        messages = str(conversation.get("channel_values", {}).get("messages", []))
+
+        routing_template = ROUTING_TEMPLATE
+        routing_template = routing_template.replace("$last_user_request$", question).replace("$conversation$", messages).replace("$last_most_specialized_agent$", last_agent)
+        routing_choice = llm_ROUTING_CLASSIFIER.invoke([SystemMessage(routing_template), HumanMessage(question)]).content
+
+        choice = str(re.findall(r'<a.*?>(.*?)</a>', routing_choice)[0])"""
+
+        if self.debug:
+            code += """
+        print("Routing to agent: {}. Last used agent was {}.".format(choice, last_agent))"""
+
+        code += """
+        if choice == "undecidable":
+            pass"""
+
+        for agent in self.collaborators:
+            agent_name = agent.get("collaboratorName", "")
+            relay_param = (
+                ", messages"
+                if self.collaborator_map.get(agent_name, {}).get("relayConversationHistory", "DISABLED")
+                == "TO_COLLABORATOR"
+                else ""
+            )
+            code += f"""
+        elif choice == "{agent_name}":
+            last_agent = "{agent_name}"
+            return invoke_{agent_name}_collaborator(question{relay_param})"""
+
+        code += """
+        elif choice == "keep_previous_agent":
+            return eval(f"invoke_{last_agent}_collaborator")(question, messages)"""
+
+        return code
+
+    def translate_bedrock_to_langchain(self, output_path: str) -> dict:
+        """Translate Bedrock agent config to LangChain code."""
+        return self.translate(output_path, self.code_sections, "langchain")