tallbot.py

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