Multipage workflow improvements (#875)

This commit is contained in:
Kerem Yilmaz
2024-09-24 08:51:51 -07:00
committed by GitHub
parent c27c2a55f9
commit d67aaeffb0
3 changed files with 48 additions and 12 deletions

View File

@@ -1264,7 +1264,7 @@ class AgentDB:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def create_workflow_run_output_parameter(
async def create_or_update_workflow_run_output_parameter(
self,
workflow_run_id: str,
output_parameter_id: str,
@@ -1272,6 +1272,24 @@ class AgentDB:
) -> WorkflowRunOutputParameter:
try:
async with self.Session() as session:
# check if the workflow run output parameter already exists
# if it does, update the value
if workflow_run_output_parameter := (
await session.scalars(
select(WorkflowRunOutputParameterModel)
.filter_by(workflow_run_id=workflow_run_id)
.filter_by(output_parameter_id=output_parameter_id)
)
).first():
LOG.info(
f"Updating existing workflow run output parameter with {workflow_run_output_parameter.workflow_run_id} - {workflow_run_output_parameter.output_parameter_id}"
)
workflow_run_output_parameter.value = value
await session.commit()
await session.refresh(workflow_run_output_parameter)
return convert_to_workflow_run_output_parameter(workflow_run_output_parameter, self.debug_enabled)
# if it does not exist, create a new one
workflow_run_output_parameter = WorkflowRunOutputParameterModel(
workflow_run_id=workflow_run_id,
output_parameter_id=output_parameter_id,
@@ -1286,6 +1304,33 @@ class AgentDB:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def update_workflow_run_output_parameter(
self,
workflow_run_id: str,
output_parameter_id: str,
value: dict[str, Any] | list | str | None,
) -> WorkflowRunOutputParameter:
try:
async with self.Session() as session:
workflow_run_output_parameter = (
await session.scalars(
select(WorkflowRunOutputParameterModel)
.filter_by(workflow_run_id=workflow_run_id)
.filter_by(output_parameter_id=output_parameter_id)
)
).first()
if not workflow_run_output_parameter:
raise NotFoundError(
f"WorkflowRunOutputParameter not found for {workflow_run_id} and {output_parameter_id}"
)
workflow_run_output_parameter.value = value
await session.commit()
await session.refresh(workflow_run_output_parameter)
return convert_to_workflow_run_output_parameter(workflow_run_output_parameter, self.debug_enabled)
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def get_workflow_parameters(self, workflow_id: str) -> list[WorkflowParameter]:
try:
async with self.Session() as session: