From 419efaa395ea1660e91a6e280fa84610a3ab6873 Mon Sep 17 00:00:00 2001
From: yongtenglei <yongtengrey@outlook.com>
Date: Fri, 12 Dec 2025 19:02:37 +0800
Subject: [PATCH] Fix: Asure-OpenAI resource not found

---
 rag/llm/__init__.py   |  3 ++-
 rag/llm/chat_model.py | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
index 1e80846a8..4b36d8a05 100644
--- a/rag/llm/__init__.py
+++ b/rag/llm/__init__.py
@@ -55,6 +55,7 @@ class SupportedLiteLLMProvider(StrEnum):
     DeerAPI = "DeerAPI"
     GPUStack = "GPUStack"
     OpenAI = "OpenAI"
+    Azure_OpenAI = "Azure-OpenAI"
 
 
 FACTORY_DEFAULT_BASE_URL = {
@@ -116,7 +117,7 @@ LITELLM_PROVIDER_PREFIX = {
     SupportedLiteLLMProvider.DeerAPI: "openai/",
     SupportedLiteLLMProvider.GPUStack: "openai/",
     SupportedLiteLLMProvider.OpenAI: "openai/",
-
+    SupportedLiteLLMProvider.Azure_OpenAI: "azure/",
 }
 
 ChatModel = globals().get("ChatModel", {})
diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
index 214cd4604..9964eb7f8 100644
--- a/rag/llm/chat_model.py
+++ b/rag/llm/chat_model.py
@@ -191,7 +191,7 @@ class Base(ABC):
             except Exception as e:
                 e = await self._exceptions_async(e, attempt)
                 if e:
-                    yield e                    
+                    yield e
                     yield total_tokens
                     return
 
@@ -1219,6 +1219,7 @@ class LiteLLMBase(ABC):
         "DeerAPI",
         "GPUStack",
         "OpenAI",
+        "Azure-OpenAI",
     ]
 
     def __init__(self, key, model_name, base_url=None, **kwargs):
@@ -1244,6 +1245,9 @@ class LiteLLMBase(ABC):
         elif self.provider == SupportedLiteLLMProvider.OpenRouter:
             self.api_key = json.loads(key).get("api_key", "")
             self.provider_order = json.loads(key).get("provider_order", "")
+        elif self.provider == SupportedLiteLLMProvider.Azure_OpenAI:
+            self.api_key = json.loads(key).get("api_key", "")
+            self.api_version = json.loads(key).get("api_version", "2024-02-01")
 
     def _get_delay(self):
         return self.base_delay * random.uniform(10, 150)
@@ -1675,6 +1679,13 @@ class LiteLLMBase(ABC):
                     "api_base": self.base_url,
                 }
             )
+        elif self.provider == SupportedLiteLLMProvider.GPUStack:
+            completion_args.update(
+                {
+                    "api_base": self.base_url,
+                    "api_version": self.api_version,
+                }
+            )
 
         # Ollama deployments commonly sit behind a reverse proxy that enforces
         # Bearer auth. Ensure the Authorization header is set when an API key