556 lines
18 KiB
Python
556 lines
18 KiB
Python
import os
|
|
import inspect
|
|
import logging
|
|
import signal
|
|
import random
|
|
import uuid
|
|
|
|
import openai
|
|
import asyncio
|
|
import markdown2
|
|
import tiktoken
|
|
import duckdb
|
|
|
|
from nio import AsyncClient, RoomMessageText, MatrixRoom, Event, InviteEvent, AsyncClientConfig, MegolmEvent, GroupEncryptionError, EncryptionError, HttpClient, Api
|
|
from nio.api import MessageDirection
|
|
from nio.responses import RoomMessagesError, SyncResponse, RoomRedactError, WhoamiResponse, JoinResponse, RoomSendResponse
|
|
from nio.crypto import Olm
|
|
|
|
from configparser import ConfigParser
|
|
from datetime import datetime
|
|
from argparse import ArgumentParser
|
|
from typing import List, Dict, Union, Optional
|
|
|
|
from commands import COMMANDS
|
|
from classes import DuckDBStore
|
|
|
|
|
|
def logging(message: str, log_level: str = "info"):
|
|
caller = inspect.currentframe().f_back.f_code.co_name
|
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S:%f")
|
|
print(f"[{timestamp}] - {caller} - [{log_level.upper()}] {message}")
|
|
|
|
|
|
CONTEXT = {
|
|
"database": False,
|
|
"default_room_name": "GPTBot",
|
|
"system_message": "You are a helpful assistant.",
|
|
"max_tokens": 3000,
|
|
"max_messages": 20,
|
|
"model": "gpt-3.5-turbo",
|
|
"client": None,
|
|
"sync_token": None,
|
|
"logger": logging
|
|
}
|
|
|
|
|
|
async def gpt_query(messages: list, model: Optional[str] = None):
|
|
model = model or CONTEXT["model"]
|
|
|
|
logging(f"Querying GPT with {len(messages)} messages")
|
|
try:
|
|
response = openai.ChatCompletion.create(
|
|
model=model,
|
|
messages=messages
|
|
)
|
|
result_text = response.choices[0].message['content']
|
|
tokens_used = response.usage["total_tokens"]
|
|
logging(f"Used {tokens_used} tokens")
|
|
return result_text, tokens_used
|
|
|
|
except Exception as e:
|
|
logging(f"Error during GPT API call: {e}", "error")
|
|
return None, 0
|
|
|
|
|
|
async def fetch_last_n_messages(room_id: str, n: Optional[int] = None,
|
|
client: Optional[AsyncClient] = None, sync_token: Optional[str] = None):
|
|
messages = []
|
|
|
|
n = n or CONTEXT["max_messages"]
|
|
client = client or CONTEXT["client"]
|
|
sync_token = sync_token or CONTEXT["sync_token"]
|
|
|
|
logging(
|
|
f"Fetching last {2*n} messages from room {room_id} (starting at {sync_token})...")
|
|
|
|
response = await client.room_messages(
|
|
room_id=room_id,
|
|
start=sync_token,
|
|
limit=2*n,
|
|
)
|
|
|
|
if isinstance(response, RoomMessagesError):
|
|
logging(
|
|
f"Error fetching messages: {response.message} (status code {response.status_code})", "error")
|
|
return []
|
|
|
|
for event in response.chunk:
|
|
if len(messages) >= n:
|
|
break
|
|
if isinstance(event, MegolmEvent):
|
|
try:
|
|
event = await client.decrypt_event(event)
|
|
except (GroupEncryptionError, EncryptionError):
|
|
logging(
|
|
f"Could not decrypt message {event.event_id} in room {room_id}", "error")
|
|
continue
|
|
if isinstance(event, RoomMessageText):
|
|
if event.body.startswith("!gptbot ignoreolder"):
|
|
break
|
|
if not event.body.startswith("!"):
|
|
messages.append(event)
|
|
|
|
logging(f"Found {len(messages)} messages (limit: {n})")
|
|
|
|
# Reverse the list so that messages are in chronological order
|
|
return messages[::-1]
|
|
|
|
|
|
def truncate_messages_to_fit_tokens(messages: list, max_tokens: Optional[int] = None,
|
|
model: Optional[str] = None, system_message: Optional[str] = None):
|
|
max_tokens = max_tokens or CONTEXT["max_tokens"]
|
|
model = model or CONTEXT["model"]
|
|
system_message = system_message or CONTEXT["system_message"]
|
|
|
|
encoding = tiktoken.encoding_for_model(model)
|
|
total_tokens = 0
|
|
|
|
system_message_tokens = len(encoding.encode(system_message)) + 1
|
|
|
|
if system_message_tokens > max_tokens:
|
|
logging(
|
|
f"System message is too long to fit within token limit ({system_message_tokens} tokens) - cannot proceed", "error")
|
|
return []
|
|
|
|
total_tokens += system_message_tokens
|
|
|
|
total_tokens = len(system_message) + 1
|
|
truncated_messages = []
|
|
|
|
for message in [messages[0]] + list(reversed(messages[1:])):
|
|
content = message["content"]
|
|
tokens = len(encoding.encode(content)) + 1
|
|
if total_tokens + tokens > max_tokens:
|
|
break
|
|
total_tokens += tokens
|
|
truncated_messages.append(message)
|
|
|
|
return [truncated_messages[0]] + list(reversed(truncated_messages[1:]))
|
|
|
|
|
|
async def process_query(room: MatrixRoom, event: RoomMessageText, **kwargs):
|
|
|
|
client = kwargs.get("client") or CONTEXT["client"]
|
|
database = kwargs.get("database") or CONTEXT["database"]
|
|
system_message = kwargs.get("system_message") or CONTEXT["system_message"]
|
|
max_tokens = kwargs.get("max_tokens") or CONTEXT["max_tokens"]
|
|
|
|
await client.room_typing(room.room_id, True)
|
|
|
|
await client.room_read_markers(room.room_id, event.event_id)
|
|
|
|
last_messages = await fetch_last_n_messages(room.room_id, 20)
|
|
|
|
chat_messages = [{"role": "system", "content": system_message}]
|
|
|
|
for message in last_messages:
|
|
role = "assistant" if message.sender == client.user_id else "user"
|
|
if not message.event_id == event.event_id:
|
|
chat_messages.append({"role": role, "content": message.body})
|
|
|
|
chat_messages.append({"role": "user", "content": event.body})
|
|
|
|
# Truncate messages to fit within the token limit
|
|
truncated_messages = truncate_messages_to_fit_tokens(
|
|
chat_messages, max_tokens - 1)
|
|
|
|
response, tokens_used = await gpt_query(truncated_messages)
|
|
|
|
if response:
|
|
logging(f"Sending response to room {room.room_id}...")
|
|
|
|
# Convert markdown to HTML
|
|
|
|
message = await send_message(room, response)
|
|
|
|
if database:
|
|
logging("Logging tokens used...")
|
|
|
|
with database.cursor() as cursor:
|
|
cursor.execute(
|
|
"INSERT INTO token_usage (message_id, room_id, tokens, timestamp) VALUES (?, ?, ?, ?)",
|
|
(message.event_id, room.room_id, tokens_used, datetime.now()))
|
|
database.commit()
|
|
else:
|
|
# Send a notice to the room if there was an error
|
|
|
|
logging("Error during GPT API call - sending notice to room")
|
|
send_message(
|
|
room, "Sorry, I'm having trouble connecting to the GPT API right now. Please try again later.", True)
|
|
print("No response from GPT API")
|
|
|
|
await client.room_typing(room.room_id, False)
|
|
|
|
|
|
async def process_command(room: MatrixRoom, event: RoomMessageText, context: Optional[dict] = None):
|
|
context = context or CONTEXT
|
|
|
|
logging(
|
|
f"Received command {event.body} from {event.sender} in room {room.room_id}")
|
|
command = event.body.split()[1] if event.body.split()[1:] else None
|
|
|
|
message = await COMMANDS.get(command, COMMANDS[None])(room, event, context)
|
|
|
|
if message:
|
|
room_id, event, content = message
|
|
await send_message(context["client"].rooms[room_id], content["body"],
|
|
True if content["msgtype"] == "m.notice" else False, context["client"])
|
|
|
|
|
|
async def message_callback(room: MatrixRoom, event: RoomMessageText | MegolmEvent, **kwargs):
|
|
context = kwargs.get("context") or CONTEXT
|
|
|
|
logging(f"Received message from {event.sender} in room {room.room_id}")
|
|
|
|
if isinstance(event, MegolmEvent):
|
|
try:
|
|
event = await context["client"].decrypt_event(event)
|
|
except Exception as e:
|
|
try:
|
|
logging("Requesting new encryption keys...")
|
|
await context["client"].request_room_key(event)
|
|
except:
|
|
pass
|
|
|
|
logging(f"Error decrypting message: {e}", "error")
|
|
await send_message(room, "Sorry, I couldn't decrypt that message. Please try again later or switch to a room without encryption.", True, context["client"])
|
|
return
|
|
|
|
if event.sender == context["client"].user_id:
|
|
logging("Message is from bot itself - ignoring")
|
|
|
|
elif event.body.startswith("!gptbot"):
|
|
await process_command(room, event)
|
|
|
|
elif event.body.startswith("!"):
|
|
logging("Might be a command, but not for this bot - ignoring")
|
|
|
|
else:
|
|
await process_query(room, event, context=context)
|
|
|
|
|
|
async def room_invite_callback(room: MatrixRoom, event: InviteEvent, **kwargs):
|
|
client: AsyncClient = kwargs.get("client") or CONTEXT["client"]
|
|
|
|
if room.room_id in client.rooms:
|
|
logging(f"Already in room {room.room_id} - ignoring invite")
|
|
return
|
|
|
|
logging(f"Received invite to room {room.room_id} - joining...")
|
|
|
|
response = await client.join(room.room_id)
|
|
if isinstance(response, JoinResponse):
|
|
await send_message(room, "Hello! I'm a helpful assistant. How can I help you today?", client)
|
|
else:
|
|
logging(f"Error joining room {room.room_id}: {response}", "error")
|
|
|
|
|
|
async def send_message(room: MatrixRoom, message: str, notice: bool = False, client: Optional[AsyncClient] = None):
|
|
client = client or CONTEXT["client"]
|
|
|
|
markdowner = markdown2.Markdown(extras=["fenced-code-blocks"])
|
|
formatted_body = markdowner.convert(message)
|
|
|
|
msgtype = "m.notice" if notice else "m.text"
|
|
|
|
msgcontent = {"msgtype": msgtype, "body": message,
|
|
"format": "org.matrix.custom.html", "formatted_body": formatted_body}
|
|
|
|
content = None
|
|
|
|
if client.olm and room.encrypted:
|
|
try:
|
|
if not room.members_synced:
|
|
responses = []
|
|
responses.append(await client.joined_members(room.room_id))
|
|
|
|
if client.olm.should_share_group_session(room.room_id):
|
|
try:
|
|
event = client.sharing_session[room.room_id]
|
|
await event.wait()
|
|
except KeyError:
|
|
await client.share_group_session(
|
|
room.room_id,
|
|
ignore_unverified_devices=True,
|
|
)
|
|
|
|
if msgtype != "m.reaction":
|
|
response = client.encrypt(room.room_id, "m.room.message", msgcontent)
|
|
msgtype, content = response
|
|
|
|
except Exception as e:
|
|
logging(
|
|
f"Error encrypting message: {e} - sending unencrypted", "error")
|
|
raise
|
|
|
|
if not content:
|
|
msgtype = "m.room.message"
|
|
content = msgcontent
|
|
|
|
method, path, data = Api.room_send(
|
|
client.access_token, room.room_id, msgtype, content, uuid.uuid4()
|
|
)
|
|
|
|
return await client._send(RoomSendResponse, method, path, data, (room.room_id,))
|
|
|
|
|
|
async def accept_pending_invites(client: Optional[AsyncClient] = None):
|
|
client = client or CONTEXT["client"]
|
|
|
|
logging("Accepting pending invites...")
|
|
|
|
for room_id in list(client.invited_rooms.keys()):
|
|
logging(f"Joining room {room_id}...")
|
|
|
|
response = await client.join(room_id)
|
|
|
|
if isinstance(response, JoinResponse):
|
|
logging(response, "debug")
|
|
rooms = await client.joined_rooms()
|
|
await send_message(rooms[room_id], "Hello! I'm a helpful assistant. How can I help you today?", client)
|
|
else:
|
|
logging(f"Error joining room {room_id}: {response}", "error")
|
|
|
|
|
|
async def sync_cb(response, write_global: bool = True):
|
|
logging(
|
|
f"Sync response received (next batch: {response.next_batch})", "debug")
|
|
SYNC_TOKEN = response.next_batch
|
|
|
|
if write_global:
|
|
global CONTEXT
|
|
CONTEXT["sync_token"] = SYNC_TOKEN
|
|
|
|
|
|
async def test_callback(room: MatrixRoom, event: Event, **kwargs):
|
|
logging(
|
|
f"Received event {event.__class__.__name__} in room {room.room_id}", "debug")
|
|
|
|
|
|
async def init(config: ConfigParser):
|
|
# Set up Matrix client
|
|
try:
|
|
assert "Matrix" in config
|
|
assert "Homeserver" in config["Matrix"]
|
|
assert "AccessToken" in config["Matrix"]
|
|
except:
|
|
logging("Matrix config not found or incomplete", "critical")
|
|
exit(1)
|
|
|
|
homeserver = config["Matrix"]["Homeserver"]
|
|
access_token = config["Matrix"]["AccessToken"]
|
|
|
|
device_id, user_id = await get_device_id(access_token, homeserver)
|
|
|
|
device_id = config["Matrix"].get("DeviceID", device_id)
|
|
user_id = config["Matrix"].get("UserID", user_id)
|
|
|
|
# Set up database
|
|
if "Database" in config and config["Database"].get("Path"):
|
|
database = CONTEXT["database"] = initialize_database(
|
|
config["Database"]["Path"])
|
|
matrix_store = DuckDBStore
|
|
|
|
client_config = AsyncClientConfig(
|
|
store_sync_tokens=True, encryption_enabled=True, store=matrix_store)
|
|
|
|
else:
|
|
client_config = AsyncClientConfig(
|
|
store_sync_tokens=True, encryption_enabled=False)
|
|
|
|
client = AsyncClient(
|
|
config["Matrix"]["Homeserver"], config=client_config)
|
|
|
|
if client.config.encryption_enabled:
|
|
client.store = client.config.store(
|
|
user_id,
|
|
device_id,
|
|
database
|
|
)
|
|
assert client.store
|
|
|
|
client.olm = Olm(client.user_id, client.device_id, client.store)
|
|
client.encrypted_rooms = client.store.load_encrypted_rooms()
|
|
|
|
CONTEXT["client"] = client
|
|
|
|
CONTEXT["client"].access_token = config["Matrix"]["AccessToken"]
|
|
CONTEXT["client"].user_id = user_id
|
|
CONTEXT["client"].device_id = device_id
|
|
|
|
# Set up GPT API
|
|
try:
|
|
assert "OpenAI" in config
|
|
assert "APIKey" in config["OpenAI"]
|
|
except:
|
|
logging("OpenAI config not found or incomplete", "critical")
|
|
exit(1)
|
|
|
|
openai.api_key = config["OpenAI"]["APIKey"]
|
|
|
|
if "Model" in config["OpenAI"]:
|
|
CONTEXT["model"] = config["OpenAI"]["Model"]
|
|
|
|
if "MaxTokens" in config["OpenAI"]:
|
|
CONTEXT["max_tokens"] = int(config["OpenAI"]["MaxTokens"])
|
|
|
|
if "MaxMessages" in config["OpenAI"]:
|
|
CONTEXT["max_messages"] = int(config["OpenAI"]["MaxMessages"])
|
|
|
|
# Listen for SIGTERM
|
|
|
|
def sigterm_handler(_signo, _stack_frame):
|
|
logging("Received SIGTERM - exiting...")
|
|
exit()
|
|
|
|
signal.signal(signal.SIGTERM, sigterm_handler)
|
|
|
|
|
|
async def main(config: Optional[ConfigParser] = None, client: Optional[AsyncClient] = None):
|
|
if not client and not CONTEXT.get("client"):
|
|
await init(config)
|
|
|
|
client = client or CONTEXT["client"]
|
|
|
|
try:
|
|
assert client.user_id
|
|
except AssertionError:
|
|
logging(
|
|
"Failed to get user ID - check your access token or try setting it manually", "critical")
|
|
await client.close()
|
|
return
|
|
|
|
logging("Starting bot...")
|
|
|
|
client.add_response_callback(sync_cb, SyncResponse)
|
|
|
|
logging("Syncing...")
|
|
|
|
await client.sync(timeout=30000)
|
|
|
|
client.add_event_callback(message_callback, RoomMessageText)
|
|
client.add_event_callback(message_callback, MegolmEvent)
|
|
client.add_event_callback(room_invite_callback, InviteEvent)
|
|
client.add_event_callback(test_callback, Event)
|
|
|
|
await accept_pending_invites() # Accept pending invites
|
|
|
|
logging("Bot started")
|
|
|
|
try:
|
|
# Continue syncing events
|
|
await client.sync_forever(timeout=30000)
|
|
finally:
|
|
logging("Syncing one last time...")
|
|
await client.sync(timeout=30000)
|
|
await client.close() # Properly close the aiohttp client session
|
|
logging("Bot stopped")
|
|
|
|
|
|
def initialize_database(path: os.PathLike):
|
|
logging("Initializing database...")
|
|
database = duckdb.connect(path)
|
|
|
|
with database.cursor() as cursor:
|
|
# Get the latest migration ID if the migrations table exists
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
SELECT MAX(id) FROM migrations
|
|
"""
|
|
)
|
|
|
|
latest_migration = int(cursor.fetchone()[0])
|
|
except:
|
|
latest_migration = 0
|
|
|
|
# Version 1
|
|
|
|
if latest_migration < 1:
|
|
cursor.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS token_usage (
|
|
message_id TEXT PRIMARY KEY,
|
|
room_id TEXT NOT NULL,
|
|
tokens INTEGER NOT NULL,
|
|
timestamp TIMESTAMP NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
|
|
cursor.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS migrations (
|
|
id INTEGER NOT NULL,
|
|
timestamp TIMESTAMP NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
|
|
cursor.execute(
|
|
"INSERT INTO migrations (id, timestamp) VALUES (1, ?)",
|
|
(datetime.now(),)
|
|
)
|
|
|
|
database.commit()
|
|
|
|
return database
|
|
|
|
|
|
async def get_device_id(access_token, homeserver):
|
|
client = AsyncClient(homeserver)
|
|
client.access_token = access_token
|
|
|
|
logging(f"Obtaining device ID for access token {access_token}...", "debug")
|
|
response = await client.whoami()
|
|
if isinstance(response, WhoamiResponse):
|
|
logging(
|
|
f"Authenticated as {response.user_id}.")
|
|
user_id = response.user_id
|
|
devices = await client.devices()
|
|
device_id = devices.devices[0].id
|
|
|
|
await client.close()
|
|
|
|
return device_id, user_id
|
|
|
|
else:
|
|
logging(f"Failed to obtain device ID: {response}", "error")
|
|
|
|
await client.close()
|
|
|
|
return None, None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Parse command line arguments
|
|
parser = ArgumentParser()
|
|
parser.add_argument(
|
|
"--config", help="Path to config file (default: config.ini in working directory)", default="config.ini")
|
|
args = parser.parse_args()
|
|
|
|
# Read config file
|
|
config = ConfigParser()
|
|
config.read(args.config)
|
|
|
|
# Start bot loop
|
|
try:
|
|
asyncio.run(main(config))
|
|
except KeyboardInterrupt:
|
|
logging("Received KeyboardInterrupt - exiting...")
|
|
except SystemExit:
|
|
logging("Received SIGTERM - exiting...")
|
|
finally:
|
|
if CONTEXT["database"]:
|
|
CONTEXT["database"].close()
|