Skip to content

Add Azure OpenAI support and configuration options #211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
abozaralizadeh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Azure OpenAI support and configuration options #211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
OPENAI_API_KEY=
AZURE_OPENAI_API_KEY=
ANTHROPIC_API_KEY=
GOOGLE_API_KEY=
TAVILY_API_KEY=
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "open_deep_research"
version = "0.0.16"
version = "0.0.17"
description = "Planning, research, and report generation."
authors = [
{ name = "Lance Martin" }
Expand Down
33 changes: 33 additions & 0 deletions src/open_deep_research/configuration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ class SearchAPI(Enum):

ANTHROPIC = "anthropic"
OPENAI = "openai"
AZURE_OPENAI = "azure_openai"
TAVILY = "tavily"
NONE = "none"

Expand Down Expand Up @@ -85,6 +86,7 @@ class Configuration(BaseModel):
"options": [
{"label": "Tavily", "value": SearchAPI.TAVILY.value},
{"label": "OpenAI Native Web Search", "value": SearchAPI.OPENAI.value},
{"label": "Azure OpenAI Native Web Search", "value": SearchAPI.AZURE_OPENAI.value},
{"label": "Anthropic Native Web Search", "value": SearchAPI.ANTHROPIC.value},
{"label": "None", "value": SearchAPI.NONE.value}
]
Expand Down Expand Up @@ -210,6 +212,37 @@ class Configuration(BaseModel):
}
}
)
# Azure OpenAI Configuration
azure_endpoint: Optional[str] = Field(
default=None,
optional=True,
metadata={
"x_oap_ui_config": {
"type": "text",
"description": "Azure OpenAI endpoint URL (e.g., https://your-resource.openai.azure.com/)"
}
}
)
deployment_name: Optional[str] = Field(
default=None,
optional=True,
metadata={
"x_oap_ui_config": {
"type": "text",
"description": "Azure OpenAI deployment name"
}
}
)
api_version: Optional[str] = Field(
default="2024-02-15-preview",
optional=True,
metadata={
"x_oap_ui_config": {
"type": "text",
"description": "Azure OpenAI API version"
}
}
)
# MCP server configuration
mcp_config: Optional[MCPConfig] = Field(
default=None,
Expand Down
78 changes: 41 additions & 37 deletions src/open_deep_research/deep_researcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,10 @@
)
from open_deep_research.utils import (
anthropic_websearch_called,
azure_openai_websearch_called,
get_all_tools,
get_api_key_for_model,
get_model_config,
get_model_token_limit,
get_notes_from_tool_calls,
get_today_str,
Expand All @@ -54,7 +56,7 @@

# Initialize a configurable model that we will use throughout the agent
configurable_model = init_chat_model(
configurable_fields=("model", "max_tokens", "api_key"),
configurable_fields=("model", "max_tokens", "api_key", "azure_endpoint", "deployment_name", "api_version"),
)

async def clarify_with_user(state: AgentState, config: RunnableConfig) -> Command[Literal["write_research_brief", "__end__"]]:
Expand All @@ -78,12 +80,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 = get_model_config(
configurable.research_model,
configurable.research_model_max_tokens,
config,
configurable
)

# Configure model with structured output and retry logic
clarification_model = (
Expand Down Expand Up @@ -131,12 +133,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 = get_model_config(
configurable.research_model,
configurable.research_model_max_tokens,
config,
configurable
)

# Configure model for structured research question generation
research_model = (
Expand Down Expand Up @@ -191,12 +193,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 = get_model_config(
configurable.research_model,
configurable.research_model_max_tokens,
config,
configurable
)

# Available tools: research delegation, completion signaling, and strategic thinking
lead_researcher_tools = [ConductResearch, ResearchComplete, think_tool]
Expand Down Expand Up @@ -389,12 +391,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 = get_model_config(
configurable.research_model,
configurable.research_model_max_tokens,
config,
configurable
)

# Prepare system prompt with MCP context if available
researcher_prompt = research_system_prompt.format(
Expand Down Expand Up @@ -457,6 +459,7 @@ async def researcher_tools(state: ResearcherState, config: RunnableConfig) -> Co
has_tool_calls = bool(most_recent_message.tool_calls)
has_native_search = (
openai_websearch_called(most_recent_message) or
azure_openai_websearch_called(most_recent_message) or
anthropic_websearch_called(most_recent_message)
)

Expand Down Expand Up @@ -524,12 +527,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 = get_model_config(
configurable.compression_model,
configurable.compression_model_max_tokens,
config,
configurable
)
synthesizer_model = configurable_model.with_config(compression_model_config)

# Step 2: Prepare messages for compression
researcher_messages = state.get("researcher_messages", [])
Expand Down Expand Up @@ -624,12 +628,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 = get_model_config(
configurable.final_report_model,
configurable.final_report_model_max_tokens,
config,
configurable
)

# Step 3: Attempt report generation with token limit retry logic
max_retries = 3
Expand Down
81 changes: 74 additions & 7 deletions src/open_deep_research/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,13 +82,13 @@ 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,
tags=["langsmith:nostream"]
).with_structured_output(Summary).with_retry(
summarization_model_config = get_model_config(
configurable.summarization_model,
configurable.summarization_model_max_tokens,
config,
configurable
)
summarization_model = init_chat_model(**summarization_model_config).with_structured_output(Summary).with_retry(
stop_after_attempt=configurable.max_structured_output_retries
)

Expand Down Expand Up @@ -657,6 +657,28 @@ def openai_websearch_called(response):

return False

def azure_openai_websearch_called(response):
"""Detect if Azure OpenAI's web search functionality was used in the response.

Args:
response: The response object from Azure OpenAI's API

Returns:
True if web search was called, False otherwise
"""
# Azure OpenAI uses the same web search mechanism as OpenAI
# Check for tool outputs in the response metadata
tool_outputs = response.additional_kwargs.get("tool_outputs")
if not tool_outputs:
return False

# Look for web search calls in the tool outputs
for tool_output in tool_outputs:
if tool_output.get("type") == "web_search_call":
return True

return False


##########################
# Token Limit Exceeded Utils
Expand Down Expand Up @@ -797,6 +819,17 @@ def _check_gemini_token_limit(exception: Exception, error_str: str) -> bool:
"openai:o3-pro": 200000,
"openai:o1": 200000,
"openai:o1-pro": 200000,
"azure_openai:gpt-4.1-mini": 1047576,
"azure_openai:gpt-4.1-nano": 1047576,
"azure_openai:gpt-4.1": 1047576,
"azure_openai:gpt-4o-mini": 128000,
"azure_openai:gpt-4o": 128000,
"azure_openai:o4-mini": 200000,
"azure_openai:o3-mini": 200000,
"azure_openai:o3": 200000,
"azure_openai:o3-pro": 200000,
"azure_openai:o1": 200000,
"azure_openai:o1-pro": 200000,
"anthropic:claude-opus-4": 200000,
"anthropic:claude-sonnet-4": 200000,
"anthropic:claude-3-7-sonnet": 200000,
Expand Down Expand Up @@ -899,6 +932,8 @@ def get_api_key_for_model(model_name: str, config: RunnableConfig):
return None
if model_name.startswith("openai:"):
return api_keys.get("OPENAI_API_KEY")
elif model_name.startswith("azure_openai:"):
return api_keys.get("AZURE_OPENAI_API_KEY")
elif model_name.startswith("anthropic:"):
return api_keys.get("ANTHROPIC_API_KEY")
elif model_name.startswith("google"):
Expand All @@ -907,6 +942,8 @@ def get_api_key_for_model(model_name: str, config: RunnableConfig):
else:
if model_name.startswith("openai:"):
return os.getenv("OPENAI_API_KEY")
elif model_name.startswith("azure_openai:"):
return os.getenv("AZURE_OPENAI_API_KEY")
elif model_name.startswith("anthropic:"):
return os.getenv("ANTHROPIC_API_KEY")
elif model_name.startswith("google"):
Expand All @@ -923,3 +960,33 @@ def get_tavily_api_key(config: RunnableConfig):
return api_keys.get("TAVILY_API_KEY")
else:
return os.getenv("TAVILY_API_KEY")

def get_model_config(model_name: str, max_tokens: int, config: RunnableConfig, configurable: "Configuration") -> dict:
"""Get properly configured model parameters based on model type.

Args:
model_name: The model identifier (e.g., "azure_openai:gpt-4", "openai:gpt-4")
max_tokens: Maximum tokens for the model
config: Runtime configuration
configurable: Configuration object with Azure settings

Returns:
Dictionary with properly configured model parameters
"""
model_config = {
"model": model_name,
"max_tokens": max_tokens,
"api_key": get_api_key_for_model(model_name, config),
"tags": ["langsmith:nostream"]
}

# Only add Azure-specific parameters for Azure OpenAI models
if model_name.startswith("azure_openai:"):
if configurable.azure_endpoint:
model_config["azure_endpoint"] = configurable.azure_endpoint
if configurable.deployment_name:
model_config["deployment_name"] = configurable.deployment_name
if configurable.api_version:
model_config["api_version"] = configurable.api_version

return model_config