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.
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
from typing import List, Dict, Callable
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.tool_model import ToolDefinition
|
|
from app.repositories.tool_repository import ToolRepository
|
|
from app.services.handlers import (
|
|
consultar_estoque,
|
|
validar_cliente_venda,
|
|
avaliar_veiculo_troca,
|
|
agendar_revisao,
|
|
cancelar_pedido,
|
|
)
|
|
|
|
|
|
HANDLERS: Dict[str, Callable] = {
|
|
"consultar_estoque": consultar_estoque,
|
|
"validar_cliente_venda": validar_cliente_venda,
|
|
"avaliar_veiculo_troca": avaliar_veiculo_troca,
|
|
"agendar_revisao": agendar_revisao,
|
|
"cancelar_pedido": cancelar_pedido,
|
|
}
|
|
|
|
|
|
class ToolRegistry:
|
|
|
|
def __init__(self, db: Session):
|
|
self._tools = []
|
|
repo = ToolRepository(db)
|
|
db_tools = repo.get_all()
|
|
for db_tool in db_tools:
|
|
handler = HANDLERS.get(db_tool.name)
|
|
if not handler:
|
|
continue
|
|
self.register_tool(
|
|
name=db_tool.name,
|
|
description=db_tool.description,
|
|
parameters=db_tool.parameters,
|
|
handler=handler,
|
|
)
|
|
|
|
# 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)
|