feat: support AzureChatOpenAI under configuring azure_endpoint or AZURE_OPENAI_ENDPOINT (#237)

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

src/llms/llm.py

@@ -6,7 +6,8 @@ from typing import Any, Dict
 import os
 import httpx
 
-from langchain_openai import ChatOpenAI
+from langchain_core.language_models import BaseChatModel
+from langchain_openai import ChatOpenAI, AzureChatOpenAI
 from langchain_deepseek import ChatDeepSeek
 from typing import get_args
 
@@ -14,7 +15,7 @@ from src.config import load_yaml_config
 from src.config.agents import LLMType
 
 # Cache for LLM instances
-_llm_cache: dict[LLMType, ChatOpenAI] = {}
+_llm_cache: dict[LLMType, BaseChatModel] = {}
 
 
 def _get_config_file_path() -> str:
@@ -48,7 +49,7 @@ def _get_env_llm_conf(llm_type: str) -> Dict[str, Any]:
 
 def _create_llm_use_conf(
     llm_type: LLMType, conf: Dict[str, Any]
-) -> ChatOpenAI | ChatDeepSeek:
+) -> BaseChatModel :
     """Create LLM instance using configuration."""
     llm_type_config_keys = _get_llm_type_config_keys()
     config_key = llm_type_config_keys.get(llm_type)
@@ -86,16 +87,17 @@ def _create_llm_use_conf(
         merged_conf["http_client"] = http_client
         merged_conf["http_async_client"] = http_async_client
 
-    return (
-        ChatOpenAI(**merged_conf)
-        if llm_type != "reasoning"
-        else ChatDeepSeek(**merged_conf)
-    )
-
+    if "azure_endpoint" in merged_conf or os.getenv("AZURE_OPENAI_ENDPOINT"):
+        return AzureChatOpenAI(**merged_conf)
+    if llm_type == "reasoning":
+        return ChatDeepSeek(**merged_conf)
+    else 
+        return ChatOpenAI(**merged_conf)
+        
 
 def get_llm_by_type(
     llm_type: LLMType,
-) -> ChatOpenAI:
+) -> BaseChatModel:
     """
     Get LLM instance by type. Returns cached instance if available.
     """