From 9897d675006e66fff2dedc8798af360f710b2f33 Mon Sep 17 00:00:00 2001 From: Robert Xu Date: Tue, 18 Nov 2025 22:58:45 -0500 Subject: [PATCH] Modify to support Bedrock --- langchain/deep-research/.env.example | 3 + .../deep-research/src/agent/configuration.py | 26 +++---- langchain/deep-research/src/agent/graph.py | 74 ++++++++++--------- langchain/deep-research/src/agent/utils.py | 68 +++++++++++++++-- 4 files changed, 116 insertions(+), 55 deletions(-) diff --git a/langchain/deep-research/.env.example b/langchain/deep-research/.env.example index 778272b..058ecaa 100644 --- a/langchain/deep-research/.env.example +++ b/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= diff --git a/langchain/deep-research/src/agent/configuration.py b/langchain/deep-research/src/agent/configuration.py index 66e08c3..7faf25f 100644 --- a/langchain/deep-research/src/agent/configuration.py +++ b/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" } } diff --git a/langchain/deep-research/src/agent/graph.py b/langchain/deep-research/src/agent/graph.py index 18ddbb0..2847aa2 100644 --- a/langchain/deep-research/src/agent/graph.py +++ b/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 diff --git a/langchain/deep-research/src/agent/utils.py b/langchain/deep-research/src/agent/utils.py index a3271d2..ad4643f 100644 --- a/langchain/deep-research/src/agent/utils.py +++ b/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")