tallbot.py

   1 #!/usr/bin/env python3
   2 # tall.bot - a simple Discord Wikipedia bot
   3 # Armaan Bhojwani 2021
   4
   5 import logging
   6 import sqlite3
   7 import sys
   8 import time
   9 import wikipedia
  10 import discord
  11
  12 def wiki_sum(search, sen, pid):
  13     try:
  14         logging.info(str(pid) + " | " + str(int(time.time())) + " | " + str(search))
  15         log_query(search, pid)
  16         return "<Wikipedia> " + wikipedia.summary(search,
  17                                                   sentences=sen,
  18                                                   auto_suggest=False)
  19     except Exception as e:
  20         logging.debug(str(pid) + " | " + str(search) + " | " + str(e))
  21         return str(e)
  22
  23 def create_query(inp):
  24     key1 = ["What", "what", "who", "Who", "what's", "What's", "whats", "Whats",
  25             "Whatsa", "whatsa"]
  26     key2 = ["is", "are", "were", "was"]
  27     key3 = ["a", "an", "the"]
  28
  29     x = inp.split()
  30     if x[0] in key1 and x[1] in key2:
  31         x = x[2:]
  32         if x[0] in key3:
  33             x = x[1:]
  34     else:
  35         return False
  36
  37     return " ".join(x)
  38
  39 def tell_more(inp, pid):
  40     def find_query(inp):
  41         key1 = ["Tell me more", "tell me more"]
  42         if inp[:12] in key1:
  43             return inp[11:]
  44         else:
  45             return False
  46
  47     def query_lookup(pid):
  48         old = int(time.time()) - 150
  49         c.execute("SELECT query FROM queries WHERE pid = (?) AND time > (?)",
  50                   (str(pid), old))
  51
  52         (query, *_) = c.fetchone()
  53         return str(query)
  54
  55     if find_query(inp):
  56         return query_lookup(pid)
  57     else:
  58         return False
  59
  60 def prep_database():
  61     c.execute("""CREATE TABLE IF NOT EXISTS queries(
  62               time INTEGER,
  63               pid TEXT,
  64               query TEXT)""")
  65     conn.commit()
  66
  67 def log_query(search, pid):
  68     c.execute("DELETE FROM queries WHERE pid = (?)", (str(pid),))
  69     c.execute("""INSERT INTO queries
  70               (time, pid, query) VALUES(?, ?, ?)""",
  71               (int(time.time()), str(pid), str(search)))
  72     conn.commit()
  73
  74 class TallBot(discord.Client):
  75     async def on_ready(self):
  76         logging.info('Logged in as ', self.user)
  77
  78     async def on_message(self, message):
  79         if message.author == self.user:
  80             return
  81
  82         query = create_query(message.content)
  83         more_info = tell_more(message.content, message.author)
  84         if query:
  85             await message.channel.send(wiki_sum(query, 2, message.author))
  86         elif more_info:
  87             await message.channel.send(wiki_sum(more_info, 8, message.author))
  88
  89 def main():
  90     prep_database()
  91     TallBot().run(sys.argv[1])
  92
  93 if __name__ == "__main__":
  94     conn = sqlite3.connect("queries.db")
  95     c = conn.cursor()
  96     main()