From 46f3a987590280c67db5f578f1018269d55dd77e Mon Sep 17 00:00:00 2001
From: Hirokazu Kiyomaru <h.kiyomaru@gmail.com>
Date: Fri, 9 Feb 2024 10:56:21 +0900
Subject: [PATCH] use engine when api type is azure

---
 llm_judge/common.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/llm_judge/common.py b/llm_judge/common.py
index d4e8ded..789bf12 100644
--- a/llm_judge/common.py
+++ b/llm_judge/common.py
@@ -59,12 +59,16 @@ def judge(self, **kwargs):
         ]
         for _ in range(API_MAX_RETRY):
             try:
-                response = openai.ChatCompletion.create(
-                    model=self.model,
-                    messages=messages,
-                    temperature=0,
-                    max_tokens=2048,
-                )
+                params = {
+                    "messages": messages,
+                    "temperature": 0,
+                    "max_tokens": 2048,
+                }
+                if openai.api_type == "azure":
+                    params["engine"] = self.model
+                else:
+                    params["model"] = self.model
+                response = openai.ChatCompletion.create(**params)
                 return response["choices"][0]["message"]["content"]
             except openai.error.OpenAIError as e:
                 logger.warning(f"OpenAI API error: {e}")