forust
0506aaaac8
- ruff check --fix: 210 auto-fixed errors (import sorting, trailing whitespace, unused imports, f-string fixups, deprecated annotations) - ruff format: 104 files reformatted to consistent style - 268 non-auto-fixable issues remain (S113 requests timeout, etc.)
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import asyncio
|
|
|
|
from utils.config import cohere_key
|
|
from utils.scripts import import_library
|
|
|
|
cohere = import_library('cohere')
|
|
|
|
import cohere
|
|
|
|
co = cohere.Client(cohere_key)
|
|
|
|
from pyrogram import Client, enums, filters
|
|
from pyrogram.errors import MessageTooLong
|
|
from pyrogram.types import Message
|
|
from utils.db import db
|
|
from utils.misc import modules_help, prefix
|
|
from utils.rentry import paste as rentry_paste
|
|
from utils.scripts import format_exc
|
|
|
|
|
|
@Client.on_message(filters.command('cohere', prefix) & filters.me)
|
|
async def cohere(c: Client, message: Message):
|
|
try:
|
|
user_id = message.from_user.id
|
|
chat_history = db.get_chat_history(user_id)
|
|
|
|
if len(message.command) > 1:
|
|
prompt = message.text.split(maxsplit=1)[1]
|
|
elif message.reply_to_message:
|
|
prompt = message.reply_to_message.text
|
|
else:
|
|
await message.edit_text(f'<b>Usage: </b><code>{prefix}cohere [prompt/reply to message]</code>')
|
|
return
|
|
|
|
db.add_chat_history(user_id, {'role': 'USER', 'message': prompt})
|
|
|
|
await message.edit_text('<code>Umm, lemme think...</code>')
|
|
|
|
response = co.chat_stream(
|
|
chat_history=chat_history,
|
|
model='command-r-plus',
|
|
message=prompt,
|
|
temperature=0.8,
|
|
tools=[{'name': 'internet_search'}],
|
|
connectors=[],
|
|
prompt_truncation='OFF',
|
|
)
|
|
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"
|
|
|
|
db.add_chat_history(user_id, {'role': 'CHATBOT', 'message': output})
|
|
|
|
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}',
|
|
parse_mode=enums.ParseMode.MARKDOWN,
|
|
disable_web_page_preview=True,
|
|
)
|
|
|
|
except IndexError:
|
|
references = ''
|
|
await message.edit_text(
|
|
f'**Question:**`{prompt}`\n**Answer:** {output}\n',
|
|
parse_mode=enums.ParseMode.MARKDOWN,
|
|
disable_web_page_preview=True,
|
|
)
|
|
except MessageTooLong:
|
|
await message.edit_text('<code>Output is too long... Pasting to rentry...</code>')
|
|
try:
|
|
output = output + '\n\n' + references if references else output
|
|
rentry_url, edit_code = await rentry_paste(text=output, return_edit=True)
|
|
except RuntimeError:
|
|
await message.edit_text('<b>Error:</b> <code>Failed to paste to rentry</code>')
|
|
return
|
|
await c.send_message(
|
|
'me',
|
|
f"Here's your edit code for Url: {rentry_url}\nEdit code: <code>{edit_code}</code>",
|
|
disable_web_page_preview=True,
|
|
)
|
|
await message.edit_text(
|
|
f'<b>Output:</b> {rentry_url}\n<b>Note:</b> <code>Edit Code has been sent to your saved messages</code>',
|
|
disable_web_page_preview=True,
|
|
)
|
|
|
|
except Exception as e:
|
|
await message.edit_text(f'An error occurred: {format_exc(e)}')
|
|
|
|
|
|
modules_help['cohere'] = {
|
|
'cohere': 'Chat with cohere ai' + '\nSupports Chat History\n' + 'Supports real time internet search'
|
|
}
|