From 2eeb2a3a97ebc8f19cd0c930f627df6c6925332c Mon Sep 17 00:00:00 2001 From: ykeremy Date: Mon, 25 Mar 2024 07:50:12 +0000 Subject: [PATCH 1/3] =?UTF-8?q?=F0=9F=94=84=20synced=20local=20'skyvern/'?= =?UTF-8?q?=20with=20remote=20'skyvern/'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- skyvern/forge/sdk/db/client.py | 4 +- skyvern/forge/sdk/prompting.py | 18 ++++++ skyvern/forge/sdk/workflow/models/block.py | 74 +++++++++++++++++++++- skyvern/forge/sdk/workflow/models/yaml.py | 15 ++++- skyvern/forge/sdk/workflow/service.py | 20 +++++- 5 files changed, 127 insertions(+), 4 deletions(-) diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index a8128b424..645004f38 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -692,7 +692,9 @@ async def update_workflow( ) -> Workflow: try: async with self.Session() as session: - if workflow := await session.scalars(select(WorkflowModel).filter_by(workflow_id=workflow_id).first()): + if workflow := ( + await session.scalars(select(WorkflowModel).filter_by(workflow_id=workflow_id)) + ).first(): if title: workflow.title = title if description: diff --git a/skyvern/forge/sdk/prompting.py b/skyvern/forge/sdk/prompting.py index a23f8c67a..9aa9f947a 100644 --- a/skyvern/forge/sdk/prompting.py +++ b/skyvern/forge/sdk/prompting.py @@ -96,3 +96,21 @@ def load_prompt(self, template: str, **kwargs: Any) -> str: except Exception: LOG.error("Failed to load prompt.", template=template, kwargs_keys=kwargs.keys(), exc_info=True) raise + + def load_prompt_from_string(self, template: str, **kwargs: Any) -> str: + """ + Load and populate the specified template from a string. + + Args: + template (str): The template string to load. + **kwargs: The arguments to populate the template with. + + Returns: + str: The populated template. + """ + try: + jinja_template = self.env.from_string(template) + return jinja_template.render(**kwargs) + except Exception: + LOG.error("Failed to load prompt from string.", template=template, kwargs_keys=kwargs.keys(), exc_info=True) + raise diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 00ba04ef9..0a09ccddb 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -1,4 +1,5 @@ import abc +import json from enum import StrEnum from typing import Annotated, Any, Literal, Union @@ -12,6 +13,8 @@ UnexpectedTaskStatus, ) from skyvern.forge import app +from skyvern.forge.prompts import prompt_engine +from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory from skyvern.forge.sdk.schemas.tasks import TaskStatus from skyvern.forge.sdk.workflow.context_manager import WorkflowRunContext from skyvern.forge.sdk.workflow.models.parameter import ( @@ -28,6 +31,7 @@ class BlockType(StrEnum): TASK = "task" FOR_LOOP = "for_loop" CODE = "code" + TEXT_PROMPT = "text_prompt" class Block(BaseModel, abc.ABC): @@ -345,5 +349,73 @@ async def execute(self, workflow_run_id: str, **kwargs: dict) -> OutputParameter return None -BlockSubclasses = Union[ForLoopBlock, TaskBlock, CodeBlock] +class TextPromptBlock(Block): + block_type: Literal[BlockType.TEXT_PROMPT] = BlockType.TEXT_PROMPT + + llm_key: str + prompt: str + parameters: list[PARAMETER_TYPE] = [] + json_schema: dict[str, Any] | None = None + + def get_all_parameters( + self, + ) -> list[PARAMETER_TYPE]: + return self.parameters + + async def send_prompt(self, prompt: str, parameter_values: dict[str, Any]) -> dict[str, Any]: + llm_api_handler = LLMAPIHandlerFactory.get_llm_api_handler(self.llm_key) + if not self.json_schema: + self.json_schema = { + "type": "object", + "properties": { + "llm_response": { + "type": "string", + "description": "Your response to the prompt", + } + }, + } + + prompt = prompt_engine.load_prompt_from_string(prompt, **parameter_values) + prompt += ( + "\n\n" + + "Please respond to the prompt above using the following JSON definition:\n\n" + + "```json\n" + + json.dumps(self.json_schema, indent=2) + + "\n```\n\n" + ) + LOG.info("TextPromptBlock: Sending prompt to LLM", prompt=prompt, llm_key=self.llm_key) + response = await llm_api_handler(prompt=prompt) + LOG.info("TextPromptBlock: Received response from LLM", response=response) + return response + + async def execute(self, workflow_run_id: str, **kwargs: dict) -> OutputParameter | None: + # get workflow run context + workflow_run_context = self.get_workflow_run_context(workflow_run_id) + # get all parameters into a dictionary + parameter_values = {} + for parameter in self.parameters: + value = workflow_run_context.get_value(parameter.key) + secret_value = workflow_run_context.get_original_secret_value_or_none(value) + if secret_value is not None: + parameter_values[parameter.key] = secret_value + else: + parameter_values[parameter.key] = value + + response = await self.send_prompt(self.prompt, parameter_values) + if self.output_parameter: + await workflow_run_context.register_output_parameter_value_post_execution( + parameter=self.output_parameter, + value=response, + ) + await app.DATABASE.create_workflow_run_output_parameter( + workflow_run_id=workflow_run_id, + output_parameter_id=self.output_parameter.output_parameter_id, + value=response, + ) + return self.output_parameter + + return None + + +BlockSubclasses = Union[ForLoopBlock, TaskBlock, CodeBlock, TextPromptBlock] BlockTypeVar = Annotated[BlockSubclasses, Field(discriminator="block_type")] diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py index df5f35e95..fefb9ad3c 100644 --- a/skyvern/forge/sdk/workflow/models/yaml.py +++ b/skyvern/forge/sdk/workflow/models/yaml.py @@ -94,10 +94,23 @@ class CodeBlockYAML(BlockYAML): parameter_keys: list[str] | None = None +class TextPromptBlockYAML(BlockYAML): + # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error: + # Parameter 1 of Literal[...] cannot be of type "Any" + # This pattern already works in block.py but since the BlockType is not defined in this file, mypy is not able + # to infer the type of the parameter_type attribute. + block_type: Literal[BlockType.TEXT_PROMPT] = BlockType.TEXT_PROMPT # type: ignore + + llm_key: str + prompt: str + parameter_keys: list[str] | None = None + json_schema: dict[str, Any] | None = None + + PARAMETER_YAML_SUBCLASSES = AWSSecretParameterYAML | WorkflowParameterYAML | ContextParameterYAML | OutputParameterYAML PARAMETER_YAML_TYPES = Annotated[PARAMETER_YAML_SUBCLASSES, Field(discriminator="parameter_type")] -BLOCK_YAML_SUBCLASSES = TaskBlockYAML | ForLoopBlockYAML | CodeBlockYAML +BLOCK_YAML_SUBCLASSES = TaskBlockYAML | ForLoopBlockYAML | CodeBlockYAML | TextPromptBlockYAML BLOCK_YAML_TYPES = Annotated[BLOCK_YAML_SUBCLASSES, Field(discriminator="block_type")] diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 352f79b9e..c2c9c160c 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -21,7 +21,14 @@ from skyvern.forge.sdk.models import Step from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateParameterKeys -from skyvern.forge.sdk.workflow.models.block import BlockType, BlockTypeVar, CodeBlock, ForLoopBlock, TaskBlock +from skyvern.forge.sdk.workflow.models.block import ( + BlockType, + BlockTypeVar, + CodeBlock, + ForLoopBlock, + TaskBlock, + TextPromptBlock, +) from skyvern.forge.sdk.workflow.models.parameter import ( AWSSecretParameter, OutputParameter, @@ -714,4 +721,15 @@ async def block_yaml_to_block(block_yaml: BLOCK_YAML_TYPES, parameters: dict[str else [], output_parameter=output_parameter, ) + elif block_yaml.block_type == BlockType.TEXT_PROMPT: + return TextPromptBlock( + label=block_yaml.label, + llm_key=block_yaml.llm_key, + prompt=block_yaml.prompt, + parameters=[parameters[parameter_key] for parameter_key in block_yaml.parameter_keys] + if block_yaml.parameter_keys + else [], + json_schema=block_yaml.json_schema, + output_parameter=output_parameter, + ) raise ValueError(f"Invalid block type {block_yaml.block_type}") From f06def52788ffbf80f39ec6896040f0faed06fd2 Mon Sep 17 00:00:00 2001 From: ykeremy Date: Mon, 25 Mar 2024 07:50:12 +0000 Subject: [PATCH 2/3] =?UTF-8?q?=F0=9F=94=84=20synced=20local=20'pyproject.?= =?UTF-8?q?toml'=20with=20remote=20'pyproject.toml'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 74bcc3bdb..cd2cd9326 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ curlify = "^2.2.1" typer = "^0.9.0" types-toml = "^0.10.8.7" apscheduler = "^3.10.4" +httpx = "^0.27.0" [tool.poetry.group.dev.dependencies] From 7a408b5b3b77f8aa41e67a0bf9b4b46e14bd0b6c Mon Sep 17 00:00:00 2001 From: ykeremy Date: Mon, 25 Mar 2024 07:50:12 +0000 Subject: [PATCH 3/3] =?UTF-8?q?=F0=9F=94=84=20synced=20local=20'poetry.loc?= =?UTF-8?q?k'=20with=20remote=20'poetry.lock'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- poetry.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/poetry.lock b/poetry.lock index c10bedc18..a767b3dca 100644 --- a/poetry.lock +++ b/poetry.lock @@ -6880,4 +6880,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = "^3.11,<3.12" -content-hash = "402c47a5e38eef5bbd38a63cc8116f11b4abf0d676c9b411a2d7f425748e9c4c" +content-hash = "9c2a8d3c2c9b239c6338f53485f9eace6d3eac112fa9246d9bc0a83c92f61a1d"