added chatgpt functionality to get category and language-specific pic…

…k up lines + set up a dummy main menu + added gpt generated pick-up lines to database as a separate collection

requirements.txt

@@ -1,2 +1,3 @@
 pymongo==4.6.0
-python-dotenv==0.16.0
+python-dotenv==0.16.0
+openai==0.28.1

src/pyrizz/__main__.py

@@ -0,0 +1,49 @@
+import pyrizz
+
+"""Main function for PyRizz."""
+
+#ASCII art
+
+print(" __    " + "    ____  _            ")
+print("|__)   " + "   / __ \(_)_______    ")
+print("|   \/ " + "  / /_/ / /_  /_  /    ")
+print("    /  " + " / _, _/ / / /_/ /_    ")
+print("       " + "/_/ |_/_/ /___/___/    ")
+
+print("\n\n")
+
+def main():
+    print("Welcome to PyRizz! Your journey to getting a date begins here...\n")
+
+    while True:
+
+        print("What would you like to do today?\n")
+        print("1. Get a random pick-up line hand-picked by the devs with a guaranteed 100% success rate.")
+        print("2. Get a category-specific random pick-up line hand-picked by the devs with a guaranteed 100% success rate.")
+        print("3. Have AI generate a pick-up line in your chosen category / language with a 50% success rate.")
+        print("4. Have AI rate your pick-up line out of 10. Test it on AI before trying it on a human! ;)")
+        print("5. Insert your own pick-up line to our database.\n")
+        print("!! Type Q to quit !!\n")
+
+        print("Enter your choice: ")
+        user_input = input("> ")
+
+        if user_input == "3":
+            print("Enter a category / language: ")
+            category = input("> ")
+            print("\n" + pyrizz.get_ai_line(category), end = "\n\n")
+
+        elif user_input == "q" or user_input == "Q":
+            break
+
+        print("Would you like to do something else? (y/n)")
+        user_cont = input("> ")
+
+        if user_cont == "n":
+            break
+
+        if user_cont == "y":
+            continue
+
+if __name__ == "__main__":
+    main()

src/pyrizz/pyrizz.py

@@ -2,16 +2,18 @@
 import os
 from dotenv import load_dotenv
 from random import randint
+import openai
 
 load_dotenv()
 
 client = MongoClient(os.getenv('MONGO_URI'))
+openai.api_key = os.getenv('OPENAI_API_KEY')
 
 # Checks if the connection has been made, else make an error printout
 try:
     client.admin.command('ping')          
     database = client[os.getenv('MONGO_DBNAME')]          
-    print('* Connected to MongoDB!')         
+    print('*****')         
 
 except Exception as err:
     print('* "Failed to connect to MongoDB at', os.getenv('MONGO_URI'))
@@ -25,4 +27,26 @@ def get_random_line() -> None:
 
     line = collection.find_one({})["lines"][random_number]
 
-    print(line)
+    print(line)
+
+def get_ai_line(category) -> str:
+    response = openai.ChatCompletion.create(
+        model = os.getenv('OPENAI_MODEL'),
+        messages =
+            [{"role": "user", "content": f"I need a {category} pick-up line."},]
+    )
+
+    message = response.choices[0]['message']
+    ai_line = "{}".format(message['content'])
+
+    collection = database['ai_generated']
+
+    lines = collection.find_one({})["lines"]
+
+    lines.append(ai_line)
+
+    collection.insert_one({
+        'lines': lines
+    })
+
+    return ai_line

src/scripts/extract.py

@@ -16,5 +16,4 @@
     d['lines'].append(line)
 
 with open(f'{PROJECT_ROOT}/src/data/lines.json', 'w') as file:
-    json.dump(d, file)
-
+    json.dump(d, file)

src/scripts/mongo-seed.py

@@ -27,10 +27,15 @@
     data = json.load(file)
     lines = data['lines']
 
-collection = database['lines']
+collection_dev_lines = database['dev_lines']
+collection_ai_generated = database['ai_generated']
 
-collection.insert_one({
+collection_dev_lines.insert_one({
     'lines': lines
 })
 
+collection_ai_generated.insert_one({
+    'lines': []
+})
+
 client.close()