From 0d4ddae5cb500b05a60871c19969ff456a22670d Mon Sep 17 00:00:00 2001
From: yym68686 <yym68686@outlook.com>
Date: Mon, 27 May 2024 20:48:33 +0800
Subject: [PATCH] Fix the bug where Azure is not compatible with function role

---
 main.py                          | 4 ++--
 setup.py                         | 2 +-
 src/ModelMerge/models/chatgpt.py | 8 ++++++++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/main.py b/main.py
index 713d876..f3ba9e6 100644
--- a/main.py
+++ b/main.py
@@ -28,10 +28,10 @@
 # for text in bot.ask_stream("英伟达最早支持杜比视界的显卡是哪一代"):
 # for text in bot.ask_stream("100个斐波纳切数列的和是多少"):
 # for text in bot.ask_stream("上海有哪些好玩的地方？"):
-for text in bot.ask_stream("https://arxiv.org/abs/2404.02041 这篇论文讲了啥？"):
+# for text in bot.ask_stream("https://arxiv.org/abs/2404.02041 这篇论文讲了啥？"):
 # for text in bot.ask_stream("今天伊朗总统目前的情况怎么样？"):
 # for text in bot.ask_stream("just say test"):
-# for text in bot.ask_stream("hi"):
+for text in bot.ask_stream("画一只猫猫"):
     print(text, end="")
 
 # print("\n bot tokens usage", bot.tokens_usage)
\ No newline at end of file
diff --git a/setup.py b/setup.py
index c4c2ad1..ff06a72 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@
 
 setup(
     name="modelmerge",
-    version="0.6.0",
+    version="0.6.1",
     description="modelmerge is a multi-large language model API aggregator.",
     long_description=Path.open(Path("README.md"), encoding="utf-8").read(),
     long_description_content_type="text/markdown",
diff --git a/src/ModelMerge/models/chatgpt.py b/src/ModelMerge/models/chatgpt.py
index e82cbc2..6dc1ab7 100644
--- a/src/ModelMerge/models/chatgpt.py
+++ b/src/ModelMerge/models/chatgpt.py
@@ -360,6 +360,14 @@ def ask_stream(
                                 "role": mess["role"],
                                 "content": mess["content"][0]["text"]
                             }
+                if "'function' is not an allowed role" in response.text:
+                    if json_post["messages"][-1]["role"] == "function":
+                        mess = json_post["messages"][-1]
+                        json_post["messages"][-1] = {
+                            "role": "assistant",
+                            "name": mess["name"],
+                            "content": mess["content"]
+                        }
                 else:
                     if "function_call" in json_post:
                         del json_post["function_call"]