Feat/add azure dalle tool (#2276)
Co-authored-by: lux@njuelectronics.com <lux@njuelectronics.com> Co-authored-by: crazywoola <427733928@qq.com>pull/2281/head
呆萌闷油瓶
2 years ago
committed by
GitHub
parent
76cc19f525
commit
c97b7f6748
Binary file not shown.
|
After Width: | Height: | Size: 153 KiB |
@ -0,0 +1,23 @@
|
||||
from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
|
||||
from core.tools.provider.builtin.azuredalle.tools.dalle3 import DallE3Tool
|
||||
from core.tools.errors import ToolProviderCredentialValidationError
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
class AzureDALLEProvider(BuiltinToolProviderController):
|
||||
def _validate_credentials(self, credentials: Dict[str, Any]) -> None:
|
||||
try:
|
||||
DallE3Tool().fork_tool_runtime(
|
||||
meta={
|
||||
"credentials": credentials,
|
||||
}
|
||||
).invoke(
|
||||
user_id='',
|
||||
tool_paramters={
|
||||
"prompt": "cute girl, blue eyes, white hair, anime style",
|
||||
"size": "square",
|
||||
"n": 1
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
raise ToolProviderCredentialValidationError(str(e))
|
||||
@ -0,0 +1,66 @@
|
||||
from typing import Any, Dict, List, Union
|
||||
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
from core.tools.tool.builtin_tool import BuiltinTool
|
||||
|
||||
from base64 import b64decode
|
||||
from os.path import join
|
||||
|
||||
from openai import AzureOpenAI
|
||||
|
||||
class DallE3Tool(BuiltinTool):
|
||||
def _invoke(self,
|
||||
user_id: str,
|
||||
tool_paramters: Dict[str, Any],
|
||||
) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]:
|
||||
"""
|
||||
invoke tools
|
||||
"""
|
||||
client = AzureOpenAI(
|
||||
api_version=self.runtime.credentials['azure_openai_api_version'],
|
||||
azure_endpoint=self.runtime.credentials['azure_openai_base_url'],
|
||||
api_key=self.runtime.credentials['azure_openai_api_key'],
|
||||
)
|
||||
|
||||
SIZE_MAPPING = {
|
||||
'square': '1024x1024',
|
||||
'vertical': '1024x1792',
|
||||
'horizontal': '1792x1024',
|
||||
}
|
||||
|
||||
# prompt
|
||||
prompt = tool_paramters.get('prompt', '')
|
||||
if not prompt:
|
||||
return self.create_text_message('Please input prompt')
|
||||
# get size
|
||||
size = SIZE_MAPPING[tool_paramters.get('size', 'square')]
|
||||
# get n
|
||||
n = tool_paramters.get('n', 1)
|
||||
# get quality
|
||||
quality = tool_paramters.get('quality', 'standard')
|
||||
if quality not in ['standard', 'hd']:
|
||||
return self.create_text_message('Invalid quality')
|
||||
# get style
|
||||
style = tool_paramters.get('style', 'vivid')
|
||||
if style not in ['natural', 'vivid']:
|
||||
return self.create_text_message('Invalid style')
|
||||
|
||||
# call openapi dalle3
|
||||
model=self.runtime.credentials['azure_openai_api_model_name']
|
||||
response = client.images.generate(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
size=size,
|
||||
n=n,
|
||||
style=style,
|
||||
quality=quality,
|
||||
response_format='b64_json'
|
||||
)
|
||||
|
||||
result = []
|
||||
|
||||
for image in response.data:
|
||||
result.append(self.create_blob_message(blob=b64decode(image.b64_json),
|
||||
meta={ 'mime_type': 'image/png' },
|
||||
save_as=self.VARIABLE_KEY.IMAGE.value))
|
||||
|
||||
return result
|
||||
Loading…
Reference in New Issue