Update "no OpenAPI API key" warning

loopgpt/__init__.py

@@ -11,28 +11,18 @@
 from loopgpt.memory import from_config as memory_from_config
 from loopgpt.embeddings import from_config as embedding_provider_from_config
 
-from loopgpt.logger import logger
-from colorama import Fore, Style
-
+from loopgpt.utils.openai_key import check_openai_key
 from dotenv import load_dotenv
 
-import os
+import sys
 
 
 load_dotenv()
 
 
-def check_openai_key():
-    if "OPENAI_API_KEY" not in os.environ:
-        logger.warn(
-            f"{Fore.RED}WARNING: OpenAI API Key not found. Please set the `OPENAI_API_KEY` environment variable. "
-            f"LoopGPT cannot work without it. "
-            f"See https://github.com/farizrahman4u/loopgpt#-requirements for more details{Style.RESET_ALL}"
-        )
-
-
 def from_config(config):
     return globals()[config["type"] + "_from_config"](config)
 
 
-check_openai_key()
+if "pytest" not in sys.modules:
+    check_openai_key()

loopgpt/agent.py

@@ -5,7 +5,7 @@
     DEFAULT_AGENT_DESCRIPTION,
     NEXT_PROMPT,
     NEXT_PROMPT_SMALL,
-    AgentStates
+    AgentStates,
 )
 from loopgpt.memory import from_config as memory_from_config
 from loopgpt.models import OpenAIModel, from_config as model_from_config
@@ -132,15 +132,17 @@ def _get_compressed_history(self):
         user_msgs = [i for i in range(len(hist)) if hist[i]["role"] == "user"]
         hist = [hist[i] for i in range(len(hist)) if i not in user_msgs]
         return hist
-    
+
     def get_full_message(self, message: Optional[str]):
         if self.state == AgentStates.START:
             return self.init_prompt + "\n\n" + (message or "")
         else:
             return self.next_prompt + "\n\n" + (message or "")
 
     @spinner
-    def chat(self, message: Optional[str] = None, run_tool=False) -> Optional[Union[str, Dict]]:
+    def chat(
+        self, message: Optional[str] = None, run_tool=False
+    ) -> Optional[Union[str, Dict]]:
         if self.state == AgentStates.STOP:
             raise ValueError(
                 "This agent has completed its tasks. It will not accept any more messages."

loopgpt/constants.py

@@ -86,12 +86,14 @@
     "Always execute plans to completion",
 ]
 
+
 class AgentStates:
     START = "START"
     IDLE = "IDLE"
     TOOL_STAGED = "TOOL_STAGED"
     STOP = "STOP"
 
+
 # SPINNER
 SPINNER_ENABLED = True
 SPINNER_START_DELAY = 2

loopgpt/loops/cli.py

@@ -13,7 +13,9 @@ def main():
         help="Action to perform. If no file is specified, a new agent is created.",
     )
     parser.add_argument("filename", nargs="?", help="Agent state JSON.", default=None)
-    parser.add_argument("--model", help="Model to use, uses gpt-3.5-turbo by default.", default=None)
+    parser.add_argument(
+        "--model", help="Model to use, uses gpt-3.5-turbo by default.", default=None
+    )
     parser.add_argument(
         "--readonly",
         help="Read only mode. Does not write agent state to disk.",

loopgpt/models/openai_.py

@@ -1,6 +1,7 @@
 from typing import *
 from loopgpt.logger import logger
 from loopgpt.models.base import BaseModel
+
 import tiktoken
 import time
 import os
@@ -10,9 +11,11 @@ def _getkey(key: Optional[str] = None):
     key = key or os.getenv("OPENAI_API_KEY")
     if key is None:
         raise ValueError(
-            "OpenAI API Key not found. Please set the `OPENAI_API_KEY` environment variable. "
-            "See https://github.com/farizrahman4u/loopgpt#-requirements for more details"
+            f"OpenAI API Key not found in the current working directory: {os.getcwd()}. "
+            "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. "
+            "See https://github.com/farizrahman4u/loopgpt#setup-your-openai-api-key- for more details"
         )
+    return key
 
 
 class OpenAIModel(BaseModel):

loopgpt/tools/filesystem.py

@@ -106,7 +106,7 @@ def resp(self):
         return {
             "result": "list of files and directories",
         }
-    
+
     @property
     def desc(self):
         return "List files and directories in a given path. Directories end with a trailing slash."
@@ -120,7 +120,14 @@ def run(self, path, recursive, show_hidden=False, exclude_dirs=False):
                         if not exclude_dirs:
                             entries_list.append(f"{entry.name}/")
                         if recursive:
-                            entries_list.extend(self.run(os.path.join(path, entry.name), recursive, show_hidden, exclude_dirs)["result"])
+                            entries_list.extend(
+                                self.run(
+                                    os.path.join(path, entry.name),
+                                    recursive,
+                                    show_hidden,
+                                    exclude_dirs,
+                                )["result"]
+                            )
                     else:
                         entries_list.append(entry.name)
         return {"result": entries_list}

loopgpt/utils/openai_key.py

@@ -0,0 +1,21 @@
+import os
+import sys
+
+from dotenv import load_dotenv
+
+
+def check_openai_key():
+    if "OPENAI_API_KEY" not in os.environ:
+        yn = input(
+            "OpenAI API key not found. Would you like to add your key to `.env` now? (y/n): "
+        )
+        if yn.lower() == "y":
+            key = input("Please enter your OpenAI API key: ")
+            with open(".env", "w") as f:
+                f.write(f'OPENAI_API_KEY = "{key}"')
+                print("Key added to `.env`")
+            load_dotenv()
+        else:
+            print(
+                "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. LoopGPT cannot work without it."
+            )

tests/test_keys.py

@@ -0,0 +1,19 @@
+import builtins
+import os
+
+
+def test_add_key_prompt(monkeypatch):
+    # Set up mock user input
+    responses = iter(["y", "my-api-key"])
+    monkeypatch.setattr(builtins, "input", lambda _: next(responses))
+
+    from loopgpt.utils.openai_key import check_openai_key
+
+    # Run the function
+    check_openai_key()
+
+    # Assert that the .env file was created with the correct key value
+    with open(".env", "r") as f:
+        assert f.read().strip() == 'OPENAI_API_KEY = "my-api-key"'
+
+    os.remove(".env")