diff --git a/autogpts/forge/forge/sdk/db.py b/autogpts/forge/forge/sdk/db.py index 1a3d98fd..0b35139b 100644 --- a/autogpts/forge/forge/sdk/db.py +++ b/autogpts/forge/forge/sdk/db.py @@ -54,6 +54,7 @@ class StepModel(Base): name = Column(String) input = Column(String) status = Column(String) + output = Column(String) is_last = Column(Boolean, default=False) created_at = Column(DateTime, default=datetime.datetime.utcnow) modified_at = Column( @@ -61,6 +62,7 @@ class StepModel(Base): ) additional_input = Column(JSON) + additional_output = Column(JSON) artifacts = relationship("ArtifactModel", back_populates="step") @@ -111,9 +113,11 @@ def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step: name=step_model.name, input=step_model.input, status=status, + output=step_model.output, artifacts=step_artifacts, is_last=step_model.is_last == 1, additional_input=step_model.additional_input, + additional_output=step_model.additional_output, ) @@ -280,7 +284,7 @@ class AgentDB: LOG.error(f"Unexpected error while getting task: {e}") raise - async def get_step(self, task_id: int, step_id: int) -> Step: + async def get_step(self, task_id: str, step_id: str) -> Step: if self.debug_enabled: LOG.debug(f"Getting step with task_id: {task_id} and step_id: {step_id}") try: @@ -311,8 +315,10 @@ class AgentDB: self, task_id: str, step_id: str, - status: str, - additional_input: Optional[Dict[str, Any]] = {}, + status: Optional[str] = None, + output: Optional[str] = None, + additional_input: Optional[Dict[str, Any]] = None, + additional_output: Optional[Dict[str, Any]] = None, ) -> Step: if self.debug_enabled: LOG.debug(f"Updating step with task_id: {task_id} and step_id: {step_id}") @@ -323,8 +329,14 @@ class AgentDB: .filter_by(task_id=task_id, step_id=step_id) .first() ): - step.status = status - step.additional_input = additional_input + if status is not None: + step.status = status + if additional_input is not None: + step.additional_input = additional_input + if output is not None: + step.output = output + if additional_output is not None: + step.additional_output = additional_output session.commit() return await self.get_step(task_id, step_id) else: diff --git a/autogpts/forge/forge/sdk/schema.py b/autogpts/forge/forge/sdk/schema.py index de9919fa..1dc2e25e 100644 --- a/autogpts/forge/forge/sdk/schema.py +++ b/autogpts/forge/forge/sdk/schema.py @@ -156,7 +156,7 @@ class Step(StepRequestBody): description="Output of the task step.", example="I am going to use the write_to_file command and write Washington to a file called output.txt