Skip to content

AIAction

graphorchestrator.ai.ai_action

AIActionBase

Bases: ABC

Abstract base class for defining AI actions within the graph orchestration framework.

This class serves as a template for creating specialized nodes that can interact with and modify the state of a graph using an underlying AI model.

Attributes:

Name Type Description
config dict

A dictionary containing configuration parameters for the AI model.
_model_built bool

A flag indicating whether the AI model has been built.
model Any

The AI model instance.
is_node_action bool

A flag indicating that this object is an AI node action.
__name__ str

The name of the AI action, which defaults to the class name.
Source code in graphorchestrator\ai\ai_action.py 
 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
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class AIActionBase(ABC):
    """Abstract base class for defining AI actions within the graph orchestration framework.

    This class serves as a template for creating specialized nodes that can interact
    with and modify the state of a graph using an underlying AI model.

    Attributes:
        config (dict): A dictionary containing configuration parameters for the AI model.
        _model_built (bool): A flag indicating whether the AI model has been built.
        model (Any): The AI model instance.
        is_node_action (bool): A flag indicating that this object is an AI node action.
        __name__ (str): The name of the AI action, which defaults to the class name.
    """

    def __init__(self, config: dict) -> None:
        """
        Initializes an AIActionBase instance.

        Args:
            config (dict): Configuration parameters for the AI model.
        """
        # Store the provided configuration
        self.config: dict = config
        # Initialize flags and model to default states
        self._model_built = False
        self.model: Any = None
        # Identify the instance as an AI node action
        self.is_node_action = True
        # Set the name to the class name
        self.__name__ = self.__class__.__name__

        # Get the logger instance
        log = GraphLogger.get()
        # Use the class name as the node label
        node_label = self.__name__

        # Log the initialization of the AIActionBase instance
        log.info(
            **wrap_constants(
                message="AIActionBase initialized",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.ACTION: "node_created",  # Action being taken: node creation
                    LC.NODE_ID: node_label,
                    LC.NODE_TYPE: "AINode",
                    LC.CUSTOM: {"config": self.config},  # Full config included here
                }
            )
        )

    @abstractmethod
    def build_model(self) -> None:
        """Build and configure the AI model.

        This method should be implemented by subclasses to handle the creation and
        configuration of the AI model. It is typically called before processing any state.

        Subclasses must set the following attributes:
            self.model: The constructed AI model instance.
            self._model_built: Set to True to indicate the model is built.

        Raises:
            NotImplementedError: If the method is not implemented by a subclass.
        """
        raise NotImplementedError

    @abstractmethod
    async def process_state(self, state: State) -> State:
        """Process the state using the AI model.

        This is the main method where the AI model logic is applied to the current state.
        Subclasses must implement this method to define how the AI model modifies the state.

        Args:
            state (State): The current state to be processed.

        Returns:
            State: The new state after processing.
        """
        raise NotImplementedError

    async def __call__(self, state: State) -> State:
        """
        Invokes the AI action's processing logic.
        """
        # Get the logger instance and the node label
        log = GraphLogger.get()
        node_label = getattr(self, "__name__", self.__class__.__name__)

        # If the model has not been built yet, build it
        if not self._model_built:
            self.build_model()

            # Log that the AI model has been built
            log.info(
                **wrap_constants(
                    message="AI model built",
                    **{
                        LC.EVENT_TYPE: "node",
                        LC.NODE_ID: node_label,
                        LC.NODE_TYPE: "AINode",
                        LC.ACTION: "build_model",
                        LC.CUSTOM: {"config": self.config},
                    }
                )
            )

        # Log the start of the AI node execution
        log.info(
            **wrap_constants(
                message="AI node execution started",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: node_label,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "execute_start",
                    LC.INPUT_SIZE: len(state.messages),
                }
            )
        )

        # Process the state using the AI model
        result_or_coro = self.process_state(state)
        # Handle coroutines if needed
        result = (
            await result_or_coro
            if isinstance(result_or_coro, abc.Awaitable)
            else result_or_coro
        )

        # Validate the output
        # Ensure that the output is a State object; if not, raise an error
        if not isinstance(result, State):
            log.error(
                **wrap_constants(
                    message="AI action returned non-State object",
                    **{
                        LC.EVENT_TYPE: "node",
                        LC.NODE_ID: node_label,
                        LC.NODE_TYPE: "AINode",
                        LC.ACTION: "invalid_output",
                        LC.CUSTOM: {"actual_type": str(type(result))},
                    }
                )
            )
            raise InvalidAIActionOutput(result)

        # Log the completion of the AI node execution
        log.info(
            **wrap_constants(
                message="AI node execution completed",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: node_label,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "execute_end",
                    LC.OUTPUT_SIZE: len(result.messages),
                    LC.SUCCESS: True,
                }
            )
        )

        # Return the processed state
        return result

__call__(state) async

