Update Akeno/plugins/cohere.py

Akeno/plugins/cohere.py

@@ -5,13 +5,13 @@ from Akeno.utils.database import db
 
 from pyrogram import Client, filters, enums
 from pyrogram.types import Message
-from cohere import Client as Client_cohere
 from config import cohere_key, CMD_HANDLER
+import cohere
 
-co_pro = Client_cohere(cohere_key)
+co = cohere.Client(api_key=cohere_key)
 
 @Akeno(filters.command("cohere", CMD_HANDLER) & filters.me)
-async def cohere(c: Client, message: Message):
+async def coheres_(c: Client, message: Message):
     try:
         user_id = message.from_user.id
         chat_history = await db._get_cohere_chat_from_db(user_id)
@@ -26,60 +26,14 @@ async def cohere(c: Client, message: Message):
             return
         chat_history.append({"role": "USER", "message": prompt})
         await message.edit_text("<code>Umm, lemme think...</code>")
-        response = co_pro.chat_stream(
+        response = co.chat(
             chat_history=chat_history,
             model="command-r-plus",
-            message=prompt,
-            temperature=0.3,
-            tools=[{"name": "internet_search"}],
-            connectors=[],
-            prompt_truncation="OFF",
+            message=prompt
         )
-        output = ""
-        tool_message = ""
-        data = []
-        for event in response:
-            if event.event_type == "tool-calls-chunk":
-                if event.tool_call_delta and event.tool_call_delta.text is None:
-                    tool_message += ""
-                else:
-                    tool_message += event.text
-            if event.event_type == "search-results":
-                data.append(event.documents)
-            if event.event_type == "text-generation":
-                output += event.text
-        if output == "":
-            output = "I can't seem to find an answer to that"
+        output = response.text
         chat_history.append({"role": "CHATBOT", "message": output})
         await db._update_cohere_chat_in_db(user_id, chat_history)
-        await message.edit_text(f"<code>{tool_message}</code>")
-        await asyncio.sleep(5)
-        try:
-            data = data[0]
-            references = ""
-            reference_dict = {}
-            for item in data:
-                title = item["title"]
-                url = item["url"]
-                if title not in reference_dict:
-                    reference_dict[title] = url
-
-            i = 1
-            for title, url in reference_dict.items():
-                references += f"**{i}.** [{title}]({url})\n"
-                i += 1
-
-            await message.edit_text(
-                f"**Question:**`{prompt}`\n**Answer:** {output}\n\n**References:**\n{references}",
-                disable_web_page_preview=True
-            )
-
-        except IndexError:
-            references = ""
-            await message.edit_text(
-                f"**Question:**`{prompt}`\n**Answer:** {output}\n",
-                disable_web_page_preview=True
-            )
-
+        await message.edit_text(output, disable_web_page_preview=True)
     except Exception as e:
         await message.edit_text(f"An error occurred: {e}")