From 4ed3b9fbcc72f6f60c9960d34f5d574a5ada2169 Mon Sep 17 00:00:00 2001
From: harry <harry@wangxutech.com>
Date: Mon, 25 Mar 2024 15:20:00 +0800
Subject: [PATCH] supported azure openai

---
 app/services/llm.py | 23 ++++++++++++++++++-----
 config.example.toml | 19 ++++++++++++++-----
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/app/services/llm.py b/app/services/llm.py
index 97f3fff..764d2f3 100644
--- a/app/services/llm.py
+++ b/app/services/llm.py
@@ -5,7 +5,7 @@ from typing import List
 import g4f
 from loguru import logger
 from openai import OpenAI
-
+from openai import AzureOpenAI
 from app.config import config
 
 
@@ -23,6 +23,7 @@ def _generate_response(prompt: str) -> str:
             messages=[{"role": "user", "content": prompt}],
         )
     else:
+        api_version = ""  # for azure
         if llm_provider == "moonshot":
             api_key = config.app.get("moonshot_api_key")
             model_name = config.app.get("moonshot_model_name")
@@ -37,6 +38,11 @@ def _generate_response(prompt: str) -> str:
             api_key = config.app.get("oneapi_api_key")
             model_name = config.app.get("oneapi_model_name")
             base_url = config.app.get("oneapi_base_url", "")
+        elif llm_provider == "azure":
+            api_key = config.app.get("azure_api_key")
+            model_name = config.app.get("azure_model_name")
+            base_url = config.app.get("azure_base_url", "")
+            api_version = config.app.get("azure_api_version", "2024-02-15-preview")
         else:
             raise ValueError("llm_provider is not set, please set it in the config.toml file.")
 
@@ -47,10 +53,17 @@ def _generate_response(prompt: str) -> str:
         if not base_url:
             raise ValueError(f"{llm_provider}: base_url is not set, please set it in the config.toml file.")
 
-        client = OpenAI(
-            api_key=api_key,
-            base_url=base_url,
-        )
+        if llm_provider == "azure":
+            client = AzureOpenAI(
+                api_key=api_key,
+                api_version=api_version,
+                azure_endpoint=base_url,
+            )
+        else:
+            client = OpenAI(
+                api_key=api_key,
+                base_url=base_url,
+            )
 
         response = client.chat.completions.create(
             model=model_name,
diff --git a/config.example.toml b/config.example.toml
index e2aafb6..1c4d685 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -7,31 +7,40 @@
     pexels_api_keys = []
 
     # 如果你没有 OPENAI API Key，可以使用 g4f 代替，或者使用国内的 Moonshot API
-    llm_provider="openai" # "openai" or "moonshot" or "oneapi" or "g4f"
+    llm_provider="openai" # "openai" or "moonshot" or "oneapi" or "g4f" or "azure"
 
-    # OpenAI API Key
+    ########## OpenAI API Key
     # Visit https://openai.com/api/ for details on obtaining an API key.
     openai_api_key = ""
     openai_base_url = ""
     openai_model_name = "gpt-4-turbo-preview"
 
-    # Moonshot API Key
+    ########## Moonshot API Key
     # Visit https://platform.moonshot.cn/console/api-keys to get your API key.
     moonshot_api_key=""
     moonshot_base_url = "https://api.moonshot.cn/v1"
     moonshot_model_name = "moonshot-v1-8k"
 
-    # OneAPI API Key
+    ########## OneAPI API Key
     # Visit https://github.com/songquanpeng/one-api to get your API key
     oneapi_api_key=""
     oneapi_base_url=""
     oneapi_model_name=""
 
-    # G4F
+    ########## G4F
     # Visit https://github.com/xtekky/gpt4free to get more details
     # Supported model list: https://github.com/xtekky/gpt4free/blob/main/g4f/models.py
     g4f_model_name = "gpt-3.5-turbo-16k-0613"
 
+    ########## Azure API Key
+    # Visit https://learn.microsoft.com/zh-cn/azure/ai-services/openai/ to get more details
+    # API documentation: https://learn.microsoft.com/zh-cn/azure/ai-services/openai/reference
+    azure_api_key = ""
+    azure_base_url=""
+    azure_model_name="gpt-35-turbo" # replace with your model deployment name
+    azure_api_version = "2024-02-15-preview"
+
+
     # Subtitle Provider, "edge" or "whisper"
     # If empty, the subtitle will not be generated
     subtitle_provider = "edge"