diff --git a/skyvern/forge/prompts/skyvern/task_v2_summary.j2 b/skyvern/forge/prompts/skyvern/task_v2_summary.j2 index 4b54e7c2..239187de 100644 --- a/skyvern/forge/prompts/skyvern/task_v2_summary.j2 +++ b/skyvern/forge/prompts/skyvern/task_v2_summary.j2 @@ -13,7 +13,10 @@ User goal: {{ user_goal }} ```{% if extracted_information_schema %} -The ideal schema for the "output" field: {{ extracted_information_schema }}{% endif %} +The ideal schema for the "output" field: +``` +{{ extracted_information_schema }} +```{% endif %} Task history (the earliest task is the first in the list and the latest is the last in the list): ``` diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index b07661ce..e4747ead 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -2251,6 +2251,8 @@ class AgentDB: totp_identifier: str | None = None, totp_verification_url: str | None = None, webhook_callback_url: str | None = None, + extracted_information_schema: dict | list | str | None = None, + error_code_mapping: dict | None = None, ) -> TaskV2: async with self.Session() as session: new_task_v2 = TaskV2Model( @@ -2263,6 +2265,8 @@ class AgentDB: totp_identifier=totp_identifier, totp_verification_url=totp_verification_url, webhook_callback_url=webhook_callback_url, + extracted_information_schema=extracted_information_schema, + error_code_mapping=error_code_mapping, organization_id=organization_id, ) session.add(new_task_v2) diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index a053b926..415d41a9 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -1133,6 +1133,8 @@ async def run_task_v2( proxy_location=data.proxy_location, publish_workflow=data.publish_workflow, create_task_run=True, + extracted_information_schema=data.extracted_information_schema, + error_code_mapping=data.error_code_mapping, ) except LLMProviderError: LOG.error("LLM failure to initialize task v2", exc_info=True) diff --git a/skyvern/forge/sdk/services/task_v2_service.py b/skyvern/forge/sdk/services/task_v2_service.py index cf008e49..27f9a05b 100644 --- a/skyvern/forge/sdk/services/task_v2_service.py +++ b/skyvern/forge/sdk/services/task_v2_service.py @@ -96,6 +96,8 @@ async def initialize_task_v2( publish_workflow: bool = False, parent_workflow_run_id: str | None = None, create_task_run: bool = False, + extracted_information_schema: dict | list | str | None = None, + error_code_mapping: dict | None = None, ) -> TaskV2: task_v2 = await app.DATABASE.create_task_v2( prompt=user_prompt, @@ -104,6 +106,8 @@ async def initialize_task_v2( totp_identifier=totp_identifier, webhook_callback_url=webhook_callback_url, proxy_location=proxy_location, + extracted_information_schema=extracted_information_schema, + error_code_mapping=error_code_mapping, ) # set task_v2_id in context context = skyvern_context.current()