Invokes the AI action's processing logic.

Source code in graphorchestrator\ai\ai_action.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
async def __call__(self, state: State) -> State:
    """
    Invokes the AI action's processing logic.
    """
    # Get the logger instance and the node label
    log = GraphLogger.get()
    node_label = getattr(self, "__name__", self.__class__.__name__)

    # If the model has not been built yet, build it
    if not self._model_built:
        self.build_model()

        # Log that the AI model has been built
        log.info(
            **wrap_constants(
                message="AI model built",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: node_label,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "build_model",
                    LC.CUSTOM: {"config": self.config},
                }
            )
        )

    # Log the start of the AI node execution
    log.info(
        **wrap_constants(
            message="AI node execution started",
            **{
                LC.EVENT_TYPE: "node",
                LC.NODE_ID: node_label,
                LC.NODE_TYPE: "AINode",
                LC.ACTION: "execute_start",
                LC.INPUT_SIZE: len(state.messages),
            }
        )
    )

    # Process the state using the AI model
    result_or_coro = self.process_state(state)
    # Handle coroutines if needed
    result = (
        await result_or_coro
        if isinstance(result_or_coro, abc.Awaitable)
        else result_or_coro
    )

    # Validate the output
    # Ensure that the output is a State object; if not, raise an error
    if not isinstance(result, State):
        log.error(
            **wrap_constants(
                message="AI action returned non-State object",
                **{
                    LC.EVENT_TYPE: "node",
                    LC.NODE_ID: node_label,
                    LC.NODE_TYPE: "AINode",
                    LC.ACTION: "invalid_output",
                    LC.CUSTOM: {"actual_type": str(type(result))},
                }
            )
        )
        raise InvalidAIActionOutput(result)

    # Log the completion of the AI node execution
    log.info(
        **wrap_constants(
            message="AI node execution completed",
            **{
                LC.EVENT_TYPE: "node",
                LC.NODE_ID: node_label,
                LC.NODE_TYPE: "AINode",
                LC.ACTION: "execute_end",
                LC.OUTPUT_SIZE: len(result.messages),
                LC.SUCCESS: True,
            }
        )
    )

    # Return the processed state
    return result

__init__(config)

Initializes an AIActionBase instance.

Parameters:

Name Type Description Default
config dict

Configuration parameters for the AI model.

 required
Source code in graphorchestrator\ai\ai_action.py 
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
def __init__(self, config: dict) -> None:
    """
    Initializes an AIActionBase instance.

    Args:
        config (dict): Configuration parameters for the AI model.
    """
    # Store the provided configuration
    self.config: dict = config
    # Initialize flags and model to default states
    self._model_built = False
    self.model: Any = None
    # Identify the instance as an AI node action
    self.is_node_action = True
    # Set the name to the class name
    self.__name__ = self.__class__.__name__

    # Get the logger instance
    log = GraphLogger.get()
    # Use the class name as the node label
    node_label = self.__name__

    # Log the initialization of the AIActionBase instance
    log.info(
        **wrap_constants(
            message="AIActionBase initialized",
            **{
                LC.EVENT_TYPE: "node",
                LC.ACTION: "node_created",  # Action being taken: node creation
                LC.NODE_ID: node_label,
                LC.NODE_TYPE: "AINode",
                LC.CUSTOM: {"config": self.config},  # Full config included here
            }
        )
    )

build_model() abstractmethod

Build and configure the AI model.

This method should be implemented by subclasses to handle the creation and configuration of the AI model. It is typically called before processing any state.

Subclasses must set the following attributes

self.model: The constructed AI model instance. self._model_built: Set to True to indicate the model is built.

Raises:

Type Description
NotImplementedError

If the method is not implemented by a subclass.
Source code in graphorchestrator\ai\ai_action.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@abstractmethod
def build_model(self) -> None:
    """Build and configure the AI model.

    This method should be implemented by subclasses to handle the creation and
    configuration of the AI model. It is typically called before processing any state.

    Subclasses must set the following attributes:
        self.model: The constructed AI model instance.
        self._model_built: Set to True to indicate the model is built.

    Raises:
        NotImplementedError: If the method is not implemented by a subclass.
    """
    raise NotImplementedError

process_state(state) abstractmethod async

Process the state using the AI model.

This is the main method where the AI model logic is applied to the current state. Subclasses must implement this method to define how the AI model modifies the state.

Parameters:

Name Type Description Default
state State

The current state to be processed.

 required

Returns:

Name Type Description
State State

The new state after processing.
Source code in graphorchestrator\ai\ai_action.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@abstractmethod
async def process_state(self, state: State) -> State:
    """Process the state using the AI model.

    This is the main method where the AI model logic is applied to the current state.
    Subclasses must implement this method to define how the AI model modifies the state.

    Args:
        state (State): The current state to be processed.

    Returns:
        State: The new state after processing.
    """
    raise NotImplementedError