From 382cf1a3eb91aa37c115350c77def3c964469f79 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Tue, 17 Sep 2024 21:39:22 -0400 Subject: [PATCH] fix(agents-api): fix create_agent and create_or_update_agent query --- .env.example | 1 + .../task_steps/raise_complete_async.py | 2 +- .../activities/task_steps/transition_step.py | 2 +- .../models/agent/create_or_update_agent.py | 26 +++++++++++++++++- .../session/create_or_update_session.py | 1 + .../models/session/create_session.py | 1 + .../models/session/prepare_session_data.py | 27 +++++++++++++++++++ .../agents_api/routers/agents/create_agent.py | 1 + .../routers/agents/create_or_update_agent.py | 1 + .../sessions/create_or_update_session.py | 11 +++----- .../workflows/task_execution/__init__.py | 6 ++--- 11 files changed, 66 insertions(+), 13 deletions(-) diff --git a/.env.example b/.env.example index fc2bdd7fb..e3cda1a97 100644 --- a/.env.example +++ b/.env.example @@ -6,6 +6,7 @@ COZO_AUTH_TOKEN= TEMPORAL_POSTGRES_PASSWORD= LITELLM_POSTGRES_PASSWORD= LITELLM_MASTER_KEY= +LITELLM_SALT_KEY= LITELLM_REDIS_PASSWORD= EMBEDDING_SERVICE_BASE=http://text-embeddings-inference- # Use the 'gpu' profile to run on GPU diff --git a/agents-api/agents_api/activities/task_steps/raise_complete_async.py b/agents-api/agents_api/activities/task_steps/raise_complete_async.py index 0748f08b3..3944c914f 100644 --- a/agents-api/agents_api/activities/task_steps/raise_complete_async.py +++ b/agents-api/agents_api/activities/task_steps/raise_complete_async.py @@ -13,7 +13,7 @@ async def raise_complete_async(context: StepContext, output: StepOutcome) -> Non # TODO: Create a transtition to "wait" and save the captured_token to the transition captured_token = activity.info().task_token - captured_token = captured_token.decode('latin-1') + captured_token = captured_token.decode("latin-1") transition_info = CreateTransitionRequest( current=context.cursor, type="wait", 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 cd3c09724..04ce6e8fd 100644 --- a/agents-api/agents_api/activities/task_steps/transition_step.py +++ b/agents-api/agents_api/activities/task_steps/transition_step.py @@ -29,4 +29,4 @@ async def transition_step( transition_step = activity.defn(name="transition_step")( transition_step if not testing else mock_transition_step -) \ No newline at end of file +) diff --git a/agents-api/agents_api/models/agent/create_or_update_agent.py b/agents-api/agents_api/models/agent/create_or_update_agent.py index 64b008d44..156c8ae61 100644 --- a/agents-api/agents_api/models/agent/create_or_update_agent.py +++ b/agents-api/agents_api/models/agent/create_or_update_agent.py @@ -80,6 +80,29 @@ def create_or_update_agent( } ) + # TODO: remove this + ### # Create default agent settings + ### # Construct a query to insert default settings for the new agent + ### default_settings_query = f""" + ### %if {{ + ### len[count(agent_id)] := + ### *agent_default_settings{{agent_id}}, + ### agent_id = to_uuid($agent_id) + + ### ?[should_create] := len[count], count > 0 + ### }} + ### %then {{ + ### ?[{settings_cols}] <- $settings_vals + + ### :put agent_default_settings {{ + ### {settings_cols} + ### }} + ### }} + ### """ + + # FIXME: This create or update query will overwrite the settings + # Need to find a way to only run the insert query if the agent_default_settings + # Create default agent settings # Construct a query to insert default settings for the new agent default_settings_query = f""" @@ -89,6 +112,7 @@ def create_or_update_agent( {settings_cols} }} """ + # create the agent # Construct a query to insert the new agent record into the agents table agent_query = """ @@ -127,7 +151,7 @@ def create_or_update_agent( queries = [ verify_developer_id_query(developer_id), - default_settings_query if default_settings else None, + default_settings_query, agent_query, ] diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py index 60c0b7f71..6c81b2563 100644 --- a/agents-api/agents_api/models/session/create_or_update_session.py +++ b/agents-api/agents_api/models/session/create_or_update_session.py @@ -26,6 +26,7 @@ @rewrap_exceptions( { + AssertionError: partialclass(HTTPException, status_code=400), QueryException: partialclass(HTTPException, status_code=400), ValidationError: partialclass(HTTPException, status_code=400), TypeError: partialclass(HTTPException, status_code=400), diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py index a83837ffd..32a74533a 100644 --- a/agents-api/agents_api/models/session/create_session.py +++ b/agents-api/agents_api/models/session/create_session.py @@ -27,6 +27,7 @@ @rewrap_exceptions( { + AssertionError: partialclass(HTTPException, status_code=400), QueryException: partialclass(HTTPException, status_code=400), ValidationError: partialclass(HTTPException, status_code=400), TypeError: partialclass(HTTPException, status_code=400), diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py index 49cdc8cbc..e8cfb7fc7 100644 --- a/agents-api/agents_api/models/session/prepare_session_data.py +++ b/agents-api/agents_api/models/session/prepare_session_data.py @@ -127,6 +127,33 @@ def prepare_session_data( "instructions": instructions, } + # Version where we don't have default settings + agent_data[collect(record)] := + agents[agent_ids], + agent_id in agent_ids, + *agents{ + agent_id, + model, + name, + about, + created_at, + updated_at, + metadata, + instructions, + }, + not settings_data[agent_id, _], + record = { + "id": agent_id, + "name": name, + "model": model, + "about": about, + "created_at": created_at, + "updated_at": updated_at, + "metadata": metadata, + "default_settings": {}, + "instructions": instructions, + } + user_data[collect(record)] := users[user_ids], user_id in user_ids, diff --git a/agents-api/agents_api/routers/agents/create_agent.py b/agents-api/agents_api/routers/agents/create_agent.py index d1cac0d6b..96e12e2f2 100644 --- a/agents-api/agents_api/routers/agents/create_agent.py +++ b/agents-api/agents_api/routers/agents/create_agent.py @@ -19,6 +19,7 @@ async def create_agent( x_developer_id: Annotated[UUID4, Depends(get_developer_id)], data: CreateAgentRequest, ) -> ResourceCreatedResponse: + # TODO: Validate model name agent = models.agent.create_agent( developer_id=x_developer_id, data=data, diff --git a/agents-api/agents_api/routers/agents/create_or_update_agent.py b/agents-api/agents_api/routers/agents/create_or_update_agent.py index fc2fa5563..f91e8aa10 100644 --- a/agents-api/agents_api/routers/agents/create_or_update_agent.py +++ b/agents-api/agents_api/routers/agents/create_or_update_agent.py @@ -21,6 +21,7 @@ async def create_or_update_agent( data: CreateOrUpdateAgentRequest, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> ResourceCreatedResponse: + # TODO: Validate model name agent = models.agent.create_or_update_agent( developer_id=x_developer_id, agent_id=agent_id, diff --git a/agents-api/agents_api/routers/sessions/create_or_update_session.py b/agents-api/agents_api/routers/sessions/create_or_update_session.py index 8ed9ff7ba..37e928a16 100644 --- a/agents-api/agents_api/routers/sessions/create_or_update_session.py +++ b/agents-api/agents_api/routers/sessions/create_or_update_session.py @@ -7,7 +7,7 @@ from ...autogen.openapi_model import ( CreateOrUpdateSessionRequest, - ResourceCreatedResponse, + ResourceUpdatedResponse, ) from ...dependencies.developer_id import get_developer_id from ...models.session.create_or_update_session import ( @@ -21,14 +21,11 @@ async def create_or_update_session( x_developer_id: Annotated[UUID4, Depends(get_developer_id)], session_id: UUID, data: CreateOrUpdateSessionRequest, -) -> ResourceCreatedResponse: - session = create_session_query( +) -> ResourceUpdatedResponse: + session_updated = create_session_query( developer_id=x_developer_id, session_id=session_id, data=data, ) - return ResourceCreatedResponse( - id=session.id, - created_at=session.created_at, - ) + return session_updated diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py index b5f0b0a29..2ff45e55c 100644 --- a/agents-api/agents_api/workflows/task_execution/__init__.py +++ b/agents-api/agents_api/workflows/task_execution/__init__.py @@ -416,10 +416,10 @@ async def run( task_steps.prompt_step, context, schedule_to_close_timeout=timedelta( - seconds=30 if debug or testing else 600), + seconds=30 if debug or testing else 600 + ), ) - state = PartialTransition( - output=new_response.output, type="resume") + state = PartialTransition(output=new_response.output, type="resume") # case PromptStep(), StepOutcome( # output=response