Skip to content

tool_callback_handler

ToolCallbackHandler

Source code in wizard_ai/conversational_engine/tool_callback_handler.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class ToolCallbackHandler:

    def __init__(
        self,
        chat_id: str,
        tools: list = None,
        rabbitmq_producer: RabbitMQProducer = None,
        queue: MessageQueues = None
    ) -> None:
        self.chat_id = chat_id
        self.tools = tools
        self.rabbitmq_client = rabbitmq_producer
        self.queue = queue

    def on_tool_start(
        self,
        tool: FormTool,
        tool_input: str
    ) -> Any:
        """Run when tool starts running."""
        if not self.rabbitmq_client:
            return

        try:
            tool_start_message = tool.get_tool_start_message(tool_input)
        except BaseException:
            tool_start_message = f"{tool.name}: {tool_input}"

        self.rabbitmq_client.publish(
            queue=self.queue,
            message=json.dumps({
                "chat_id": self.chat_id,
                "type": MessageType.TOOL_START.value,
                "content": tool_start_message
            })
        )

    def on_tool_end(
        self,
        tool: FormTool,
        tool_output: str
    ) -> Any:
        """Run when tool ends running."""

        if not self.rabbitmq_client:
            return

        self.rabbitmq_client.publish(
            queue=self.queue,
            message=json.dumps({
                "chat_id": self.chat_id,
                "type": MessageType.TOOL_END.value,
                "content": tool_output
            })
        )

on_tool_end(tool, tool_output)

Run when tool ends running.

Source code in wizard_ai/conversational_engine/tool_callback_handler.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def on_tool_end(
    self,
    tool: FormTool,
    tool_output: str
) -> Any:
    """Run when tool ends running."""

    if not self.rabbitmq_client:
        return

    self.rabbitmq_client.publish(
        queue=self.queue,
        message=json.dumps({
            "chat_id": self.chat_id,
            "type": MessageType.TOOL_END.value,
            "content": tool_output
        })
    )

on_tool_start(tool, tool_input)

Run when tool starts running.

Source code in wizard_ai/conversational_engine/tool_callback_handler.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def on_tool_start(
    self,
    tool: FormTool,
    tool_input: str
) -> Any:
    """Run when tool starts running."""
    if not self.rabbitmq_client:
        return

    try:
        tool_start_message = tool.get_tool_start_message(tool_input)
    except BaseException:
        tool_start_message = f"{tool.name}: {tool_input}"

    self.rabbitmq_client.publish(
        queue=self.queue,
        message=json.dumps({
            "chat_id": self.chat_id,
            "type": MessageType.TOOL_START.value,
            "content": tool_start_message
        })
    )