-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
148 lines (131 loc) · 5.49 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import logging
from telegram import Update
from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler
import os
from pymongo import MongoClient
from gpt_utils import AnswerGenerator
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
global client
# crate a db entry with 10 starting tokens
n_start_tokens = 10
# connect to collection
col = client[os.environ["MONGO_INITDB_DATABASE"]][os.environ["db-collection"]]
# check if person exists or not
res = list(col.find({"username": update.message.from_user.username}))
# user already has info
if res:
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f"You already received your free tokens so no more free tokens for you lol\n\n"
f"Btw you can check your /balance to see how many tokens you have left\n\n"
f"Type /help for getting the format of question",
)
# no info about the user
else:
# create entry in collection
col.insert_one(
{"username": update.message.from_user.username, "tokens": n_start_tokens}
)
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f"Hello!\nYou have {n_start_tokens} tokens which you can use to ask me questions!",
)
async def balance(update: Update, context: ContextTypes.DEFAULT_TYPE):
global client
# connect to collection
col = client[os.environ["MONGO_INITDB_DATABASE"]][os.environ["db-collection"]]
# get info about user
res = col.find({"username": update.message.from_user.username})
tokens = list(res)[0]["tokens"]
await context.bot.send_message(
chat_id=update.effective_chat.id, text=f"You have {tokens} tokens left!"
)
async def question(update: Update, context: ContextTypes.DEFAULT_TYPE):
global client, answer_generator
# log the input of user
logging.log(level=logging.INFO, msg=" ".join(context.args))
# connect to collection
col = client[os.environ["MONGO_INITDB_DATABASE"]][os.environ["db-collection"]]
# get info about user
res = col.find({"username": update.message.from_user.username})
entry = list(res)[0]
# check if user can ask questions
if entry["tokens"] == 0:
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f"Sorry(\nYou ran out of tokens and you cannot ask questions anymore",
)
else:
# get the input from command
text = " ".join(context.args)
if text.endswith("?"):
# subtract one token and form response
col.update_one(
{"username": update.message.from_user.username},
{"$inc": {"tokens": -1}},
)
# get user the balance
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f"Your balance is now {entry['tokens']-1}",
)
# query gpt-3 for answer
await context.bot.send_message(
chat_id=update.effective_chat.id, text=answer_generator(text)
)
else:
# notify user about wrong input
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f"What you asked is not a question. Place write your query as an example:\n\n"
f"/question What hyperparameters does SVM have?",
)
async def help_function(update: Update, context: ContextTypes.DEFAULT_TYPE):
await context.bot.send_message(
chat_id=update.effective_chat.id,
text="To ask the question type /question <YOUR QUESTION>. Example:\n\n /question What hyperparameters does SVM have?"
)
def connect_init_db():
# connect to main bot db
client = MongoClient(
f"mongodb://{os.environ['MONGO_INITDB_ROOT_USERNAME']}:{os.environ['MONGO_INITDB_ROOT_PASSWORD']}@db",
27017,
)
db = client[os.environ["MONGO_INITDB_DATABASE"]]
# create collection if needed
if os.environ["db-collection"] not in db.list_collection_names():
col = db[os.environ["db-collection"]]
# make devs have many tokens
col.insert_many(
[
{"username": "matsusha", "tokens": 1488},
{"username": "Dmmc123", "tokens": 1337},
]
)
logging.log(
level=logging.INFO, msg=f"Created collection {os.environ['db-collection']}!"
)
return client
if __name__ == "__main__":
# connect to openai
answer_generator = AnswerGenerator(completion_model="text-davinci-002")
logging.log(level=logging.INFO, msg="Instantiated Answer Generator")
# connect to mongo
client = connect_init_db()
logging.log(level=logging.INFO, msg="Successfully connected to MongoDB!")
# initialize the bot app
application = ApplicationBuilder().token(os.environ["bot-api-token"]).build()
# make handlers
start_handler = CommandHandler("start", start)
balance_handler = CommandHandler("balance", balance)
question_handler = CommandHandler("question", question)
help_handler = CommandHandler('help', help_function)
application.add_handler(start_handler)
application.add_handler(balance_handler)
application.add_handler(question_handler)
application.add_handler(help_handler)
# start accepting requests
application.run_polling()