Added: org_id config (#63)

signebedi · May 14, 2023 · 9d051e7 · 9d051e7
1 parent 62258bc
commit 9d051e7
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 0 deletions.
diff --git a/assets/gptty.ini.example b/assets/gptty.ini.example
@@ -1,5 +1,6 @@
 [main]
 api_key=SOME_KEY_HERE
+org_id=org-ADD_YOUR_ORG_ID
 your_name=question
 gpt_name=response
 output_file=output.txt

diff --git a/gptty/config.py b/gptty/config.py
@@ -39,6 +39,7 @@ def get_config_data(config_file='gptty.ini'):
 
     config['DEFAULT'] = {
         'api_key': "",
+        'org_id': "",
         'your_name': 'question',
         'gpt_name': 'response',
         'output_file': 'output.txt',
@@ -56,6 +57,7 @@ def get_config_data(config_file='gptty.ini'):
 
     parsed_data = {
         'api_key': config.get('main', 'api_key', fallback="",),
+        'org_id': config.get('main', 'org_id', fallback="",),
         'your_name': config.get('main', 'your_name', fallback='question'),
         'gpt_name': config.get('main', 'gpt_name', fallback='response'),
         'output_file': config.get('main', 'output_file', fallback='output.txt'),

diff --git a/gptty/gptty.py b/gptty/gptty.py
@@ -216,6 +216,7 @@ async def create_chat_room(configs=get_config_data(), log_responses:bool=True, c
         df = pd.DataFrame(columns=['timestamp','tag','question','response'])
 
     try:
+        openai.organization = configs['org_id'].rstrip('\n')
         openai.api_key = configs['api_key'].rstrip('\n')
     except:
         click.echo(f"{RED}FAILED to initialize connection to OpenAI. Have you added an API token? See gptty docs <https://github.com/signebedi/gptty#configuration> or <https://platform.openai.com/account/api-keys> for more information.")

diff --git a/tests/test_config.py b/tests/test_config.py
@@ -7,6 +7,7 @@ class TestConfig(unittest.TestCase):
     def test_default_config(self):
         default_config_data = get_config_data('test/test_default_gptty.ini')
         self.assertEqual(default_config_data['api_key'], "")
+        self.assertEqual(default_config_data['org_id'], "")
         self.assertEqual(default_config_data['your_name'], 'question')
         self.assertEqual(default_config_data['gpt_name'], 'response')
         self.assertEqual(default_config_data['output_file'], 'output.txt')
@@ -21,6 +22,7 @@ def test_default_config(self):
     def test_custom_config(self):
         custom_config_data = get_config_data(config_file='tests/test_gptty.ini')
         self.assertEqual(custom_config_data['api_key'], "ANOTHER_KEY_HERE")
+        self.assertEqual(default_config_data['org_id'], "org-536JCU1SlmsQZB0i734dCsGC")
         self.assertEqual(custom_config_data['your_name'], 'custom_question')
         self.assertEqual(custom_config_data['gpt_name'], 'custom_response')
         self.assertEqual(custom_config_data['output_file'], 'custom_output.txt')

diff --git a/tests/test_default_gptty.ini b/tests/test_default_gptty.ini
@@ -1,5 +1,6 @@
 [main]
 api_key=
+org_id=
 your_name=question
 gpt_name=response
 output_file=output.txt

diff --git a/tests/test_gptty.ini b/tests/test_gptty.ini
@@ -1,5 +1,6 @@
 [main]
 api_key=ANOTHER_KEY_HERE
+org_id=org-536JCU1SlmsQZB0i734dCsGC
 your_name=custom_question
 gpt_name=custom_response
 output_file=custom_output.txt