Merge pull request #21 from ericmjl/update-openai

refactor: change openai api key assignment and model list retrieval

llamabot/bot/__init__.py

@@ -3,7 +3,6 @@
 import os
 import warnings
 
-import openai
 import panel as pn
 from dotenv import load_dotenv
 
@@ -26,7 +25,6 @@
         "No OpenAI API key found. Please set OPENAI_API_KEY in your environment.",
         category=RuntimeWarning,
     )
-openai.api_key = api_key
 
 
 __all__ = ["SimpleBot", "ChatBot", "QueryBot", "ImageBot"]

llamabot/bot/imagebot.py

@@ -5,6 +5,7 @@
 import requests
 from pathlib import Path
 from typing import Union
+from langchain.schema import AIMessage
 
 
 class ImageBot:
@@ -47,18 +48,8 @@ def __call__(self, prompt: str, save_path: Path = None) -> Union[str, Path]:
 
         image_data = requests.get(image_url).content
 
-        from llamabot import SimpleBot
-
-        bot = SimpleBot(
-            "You are a helpful filenaming assistant. "
-            "Filenames should use underscores instead of spaces, "
-            "and should be all lowercase. "
-            "Exclude the file extension. "
-            "Give me a compact filename for the following prompt:"
-        )
-        response = bot(prompt)
-        filename = response.message
         if not save_path:
+            filename = filename_bot(prompt).content
             save_path = Path(f"{filename}.jpg")
         with open(save_path, "wb") as file:
             file.write(image_data)
@@ -75,3 +66,22 @@ def is_running_in_jupyter() -> bool:
         return True
     except NameError:
         return False
+
+
+def filename_bot(image_prompt: str) -> AIMessage:
+    """Generate a filename from an image prompt.
+
+    :param image_prompt: The image prompt to generate a filename from.
+    :return: The filename generated from the image prompt within an AIMessage object.
+    """
+    from llamabot import SimpleBot
+
+    bot = SimpleBot(
+        "You are a helpful filenaming assistant. "
+        "Filenames should use underscores instead of spaces, "
+        "and should be all lowercase. "
+        "Exclude the file extension. "
+        "Give me a compact filename for the following prompt:"
+    )
+    response = bot(image_prompt)
+    return response

llamabot/cli/configure.py

@@ -1,11 +1,14 @@
 """LlamaBot configuration."""
-import openai
+from openai import OpenAI
+import os
+
 import typer
 from prompt_toolkit import prompt
 from prompt_toolkit.completion import WordCompleter
-import os
 from .utils import configure_environment_variable
 
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
 app = typer.Typer()
 
 
@@ -43,9 +46,7 @@ def default_model(model_name=None):
 
     load_dotenv(llamabotrc_path)
 
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-
-    model_list = openai.Model.list()["data"]
+    model_list = client.models.list()["data"]
     available_models = [x["id"] for x in model_list if "gpt" in x["id"]]
     available_models.sort()
 

tests/bot/test_imagebot.py

@@ -1,6 +1,8 @@
 """Test the ImageBot class."""
 from llamabot import ImageBot, SimpleBot
 import requests
+from langchain.schema import AIMessage
+from pathlib import Path
 
 
 def test_initialization_defaults():
@@ -80,3 +82,48 @@ def test_call_outside_jupyter(mocker, tmp_path):
     result = bot("test prompt", tmp_path / "test_prompt.jpg")
     assert result == tmp_path / "test_prompt.jpg"
     requests.get.assert_called_with("http://image.url")
+
+
+def test_call_outside_jupyter_no_save_path(mocker, tmp_path):
+    """Test the call method when not running in a Jupyter notebook
+    and no save path is provided.
+
+    This test checks that the call method
+    saves the image to a generated filename based on the prompt
+    when not running in a Jupyter notebook and no save path is provided.
+
+    :param mocker: The pytest-mock fixture.
+    """
+    # Mock the is_running_in_jupyter method
+    mocker.patch("llamabot.bot.imagebot.is_running_in_jupyter", return_value=False)
+
+    # Instantiate ImageBot
+    bot = ImageBot()
+
+    # Mock the client's generate method on the instance to return the desired URL
+    mock_response = mocker.MagicMock()
+    mock_response.data = [mocker.MagicMock(url="http://image.url")]
+    bot.client = mocker.MagicMock()
+    bot.client.images.generate.return_value = mock_response
+
+    # Mock requests.get to return a mock response with content
+    mock_get_response = mocker.MagicMock()
+    mock_get_response.content = b"image_data"
+    mocker.patch("requests.get", return_value=mock_get_response)
+
+    # Mock the filename_bot method to return a generated filename
+    generated_filename = str(tmp_path / "generated_filename")
+    mocker.patch(
+        "llamabot.bot.imagebot.filename_bot",
+        return_value=AIMessage(content=generated_filename),
+    )
+
+    # Call the method and perform the assertion
+    result = bot("test prompt")
+    assert result == Path(f"{generated_filename}.jpg")
+    requests.get.assert_called_with("http://image.url")
+
+    # Check if the file was saved correctly
+    with open(result, "rb") as file:
+        saved_content = file.read()
+    assert saved_content == b"image_data"

tests/prompt_library/test_python_prompt_library.py

@@ -9,6 +9,7 @@
 
 
 @given(st.text())
+@settings(deadline=None)
 def test_codebot_instance(input_text: str):
     """Test that codebot returns a SimpleBot instance.