From f84efa35cb2cc213828677fa904fe41b06dcba81 Mon Sep 17 00:00:00 2001
From: hibobmaster <32976627+hibobmaster@users.noreply.github.com>
Date: Fri, 5 Jan 2024 21:27:44 +0800
Subject: [PATCH] Make gptbot more compatible

---
 CHANGELOG.md  |  3 +++
 src/bot.py    |  2 +-
 src/gptbot.py | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d7a3f1..48343ae 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.3.2
+- Make gptbot more compatible
+
 ## v1.3.1
 - Expose more stable diffusion webui api parameters
 
diff --git a/src/bot.py b/src/bot.py
index 5d1a2ac..d618457 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -255,7 +255,7 @@ async def message_callback(
                                 "channel_id": channel_id,
                             },
                         )
-                        response = await self.chatbot.ask_async(
+                        response = await self.chatbot.ask_async_v2(
                             prompt=prompt, convo_id=user_id
                         )
                         await self.send_message(channel_id, f"{response}", root_id)
diff --git a/src/gptbot.py b/src/gptbot.py
index 31e9c72..c9cfed4 100644
--- a/src/gptbot.py
+++ b/src/gptbot.py
@@ -250,6 +250,54 @@ async def ask_async(
         full_response: str = "".join([r async for r in response])
         return full_response
 
+    async def ask_async_v2(
+        self,
+        prompt: str,
+        role: str = "user",
+        convo_id: str = "default",
+        model: str = None,
+        pass_history: bool = True,
+        **kwargs,
+    ) -> str:
+        # Make conversation if it doesn't exist
+        if convo_id not in self.conversation:
+            self.reset(convo_id=convo_id, system_prompt=self.system_prompt)
+        self.add_to_conversation(prompt, "user", convo_id=convo_id)
+        self.__truncate_conversation(convo_id=convo_id)
+        # Get response
+        response = await self.aclient.post(
+            url=self.api_url,
+            headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
+            json={
+                "model": model or self.engine,
+                "messages": self.conversation[convo_id] if pass_history else [prompt],
+                # kwargs
+                "temperature": kwargs.get("temperature", self.temperature),
+                "top_p": kwargs.get("top_p", self.top_p),
+                "presence_penalty": kwargs.get(
+                    "presence_penalty",
+                    self.presence_penalty,
+                ),
+                "frequency_penalty": kwargs.get(
+                    "frequency_penalty",
+                    self.frequency_penalty,
+                ),
+                "n": kwargs.get("n", self.reply_count),
+                "user": role,
+                "max_tokens": min(
+                    self.get_max_tokens(convo_id=convo_id),
+                    kwargs.get("max_tokens", self.max_tokens),
+                ),
+            },
+            timeout=kwargs.get("timeout", self.timeout),
+        )
+        resp = response.json()
+        full_response = resp["choices"][0]["message"]["content"]
+        self.add_to_conversation(
+            full_response, resp["choices"][0]["message"]["role"], convo_id=convo_id
+        )
+        return full_response
+
     def reset(self, convo_id: str = "default", system_prompt: str = None) -> None:
         """
         Reset the conversation