65 lines
1.6 KiB
Python
Executable File
65 lines
1.6 KiB
Python
Executable File
#!/bin/python3.10
|
|
import sys
|
|
import os
|
|
import random
|
|
import re
|
|
|
|
import discord
|
|
from ollama import AsyncClient
|
|
from discord.ext import commands
|
|
from dotenv import load_dotenv
|
|
|
|
description = """
|
|
An example bot to showcase the discord.ext.commands extension module.
|
|
There are a number of utility commands being showcased here.
|
|
"""
|
|
|
|
load_dotenv()
|
|
client = AsyncClient()
|
|
|
|
bot = commands.Bot(
|
|
command_prefix=commands.when_mentioned_or("!"),
|
|
description=description,
|
|
intents=discord.Intents.all(),
|
|
)
|
|
|
|
@bot.event
|
|
async def on_ready():
|
|
print(f"Logged in as {bot.user} (ID: {bot.user.id})")
|
|
|
|
def extract_response(input_string):
|
|
# Find the position of the closing tag
|
|
end_index = input_string.find('</think>')
|
|
|
|
if end_index == -1:
|
|
raise ValueError("Closing tag </think> not found")
|
|
|
|
# Extract the response part of the string
|
|
response = input_string[end_index + len('</think>'):].strip()
|
|
|
|
return response
|
|
|
|
@bot.event
|
|
async def on_message(message: discord.Message):
|
|
msg = ""
|
|
# Make sure we won't be replying to ourselves.
|
|
if message.author.id == bot.user.id:
|
|
return
|
|
|
|
if f"""<@{bot.user.id}>""" in message.content:
|
|
async with message.channel.typing():
|
|
messages = [
|
|
{
|
|
'role':'user',
|
|
'content': message.content
|
|
}
|
|
]
|
|
|
|
response = await client.chat('marx', messages=messages)
|
|
print(response['message']['content'])
|
|
msg = extract_response(response['message']['content'])
|
|
|
|
await message.channel.send(msg)
|
|
|
|
bot.run(os.environ.get("DISCORD_TOKEN"))
|