tallbot.py

   1 #!/usr/bin/env python3
   2 # tall.bot - a simple Discord Wikipedia bot
   3 # Armaan Bhojwani 2020
   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         return c.fetchone()
  52
  53     if find_query(inp):
  54         return query_lookup(pid)
  55     else:
  56         return False
  57
  58 def prep_database():
  59     c.execute("""CREATE TABLE IF NOT EXISTS queries(
  60               time INTEGER,
  61               pid TEXT,
  62               query TEXT)""")
  63     conn.commit()
  64
  65 def log_query(search, pid):
  66     c.execute("DELETE FROM queries WHERE pid = (?)", (str(pid),))
  67     c.execute("""INSERT INTO queries
  68               (time, pid, query) VALUES(?, ?, ?)""",
  69               (int(time.time()), str(pid), str(search)))
  70     conn.commit()
  71
  72 class TallBot(discord.Client):
  73     async def on_ready(self):
  74         logging.info('Logged in as ', self.user)
  75
  76     async def on_message(self, message):
  77         if message.author == self.user:
  78             return
  79
  80         query = create_query(message.content)
  81         more_info = tell_more(message.content, message.author)
  82         if query:
  83             await message.channel.send(wiki_sum(query, 2, message.author))
  84         elif more_info:
  85             await message.channel.send(wiki_sum(more_info, 8, message.author))
  86
  87 def main():
  88     prep_database()
  89     TallBot().run(sys.argv[1])
  90
  91 if __name__ == "__main__":
  92     conn = sqlite3.connect("queries.db")
  93     c = conn.cursor()
  94     main()