diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 8ac157c94f5..01e7964a2ee 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -597,6 +597,7 @@ jobs: tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/deps tests/unit/proxy/config_resolvers + tests/unit/proxy/workflows tests/unit/proxy/utils python-version: "3.12" workers: 4 diff --git a/litellm/proxy/workflows/monty/compiler.py b/litellm/proxy/workflows/monty/compiler.py index 1956f6b6c56..53bf6222ab4 100644 --- a/litellm/proxy/workflows/monty/compiler.py +++ b/litellm/proxy/workflows/monty/compiler.py @@ -2,6 +2,7 @@ import ast import hashlib import keyword import re +from itertools import chain from typing import Final from jsonschema import Draft202012Validator @@ -49,18 +50,48 @@ def _safe_name(name: str) -> bool: return name.isidentifier() and not keyword.iskeyword(name) and "__" not in name and not name.startswith(_PREFIX) +def _module_statement_allowed(statement: ast.stmt) -> bool: + if isinstance(statement, ast.FunctionDef): + return True + if isinstance(statement, ast.Expr) and isinstance(statement.value, ast.Constant): + return isinstance(statement.value.value, str) + return ( + isinstance(statement, ast.ImportFrom) + and statement.module == "__future__" + and statement.level == 0 + and all(alias.name == "annotations" and alias.asname is None for alias in statement.names) + ) + + +def _function_error(node: ast.FunctionDef, functions: tuple[ast.FunctionDef, ...]) -> MontyFailure | None: + if node not in functions: + return _invalid("nested functions and closures are forbidden", node) + if node.decorator_list: + return _invalid("decorators are forbidden", node) + if not _safe_name(node.name) or node.name in _GENERATED_GLOBALS: + return _invalid("reserved function name", node) + return None + + +def _node_error(node: ast.AST, tree: ast.Module, functions: tuple[ast.FunctionDef, ...]) -> MontyFailure | None: + if isinstance(node, (ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda, ast.Global, ast.Nonlocal)): + return _invalid(f"forbidden {type(node).__name__}", node) + if isinstance(node, ast.FunctionDef): + return _function_error(node, functions) + if isinstance(node, (ast.Import, ast.ImportFrom)) and node not in tree.body: + return _invalid("imports inside functions are forbidden", node) + if isinstance(node, ast.Name) and (not _safe_name(node.id) or node.id in _FORBIDDEN): + return _invalid("forbidden builtin or reserved name", node) + if isinstance(node, ast.arg) and not _safe_name(node.arg): + return _invalid("reserved parameter name", node) + if isinstance(node, ast.Attribute) and not _safe_name(node.attr): + return _invalid("dunder or reserved attribute access", node) + return None + + def _source_error(tree: ast.Module, program: PythonProgram) -> MontyFailure | None: for statement in tree.body: - if ( - isinstance(statement, ast.Expr) - and isinstance(statement.value, ast.Constant) - and isinstance(statement.value.value, str) - ): - continue - if isinstance(statement, ast.ImportFrom) and statement.module == "__future__" and statement.level == 0: - if all(alias.name == "annotations" and alias.asname is None for alias in statement.names): - continue - if not isinstance(statement, ast.FunctionDef): + if not _module_statement_allowed(statement): return _invalid("module body permits only function definitions and future annotations", statement) functions: Final = tuple(node for node in tree.body if isinstance(node, ast.FunctionDef)) if len({node.name for node in functions}) != len(functions): @@ -73,27 +104,9 @@ def _source_error(tree: ast.Module, program: PythonProgram) -> MontyFailure | No return _invalid("entry parameters must be positional without defaults or variadics", entries[0]) if tuple(arg.arg for arg in args.args) != program.parameter_names: return _invalid("entry parameter order must match parameter_names", entries[0]) - for node in ast.walk(tree): - if isinstance(node, (ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda, ast.Global, ast.Nonlocal)): - return _invalid(f"forbidden {type(node).__name__}", node) - if isinstance(node, ast.FunctionDef): - if node not in functions: - return _invalid("nested functions and closures are forbidden", node) - if node.decorator_list: - return _invalid("decorators are forbidden", node) - if not _safe_name(node.name) or node.name in _GENERATED_GLOBALS: - return _invalid("reserved function name", node) - if node.args.defaults or any(default is not None for default in node.args.kw_defaults): - return _invalid("function defaults are forbidden", node) - if isinstance(node, (ast.Import, ast.ImportFrom)) and node not in tree.body: - return _invalid("imports inside functions are forbidden", node) - if isinstance(node, ast.Name) and (not _safe_name(node.id) or node.id in _FORBIDDEN): - return _invalid("forbidden builtin or reserved name", node) - if isinstance(node, ast.arg) and not _safe_name(node.arg): - return _invalid("reserved parameter name", node) - if isinstance(node, ast.Attribute) and not _safe_name(node.attr): - return _invalid("dunder or reserved attribute access", node) - return None + return next( + (failure for node in ast.walk(tree) if (failure := _node_error(node, tree, functions)) is not None), None + ) def _mapping(value: JsonValue) -> JsonSchema: @@ -117,9 +130,20 @@ def _resolve(schema: JsonSchema, root: JsonSchema) -> JsonSchema: return descend(root, tuple(ref[2:].split("/"))) -def _schema_error(value: JsonValue) -> str | None: +def _reference_exists(root: JsonValue, parts: tuple[str, ...]) -> bool: + if not parts: + return isinstance(root, (dict, bool)) + part: Final = parts[0].replace("~1", "/").replace("~0", "~") + if isinstance(root, dict) and part in root: + return _reference_exists(root[part], parts[1:]) + if isinstance(root, list) and part.isdigit() and int(part) < len(root): + return _reference_exists(root[int(part)], parts[1:]) + return False + + +def _schema_error(value: JsonValue, root: JsonSchema) -> str | None: if isinstance(value, list): - return next((error for item in value if (error := _schema_error(item)) is not None), None) + return next((error for item in value if (error := _schema_error(item, root)) is not None), None) if not isinstance(value, dict): return None for key, item in value.items(): @@ -127,9 +151,37 @@ def _schema_error(value: JsonValue) -> str | None: return f"unsupported schema keyword {key}" if key == "$ref" and (not isinstance(item, str) or not item.startswith("#/")): return "only local #/ JSON references are supported" - if key not in {"const", "enum", "default", "examples"}: - if (error := _schema_error(item)) is not None: - return error + if key == "$ref" and isinstance(item, str) and not _reference_exists(root, tuple(item[2:].split("/"))): + return "local JSON reference target does not exist" + if key not in {"const", "enum", "default", "examples"} and (error := _schema_error(item, root)) is not None: + return error + return None + + +def _record_shape_error(schema: JsonSchema, root: JsonSchema) -> str | None: + if schema.get("x-liteflow-record") is not True: + return None + if schema not in tuple(record for _, record in _records(root)): + return "records must be defined in root $defs" + if schema.get("type") != "object" or not isinstance(schema.get("properties"), dict): + return "records require fixed object properties" + if schema.get("additionalProperties") is not False: + return "records require additionalProperties false" + if any(not _safe_name(name) for name in _mapping(schema.get("properties"))): + return "record field names must be safe identifiers" + return None + + +def _record_union_error(schema: JsonSchema, root: JsonSchema, seen: frozenset[str]) -> str | None: + for keyword_name in ("anyOf", "oneOf", "allOf"): + if not isinstance(alternatives := schema.get(keyword_name), list): + continue + if not any(_contains_record(_mapping(item), root, seen) for item in alternatives): + continue + if keyword_name != "anyOf" or len(alternatives) != 2: + return "records in unions are supported only as Optional X" + if sum(_mapping(item).get("type") == "null" for item in alternatives) != 1: + return "records in unions are supported only as Optional X" return None @@ -139,34 +191,11 @@ def _record_error(schema: JsonSchema, root: JsonSchema, seen: frozenset[str] = f if ref in seen: return "recursive schema references are unsupported" return _record_error(_resolve(schema, root), root, seen | {ref}) - if schema.get("x-liteflow-record") is True: - if schema.get("type") != "object" or not isinstance(schema.get("properties"), dict): - return "records require fixed object properties" - if schema.get("additionalProperties") is not False: - return "records require additionalProperties false" - if any(not _safe_name(name) for name in _mapping(schema.get("properties"))): - return "record field names must be safe identifiers" - for keyword_name in ("anyOf", "oneOf", "allOf"): - alternatives: Final = schema.get(keyword_name) - if isinstance(alternatives, list): - if any(_contains_record(_mapping(item), root, seen) for item in alternatives): - if keyword_name != "anyOf" or len(alternatives) != 2: - return "records in unions are supported only as Optional X" - if sum(_mapping(item).get("type") == "null" for item in alternatives) != 1: - return "records in unions are supported only as Optional X" - for key, value in schema.items(): - if key in {"properties", "$defs"} and isinstance(value, dict): - for child in value.values(): - if (error := _record_error(_mapping(child), root, seen)) is not None: - return error - if key in {"anyOf", "oneOf", "allOf", "prefixItems"} and isinstance(value, list): - for child_schema in value: - if (error := _record_error(_mapping(child_schema), root, seen)) is not None: - return error - if key in {"items", "additionalProperties"} and isinstance(value, dict): - if (error := _record_error(value, root, seen)) is not None: - return error - return None + if (shape_error := _record_shape_error(schema, root)) is not None: + return shape_error + if (union_error := _record_union_error(schema, root, seen)) is not None: + return union_error + return next((error for child in _children(schema) if (error := _record_error(child, root, seen)) is not None), None) def _contains_record(schema: JsonSchema, root: JsonSchema, seen: frozenset[str]) -> bool: @@ -188,7 +217,7 @@ def _children(schema: JsonSchema) -> tuple[JsonSchema, ...]: return (value,) return () - return tuple(child for group in (children_of(key, value) for key, value in schema.items()) for child in group) + return tuple(chain.from_iterable(children_of(key, value) for key, value in schema.items())) def _records(root: JsonSchema) -> tuple[tuple[str, JsonSchema], ...]: @@ -306,6 +335,39 @@ def _source(program: PythonProgram) -> str: ) +def _schemas_error(program: PythonProgram, tree: ast.Module) -> MontyFailure | None: + function_names: Final = frozenset(node.name for node in tree.body if isinstance(node, ast.FunctionDef)) + for schema in (program.input_schema, program.output_schema): + if (error := _schema_error(schema, schema)) is not None: + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(error)) + Draft202012Validator.check_schema(schema) + if (record_error := _record_error(schema, schema)) is not None: + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(record_error)) + for name, _ in _records(schema): + if not _safe_name(name) or name in _GENERATED_GLOBALS or name in function_names: + return MontyFailure( + MontyErrorKind.UNSUPPORTED_SCHEMA, "record class name must be a distinct safe identifier" + ) + input_records: Final = dict(_records(program.input_schema)) + if any(name in input_records and input_records[name] != schema for name, schema in _records(program.output_schema)): + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, "conflicting record definitions") + return _input_schema_error(program) + + +def _input_schema_error(program: PythonProgram) -> MontyFailure | None: + inputs: Final = program.input_schema + required: Final = inputs.get("required") + if ( + inputs.get("type") != "object" + or set(_mapping(inputs.get("properties"))) != set(program.parameter_names) + or inputs.get("additionalProperties") is not False + or not isinstance(required, list) + or set(required) != set(program.parameter_names) + ): + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, "input schema must require exactly the named parameters") + return None + + def compile_program(program: PythonProgram, limits: MontyLimits) -> CompiledProgram | MontyFailure: try: encoded: Final = program.source.encode("utf-8") @@ -318,37 +380,8 @@ def compile_program(program: PythonProgram, limits: MontyLimits) -> CompiledProg tree: Final = ast.parse(program.source) if (failure := _source_error(tree, program)) is not None: return failure - for schema in (program.input_schema, program.output_schema): - if (error := _schema_error(schema)) is not None: - return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(error)) - Draft202012Validator.check_schema(schema) - if (record_error := _record_error(schema, schema)) is not None: - return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(record_error)) - for name, record in _records(schema): - if ( - not _safe_name(name) - or name in _GENERATED_GLOBALS - or name in {node.name for node in tree.body if isinstance(node, ast.FunctionDef)} - ): - return MontyFailure( - MontyErrorKind.UNSUPPORTED_SCHEMA, "record class name must be a distinct safe identifier" - ) - input_records: Final = dict(_records(program.input_schema)) - if any( - name in input_records and input_records[name] != schema for name, schema in _records(program.output_schema) - ): - return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, "conflicting record definitions") - inputs: Final = program.input_schema - if ( - inputs.get("type") != "object" - or set(_mapping(inputs.get("properties"))) != set(program.parameter_names) - or inputs.get("additionalProperties") is not False - or not isinstance(inputs.get("required"), list) - or set(inputs["required"]) != set(program.parameter_names) - ): - return MontyFailure( - MontyErrorKind.UNSUPPORTED_SCHEMA, "input schema must require exactly the named parameters" - ) + if (schema_failure := _schemas_error(program, tree)) is not None: + return schema_failure return CompiledProgram(program, _source(program), program.input_schema, program.output_schema) except SyntaxError as exc: return MontyFailure( diff --git a/litellm/proxy/workflows/monty/models.py b/litellm/proxy/workflows/monty/models.py index b41a2291719..62abae01e72 100644 --- a/litellm/proxy/workflows/monty/models.py +++ b/litellm/proxy/workflows/monty/models.py @@ -1,22 +1,32 @@ +import keyword from dataclasses import dataclass from enum import StrEnum -from typing import Annotated, Literal +from typing import Annotated, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, StringConstraints -JsonSchema = dict[str, JsonValue] +JsonSchema: TypeAlias = dict[str, JsonValue] + + +def _identifier(value: str) -> str: + if not value.isidentifier() or keyword.iskeyword(value): + raise ValueError("must be a Python identifier") + return value + + +ProgramName: TypeAlias = Annotated[str, StringConstraints(min_length=1, max_length=128), AfterValidator(_identifier)] class PythonProgram(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = ConfigDict(frozen=True, extra="forbid") contract_version: Literal["1"] source: str - function_name: str - parameter_names: Annotated[tuple[str, ...], Field(min_length=1)] + function_name: ProgramName + parameter_names: Annotated[tuple[ProgramName, ...], Field(min_length=1, max_length=32)] input_schema: JsonSchema output_schema: JsonSchema - source_sha256: str + source_sha256: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] class MontyErrorKind(StrEnum): @@ -60,17 +70,17 @@ class CompiledProgram: class MontyLimits(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) - max_source_bytes: int = Field(default=64 * 1024, gt=0) - max_input_bytes: int = Field(default=1024 * 1024, gt=0) - max_output_bytes: int = Field(default=1024 * 1024, gt=0) - max_input_depth: int = Field(default=64, gt=0) - max_output_depth: int = Field(default=64, gt=0) - max_memory: int = Field(default=64 * 1024 * 1024, gt=0) - worker_address_space_bytes: int = Field(default=512 * 1024 * 1024, gt=0) - feed_seconds: float = Field(default=5, gt=0) - max_recursion_depth: int = Field(default=200, gt=0) - max_stdout_bytes: int = Field(default=64 * 1024, gt=0) - max_concurrent_workers: int = Field(default=4, gt=0) - max_queued_waiters: int = Field(default=16, ge=0) + max_source_bytes: int = Field(default=64 * 1024, gt=0, le=1024 * 1024) + max_input_bytes: int = Field(default=1024 * 1024, gt=0, le=16 * 1024 * 1024) + max_output_bytes: int = Field(default=1024 * 1024, gt=0, le=16 * 1024 * 1024) + max_input_depth: int = Field(default=64, gt=0, le=128) + max_output_depth: int = Field(default=64, gt=0, le=128) + max_memory: int = Field(default=64 * 1024 * 1024, gt=0, le=1024 * 1024 * 1024) + worker_address_space_bytes: int = Field(default=512 * 1024 * 1024, gt=0, le=8 * 1024 * 1024 * 1024) + feed_seconds: float = Field(default=5, gt=0, le=60) + max_recursion_depth: int = Field(default=200, gt=0, le=1000) + max_stdout_bytes: int = Field(default=64 * 1024, gt=0, le=1024 * 1024) + max_concurrent_workers: int = Field(default=4, gt=0, le=64) + max_queued_waiters: int = Field(default=16, gt=0, le=1024) diff --git a/litellm/proxy/workflows/monty/runtime.py b/litellm/proxy/workflows/monty/runtime.py index e4f81c354c4..ed0fefa0f7a 100644 --- a/litellm/proxy/workflows/monty/runtime.py +++ b/litellm/proxy/workflows/monty/runtime.py @@ -17,14 +17,15 @@ from typing import TYPE_CHECKING, Final from jsonschema import Draft202012Validator from jsonschema.exceptions import ValidationError from pydantic import JsonValue, TypeAdapter +from pydantic import ValidationError as PydanticValidationError from .compiler import escape_diagnostics from .models import CompiledProgram, JsonSchema, MontyErrorKind, MontyFailure, MontyLimits, MontySuccess if TYPE_CHECKING: - from pydantic_monty import AsyncMonty, AsyncMontySession + from pydantic_monty import AsyncMonty, AsyncMontySession, CollectString, MontyRuntimeError -_JSON: Final = TypeAdapter(JsonValue) +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @dataclass(frozen=True, slots=True) @@ -43,12 +44,12 @@ def installed_binary() -> BinaryArtifact | MontyFailure: if file.hash is None or file.hash.mode != "sha256": return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "runtime binary has no SHA-256 RECORD hash") digest: Final = base64.urlsafe_b64decode(file.hash.value + "=" * (-len(file.hash.value) % 4)).hex() - return BinaryArtifact(Path(distribution.locate_file(file)).resolve(), digest) + return BinaryArtifact(Path(str(distribution.locate_file(file))).resolve(), digest) except (importlib.metadata.PackageNotFoundError, OSError, ValueError) as exc: return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc))) -def _json_issue(value: object, depth: int, max_depth: int) -> MontyErrorKind | None: +def _json_issue(value: JsonValue, depth: int, max_depth: int) -> MontyErrorKind | None: if depth > max_depth: return MontyErrorKind.INPUT_TOO_LARGE if value is None or type(value) in (str, int, bool): @@ -61,7 +62,12 @@ def _json_issue(value: object, depth: int, max_depth: int) -> MontyErrorKind | N if any(type(key) is not str for key in value): return MontyErrorKind.INPUT_INVALID return next( - (issue for item in value.values() if (issue := _json_issue(item, depth + 1, max_depth)) is not None), None + ( + child_issue + for item in value.values() + if (child_issue := _json_issue(item, depth + 1, max_depth)) is not None + ), + None, ) return MontyErrorKind.INPUT_INVALID @@ -71,20 +77,26 @@ def validate_json( ) -> JsonValue | MontyFailure: invalid: Final = MontyErrorKind.OUTPUT_INVALID if output else MontyErrorKind.INPUT_INVALID too_large: Final = MontyErrorKind.OUTPUT_TOO_LARGE if output else MontyErrorKind.INPUT_TOO_LARGE - issue: Final = _json_issue(value, 0, max_depth) - if issue is not None: - return MontyFailure( - too_large if issue == MontyErrorKind.INPUT_TOO_LARGE else invalid, "JSON depth or type invalid" - ) try: - encoded: Final = json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8") + parsed: Final = _JSON.validate_python(value, strict=True) + issue: Final = _json_issue(parsed, 0, max_depth) + if issue is not None: + return MontyFailure( + too_large if issue == MontyErrorKind.INPUT_TOO_LARGE else invalid, "JSON depth or type invalid" + ) + encoded: Final = json.dumps(parsed, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8") if len(encoded) > max_bytes: return MontyFailure(too_large, "JSON byte limit exceeded") - parsed: Final = _JSON.validate_python(value, strict=True) - Draft202012Validator(schema).validate(parsed) - return parsed + validator: Final = Draft202012Validator(schema) + validator.validate(parsed) # pyright: ignore[reportUnknownMemberType] # upstream kwargs are untyped + except PydanticValidationError as exc: + if '"type":"recursion_loop"' in exc.json(): + return MontyFailure(too_large, "JSON depth limit exceeded") + return MontyFailure(invalid, escape_diagnostics(str(exc))[:4096]) except (ValidationError, ValueError, OverflowError, RecursionError) as exc: return MontyFailure(invalid, escape_diagnostics(str(exc))[:4096]) + else: + return parsed def _kill(session: AsyncMontySession) -> None: @@ -97,6 +109,40 @@ def _kill(session: AsyncMontySession) -> None: return +def _worker_limit(session: AsyncMontySession, limits: MontyLimits) -> MontyFailure | None: + if sys.platform != "linux": + return None + import resource + + pid: Final = session.worker_pid + if pid is None: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "worker PID unavailable") + try: + resource.prlimit( + pid, resource.RLIMIT_AS, (limits.worker_address_space_bytes, limits.worker_address_space_bytes) + ) + except (OSError, ValueError) as exc: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc))) + return None + + +def _runtime_failure(exc: MontyRuntimeError, collector: CollectString, limits: MontyLimits) -> MontyFailure: + error: Final = exc.exception() + message: Final = escape_diagnostics(exc.display())[:4096] + if isinstance(error, TimeoutError): + return MontyFailure(MontyErrorKind.TIMEOUT, message) + if isinstance(error, MemoryError): + kind: Final = ( + MontyErrorKind.STDOUT_LIMIT + if len(collector.output.encode("utf-8")) >= limits.max_stdout_bytes and not exc.traceback() + else MontyErrorKind.MEMORY_LIMIT + ) + return MontyFailure(kind, message) + if isinstance(error, RecursionError): + return MontyFailure(MontyErrorKind.RECURSION_LIMIT, message) + return MontyFailure(MontyErrorKind.RUNTIME_ERROR, message) + + class MontyRuntime: def __init__(self, limits: MontyLimits, *, binary: BinaryArtifact | None = None) -> None: self.limits = limits @@ -127,9 +173,10 @@ class MontyRuntime: ) await pool.__aenter__() self._pool = pool - return None except (ImportError, OSError, ValueError, RuntimeError) as exc: return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc))) + else: + return None async def close(self) -> None: pool: Final = self._pool @@ -189,51 +236,22 @@ class MontyRuntime: try: async with session: try: - if sys.platform == "linux": - import resource - - pid: Final = session.worker_pid - if pid is None: - return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "worker PID unavailable") - try: - resource.prlimit(pid, resource.RLIMIT_AS, (self.limits.worker_address_space_bytes,) * 2) - except (OSError, ValueError) as exc: - return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc))) - result: Final[object] = await session.feed_run( - compiled.sandbox_source, - inputs={"_liteflow_input": input}, - print_callback=collector, - ) - if not isinstance(result, list) or len(result) != 2 or result[0] is not True: - return MontyFailure(MontyErrorKind.OUTPUT_INVALID, "return value is not strict JSON") - output: Final = validate_json( - result[1], - compiled.output_schema, - max_bytes=self.limits.max_output_bytes, - max_depth=self.limits.max_output_depth, - output=True, - ) - if isinstance(output, MontyFailure): - return output - if time.monotonic() >= deadline: - return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired during validation") - return MontySuccess( - output, escape_diagnostics(collector.output), int((time.monotonic() - started) * 1000) + if (limit_failure := _worker_limit(session, self.limits)) is not None: + return limit_failure + result: Final = _JSON.validate_python( + await session.feed_run( + compiled.sandbox_source, + inputs={"_liteflow_input": input}, + print_callback=collector, + ), + strict=True, ) + return self._output(result, compiled, collector.output, deadline, started) except asyncio.CancelledError: _kill(session) raise except MontyRuntimeError as exc: - error: Final = exc.exception() - message: Final = escape_diagnostics(exc.display())[:4096] - if isinstance(error, TimeoutError): - return MontyFailure(MontyErrorKind.TIMEOUT, message) - if isinstance(error, MemoryError): - kind: Final = MontyErrorKind.STDOUT_LIMIT if "bytes >" in str(error) else MontyErrorKind.MEMORY_LIMIT - return MontyFailure(kind, message) - if isinstance(error, RecursionError): - return MontyFailure(MontyErrorKind.RECURSION_LIMIT, message) - return MontyFailure(MontyErrorKind.RUNTIME_ERROR, message) + return _runtime_failure(exc, collector, self.limits) except MontyCrashedError as exc: return MontyFailure( MontyErrorKind.TIMEOUT if exc.timed_out else MontyErrorKind.WORKER_CRASHED, @@ -244,6 +262,28 @@ class MontyRuntime: except (ValueError, OSError) as exc: return MontyFailure(MontyErrorKind.RUNTIME_ERROR, escape_diagnostics(str(exc))[:4096]) + def _output( + self, result: object, compiled: CompiledProgram, stdout: str, deadline: float, started: float + ) -> MontySuccess | MontyFailure: + try: + envelope: Final = _JSON.validate_python(result, strict=True) + except ValueError: + return MontyFailure(MontyErrorKind.OUTPUT_INVALID, "return value is not strict JSON") + if not isinstance(envelope, list) or len(envelope) != 2 or envelope[0] is not True: + return MontyFailure(MontyErrorKind.OUTPUT_INVALID, "return value is not strict JSON") + output: Final = validate_json( + envelope[1], + compiled.output_schema, + max_bytes=self.limits.max_output_bytes, + max_depth=self.limits.max_output_depth, + output=True, + ) + if isinstance(output, MontyFailure): + return output + if time.monotonic() >= deadline: + return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired during validation") + return MontySuccess(output, escape_diagnostics(stdout), int((time.monotonic() - started) * 1000)) + async def evaluate_predicate( self, compiled: CompiledProgram, input: object, *, deadline_monotonic: float ) -> bool | MontyFailure: diff --git a/pyproject.toml b/pyproject.toml index cb3e475b1d9..bb7ac808b90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -258,6 +258,7 @@ admin-mcp = [ "litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/31fb096904ae05684c0779110a93b135b5601182.tar.gz#sha256=5803a20a00a18177f4591706956f31be311879543049b8b38143ad90bb1cb0a6 ; python_version >= '3.12'", ] proxy-dev = [ + "pydantic-monty==1.4.0", { include-group = "admin-mcp" }, "prisma==0.11.0", "hypercorn==0.17.3", diff --git a/tests/integration/sandbox/test_monty_runtime.py b/tests/integration/sandbox/test_monty_runtime.py index b3a6feef70b..c58e5f4531c 100644 --- a/tests/integration/sandbox/test_monty_runtime.py +++ b/tests/integration/sandbox/test_monty_runtime.py @@ -1,14 +1,17 @@ import asyncio import hashlib import os -import signal import time from collections.abc import AsyncIterator +from dataclasses import replace from pathlib import Path from typing import Final +import psutil import pytest import pytest_asyncio +from openai.types.chat import ChatCompletionMessageParam +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.proxy.workflows.monty.compiler import compile_program from litellm.proxy.workflows.monty.models import ( @@ -21,6 +24,7 @@ from litellm.proxy.workflows.monty.models import ( PythonProgram, ) from litellm.proxy.workflows.monty.runtime import BinaryArtifact, MontyRuntime, installed_binary +from tests.integration._support.process import stop_processes pytest.importorskip("pydantic_monty") @@ -53,7 +57,7 @@ def _compile( @pytest_asyncio.fixture async def runtime() -> AsyncIterator[MontyRuntime]: - instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1, max_queued_waiters=0)) + instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1, max_queued_waiters=1)) assert await instance.start() is None try: yield instance @@ -256,7 +260,7 @@ async def test_expired_deadline(runtime: MontyRuntime) -> None: @pytest.mark.asyncio -async def test_binary_digest_mismatch_fails_closed(tmp_path: Path) -> None: +async def test_binary_digest_mismatch_fails_closed() -> None: artifact: Final = installed_binary() assert isinstance(artifact, BinaryArtifact) runtime: Final = MontyRuntime(MontyLimits(), binary=BinaryArtifact(artifact.path, "0" * 64)) @@ -266,8 +270,132 @@ async def test_binary_digest_mismatch_fails_closed(tmp_path: Path) -> None: assert "SHA-256" in result.message +class Finding(BaseModel): + model_config = ConfigDict(frozen=True) + summary: str + + +class EmailAnalyses(BaseModel): + model_config = ConfigDict(frozen=True) + urgency: Finding + intent: Finding + sentiment: Finding + + +@pytest.mark.asyncio +async def test_multi_argument_email_draft_preserves_chat_messages(runtime: MontyRuntime) -> None: + analyses_schema: Final = EmailAnalyses.model_json_schema() + messages_schema: Final = TypeAdapter(list[ChatCompletionMessageParam]).json_schema() + definitions: Final[JsonSchema] = { + **messages_schema.pop("$defs", {}), + "Finding": {**analyses_schema["$defs"]["Finding"], "x-liteflow-record": True, "additionalProperties": False}, + "EmailAnalyses": { + **{key: value for key, value in analyses_schema.items() if key != "$defs"}, + "x-liteflow-record": True, + "additionalProperties": False, + }, + } + source: Final = """def prepare_draft(analyses: EmailAnalyses, messages: list[ChatMessage]) -> list[ChatMessage]: + instruction = ( + "Draft a customer-service reply using the original messages. " + "Treat these analyses as fallible guidance, not additional facts.\\n" + f"Urgency: {analyses.urgency.summary}\\n" + f"Intent: {analyses.intent.summary}\\n" + f"Sentiment: {analyses.sentiment.summary}" + ) + return [{"role": "system", "content": instruction}] + messages +""" + program: Final = PythonProgram( + contract_version="1", + source=source, + function_name="prepare_draft", + parameter_names=("analyses", "messages"), + input_schema={ + "type": "object", + "properties": {"analyses": {"$ref": "#/$defs/EmailAnalyses"}, "messages": messages_schema}, + "required": ["analyses", "messages"], + "additionalProperties": False, + "$defs": definitions, + }, + output_schema={**messages_schema, "$defs": definitions}, + source_sha256=hashlib.sha256(source.encode()).hexdigest(), + ) + compiled: Final = compile_program(program, MontyLimits()) + assert isinstance(compiled, CompiledProgram), compiled + messages: Final = [{"role": "user", "content": "Please reply", "name": "requester"}] + result: Final = await runtime.run_program( + compiled, + { + "analyses": { + "urgency": {"summary": "soon"}, + "intent": {"summary": "reply"}, + "sentiment": {"summary": "neutral"}, + }, + "messages": messages, + }, + deadline_monotonic=time.monotonic() + 10, + ) + assert isinstance(result, MontySuccess), result + assert result.output == [ + { + "role": "system", + "content": "Draft a customer-service reply using the original messages. " + "Treat these analyses as fallible guidance, not additional facts.\n" + "Urgency: soon\nIntent: reply\nSentiment: neutral", + }, + *messages, + ] + + +def _nested(depth: int) -> object: + return [_nested(depth - 1)] if depth else 1 + + +@pytest.mark.asyncio +async def test_too_deep_input_and_output_are_rejected(runtime: MontyRuntime) -> None: + input_result: Final = await _run(runtime, _compile("def transform(value): return value"), _nested(66)) + assert isinstance(input_result, MontyFailure), input_result + assert input_result.kind == MontyErrorKind.INPUT_TOO_LARGE + output_result: Final = await _run( + runtime, _compile("def transform(value):\n for i in range(66): value = [value]\n return value"), 1 + ) + assert isinstance(output_result, MontyFailure), output_result + assert output_result.kind == MontyErrorKind.OUTPUT_TOO_LARGE + + +@pytest.mark.asyncio +async def test_invalid_bound_input_is_rejected(runtime: MontyRuntime) -> None: + result: Final = await _run( + runtime, _compile("def transform(value): return value", argument={"type": "integer"}), "wrong" + ) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.INPUT_INVALID + + +@pytest.mark.asyncio +async def test_sandbox_denies_file_access_even_if_compiler_is_bypassed(runtime: MontyRuntime) -> None: + compiled: Final = _compile("def transform(value): return 42") + unfiltered: Final = replace(compiled, sandbox_source="open('/etc/passwd').read()") + result: Final = await _run(runtime, unfiltered, None) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.RUNTIME_ERROR + + +@pytest.mark.asyncio +async def test_host_deadline_cancels_running_worker_and_releases_slot(runtime: MontyRuntime) -> None: + compiled: Final = _compile("def transform(value):\n while True: pass") + result: Final = await runtime.run_program(compiled, {"value": None}, deadline_monotonic=time.monotonic() + 0.05) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.TIMEOUT + next_result: Final = await _run(runtime, _compile("def transform(value): return value"), 4) + assert isinstance(next_result, MontySuccess), next_result + assert next_result.output == 4 + + def _worker_pids() -> frozenset[int]: - children: Final = Path(f"/proc/{os.getpid()}/task/{os.getpid()}/children").read_text().split() + children: Final = " ".join( + path.read_text() for path in Path(f"/proc/{os.getpid()}/task").glob("*/children") + ).split() return frozenset(int(pid) for pid in children if Path(f"/proc/{pid}/comm").read_text().strip() == "monty") @@ -285,9 +413,14 @@ async def test_cancellation_kills_worker_releases_slot_and_queue_is_bounded(runt previous: Final = _worker_pids() task: Final = asyncio.create_task(_run(runtime, _compile("def transform(value):\n while True: pass"), None)) pid: Final = await _active_worker(previous) + waiter: Final = asyncio.create_task(_run(runtime, _compile("def transform(value): return value"), 1)) + await asyncio.sleep(0) full: Final = await _run(runtime, _compile("def transform(value): return value"), 1) assert isinstance(full, MontyFailure) assert full.kind == MontyErrorKind.QUEUE_FULL + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -302,7 +435,7 @@ async def test_killed_worker_fails_and_next_run_succeeds(runtime: MontyRuntime) previous: Final = _worker_pids() task: Final = asyncio.create_task(_run(runtime, _compile("def transform(value):\n while True: pass"), None)) pid: Final = await _active_worker(previous) - os.kill(pid, signal.SIGKILL) + await asyncio.to_thread(stop_processes, (psutil.Process(pid),), 0.1) result: Final = await task assert isinstance(result, MontyFailure), result assert result.kind == MontyErrorKind.WORKER_CRASHED diff --git a/tests/unit/proxy/workflows/monty/test_compiler.py b/tests/unit/proxy/workflows/monty/test_compiler.py index e38f2327f50..f9b62812735 100644 --- a/tests/unit/proxy/workflows/monty/test_compiler.py +++ b/tests/unit/proxy/workflows/monty/test_compiler.py @@ -55,16 +55,12 @@ def _program(source: str, *, output: JsonSchema | None = None) -> PythonProgram: "def transform(*, value): return value", "def transform(value, /): return value", "def transform(value):\n import json\n return value", - "def transform(value): return", "def transform(value): return value\ndef isinstance(x,y): return True", ], ) def test_rejects_unsupported_source_without_running_it(source: str) -> None: program: Final = _program(source) result: Final = compile_program(program, MontyLimits()) - if source == "def transform(value): return": - assert isinstance(result, CompiledProgram) - return assert isinstance(result, MontyFailure) assert result.kind == MontyErrorKind.INVALID_PROGRAM assert "line" in result.message @@ -103,6 +99,7 @@ def test_rejects_dangerous_builtin_references(builtin: str) -> None: {"$id": "https://example.invalid/schema"}, {"$dynamicRef": "#/x"}, {"$schema": "https://example.invalid/schema"}, + {"$ref": "#/$defs/Missing"}, {"type": "not-a-json-type"}, {"$defs": {"Bad": {"type": "object", "x-liteflow-record": True, "properties": {}}}}, { diff --git a/tests/unit/proxy/workflows/monty/test_models.py b/tests/unit/proxy/workflows/monty/test_models.py new file mode 100644 index 00000000000..cb059f63e3d --- /dev/null +++ b/tests/unit/proxy/workflows/monty/test_models.py @@ -0,0 +1,55 @@ +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.proxy.workflows.monty.models import MontyLimits, PythonProgram + + +def _artifact() -> PythonProgram: + return PythonProgram( + contract_version="1", + source="def transform(value): return value", + function_name="transform", + parameter_names=("value",), + input_schema={}, + output_schema={}, + source_sha256="a" * 64, + ) + + +def test_program_rejects_unknown_fields() -> None: + with pytest.raises(ValidationError, match="extra_forbidden"): + PythonProgram.model_validate({**_artifact().model_dump(), "credentials": "not-allowed"}) + + +def test_limits_reject_unknown_fields() -> None: + with pytest.raises(ValidationError, match="extra_forbidden"): + MontyLimits.model_validate({"mount": "/"}) + + +@pytest.mark.parametrize("field", tuple(MontyLimits.model_fields)) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf"), -1, 0, 10**20]) +def test_every_limit_is_positive_finite_and_bounded(field: str, value: float | int) -> None: + with pytest.raises(ValidationError): + MontyLimits.model_validate({field: value}) + + +@pytest.mark.parametrize( + "field,value", + [ + ("function_name", "not valid"), + ("function_name", "class"), + ("function_name", "x" * 129), + ("parameter_names", ("bad-name",)), + ("parameter_names", ("x" * 129,)), + ("parameter_names", tuple(f"p{i}" for i in range(33))), + ("source_sha256", "A" * 64), + ("source_sha256", "a" * 63), + ("source_sha256", "g" * 64), + ], +) +def test_artifact_names_count_and_digest_are_bounded(field: str, value: object) -> None: + payload: Final = {**_artifact().model_dump(), field: value} + with pytest.raises(ValidationError): + PythonProgram.model_validate(payload) diff --git a/uv.lock b/uv.lock index 6d4711abd26..3dc5ec031b1 100644 --- a/uv.lock +++ b/uv.lock @@ -5261,6 +5261,7 @@ proxy-dev = [ { name = "opentelemetry-sdk" }, { name = "prisma" }, { name = "prometheus-client" }, + { name = "pydantic-monty" }, { name = "semver" }, { name = "sentry-sdk" }, ] @@ -5496,6 +5497,7 @@ proxy-dev = [ { name = "opentelemetry-sdk", specifier = "==1.45.1" }, { name = "prisma", specifier = "==0.11.0" }, { name = "prometheus-client", specifier = "==0.20.0" }, + { name = "pydantic-monty", specifier = "==1.4.0" }, { name = "semver", specifier = ">=3.0.4,<4" }, { name = "sentry-sdk", specifier = "==2.21.0" }, ]