From 103b920cb777d0739f6b2d50d1cfdd1a6c06f606 Mon Sep 17 00:00:00 2001
From: pavel-pythagora <pavel@pythagora.ai>
Date: Thu, 28 Sep 2023 13:29:21 +0300
Subject: [PATCH] Add arguments for extension

---
 pilot/.env.example   | 2 +-
 pilot/main.py        | 3 +++
 pilot/utils/files.py | 2 +-
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/pilot/.env.example b/pilot/.env.example
index bd6f817..db2ecd2 100644
--- a/pilot/.env.example
+++ b/pilot/.env.example
@@ -1,7 +1,7 @@
 # OPENAI or AZURE or OPENROUTER
 ENDPOINT=OPENAI
 
-OPENAI_ENDPOINT=
+OPENAI_ENDPOINT=https://api.openai.com/v1/chat/completions
 OPENAI_API_KEY=
 
 AZURE_API_KEY=
diff --git a/pilot/main.py b/pilot/main.py
index ea485f4..36076cb 100644
--- a/pilot/main.py
+++ b/pilot/main.py
@@ -2,6 +2,7 @@
 from __future__ import print_function, unicode_literals
 import builtins
 import json
+import os
 
 import sys
 import traceback
@@ -76,6 +77,8 @@ if __name__ == "__main__":
     try:
         args = init()
         builtins.print, ipc_client_instance = get_custom_print(args)
+        if '--api-key' in args:
+          os.environ["OPENAI_API_KEY"] = args['--api-key']
         if '--get-created-apps-with-steps' in args:
             print({ 'db_data': get_created_apps_with_steps() }, type='info')
         else:
diff --git a/pilot/utils/files.py b/pilot/utils/files.py
index b77b5ae..6536e9c 100644
--- a/pilot/utils/files.py
+++ b/pilot/utils/files.py
@@ -20,7 +20,7 @@ def setup_workspace(args):
 
         return args['workspace']
 
-    root = get_parent_folder('pilot')
+    root = args['root'] or get_parent_folder('pilot')
     create_directory(root, 'workspace')
     project_path = create_directory(os.path.join(root, 'workspace'), args['name'])
     create_directory(project_path, 'tests')