diff --git a/vision_agent/llm/__init__.py b/vision_agent/llm/__init__.py
index dd5f5c54..e482f69d 100644
--- a/vision_agent/llm/__init__.py
+++ b/vision_agent/llm/__init__.py
@@ -1 +1 @@
-from .llm import LLM, OpenAILLM
+from .llm import LLM, AzureOpenAILLM, OpenAILLM
diff --git a/vision_agent/lmm/__init__.py b/vision_agent/lmm/__init__.py
index 26bc23c1..9c7ace7a 100644
--- a/vision_agent/lmm/__init__.py
+++ b/vision_agent/lmm/__init__.py
@@ -1 +1 @@
-from .lmm import LMM, LLaVALMM, OpenAILMM, get_lmm
+from .lmm import LMM, AzureOpenAILMM, LLaVALMM, OpenAILMM, get_lmm