Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dalle openai fix #953

Merged
merged 6 commits into from
Aug 2, 2023
Merged

Dalle openai fix #953

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
29d7e28
Dalle open ai api key fix
Aug 2, 2023
1fec770
Changes
Aug 2, 2023
04ce082
Update dalle_image_gen.py
Tarraann Aug 2, 2023
b3f818f
Merge branch 'main' into dalle_openai_fix
Aug 2, 2023
e621690
Changes
Aug 2, 2023
4fa10eb
Update test_dalle_image_gen.py
Tarraann Aug 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion superagi/tools/image_generation/dalle_image_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
from superagi.image_llms.openai_dalle import OpenAiDalle
from superagi.llms.base_llm import BaseLlm
from superagi.resource_manager.file_manager import FileManager
from superagi.models.toolkit import Toolkit
from superagi.models.configuration import Configuration
from superagi.tools.base_tool import BaseTool

class DalleImageGenInput(BaseModel):
Expand All @@ -30,6 +32,7 @@ class DalleImageGenTool(BaseTool):
args_schema: Type[BaseModel] = DalleImageGenInput
description: str = "Generate Images using Dalle"
agent_id: int = None
agent_execution_id: int = None
resource_manager: Optional[FileManager] = None

class Config:
Expand All @@ -48,9 +51,19 @@ def _execute(self, prompt: str, image_names: list, size: int = 512, num: int = 2
Returns:
Image generated successfully message if image is generated or error message.
"""
session = self.toolkit_config.session
toolkit = session.query(Toolkit).filter(Toolkit.id == self.toolkit_config.toolkit_id).first()
organisation_id = toolkit.organisation_id
if size not in [256, 512, 1024]:
size = min([256, 512, 1024], key=lambda x: abs(x - size))
response = OpenAiDalle(api_key=self.get_tool_config("OPENAI_API_KEY"), number_of_results=num).generate_image(
api_key = self.get_tool_config("OPENAI_API_KEY")
if api_key is None:
model_source = Configuration.fetch_configuration(session, organisation_id, "model_source")
if model_source != "OpenAi":
return "Enter your OpenAi api key in the configuration"
api_key = Configuration.fetch_configuration(session, organisation_id, "model_api_key")

response = OpenAiDalle(api_key=api_key, number_of_results=num).generate_image(
prompt, size)
response = response.__dict__
response = response['_previous']['data']
Expand Down
36 changes: 32 additions & 4 deletions tests/unit_tests/tools/image_generation/test_dalle_image_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
from unittest.mock import patch, MagicMock

from superagi.tools.image_generation.dalle_image_gen import DalleImageGenTool


@patch('superagi.tools.image_generation.dalle_image_gen.OpenAiDalle')
@patch('superagi.tools.image_generation.dalle_image_gen.requests')
def test_execute_dalle_image_gen_tool(mock_requests, mock_dalle):
@patch('superagi.tools.image_generation.dalle_image_gen.Configuration')
def test_execute_dalle_image_gen_tool(mock_config, mock_requests, mock_dalle):
# Arrange
tool = DalleImageGenTool()
tool.toolkit_config = MagicMock(toolkit_id=1)
tool.toolkit_config.get_tool_config = MagicMock(return_value="test_api_key")
tool.toolkit_config.session = MagicMock()
tool.toolkit_config.session.query.return_value.filter.return_value.first.return_value = MagicMock(organisation_id="test_org_id")
tool.resource_manager = MagicMock()
mock_config.fetch_configuration = MagicMock(side_effect=("OpenAi", "test_api_key"))
mock_dalle_instance = mock_dalle.return_value
mock_dalle_instance.generate_image.return_value = MagicMock(
_previous=MagicMock(data=[
Expand All @@ -22,7 +26,7 @@ def test_execute_dalle_image_gen_tool(mock_requests, mock_dalle):
size = 512
num = 2
image_names = ["image1.png", "image2.png"]

# Act
result = tool._execute(prompt, image_names, size, num)

Expand All @@ -31,4 +35,28 @@ def test_execute_dalle_image_gen_tool(mock_requests, mock_dalle):
mock_dalle.assert_called_once_with(api_key="test_api_key", number_of_results=num)
mock_dalle_instance.generate_image.assert_called_once_with(prompt, size)
tool.resource_manager.write_binary_file.assert_any_call("image1.png", b"test_image_data")
tool.resource_manager.write_binary_file.assert_any_call("image2.png", b"test_image_data")
tool.resource_manager.write_binary_file.assert_any_call("image2.png", b"test_image_data")


@patch('superagi.tools.image_generation.dalle_image_gen.OpenAiDalle')
@patch('superagi.tools.image_generation.dalle_image_gen.requests')
@patch('superagi.tools.image_generation.dalle_image_gen.Configuration')
def test_execute_dalle_image_gen_tool_invalid_api_key(mock_config, mock_requests, mock_dalle):
tool = DalleImageGenTool()
tool.toolkit_config = MagicMock(toolkit_id=1)
tool.toolkit_config.get_tool_config = MagicMock(return_value=None)
tool.toolkit_config.session = MagicMock()
tool.toolkit_config.session.query.return_value.filter.return_value.first.return_value = MagicMock(organisation_id="test_org_id")
tool.resource_manager = MagicMock()
mock_config.fetch_configuration = MagicMock(return_value="notOpenAi")

prompt = "test_prompt"
size = 512
num = 2
image_names = ["image1.png", "image2.png"]

# Act
result = tool._execute(prompt, image_names, size, num)

# Assert
assert result == "Enter your OpenAi api key in the configuration"