You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
from typing import List
|
|
from app.models.tool_model import ToolDefinition
|
|
# from app.integrations.estoque_service import consultar_estoque
|
|
|
|
|
|
class ToolRegistry:
|
|
|
|
def __init__(self):
|
|
self._tools = []
|
|
"""
|
|
# Registrando manualmente por enquanto, porém isso poderia vir de um BD.
|
|
self.register_tool(
|
|
name="consultar_estoque",
|
|
description="Consulta veículos disponíveis até determinado preço", # Descrição que ajuda o modelo Gemini a entender quado selecionar essa Tool.
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"preco_max": {
|
|
"type": "number",
|
|
"description": "Preço máximo do veículo"
|
|
}
|
|
},
|
|
"required": ["preco_max"]
|
|
},
|
|
handler=consultar_estoque
|
|
)
|
|
"""
|
|
|
|
# O método abaixo serve para adicionar uma nova tool ao sistema (podendo inserir a lógica de adcicionar ao BD futuramente).
|
|
def register_tool(self, name, description, parameters, handler):
|
|
self._tools.append(
|
|
ToolDefinition(
|
|
name=name,
|
|
description=description,
|
|
parameters=parameters,
|
|
handler=handler
|
|
)
|
|
)
|
|
|
|
|
|
# Retorna todas as tools registradas, isso será usado pelo LLMService para enviar as ferramentas para o modelo
|
|
def get_tools(self) -> List[ToolDefinition]:
|
|
return self._tools
|
|
|
|
|
|
"""
|
|
Essa função é responsável por executar a tool solicitada pelo modelo.
|
|
Parâmetros:
|
|
- name: nome da função que o Gemini decidiu chamar
|
|
- arguments: argumentos extraídos da mensagem do usuário
|
|
"""
|
|
async def execute(self, name: str, arguments: dict):
|
|
tool = next((t for t in self._tools if t.name == name), None) # Procura dentro da lista de tools aquela que tem o mesmo nome
|
|
|
|
if not tool:
|
|
raise Exception(f"Tool {name} não encontrada.")
|
|
|
|
return await tool.handler(**arguments)
|