Reduce amount of information returned

[tall.bot.git] / tallbot.py

diff --git a/tallbot.py b/tallbot.py
index e097ba78234fbc75146932e3f8af73d2ad11cc98..fba1c1639eb445f2ec52de26a9f36e087525e174 100755 (executable)
--- a/tallbot.py
+++ b/tallbot.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 # tall.bot - a simple Discord Wikipedia bot
-# Armaan Bhojwani 2020
+# Armaan Bhojwani 2021
 
 import logging
+import sqlite3
 import sys
 import time
 import wikipedia
@@ -10,16 +11,18 @@ import discord
 
 def wiki_sum(search, sen, pid):
     try:
-        logging.info(str(pid) + " | " + str(int(time.time())) + " | " + search)
+        logging.info(str(pid) + " | " + str(int(time.time())) + " | " + str(search))
+        log_query(search, pid)
         return "<Wikipedia> " + wikipedia.summary(search,
                                                   sentences=sen,
                                                   auto_suggest=False)
     except Exception as e:
-        logging.debug(str(pid) + " | " + search + " | " + str(e))
+        logging.debug(str(pid) + " | " + str(search) + " | " + str(e))
         return str(e)
 
 def create_query(inp):
-    key1 = ["What", "what", "who", "Who"]
+    key1 = ["What", "what", "who", "Who", "what's", "What's", "whats", "Whats",
+            "Whatsa", "whatsa"]
     key2 = ["is", "are", "were", "was"]
     key3 = ["a", "an", "the"]
 
@@ -33,6 +36,41 @@ def create_query(inp):
 
     return " ".join(x)
 
+def tell_more(inp, pid):
+    def find_query(inp):
+        key1 = ["Tell me more", "tell me more"]
+        if inp[:12] in key1:
+            return inp[11:]
+        else:
+            return False
+
+    def query_lookup(pid):
+        old = int(time.time()) - 150
+        c.execute("SELECT query FROM queries WHERE pid = (?) AND time > (?)",
+                  (str(pid), old))
+
+        (query, *_) = c.fetchone()
+        return str(query)
+
+    if find_query(inp):
+        return query_lookup(pid)
+    else:
+        return False
+
+def prep_database():
+    c.execute("""CREATE TABLE IF NOT EXISTS queries(
+              time INTEGER,
+              pid TEXT,
+              query TEXT)""")
+    conn.commit()
+
+def log_query(search, pid):
+    c.execute("DELETE FROM queries WHERE pid = (?)", (str(pid),))
+    c.execute("""INSERT INTO queries
+              (time, pid, query) VALUES(?, ?, ?)""",
+              (int(time.time()), str(pid), str(search)))
+    conn.commit()
+
 class TallBot(discord.Client):
     async def on_ready(self):
         logging.info('Logged in as ', self.user)
@@ -41,12 +79,18 @@ class TallBot(discord.Client):
         if message.author == self.user:
             return
 
-        key = create_query(message.content)
-        if key:
-            await message.channel.send(wiki_sum(key, 2, message.author))
+        query = create_query(message.content)
+        more_info = tell_more(message.content, message.author)
+        if query:
+            await message.channel.send(wiki_sum(query, 1, message.author))
+        elif more_info:
+            await message.channel.send(wiki_sum(more_info, 6, message.author))
 
 def main():
+    prep_database()
     TallBot().run(sys.argv[1])
 
 if __name__ == "__main__":
+    conn = sqlite3.connect("queries.db")
+    c = conn.cursor()
     main()