Merge pull request #112 from aws-samples/rxu/deep-research

Modify to support Bedrock

Author: xuro-langchain

langchain/deep-research/.env.example

@@ -1,4 +1,7 @@
 OPENAI_API_KEY=
+BEDROCK_AWS_ACCESS_KEY_ID=
+BEDROCK_AWS_SECRET_ACCESS_KEY=
+BEDROCK_AWS_REGION=
 ANTHROPIC_API_KEY=
 GOOGLE_API_KEY=
 TAVILY_API_KEY=

langchain/deep-research/src/agent/configuration.py

@@ -40,7 +40,7 @@ class Configuration(BaseModel):
 
     # General Configuration
     max_structured_output_retries: int = Field(
-        default=3,
+        default=2,
         metadata={
             "x_oap_ui_config": {
                 "type": "number",
@@ -92,11 +92,11 @@ class Configuration(BaseModel):
         }
     )
     max_researcher_iterations: int = Field(
-        default=6,
+        default=3,
         metadata={
             "x_oap_ui_config": {
                 "type": "slider",
-                "default": 4,
+                "default": 3,
                 "min": 1,
                 "max": 10,
                 "step": 1,
@@ -105,11 +105,11 @@ class Configuration(BaseModel):
         }
     )
     max_react_tool_calls: int = Field(
-        default=10,
+        default=4,
         metadata={
             "x_oap_ui_config": {
                 "type": "slider",
-                "default": 10,
+                "default": 4,
                 "min": 1,
                 "max": 30,
                 "step": 1,
@@ -119,11 +119,11 @@ class Configuration(BaseModel):
     )
     # Model Configuration
     summarization_model: str = Field(
-        default="openai:gpt-4.1-nano",
+        default="bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
         metadata={
             "x_oap_ui_config": {
                 "type": "text",
-                "default": "openai:gpt-4.1-nano",
+                "default": "bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
                 "description": "Model for summarizing research results from Tavily search results"
             }
         }
@@ -151,11 +151,11 @@ class Configuration(BaseModel):
         }
     )
     research_model: str = Field(
-        default="openai:gpt-4.1-nano",
+        default="bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
         metadata={
             "x_oap_ui_config": {
                 "type": "text",
-                "default": "openai:gpt-4.1-nano",
+                "default": "bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
                 "description": "Model for conducting research. NOTE: Make sure your Researcher Model supports the selected search API."
             }
         }
@@ -171,11 +171,11 @@ class Configuration(BaseModel):
         }
     )
     compression_model: str = Field(
-        default="openai:gpt-4.1-nano",
+        default="bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
         metadata={
             "x_oap_ui_config": {
                 "type": "text",
-                "default": "openai:gpt-4.1-nano",
+                "default": "bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
                 "description": "Model for compressing research findings from sub-agents. NOTE: Make sure your Compression Model supports the selected search API."
             }
         }
@@ -191,11 +191,11 @@ class Configuration(BaseModel):
         }
     )
     final_report_model: str = Field(
-        default="openai:gpt-4.1-nano",
+        default="bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
         metadata={
             "x_oap_ui_config": {
                 "type": "text",
-                "default": "openai:gpt-4.1-nano",
+                "default": "bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0",
                 "description": "Model for writing the final report from all research findings"
             }
         }

langchain/deep-research/src/agent/graph.py

@@ -41,6 +41,7 @@
 )
 from agent.utils import (
     anthropic_websearch_called,
+    build_model_config,
     get_all_tools,
     get_api_key_for_model,
     get_model_token_limit,
@@ -78,12 +79,12 @@ async def clarify_with_user(state: AgentState, config: RunnableConfig) -> Comman
 
     # Step 2: Prepare the model for structured clarification analysis
     messages = state["messages"]
-    model_config = {
-        "model": configurable.research_model,
-        "max_tokens": configurable.research_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.research_model, config),
-        "tags": ["langsmith:nostream"]
-    }
+    model_config = build_model_config(
+        configurable.research_model,
+        configurable.research_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
 
     # Configure model with structured output and retry logic
     clarification_model = (
@@ -131,12 +132,12 @@ async def write_research_brief(state: AgentState, config: RunnableConfig) -> Com
     """
     # Step 1: Set up the research model for structured output
     configurable = Configuration.from_runnable_config(config)
-    research_model_config = {
-        "model": configurable.research_model,
-        "max_tokens": configurable.research_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.research_model, config),
-        "tags": ["langsmith:nostream"]
-    }
+    research_model_config = build_model_config(
+        configurable.research_model,
+        configurable.research_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
 
     # Configure model for structured research question generation
     research_model = (
@@ -191,12 +192,12 @@ async def supervisor(state: SupervisorState, config: RunnableConfig) -> Command[
     """
     # Step 1: Configure the supervisor model with available tools
     configurable = Configuration.from_runnable_config(config)
-    research_model_config = {
-        "model": configurable.research_model,
-        "max_tokens": configurable.research_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.research_model, config),
-        "tags": ["langsmith:nostream"]
-    }
+    research_model_config = build_model_config(
+        configurable.research_model,
+        configurable.research_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
 
     # Available tools: research delegation, completion signaling, and strategic thinking
     lead_researcher_tools = [ConductResearch, ResearchComplete, think_tool]
@@ -389,12 +390,12 @@ async def researcher(state: ResearcherState, config: RunnableConfig) -> Command[
         )
 
     # Step 2: Configure the researcher model with tools
-    research_model_config = {
-        "model": configurable.research_model,
-        "max_tokens": configurable.research_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.research_model, config),
-        "tags": ["langsmith:nostream"]
-    }
+    research_model_config = build_model_config(
+        configurable.research_model,
+        configurable.research_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
 
     # Prepare system prompt with MCP context if available
     researcher_prompt = research_system_prompt.format(
@@ -524,12 +525,13 @@ async def compress_research(state: ResearcherState, config: RunnableConfig):
     """
     # Step 1: Configure the compression model
     configurable = Configuration.from_runnable_config(config)
-    synthesizer_model = configurable_model.with_config({
-        "model": configurable.compression_model,
-        "max_tokens": configurable.compression_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.compression_model, config),
-        "tags": ["langsmith:nostream"]
-    })
+    compression_model_config = build_model_config(
+        configurable.compression_model,
+        configurable.compression_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
+    synthesizer_model = configurable_model.with_config(compression_model_config)
 
     # Step 2: Prepare messages for compression
     researcher_messages = state.get("researcher_messages", [])
@@ -624,12 +626,12 @@ async def final_report_generation(state: AgentState, config: RunnableConfig):
 
     # Step 2: Configure the final report generation model
     configurable = Configuration.from_runnable_config(config)
-    writer_model_config = {
-        "model": configurable.final_report_model,
-        "max_tokens": configurable.final_report_model_max_tokens,
-        "api_key": get_api_key_for_model(configurable.final_report_model, config),
-        "tags": ["langsmith:nostream"]
-    }
+    writer_model_config = build_model_config(
+        configurable.final_report_model,
+        configurable.final_report_model_max_tokens,
+        config,
+        tags=["langsmith:nostream"]
+    )
 
     # Step 3: Attempt report generation with token limit retry logic
     max_retries = 3

langchain/deep-research/src/agent/utils.py

@@ -33,6 +33,22 @@
 from agent.prompts import summarize_webpage_prompt
 from agent.state import ResearchComplete, Summary
 
+##########################
+# AWS Credentials Setup
+##########################
+def setup_bedrock_credentials():
+    """Set up AWS credentials from BEDROCK_* environment variables if standard AWS_* vars are not set."""
+    # Only set if standard AWS credentials are not already set
+    if not os.getenv("AWS_ACCESS_KEY_ID") and os.getenv("BEDROCK_AWS_ACCESS_KEY_ID"):
+        os.environ["AWS_ACCESS_KEY_ID"] = os.getenv("BEDROCK_AWS_ACCESS_KEY_ID")
+    if not os.getenv("AWS_SECRET_ACCESS_KEY") and os.getenv("BEDROCK_AWS_SECRET_ACCESS_KEY"):
+        os.environ["AWS_SECRET_ACCESS_KEY"] = os.getenv("BEDROCK_AWS_SECRET_ACCESS_KEY")
+    if not os.getenv("AWS_DEFAULT_REGION") and os.getenv("BEDROCK_AWS_REGION"):
+        os.environ["AWS_DEFAULT_REGION"] = os.getenv("BEDROCK_AWS_REGION")
+    # Also set AWS_REGION if not set
+    if not os.getenv("AWS_REGION") and os.getenv("BEDROCK_AWS_REGION"):
+        os.environ["AWS_REGION"] = os.getenv("BEDROCK_AWS_REGION")
+
 ##########################
 # Tavily Search Tool Utils
 ##########################
@@ -82,12 +98,14 @@ async def tavily_search(
     max_char_to_include = configurable.max_content_length
 
     # Initialize summarization model with retry logic
-    model_api_key = get_api_key_for_model(configurable.summarization_model, config)
-    summarization_model = init_chat_model(
-        model=configurable.summarization_model,
-        max_tokens=configurable.summarization_model_max_tokens,
-        api_key=model_api_key,
+    summarization_model_config = build_model_config(
+        configurable.summarization_model,
+        configurable.summarization_model_max_tokens,
+        config,
         tags=["langsmith:nostream"]
+    )
+    summarization_model = init_chat_model(
+        **summarization_model_config
     ).with_structured_output(Summary).with_retry(
         stop_after_attempt=configurable.max_structured_output_retries
     )
@@ -890,9 +908,20 @@ def get_config_value(value):
         return value.value
 
 def get_api_key_for_model(model_name: str, config: RunnableConfig):
-    """Get API key for a specific model from environment or config."""
+    """Get API key for a specific model from environment or config.
+    
+    For Bedrock models, this sets up AWS credentials from BEDROCK_* env vars and returns None
+    (Bedrock uses AWS credentials instead of API keys).
+    """
     should_get_from_config = os.getenv("GET_API_KEYS_FROM_CONFIG", "false")
     model_name = model_name.lower()
+    
+    # Bedrock models don't use API keys, they use AWS credentials
+    # Set up AWS credentials from BEDROCK_* env vars if this is a Bedrock model
+    if model_name.startswith("bedrock:"):
+        setup_bedrock_credentials()
+        return None
+    
     if should_get_from_config.lower() == "true":
         api_keys = config.get("configurable", {}).get("apiKeys", {})
         if not api_keys:
@@ -913,6 +942,33 @@ def get_api_key_for_model(model_name: str, config: RunnableConfig):
             return os.getenv("GOOGLE_API_KEY")
         return None
 
+def build_model_config(model_name: str, max_tokens: int, config: RunnableConfig, tags: Optional[List[str]] = None):
+    """Build a model configuration dictionary, excluding api_key for Bedrock models.
+    
+    Args:
+        model_name: The model identifier
+        max_tokens: Maximum tokens for the model
+        config: Runtime configuration
+        tags: Optional list of tags to include
+        
+    Returns:
+        Dictionary with model configuration, excluding api_key for Bedrock models
+    """
+    api_key = get_api_key_for_model(model_name, config)
+    model_config = {
+        "model": model_name,
+        "max_tokens": max_tokens,
+    }
+    
+    # Only include api_key if it's not None (i.e., not a Bedrock model)
+    if api_key is not None:
+        model_config["api_key"] = api_key
+    
+    if tags:
+        model_config["tags"] = tags
+    
+    return model_config
+
 def get_tavily_api_key(config: RunnableConfig):
     """Get Tavily API key from environment or config."""
     should_get_from_config = os.getenv("GET_API_KEYS_FROM_CONFIG", "false")