From ae1333cc70fff14bacbab6c1a0a96f63c4db92fc Mon Sep 17 00:00:00 2001 From: Farid Rener Date: Fri, 13 Dec 2024 14:51:39 -0500 Subject: [PATCH 1/5] Add CodeNode --- apps/pipelines/nodes/nodes.py | 99 +++++++++++++++++++++++++- apps/pipelines/tests/test_code_node.py | 88 +++++++++++++++++++++++ apps/pipelines/tests/utils.py | 12 ++++ requirements/requirements.in | 1 + requirements/requirements.txt | 2 + 5 files changed, 201 insertions(+), 1 deletion(-) create mode 100644 apps/pipelines/tests/test_code_node.py diff --git a/apps/pipelines/nodes/nodes.py b/apps/pipelines/nodes/nodes.py index 2d68811c3..9cf93e2ea 100644 --- a/apps/pipelines/nodes/nodes.py +++ b/apps/pipelines/nodes/nodes.py @@ -1,4 +1,6 @@ +import datetime import json +import time from typing import Literal import tiktoken @@ -14,13 +16,14 @@ from pydantic.config import ConfigDict from pydantic_core import PydanticCustomError from pydantic_core.core_schema import FieldValidationInfo +from RestrictedPython import compile_restricted_function, safe_builtins, safe_globals from apps.assistants.models import OpenAiAssistant from apps.channels.datamodels import Attachment from apps.chat.conversation import compress_chat_history, compress_pipeline_chat_history from apps.chat.models import ChatMessageType from apps.experiments.models import ExperimentSession, ParticipantData -from apps.pipelines.exceptions import PipelineNodeBuildError +from apps.pipelines.exceptions import PipelineNodeBuildError, PipelineNodeRunError from apps.pipelines.models import PipelineChatHistory, PipelineChatHistoryTypes from apps.pipelines.nodes.base import NodeSchema, OptionsSource, PipelineNode, PipelineState, UiSchema, Widgets from apps.pipelines.tasks import send_email_from_pipeline @@ -622,3 +625,97 @@ def _get_assistant_runnable(self, assistant: OpenAiAssistant, session: Experimen return AgentAssistantChat(adapter=adapter, history_manager=history_manager) else: return AssistantChat(adapter=adapter, history_manager=history_manager) + + +class CodeNode(PipelineNode): + """Runs python""" + + model_config = ConfigDict(json_schema_extra=NodeSchema(label="Python Node")) + code: str = Field( + description="The code to run", + json_schema_extra=UiSchema(widget=Widgets.expandable_text), # TODO: add a code widget + ) + + @field_validator("code") + def validate_code(cls, value, info: FieldValidationInfo): + if not value: + value = "return input" + + byte_code = compile_restricted_function( + "input,shared_state", + value, + name="main", + filename="", + ) + + if byte_code.errors: + raise PydanticCustomError("invalid_code", "{errors}", {"errors": "\n".join(byte_code.errors)}) + return value + + def _process(self, input: str, state: PipelineState, node_id: str) -> PipelineState: + function_name = "main" + function_args = "input" + byte_code = compile_restricted_function( + function_args, + self.code, + name=function_name, + filename="", + ) + + custom_locals = {} + custom_globals = self._get_custom_globals() + exec(byte_code.code, custom_globals, custom_locals) + + try: + result = str(custom_locals[function_name](input)) + except Exception as exc: + raise PipelineNodeRunError(exc) from exc + return PipelineState.from_node_output(node_id=node_id, output=result) + + def _get_custom_globals(self): + from RestrictedPython.Eval import ( + default_guarded_getitem, + default_guarded_getiter, + ) + + custom_globals = safe_globals.copy() + custom_globals.update( + { + "__builtins__": self._get_custom_builtins(), + "json": json, + "datetime": datetime, + "time": time, + "_getitem_": default_guarded_getitem, + "_getiter_": default_guarded_getiter, + "_write_": lambda x: x, + } + ) + return custom_globals + + def _get_custom_builtins(self): + allowed_modules = { + "json", + "re", + "datetime", + "time", + } + custom_builtins = safe_builtins.copy() + custom_builtins.update( + { + "min": min, + "max": max, + "sum": sum, + "abs": abs, + "all": all, + "any": any, + "datetime": datetime, + } + ) + + def guarded_import(name, *args, **kwargs): + if name not in allowed_modules: + raise ImportError(f"Importing '{name}' is not allowed") + return __import__(name, *args, **kwargs) + + custom_builtins["__import__"] = guarded_import + return custom_builtins diff --git a/apps/pipelines/tests/test_code_node.py b/apps/pipelines/tests/test_code_node.py new file mode 100644 index 000000000..8edcc1a07 --- /dev/null +++ b/apps/pipelines/tests/test_code_node.py @@ -0,0 +1,88 @@ +import json +from unittest import mock + +import pytest + +from apps.pipelines.exceptions import PipelineNodeBuildError, PipelineNodeRunError +from apps.pipelines.nodes.base import PipelineState +from apps.pipelines.tests.utils import ( + code_node, + create_runnable, + end_node, + start_node, +) +from apps.utils.factories.pipelines import PipelineFactory +from apps.utils.pytest import django_db_with_data + + +@pytest.fixture() +def pipeline(): + return PipelineFactory() + + +EXTRA_FUNCTION = """ +def other(foo): + return f"other {foo}" + +return other(input) +""" + +IMPORTS = """ +import json +import datetime +import re +import time +return json.loads(input) +""" + + +@django_db_with_data(available_apps=("apps.service_providers",)) +@mock.patch("apps.pipelines.nodes.base.PipelineNode.logger", mock.Mock()) +@pytest.mark.parametrize( + ("code", "input", "output"), + [ + ("return f'Hello, {input}!'", "World", "Hello, World!"), + ("", "foo", "foo"), # No code just returns the input + (EXTRA_FUNCTION, "blah", "other blah"), # Calling a separate function is possible + ("'foo'", "", "None"), # No return value will return "None" + (IMPORTS, json.dumps({"a": "b"}), str(json.loads('{"a": "b"}'))), # Importing json will work + ], +) +def test_code_node(pipeline, code, input, output): + nodes = [ + start_node(), + code_node(code), + end_node(), + ] + assert create_runnable(pipeline, nodes).invoke(PipelineState(messages=[input]))["messages"][-1] == output + + +@django_db_with_data(available_apps=("apps.service_providers",)) +@mock.patch("apps.pipelines.nodes.base.PipelineNode.logger", mock.Mock()) +def test_code_node_syntax_error(pipeline): + nodes = [ + start_node(), + code_node("this{}"), + end_node(), + ] + with pytest.raises(PipelineNodeBuildError, match="SyntaxError: invalid syntax at statement: 'this{}'"): + create_runnable(pipeline, nodes).invoke(PipelineState(messages=["World"]))["messages"][-1] + + +@django_db_with_data(available_apps=("apps.service_providers",)) +@mock.patch("apps.pipelines.nodes.base.PipelineNode.logger", mock.Mock()) +@pytest.mark.parametrize( + ("code", "input", "error"), + [ + ("import collections", "", "Importing 'collections' is not allowed"), + ("return f'Hello, {blah}!'", "", "name 'blah' is not defined"), + ], +) +def test_code_node_runtime_errors(pipeline, code, input, error): + nodes = [ + start_node(), + code_node(code), + end_node(), + ] + with pytest.raises(PipelineNodeRunError, match=error): + create_runnable(pipeline, nodes).invoke(PipelineState(messages=[input]))["messages"][-1] diff --git a/apps/pipelines/tests/utils.py b/apps/pipelines/tests/utils.py index d2513921d..8ded20c8a 100644 --- a/apps/pipelines/tests/utils.py +++ b/apps/pipelines/tests/utils.py @@ -174,3 +174,15 @@ def extract_structured_data_node(provider_id: str, provider_model_id: str, data_ "data_schema": data_schema, }, } + + +def code_node(code: str | None = None): + if code is None: + code = "return f'Hello, {input}!'" + return { + "id": str(uuid4()), + "type": nodes.CodeNode.__name__, + "params": { + "code": code, + }, + } diff --git a/requirements/requirements.in b/requirements/requirements.in index d734a7d3b..b0a4a53a8 100644 --- a/requirements/requirements.in +++ b/requirements/requirements.in @@ -47,6 +47,7 @@ psycopg[binary] pyTelegramBotAPI==4.12.0 pydantic pydub # Audio transcription +RestrictedPython sentry-sdk slack-bolt taskbadger diff --git a/requirements/requirements.txt b/requirements/requirements.txt index d43f13dd9..41485e94e 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -412,6 +412,8 @@ requests-oauthlib==1.3.1 # via django-allauth requests-toolbelt==1.0.0 # via langsmith +restrictedpython==7.4 + # via -r requirements.in rich==13.6.0 # via typer rpds-py==0.12.0 From 914e864961dbb481678604ca791445944b0cef77 Mon Sep 17 00:00:00 2001 From: Farid Rener Date: Fri, 13 Dec 2024 22:25:10 -0500 Subject: [PATCH 2/5] Add CodeMirror to Python Node --- apps/pipelines/nodes/base.py | 1 + apps/pipelines/nodes/nodes.py | 2 +- apps/pipelines/tests/data/CodeNode.json | 18 ++ .../apps/pipeline/nodes/widgets.tsx | 141 +++++++++- package-lock.json | 250 ++++++++++++++++++ package.json | 3 + 6 files changed, 404 insertions(+), 11 deletions(-) create mode 100644 apps/pipelines/tests/data/CodeNode.json diff --git a/apps/pipelines/nodes/base.py b/apps/pipelines/nodes/base.py index c5c52ecaa..3e971da90 100644 --- a/apps/pipelines/nodes/base.py +++ b/apps/pipelines/nodes/base.py @@ -121,6 +121,7 @@ def logger(self): class Widgets(StrEnum): expandable_text = "expandable_text" + code = "code" toggle = "toggle" select = "select" float = "float" diff --git a/apps/pipelines/nodes/nodes.py b/apps/pipelines/nodes/nodes.py index 9cf93e2ea..1350f68f0 100644 --- a/apps/pipelines/nodes/nodes.py +++ b/apps/pipelines/nodes/nodes.py @@ -633,7 +633,7 @@ class CodeNode(PipelineNode): model_config = ConfigDict(json_schema_extra=NodeSchema(label="Python Node")) code: str = Field( description="The code to run", - json_schema_extra=UiSchema(widget=Widgets.expandable_text), # TODO: add a code widget + json_schema_extra=UiSchema(widget=Widgets.code), ) @field_validator("code") diff --git a/apps/pipelines/tests/data/CodeNode.json b/apps/pipelines/tests/data/CodeNode.json new file mode 100644 index 000000000..c58a1ed06 --- /dev/null +++ b/apps/pipelines/tests/data/CodeNode.json @@ -0,0 +1,18 @@ +{ + "description": "Runs python", + "properties": { + "code": { + "description": "The code to run", + "title": "Code", + "type": "string", + "ui:widget": "code" + } + }, + "required": [ + "code" + ], + "title": "CodeNode", + "type": "object", + "ui:flow_node_type": "pipelineNode", + "ui:label": "Python Node" +} \ No newline at end of file diff --git a/assets/javascript/apps/pipeline/nodes/widgets.tsx b/assets/javascript/apps/pipeline/nodes/widgets.tsx index b86ad6bda..636cd4182 100644 --- a/assets/javascript/apps/pipeline/nodes/widgets.tsx +++ b/assets/javascript/apps/pipeline/nodes/widgets.tsx @@ -3,13 +3,17 @@ import React, { ChangeEventHandler, ReactNode, useId, + useEffect, } from "react"; import { useState } from "react"; -import {TypedOption} from "../types/nodeParameterValues"; +import CodeMirror from '@uiw/react-codemirror'; +import { python } from "@codemirror/lang-python"; +import { githubLight, githubDark } from "@uiw/codemirror-theme-github"; +import { TypedOption } from "../types/nodeParameterValues"; import usePipelineStore from "../stores/pipelineStore"; -import {classNames, concatenate, getCachedData, getSelectOptions} from "../utils"; -import {NodeParams, PropertySchema} from "../types/nodeParams"; -import {Node, useUpdateNodeInternals} from "reactflow"; +import { classNames, concatenate, getCachedData, getSelectOptions } from "../utils"; +import { NodeParams, PropertySchema } from "../types/nodeParams"; +import { Node, useUpdateNodeInternals } from "reactflow"; import DOMPurify from 'dompurify'; export function getWidget(name: string) { @@ -22,10 +26,12 @@ export function getWidget(name: string) { return RangeWidget case "expandable_text": return ExpandableTextWidget + case "code": + return CodeWidget case "select": return SelectWidget - case "multiselect": - return MultiSelectWidget + case "multiselect": + return MultiSelectWidget case "llm_provider_model": return LlmWidget case "history": @@ -140,7 +146,7 @@ function SelectWidget(props: WidgetParams) { setLink(selectedOption?.edit_url); props.updateParamValue(event); }; - + return
@@ -164,7 +170,7 @@ function SelectWidget(props: WidgetParams) {
)} - + } @@ -195,8 +201,8 @@ function MultiSelectWidget(props: WidgetParams) { selectedValues = selectedValues.filter((tool) => tool !== event.target.name) } setNode(props.nodeId, (old) => { - return getNewNodeData(old, selectedValues); - } + return getNewNodeData(old, selectedValues); + } ); }; @@ -220,6 +226,121 @@ function MultiSelectWidget(props: WidgetParams) { ) } +export function CodeWidget(props: WidgetParams) { + const [isDarkMode, setIsDarkMode] = useState(false); + const setNode = usePipelineStore((state) => state.setNode); + const onChangeCallback = (value: string) => { + setNode(props.nodeId, (old) => ({ + ...old, + data: { + ...old.data, + params: { + ...old.data.params, + [props.name]: value, + }, + }, + })); + }; + + useEffect(() => { + // Set dark / light mode + const mediaQuery: MediaQueryList = window.matchMedia("(prefers-color-scheme: dark)"); + const handleChange = (event: MediaQueryListEvent): void => { + setIsDarkMode(event.matches); + }; + setIsDarkMode(mediaQuery.matches); + + mediaQuery.addEventListener("change", handleChange); + return () => mediaQuery.removeEventListener("change", handleChange); + }, []); + + const modalId = useId(); + const openModal = () => (document.getElementById(modalId) as HTMLDialogElement)?.showModal() + const label = ( + <> + {props.label} +
+ +
+ + ) + return ( + <> + +
+ +
+
+ + + + ); +} + + +export function CodeModal( + { modalId, humanName, value, onChange, isDarkMode }: { + modalId: string; + humanName: string; + value: string; + onChange: (value: string) => void; + isDarkMode: boolean; + }) { + return ( + +
+
+ +
+
+

+ {humanName} +

+ +
+
+
+ {/* Allows closing the modal by clicking outside of it */} + +
+ + ); +} + + export function TextModal( {modalId, humanName, name, value, onChange}: { diff --git a/package-lock.json b/package-lock.json index 64224b2be..9a7540290 100644 --- a/package-lock.json +++ b/package-lock.json @@ -8,6 +8,9 @@ "name": "open-chat-studio", "version": "1.0.0", "dependencies": { + "@codemirror/lang-python": "^6.1.6", + "@uiw/codemirror-theme-github": "^4.23.6", + "@uiw/react-codemirror": "^4.23.6", "alertifyjs": "^1.14.0", "alpinejs": "^3.14.3", "axios": "^1.7.7", @@ -1776,6 +1779,108 @@ "node": ">=6.9.0" } }, + "node_modules/@codemirror/autocomplete": { + "version": "6.18.3", + "resolved": "https://registry.npmjs.org/@codemirror/autocomplete/-/autocomplete-6.18.3.tgz", + "integrity": "sha512-1dNIOmiM0z4BIBwxmxEfA1yoxh1MF/6KPBbh20a5vphGV0ictKlgQsbJs6D6SkR6iJpGbpwRsa6PFMNlg9T9pQ==", + "dependencies": { + "@codemirror/language": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.17.0", + "@lezer/common": "^1.0.0" + }, + "peerDependencies": { + "@codemirror/language": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0", + "@lezer/common": "^1.0.0" + } + }, + "node_modules/@codemirror/commands": { + "version": "6.7.1", + "resolved": "https://registry.npmjs.org/@codemirror/commands/-/commands-6.7.1.tgz", + "integrity": "sha512-llTrboQYw5H4THfhN4U3qCnSZ1SOJ60ohhz+SzU0ADGtwlc533DtklQP0vSFaQuCPDn3BPpOd1GbbnUtwNjsrw==", + "dependencies": { + "@codemirror/language": "^6.0.0", + "@codemirror/state": "^6.4.0", + "@codemirror/view": "^6.27.0", + "@lezer/common": "^1.1.0" + } + }, + "node_modules/@codemirror/lang-python": { + "version": "6.1.6", + "resolved": "https://registry.npmjs.org/@codemirror/lang-python/-/lang-python-6.1.6.tgz", + "integrity": "sha512-ai+01WfZhWqM92UqjnvorkxosZ2aq2u28kHvr+N3gu012XqY2CThD67JPMHnGceRfXPDBmn1HnyqowdpF57bNg==", + "dependencies": { + "@codemirror/autocomplete": "^6.3.2", + "@codemirror/language": "^6.8.0", + "@codemirror/state": "^6.0.0", + "@lezer/common": "^1.2.1", + "@lezer/python": "^1.1.4" + } + }, + "node_modules/@codemirror/language": { + "version": "6.10.6", + "resolved": "https://registry.npmjs.org/@codemirror/language/-/language-6.10.6.tgz", + "integrity": "sha512-KrsbdCnxEztLVbB5PycWXFxas4EOyk/fPAfruSOnDDppevQgid2XZ+KbJ9u+fDikP/e7MW7HPBTvTb8JlZK9vA==", + "dependencies": { + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.23.0", + "@lezer/common": "^1.1.0", + "@lezer/highlight": "^1.0.0", + "@lezer/lr": "^1.0.0", + "style-mod": "^4.0.0" + } + }, + "node_modules/@codemirror/lint": { + "version": "6.8.4", + "resolved": "https://registry.npmjs.org/@codemirror/lint/-/lint-6.8.4.tgz", + "integrity": "sha512-u4q7PnZlJUojeRe8FJa/njJcMctISGgPQ4PnWsd9268R4ZTtU+tfFYmwkBvgcrK2+QQ8tYFVALVb5fVJykKc5A==", + "dependencies": { + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.35.0", + "crelt": "^1.0.5" + } + }, + "node_modules/@codemirror/search": { + "version": "6.5.8", + "resolved": "https://registry.npmjs.org/@codemirror/search/-/search-6.5.8.tgz", + "integrity": "sha512-PoWtZvo7c1XFeZWmmyaOp2G0XVbOnm+fJzvghqGAktBW3cufwJUWvSCcNG0ppXiBEM05mZu6RhMtXPv2hpllig==", + "dependencies": { + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0", + "crelt": "^1.0.5" + } + }, + "node_modules/@codemirror/state": { + "version": "6.5.0", + "resolved": "https://registry.npmjs.org/@codemirror/state/-/state-6.5.0.tgz", + "integrity": "sha512-MwBHVK60IiIHDcoMet78lxt6iw5gJOGSbNbOIVBHWVXIH4/Nq1+GQgLLGgI1KlnN86WDXsPudVaqYHKBIx7Eyw==", + "dependencies": { + "@marijn/find-cluster-break": "^1.0.0" + } + }, + "node_modules/@codemirror/theme-one-dark": { + "version": "6.1.2", + "resolved": "https://registry.npmjs.org/@codemirror/theme-one-dark/-/theme-one-dark-6.1.2.tgz", + "integrity": "sha512-F+sH0X16j/qFLMAfbciKTxVOwkdAS336b7AXTKOZhy8BR3eH/RelsnLgLFINrpST63mmN2OuwUt0W2ndUgYwUA==", + "dependencies": { + "@codemirror/language": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0", + "@lezer/highlight": "^1.0.0" + } + }, + "node_modules/@codemirror/view": { + "version": "6.35.3", + "resolved": "https://registry.npmjs.org/@codemirror/view/-/view-6.35.3.tgz", + "integrity": "sha512-ScY7L8+EGdPl4QtoBiOzE4FELp7JmNUsBvgBcCakXWM2uiv/K89VAzU3BMDscf0DsACLvTKePbd5+cFDTcei6g==", + "dependencies": { + "@codemirror/state": "^6.5.0", + "style-mod": "^4.1.0", + "w3c-keyname": "^2.2.4" + } + }, "node_modules/@discoveryjs/json-ext": { "version": "0.5.7", "resolved": "https://registry.npmjs.org/@discoveryjs/json-ext/-/json-ext-0.5.7.tgz", @@ -2109,6 +2214,42 @@ "integrity": "sha512-M5UknZPHRu3DEDWoipU6sE8PdkZ6Z/S+v4dD+Ke8IaNlpdSQah50lz1KtcFBa2vsdOnwbbnxJwVM4wty6udA5w==", "license": "MIT" }, + "node_modules/@lezer/common": { + "version": "1.2.3", + "resolved": "https://registry.npmjs.org/@lezer/common/-/common-1.2.3.tgz", + "integrity": "sha512-w7ojc8ejBqr2REPsWxJjrMFsA/ysDCFICn8zEOR9mrqzOu2amhITYuLD8ag6XZf0CFXDrhKqw7+tW8cX66NaDA==" + }, + "node_modules/@lezer/highlight": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@lezer/highlight/-/highlight-1.2.1.tgz", + "integrity": "sha512-Z5duk4RN/3zuVO7Jq0pGLJ3qynpxUVsh7IbUbGj88+uV2ApSAn6kWg2au3iJb+0Zi7kKtqffIESgNcRXWZWmSA==", + "dependencies": { + "@lezer/common": "^1.0.0" + } + }, + "node_modules/@lezer/lr": { + "version": "1.4.2", + "resolved": "https://registry.npmjs.org/@lezer/lr/-/lr-1.4.2.tgz", + "integrity": "sha512-pu0K1jCIdnQ12aWNaAVU5bzi7Bd1w54J3ECgANPmYLtQKP0HBj2cE/5coBD66MT10xbtIuUr7tg0Shbsvk0mDA==", + "dependencies": { + "@lezer/common": "^1.0.0" + } + }, + "node_modules/@lezer/python": { + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@lezer/python/-/python-1.1.15.tgz", + "integrity": "sha512-aVQ43m2zk4FZYedCqL0KHPEUsqZOrmAvRhkhHlVPnDD1HODDyyQv5BRIuod4DadkgBEZd53vQOtXTonNbEgjrQ==", + "dependencies": { + "@lezer/common": "^1.2.0", + "@lezer/highlight": "^1.0.0", + "@lezer/lr": "^1.0.0" + } + }, + "node_modules/@marijn/find-cluster-break": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/@marijn/find-cluster-break/-/find-cluster-break-1.0.2.tgz", + "integrity": "sha512-l0h88YhZFyKdXIFNfSWpyjStDjGHwZ/U7iobcK1cQQD8sejsONdQtTVU+1wVN1PBw40PiiHB1vA5S7VTfQiP9g==" + }, "node_modules/@nicolo-ribaudo/chokidar-2": { "version": "2.1.8-no-fsevents.3", "resolved": "https://registry.npmjs.org/@nicolo-ribaudo/chokidar-2/-/chokidar-2-2.1.8-no-fsevents.3.tgz", @@ -3188,6 +3329,86 @@ "url": "https://opencollective.com/eslint" } }, + "node_modules/@uiw/codemirror-extensions-basic-setup": { + "version": "4.23.6", + "resolved": "https://registry.npmjs.org/@uiw/codemirror-extensions-basic-setup/-/codemirror-extensions-basic-setup-4.23.6.tgz", + "integrity": "sha512-bvtq8IOvdkLJMhoJBRGPEzU51fMpPDwEhcAHp9xCR05MtbIokQgsnLXrmD1aZm6e7s/3q47H+qdSfAAkR5MkLA==", + "dependencies": { + "@codemirror/autocomplete": "^6.0.0", + "@codemirror/commands": "^6.0.0", + "@codemirror/language": "^6.0.0", + "@codemirror/lint": "^6.0.0", + "@codemirror/search": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0" + }, + "funding": { + "url": "https://jaywcjlove.github.io/#/sponsor" + }, + "peerDependencies": { + "@codemirror/autocomplete": ">=6.0.0", + "@codemirror/commands": ">=6.0.0", + "@codemirror/language": ">=6.0.0", + "@codemirror/lint": ">=6.0.0", + "@codemirror/search": ">=6.0.0", + "@codemirror/state": ">=6.0.0", + "@codemirror/view": ">=6.0.0" + } + }, + "node_modules/@uiw/codemirror-theme-github": { + "version": "4.23.6", + "resolved": "https://registry.npmjs.org/@uiw/codemirror-theme-github/-/codemirror-theme-github-4.23.6.tgz", + "integrity": "sha512-p74sbyuDo7JsYlGWNtyNyJSyr9N4niCPGqYreQhoAk2C6c4JXi0sGwZRlHd8KyynsxCdSA9L727SADOuukD1Ug==", + "dependencies": { + "@uiw/codemirror-themes": "4.23.6" + }, + "funding": { + "url": "https://jaywcjlove.github.io/#/sponsor" + } + }, + "node_modules/@uiw/codemirror-themes": { + "version": "4.23.6", + "resolved": "https://registry.npmjs.org/@uiw/codemirror-themes/-/codemirror-themes-4.23.6.tgz", + "integrity": "sha512-0dpuLQW+V6zrKvfvor/eo71V3tpr2L2Hsu8QZAdtSzksjWABxTOzH3ShaBRxCEsrz6sU9sa9o7ShwBMMDz59bQ==", + "dependencies": { + "@codemirror/language": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0" + }, + "funding": { + "url": "https://jaywcjlove.github.io/#/sponsor" + }, + "peerDependencies": { + "@codemirror/language": ">=6.0.0", + "@codemirror/state": ">=6.0.0", + "@codemirror/view": ">=6.0.0" + } + }, + "node_modules/@uiw/react-codemirror": { + "version": "4.23.6", + "resolved": "https://registry.npmjs.org/@uiw/react-codemirror/-/react-codemirror-4.23.6.tgz", + "integrity": "sha512-caYKGV6TfGLRV1HHD3p0G3FiVzKL1go7wes5XT2nWjB0+dTdyzyb81MKRSacptgZcotujfNO6QXn65uhETRAMw==", + "dependencies": { + "@babel/runtime": "^7.18.6", + "@codemirror/commands": "^6.1.0", + "@codemirror/state": "^6.1.1", + "@codemirror/theme-one-dark": "^6.0.0", + "@uiw/codemirror-extensions-basic-setup": "4.23.6", + "codemirror": "^6.0.0" + }, + "funding": { + "url": "https://jaywcjlove.github.io/#/sponsor" + }, + "peerDependencies": { + "@babel/runtime": ">=7.11.0", + "@codemirror/state": ">=6.0.0", + "@codemirror/theme-one-dark": ">=6.0.0", + "@codemirror/view": ">=6.0.0", + "codemirror": ">=6.0.0", + "react": ">=16.8.0", + "react-dom": ">=16.8.0" + } + }, "node_modules/@vue/reactivity": { "version": "3.1.5", "resolved": "https://registry.npmjs.org/@vue/reactivity/-/reactivity-3.1.5.tgz", @@ -4114,6 +4335,20 @@ "node": ">=6" } }, + "node_modules/codemirror": { + "version": "6.0.1", + "resolved": "https://registry.npmjs.org/codemirror/-/codemirror-6.0.1.tgz", + "integrity": "sha512-J8j+nZ+CdWmIeFIGXEFbFPtpiYacFMDR8GlHK3IyHQJMCaVRfGx9NT+Hxivv1ckLWPvNdZqndbr/7lVhrf/Svg==", + "dependencies": { + "@codemirror/autocomplete": "^6.0.0", + "@codemirror/commands": "^6.0.0", + "@codemirror/language": "^6.0.0", + "@codemirror/lint": "^6.0.0", + "@codemirror/search": "^6.0.0", + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0" + } + }, "node_modules/color-convert": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/color-convert/-/color-convert-2.0.1.tgz", @@ -4223,6 +4458,11 @@ } } }, + "node_modules/crelt": { + "version": "1.0.6", + "resolved": "https://registry.npmjs.org/crelt/-/crelt-1.0.6.tgz", + "integrity": "sha512-VQ2MBenTq1fWZUH9DJNGti7kKv6EeAuYr3cLwxUWhIu1baTaXh4Ib5W2CqHVqib4/MqbYGJqiL3Zb8GJZr3l4g==" + }, "node_modules/cross-spawn": { "version": "7.0.6", "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-7.0.6.tgz", @@ -8612,6 +8852,11 @@ "webpack": "^5.27.0" } }, + "node_modules/style-mod": { + "version": "4.1.2", + "resolved": "https://registry.npmjs.org/style-mod/-/style-mod-4.1.2.tgz", + "integrity": "sha512-wnD1HyVqpJUI2+eKZ+eo1UwghftP6yuFheBqqe+bWCotBjC2K1YnteJILRMs3SM4V/0dLEW1SC27MWP5y+mwmw==" + }, "node_modules/sucrase": { "version": "3.35.0", "resolved": "https://registry.npmjs.org/sucrase/-/sucrase-3.35.0.tgz", @@ -9206,6 +9451,11 @@ "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==", "license": "MIT" }, + "node_modules/w3c-keyname": { + "version": "2.2.8", + "resolved": "https://registry.npmjs.org/w3c-keyname/-/w3c-keyname-2.2.8.tgz", + "integrity": "sha512-dpojBhNsCNN7T82Tm7k26A6G9ML3NkhDsnw9n/eoxSRlVBB4CEtIQ/KTCLI2Fwf3ataSXRhYFkQi3SlnFwPvPQ==" + }, "node_modules/watchpack": { "version": "2.4.2", "resolved": "https://registry.npmjs.org/watchpack/-/watchpack-2.4.2.tgz", diff --git a/package.json b/package.json index bca33128c..aa13fba80 100644 --- a/package.json +++ b/package.json @@ -48,6 +48,9 @@ "webpack-cli": "^5.1.4" }, "dependencies": { + "@codemirror/lang-python": "^6.1.6", + "@uiw/codemirror-theme-github": "^4.23.6", + "@uiw/react-codemirror": "^4.23.6", "alertifyjs": "^1.14.0", "alpinejs": "^3.14.3", "axios": "^1.7.7", From d42d5a256df5ac60c0127c2285e84e5c675ce3b4 Mon Sep 17 00:00:00 2001 From: Farid Rener Date: Sat, 14 Dec 2024 21:22:01 -0500 Subject: [PATCH 3/5] Expose the 'main' function to the user --- apps/pipelines/nodes/nodes.py | 56 +++++++++++++++++-------- apps/pipelines/tests/data/CodeNode.json | 4 +- apps/pipelines/tests/test_code_node.py | 53 +++++++++++++++-------- 3 files changed, 76 insertions(+), 37 deletions(-) diff --git a/apps/pipelines/nodes/nodes.py b/apps/pipelines/nodes/nodes.py index 1350f68f0..30f86e8cf 100644 --- a/apps/pipelines/nodes/nodes.py +++ b/apps/pipelines/nodes/nodes.py @@ -1,4 +1,5 @@ import datetime +import inspect import json import time from typing import Literal @@ -16,7 +17,7 @@ from pydantic.config import ConfigDict from pydantic_core import PydanticCustomError from pydantic_core.core_schema import FieldValidationInfo -from RestrictedPython import compile_restricted_function, safe_builtins, safe_globals +from RestrictedPython import compile_restricted, safe_builtins, safe_globals from apps.assistants.models import OpenAiAssistant from apps.channels.datamodels import Attachment @@ -627,11 +628,19 @@ def _get_assistant_runnable(self, assistant: OpenAiAssistant, session: Experimen return AssistantChat(adapter=adapter, history_manager=history_manager) +DEFAULT_FUNCTION = """# You must define a main function, which takes the node input as a string. +# Return a string to pass to the next node. +def main(input: str) -> str: + return input +""" + + class CodeNode(PipelineNode): """Runs python""" model_config = ConfigDict(json_schema_extra=NodeSchema(label="Python Node")) code: str = Field( + default=DEFAULT_FUNCTION, description="The code to run", json_schema_extra=UiSchema(widget=Widgets.code), ) @@ -639,34 +648,47 @@ class CodeNode(PipelineNode): @field_validator("code") def validate_code(cls, value, info: FieldValidationInfo): if not value: - value = "return input" - - byte_code = compile_restricted_function( - "input,shared_state", - value, - name="main", - filename="", - ) + value = DEFAULT_FUNCTION + try: + byte_code = compile_restricted( + value, + filename="", + mode="exec", + ) + custom_locals = {} + exec(byte_code, {}, custom_locals) - if byte_code.errors: - raise PydanticCustomError("invalid_code", "{errors}", {"errors": "\n".join(byte_code.errors)}) + try: + main = custom_locals["main"] + except KeyError: + raise SyntaxError("You must define a 'main' function") + + for name, item in custom_locals.items(): + if name != "main" and inspect.isfunction(item): + raise SyntaxError( + "You can only define a single function, 'main' at the top level. " + "You may use nested functions inside that function if required" + ) + + if len(inspect.signature(main).parameters) > 1: + raise SyntaxError("The main function should take a single argument as input") + + except SyntaxError as exc: + raise PydanticCustomError("invalid_code", "{error}", {"error": exc.msg}) return value def _process(self, input: str, state: PipelineState, node_id: str) -> PipelineState: function_name = "main" - function_args = "input" - byte_code = compile_restricted_function( - function_args, + byte_code = compile_restricted( self.code, - name=function_name, filename="", + mode="exec", ) custom_locals = {} custom_globals = self._get_custom_globals() - exec(byte_code.code, custom_globals, custom_locals) - try: + exec(byte_code, custom_globals, custom_locals) result = str(custom_locals[function_name](input)) except Exception as exc: raise PipelineNodeRunError(exc) from exc diff --git a/apps/pipelines/tests/data/CodeNode.json b/apps/pipelines/tests/data/CodeNode.json index c58a1ed06..7825da861 100644 --- a/apps/pipelines/tests/data/CodeNode.json +++ b/apps/pipelines/tests/data/CodeNode.json @@ -2,15 +2,13 @@ "description": "Runs python", "properties": { "code": { + "default": "# You must define a main function, which takes the node input as a string.\n# Return a string to pass to the next node.\ndef main(input: str) -> str:\n return input\n", "description": "The code to run", "title": "Code", "type": "string", "ui:widget": "code" } }, - "required": [ - "code" - ], "title": "CodeNode", "type": "object", "ui:flow_node_type": "pipelineNode", diff --git a/apps/pipelines/tests/test_code_node.py b/apps/pipelines/tests/test_code_node.py index 8edcc1a07..fda70633b 100644 --- a/apps/pipelines/tests/test_code_node.py +++ b/apps/pipelines/tests/test_code_node.py @@ -20,19 +20,13 @@ def pipeline(): return PipelineFactory() -EXTRA_FUNCTION = """ -def other(foo): - return f"other {foo}" - -return other(input) -""" - IMPORTS = """ import json import datetime import re import time -return json.loads(input) +def main(input): + return json.loads(input) """ @@ -41,11 +35,11 @@ def other(foo): @pytest.mark.parametrize( ("code", "input", "output"), [ - ("return f'Hello, {input}!'", "World", "Hello, World!"), + ("def main(input):\n\treturn f'Hello, {input}!'", "World", "Hello, World!"), ("", "foo", "foo"), # No code just returns the input - (EXTRA_FUNCTION, "blah", "other blah"), # Calling a separate function is possible - ("'foo'", "", "None"), # No return value will return "None" + ("def main(input):\n\t'foo'", "", "None"), # No return value will return "None" (IMPORTS, json.dumps({"a": "b"}), str(json.loads('{"a": "b"}'))), # Importing json will work + ("def main(blah):\n\treturn f'Hello, {blah}!'", "World", "Hello, World!"), # Renaming the argument works ], ) def test_code_node(pipeline, code, input, output): @@ -57,16 +51,41 @@ def test_code_node(pipeline, code, input, output): assert create_runnable(pipeline, nodes).invoke(PipelineState(messages=[input]))["messages"][-1] == output +EXTRA_FUNCTION = """ +def other(foo): + return f"other {foo}" + +def main(input): + return other(input) +""" + + @django_db_with_data(available_apps=("apps.service_providers",)) @mock.patch("apps.pipelines.nodes.base.PipelineNode.logger", mock.Mock()) -def test_code_node_syntax_error(pipeline): +@pytest.mark.parametrize( + ("code", "input", "error"), + [ + ("this{}", "", "SyntaxError: invalid syntax at statement: 'this{}"), + ( + EXTRA_FUNCTION, + "", + ( + "You can only define a single function, 'main' at the top level. " + "You may use nested functions inside that function if required" + ), + ), + ("def other(input):\n\treturn input", "", "You must define a 'main' function"), + ("def main(input, others):\n\treturn input", "", "The main function should take a single argument as input"), + ], +) +def test_code_node_build_errors(pipeline, code, input, error): nodes = [ start_node(), - code_node("this{}"), + code_node(code), end_node(), ] - with pytest.raises(PipelineNodeBuildError, match="SyntaxError: invalid syntax at statement: 'this{}'"): - create_runnable(pipeline, nodes).invoke(PipelineState(messages=["World"]))["messages"][-1] + with pytest.raises(PipelineNodeBuildError, match=error): + create_runnable(pipeline, nodes).invoke(PipelineState(messages=[input]))["messages"][-1] @django_db_with_data(available_apps=("apps.service_providers",)) @@ -74,8 +93,8 @@ def test_code_node_syntax_error(pipeline): @pytest.mark.parametrize( ("code", "input", "error"), [ - ("import collections", "", "Importing 'collections' is not allowed"), - ("return f'Hello, {blah}!'", "", "name 'blah' is not defined"), + ("import collections\ndef main(input):\n\treturn input", "", "Importing 'collections' is not allowed"), + ("def main(input):\n\treturn f'Hello, {blah}!'", "", "name 'blah' is not defined"), ], ) def test_code_node_runtime_errors(pipeline, code, input, error): From 363d5648a7beb12404c01aeaecbf174c0e35c74c Mon Sep 17 00:00:00 2001 From: Farid Rener Date: Sat, 14 Dec 2024 21:25:59 -0500 Subject: [PATCH 4/5] Show input errors in code modal --- assets/javascript/apps/pipeline/nodes/widgets.tsx | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/assets/javascript/apps/pipeline/nodes/widgets.tsx b/assets/javascript/apps/pipeline/nodes/widgets.tsx index 636cd4182..7842c3f0d 100644 --- a/assets/javascript/apps/pipeline/nodes/widgets.tsx +++ b/assets/javascript/apps/pipeline/nodes/widgets.tsx @@ -290,6 +290,7 @@ export function CodeWidget(props: WidgetParams) { value={concatenate(props.paramValue)} onChange={onChangeCallback} isDarkMode={isDarkMode} + inputError={props.inputError} /> ); @@ -297,12 +298,13 @@ export function CodeWidget(props: WidgetParams) { export function CodeModal( - { modalId, humanName, value, onChange, isDarkMode }: { + { modalId, humanName, value, onChange, isDarkMode, inputError }: { modalId: string; humanName: string; value: string; onChange: (value: string) => void; isDarkMode: boolean; + inputError: string | undefined; }) { return ( +
+ {inputError} +
{/* Allows closing the modal by clicking outside of it */} From 0e00b338733b4beef7cb348d1f6be506ac65b762 Mon Sep 17 00:00:00 2001 From: Farid Rener Date: Sun, 15 Dec 2024 11:48:22 -0500 Subject: [PATCH 5/5] Add **kwargs to default function signature --- apps/pipelines/nodes/nodes.py | 6 +++--- apps/pipelines/tests/data/CodeNode.json | 2 +- apps/pipelines/tests/test_code_node.py | 23 +++++++++++++++-------- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/apps/pipelines/nodes/nodes.py b/apps/pipelines/nodes/nodes.py index 30f86e8cf..629919bd5 100644 --- a/apps/pipelines/nodes/nodes.py +++ b/apps/pipelines/nodes/nodes.py @@ -630,7 +630,7 @@ def _get_assistant_runnable(self, assistant: OpenAiAssistant, session: Experimen DEFAULT_FUNCTION = """# You must define a main function, which takes the node input as a string. # Return a string to pass to the next node. -def main(input: str) -> str: +def main(input: str, **kwargs) -> str: return input """ @@ -670,8 +670,8 @@ def validate_code(cls, value, info: FieldValidationInfo): "You may use nested functions inside that function if required" ) - if len(inspect.signature(main).parameters) > 1: - raise SyntaxError("The main function should take a single argument as input") + if list(inspect.signature(main).parameters) != ["input", "kwargs"]: + raise SyntaxError("The main function should have the signature main(input, **kwargs) only.") except SyntaxError as exc: raise PydanticCustomError("invalid_code", "{error}", {"error": exc.msg}) diff --git a/apps/pipelines/tests/data/CodeNode.json b/apps/pipelines/tests/data/CodeNode.json index 7825da861..e6a926041 100644 --- a/apps/pipelines/tests/data/CodeNode.json +++ b/apps/pipelines/tests/data/CodeNode.json @@ -2,7 +2,7 @@ "description": "Runs python", "properties": { "code": { - "default": "# You must define a main function, which takes the node input as a string.\n# Return a string to pass to the next node.\ndef main(input: str) -> str:\n return input\n", + "default": "# You must define a main function, which takes the node input as a string.\n# Return a string to pass to the next node.\ndef main(input: str, **kwargs) -> str:\n return input\n", "description": "The code to run", "title": "Code", "type": "string", diff --git a/apps/pipelines/tests/test_code_node.py b/apps/pipelines/tests/test_code_node.py index fda70633b..b8fe29b49 100644 --- a/apps/pipelines/tests/test_code_node.py +++ b/apps/pipelines/tests/test_code_node.py @@ -25,7 +25,7 @@ def pipeline(): import datetime import re import time -def main(input): +def main(input, **kwargs): return json.loads(input) """ @@ -35,11 +35,10 @@ def main(input): @pytest.mark.parametrize( ("code", "input", "output"), [ - ("def main(input):\n\treturn f'Hello, {input}!'", "World", "Hello, World!"), + ("def main(input, **kwargs):\n\treturn f'Hello, {input}!'", "World", "Hello, World!"), ("", "foo", "foo"), # No code just returns the input - ("def main(input):\n\t'foo'", "", "None"), # No return value will return "None" + ("def main(input, **kwargs):\n\t'foo'", "", "None"), # No return value will return "None" (IMPORTS, json.dumps({"a": "b"}), str(json.loads('{"a": "b"}'))), # Importing json will work - ("def main(blah):\n\treturn f'Hello, {blah}!'", "World", "Hello, World!"), # Renaming the argument works ], ) def test_code_node(pipeline, code, input, output): @@ -55,7 +54,7 @@ def test_code_node(pipeline, code, input, output): def other(foo): return f"other {foo}" -def main(input): +def main(input, **kwargs): return other(input) """ @@ -75,7 +74,11 @@ def main(input): ), ), ("def other(input):\n\treturn input", "", "You must define a 'main' function"), - ("def main(input, others):\n\treturn input", "", "The main function should take a single argument as input"), + ( + "def main(input, others, **kwargs):\n\treturn input", + "", + r"The main function should have the signature main\(input, \*\*kwargs\) only\.", + ), ], ) def test_code_node_build_errors(pipeline, code, input, error): @@ -93,8 +96,12 @@ def test_code_node_build_errors(pipeline, code, input, error): @pytest.mark.parametrize( ("code", "input", "error"), [ - ("import collections\ndef main(input):\n\treturn input", "", "Importing 'collections' is not allowed"), - ("def main(input):\n\treturn f'Hello, {blah}!'", "", "name 'blah' is not defined"), + ( + "import collections\ndef main(input, **kwargs):\n\treturn input", + "", + "Importing 'collections' is not allowed", + ), + ("def main(input, **kwargs):\n\treturn f'Hello, {blah}!'", "", "name 'blah' is not defined"), ], ) def test_code_node_runtime_errors(pipeline, code, input, error):