From 3372a3b2224fa9c29cde7d587be2f501c31395d1 Mon Sep 17 00:00:00 2001
From: Tamerlan Beyimjan <beyimjan@gmail.com>
Date: Tue, 2 Apr 2024 04:48:04 +0500
Subject: [PATCH] Add OpenAI org ID prompt to config setup

---
 shell_whiz/cli/commands/config.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/shell_whiz/cli/commands/config.py b/shell_whiz/cli/commands/config.py
index f8ecea4..27c7d74 100644
--- a/shell_whiz/cli/commands/config.py
+++ b/shell_whiz/cli/commands/config.py
@@ -15,14 +15,22 @@ def config() -> None:
     rich.print(
         "Visit https://platform.openai.com/api-keys to get your API key."
     )
+    openai_api_key = questionary.text(
+        "OpenAI API key",
+        default=os.environ.get("OPENAI_API_KEY", ""),
+        validate=lambda text: len(text) > 0,
+    ).unsafe_ask()
+
+    rich.print(
+        "\n[bold italic](Optional)[/] Leave blank if you are not a member of multiple organizations."
+    )
+    openai_org_id = questionary.text(
+        "Organization ID", default=os.environ.get("OPENAI_ORG_ID", "")
+    ).unsafe_ask()
 
     try:
         config = ConfigModel(
-            openai_api_key=questionary.text(
-                "OpenAI API key",
-                default=os.environ.get("OPENAI_API_KEY", ""),
-                validate=lambda text: len(text) > 0,
-            ).unsafe_ask()
+            openai_api_key=openai_api_key, openai_org_id=openai_org_id or None
         )
     except pydantic.ValidationError:
         rich.print(
@@ -30,6 +38,9 @@ def config() -> None:
         )
         raise typer.Exit(1)
 
+    print(config)
+    return
+
     try:
         Config.write(config)
     except ConfigError as e: