|
|
from enum import Enum |
|
|
from pydantic import Field |
|
|
from typing import Union, Optional, List |
|
|
from ..core.module import BaseModule |
|
|
from ..core.message import Message, MessageType |
|
|
from ..models.base_model import LLMOutputParser |
|
|
|
|
|
|
|
|
class TrajectoryState(str, Enum): |
|
|
""" |
|
|
Enum representing the status of a trajectory step. |
|
|
""" |
|
|
COMPLETED = "COMPLETED" |
|
|
FAILED = "FAILED" |
|
|
|
|
|
|
|
|
class TrajectoryStep(BaseModule): |
|
|
|
|
|
message: Message = None |
|
|
status: TrajectoryState |
|
|
error: Optional[str] = None |
|
|
|
|
|
|
|
|
class Environment(BaseModule): |
|
|
|
|
|
""" |
|
|
Responsible for storing and managing intermediate states of execution. |
|
|
""" |
|
|
trajectory: List[TrajectoryStep] = Field(default_factory=list) |
|
|
task_execution_history: List[str] = Field(default_factory=list) |
|
|
execution_data: dict = Field(default_factory=dict) |
|
|
|
|
|
def update(self, message: Message, state: TrajectoryState = None, error: str = None, **kwargs): |
|
|
""" |
|
|
Add a message to the shared memory and optionally to a specific task's message list. |
|
|
|
|
|
Args: |
|
|
message (Message): The message to be added. |
|
|
task_name (str, optional): The name of the task this message is related to. If None, the message is considered global. |
|
|
""" |
|
|
state = state or TrajectoryState.COMPLETED |
|
|
step = TrajectoryStep(message=message, status=state, error=error) |
|
|
self.trajectory.append(step) |
|
|
self.update_task_execution_history(message=message) |
|
|
self.update_execution_data(message=message) |
|
|
|
|
|
def update_task_execution_history(self, message: Message): |
|
|
if message.wf_task is not None and message.msg_type in [MessageType.RESPONSE]: |
|
|
|
|
|
if not self.task_execution_history or message.wf_task != self.task_execution_history[-1]: |
|
|
self.task_execution_history.append(message.wf_task) |
|
|
|
|
|
def update_execution_data(self, message: Message): |
|
|
if isinstance(message.content, LLMOutputParser): |
|
|
data = message.content.get_structured_data() |
|
|
self.execution_data.update(data) |
|
|
if isinstance(message.content, dict): |
|
|
data = message.content |
|
|
self.execution_data.update(data) |
|
|
|
|
|
def update_execution_data_from_context_extraction(self, extracted_data: dict): |
|
|
for key, value in extracted_data.items(): |
|
|
if key not in self.execution_data: |
|
|
self.execution_data[key] = value |
|
|
|
|
|
def get_task_messages(self, tasks: Union[str, List[str]], n: int = None, include_inputs: bool = False, **kwargs) -> List[Message]: |
|
|
""" |
|
|
Retrieve all messages related to specified tasks |
|
|
|
|
|
Returns: |
|
|
List[Message]: A list of messages related to the task. |
|
|
""" |
|
|
if isinstance(tasks, str): |
|
|
tasks = [tasks] |
|
|
message_list = [] |
|
|
for step in self.trajectory: |
|
|
message = step.message |
|
|
if message.wf_task is not None and message.wf_task in tasks: |
|
|
message_list.append(message) |
|
|
if include_inputs and message.msg_type == MessageType.INPUT and message not in message_list: |
|
|
message_list.append(message) |
|
|
message_list = message_list if n is None else message_list[-n:] |
|
|
return message_list |
|
|
|
|
|
def get(self, n: int=None) -> List[Message]: |
|
|
""" |
|
|
return the most recent n messages |
|
|
""" |
|
|
assert n is None or n>=0, "n must be None or a positive int" |
|
|
all_messages = [step.message for step in self.trajectory] |
|
|
messages = all_messages if n is None else all_messages[-n:] |
|
|
return messages |
|
|
|
|
|
def get_last_executed_task(self) -> str: |
|
|
if self.task_execution_history: |
|
|
return self.task_execution_history[-1] |
|
|
return None |
|
|
|
|
|
def get_all_execution_data(self) -> dict: |
|
|
return self.execution_data |
|
|
|
|
|
def get_execution_data(self, params: Union[str, List[str]]) -> dict: |
|
|
if isinstance(params, str): |
|
|
params = [params] |
|
|
data = {} |
|
|
for param in params: |
|
|
if param not in self.execution_data: |
|
|
raise KeyError(f"Couldn't find execution data with key '{param}'. Available execution data: {list(self.execution_data.keys())}") |
|
|
data[param] = self.execution_data[param] |
|
|
return data |
|
|
|
|
|
|
