Skip to content

Add Azure Open AI support #158

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
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add Azure Open AI support #158

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bf85613
Add Azure OPenAI support
RamananVr Mar 9, 2025
0daf783
Merge branch 'main' of https://github.com/virattt/ai-hedge-fund
RamananVr Mar 9, 2025
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
29 changes: 27 additions & 2 deletions src/llm/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from langchain_deepseek import ChatDeepSeek
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_openai import ChatOpenAI, AzureChatOpenAI
from enum import Enum
from pydantic import BaseModel
from typing import Tuple
Expand All @@ -12,6 +12,7 @@
class ModelProvider(str, Enum):
"""Enum for supported LLM providers"""
ANTHROPIC = "Anthropic"
AZURE_OPENAI = "AzureOpenAI"
DEEPSEEK = "DeepSeek"
GEMINI = "Gemini"
GROQ = "Groq"
Expand Down Expand Up @@ -89,6 +90,12 @@ def is_gemini(self) -> bool:
model_name="gpt-4.5-preview",
provider=ModelProvider.OPENAI
),
# Add Azure OpenAI models
LLMModel(
display_name="[azure] gpt-4o",
model_name="gpt-4o", # This should match your Azure deployment name
provider=ModelProvider.AZURE_OPENAI
),
LLMModel(
display_name="[openai] gpt-4o",
model_name="gpt-4o",
Expand All @@ -103,7 +110,7 @@ def is_gemini(self) -> bool:
display_name="[openai] o3-mini",
model_name="o3-mini",
provider=ModelProvider.OPENAI
),
)
]

# Create LLM_ORDER in the format expected by the UI
Expand Down Expand Up @@ -135,6 +142,24 @@ def get_model(model_name: str, model_provider: ModelProvider) -> ChatOpenAI | Ch
print(f"API Key Error: Please make sure ANTHROPIC_API_KEY is set in your .env file.")
raise ValueError("Anthropic API key not found. Please make sure ANTHROPIC_API_KEY is set in your .env file.")
return ChatAnthropic(model=model_name, api_key=api_key)
elif model_provider == ModelProvider.AZURE_OPENAI:
api_key = os.getenv("AZURE_OPENAI_API_KEY")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
api_version = os.getenv("AZURE_OPENAI_API_VERSION")
if not api_key or not endpoint:
print(f"API Key Error: Please make sure AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT are set in your .env file.")
raise ValueError("Azure OpenAI settings not found. Please set AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT in your .env file.")

# Default API version if not specified
if not api_version:
api_version = "2023-05-15"

return AzureChatOpenAI(
azure_deployment=model_name,
openai_api_version=api_version,
azure_endpoint=endpoint,
api_key=api_key
)
elif model_provider == ModelProvider.DEEPSEEK:
api_key = os.getenv("DEEPSEEK_API_KEY")
if not api_key:
Expand Down