From 73d0f71d95578dadc22b2884cc88341c632e65d9 Mon Sep 17 00:00:00 2001
From: ciuzaak <as460230379@hotmail.com>
Date: Mon, 15 May 2023 00:48:58 +0800
Subject: [PATCH] feat: get config from environment variables (#9)

---
 README.md          | 13 +++++++++++--
 config/__init__.py | 26 +++++++++++++++++---------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index b754ac6..75e9b6e 100644
--- a/README.md
+++ b/README.md
@@ -31,12 +31,21 @@ If you only have access to one of the models, you can still continue to use this
 
 1. Clone this repository.
 
-2. Fill in `config/config.yml` with reference to `config/config.example.yml`.
+2. Configure the bot in the following two ways:
+   1. Create `config/config.yml` and fill in the information with reference to `config/config.example.yml`.
+   2. Set environment variables:
+
+      ```bash
+      export BOT_TOKEN="your bot token"
+      export USER_IDS="user_id1,user_id2,..."
+      export CLAUDE_API="your claude api" # ignore it if you don't want to use claude
+      export BARD_API="your bard api" # ignore it if you don't want to use bard
+      ```
 
     - [How to obtain telegram bot token](https://core.telegram.org/bots/tutorial#obtain-your-bot-token)
     - [How to obtain telegram user id](https://bigone.zendesk.com/hc/en-us/articles/360008014894-How-to-get-the-Telegram-user-ID-)
 
-3. Start the bot by:
+3. Start the bot in the following two ways:
 
     - **Docker** (with docker engine and docker-compose pre-installed):
 
diff --git a/config/__init__.py b/config/__init__.py
index cf53d51..a7d8d3d 100644
--- a/config/__init__.py
+++ b/config/__init__.py
@@ -1,14 +1,22 @@
-from yaml import safe_load
+from os import getenv, path
 
-# load yaml config
-with open('config/config.yml', 'r') as f:
-    config_yaml = safe_load(f)
+from yaml import safe_load
 
-# config parameters
-bot_token = config_yaml['telegram']['bot_token']
-user_ids = config_yaml['telegram']['user_ids']
-claude_api = config_yaml['claude']['api']
-bard_api = config_yaml['bard']['api']
+if not path.exists('config/config.yml'):
+    # load env vars
+    bot_token = getenv('BOT_TOKEN')
+    user_ids = [int(user_id) for user_id in getenv('USER_IDS').split(',')]
+    claude_api = getenv('CLAUDE_API')
+    bard_api = getenv('BARD_API')
+else:
+    # load yaml config
+    with open('config/config.yml', 'r') as f:
+        config_yaml = safe_load(f)
+    # config parameters
+    bot_token = config_yaml['telegram']['bot_token']
+    user_ids = config_yaml['telegram']['user_ids']
+    claude_api = config_yaml['claude']['api']
+    bard_api = config_yaml['bard']['api']
 
 assert bot_token is not None and user_ids is not None
 assert claude_api is not None or bard_api is not None