71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
import discord
|
|
import requests
|
|
import json
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# Replace with your bot token
|
|
BOT_TOKEN = os.getenv('DISCORD_TOKEN')
|
|
|
|
# Ollama configuration
|
|
OLLAMA_API_URL = 'http://192.168.1.159:11434/api/generate' # Adjust if your Ollama setup is different
|
|
|
|
# Set up the Discord client
|
|
intents = discord.Intents.default()
|
|
intents.messages = True
|
|
intents.message_content = True
|
|
|
|
client = discord.Client(intents=intents)
|
|
|
|
|
|
# Function to query Ollama
|
|
def query_ollama(prompt):
|
|
payload = {
|
|
"prompt": prompt,
|
|
"model": "nollama/mythomax-l2-13b:Q4_K_M" # Replace with your Ollama model
|
|
}
|
|
try:
|
|
response = requests.post(OLLAMA_API_URL, json=payload, stream=True)
|
|
if response.status_code == 200:
|
|
collected_response = ""
|
|
# Stream and parse each line of JSON from the response
|
|
for line in response.iter_lines(decode_unicode=True):
|
|
if line.strip(): # Skip empty lines
|
|
try:
|
|
data = json.loads(line) # Parse each line as JSON
|
|
collected_response += data.get("response", "")
|
|
if data.get("done", False):
|
|
break
|
|
except json.JSONDecodeError as e:
|
|
print(f"Error decoding JSON line: {line}, Error: {e}")
|
|
return collected_response.strip() or "No response from model."
|
|
else:
|
|
return f"Error: {response.status_code} - {response.text}"
|
|
except requests.RequestException as e:
|
|
return f"Error connecting to Ollama: {str(e)}"
|
|
|
|
|
|
# Event for when the bot is ready
|
|
@client.event
|
|
async def on_ready():
|
|
print(f'We have logged in as {client.user}')
|
|
|
|
|
|
# Event for when a message is sent
|
|
@client.event
|
|
async def on_message(message):
|
|
# Ignore the bot's own messages
|
|
if message.author == client.user:
|
|
return
|
|
|
|
# Respond to all messages except those in DMs
|
|
if not isinstance(message.channel, discord.DMChannel):
|
|
response = query_ollama(message.content.strip())
|
|
await message.channel.send(response)
|
|
|
|
# Run the bot
|
|
client.run(BOT_TOKEN)
|