Skip to content

Commit

Permalink
Merge pull request TheExplainthis#21 from TheExplainthis/develop
Browse files Browse the repository at this point in the history
Add MongoDB
  • Loading branch information
TheExplainthis authored Mar 23, 2023
2 parents f3f7467 + 4300f76 commit a316a11
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 7 deletions.
15 changes: 11 additions & 4 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,18 @@
from src.models import OpenAIModel
from src.memory import Memory
from src.logger import logger
from src.storage import Storage
from src.storage import Storage, FileStorage, MongoStorage
from src.utils import get_role_and_content
from src.service.youtube import Youtube, YoutubeTranscriptReader
from src.service.website import Website, WebsiteReader
from src.mongodb import mongodb

load_dotenv('.env')

app = Flask(__name__)
line_bot_api = LineBotApi(os.getenv('LINE_CHANNEL_ACCESS_TOKEN'))
handler = WebhookHandler(os.getenv('LINE_CHANNEL_SECRET'))
storage = Storage('db.json')
storage = None
youtube = Youtube(step=4)
website = Website()

Expand Down Expand Up @@ -62,8 +63,9 @@ def handle_text_message(event):
if not is_successful:
raise ValueError('Invalid API token')
model_management[user_id] = model
api_keys[user_id] = api_key
storage.save(api_keys)
storage.save({
user_id: api_key
})
msg = TextSendMessage(text='Token 有效,註冊成功')

elif text.startswith('/指令說明'):
Expand Down Expand Up @@ -180,6 +182,11 @@ def home():


if __name__ == "__main__":
if os.getenv('USE_MONGO'):
mongodb.connect_to_database()
storage = Storage(MongoStorage(mongodb.db))
else:
storage = Storage(FileStorage('db.json'))
try:
data = storage.load()
for user_id in data.keys():
Expand Down
23 changes: 23 additions & 0 deletions src/mongodb.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
import os

from pymongo import MongoClient


class MongoDB():
"""
Environment Variables:
MONGODB__PATH
MONGODB__DBNAME
"""
client: None
db: None

def connect_to_database(self, mongo_path=None, db_name=None):
mongo_path = mongo_path or os.getenv('MONGODB__PATH')
db_name = db_name or os.getenv('MONGODB__DBNAME')
self.client = MongoClient(mongo_path)
assert self.client.config.command('ping')['ok'] == 1.0
self.db = self.client[db_name]


mongodb = MongoDB()
45 changes: 42 additions & 3 deletions src/storage.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,54 @@
import json
import datetime


class Storage():
class FileStorage:
def __init__(self, file_name):
self.fine_name = file_name
self.history = {}

def save(self, data):
self.history.update(data)
with open(self.fine_name, 'w', newline='') as f:
json.dump(data, f)
json.dump(self.history, f)

def load(self):
with open(self.fine_name, newline='') as jsonfile:
data = json.load(jsonfile)
return data
self.history = data
return self.history


class MongoStorage:
def __init__(self, db):
self.db = db

def save(self, data):
user_id, api_key = list(data.items())[0]
self.db['api_key'].update_one({
'user_id': user_id
}, {
'$set': {
'user_id': user_id,
'api_key': api_key,
'created_at': datetime.datetime.utcnow()
}
}, upsert=True)

def load(self):
data = list(self.db['api_key'].find())
res = {}
for i in range(len(data)):
res[data[i]['user_id']] = data[i]['api_key']
return res


class Storage:
def __init__(self, storage):
self.storage = storage

def save(self, data):
self.storage.save(data)

def load(self):
return self.storage.load()

0 comments on commit a316a11

Please sign in to comment.