Overview

graphorchestrator.nodes.nodes.AINode

Bases: ProcessingNode

A node that represents an AI model.

This node wraps an AI model action.

Source code in graphorchestrator\nodes\nodes.py

class AINode(ProcessingNode):
    """
    A node that represents an AI model.

    This node wraps an AI model action.
    """

    def __init__(
        self,
        node_id: str,
        description: str,
        model_action: Callable[[State], State],
        response_format: Optional[str] = None,
        response_parser: Optional[Callable[[State], Any]] = None,
    ) -> None:
        super().__init__(node_id, model_action)
        self.description = description
        self.response_format = response_format
        self.response_parser = response_parser

        GraphLogger.get().info(
            **wrap_constants(
                message="AINode created",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: self.node_id,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "node_created",
                    LC.CUSTOM: {
                        "description": description,
                        "response_format": response_format,
                        "has_parser": bool(response_parser),
                    },
                },
            )
        )

    async def execute(self, state: State) -> State:
        """
        Executes the AI model action.

        Args:
            state (State): The input state for the node.

        Returns:
            State: The state after executing the model action.
        """
        log = GraphLogger.get()

        log.info(
            **wrap_constants(
                message="AINode execution started",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: self.node_id,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "execute_start",
                    LC.INPUT_SIZE: len(state.messages),
                },
            )
        )

        result = await self.func(state)

        if not isinstance(result, State):
            log.error(
                **wrap_constants(
                    message="AINode returned invalid output",
                    **{
                        LC.EVENT_TYPE: "node",
                        LC.NODE_ID: self.node_id,
                        LC.NODE_TYPE: "AINode",
                        LC.ACTION: "invalid_output",
                        LC.CUSTOM: {"result_type": str(type(result))},
                    },
                )
            )
            raise InvalidAIActionOutput(result)

        log.info(
            **wrap_constants(
                message="AINode execution completed",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: self.node_id,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "execute_end",
                    LC.OUTPUT_SIZE: len(result.messages),
                    LC.SUCCESS: True,
                },
            )
        )

        return result

execute(state) async

Executes the AI model action.

Parameters:

Name	Type	Description	Default
state	State	The input state for the node.	required

Returns:

Name	Type	Description
State	State	The state after executing the model action.

Source code in graphorchestrator\nodes\nodes.py

async def execute(self, state: State) -> State:
    """
    Executes the AI model action.

    Args:
        state (State): The input state for the node.

    Returns:
        State: The state after executing the model action.
    """
    log = GraphLogger.get()

    log.info(
        **wrap_constants(
            message="AINode execution started",
            **{
                LC.EVENT_TYPE: "node",
                LC.NODE_ID: self.node_id,
                LC.NODE_TYPE: "AINode",
                LC.ACTION: "execute_start",
                LC.INPUT_SIZE: len(state.messages),
            },
        )
    )

    result = await self.func(state)

    if not isinstance(result, State):
        log.error(
            **wrap_constants(
                message="AINode returned invalid output",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: self.node_id,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "invalid_output",
                    LC.CUSTOM: {"result_type": str(type(result))},
                },
            )
        )
        raise InvalidAIActionOutput(result)

    log.info(
        **wrap_constants(
            message="AINode execution completed",
            **{
                LC.EVENT_TYPE: "node",
                LC.NODE_ID: self.node_id,
                LC.NODE_TYPE: "AINode",
                LC.ACTION: "execute_end",
                LC.OUTPUT_SIZE: len(result.messages),
                LC.SUCCESS: True,
            },
        )
    )

    return result