From 505a25d53d0e8d05c6ad95321d274319b2b3a0e7 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Fri, 27 Dec 2024 22:25:30 +0530 Subject: [PATCH] fix(agents-api): Fix tests for workflows Signed-off-by: Diwank Singh Tomer --- .../agents_api/activities/execute_system.py | 4 +- .../activities/task_steps/base_evaluate.py | 11 +----- .../activities/task_steps/evaluate_step.py | 11 +----- .../activities/task_steps/for_each_step.py | 11 +----- .../activities/task_steps/get_value_step.py | 11 +----- .../activities/task_steps/if_else_step.py | 11 +----- .../activities/task_steps/log_step.py | 9 +---- .../activities/task_steps/map_reduce_step.py | 11 +----- .../activities/task_steps/pg_query_step.py | 12 +----- .../activities/task_steps/return_step.py | 11 +----- .../activities/task_steps/set_value_step.py | 11 +----- .../activities/task_steps/switch_step.py | 9 +---- .../activities/task_steps/transition_step.py | 12 ++---- .../task_steps/wait_for_input_step.py | 9 +---- .../activities/task_steps/yield_step.py | 13 +------ agents-api/agents_api/app.py | 38 ++++++++++--------- agents-api/tests/fixtures.py | 11 ++++-- agents-api/tests/test_execution_workflow.py | 24 ------------ 18 files changed, 47 insertions(+), 182 deletions(-) diff --git a/agents-api/agents_api/activities/execute_system.py b/agents-api/agents_api/activities/execute_system.py index e8fdb06a8..3b2c4f58a 100644 --- a/agents-api/agents_api/activities/execute_system.py +++ b/agents-api/agents_api/activities/execute_system.py @@ -9,7 +9,7 @@ from fastapi.background import BackgroundTasks from temporalio import activity -from ..app import lifespan +from ..app import app, lifespan from ..autogen.openapi_model import ( ChatInput, CreateDocRequest, @@ -29,7 +29,7 @@ process_pool_executor = ProcessPoolExecutor() -@lifespan(container) +@lifespan(app, container) # Both are needed because we are using the routes @beartype async def execute_system( context: StepContext, diff --git a/agents-api/agents_api/activities/task_steps/base_evaluate.py b/agents-api/agents_api/activities/task_steps/base_evaluate.py index 3bb04e390..a23db0eaf 100644 --- a/agents-api/agents_api/activities/task_steps/base_evaluate.py +++ b/agents-api/agents_api/activities/task_steps/base_evaluate.py @@ -13,7 +13,6 @@ from temporalio import activity # noqa: E402 from thefuzz import fuzz # noqa: E402 -from ...env import testing # noqa: E402 from ..utils import get_evaluator # noqa: E402 @@ -62,6 +61,7 @@ def _recursive_evaluate(expr, evaluator: SimpleEval): raise ValueError(f"Invalid expression: {expr}") +@activity.defn @beartype async def base_evaluate( exprs: Any, @@ -100,12 +100,3 @@ async def base_evaluate( # Recursively evaluate the expression result = _recursive_evaluate(exprs, evaluator) return result - - -# Note: This is here just for clarity. We could have just imported base_evaluate directly -# They do the same thing, so we dont need to mock the base_evaluate function -mock_base_evaluate = base_evaluate - -base_evaluate = activity.defn(name="base_evaluate")( - base_evaluate if not testing else mock_base_evaluate -) diff --git a/agents-api/agents_api/activities/task_steps/evaluate_step.py b/agents-api/agents_api/activities/task_steps/evaluate_step.py index 08fa6cd55..6012f8d44 100644 --- a/agents-api/agents_api/activities/task_steps/evaluate_step.py +++ b/agents-api/agents_api/activities/task_steps/evaluate_step.py @@ -5,9 +5,9 @@ from ...activities.utils import simple_eval_dict from ...common.protocol.tasks import StepContext, StepOutcome -from ...env import testing +@activity.defn @beartype async def evaluate_step( context: StepContext, @@ -31,12 +31,3 @@ async def evaluate_step( except BaseException as e: activity.logger.error(f"Error in evaluate_step: {e}") return StepOutcome(error=str(e) or repr(e)) - - -# Note: This is here just for clarity. We could have just imported evaluate_step directly -# They do the same thing, so we dont need to mock the evaluate_step function -mock_evaluate_step = evaluate_step - -evaluate_step = activity.defn(name="evaluate_step")( - evaluate_step if not testing else mock_evaluate_step -) diff --git a/agents-api/agents_api/activities/task_steps/for_each_step.py b/agents-api/agents_api/activities/task_steps/for_each_step.py index ca84eb75d..4c8495ad3 100644 --- a/agents-api/agents_api/activities/task_steps/for_each_step.py +++ b/agents-api/agents_api/activities/task_steps/for_each_step.py @@ -6,10 +6,10 @@ StepContext, StepOutcome, ) -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def for_each_step(context: StepContext) -> StepOutcome: try: @@ -23,12 +23,3 @@ async def for_each_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in for_each_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported if_else_step directly -# They do the same thing, so we dont need to mock the if_else_step function -mock_if_else_step = for_each_step - -for_each_step = activity.defn(name="for_each_step")( - for_each_step if not testing else mock_if_else_step -) diff --git a/agents-api/agents_api/activities/task_steps/get_value_step.py b/agents-api/agents_api/activities/task_steps/get_value_step.py index feeb71bbf..f7f285115 100644 --- a/agents-api/agents_api/activities/task_steps/get_value_step.py +++ b/agents-api/agents_api/activities/task_steps/get_value_step.py @@ -2,24 +2,15 @@ from temporalio import activity from ...common.protocol.tasks import StepContext, StepOutcome -from ...env import testing # TODO: We should use this step to query the parent workflow and get the value from the workflow context # SCRUM-1 +@activity.defn @beartype async def get_value_step( context: StepContext, ) -> StepOutcome: key: str = context.current_step.get # noqa: F841 raise NotImplementedError("Not implemented yet") - - -# Note: This is here just for clarity. We could have just imported get_value_step directly -# They do the same thing, so we dont need to mock the get_value_step function -mock_get_value_step = get_value_step - -get_value_step = activity.defn(name="get_value_step")( - get_value_step if not testing else mock_get_value_step -) diff --git a/agents-api/agents_api/activities/task_steps/if_else_step.py b/agents-api/agents_api/activities/task_steps/if_else_step.py index ec4368640..d9997b492 100644 --- a/agents-api/agents_api/activities/task_steps/if_else_step.py +++ b/agents-api/agents_api/activities/task_steps/if_else_step.py @@ -6,10 +6,10 @@ StepContext, StepOutcome, ) -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def if_else_step(context: StepContext) -> StepOutcome: # NOTE: This activity is only for logging, so we just evaluate the expression @@ -27,12 +27,3 @@ async def if_else_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in if_else_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported if_else_step directly -# They do the same thing, so we dont need to mock the if_else_step function -mock_if_else_step = if_else_step - -if_else_step = activity.defn(name="if_else_step")( - if_else_step if not testing else mock_if_else_step -) diff --git a/agents-api/agents_api/activities/task_steps/log_step.py b/agents-api/agents_api/activities/task_steps/log_step.py index f54018683..c83fdca8f 100644 --- a/agents-api/agents_api/activities/task_steps/log_step.py +++ b/agents-api/agents_api/activities/task_steps/log_step.py @@ -7,9 +7,9 @@ StepOutcome, ) from ...common.utils.template import render_template -from ...env import testing +@activity.defn @beartype async def log_step(context: StepContext) -> StepOutcome: # NOTE: This activity is only for logging, so we just evaluate the expression @@ -30,10 +30,3 @@ async def log_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in log_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported log_step directly -# They do the same thing, so we dont need to mock the log_step function -mock_log_step = log_step - -log_step = activity.defn(name="log_step")(log_step if not testing else mock_log_step) diff --git a/agents-api/agents_api/activities/task_steps/map_reduce_step.py b/agents-api/agents_api/activities/task_steps/map_reduce_step.py index c39bace20..600f98615 100644 --- a/agents-api/agents_api/activities/task_steps/map_reduce_step.py +++ b/agents-api/agents_api/activities/task_steps/map_reduce_step.py @@ -8,10 +8,10 @@ StepContext, StepOutcome, ) -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def map_reduce_step(context: StepContext) -> StepOutcome: try: @@ -26,12 +26,3 @@ async def map_reduce_step(context: StepContext) -> StepOutcome: except BaseException as e: logging.error(f"Error in map_reduce_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported if_else_step directly -# They do the same thing, so we dont need to mock the if_else_step function -mock_if_else_step = map_reduce_step - -map_reduce_step = activity.defn(name="map_reduce_step")( - map_reduce_step if not testing else mock_if_else_step -) diff --git a/agents-api/agents_api/activities/task_steps/pg_query_step.py b/agents-api/agents_api/activities/task_steps/pg_query_step.py index cdbaa911c..2c081cb15 100644 --- a/agents-api/agents_api/activities/task_steps/pg_query_step.py +++ b/agents-api/agents_api/activities/task_steps/pg_query_step.py @@ -5,10 +5,11 @@ from ... import queries from ...app import lifespan -from ...env import pg_dsn, testing +from ...env import pg_dsn from ..container import container +@activity.defn @lifespan(container) @beartype async def pg_query_step( @@ -21,12 +22,3 @@ async def pg_query_step( module = getattr(queries, module_name) query = getattr(module, name) return await query(**values, connection_pool=container.state.postgres_pool) - - -# Note: This is here just for clarity. We could have just imported pg_query_step directly -# They do the same thing, so we dont need to mock the pg_query_step function -mock_pg_query_step = pg_query_step - -pg_query_step = activity.defn(name="pg_query_step")( - pg_query_step if not testing else mock_pg_query_step -) diff --git a/agents-api/agents_api/activities/task_steps/return_step.py b/agents-api/agents_api/activities/task_steps/return_step.py index f15354536..71b281f11 100644 --- a/agents-api/agents_api/activities/task_steps/return_step.py +++ b/agents-api/agents_api/activities/task_steps/return_step.py @@ -6,10 +6,10 @@ StepContext, StepOutcome, ) -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def return_step(context: StepContext) -> StepOutcome: try: @@ -24,12 +24,3 @@ async def return_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in log_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported return_step directly -# They do the same thing, so we dont need to mock the return_step function -mock_return_step = return_step - -return_step = activity.defn(name="return_step")( - return_step if not testing else mock_return_step -) diff --git a/agents-api/agents_api/activities/task_steps/set_value_step.py b/agents-api/agents_api/activities/task_steps/set_value_step.py index 96db5d0d1..a8ef06ce2 100644 --- a/agents-api/agents_api/activities/task_steps/set_value_step.py +++ b/agents-api/agents_api/activities/task_steps/set_value_step.py @@ -5,12 +5,12 @@ from ...activities.utils import simple_eval_dict from ...common.protocol.tasks import StepContext, StepOutcome -from ...env import testing # TODO: We should use this step to signal to the parent workflow and set the value on the workflow context # SCRUM-2 +@activity.defn @beartype async def set_value_step( context: StepContext, @@ -29,12 +29,3 @@ async def set_value_step( except BaseException as e: activity.logger.error(f"Error in set_value_step: {e}") return StepOutcome(error=str(e) or repr(e)) - - -# Note: This is here just for clarity. We could have just imported set_value_step directly -# They do the same thing, so we dont need to mock the set_value_step function -mock_set_value_step = set_value_step - -set_value_step = activity.defn(name="set_value_step")( - set_value_step if not testing else mock_set_value_step -) diff --git a/agents-api/agents_api/activities/task_steps/switch_step.py b/agents-api/agents_api/activities/task_steps/switch_step.py index 100d8020a..82c814bb1 100644 --- a/agents-api/agents_api/activities/task_steps/switch_step.py +++ b/agents-api/agents_api/activities/task_steps/switch_step.py @@ -6,10 +6,10 @@ StepContext, StepOutcome, ) -from ...env import testing from ..utils import get_evaluator +@activity.defn @beartype async def switch_step(context: StepContext) -> StepOutcome: try: @@ -34,10 +34,3 @@ async def switch_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in switch_step: {e}") return StepOutcome(error=str(e)) - - -mock_switch_step = switch_step - -switch_step = activity.defn(name="switch_step")( - switch_step if not testing else mock_switch_step -) diff --git a/agents-api/agents_api/activities/task_steps/transition_step.py b/agents-api/agents_api/activities/task_steps/transition_step.py index c44fa05d0..4b258b8bd 100644 --- a/agents-api/agents_api/activities/task_steps/transition_step.py +++ b/agents-api/agents_api/activities/task_steps/transition_step.py @@ -9,11 +9,7 @@ from ...autogen.openapi_model import CreateTransitionRequest, Transition from ...clients.temporal import get_workflow_handle from ...common.protocol.tasks import ExecutionInput, StepContext -from ...env import ( - temporal_activity_after_retry_timeout, - testing, - transition_requests_per_minute, -) +from ...env import temporal_activity_after_retry_timeout, transition_requests_per_minute from ...exceptions import LastErrorInput, TooManyRequestsError from ...queries.executions.create_execution_transition import ( create_execution_transition, @@ -74,9 +70,7 @@ async def transition_step( return transition +# NOTE: Here because needed by a different step original_transition_step = transition_step -mock_transition_step = transition_step -transition_step = activity.defn(name="transition_step")( - transition_step if not testing else mock_transition_step -) +transition_step = activity.defn(transition_step) diff --git a/agents-api/agents_api/activities/task_steps/wait_for_input_step.py b/agents-api/agents_api/activities/task_steps/wait_for_input_step.py index a3cb00f67..ac4bac9d6 100644 --- a/agents-api/agents_api/activities/task_steps/wait_for_input_step.py +++ b/agents-api/agents_api/activities/task_steps/wait_for_input_step.py @@ -3,10 +3,10 @@ from ...autogen.openapi_model import WaitForInputStep from ...common.protocol.tasks import StepContext, StepOutcome -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def wait_for_input_step(context: StepContext) -> StepOutcome: try: @@ -21,10 +21,3 @@ async def wait_for_input_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in wait_for_input_step: {e}") return StepOutcome(error=str(e)) - - -mock_wait_for_input_step = wait_for_input_step - -wait_for_input_step = activity.defn(name="wait_for_input_step")( - wait_for_input_step if not testing else mock_wait_for_input_step -) diff --git a/agents-api/agents_api/activities/task_steps/yield_step.py b/agents-api/agents_api/activities/task_steps/yield_step.py index 18e5383cc..2136da763 100644 --- a/agents-api/agents_api/activities/task_steps/yield_step.py +++ b/agents-api/agents_api/activities/task_steps/yield_step.py @@ -1,14 +1,12 @@ -from typing import Callable - from beartype import beartype from temporalio import activity from ...autogen.openapi_model import TransitionTarget, YieldStep from ...common.protocol.tasks import ExecutionInput, StepContext, StepOutcome -from ...env import testing from .base_evaluate import base_evaluate +@activity.defn @beartype async def yield_step(context: StepContext) -> StepOutcome: try: @@ -39,12 +37,3 @@ async def yield_step(context: StepContext) -> StepOutcome: except BaseException as e: activity.logger.error(f"Error in yield_step: {e}") return StepOutcome(error=str(e)) - - -# Note: This is here just for clarity. We could have just imported yield_step directly -# They do the same thing, so we dont need to mock the yield_step function -mock_yield_step: Callable[[StepContext], StepOutcome] = yield_step - -yield_step: Callable[[StepContext], StepOutcome] = activity.defn(name="yield_step")( - yield_step if not testing else mock_yield_step -) diff --git a/agents-api/agents_api/app.py b/agents-api/agents_api/app.py index 122de41b2..38582d85d 100644 --- a/agents-api/agents_api/app.py +++ b/agents-api/agents_api/app.py @@ -21,39 +21,43 @@ class ObjectWithState(Protocol): # TODO: This currently doesn't use .env variables, but we should move to using them @asynccontextmanager -async def lifespan(app: FastAPI | ObjectWithState): +async def lifespan(*containers: list[FastAPI | ObjectWithState]): # INIT POSTGRES # pg_dsn = os.environ.get("PG_DSN") - if not getattr(app.state, "postgres_pool", None): - app.state.postgres_pool = await create_db_pool(pg_dsn) + for container in containers: + if not getattr(container.state, "postgres_pool", None): + container.state.postgres_pool = await create_db_pool(pg_dsn) # INIT S3 # s3_access_key = os.environ.get("S3_ACCESS_KEY") s3_secret_key = os.environ.get("S3_SECRET_KEY") s3_endpoint = os.environ.get("S3_ENDPOINT") - if not getattr(app.state, "s3_client", None): - session = get_session() - app.state.s3_client = await session.create_client( - "s3", - aws_access_key_id=s3_access_key, - aws_secret_access_key=s3_secret_key, - endpoint_url=s3_endpoint, - ).__aenter__() + for container in containers: + if not getattr(container.state, "s3_client", None): + session = get_session() + container.state.s3_client = await session.create_client( + "s3", + aws_access_key_id=s3_access_key, + aws_secret_access_key=s3_secret_key, + endpoint_url=s3_endpoint, + ).__aenter__() try: yield finally: # CLOSE POSTGRES # - if getattr(app.state, "postgres_pool", None): - await app.state.postgres_pool.close() - app.state.postgres_pool = None + for container in containers: + if getattr(container.state, "postgres_pool", None): + await container.state.postgres_pool.close() + container.state.postgres_pool = None # CLOSE S3 # - if getattr(app.state, "s3_client", None): - await app.state.s3_client.close() - app.state.s3_client = None + for container in containers: + if getattr(container.state, "s3_client", None): + await container.state.s3_client.close() + container.state.s3_client = None app: FastAPI = FastAPI( diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py index 9a5bbb058..cb9e40a91 100644 --- a/agents-api/tests/fixtures.py +++ b/agents-api/tests/fixtures.py @@ -52,7 +52,12 @@ @fixture(scope="global") def pg_dsn(): with get_pg_dsn() as pg_dsn: - yield pg_dsn + os.environ["PG_DSN"] = pg_dsn + + try: + yield pg_dsn + finally: + del os.environ["PG_DSN"] @fixture(scope="global") @@ -376,9 +381,7 @@ async def test_tool( @fixture(scope="global") -def client(dsn=pg_dsn): - os.environ["PG_DSN"] = dsn - +def client(_dsn=pg_dsn): with TestClient(app=app) as client: yield client diff --git a/agents-api/tests/test_execution_workflow.py b/agents-api/tests/test_execution_workflow.py index 4a525e571..04f19d338 100644 --- a/agents-api/tests/test_execution_workflow.py +++ b/agents-api/tests/test_execution_workflow.py @@ -19,7 +19,6 @@ from agents_api.routers.tasks.create_task_execution import start_execution from .fixtures import ( - client, pg_dsn, s3_client, test_agent, @@ -75,7 +74,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -120,7 +118,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -165,7 +162,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -217,7 +213,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -270,7 +265,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -317,7 +311,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -371,7 +364,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -424,7 +416,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -480,7 +471,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={}) @@ -543,7 +533,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -605,7 +594,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -681,7 +669,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -738,7 +725,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -800,7 +786,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -867,7 +852,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -939,7 +923,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -990,7 +973,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -1050,7 +1032,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -1100,7 +1081,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -1152,7 +1132,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) data = CreateExecutionRequest(input={"test": "input"}) @@ -1208,7 +1187,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) mock_model_response = ModelResponse( @@ -1267,7 +1245,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) mock_model_response = ModelResponse( @@ -1331,7 +1308,6 @@ async def _( developer_id=test_developer_id, agent=test_agent, _s3_client=s3_client, # Adding coz blob store might be used - _app_client=client, ): pool = await create_db_pool(dsn=dsn) mock_model_response = ModelResponse(