diff --git a/main.py b/main.py new file mode 100644 index 0000000..70d26b7 --- /dev/null +++ b/main.py @@ -0,0 +1,70 @@ +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)