mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-08-17 02:10:31 +03:00
chore: tools (#2575)
# Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate):
This commit is contained in:
parent
fb46bb9341
commit
03ebbe70fe
@ -24,6 +24,7 @@ from modules.chat.dto.outputs import GetChatHistoryOutput
|
||||
from modules.chat.service.chat_service import ChatService
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
from modules.tools import ImageGeneratorTool
|
||||
|
||||
|
||||
class AgentState(TypedDict):
|
||||
@ -36,56 +37,6 @@ logger = get_logger(__name__)
|
||||
|
||||
chat_service = ChatService()
|
||||
|
||||
|
||||
class ImageGenerationInput(BaseModelV1):
|
||||
query: str = FieldV1(
|
||||
...,
|
||||
title="description",
|
||||
description="A detailled prompt to generate the image from. Takes into account the history of the chat.",
|
||||
)
|
||||
|
||||
|
||||
class ImageGeneratorTool(BaseTool):
|
||||
name = "image-generator"
|
||||
description = "useful for when you need to generate an image from a prompt."
|
||||
args_schema: Type[BaseModel] = ImageGenerationInput
|
||||
return_direct = True
|
||||
|
||||
def _run(
|
||||
self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
|
||||
) -> str:
|
||||
client = OpenAI()
|
||||
|
||||
response = client.images.generate(
|
||||
model="dall-e-3",
|
||||
prompt=query,
|
||||
size="1024x1024",
|
||||
quality="standard",
|
||||
n=1,
|
||||
)
|
||||
image_url = response.data[0].url
|
||||
revised_prompt = response.data[0].revised_prompt
|
||||
# Make the url a markdown image
|
||||
return f"{revised_prompt} \n ![Generated Image]({image_url}) "
|
||||
|
||||
async def _arun(
|
||||
self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
|
||||
) -> str:
|
||||
"""Use the tool asynchronously."""
|
||||
client = OpenAI()
|
||||
response = await run_manager.run_async(
|
||||
client.images.generate,
|
||||
model="dall-e-3",
|
||||
prompt=query,
|
||||
size="1024x1024",
|
||||
quality="standard",
|
||||
n=1,
|
||||
)
|
||||
image_url = response.data[0].url
|
||||
# Make the url a markdown image
|
||||
return f"![Generated Image]({image_url})"
|
||||
|
||||
|
||||
class GPT4Brain(KnowledgeBrainQA):
|
||||
"""This is the Notion brain class. it is a KnowledgeBrainQA has the data is stored locally.
|
||||
It is going to call the Data Store internally to get the data.
|
||||
|
2
backend/modules/tools/__init__.py
Normal file
2
backend/modules/tools/__init__.py
Normal file
@ -0,0 +1,2 @@
|
||||
from .image_generator import ImageGeneratorTool
|
||||
|
61
backend/modules/tools/image_generator.py
Normal file
61
backend/modules/tools/image_generator.py
Normal file
@ -0,0 +1,61 @@
|
||||
from typing import Optional, Type
|
||||
|
||||
from langchain.callbacks.manager import (
|
||||
AsyncCallbackManagerForToolRun,
|
||||
CallbackManagerForToolRun,
|
||||
)
|
||||
from langchain.pydantic_v1 import BaseModel as BaseModelV1
|
||||
from langchain.pydantic_v1 import Field as FieldV1
|
||||
from langchain.tools import BaseTool
|
||||
from langchain_core.tools import BaseTool
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class ImageGenerationInput(BaseModelV1):
|
||||
query: str = FieldV1(
|
||||
...,
|
||||
title="description",
|
||||
description="A detailled prompt to generate the image from. Takes into account the history of the chat.",
|
||||
)
|
||||
|
||||
|
||||
class ImageGeneratorTool(BaseTool):
|
||||
name = "image-generator"
|
||||
description = "useful for when you need to generate an image from a prompt."
|
||||
args_schema: Type[BaseModel] = ImageGenerationInput
|
||||
return_direct = True
|
||||
|
||||
def _run(
|
||||
self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
|
||||
) -> str:
|
||||
client = OpenAI()
|
||||
|
||||
response = client.images.generate(
|
||||
model="dall-e-3",
|
||||
prompt=query,
|
||||
size="1024x1024",
|
||||
quality="standard",
|
||||
n=1,
|
||||
)
|
||||
image_url = response.data[0].url
|
||||
revised_prompt = response.data[0].revised_prompt
|
||||
# Make the url a markdown image
|
||||
return f"{revised_prompt} \n ![Generated Image]({image_url}) "
|
||||
|
||||
async def _arun(
|
||||
self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
|
||||
) -> str:
|
||||
"""Use the tool asynchronously."""
|
||||
client = OpenAI()
|
||||
response = await run_manager.run_async(
|
||||
client.images.generate,
|
||||
model="dall-e-3",
|
||||
prompt=query,
|
||||
size="1024x1024",
|
||||
quality="standard",
|
||||
n=1,
|
||||
)
|
||||
image_url = response.data[0].url
|
||||
# Make the url a markdown image
|
||||
return f"![Generated Image]({image_url})"
|
Loading…
Reference in New Issue
Block a user