diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 41198887cab..958546dde00 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -547,6 +547,7 @@ jobs: - shard: proxy-endpoints test-path: |- + tests/unit/proxy/workflows tests/unit/proxy/management_endpoints tests/unit/proxy/management_helpers tests/unit/proxy/list_api diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106010000_workflow_definitions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106010000_workflow_definitions/migration.sql new file mode 100644 index 00000000000..d521c665597 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106010000_workflow_definitions/migration.sql @@ -0,0 +1,52 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE "LiteLLM_WorkflowDefinition" ( + "workflow_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "etag" TEXT NOT NULL, + "digest" TEXT NOT NULL, + "bundle_json" TEXT NOT NULL, + "contributors" TEXT[], + "published_snapshot_id" TEXT, + "published_etag" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + CONSTRAINT "LiteLLM_WorkflowDefinition_pkey" PRIMARY KEY ("workflow_id") +); + +CREATE TABLE "LiteLLM_WorkflowSnapshot" ( + "snapshot_id" TEXT NOT NULL, + "workflow_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "digest" TEXT NOT NULL, + "bundle_json" TEXT NOT NULL, + "contributors" TEXT[], + "published_by" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT "LiteLLM_WorkflowSnapshot_pkey" PRIMARY KEY ("snapshot_id"), + CONSTRAINT "LiteLLM_WorkflowSnapshot_workflow_id_fkey" FOREIGN KEY ("workflow_id") + REFERENCES "LiteLLM_WorkflowDefinition"("workflow_id") ON DELETE RESTRICT ON UPDATE CASCADE +); + +CREATE UNIQUE INDEX "LiteLLM_WorkflowDefinition_team_id_name_key" ON "LiteLLM_WorkflowDefinition"("team_id", "name"); +CREATE INDEX "LiteLLM_WorkflowDefinition_team_id_workflow_id_idx" ON "LiteLLM_WorkflowDefinition"("team_id", "workflow_id"); +CREATE INDEX "LiteLLM_WorkflowSnapshot_workflow_id_created_at_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("workflow_id", "created_at", "snapshot_id"); +CREATE INDEX "LiteLLM_WorkflowSnapshot_team_id_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("team_id", "snapshot_id"); + +ALTER TABLE "LiteLLM_WorkflowRun" + ADD COLUMN "workflow_id" TEXT, + ADD COLUMN "snapshot_id" TEXT, + ADD COLUMN "team_id" TEXT, + ADD COLUMN "engine_owned" BOOLEAN NOT NULL DEFAULT false, + ADD COLUMN "request_digest" TEXT, + ADD COLUMN "idempotency_key" TEXT, + ADD COLUMN "lease_token" TEXT, + ADD COLUMN "lease_expires_at" TIMESTAMP(3), + ADD COLUMN "cancel_requested" BOOLEAN NOT NULL DEFAULT false; + +CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_principal_idempotency_key" ON "LiteLLM_WorkflowRun"("created_by", "team_id", "workflow_id", "idempotency_key"); +CREATE INDEX "LiteLLM_WorkflowRun_workflow_id_team_id_created_at_run_id_idx" ON "LiteLLM_WorkflowRun"("workflow_id", "team_id", "created_at", "run_id"); +CREATE INDEX "LiteLLM_WorkflowRun_engine_owned_status_lease_expires_at_idx" ON "LiteLLM_WorkflowRun"("engine_owned", "status", "lease_expires_at"); +CREATE INDEX "LiteLLM_WorkflowRun_team_id_engine_owned_status_idx" ON "LiteLLM_WorkflowRun"("team_id", "engine_owned", "status"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 818f2dc4858..c0dc2eb960b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1975,10 +1975,53 @@ model LiteLLM_ShadowEvalFunnel { // One instance of work being done. `status` is a materialized cache of the // latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowDefinition { + workflow_id String @id @default(uuid()) + team_id String + name String + etag String + digest String + bundle_json String + contributors String[] + published_snapshot_id String? + published_etag String? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + snapshots LiteLLM_WorkflowSnapshot[] + + @@unique([team_id, name]) + @@index([team_id, workflow_id]) +} + +model LiteLLM_WorkflowSnapshot { + snapshot_id String @id @default(uuid()) + workflow_id String + team_id String + name String + digest String + bundle_json String + contributors String[] + published_by String? + created_at DateTime @default(now()) + definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict) + + @@index([workflow_id, created_at, snapshot_id]) + @@index([team_id, snapshot_id]) +} + model LiteLLM_WorkflowRun { run_id String @id @default(uuid()) session_id String @unique @default(uuid()) workflow_type String + workflow_id String? + snapshot_id String? + team_id String? + engine_owned Boolean @default(false) + request_digest String? + idempotency_key String? + lease_token String? + lease_expires_at DateTime? + cancel_requested Boolean @default(false) status String @default("pending") created_by String? // user_id of the key that created this run; null = created by master key created_at DateTime @default(now()) @@ -1994,6 +2037,10 @@ model LiteLLM_WorkflowRun { @@index([session_id]) @@index([created_at]) @@index([created_by]) + @@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key") + @@index([workflow_id, team_id, created_at, run_id]) + @@index([engine_owned, status, lease_expires_at]) + @@index([team_id, engine_owned, status]) } // Append-only log of state transitions. Never mutate rows here. diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 0f98a8fce0e..469b4d335f8 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -355,6 +355,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_ShadowEvalJob", "LiteLLM_ShadowEvalAttempt", "LiteLLM_ShadowEvalFunnel", + "LiteLLM_WorkflowDefinition", + "LiteLLM_WorkflowSnapshot", "LiteLLM_WorkflowRun", "LiteLLM_WorkflowEvent", "LiteLLM_WorkflowMessage", diff --git a/litellm/models/workflow_definition.py b/litellm/models/workflow_definition.py new file mode 100644 index 00000000000..8a0a7907a76 --- /dev/null +++ b/litellm/models/workflow_definition.py @@ -0,0 +1,170 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Annotated, Literal, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, StringConstraints + +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy.workflows.monty.models import PythonProgram + +WorkflowName: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$")] +Digest: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + + +class ValueBinding(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + value: JsonValue + + +class ReferenceBinding(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + source: str = Field(alias="from", min_length=1, max_length=128) + output: Literal["messages", "result"] + + +Binding: TypeAlias = ValueBinding | ReferenceBinding + + +class WorkflowRetry(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + max_attempts: int = Field(default=1, ge=1, le=3, strict=True) + backoff_seconds: float = Field(default=0, ge=0, le=30) + + +class StructuredSchema(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + name: str = Field(min_length=1, max_length=128) + strict: Literal[True] + schema_: Mapping[str, JsonValue] = Field(alias="schema") + + +class StructuredOutput(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["json_schema"] + json_schema: StructuredSchema + + +class ChatCompletionNode(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + id: WorkflowName + type: Literal["chat_completion"] + model: str = Field(min_length=1, max_length=256) + messages: Binding + timeout_seconds: float = Field(gt=0, le=300) + retry: WorkflowRetry + temperature: float | None = Field(default=None, ge=0, le=2) + top_p: float | None = Field(default=None, ge=0, le=1) + max_tokens: int | None = Field(default=None, gt=0, strict=True) + max_completion_tokens: int | None = Field(default=None, gt=0, strict=True) + seed: int | None = Field(default=None, strict=True) + stop: str | tuple[str, ...] | None = None + response_format: StructuredOutput | None = None + + +class PythonNode(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + id: WorkflowName + type: Literal["python"] + program: PythonProgram + inputs: Mapping[str, Binding] + timeout_seconds: float = Field(gt=0, le=300) + retry: WorkflowRetry + + +WorkflowNode: TypeAlias = Annotated[ChatCompletionNode | PythonNode, Field(discriminator="type")] + + +class WorkflowBundle(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + contract_version: Literal["1"] + name: WorkflowName + timeout_seconds: float = Field(gt=0, le=1800) + input_schema: Mapping[str, JsonValue] + output: ReferenceBinding + nodes: tuple[WorkflowNode, ...] = Field(min_length=1, max_length=64) + + +class WorkflowDraft(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + workflow_id: str + team_id: str + etag: str + digest: Digest + bundle: WorkflowBundle + published_snapshot_id: str | None + + +@dataclass(frozen=True, slots=True) +class WorkflowDefinition: + draft: WorkflowDraft + contributors: tuple[str, ...] + + +class WorkflowSnapshot(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + snapshot_id: str + workflow_id: str + team_id: str + name: str + digest: Digest + bundle: WorkflowBundle + contributors: tuple[str, ...] + published_by: str | None + created_at: datetime + + +class DefinitionPage(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + workflows: tuple[WorkflowDraft, ...] + next_cursor: str | None + + +DefinitionErrorCode: TypeAlias = Literal[ + "not_found", "forbidden", "team_required", "invalid_bundle", "conflict", "stale_draft", "precondition_required" +] + + +@dataclass(frozen=True, slots=True) +class DefinitionError: + code: DefinitionErrorCode + message: str + param: str | None = None + + +class WorkflowDefinitionStore(Protocol): + async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None: ... + + async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None: ... + + async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage: ... + + async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None: ... + + async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None: ... + + async def create( + self, bundle: WorkflowBundle, team_id: str, contributor: str + ) -> WorkflowDraft | DefinitionError: ... + + async def replace( + self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str + ) -> WorkflowDraft | DefinitionError: ... + + async def publish( + self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str + ) -> WorkflowSnapshot | DefinitionError: ... + + async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: ... diff --git a/litellm/models/workflow_execution.py b/litellm/models/workflow_execution.py new file mode 100644 index 00000000000..28cce6dfb1b --- /dev/null +++ b/litellm/models/workflow_execution.py @@ -0,0 +1,171 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Protocol +from uuid import NAMESPACE_URL, uuid5 + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + + +class RunState(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class RunInput(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256) + end_user_id: str | None = None + deadline: datetime + draft: bool + + +class EngineRun(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + run_id: str + session_id: str + workflow_id: str + snapshot_id: str + team_id: str + created_by: str + status: RunState + created_at: datetime + updated_at: datetime + input: RunInput + output: JsonValue = None + metadata: Mapping[str, JsonValue] | None = None + request_digest: str + cancel_requested: bool = False + lease_token: str | None = None + lease_expires_at: datetime | None = None + + +class StepTrace(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + step_id: str + run_id: str + node_id: str + state: RunState + input: JsonValue + output: JsonValue = None + attempts: int = Field(ge=0) + deadline: datetime + started_at: datetime + finished_at: datetime | None = None + error: str | None = None + prompt_tokens: int = 0 + completion_tokens: int = 0 + response_cost: float | None = None + request_id: str | None = None + call_id: str | None = None + request_ids: tuple[str, ...] = () + + +class RunSteps(BaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + snapshot_id: str + state: RunState + cancel_requested: bool + steps: tuple[StepTrace, ...] + error: str | None = None + + +class RunView(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + run_id: str + session_id: str + workflow_id: str + snapshot_id: str + team_id: str + status: RunState + created_at: datetime + updated_at: datetime + cancel_requested: bool + + +class RunPage(BaseModel): + model_config = ConfigDict(frozen=True) + + runs: tuple[RunView, ...] + count: int + next_cursor: str | None + + +@dataclass(frozen=True, slots=True) +class ChildCallResult: + output: JsonValue + prompt_tokens: int + completion_tokens: int + request_id: str | None + response_cost: float | None + + +@dataclass(frozen=True, slots=True) +class ExecutionError: + code: str + message: str + receipt: ChildCallResult | None = None + + +def workflow_call_id(run_id: str, node_id: str, attempt: int) -> str: + return str(uuid5(NAMESPACE_URL, f"litellm/workflows/{run_id}/{node_id}/{attempt}")) + + +@dataclass(frozen=True, slots=True) +class RunAdmission: + run: EngineRun + created: bool + + +class WorkflowExecutionStore(Protocol): + async def find_idempotent( + self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str + ) -> EngineRun | None: ... + + async def admit( + self, + *, + workflow_id: str, + snapshot_id: str, + team_id: str, + key_hash: str, + input: RunInput, + digest: str, + idempotency_key: str | None, + ) -> RunAdmission | ExecutionError: ... + + async def get(self, run_id: str) -> EngineRun | None: ... + + async def recover(self, limit: int) -> tuple[EngineRun, ...]: ... + + async def claim(self, run_id: str, token: str, lease_expires_at: datetime) -> EngineRun | None: ... + + async def renew(self, run_id: str, token: str, lease_expires_at: datetime) -> bool: ... + + async def checkpoint(self, run_id: str, token: str, step: StepTrace) -> bool: ... + + async def steps(self, run_id: str) -> tuple[StepTrace, ...]: ... + + async def finish(self, run_id: str, token: str, state: RunState, output: JsonValue, error: str | None) -> bool: ... + + async def cancel(self, run_id: str) -> None: ... + + async def list_runs( + self, + workflow_id: str, + team_id: str, + key_hash: str | None, + state: RunState | None, + limit: int, + cursor: str | None, + ) -> RunPage | ExecutionError: ... diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..6cb6e61864d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -660,6 +660,21 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_UPDATE.value, ) + workflow_definition_routes = ( + "/v1/workflows", + "/v1/workflows/{workflow_id}", + "/v1/workflows/{workflow_id}/validate", + "/v1/workflows/{workflow_id}/publish", + ) + + workflow_execution_routes = ( + "/v1/workflows/{workflow_id}/test", + "/v1/workflows/{workflow_id}/runs", + "/v1/workflows/runs/{run_id}/steps", + "/v1/workflows/runs/{run_id}/definition", + "/v1/workflows/runs/{run_id}/cancel", + ) + management_routes = ( [ # user @@ -723,6 +738,8 @@ class LiteLLMRoutes(enum.Enum): + mcp_management_routes + list(agent_management_routes) + list(skill_hub_routes) + + list(workflow_definition_routes) + + list(workflow_execution_routes) ) spend_tracking_routes = [ @@ -966,6 +983,8 @@ class LiteLLMRoutes(enum.Enum): # inside agent_endpoints/endpoints.py *agent_management_routes, *skill_hub_routes, + *workflow_definition_routes, + *workflow_execution_routes, ] # routes that manage their own allowed/disallowed logic ## Org Admin Routes ## diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 125bfab4edf..39647cc8fde 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -603,8 +603,10 @@ class RouteChecks: if not isinstance(route, str): return False - def _placeholder_to_regex(match: re.Match) -> str: + def _placeholder_to_regex(match: re.Match[str]) -> str: placeholder: Final = match.group(0).strip("{}") + if placeholder == "workflow_id" and match.string.startswith("/v1/workflows/"): + return r"(?!runs(?:/|$))[^/]+" if not placeholder.endswith(":path"): return r"[^/]+" # A ":path" placeholder takes whatever the router's own path diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ba253b44ea3..2f9a1300244 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -4871,7 +4871,12 @@ async def _run_post_custom_auth_checks( async def authorize_internal_virtual_key( - key_hash: str, request: Request, request_data: dict[str, object] + key_hash: str, + request: Request, + request_data: dict[str, object], + *, + route: str | None = None, + expected_team_id: str | None = None, ) -> UserAPIKeyAuth: """Authorize a server-owned job against its persisted virtual-key assignment, never a client-supplied bearer hash.""" from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache @@ -4880,15 +4885,17 @@ async def authorize_internal_virtual_key( await IdentityStore(prisma_client, user_api_key_cache, proxy_logging_obj=proxy_logging_obj).resolve( hashed_token=key_hash ) - ) - route: Final = get_request_route(request=request) - await pre_db_read_auth_checks(request_data=request_data, request=request, route=route) + ).model_copy(deep=True) + if expected_team_id is not None and identity.team_id != expected_team_id: + raise HTTPException(status_code=403, detail="The execution key no longer belongs to its original team") + execution_route: Final = route if route is not None else get_request_route(request=request) + await pre_db_read_auth_checks(request_data=request_data, request=request, route=execution_route) auth: Final = await validate_resolved_virtual_key( request=request, request_data=request_data, valid_token=identity, api_key=key_hash, - route=route, + route=execution_route, start_time=datetime.now(timezone.utc), parent_otel_span=None, end_user_id=None, @@ -4900,13 +4907,13 @@ async def authorize_internal_virtual_key( user_api_key_auth_obj=auth, request=request, request_data=request_data, - route=route, + route=execution_route, api_key=key_hash, force_virtual_key_checks=True, ) if recovered is not None: return recovered _seed_request_destinations(auth, request) - auth.request_route = route + auth.request_route = execution_route request.state.principal = _resolve_request_principal(request, auth) return auth diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 323cf57dfec..908d0fa8969 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -1,4 +1,4 @@ -from typing import Annotated, TypeAlias +from typing import Annotated, TypeAlias, cast # noqa: TID251 # Prisma writer delegates are generated dynamically from fastapi import Depends @@ -7,12 +7,16 @@ from litellm.models.gateway_activity import GatewayActivityStore from litellm.models.plugin_upload import PluginUploadStore from litellm.models.skill_submission import SkillSubmissionStore from litellm.models.user_banner import UserBannerStore +from litellm.models.workflow_definition import WorkflowDefinitionStore +from litellm.models.workflow_execution import WorkflowExecutionStore from litellm.proxy.deps.database import DatabaseDep from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository from litellm.repositories.gateway_activity_repository import GatewayActivityRepository from litellm.repositories.plugin_upload_repository import PluginUploadRepository from litellm.repositories.skill_submission_repository import SkillSubmissionRepository from litellm.repositories.user_banner_repository import UserBannerRepository +from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository +from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: @@ -54,3 +58,33 @@ async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmiss SkillSubmissionRepositoryDep: TypeAlias = Annotated[ SkillSubmissionStore | None, Depends(get_skill_submission_repository) ] + + +async def get_workflow_definition_repository(database: DatabaseDep) -> WorkflowDefinitionStore | None: + return ( + None + if database is None + else WorkflowDefinitionRepository( + cast(DefinitionDatabase, database.prisma.writer_db) # cast-ok: Prisma delegates implement the protocol + ) + ) + + +WorkflowDefinitionRepositoryDep: TypeAlias = Annotated[ + WorkflowDefinitionStore | None, Depends(get_workflow_definition_repository) +] + + +async def get_workflow_execution_repository(database: DatabaseDep) -> WorkflowExecutionStore | None: + return ( + None + if database is None + else WorkflowExecutionRepository( + cast(ExecutionDatabase, database.prisma.writer_db) # cast-ok: generated Prisma transaction delegates + ) + ) + + +WorkflowExecutionRepositoryDep: TypeAlias = Annotated[ + WorkflowExecutionStore | None, Depends(get_workflow_execution_repository) +] diff --git a/litellm/proxy/management_endpoints/workflow_definition_endpoints.py b/litellm/proxy/management_endpoints/workflow_definition_endpoints.py new file mode 100644 index 00000000000..f134db68f64 --- /dev/null +++ b/litellm/proxy/management_endpoints/workflow_definition_endpoints.py @@ -0,0 +1,246 @@ +from typing import Annotated, Final, Literal, NoReturn, TypeAlias, TypeVar + +from fastapi import APIRouter, Depends, Header, HTTPException, Query +from fastapi.responses import JSONResponse +from pydantic import BaseModel, ConfigDict +from typing_extensions import assert_never + +from litellm.models.workflow_definition import ( + DefinitionError, + DefinitionErrorCode, + DefinitionPage, + Digest, + WorkflowBundle, + WorkflowDefinitionStore, + WorkflowDraft, +) +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep +from litellm.proxy.workflows.definitions import ( + authorize_team, + authorized_draft, + check_etag, + requested_team, + validate_bundle, +) + +router: Final = APIRouter(tags=["workflow management"]) +_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] +_Team: TypeAlias = Annotated[str | None, Query(min_length=1, max_length=256)] +_ETag: TypeAlias = Annotated[str | None, Header(alias="If-Match")] +_Result: Final = TypeVar("_Result") + + +class WorkflowRegistration(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + bundle: WorkflowBundle + + +class WorkflowPublication(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + digest: Digest + + +class WorkflowValidationRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class WorkflowValidationResult(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + valid: Literal[True] = True + workflow_id: str + etag: str + digest: str + + +class WorkflowPublicationResult(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + workflow_id: str + snapshot_id: str + model: str + digest: str + + +def _error_status(code: DefinitionErrorCode) -> int: + match code: + case "not_found": + return 404 + case "forbidden": + return 403 + case "team_required" | "invalid_bundle": + return 400 + case "conflict": + return 409 + case "stale_draft": + return 412 + case "precondition_required": + return 412 + case _: + assert_never(code) + + +def raise_public(error: DefinitionError) -> NoReturn: + raise HTTPException( + status_code=_error_status(error.code), + detail={ + "error": { + "type": "workflow_definition_error", + "code": error.code, + "message": error.message, + "param": error.param, + } + }, + ) + + +def _result(result: _Result | DefinitionError) -> _Result: + if isinstance(result, DefinitionError): + raise_public(result) + return result + + +def _store(store: WorkflowDefinitionStore | None) -> WorkflowDefinitionStore: + if store is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + return store + + +def _author(caller: UserAPIKeyAuth) -> str: + if not caller.user_id or caller.user_role not in { + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.ORG_ADMIN, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.TEAM, + }: + raise_public(DefinitionError("forbidden", "Workflow authoring requires a team member identity")) + return caller.user_id + + +def _draft_response(draft: WorkflowDraft, status_code: int = 200) -> JSONResponse: + return JSONResponse( + content=draft.model_dump(mode="json", by_alias=True, exclude_unset=True), + status_code=status_code, + headers={"ETag": draft.etag}, + ) + + +@router.post("/v1/workflows", status_code=201, response_model=WorkflowDraft, response_model_exclude_unset=True) +async def register_workflow( + data: WorkflowRegistration, + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, +) -> JSONResponse: + """Create a team draft from a sequential LiteFlow v1 bundle.""" + contributor: Final = _author(caller) + store: Final = _store(repository) + scope: Final = _result(requested_team(caller, team_id)) + _result(await authorize_team(store, caller, scope, "write")) + if (failure := validate_bundle(data.bundle)) is not None: + raise_public(failure) + draft: Final = _result(await store.create(data.bundle, scope, contributor)) + return _draft_response(draft, 201) + + +@router.get("/v1/workflows", response_model=DefinitionPage, response_model_exclude_unset=True) +async def list_workflow_definitions( + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, + limit: Annotated[int, Query(ge=1, le=100)] = 50, + cursor: Annotated[ + str | None, Query(pattern=r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$") + ] = None, +) -> DefinitionPage: + """List readable team drafts in immutable workflow ID order.""" + store: Final = _store(repository) + scope: Final = _result(requested_team(caller, team_id)) + _result(await authorize_team(store, caller, scope, "read")) + return await store.list_drafts(scope, limit, cursor) + + +@router.get("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True) +async def get_workflow_definition( + workflow_id: str, + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, +) -> JSONResponse: + """Read the current draft and its publication pointer.""" + draft: Final = _result(await authorized_draft(_store(repository), caller, workflow_id, team_id)) + return _draft_response(draft) + + +@router.put("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True) +async def replace_workflow_definition( + workflow_id: str, + data: WorkflowRegistration, + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, + if_match: _ETag = None, +) -> JSONResponse: + """Replace a draft only when If-Match identifies its current revision.""" + contributor: Final = _author(caller) + store: Final = _store(repository) + current: Final = _result(await authorized_draft(store, caller, workflow_id, team_id)) + _result(await authorize_team(store, caller, current.team_id, "write")) + if (precondition := check_etag(current, if_match)) is not None: + raise_public(precondition) + if (failure := validate_bundle(data.bundle)) is not None: + raise_public(failure) + draft: Final = _result(await store.replace(workflow_id, current.team_id, current.etag, data.bundle, contributor)) + return _draft_response(draft) + + +@router.post("/v1/workflows/{workflow_id}/validate", response_model=WorkflowValidationResult) +async def validate_workflow_definition( + workflow_id: str, + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, + if_match: _ETag = None, + data: WorkflowValidationRequest | None = None, +) -> WorkflowValidationResult: + """Validate a draft without model calls or Python execution.""" + _author(caller) + store: Final = _store(repository) + draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id)) + _result(await authorize_team(store, caller, draft.team_id, "write")) + if (precondition := check_etag(draft, if_match)) is not None: + raise_public(precondition) + if (failure := validate_bundle(draft.bundle)) is not None: + raise_public(failure) + return WorkflowValidationResult(workflow_id=workflow_id, etag=draft.etag, digest=draft.digest) + + +@router.post("/v1/workflows/{workflow_id}/publish", response_model=WorkflowPublicationResult) +async def publish_workflow_definition( + workflow_id: str, + data: WorkflowPublication, + repository: WorkflowDefinitionRepositoryDep, + caller: _Caller, + team_id: _Team = None, + if_match: _ETag = None, +) -> WorkflowPublicationResult: + """Publish the reviewed revision as an immutable snapshot.""" + publisher: Final = _author(caller) + store: Final = _store(repository) + draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id)) + _result(await authorize_team(store, caller, draft.team_id, "publish")) + if (precondition := check_etag(draft, if_match)) is not None: + raise_public(precondition) + if (failure := validate_bundle(draft.bundle)) is not None: + raise_public(failure) + snapshot: Final = _result(await store.publish(workflow_id, draft.team_id, draft.etag, data.digest, publisher)) + return WorkflowPublicationResult( + workflow_id=workflow_id, + snapshot_id=snapshot.snapshot_id, + model=f"workflow/{snapshot.name}", + digest=snapshot.digest, + ) diff --git a/litellm/proxy/management_endpoints/workflow_execution_endpoints.py b/litellm/proxy/management_endpoints/workflow_execution_endpoints.py new file mode 100644 index 00000000000..a4ec49a58bf --- /dev/null +++ b/litellm/proxy/management_endpoints/workflow_execution_endpoints.py @@ -0,0 +1,424 @@ +import asyncio +import hashlib +import json +from collections.abc import Mapping +from datetime import timedelta +from typing import Annotated, Final, Literal, NoReturn, TypeAlias + +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, Response +from fastapi.responses import JSONResponse +from openai.types.chat import ChatCompletionMessageParam +from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError + +from litellm.models.workflow_definition import DefinitionError, WorkflowDefinitionStore, WorkflowSnapshot +from litellm.models.workflow_execution import ( + EngineRun, + ExecutionError, + RunInput, + RunPage, + RunState, + RunSteps, + WorkflowExecutionStore, +) +from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep +from litellm.proxy.management.teams.authz import is_team_admin +from litellm.proxy.management_endpoints.workflow_definition_endpoints import raise_public +from litellm.proxy.workflows.definitions import ( + authorize_team, + authorized_draft, + freeze_draft, + published_workflow, +) +from litellm.proxy.workflows.monty.models import MontyFailure +from litellm.proxy.workflows.monty.runtime import validate_json +from litellm.proxy.workflows.runner import utc_now + +router: Final = APIRouter(tags=["workflow management"]) +_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] +_Idempotency: TypeAlias = Annotated[str | None, Header(alias="Idempotency-Key", min_length=1, max_length=256)] +_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam]) + + +class TestRunRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256) + authorize_billing: StrictBool + + +class AcceptedRun(BaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + snapshot_id: str + trace_url: str + + +class EmptyRunAction(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class WorkflowInvocation(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + model: str + messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256) + stream: Literal[False] = False + user: str | None = Field(default=None, max_length=256) + metadata: Mapping[str, JsonValue] | None = None + + +def _error(error: ExecutionError, run_id: str | None = None) -> NoReturn: + code: Final = { + "not_found": 404, + "forbidden": 403, + "idempotency_conflict": 409, + "capacity_exceeded": 429, + "wait_timeout": 504, + "execution_failed": 422, + "database_unavailable": 503, + }.get(error.code, 400) + raise HTTPException( + status_code=code, + detail={"error": {"type": "workflow_error", "code": error.code, "message": error.message}}, + headers={"x-litellm-workflow-run-id": run_id} if run_id is not None else None, + ) + + +def _stores( + definitions: WorkflowDefinitionStore | None, runs: WorkflowExecutionStore | None +) -> tuple[WorkflowDefinitionStore, WorkflowExecutionStore]: + if definitions is None or runs is None: + _error(ExecutionError("database_unavailable", "Workflow execution requires a database")) + return definitions, runs + + +def _execution_key(caller: UserAPIKeyAuth, team_id: str) -> str: + if not caller.token or not caller.via_virtual_key or caller.is_session_token or caller.team_id != team_id: + _error(ExecutionError("forbidden", "Execution requires a virtual key assigned to the workflow team")) + return caller.token + + +def _idempotency_key(operation: str, key: str | None) -> str | None: + return None if key is None else hashlib.sha256(f"{operation}:{key}".encode()).hexdigest() + + +def _digest(messages: tuple[Mapping[str, JsonValue], ...], user: str | None, draft_etag: str | None) -> str: + return hashlib.sha256( + json.dumps( + {"messages": tuple(dict(message) for message in messages), "user": user, "etag": draft_etag}, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode() + ).hexdigest() + + +def _validate_messages(messages: tuple[Mapping[str, JsonValue], ...]) -> None: + bounded: Final = validate_json([dict(message) for message in messages], max_bytes=1024 * 1024, max_depth=64) + if isinstance(bounded, MontyFailure): + _error(ExecutionError("invalid_input", "Workflow messages exceed the JSON size or depth limit")) + try: + _MESSAGES.validate_python(bounded, strict=True) + except ValidationError: + _error(ExecutionError("invalid_input", "Expected standard Chat Completions messages")) + + +async def _recover_admission( + runs: WorkflowExecutionStore, + workflow_id: str, + team_id: str, + key_hash: str, + key: str | None, + digest: str, +) -> EngineRun | None: + if key is None: + return None + existing: Final = await runs.find_idempotent(workflow_id, team_id, key_hash, key) + if existing is not None and existing.request_digest != digest: + _error(ExecutionError("idempotency_conflict", "Idempotency key was used with different input")) + return existing + + +async def _admit( + runs: WorkflowExecutionStore, + snapshot: WorkflowSnapshot, + caller: UserAPIKeyAuth, + messages: tuple[Mapping[str, JsonValue], ...], + user: str | None, + key: str | None, + digest: str, + *, + draft: bool, +) -> EngineRun: + _validate_messages(messages) + admitted: Final = await runs.admit( + workflow_id=snapshot.workflow_id, + snapshot_id=snapshot.snapshot_id, + team_id=snapshot.team_id, + key_hash=_execution_key(caller, snapshot.team_id), + input=RunInput( + messages=messages, + end_user_id=user, + deadline=utc_now() + timedelta(seconds=snapshot.bundle.timeout_seconds), + draft=draft, + ), + digest=digest, + idempotency_key=key, + ) + if isinstance(admitted, ExecutionError): + _error(admitted) + return admitted.run + + +def _accepted(run: EngineRun) -> JSONResponse: + return JSONResponse( + status_code=202, + content=AcceptedRun( + run_id=run.run_id, + snapshot_id=run.snapshot_id, + trace_url=f"/v1/workflows/runs/{run.run_id}/steps", + ).model_dump(), + headers={"x-litellm-workflow-run-id": run.run_id}, + ) + + +@router.post("/v1/workflows/{workflow_id}/test", status_code=202, response_model=AcceptedRun) +async def test_workflow( + workflow_id: str, + data: TestRunRequest, + caller: _Caller, + definitions: WorkflowDefinitionRepositoryDep, + executions: WorkflowExecutionRepositoryDep, + if_match: Annotated[str | None, Header(alias="If-Match")] = None, + idempotency_key: _Idempotency = None, +) -> JSONResponse: + """Start a paid test of the exact draft revision and return its run ID immediately.""" + store, runs = _stores(definitions, executions) + if not data.authorize_billing: + _error(ExecutionError("billing_authorization_required", "Set authorize_billing to true for a paid test run")) + draft: Final = await authorized_draft(store, caller, workflow_id) + if isinstance(draft, DefinitionError): + raise_public(draft) + permission: Final = await authorize_team(store, caller, draft.team_id, "write") + if isinstance(permission, DefinitionError): + raise_public(permission) + key_hash: Final = _execution_key(caller, draft.team_id) + digest: Final = _digest(data.messages, caller.end_user_id, if_match) + key: Final = _idempotency_key("test", idempotency_key) + previous: Final = await _recover_admission(runs, workflow_id, draft.team_id, key_hash, key, digest) + if previous is not None: + return _accepted(previous) + snapshot: Final = await freeze_draft(store, caller, workflow_id, if_match) + if isinstance(snapshot, DefinitionError): + raise_public(snapshot) + return _accepted(await _admit(runs, snapshot, caller, data.messages, caller.end_user_id, key, digest, draft=True)) + + +async def _read_run( + store: WorkflowDefinitionStore, runs: WorkflowExecutionStore, caller: UserAPIKeyAuth, run_id: str +) -> EngineRun: + run: Final = await runs.get(run_id) + if run is None: + _error(ExecutionError("not_found", "Run not found")) + team: Final = await store.get_team(run.team_id) + if team is None or team.blocked: + _error(ExecutionError("not_found", "Run not found")) + owns_run: Final = bool(caller.token) and caller.token == run.created_by and caller.team_id == run.team_id + if not owns_run and not user_api_key_has_admin_view(caller) and not is_team_admin(caller, team): + _error(ExecutionError("not_found", "Run not found")) + return run + + +@router.get("/v1/workflows/runs/{run_id}/steps") +async def get_workflow_steps( + run_id: str, + caller: _Caller, + definitions: WorkflowDefinitionRepositoryDep, + executions: WorkflowExecutionRepositoryDep, +) -> RunSteps: + """Read run progress, typed node inputs and outputs, and linked model-call receipts.""" + store, runs = _stores(definitions, executions) + run: Final = await _read_run(store, runs, caller, run_id) + error: Final = (run.metadata or {}).get("error") + return RunSteps( + run_id=run_id, + snapshot_id=run.snapshot_id, + state=run.status, + cancel_requested=run.cancel_requested, + steps=await runs.steps(run_id), + error=error if isinstance(error, str) else None, + ) + + +@router.get("/v1/workflows/runs/{run_id}/definition") +async def get_run_definition( + run_id: str, + caller: _Caller, + definitions: WorkflowDefinitionRepositoryDep, + executions: WorkflowExecutionRepositoryDep, +) -> WorkflowSnapshot: + """Read the immutable workflow definition captured when this run was admitted.""" + store, runs = _stores(definitions, executions) + run: Final = await _read_run(store, runs, caller, run_id) + snapshot: Final = await store.get_snapshot(run.snapshot_id, run.team_id) + if snapshot is None or snapshot.workflow_id != run.workflow_id: + _error(ExecutionError("not_found", "Pinned workflow definition not found")) + return snapshot + + +@router.post("/v1/workflows/runs/{run_id}/cancel", status_code=202, response_model=AcceptedRun) +async def cancel_workflow_run( + run_id: str, + caller: _Caller, + definitions: WorkflowDefinitionRepositoryDep, + executions: WorkflowExecutionRepositoryDep, + data: EmptyRunAction | None = None, +) -> JSONResponse: + """Request cancellation and return the run ID to poll until execution stops.""" + store, runs = _stores(definitions, executions) + run: Final = await _read_run(store, runs, caller, run_id) + permission: Final = await authorize_team(store, caller, run.team_id, "write") + if isinstance(permission, DefinitionError): + raise_public(permission) + await runs.cancel(run_id) + return _accepted(run) + + +async def list_engine_runs( + *, + workflow_id: str, + team_id: str | None, + status: str | None, + limit: int, + after: str | None, + caller: UserAPIKeyAuth, + definitions: WorkflowDefinitionStore | None, + executions: WorkflowExecutionStore | None, +) -> RunPage: + store, runs = _stores(definitions, executions) + definition: Final = await store.get_draft(workflow_id, team_id) + if definition is None: + _error(ExecutionError("not_found", "Workflow not found")) + permission: Final = await authorize_team(store, caller, definition.draft.team_id, "invoke") + if isinstance(permission, DefinitionError): + raise_public(permission) + may_read_team: Final = user_api_key_has_admin_view(caller) or is_team_admin(caller, permission) + if not may_read_team and not caller.token: + _error(ExecutionError("forbidden", "A run reader identity is required")) + try: + state: Final = RunState(status) if status is not None else None + except ValueError: + _error(ExecutionError("invalid_status", "Unknown workflow run status")) + result: Final = await runs.list_runs( + workflow_id, + definition.draft.team_id, + None if may_read_team else caller.token, + state, + min(limit, 100), + after, + ) + if isinstance(result, ExecutionError): + _error(result) + return result + + +@router.get("/v1/workflows/{workflow_id}/runs") +async def list_definition_runs( + workflow_id: str, + caller: _Caller, + definitions: WorkflowDefinitionRepositoryDep, + executions: WorkflowExecutionRepositoryDep, + team_id: str | None = None, + status: RunState | None = None, + limit: Annotated[int, Query(ge=1, le=100)] = 50, + after: Annotated[str | None, Query(max_length=512)] = None, +) -> RunPage: + return await list_engine_runs( + workflow_id=workflow_id, + team_id=team_id, + status=status.value if status is not None else None, + limit=limit, + after=after, + caller=caller, + definitions=definitions, + executions=executions, + ) + + +async def _wait_for_run(runs: WorkflowExecutionStore, run_id: str) -> EngineRun: + while True: + current: EngineRun | None = await runs.get(run_id) + if current is None: + _error(ExecutionError("not_found", "Run not found"), run_id) + if current.status in {RunState.COMPLETED, RunState.FAILED, RunState.CANCELLED}: + return current + await asyncio.sleep(0.2) + + +async def invoke_workflow( + request: Request, + response: Response, + caller: UserAPIKeyAuth, + raw: object, + definitions: WorkflowDefinitionStore | None, + executions: WorkflowExecutionStore | None, +) -> Mapping[str, JsonValue] | JSONResponse: + store, runs = _stores(definitions, executions) + try: + data: Final = WorkflowInvocation.model_validate(raw) + except ValidationError: + _error(ExecutionError("unsupported_request", "Use model, messages, user, metadata and stream=false")) + snapshot: Final = await published_workflow(store, caller, data.model.removeprefix("workflow/")) + if isinstance(snapshot, DefinitionError): + raise_public(snapshot) + key: Final = request.headers.get("Idempotency-Key") + if key is not None and not 1 <= len(key) <= 256: + _error(ExecutionError("invalid_idempotency_key", "Idempotency-Key must contain 1 to 256 characters")) + admitted: Final = await _admit( + runs, + snapshot, + caller, + data.messages, + data.user or caller.end_user_id, + _idempotency_key("invoke", key), + _digest(data.messages, data.user or caller.end_user_id, None), + draft=False, + ) + response.headers["x-litellm-workflow-run-id"] = admitted.run_id # rebind-ok: FastAPI response header + if request.headers.get("Prefer") == "respond-async": + return _accepted(admitted) + try: + completed: Final = await asyncio.wait_for(_wait_for_run(runs, admitted.run_id), timeout=60) + except (TimeoutError, asyncio.TimeoutError): + _error( + ExecutionError("wait_timeout", "The run is still executing. Inspect its trace or retry with the same key"), + admitted.run_id, + ) + if completed.status is not RunState.COMPLETED: + _error( + ExecutionError("execution_failed", "Workflow execution did not complete. Inspect the run trace"), + admitted.run_id, + ) + traces: Final = await runs.steps(admitted.run_id) + prompt_tokens: Final = sum(step.prompt_tokens for step in traces) + completion_tokens: Final = sum(step.completion_tokens for step in traces) + content: Final = ( + completed.output + if isinstance(completed.output, str) + else json.dumps(completed.output, ensure_ascii=False, allow_nan=False) + ) + return { + "id": f"chatcmpl-workflow-{completed.run_id}", + "object": "chat.completion", + "created": int(completed.created_at.timestamp()), + "model": data.model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 649941bc233..3d61beb166f 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( user_api_key_has_admin_view, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep from litellm.repositories.table_repositories import ( WorkflowEventRepository, WorkflowMessageRepository, @@ -201,6 +202,8 @@ async def _require_run( prisma_client: object, run_id: str, user_api_key_dict: UserAPIKeyAuth | None = None, + *, + mutation: bool = False, ) -> _RunRow: """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id}) @@ -210,6 +213,8 @@ async def _require_run( caller: Final = _caller_key(user_api_key_dict) if not caller or run.created_by != caller: raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + if mutation and getattr(run, "engine_owned", False) is True: + raise HTTPException(status_code=403, detail="Gateway executions cannot be changed through tracking endpoints") return run @@ -263,11 +268,30 @@ async def list_workflow_runs( status: str | None = Query(None), limit: int = Query(50, ge=1, le=250), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + workflow_id: str | None = None, + team_id: str | None = None, + after: str | None = None, + definitions: WorkflowDefinitionRepositoryDep = None, + executions: WorkflowExecutionRepositoryDep = None, ): """List workflow runs. Filter by workflow_type and/or status. Non-admin callers only see runs created by their own API key. """ + if workflow_id is not None: + from litellm.proxy.management_endpoints.workflow_execution_endpoints import list_engine_runs + + return await list_engine_runs( + workflow_id=workflow_id, + team_id=team_id, + status=status, + limit=limit, + after=after, + caller=user_api_key_dict, + definitions=definitions, + executions=executions, + ) + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -360,7 +384,7 @@ async def update_workflow_run( raise HTTPException(status_code=400, detail="No fields to update") # Enforce ownership before writing. - await _require_run(prisma_client, run_id, user_api_key_dict) + await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True) try: run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update( @@ -398,7 +422,7 @@ async def append_workflow_event( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - await _require_run(prisma_client, run_id, user_api_key_dict) + await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True) new_status: Final = _EVENT_STATUS_MAP.get(data.event_type) @@ -487,7 +511,7 @@ async def append_workflow_message( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - await _require_run(prisma_client, run_id, user_api_key_dict) + await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True) for attempt in range(_MAX_SEQUENCE_RETRIES): try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ffd2dc305e7..13fa99aa117 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -841,6 +841,12 @@ from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) +from litellm.proxy.management_endpoints.workflow_definition_endpoints import ( + router as workflow_definition_router, +) +from litellm.proxy.management_endpoints.workflow_execution_endpoints import ( + router as workflow_execution_router, +) from litellm.proxy.management_endpoints.workflow_management_endpoints import ( router as workflow_management_router, ) @@ -1983,6 +1989,9 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener async with AsyncExitStack() as admin_mcp_stack: try: await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) + from litellm.proxy.workflows.lifespan import workflow_lifespan + + await admin_mcp_stack.enter_async_context(workflow_lifespan(database)) gc.freeze() verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) yield state @@ -13089,6 +13098,28 @@ async def chat_completion( global general_settings, user_debug, proxy_logging_obj, llm_model_list global user_temperature, user_request_timeout, user_max_tokens, user_api_base data: Final = await read_request_body(request=request) + workflow_data: Final = as_str_mapping(_UNTYPED_VALUE.validate_python(data)) + workflow_model: Final = workflow_data.get("model") if workflow_data is not None else None + if isinstance(workflow_model, str) and workflow_model.startswith("workflow/"): + from litellm.proxy.deps.database import get_database + from litellm.proxy.deps.repositories import ( + get_workflow_definition_repository, + get_workflow_execution_repository, + ) + from litellm.proxy.management_endpoints.workflow_execution_endpoints import invoke_workflow + + workflow_database: Final = await get_database(request) + try: + return await invoke_workflow( + request, + fastapi_response, + user_api_key_dict, + workflow_data, + await get_workflow_definition_repository(workflow_database), + await get_workflow_execution_repository(workflow_database), + ) + finally: + await release_unbound_budget_reservation(user_api_key_dict.budget_reservation) if user_api_key_dict is not None: if not isinstance(data.get("metadata"), dict): # Covers both missing and JSON-string metadata (multipart / @@ -22365,6 +22396,8 @@ app.include_router(model_access_group_management_router) app.include_router(auto_router_management_router) app.include_router(tag_management_router) app.include_router(workflow_management_router) +app.include_router(workflow_execution_router) +app.include_router(workflow_definition_router) app.include_router(memory_router) app.include_router(lens_router) app.include_router(plugin_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 818f2dc4858..c0dc2eb960b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1975,10 +1975,53 @@ model LiteLLM_ShadowEvalFunnel { // One instance of work being done. `status` is a materialized cache of the // latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowDefinition { + workflow_id String @id @default(uuid()) + team_id String + name String + etag String + digest String + bundle_json String + contributors String[] + published_snapshot_id String? + published_etag String? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + snapshots LiteLLM_WorkflowSnapshot[] + + @@unique([team_id, name]) + @@index([team_id, workflow_id]) +} + +model LiteLLM_WorkflowSnapshot { + snapshot_id String @id @default(uuid()) + workflow_id String + team_id String + name String + digest String + bundle_json String + contributors String[] + published_by String? + created_at DateTime @default(now()) + definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict) + + @@index([workflow_id, created_at, snapshot_id]) + @@index([team_id, snapshot_id]) +} + model LiteLLM_WorkflowRun { run_id String @id @default(uuid()) session_id String @unique @default(uuid()) workflow_type String + workflow_id String? + snapshot_id String? + team_id String? + engine_owned Boolean @default(false) + request_digest String? + idempotency_key String? + lease_token String? + lease_expires_at DateTime? + cancel_requested Boolean @default(false) status String @default("pending") created_by String? // user_id of the key that created this run; null = created by master key created_at DateTime @default(now()) @@ -1994,6 +2037,10 @@ model LiteLLM_WorkflowRun { @@index([session_id]) @@index([created_at]) @@index([created_by]) + @@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key") + @@index([workflow_id, team_id, created_at, run_id]) + @@index([engine_owned, status, lease_expires_at]) + @@index([team_id, engine_owned, status]) } // Append-only log of state transitions. Never mutate rows here. diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8f0584bd0d0..90acac978ed 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -468,7 +468,7 @@ async def release_or_invalidate_budget_reservation( budget_reservation["finalized"] = True -async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None: +async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object] | None) -> None: """Release a reservation no logging callback took ownership of, once the request ended. A handler whose litellm call never builds a logging object (batch cancel, file diff --git a/litellm/proxy/workflows/LITEFLOW.md b/litellm/proxy/workflows/LITEFLOW.md new file mode 100644 index 00000000000..a6b459161b8 --- /dev/null +++ b/litellm/proxy/workflows/LITEFLOW.md @@ -0,0 +1,87 @@ +# LiteFlow gateway workflows + +The gateway accepts LiteFlow compiled contract version `1` and executes sequential `chat_completion` and `python` nodes. Python nodes use the Monty adapter. Conditional, parallel and streaming execution are not supported + +Workflow definitions belong to a team. Authors register and edit drafts. A team administrator who did not contribute to the draft publishes an immutable snapshot. Invocation requires a virtual key assigned to that team, with access to the workflow model and each child model + +## SDK example + +This contract matches LiteFlow PRs [3](https://github.com/BerriAI/liteflow/pull/3) and [5](https://github.com/BerriAI/liteflow/pull/5). Configure `gateway-model` on the proxy and provide two different team members' keys + +```python +import os +import time + +from liteflow import ChatCompletion, GatewayClient, Input, Workflow + +workflow = Workflow( + name="welcome", + nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)], +) +messages = [{"role": "user", "content": "Write one short welcome sentence."}] + +with GatewayClient( + base_url=os.environ["PROXY_BASE_URL"], + api_key=os.environ["WORKFLOW_AUTHOR_KEY"], + team_id=os.environ["WORKFLOW_TEAM_ID"], +) as author, GatewayClient( + base_url=os.environ["PROXY_BASE_URL"], + api_key=os.environ["WORKFLOW_PUBLISHER_KEY"], + team_id=os.environ["WORKFLOW_TEAM_ID"], +) as publisher: + draft = author.register(workflow) + author.validate(draft) + publisher.publish(draft) + accepted = author.test(draft, messages=messages, authorize_billing=True) + for _ in range(300): + trace = author.inspect(accepted["run_id"]) + print(trace) + if trace.state.value in ("completed", "failed", "cancelled"): + break + time.sleep(1) + result = author.run("welcome", messages=messages) + print(result.run_id, result.output) +``` + +## API contract + +| Operation | Endpoint | Contract | +| --- | --- | --- | +| Register | `POST /v1/workflows` | `{"bundle": }`, optional `team_id` query, returns `201` | +| List definitions | `GET /v1/workflows` | `team_id`, `limit`, `cursor`; returns `workflows`, `next_cursor` | +| Retrieve draft | `GET /v1/workflows/{workflow_id}` | Returns the draft and `ETag` header | +| Replace draft | `PUT /v1/workflows/{workflow_id}` | Compiled bundle and exact `If-Match` | +| Validate | `POST /v1/workflows/{workflow_id}/validate` | Exact `If-Match`, no body required | +| Publish | `POST /v1/workflows/{workflow_id}/publish` | Exact `If-Match` and `{"digest": ""}` | +| Test draft | `POST /v1/workflows/{workflow_id}/test` | Exact `If-Match`, `messages`, `authorize_billing: true`; returns `202` | +| Invoke published | `POST /v1/chat/completions` | `model: "workflow/"`, `messages`, `stream: false` | +| List runs | `GET /v1/workflows/{workflow_id}/runs` | `limit`, `status`, `after`; returns `runs`, `count`, `next_cursor` | +| Inspect | `GET /v1/workflows/runs/{run_id}/steps` | Run state and typed step traces | +| Pinned definition | `GET /v1/workflows/runs/{run_id}/definition` | Exact snapshot used by that run | +| Cancel | `POST /v1/workflows/runs/{run_id}/cancel` | Returns `202`; poll until terminal state | + +Draft responses contain `name`, `workflow_id`, `team_id`, `etag`, `digest`, `bundle` and `published_snapshot_id`. Missing or stale `If-Match` returns `412`. Validation checks the compiled graph and Python program without executing model or Python nodes + +Test runs and chat invocations accept `Idempotency-Key`. Retrying the same operation with the same key and input returns the original run. Reusing that key with different input returns `409`. A test run pins the exact draft revision even if the draft changes later + +For live UI progress, send `Prefer: respond-async` on a published chat invocation. The gateway returns `202` with `run_id`, `snapshot_id`, `trace_url` and the `x-litellm-workflow-run-id` header. The ordinary SDK call waits for a Chat Completions response carrying the same run-ID header. A synchronous wait that reaches 60 seconds returns `504` with the run ID so the caller can continue polling + +## Trace and spend data + +Draw the pinned bundle's `nodes` and resolve each binding's `from` reference. Match `step.node_id` to `node.id`. Run and node states are `pending`, `running`, `completed`, `failed` and `cancelled` + +Each step exposes typed `input` and `output`, `attempts`, `started_at`, `finished_at`, `deadline`, sanitized `error`, token counts and `response_cost`. A run-level `error` explains failures before any node starts + +Model steps carry a deterministic `call_id` before dispatch. `request_id` and `request_ids` link to actual `LiteLLM_SpendLogs.request_id` values. Spend writes are asynchronous, so receipt IDs can be absent on the first read and appear on later polls. The trace reader sums persisted attempts once and does not duplicate earlier checkpoints + +Each child uses the standard gateway authorization, hooks, guardrails, limits and billing path. The workflow wrapper emits no additional bill. Python nodes have no model spend receipt. Known costs remain visible when output validation fails + +Definitions and snapshots use separate tables. Execution reuses `LiteLLM_WorkflowRun` and `LiteLLM_WorkflowEvent`. Existing tracking mutations cannot overwrite engine-owned runs + +## Execution limits and recovery + +The initial runner allows four active runs per worker and 64 pending or running runs per team. Database leases prevent two workers from executing the same run. Completed node checkpoints survive a restart. Recovery fails an interrupted model attempt whose outcome is uncertain instead of replaying a potentially billable request + +Cancellation is cooperative. It stops downstream nodes but cannot undo a model request already sent. Trace receipts can still arrive after cancellation + +Ordinary virtual keys can read their own runs. Team administrators and proxy administrators can inspect team runs. Definition access does not grant another caller access to private run inputs or outputs diff --git a/litellm/proxy/workflows/definitions.py b/litellm/proxy/workflows/definitions.py new file mode 100644 index 00000000000..c956574963d --- /dev/null +++ b/litellm/proxy/workflows/definitions.py @@ -0,0 +1,212 @@ +from collections.abc import Mapping +from typing import Final, Literal + +from jsonschema import Draft202012Validator +from jsonschema.exceptions import SchemaError +from openai.types.chat import ChatCompletionMessageParam +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.workflow_definition import ( + Binding, + ChatCompletionNode, + DefinitionError, + PythonNode, + ValueBinding, + WorkflowBundle, + WorkflowDefinitionStore, + WorkflowDraft, + WorkflowSnapshot, +) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view +from litellm.proxy.management.teams.authz import is_team_admin, is_team_member +from litellm.proxy.workflows.monty.compiler import compile_program +from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits + +_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam]) +_LIMITS: Final = MontyLimits() +_WRITER_ROLES: Final = frozenset( + {LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.ORG_ADMIN, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM} +) + + +def requested_team(caller: UserAPIKeyAuth, team_id: str | None) -> str | DefinitionError: + selected: Final = team_id if team_id is not None else caller.team_id + if not selected: + return DefinitionError("team_required", "Specify team_id for this workflow", "team_id") + if ( + caller.team_id + and caller.team_id != selected + and not caller.is_session_token + and not user_api_key_has_admin_view(caller) + ): + return DefinitionError("not_found", "Workflow not found") + return selected + + +async def authorize_team( + store: WorkflowDefinitionStore, + caller: UserAPIKeyAuth, + team_id: str, + action: Literal["read", "write", "publish", "invoke"], +) -> LiteLLM_TeamTable | DefinitionError: + scope: Final = requested_team(caller, team_id) + if isinstance(scope, DefinitionError): + return scope + team: Final = await store.get_team(scope) + if team is None or team.blocked: + return DefinitionError("not_found", "Workflow not found") + if action == "invoke" and (caller.team_id == scope or user_api_key_has_admin_view(caller)): + return team + if not user_api_key_has_admin_view(caller) and not is_team_member(caller, team): + return DefinitionError("not_found", "Workflow not found") + if action == "read": + return team + if not caller.user_id or caller.user_role not in _WRITER_ROLES: + return DefinitionError("forbidden", "Workflow authoring requires a team member identity") + if action == "publish" and caller.user_role != LitellmUserRoles.PROXY_ADMIN and not is_team_admin(caller, team): + return DefinitionError("forbidden", "Publishing requires a team administrator") + return team + + +async def authorized_draft( + store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, workflow_id: str, team_id: str | None = None +) -> WorkflowDraft | DefinitionError: + key_team: Final = None if caller.is_session_token or user_api_key_has_admin_view(caller) else caller.team_id + definition: Final = await store.get_draft(workflow_id, team_id if team_id is not None else key_team) + if definition is None: + return DefinitionError("not_found", "Workflow not found") + access: Final = await authorize_team(store, caller, definition.draft.team_id, "read") + return access if isinstance(access, DefinitionError) else definition.draft + + +async def published_workflow( + store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, name: str, team_id: str | None = None +) -> WorkflowSnapshot | DefinitionError: + scope: Final = requested_team(caller, team_id) + if isinstance(scope, DefinitionError): + return scope + access: Final = await authorize_team(store, caller, scope, "invoke") + if isinstance(access, DefinitionError): + return access + snapshot: Final = await store.get_published(name, scope) + return snapshot if snapshot is not None else DefinitionError("not_found", "Published workflow not found") + + +async def snapshot_by_id( + store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, snapshot_id: str +) -> WorkflowSnapshot | DefinitionError: + key_team: Final = None if caller.is_session_token or user_api_key_has_admin_view(caller) else caller.team_id + snapshot: Final = await store.get_snapshot(snapshot_id, key_team) + if snapshot is None: + return DefinitionError("not_found", "Workflow snapshot not found") + access: Final = await authorize_team(store, caller, snapshot.team_id, "read") + return access if isinstance(access, DefinitionError) else snapshot + + +def check_etag(draft: WorkflowDraft, etag: str | None) -> DefinitionError | None: + if etag is None: + return DefinitionError("precondition_required", "If-Match is required", "If-Match") + if draft.etag != etag: + return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match") + return None + + +async def freeze_draft( + store: WorkflowDefinitionStore, + caller: UserAPIKeyAuth, + workflow_id: str, + etag: str | None, + team_id: str | None = None, +) -> WorkflowSnapshot | DefinitionError: + draft: Final = await authorized_draft(store, caller, workflow_id, team_id) + if isinstance(draft, DefinitionError): + return draft + access: Final = await authorize_team(store, caller, draft.team_id, "write") + if isinstance(access, DefinitionError): + return access + if (failure := check_etag(draft, etag)) is not None: + return failure + if (invalid := validate_bundle(draft.bundle)) is not None: + return invalid + return await store.freeze_draft(workflow_id, draft.team_id, draft.etag) + + +def _local_references(value: JsonValue | Mapping[str, JsonValue], depth: int = 0) -> bool: + if depth > 64: + return False + if isinstance(value, Mapping): + if any(keyword in value for keyword in ("$id", "$dynamicRef", "$recursiveRef", "$dynamicAnchor")): + return False + reference: Final = value.get("$ref") + if reference is not None and (not isinstance(reference, str) or not reference.startswith("#/")): + return False + return all(_local_references(child, depth + 1) for child in value.values()) + if isinstance(value, list): + return all(_local_references(child, depth + 1) for child in value) + return True + + +def _schema_error(schema: Mapping[str, JsonValue], param: str) -> DefinitionError | None: + if not _local_references(schema): + return DefinitionError("invalid_bundle", "Schemas must use local references and bounded nesting", param) + try: + Draft202012Validator.check_schema(schema) + except (SchemaError, RecursionError): + return DefinitionError("invalid_bundle", "Invalid JSON schema", param) + return None + + +def _binding_error(binding: Binding, previous: frozenset[str], param: str) -> DefinitionError | None: + if isinstance(binding, ValueBinding): + return None + if binding.source == "$input" and binding.output == "messages": + return None + if binding.source in previous and binding.output == "result": + return None + return DefinitionError("invalid_bundle", "Bindings must reference input messages or an earlier node result", param) + + +def _node_error(node: ChatCompletionNode | PythonNode, previous: frozenset[str]) -> DefinitionError | None: + if node.id in previous or node.id == "runs": + return DefinitionError("invalid_bundle", "Node IDs must be unique and cannot be reserved", f"nodes.{node.id}") + if isinstance(node, ChatCompletionNode): + if node.model.startswith("workflow/") or "://" in node.model: + return DefinitionError("invalid_bundle", "Use a configured gateway model group", f"nodes.{node.id}.model") + if (failure := _binding_error(node.messages, previous, f"nodes.{node.id}.messages")) is not None: + return failure + if isinstance(node.messages, ValueBinding): + try: + _MESSAGES.validate_python(node.messages.value, strict=True) + except ValidationError: + return DefinitionError("invalid_bundle", "Invalid chat messages", f"nodes.{node.id}.messages") + if node.response_format is not None: + return _schema_error(node.response_format.json_schema.schema_, f"nodes.{node.id}.response_format") + return None + if set(node.inputs) != set(node.program.parameter_names): + return DefinitionError( + "invalid_bundle", "Bind exactly the Python program parameters", f"nodes.{node.id}.inputs" + ) + for name, binding in node.inputs.items(): + if (failure := _binding_error(binding, previous, f"nodes.{node.id}.inputs.{name}")) is not None: + return failure + compiled: Final = compile_program(node.program, _LIMITS) + if isinstance(compiled, MontyFailure): + return DefinitionError("invalid_bundle", compiled.message[:1024], f"nodes.{node.id}.program") + return None + + +def validate_bundle(bundle: WorkflowBundle) -> DefinitionError | None: + if bundle.name == "runs": + return DefinitionError("invalid_bundle", "The workflow name is reserved", "bundle.name") + if len(bundle.model_dump_json(by_alias=True).encode()) > 1024 * 1024: + return DefinitionError("invalid_bundle", "Workflow bundles cannot exceed 1 MiB", "bundle") + if (failure := _schema_error(bundle.input_schema, "bundle.input_schema")) is not None: + return failure + for index, node in enumerate(bundle.nodes): + previous: frozenset[str] = frozenset(prior.id for prior in bundle.nodes[:index]) + if (node_failure := _node_error(node, previous)) is not None: + return node_failure + if bundle.output.source != bundle.nodes[-1].id or bundle.output.output != "result": + return DefinitionError("invalid_bundle", "The workflow output must be the final node result", "bundle.output") + return None diff --git a/litellm/proxy/workflows/gateway.py b/litellm/proxy/workflows/gateway.py new file mode 100644 index 00000000000..698c0320fcc --- /dev/null +++ b/litellm/proxy/workflows/gateway.py @@ -0,0 +1,474 @@ +import asyncio +import json +import math +import re +import sys +import time +from collections.abc import Mapping +from contextvars import Context +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final, Literal, Protocol + +from fastapi import Request, Response +from openai.types.chat import ChatCompletionMessageParam +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic_core import from_json +from starlette.types import Message + +from litellm.exceptions import BudgetExceededError +from litellm.models.workflow_definition import ChatCompletionNode +from litellm.models.workflow_execution import ChildCallResult, ExecutionError, workflow_call_id +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import authorize_internal_virtual_key +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import ( + safe_set_request_parsed_body, # pyright: ignore[reportUnknownVariableType] # legacy helper accepts an untyped body +) +from litellm.proxy.spend_tracking.budget_reservation import ( + release_budget_reservation_on_cancel, # pyright: ignore[reportUnknownVariableType] # legacy reservation dictionary + release_unbound_budget_reservation, +) +from litellm.proxy.workflows.monty.models import MontyErrorKind +from litellm.proxy.workflows.monty.validation import ValidatorProcess +from litellm.repositories.table_repositories import SpendLogsRepository + +if sys.version_info >= (3, 11): + from asyncio import timeout +else: + from async_timeout import timeout + +_ROUTE: Final = "/v1/chat/completions" +_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam]) +_JSON: Final = TypeAdapter[JsonValue](JsonValue, config=ConfigDict(allow_inf_nan=False)) +_MAX_OUTPUT_BYTES: Final = 1024 * 1024 +_RECEIPT_LOOKUP_SECONDS: Final = 0.1 + + +@dataclass(frozen=True, slots=True) +class _NodeCall: + key_hash: str + team_id: str + node: ChatCompletionNode + messages: JsonValue + workflow_id: str + run_id: str + session_id: str + attempt: int + deadline: datetime + end_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class _RequestContext: + request: Request + data: dict[str, object] # mutable-ok: the existing auth and processor APIs enrich this attempt-owned body + response: Response + billing: asyncio.Future[ChildCallResult | None] + + +class _AssistantMessage(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + role: Literal["assistant"] + content: str | None + refusal: str | None = None + + +class _Choice(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + message: _AssistantMessage + finish_reason: str | None = None + + +class _Usage(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + prompt_tokens: int = Field(ge=0, strict=True) + completion_tokens: int = Field(ge=0, strict=True) + + +class _ChatResult(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + choices: tuple[_Choice, ...] = Field(min_length=1, max_length=1) + usage: _Usage + + +class _ResponseUsage(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + usage: _Usage + + +class _FinalizedBill(BaseModel): + model_config = ConfigDict(frozen=True, allow_inf_nan=False) + + cost_finalized: Literal[True] + prompt_tokens: int = Field(ge=0, strict=True) + completion_tokens: int = Field(ge=0, strict=True) + response_cost: float | None = Field(ge=0) + + +class _FailureData(BaseModel): + model_config = ConfigDict(frozen=True) + + standard_logging_object: _FinalizedBill + + +class _ProcessorFailure(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + data: _FailureData + + +class _ProxySettings(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + general_settings: Mapping[str, object] + user_model: str | None + user_temperature: float | None + user_request_timeout: float | None + user_max_tokens: int | None + user_api_base: str | None + version: str | None + + +class _Gateway(Protocol): + async def authorize(self, call: _NodeCall, context: _RequestContext) -> UserAPIKeyAuth: ... + + async def process(self, context: _RequestContext, auth: UserAPIKeyAuth) -> object: ... + + async def receipts(self, call: _NodeCall, call_id: str) -> tuple[str, ...]: ... + + async def validate_output( + self, output: JsonValue, schema: Mapping[str, JsonValue], deadline: float + ) -> ExecutionError | None: ... + + async def release(self, context: _RequestContext) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _ProxyGateway: + async def authorize(self, call: _NodeCall, context: _RequestContext) -> UserAPIKeyAuth: + return await authorize_internal_virtual_key( + key_hash=call.key_hash, + request=context.request, + request_data=context.data, + route=_ROUTE, + expected_team_id=call.team_id, + ) + + async def process(self, context: _RequestContext, auth: UserAPIKeyAuth) -> object: + from litellm.proxy import proxy_server + + processor: Final = ProxyBaseLLMRequestProcessing(data=context.data) + settings: Final = _ProxySettings.model_validate(vars(proxy_server)) + try: + try: + return await processor.base_process_llm_request( + request=context.request, + fastapi_response=context.response, + user_api_key_dict=auth, + route_type="acompletion", + proxy_logging_obj=proxy_server.proxy_logging_obj, + general_settings=dict(settings.general_settings), + proxy_config=proxy_server.proxy_config, + llm_router=proxy_server.llm_router, + user_model=settings.user_model, + user_temperature=settings.user_temperature, + user_request_timeout=settings.user_request_timeout, + user_max_tokens=settings.user_max_tokens, + user_api_base=settings.user_api_base, + version=settings.version, + ) + except Exception as exc: # noqa: BLE001 # route arbitrary provider and hook failures through accounting + try: + await processor.handle_llm_api_exception( + e=exc, user_api_key_dict=auth, proxy_logging_obj=proxy_server.proxy_logging_obj + ) + except Exception as failure: # noqa: BLE001 # sanitize errors raised by failure hooks + return _execution_error(failure) + finally: + context.billing.set_result(_failure_receipt(processor)) + except asyncio.CancelledError: + if not context.billing.done(): + context.billing.set_result(_failure_receipt(processor)) + await release_budget_reservation_on_cancel(auth.budget_reservation) + await proxy_server.proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth) + raise + + async def receipts(self, call: _NodeCall, call_id: str) -> tuple[str, ...]: + from litellm.proxy import proxy_server + + if proxy_server.prisma_client is None: + return () + return await SpendLogsRepository(proxy_server.prisma_client, use_writer=True).request_ids_for_call( + call_id=call_id, key_hash=call.key_hash, team_id=call.team_id, session_id=call.session_id + ) + + async def validate_output( + self, output: JsonValue, schema: Mapping[str, JsonValue], deadline: float + ) -> ExecutionError | None: + validator: Final = ValidatorProcess() + try: + failure: Final = await validator.validate(output, dict(schema), deadline=deadline, output=True) + if failure is None: + return None + if failure.kind is MontyErrorKind.TIMEOUT: + return ExecutionError("timeout", "The node deadline expired") + return ExecutionError("output_invalid", "The model output did not satisfy the node schema") + finally: + await validator.close() + + async def release(self, context: _RequestContext) -> None: + reservation: Final[object] = getattr(context.request.state, "budget_reservation", None) + if isinstance(reservation, Mapping): + await release_unbound_budget_reservation(reservation) # pyright: ignore[reportUnknownArgumentType] # preserve the reservation identity shared with accounting + + +def _node_data(call: _NodeCall, remaining: float) -> Mapping[str, object]: + messages: Final = _MESSAGES.validate_python(call.messages, strict=True) + if not messages: + raise ValueError("Chat messages must not be empty") + trace: Final = { + "workflow_id": call.workflow_id, + "run_id": call.run_id, + "node_id": call.node.id, + "attempt": call.attempt, + "session_id": call.session_id, + } + return { + **call.node.model_dump( + mode="json", + by_alias=True, + exclude={"id", "type", "messages", "timeout_seconds", "retry"}, + exclude_none=True, + ), + "messages": messages, + "stream": False, + "num_retries": 0, + "timeout": remaining, + "metadata": {**trace, "spend_logs_metadata": trace}, + **({"user": call.end_user_id} if call.end_user_id is not None else {}), + } + + +def _internal_request(data: Mapping[str, object], call_id: str, session_id: str) -> _RequestContext: + """Adapt a server-owned call to request-aware gateway hooks without dispatching HTTP or inheriting caller state.""" + from litellm.proxy import proxy_server + + body: Final = dict(data) + encoded: Final = json.dumps(body, allow_nan=False).encode() + incoming: Final[asyncio.Queue[Message]] = asyncio.Queue(maxsize=1) + incoming.put_nowait({"type": "http.request", "body": encoded, "more_body": False}) + + async def receive() -> Message: + return await incoming.get() + + request: Final = Request( + { + "type": "http", + "app": proxy_server.app, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": _ROUTE, + "root_path": "", + "query_string": b"", + "headers": ( + (b"content-type", b"application/json"), + (b"content-length", str(len(encoded)).encode()), + (b"x-litellm-call-id", call_id.encode()), + (b"x-litellm-session-id", session_id.encode()), + ), + "server": ("localhost", 80), + "client": None, + "state": {}, + }, + receive, + ) + safe_set_request_parsed_body(request, body) + return _RequestContext( + request=request, data=body, response=Response(), billing=asyncio.get_running_loop().create_future() + ) + + +def _execution_error(exc: Exception) -> ExecutionError: + if isinstance(exc, BudgetExceededError) or ( + isinstance(exc, ProxyException) and exc.type == ProxyErrorTypes.budget_exceeded + ): + return ExecutionError("budget_exceeded", "The execution identity has insufficient budget") + status: Final = exc.code if isinstance(exc, ProxyException) else getattr(exc, "status_code", None) + if str(status) in {"401", "403"}: + return ExecutionError("authorization_failed", "The execution identity is not authorized for this node") + if str(status) == "429": + return ExecutionError("rate_limit", "The gateway could not admit this node") + if str(status) in {"400", "422"}: + return ExecutionError("request_rejected", "The gateway rejected this node") + if str(status) == "408" or str(status).startswith("5"): + return ExecutionError("provider_error", "The provider could not complete this node") + return ExecutionError("runtime_error", "The gateway could not complete this node") + + +def _response_cost(response: Response) -> float | None: + try: + cost: Final = float(response.headers["x-litellm-response-cost"]) + return cost if math.isfinite(cost) and cost >= 0 else None + except (KeyError, ValueError): + return None + + +def _failure_receipt(processor: object) -> ChildCallResult | None: + try: + billing: Final = _ProcessorFailure.model_validate(processor).data.standard_logging_object + except ValidationError: + return None + return ChildCallResult( + output=None, + prompt_tokens=billing.prompt_tokens, + completion_tokens=billing.completion_tokens, + request_id=None, + response_cost=billing.response_cost, + ) + + +def _response_receipt(raw: object, response: Response) -> ChildCallResult | None: + try: + usage: Final = _ResponseUsage.model_validate(raw).usage + except ValidationError: + return None + return ChildCallResult( + output=None, + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + request_id=None, + response_cost=_response_cost(response), + ) + + +async def _receipt(call: _NodeCall, call_id: str, gateway: _Gateway, deadline: float) -> str | None: + try: + async with timeout(min(_RECEIPT_LOOKUP_SECONDS, max(0, deadline - time.monotonic()))): + ids: Final = await gateway.receipts(call, call_id) + return ids[0] if len(ids) == 1 else None + except Exception: # noqa: BLE001 # receipt lookup failure must not discard an already billed result + return None + + +async def _read_output( + result: _ChatResult, call: _NodeCall, gateway: _Gateway, deadline: float +) -> JsonValue | ExecutionError: + message: Final = result.choices[0].message + if message.refusal or result.choices[0].finish_reason == "content_filter": + return ExecutionError("content_rejected", "The gateway rejected the node output") + if message.content is None: + return ExecutionError("invalid_response", "The model returned no assistant content") + if len(message.content.encode()) > _MAX_OUTPUT_BYTES: + return ExecutionError("output_too_large", "The model output exceeded the workflow limit") + if call.node.response_format is None: + return message.content + try: + output: Final = _JSON.validate_python(from_json(message.content, allow_inf_nan=False)) + except ValueError: + return ExecutionError("output_invalid", "The model did not return valid JSON") + failure: Final = await gateway.validate_output(output, call.node.response_format.json_schema.schema_, deadline) + return failure if failure is not None else output + + +async def _admitted_call( + call: _NodeCall, gateway: _Gateway, context: _RequestContext, deadline: float +) -> ChildCallResult | ExecutionError: + auth: Final = await gateway.authorize(call, context) + raw: Final = await gateway.process(context, auth) + if isinstance(raw, ExecutionError): + return raw + context.billing.set_result(_response_receipt(raw, context.response)) + try: + result: Final = _ChatResult.model_validate(raw) + except ValidationError: + return ExecutionError("invalid_response", "The gateway returned an invalid chat response") + output: Final = await _read_output(result, call, gateway, deadline) + if isinstance(output, ExecutionError): + return output + return ChildCallResult( + output=output, + prompt_tokens=result.usage.prompt_tokens, + completion_tokens=result.usage.completion_tokens, + request_id=None, + response_cost=_response_cost(context.response), + ) + + +async def _timed_call( + call: _NodeCall, gateway: _Gateway, context: _RequestContext, deadline: float +) -> ChildCallResult | ExecutionError: + try: + async with timeout(max(0, deadline - time.monotonic())): + return await _admitted_call(call, gateway, context, deadline) + except asyncio.TimeoutError: + return ExecutionError("timeout", "The node deadline expired") + except Exception as exc: # noqa: BLE001 # auth, provider and validator failures need sanitized error values + return _execution_error(exc) + + +async def _execute(call: _NodeCall, gateway: _Gateway) -> ChildCallResult | ExecutionError: + if call.deadline.tzinfo is None or call.deadline.utcoffset() is None: + return ExecutionError("invalid_deadline", "The node deadline must include a timezone") + if re.fullmatch(r"[0-9a-f]{64}", call.key_hash) is None or not call.team_id: + return ExecutionError("authorization_failed", "A persisted team virtual key is required") + remaining: Final = min(call.node.timeout_seconds, (call.deadline - datetime.now(timezone.utc)).total_seconds()) + if remaining <= 0: + return ExecutionError("timeout", "The node deadline expired") + deadline: Final = time.monotonic() + remaining + try: + data: Final = _node_data(call, remaining) + call_id: Final = workflow_call_id(call.run_id, call.node.id, call.attempt) + context: Final = _internal_request(data, call_id, call.session_id) + except (ValueError, TypeError): + return ExecutionError("invalid_messages", "The node input must contain valid chat messages") + try: + result: Final = await _timed_call(call, gateway, context, deadline) + billing: Final = ( + result + if isinstance(result, ChildCallResult) + else context.billing.result() + if context.billing.done() + else result.receipt + ) + if billing is None: + return result + receipt: Final = await _receipt(call, call_id, gateway, deadline) + resolved: Final = replace(billing, request_id=receipt) + return ( + ExecutionError(result.code, result.message, receipt=resolved) + if isinstance(result, ExecutionError) + else resolved + ) + finally: + await gateway.release(context) + + +async def _execute_isolated(call: _NodeCall, gateway: _Gateway) -> ChildCallResult | ExecutionError: + task: Final = Context().run(asyncio.create_task, _execute(call, gateway)) + return await task + + +async def execute_chat_node( + *, + key_hash: str, + team_id: str, + node: ChatCompletionNode, + messages: JsonValue, + workflow_id: str, + run_id: str, + session_id: str, + attempt: int, + deadline: datetime, + end_user_id: str | None, +) -> ChildCallResult | ExecutionError: + call: Final = _NodeCall( + key_hash, team_id, node, messages, workflow_id, run_id, session_id, attempt, deadline, end_user_id + ) + return await _execute_isolated(call, _ProxyGateway()) diff --git a/litellm/proxy/workflows/lifespan.py b/litellm/proxy/workflows/lifespan.py new file mode 100644 index 00000000000..7db910be0bf --- /dev/null +++ b/litellm/proxy/workflows/lifespan.py @@ -0,0 +1,29 @@ +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Final + +from litellm.proxy.deps.repositories import ( + get_workflow_definition_repository, + get_workflow_execution_repository, +) +from litellm.proxy.workflows.gateway import execute_chat_node +from litellm.proxy.workflows.runner import WorkflowRunner +from litellm.repositories.database import Database + + +@asynccontextmanager +async def workflow_lifespan(database: Database | None) -> AsyncGenerator[None]: + if database is None: + yield + return + definitions: Final = await get_workflow_definition_repository(database) + executions: Final = await get_workflow_execution_repository(database) + if definitions is None or executions is None: + yield + return + runner: Final = WorkflowRunner(executions, definitions, execute_chat_node) + await runner.start() + try: + yield + finally: + await runner.close() diff --git a/litellm/proxy/workflows/monty/AGENTS.md b/litellm/proxy/workflows/monty/AGENTS.md new file mode 100644 index 00000000000..8bd696872e0 --- /dev/null +++ b/litellm/proxy/workflows/monty/AGENTS.md @@ -0,0 +1,11 @@ +# Monty security boundary + +Never pass host callbacks, mounts, environment variables, credentials, or live SDK objects into Monty + +Never import, exec, or eval uploaded source on the gateway. Inspect it with AST only + +Never fall back to CPython execution. Missing or unverified workers fail closed + +Monty's max_memory is an allocator budget, not an RSS limit. Apply the worker address-space limit separately + +Language isolation is not an OS security boundary. An interpreter escape requires separate OS-level containment diff --git a/litellm/proxy/workflows/monty/_compat.py b/litellm/proxy/workflows/monty/_compat.py new file mode 100644 index 00000000000..05be242098f --- /dev/null +++ b/litellm/proxy/workflows/monty/_compat.py @@ -0,0 +1,9 @@ +import sys +from typing import Final + +__all__: Final = ("timeout",) + +if sys.version_info >= (3, 11): + from asyncio import timeout +else: + from async_timeout import timeout diff --git a/litellm/proxy/workflows/monty/_validator_worker.py b/litellm/proxy/workflows/monty/_validator_worker.py new file mode 100644 index 00000000000..52b422ea2bc --- /dev/null +++ b/litellm/proxy/workflows/monty/_validator_worker.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import json +import sys +from collections.abc import Mapping +from typing import Final, NoReturn, TypeAlias, cast # noqa: TID251 # stdlib decoder produces the recursive JSON union + +from jsonschema import Draft202012Validator +from jsonschema.exceptions import ValidationError +from referencing import Registry +from referencing.exceptions import NoSuchResource, Unresolvable +from referencing.jsonschema import Schema + +JsonValue: TypeAlias = None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"] +MAX_MESSAGE_BYTES: Final = 64 * 1024 * 1024 + + +def refuse_retrieval(uri: str) -> NoReturn: + raise NoSuchResource(ref=uri) + + +def closed_registry() -> Registry[Schema]: + return Registry(retrieve=refuse_retrieval) + + +def validate_request(line: bytes) -> Mapping[str, JsonValue]: + try: + request: Final = cast(JsonValue, json.loads(line)) # cast-ok: stdlib JSON decoder returns only JSON values + if not isinstance(request, dict) or not isinstance(schema := request.get("schema"), dict): + return {"valid": False, "message": "invalid validator request"} + validator: Final = Draft202012Validator(schema, registry=closed_registry()) + validator.validate(request.get("value")) # pyright: ignore[reportUnknownMemberType] # upstream kwargs are untyped + except (ValidationError, Unresolvable, ValueError, RecursionError) as exc: + return {"valid": False, "message": str(exc)[:4096]} + else: + return {"valid": True, "message": ""} + + +def main() -> None: + sys.stdout.buffer.write(b"ready\n") + sys.stdout.buffer.flush() + while line := sys.stdin.buffer.readline(MAX_MESSAGE_BYTES + 1): + if len(line) > MAX_MESSAGE_BYTES or not line.endswith(b"\n"): + return + sys.stdout.buffer.write(json.dumps(validate_request(line), ensure_ascii=True).encode() + b"\n") + sys.stdout.buffer.flush() + + +if __name__ == "__main__": + main() diff --git a/litellm/proxy/workflows/monty/compiler.py b/litellm/proxy/workflows/monty/compiler.py new file mode 100644 index 00000000000..9f8af647641 --- /dev/null +++ b/litellm/proxy/workflows/monty/compiler.py @@ -0,0 +1,546 @@ +import ast +import hashlib +import keyword +import re +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType +from typing import Final + +from jsonschema import Draft202012Validator +from jsonschema.exceptions import SchemaError, ValidationError +from pydantic import JsonValue +from referencing.exceptions import Unresolvable + +from ._validator_worker import closed_registry +from .models import CompiledProgram, JsonSchema, MontyErrorKind, MontyFailure, MontyLimits, PythonProgram + +_FORBIDDEN: Final = frozenset( + { + "eval", + "exec", + "compile", + "open", + "__import__", + "globals", + "locals", + "vars", + "getattr", + "setattr", + "delattr", + "breakpoint", + "input", + "help", + } +) +_PREFIX: Final = "_liteflow_" +_MAX_SCHEMA_DEPTH: Final = 64 +_SCHEMA_MAPS: Final = frozenset({"properties", "$defs"}) +_SCHEMA_VALUES: Final = frozenset({"additionalProperties", "items", "not"}) +_SCHEMA_LISTS: Final = frozenset({"prefixItems", "anyOf", "oneOf", "allOf"}) +_SCHEMA_ANNOTATIONS: Final = frozenset( + { + "const", + "enum", + "default", + "examples", + "title", + "description", + "deprecated", + "readOnly", + "writeOnly", + "$comment", + "format", + "x-liteflow-record", + "type", + "required", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "minLength", + "maxLength", + "minItems", + "maxItems", + "minProperties", + "maxProperties", + "$ref", + } +) +_GENERATED_GLOBALS: Final = frozenset({"dataclass", "str", "int", "float", "bool", "list", "dict", "all", "isinstance"}) +_SCHEMA_EXPANSION_BUDGET: Final = 65536 + + +class _SchemaWork: + def __init__(self) -> None: + self.units = 0 + self.checked: frozenset[tuple[int, str]] = frozenset() + self.keywords: frozenset[tuple[int, str]] = frozenset() + self.contains: Mapping[tuple[int, str], bool] = MappingProxyType({}) + + def finish_keywords(self, key: tuple[int, str]) -> None: + self.keywords = self.keywords | {key} + + def finish_record(self, key: tuple[int, str]) -> None: + self.checked = self.checked | {key} + + def finish_contains(self, key: tuple[int, str], value: bool) -> None: + self.contains = MappingProxyType(dict(self.contains) | {key: value}) + + def spend(self, amount: int = 1) -> None: + self.units += amount + if self.units > _SCHEMA_EXPANSION_BUDGET: + raise ValueError("schema expansion budget exceeded") + + +def escape_diagnostics(value: str) -> str: + return "".join( + f"\\x{ord(char):02x}" if (ord(char) < 32 and char not in "\n\t") or 127 <= ord(char) <= 159 else char + for char in value + ) + + +def _invalid(rule: str, node: ast.AST) -> MontyFailure: + return MontyFailure( + MontyErrorKind.INVALID_PROGRAM, escape_diagnostics(f"{rule} at line {getattr(node, 'lineno', 1)}") + ) + + +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 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): + return _invalid("duplicate function name", tree) + entries: Final = tuple(node for node in functions if node.name == program.function_name) + if len(entries) != 1: + return _invalid("entry function must exist exactly once", tree) + args: Final = entries[0].args + if args.posonlyargs or args.vararg or args.kwarg or args.kwonlyargs or args.defaults or args.kw_defaults: + 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]) + 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: + return value if isinstance(value, dict) else {} + + +def _resolve(schema: JsonSchema, root: JsonSchema, work: _SchemaWork) -> JsonSchema: + work.spend() + ref: Final = schema.get("$ref") + if not isinstance(ref, str): + return schema + + def _monty_schema_pointer(value: JsonValue, parts: tuple[str, ...]) -> JsonSchema: + work.spend() + if len(parts) > _MAX_SCHEMA_DEPTH: + return {} + if not parts: + return _mapping(value) + if isinstance(value, dict): + return _monty_schema_pointer(value.get(parts[0].replace("~1", "/").replace("~0", "~")), parts[1:]) + if isinstance(value, list) and parts[0].isdigit() and int(parts[0]) < len(value): + return _monty_schema_pointer(value[int(parts[0])], parts[1:]) + return {} + + return _monty_schema_pointer(root, tuple(ref[2:].split("/"))) + + +def _monty_reference_exists(root: JsonValue, parts: tuple[str, ...], work: _SchemaWork) -> bool: + work.spend() + if len(parts) > _MAX_SCHEMA_DEPTH: + return False + 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 _monty_reference_exists(root[part], parts[1:], work) + if isinstance(root, list) and part.isdigit() and int(part) < len(root): + return _monty_reference_exists(root[int(part)], parts[1:], work) + return False + + +def _monty_schema_error( + value: JsonValue, root: JsonSchema, work: _SchemaWork, depth: int = 0, seen: frozenset[str] = frozenset() +) -> str | None: + work.spend() + if depth > _MAX_SCHEMA_DEPTH: + return "schema depth limit exceeded" + if not isinstance(value, dict): + return None + for key, item in value.items(): + work.spend() + if key not in _SCHEMA_MAPS | _SCHEMA_VALUES | _SCHEMA_LISTS | _SCHEMA_ANNOTATIONS: + 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 == "$ref" + and isinstance(item, str) + and not _monty_reference_exists(root, tuple(item[2:].split("/")), work) + ): + return "local JSON reference target does not exist" + if ( + key == "$ref" + and isinstance(item, str) + and (reference_error := _reference_schema_error(value, root, work, depth, seen, item)) is not None + ): + return reference_error + return next( + ( + error + for child in _children(value, work) + if (error := _monty_schema_error(child, root, work, depth + 1, seen)) is not None + ), + None, + ) + + +def _reference_schema_error( + schema: JsonSchema, root: JsonSchema, work: _SchemaWork, depth: int, seen: frozenset[str], ref: str +) -> str | None: + if ref in seen: + return "recursive schema references are unsupported" + cache_key: Final = (id(root), ref) + if cache_key in work.keywords: + return None + result: Final = _monty_schema_error(_resolve(schema, root, work), root, work, depth + 1, seen | {ref}) + if result is None: + work.finish_keywords(cache_key) + return result + + +def _record_shape_error(schema: JsonSchema, root: JsonSchema, work: _SchemaWork) -> str | None: + if schema.get("x-liteflow-record") is not True: + return None + if schema not in tuple(record for _, record in _records(root, work)): + 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], work: _SchemaWork) -> str | None: + for keyword_name in ("anyOf", "oneOf", "allOf"): + if not isinstance(alternatives := schema.get(keyword_name), list): + continue + if not any(_monty_contains_record(_mapping(item), root, seen, work) 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 + + +def _monty_record_error( + schema: JsonSchema, root: JsonSchema, work: _SchemaWork, seen: frozenset[str] = frozenset(), depth: int = 0 +) -> str | None: + work.spend() + if depth > _MAX_SCHEMA_DEPTH: + return "schema reference depth limit exceeded" + ref: Final = schema.get("$ref") + if isinstance(ref, str): + if ref in seen: + return "recursive schema references are unsupported" + cache_key: Final = (id(root), ref) + if cache_key in work.checked: + return None + result: Final = _monty_record_error(_resolve(schema, root, work), root, work, seen | {ref}, depth + 1) + if result is None: + work.finish_record(cache_key) + return result + if (shape_error := _record_shape_error(schema, root, work)) is not None: + return shape_error + if (union_error := _record_union_error(schema, root, seen, work)) is not None: + return union_error + return next( + ( + error + for child in _children(schema, work) + if (error := _monty_record_error(child, root, work, seen, depth + 1)) is not None + ), + None, + ) + + +def _monty_contains_record( + schema: JsonSchema, root: JsonSchema, seen: frozenset[str], work: _SchemaWork, depth: int = 0 +) -> bool: + work.spend() + if depth > _MAX_SCHEMA_DEPTH: + return False + ref: Final = schema.get("$ref") + if isinstance(ref, str): + if ref in seen: + return False + cache_key: Final = (id(root), ref) + if cache_key in work.contains: + return work.contains[cache_key] + result: Final = _monty_contains_record(_resolve(schema, root, work), root, seen | {ref}, work, depth + 1) + work.finish_contains(cache_key, result) + return result + if schema.get("x-liteflow-record") is True: + return True + return any(_monty_contains_record(child, root, seen, work, depth + 1) for child in _children(schema, work)) + + +def _children(schema: JsonSchema, work: _SchemaWork) -> tuple[JsonSchema, ...]: + def children_of(key: str, value: JsonValue) -> tuple[JsonSchema, ...]: + work.spend() + if key in {"properties", "$defs"} and isinstance(value, dict): + work.spend(len(value)) + return tuple(_mapping(child) for child in value.values()) + if key in {"anyOf", "oneOf", "allOf", "prefixItems"} and isinstance(value, list): + work.spend(len(value)) + return tuple(_mapping(child) for child in value) + if key in {"items", "additionalProperties", "not"} and isinstance(value, dict): + return (value,) + return () + + return tuple(chain.from_iterable(children_of(key, value) for key, value in schema.items())) + + +def _records(root: JsonSchema, work: _SchemaWork) -> tuple[tuple[str, JsonSchema], ...]: + work.spend(len(_mapping(root.get("$defs")))) + return tuple( + (name, value) + for name, value in _mapping(root.get("$defs")).items() + if isinstance(value, dict) and value.get("x-liteflow-record") is True + ) + + +def _monty_build(schema: JsonSchema, root: JsonSchema, expression: str, work: _SchemaWork, depth: int = 0) -> str: + work.spend() + result: Final = _build_expression(schema, root, expression, work, depth) + work.spend(len(result.encode("utf-8"))) + return result + + +def _build_expression(schema: JsonSchema, root: JsonSchema, expression: str, work: _SchemaWork, depth: int) -> str: + if depth > _MAX_SCHEMA_DEPTH: + raise ValueError("schema builder depth limit exceeded") + resolved: Final = _resolve(schema, root, work) + ref: Final = schema.get("$ref") + if isinstance(ref, str) and "$ref" in resolved: + return _monty_build(resolved, root, expression, work, depth + 1) + union: Final = resolved.get("anyOf") + if isinstance(union, list) and _monty_contains_record(resolved, root, frozenset(), work): + other: Final = next(_mapping(item) for item in union if _mapping(item).get("type") != "null") + return f"(None if {expression} is None else {_monty_build(other, root, expression, work, depth + 1)})" + if resolved.get("x-liteflow-record") is True: + name: Final = next(name for name, candidate in _records(root, work) if candidate == resolved) + required: Final = resolved.get("required", []) + + def field(name: str, child: JsonValue) -> str: + access: Final = ( + f"{expression}[{name!r}]" + if isinstance(required, list) and name in required + else f"{expression}.get({name!r})" + ) + built: Final = _monty_build(_mapping(child), root, access, work, depth + 1) + return ( + f"{name}={built}" + if isinstance(required, list) and name in required + else f"{name}=(None if {access} is None else {built})" + ) + + return ( + f"{name}(" + + ", ".join(field(key, child) for key, child in _mapping(resolved.get("properties")).items()) + + ")" + ) + if resolved.get("type") == "array": + variable: Final = f"{_PREFIX}item{depth}" + return f"[{_monty_build(_mapping(resolved.get('items')), root, variable, work, depth + 1)} for {variable} in {expression}]" + if resolved.get("type") == "object" and isinstance(resolved.get("properties"), dict): + fields: Final = _mapping(resolved.get("properties")) + variable_key: Final = f"{_PREFIX}key{depth}" + variable_value: Final = f"{_PREFIX}value{depth}" + branches: Final = "".join( + f"{_monty_build(_mapping(child), root, variable_value, work, depth + 1)} if {variable_key} == {key!r} else " + for key, child in fields.items() + ) + fallback: Final = _monty_build( + _mapping(resolved.get("additionalProperties")), root, variable_value, work, depth + 1 + ) + return ( + f"{{{variable_key}: ({branches}{fallback}) for {variable_key}, {variable_value} in {expression}.items()}}" + ) + return expression + + +def _record_class(name: str, schema: JsonSchema, work: _SchemaWork) -> str: + required: Final = schema.get("required", []) + fields: Final = _mapping(schema.get("properties")) + work.spend(len(fields)) + mandatory: Final = tuple(key for key in fields if isinstance(required, list) and key in required) + optional: Final = tuple(key for key in fields if key not in mandatory) + return ( + f"@dataclass(frozen=True)\nclass {name}:\n" + + ( + "\n".join(f" {key}: object" for key in mandatory) + + "\n" + + "\n".join(f" {key}: object = None" for key in optional) + if fields + else " pass" + ) + + "\n" + ) + + +def _source(program: PythonProgram, work: _SchemaWork) -> str: + records: Final = dict(_records(program.input_schema, work) + _records(program.output_schema, work)) + classes: Final = "\n".join(_record_class(name, schema, work) for name, schema in records.items()) + record_dump: Final = "\n".join( + f" if isinstance(value, {name}):\n return {_PREFIX}dump({{" + + ", ".join(f"{field!r}: value.{field}" for field in _mapping(schema.get("properties"))) + + "})" + for name, schema in records.items() + ) + dumper: Final = f"""def {_PREFIX}dump(value): +{record_dump} + if value is None or isinstance(value, (str, int, float, bool)): + return [True, value] + if isinstance(value, list): + pairs = [{_PREFIX}dump(item) for item in value] + return [all(pair[0] for pair in pairs), [pair[1] for pair in pairs]] + if isinstance(value, dict): + pairs = {{key: {_PREFIX}dump(item) for key, item in value.items()}} + return [all(isinstance(key, str) and pair[0] for key, pair in pairs.items()), {{key: pair[1] for key, pair in pairs.items()}}] + return [False, None] +""" + arguments: Final = ", ".join( + _monty_build( + _mapping(_mapping(program.input_schema.get("properties")).get(name)), + program.input_schema, + f"{_PREFIX}input[{name!r}]", + work, + ) + for name in program.parameter_names + ) + return ( + "from dataclasses import dataclass\n" + + classes + + "\n" + + program.source + + "\n" + + dumper + + f"\n{_PREFIX}dump({program.function_name}({arguments}))" + ) + + +def _schemas_error(program: PythonProgram, tree: ast.Module, work: _SchemaWork) -> MontyFailure | None: + function_names: Final = frozenset(node.name for node in tree.body if isinstance(node, ast.FunctionDef)) + validator: Final = Draft202012Validator(Draft202012Validator.META_SCHEMA, registry=closed_registry()) + for schema in (program.input_schema, program.output_schema): + if (error := _monty_schema_error(schema, schema, work)) is not None: + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(error)) + validator.validate(schema) # pyright: ignore[reportUnknownMemberType] # upstream kwargs are untyped + if (record_error := _monty_record_error(schema, schema, work)) is not None: + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(record_error)) + for name, _ in _records(schema, work): + 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, work)) + if any( + name in input_records and input_records[name] != schema + for name, schema in _records(program.output_schema, work) + ): + 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") + if len(encoded) > limits.max_source_bytes: + return MontyFailure(MontyErrorKind.INVALID_PROGRAM, "source byte limit exceeded at line 1") + if hashlib.sha256(encoded).hexdigest() != program.source_sha256: + return MontyFailure(MontyErrorKind.INVALID_PROGRAM, "source_sha256 mismatch at line 1") + if re.search(r"(?m)^\s*#\s*///\s*script\b", program.source): + return MontyFailure(MontyErrorKind.INVALID_PROGRAM, "PEP 723 script blocks forbidden at line 1") + tree: Final = ast.parse(program.source) + if (failure := _source_error(tree, program)) is not None: + return failure + work: Final = _SchemaWork() + if (schema_failure := _schemas_error(program, tree, work)) is not None: + return schema_failure + generated: Final = _source(program, work) + work.spend(len(generated.encode("utf-8"))) + return CompiledProgram(program, generated, program.input_schema, program.output_schema) + except SyntaxError as exc: + return MontyFailure( + MontyErrorKind.INVALID_PROGRAM, escape_diagnostics(f"syntax error at line {exc.lineno}: {exc.msg}") + ) + except (SchemaError, ValidationError, Unresolvable, ValueError, RecursionError, StopIteration) as exc: + return MontyFailure(MontyErrorKind.UNSUPPORTED_SCHEMA, escape_diagnostics(str(exc))) diff --git a/litellm/proxy/workflows/monty/models.py b/litellm/proxy/workflows/monty/models.py new file mode 100644 index 00000000000..e9f79126b5f --- /dev/null +++ b/litellm/proxy/workflows/monty/models.py @@ -0,0 +1,86 @@ +import keyword +from dataclasses import dataclass +from enum import Enum +from typing import Annotated, Literal, TypeAlias + +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, StringConstraints + +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, extra="forbid") + + contract_version: Literal["1"] + source: str + function_name: ProgramName + parameter_names: Annotated[tuple[ProgramName, ...], Field(min_length=1, max_length=32)] + input_schema: JsonSchema + output_schema: JsonSchema + source_sha256: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + + +class MontyErrorKind(str, Enum): + INVALID_PROGRAM = "INVALID_PROGRAM" + UNSUPPORTED_SCHEMA = "UNSUPPORTED_SCHEMA" + INPUT_INVALID = "INPUT_INVALID" + INPUT_TOO_LARGE = "INPUT_TOO_LARGE" + OUTPUT_INVALID = "OUTPUT_INVALID" + OUTPUT_TOO_LARGE = "OUTPUT_TOO_LARGE" + RUNTIME_ERROR = "RUNTIME_ERROR" + TIMEOUT = "TIMEOUT" + MEMORY_LIMIT = "MEMORY_LIMIT" + RECURSION_LIMIT = "RECURSION_LIMIT" + STDOUT_LIMIT = "STDOUT_LIMIT" + WORKER_CRASHED = "WORKER_CRASHED" + QUEUE_FULL = "QUEUE_FULL" + PREDICATE_NOT_BOOL = "PREDICATE_NOT_BOOL" + CANCELLED = "CANCELLED" + RUNTIME_UNAVAILABLE = "RUNTIME_UNAVAILABLE" + + +@dataclass(frozen=True, slots=True) +class MontyFailure: + kind: MontyErrorKind + message: str + + +@dataclass(frozen=True, slots=True) +class MontySuccess: + output: JsonValue + stdout: str + duration_ms: int + + +@dataclass(frozen=True, slots=True) +class CompiledProgram: + program: PythonProgram + sandbox_source: str + input_schema: JsonSchema + output_schema: JsonSchema + + +class MontyLimits(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False) + + 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 new file mode 100644 index 00000000000..1165e217b88 --- /dev/null +++ b/litellm/proxy/workflows/monty/runtime.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import importlib.metadata +import json +import math +import os +import signal +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Final + +from pydantic import JsonValue, TypeAdapter +from pydantic import ValidationError as PydanticValidationError + +from ._compat import timeout +from .compiler import escape_diagnostics +from .models import CompiledProgram, JsonSchema, MontyErrorKind, MontyFailure, MontyLimits, MontySuccess +from .validation import ValidatorProcess + +if TYPE_CHECKING: + from pydantic_monty import AsyncMonty, AsyncMontySession, CollectString, MontyRuntimeError + +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +@dataclass(frozen=True, slots=True) +class BinaryArtifact: + path: Path + sha256: str + + +def installed_binary() -> BinaryArtifact | MontyFailure: + try: + distribution: Final = importlib.metadata.distribution("pydantic-monty-runtime") + candidates: Final = tuple(file for file in distribution.files or () if file.name in {"monty", "monty.exe"}) + if len(candidates) != 1: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "runtime binary missing or ambiguous") + file: Final = candidates[0] + 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(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 _monty_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): + return None + if type(value) is float: + return None if math.isfinite(value) else MontyErrorKind.INPUT_INVALID + if type(value) is list: + return next( + (issue for item in value if (issue := _monty_json_issue(item, depth + 1, max_depth)) is not None), None + ) + if type(value) is dict: + if any(type(key) is not str for key in value): + return MontyErrorKind.INPUT_INVALID + return next( + ( + child_issue + for item in value.values() + if (child_issue := _monty_json_issue(item, depth + 1, max_depth)) is not None + ), + None, + ) + return MontyErrorKind.INPUT_INVALID + + +def validate_json(value: object, *, max_bytes: int, max_depth: int, output: bool = False) -> 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 + try: + parsed: Final = _JSON.validate_python(value, strict=True) + issue: Final = _monty_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") + 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 (ValueError, OverflowError, RecursionError) as exc: + return MontyFailure(invalid, escape_diagnostics(str(exc))[:4096]) + else: + return parsed + + +def _kill(session: AsyncMontySession) -> None: + pid: Final = session.worker_pid + if pid is None: + return + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + 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 + self._binary = binary + self._pool: AsyncMonty | None = None + self._slots = asyncio.Semaphore(limits.max_concurrent_workers) + self._admitted = 0 + self._validators = tuple(ValidatorProcess() for _ in range(limits.max_concurrent_workers)) + + async def start(self) -> MontyFailure | None: + if self._pool is not None: + return None + try: + from pydantic_monty import AsyncMonty + + artifact: Final = self._binary if self._binary is not None else installed_binary() + if isinstance(artifact, MontyFailure): + return artifact + if hashlib.sha256(artifact.path.read_bytes()).hexdigest() != artifact.sha256: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "runtime binary SHA-256 mismatch") + pool: Final = AsyncMonty( + binary_path=artifact.path, + min_processes=0, + max_processes=self.limits.max_concurrent_workers, + max_checkouts_per_worker=1, + checkout_timeout=self.limits.feed_seconds, + request_timeout=self.limits.feed_seconds + 1, + feed_duration_limit_grace=0.1, + ) + await pool.__aenter__() + validators: Final = await asyncio.gather(*(worker.start() for worker in self._validators)) + if (failure := next((result for result in validators if result is not None), None)) is not None: + await asyncio.gather(*(worker.close() for worker in self._validators)) + await pool.__aexit__(None, None, None) + return failure + self._pool = pool + 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 + self._pool = None + await asyncio.gather(*(worker.close() for worker in self._validators)) + if pool is not None: + await pool.__aexit__(None, None, None) + + async def run_program( + self, compiled: CompiledProgram, input: object, *, deadline_monotonic: float + ) -> MontySuccess | MontyFailure: + started: Final = time.monotonic() + if started >= deadline_monotonic: + return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired") + if self._pool is None: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "Monty runtime is not started") + if self._admitted >= self.limits.max_concurrent_workers + self.limits.max_queued_waiters: + return MontyFailure(MontyErrorKind.QUEUE_FULL, "Monty execution queue is full") + self._admitted += 1 + try: + async with timeout(max(0, deadline_monotonic - time.monotonic())), self._slots: + worker: Final = next(worker for worker in self._validators if not worker.busy) + worker.busy = True + try: + validated: Final = await self._validate(input, compiled.input_schema, worker, deadline_monotonic) + if isinstance(validated, MontyFailure): + return validated + return await self._execute(compiled, validated, deadline_monotonic, started, worker) + except asyncio.CancelledError: + await worker.close() + raise + finally: + worker.busy = False + except asyncio.TimeoutError: + return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired") + finally: + self._admitted -= 1 + + async def _execute( + self, compiled: CompiledProgram, input: JsonValue, deadline: float, started: float, worker: ValidatorProcess + ) -> MontySuccess | MontyFailure: + from pydantic_monty import CollectString, MontyCrashedError, MontyRuntimeError, MontySyntaxError, ResourceLimits + + pool: Final = self._pool + if pool is None: + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "Monty runtime is closed") + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired") + collector: Final = CollectString(max_bytes=self.limits.max_stdout_bytes) + session: Final = pool.checkout( + script_name="workflow.py", + limits=ResourceLimits( + max_feed_duration_secs=min(self.limits.feed_seconds, remaining), + max_memory=self.limits.max_memory, + max_recursion_depth=self.limits.max_recursion_depth, + ), + os_policy={"sleep": "zero"}, + ) + try: + async with session: + try: + 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 await self._output(result, compiled, collector.output, deadline, started, worker) + except asyncio.CancelledError: + _kill(session) + raise + except MontyRuntimeError as exc: + return _runtime_failure(exc, collector, self.limits) + except MontyCrashedError as exc: + return MontyFailure( + MontyErrorKind.TIMEOUT if exc.timed_out else MontyErrorKind.WORKER_CRASHED, + escape_diagnostics(str(exc))[:4096], + ) + except MontySyntaxError as exc: + return MontyFailure(MontyErrorKind.INVALID_PROGRAM, escape_diagnostics(str(exc))[:4096]) + except (ValueError, OSError) as exc: + return MontyFailure(MontyErrorKind.RUNTIME_ERROR, escape_diagnostics(str(exc))[:4096]) + + async def _output( + self, + result: object, + compiled: CompiledProgram, + stdout: str, + deadline: float, + started: float, + worker: ValidatorProcess, + ) -> 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 = await self._validate(envelope[1], compiled.output_schema, worker, deadline, 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 _validate( + self, value: object, schema: JsonSchema, worker: ValidatorProcess, deadline: float, *, output: bool = False + ) -> JsonValue | MontyFailure: + parsed: Final = validate_json( + value, + max_bytes=self.limits.max_output_bytes if output else self.limits.max_input_bytes, + max_depth=self.limits.max_output_depth if output else self.limits.max_input_depth, + output=output, + ) + if isinstance(parsed, MontyFailure): + return parsed + failure: Final = await worker.validate(parsed, schema, deadline=deadline, output=output) + return failure if failure is not None else parsed + + async def evaluate_predicate( + self, compiled: CompiledProgram, input: object, *, deadline_monotonic: float + ) -> bool | MontyFailure: + result: Final = await self.run_program(compiled, input, deadline_monotonic=deadline_monotonic) + if isinstance(result, MontyFailure): + return result + if type(result.output) is not bool: + return MontyFailure(MontyErrorKind.PREDICATE_NOT_BOOL, "predicate must return exactly bool") + return result.output diff --git a/litellm/proxy/workflows/monty/validation.py b/litellm/proxy/workflows/monty/validation.py new file mode 100644 index 00000000000..35cdff78cf2 --- /dev/null +++ b/litellm/proxy/workflows/monty/validation.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import asyncio +import json +import sys +import time +from pathlib import Path +from typing import Final + +from pydantic import BaseModel, ConfigDict, JsonValue, StrictBool, StrictStr + +from ._compat import timeout +from ._validator_worker import MAX_MESSAGE_BYTES +from .compiler import escape_diagnostics +from .models import JsonSchema, MontyErrorKind, MontyFailure + + +class _Response(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + valid: StrictBool + message: StrictStr + + +class ValidatorProcess: + def __init__(self) -> None: + self.busy = False + self.process: asyncio.subprocess.Process | None = None + + async def start(self) -> MontyFailure | None: + if self.process is not None and self.process.returncode is None: + return None + try: + async with timeout(5): + self.process = await asyncio.create_subprocess_exec( + sys.executable, + "-I", + str(Path(__file__).with_name("_validator_worker.py")), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + env={}, + limit=65536, + ) + if self.process.stdout is None or await self.process.stdout.readline() != b"ready\n": + await self.close() + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "validator startup failed") + except asyncio.CancelledError: + await self.close() + raise + except (OSError, ValueError, asyncio.TimeoutError) as exc: + await self.close() + return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc))) + return None + + async def close(self) -> None: + process: Final = self.process + self.process = None + if process is not None: + if process.returncode is None: + try: + process.kill() + except ProcessLookupError: + pass + await process.wait() + + async def validate( + self, value: JsonValue, schema: JsonSchema, *, deadline: float, output: bool = False + ) -> MontyFailure | None: + try: + async with timeout(max(0, deadline - time.monotonic())): + if (failure := await self.start()) is not None: + return failure + process: Final = self.process + if process is None or process.stdin is None or process.stdout is None: + return MontyFailure(MontyErrorKind.WORKER_CRASHED, "validator pipes unavailable") + request: Final = json.dumps({"value": value, "schema": schema}, ensure_ascii=True).encode() + b"\n" + if len(request) > MAX_MESSAGE_BYTES: + return MontyFailure( + MontyErrorKind.OUTPUT_TOO_LARGE if output else MontyErrorKind.INPUT_TOO_LARGE, + "validator request byte limit exceeded", + ) + process.stdin.write(request) + await process.stdin.drain() + response: Final = _Response.model_validate_json(await process.stdout.readline()) + if not response.valid: + return MontyFailure( + MontyErrorKind.OUTPUT_INVALID if output else MontyErrorKind.INPUT_INVALID, + escape_diagnostics(response.message)[:4096], + ) + except asyncio.CancelledError: + await self.close() + raise + except asyncio.TimeoutError: + await self.close() + return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired during schema validation") + except (OSError, ValueError, RecursionError) as exc: + await self.close() + return MontyFailure(MontyErrorKind.WORKER_CRASHED, escape_diagnostics(str(exc))[:4096]) + return None diff --git a/litellm/proxy/workflows/runner.py b/litellm/proxy/workflows/runner.py new file mode 100644 index 00000000000..9d5782ed1db --- /dev/null +++ b/litellm/proxy/workflows/runner.py @@ -0,0 +1,369 @@ +import asyncio +import time +from collections.abc import Callable, Mapping +from datetime import datetime, timedelta, timezone +from typing import Final, Protocol +from uuid import uuid4 + +from pydantic import JsonValue + +from litellm._logging import verbose_proxy_logger +from litellm.models.workflow_definition import ( + Binding, + ChatCompletionNode, + ValueBinding, + WorkflowBundle, + WorkflowDefinitionStore, + WorkflowNode, +) +from litellm.models.workflow_execution import ( + ChildCallResult, + EngineRun, + ExecutionError, + RunState, + StepTrace, + WorkflowExecutionStore, + workflow_call_id, +) +from litellm.proxy.workflows.monty.compiler import compile_program +from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits +from litellm.proxy.workflows.monty.runtime import MontyRuntime +from litellm.proxy.workflows.monty.validation import ValidatorProcess + + +class ChatExecutor(Protocol): + async def __call__( + self, + *, + key_hash: str, + team_id: str, + workflow_id: str, + node: ChatCompletionNode, + messages: JsonValue, + run_id: str, + session_id: str, + attempt: int, + deadline: datetime, + end_user_id: str | None, + ) -> ChildCallResult | ExecutionError: ... + + +def utc_now() -> datetime: + return datetime.now(timezone.utc) + + +class SchemaValidator(Protocol): + async def __call__(self, messages: JsonValue, schema: Mapping[str, JsonValue]) -> ExecutionError | None: ... + + +async def _validate_run_input(messages: JsonValue, schema: Mapping[str, JsonValue]) -> ExecutionError | None: + validator: Final = ValidatorProcess() + try: + failure: Final = await validator.validate(messages, dict(schema), deadline=time.monotonic() + 5) + return ( + ExecutionError("invalid_input", "Workflow messages did not satisfy the pinned input schema") + if failure is not None + else None + ) + finally: + await validator.close() + + +def resolve_binding(binding: Binding, run: EngineRun, steps: tuple[StepTrace, ...]) -> JsonValue | ExecutionError: + if isinstance(binding, ValueBinding): + return binding.value + if binding.source == "$input": + return [dict(message) for message in run.input.messages] + source: Final = next((step for step in steps if step.node_id == binding.source), None) + if source is None or source.state is not RunState.COMPLETED: + return ExecutionError("missing_input", "The referenced node has no completed output") + return source.output + + +def _bound_inputs( + bindings: Mapping[str, Binding], run: EngineRun, steps: tuple[StepTrace, ...] +) -> JsonValue | ExecutionError: + resolved: Final = {name: resolve_binding(binding, run, steps) for name, binding in bindings.items()} + failure: Final = next((value for value in resolved.values() if isinstance(value, ExecutionError)), None) + if failure is not None: + return failure + return {name: value for name, value in resolved.items() if not isinstance(value, ExecutionError)} + + +async def _settle_task(task: asyncio.Task[object]) -> None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +class WorkflowRunner: + def __init__( + self, + runs: WorkflowExecutionStore, + definitions: WorkflowDefinitionStore, + execute_chat: ChatExecutor, + *, + clock: Callable[[], datetime] = utc_now, + validate_input: SchemaValidator = _validate_run_input, + ) -> None: + self.runs: Final = runs + self.definitions: Final = definitions + self._execute_chat: Final = execute_chat + self._clock: Final = clock + self._validate_input: Final = validate_input + self._monty: Final = MontyRuntime(MontyLimits()) + self._monty_start: Final = asyncio.Lock() + self._poller: asyncio.Task[None] | None = None + self._active: tuple[asyncio.Task[None], ...] = () + + async def start(self) -> None: + self._poller = asyncio.create_task(self._poll()) + + async def close(self) -> None: + if self._poller is not None: + self._poller.cancel() + await asyncio.gather(self._poller, return_exceptions=True) + for task in self._active: + task.cancel() + await asyncio.gather(*self._active, return_exceptions=True) + await self._monty.close() + + async def _poll(self) -> None: + while True: + for task in self._active: + if task.done() and not task.cancelled() and task.exception() is not None: + verbose_proxy_logger.error("Workflow execution task failed before settlement") + self._active = tuple(task for task in self._active if not task.done()) + try: + pending: tuple[EngineRun, ...] = ( + await self.runs.recover(4 - len(self._active)) if len(self._active) < 4 else () + ) + self._active = (*self._active, *(asyncio.create_task(self._run(run.run_id)) for run in pending)) + except Exception: # noqa: BLE001 # polling must survive transient database failures without logging inputs + verbose_proxy_logger.error("Workflow dispatch polling failed") + await asyncio.sleep(1) + + async def _heartbeat(self, run: EngineRun, token: str) -> ExecutionError: + while True: + await asyncio.sleep(1) + current: EngineRun | None = await self.runs.get(run.run_id) + if current is None or current.lease_token != token: + return ExecutionError("lease_lost", "Execution ownership changed") + if current.cancel_requested: + return ExecutionError("cancelled", "Cancellation requested") + if self._clock() >= run.input.deadline: + return ExecutionError("deadline_exceeded", "Workflow deadline expired") + renewed: bool = await self.runs.renew(run.run_id, token, self._clock() + timedelta(seconds=15)) + if not renewed: + return ExecutionError("lease_lost", "Execution ownership expired") + + async def _run(self, run_id: str) -> None: + token: Final = str(uuid4()) + run: Final = await self.runs.claim(run_id, token, self._clock() + timedelta(seconds=15)) + if run is None: + return + if run.cancel_requested: + await self.runs.finish(run_id, token, RunState.CANCELLED, None, "Cancellation requested") + return + execution: Final = asyncio.create_task(self._execute(run, token)) + heartbeat: Final = asyncio.create_task(self._heartbeat(run, token)) + try: + done, _ = await asyncio.wait((execution, heartbeat), return_when=asyncio.FIRST_COMPLETED) + if execution in done: + result: Final = execution.result() + await self.runs.finish( + run_id, + token, + RunState.FAILED if isinstance(result, ExecutionError) else RunState.COMPLETED, + None if isinstance(result, ExecutionError) else result, + result.message if isinstance(result, ExecutionError) else None, + ) + else: + reason: Final = heartbeat.result() + await _settle_task(execution) + await self.runs.finish( + run_id, + token, + RunState.CANCELLED if reason.code == "cancelled" else RunState.FAILED, + None, + reason.message, + ) + except asyncio.CancelledError: + await _settle_task(execution) + await self.runs.finish(run_id, token, RunState.FAILED, None, "Execution worker stopped") + raise + except Exception: # noqa: BLE001 # settle durable execution without exposing provider errors or credentials + await _settle_task(execution) + await self.runs.finish(run_id, token, RunState.FAILED, None, "Execution could not complete") + finally: + await _settle_task(heartbeat) + + async def _execute(self, run: EngineRun, token: str) -> JsonValue | ExecutionError: + snapshot: Final = await self.definitions.get_snapshot(run.snapshot_id, run.team_id) + if snapshot is None or snapshot.workflow_id != run.workflow_id: + return ExecutionError("snapshot_unavailable", "The pinned workflow definition is unavailable") + invalid: Final = await self._validate_input( + [dict(message) for message in run.input.messages], snapshot.bundle.input_schema + ) + if invalid is not None: + return invalid + steps: Final = await self.runs.steps(run.run_id) + if any(step.state is not RunState.COMPLETED for step in steps): + return ExecutionError( + "interrupted_execution", "A previous worker stopped before recording a completed node outcome" + ) + return await self._sequence(run, token, snapshot.bundle, steps, 0) + + async def _sequence( + self, run: EngineRun, token: str, bundle: WorkflowBundle, steps: tuple[StepTrace, ...], index: int + ) -> JsonValue | ExecutionError: + if self._clock() >= run.input.deadline: + return ExecutionError("deadline_exceeded", "Workflow deadline expired") + if index == len(bundle.nodes): + return resolve_binding(bundle.output, run, steps) + node: Final = bundle.nodes[index] + recorded: Final = next((step for step in steps if step.node_id == node.id), None) + if recorded is not None: + return await self._sequence(run, token, bundle, steps, index + 1) + value: Final = ( + resolve_binding(node.messages, run, steps) + if isinstance(node, ChatCompletionNode) + else _bound_inputs(node.inputs, run, steps) + ) + if isinstance(value, ExecutionError): + return value + started: Final = self._clock() + step: Final = StepTrace( + step_id=str(uuid4()), + run_id=run.run_id, + node_id=node.id, + state=RunState.RUNNING, + input=value, + attempts=1, + deadline=min(run.input.deadline, started + timedelta(seconds=node.timeout_seconds)), + started_at=started, + call_id=workflow_call_id(run.run_id, node.id, 1) if isinstance(node, ChatCompletionNode) else None, + ) + completed: Final = await self._attempt(run, token, node, step) + if isinstance(completed, ExecutionError): + return completed + return await self._sequence(run, token, bundle, (*steps, completed), index + 1) + + async def _attempt( + self, run: EngineRun, token: str, node: WorkflowNode, step: StepTrace + ) -> StepTrace | ExecutionError: + if not await self.runs.checkpoint(run.run_id, token, step): + return ExecutionError("lease_lost", "Execution ownership expired before dispatch") + result: Final = await self._recorded_node_result(run, token, node, step) + receipt: Final = result.receipt if isinstance(result, ExecutionError) else result + finished: Final = step.model_copy( + update={ + **( + { + "state": RunState.FAILED, + "finished_at": self._clock(), + "error": result.message, + } + if isinstance(result, ExecutionError) + else { + "state": RunState.COMPLETED, + "finished_at": self._clock(), + "output": result.output, + } + ), + "prompt_tokens": receipt.prompt_tokens if receipt is not None else 0, + "completion_tokens": receipt.completion_tokens if receipt is not None else 0, + "response_cost": receipt.response_cost if receipt is not None else None, + "request_id": receipt.request_id if receipt is not None else None, + } + ) + if not await self.runs.checkpoint(run.run_id, token, finished): + return ExecutionError("lease_lost", "Execution ownership expired before the checkpoint") + if not isinstance(result, ExecutionError): + return finished + retryable: Final = result.code in {"provider_error", "rate_limit", "runtime_error", "worker_crashed"} + remaining_after_backoff: Final = (step.deadline - self._clock()).total_seconds() - node.retry.backoff_seconds + if not retryable or step.attempts >= node.retry.max_attempts or remaining_after_backoff <= 0: + return result + await asyncio.sleep(node.retry.backoff_seconds) + return await self._attempt( + run, + token, + node, + step.model_copy( + update={ + "attempts": step.attempts + 1, + "call_id": ( + workflow_call_id(run.run_id, node.id, step.attempts + 1) + if isinstance(node, ChatCompletionNode) + else None + ), + } + ), + ) + + async def _recorded_node_result( + self, run: EngineRun, token: str, node: WorkflowNode, step: StepTrace + ) -> ChildCallResult | ExecutionError: + try: + return await self._timed_node_result(run, node, step) + except asyncio.CancelledError: + current: Final = await self.runs.get(run.run_id) + await self.runs.checkpoint( + run.run_id, + token, + step.model_copy( + update={ + "state": ( + RunState.CANCELLED if current is not None and current.cancel_requested else RunState.FAILED + ), + "finished_at": self._clock(), + "error": "Node execution interrupted", + } + ), + ) + raise + except Exception: # noqa: BLE001 # convert unexpected node failures into a sanitized terminal checkpoint + return ExecutionError("execution_failed", "The node could not complete") + + async def _timed_node_result( + self, run: EngineRun, node: WorkflowNode, step: StepTrace + ) -> ChildCallResult | ExecutionError: + if isinstance(node, ChatCompletionNode): + return await self._node_result(run, node, step) + try: + remaining: Final = max(0, (step.deadline - self._clock()).total_seconds()) + return await asyncio.wait_for(self._node_result(run, node, step), remaining) + except (TimeoutError, asyncio.TimeoutError): + return ExecutionError("deadline_exceeded", "Node deadline expired") + + async def _node_result( + self, run: EngineRun, node: WorkflowNode, step: StepTrace + ) -> ChildCallResult | ExecutionError: + if isinstance(node, ChatCompletionNode): + return await self._execute_chat( + key_hash=run.created_by, + team_id=run.team_id, + workflow_id=run.workflow_id, + node=node, + messages=step.input, + run_id=run.run_id, + session_id=run.session_id, + attempt=step.attempts, + deadline=step.deadline, + end_user_id=run.input.end_user_id, + ) + compiled: Final = compile_program(node.program, self._monty.limits) + if isinstance(compiled, MontyFailure): + return ExecutionError("invalid_program", compiled.message) + async with self._monty_start: + availability: Final = await self._monty.start() + if availability is not None: + return ExecutionError("runtime_unavailable", "The Monty runtime is unavailable") + outcome: Final = await self._monty.run_program( + compiled, + step.input, + deadline_monotonic=time.monotonic() + max(0, (step.deadline - self._clock()).total_seconds()), + ) + if isinstance(outcome, MontyFailure): + return ExecutionError(outcome.kind.value.lower(), outcome.message) + return ChildCallResult(outcome.output, 0, 0, None, None) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index d010455faa3..22d525faad5 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -111,6 +111,20 @@ def _delivery_id(row: Mapping[str, object]) -> str: class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]): table_name = "litellm_spendlogs" + async def request_ids_for_call( + self, *, call_id: str, key_hash: str, team_id: str, session_id: str + ) -> tuple[str, ...]: + rows: Final = await self.table.find_many( + where={ + "litellm_call_id": call_id, + "api_key": key_hash, + "team_id": team_id, + "session_id": session_id, + }, + take=2, + ) + return tuple(row.request_id for row in rows) + async def insert_with_deliveries( self, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool = True ) -> int: diff --git a/litellm/repositories/workflow_definition_repository.py b/litellm/repositories/workflow_definition_repository.py new file mode 100644 index 00000000000..e4bbd9c884a --- /dev/null +++ b/litellm/repositories/workflow_definition_repository.py @@ -0,0 +1,272 @@ +import hashlib +import json +from contextlib import AbstractAsyncContextManager +from datetime import datetime +from typing import Final, Protocol +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict + +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.workflow_definition import ( + DefinitionError, + DefinitionPage, + WorkflowBundle, + WorkflowDefinition, + WorkflowDraft, + WorkflowSnapshot, +) +from litellm.repositories.prisma_protocols import TableActions + + +class DefinitionTransaction(Protocol): + @property + def litellm_workflowdefinition(self) -> TableActions[object]: ... + + @property + def litellm_workflowsnapshot(self) -> TableActions[object]: ... + + @property + def litellm_teamtable(self) -> TableActions[object]: ... + + async def query_raw(self, query: str, *args: object) -> object: ... + + +class DefinitionDatabase(DefinitionTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[DefinitionTransaction]: ... + + +class _DefinitionRow(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + workflow_id: str + name: str + team_id: str + etag: str + digest: str + bundle_json: str + contributors: tuple[str, ...] + published_snapshot_id: str | None + published_etag: str | None + + def definition(self) -> WorkflowDefinition: + return WorkflowDefinition( + draft=WorkflowDraft( + workflow_id=self.workflow_id, + name=self.name, + team_id=self.team_id, + etag=self.etag, + digest=self.digest, + bundle=WorkflowBundle.model_validate_json(self.bundle_json), + published_snapshot_id=self.published_snapshot_id, + ), + contributors=self.contributors, + ) + + +class _SnapshotRow(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + snapshot_id: str + workflow_id: str + team_id: str + name: str + digest: str + bundle_json: str + contributors: tuple[str, ...] + published_by: str | None + created_at: datetime + + def snapshot(self) -> WorkflowSnapshot: + return WorkflowSnapshot( + snapshot_id=self.snapshot_id, + workflow_id=self.workflow_id, + team_id=self.team_id, + name=self.name, + digest=self.digest, + bundle=WorkflowBundle.model_validate_json(self.bundle_json), + contributors=self.contributors, + published_by=self.published_by, + created_at=self.created_at, + ) + + +def bundle_digest(bundle: WorkflowBundle) -> str: + return hashlib.sha256( + json.dumps( + bundle.model_dump(mode="json", by_alias=True, exclude_unset=True), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode() + ).hexdigest() + + +def _bundle_json(bundle: WorkflowBundle) -> str: + return bundle.model_dump_json(by_alias=True, exclude_unset=True) + + +def _etag() -> str: + return f'"{uuid4()}"' + + +async def _locked_definition( + transaction: DefinitionTransaction, workflow_id: str, team_id: str, etag: str +) -> _DefinitionRow | DefinitionError: + await transaction.query_raw( + 'SELECT workflow_id FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND team_id = $2 FOR UPDATE', + workflow_id, + team_id, + ) + stored: Final = await transaction.litellm_workflowdefinition.find_first( + where={"workflow_id": workflow_id, "team_id": team_id} + ) + if stored is None: + return DefinitionError("not_found", "Workflow not found") + row: Final = _DefinitionRow.model_validate(stored) + if row.etag != etag: + return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match") + return row + + +async def _snapshot( + transaction: DefinitionTransaction, row: _DefinitionRow, published_by: str | None +) -> WorkflowSnapshot: + stored: Final = await transaction.litellm_workflowsnapshot.create( + data={ + "snapshot_id": str(uuid4()), + "workflow_id": row.workflow_id, + "team_id": row.team_id, + "name": row.name, + "digest": row.digest, + "bundle_json": row.bundle_json, + "contributors": list(row.contributors), + "published_by": published_by, + } + ) + return _SnapshotRow.model_validate(stored).snapshot() + + +class WorkflowDefinitionRepository: + def __init__(self, database: DefinitionDatabase) -> None: + self._database: Final = database + + async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None: + row: Final = await self._database.litellm_teamtable.find_unique(where={"team_id": team_id}) + return None if row is None else LiteLLM_TeamTable.model_validate(row, from_attributes=True) + + async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None: + row: Final = await self._database.litellm_workflowdefinition.find_first( + where={"workflow_id": workflow_id, **({"team_id": team_id} if team_id is not None else {})} + ) + return None if row is None else _DefinitionRow.model_validate(row).definition() + + async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage: + rows: Final = await self._database.litellm_workflowdefinition.find_many( + where={"team_id": team_id, **({"workflow_id": {"gt": cursor}} if cursor is not None else {})}, + order={"workflow_id": "asc"}, + take=limit + 1, + ) + drafts: Final = tuple(_DefinitionRow.model_validate(row).definition().draft for row in rows[:limit]) + return DefinitionPage(workflows=drafts, next_cursor=drafts[-1].workflow_id if len(rows) > limit else None) + + async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None: + row: Final = await self._database.litellm_workflowsnapshot.find_first( + where={"snapshot_id": snapshot_id, **({"team_id": team_id} if team_id is not None else {})} + ) + return None if row is None else _SnapshotRow.model_validate(row).snapshot() + + async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None: + row: Final = await self._database.litellm_workflowdefinition.find_first( + where={"name": name, "team_id": team_id} + ) + if row is None: + return None + definition: Final = _DefinitionRow.model_validate(row) + if definition.published_snapshot_id is None: + return None + return await self.get_snapshot(definition.published_snapshot_id, team_id) + + async def create(self, bundle: WorkflowBundle, team_id: str, contributor: str) -> WorkflowDraft | DefinitionError: + from prisma.errors import UniqueViolationError + + try: + stored: Final = await self._database.litellm_workflowdefinition.create( + data={ + "workflow_id": str(uuid4()), + "name": bundle.name, + "team_id": team_id, + "etag": _etag(), + "digest": bundle_digest(bundle), + "bundle_json": _bundle_json(bundle), + "contributors": [contributor], + } + ) + return _DefinitionRow.model_validate(stored).definition().draft + except UniqueViolationError: + return DefinitionError("conflict", "A workflow with this name already exists in the team", "bundle.name") + + async def replace( + self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str + ) -> WorkflowDraft | DefinitionError: + async with self._database.tx() as transaction: + row: Final = await _locked_definition(transaction, workflow_id, team_id, etag) + if isinstance(row, DefinitionError): + return row + if row.name != bundle.name: + return DefinitionError("conflict", "A workflow name cannot change", "bundle.name") + contributors: Final = ( + (contributor,) + if row.etag == row.published_etag + else tuple(sorted(frozenset(row.contributors) | {contributor})) + ) + new_etag: Final = _etag() + digest: Final = bundle_digest(bundle) + await transaction.litellm_workflowdefinition.update( + where={"workflow_id": workflow_id}, + data={ + "etag": new_etag, + "digest": digest, + "bundle_json": _bundle_json(bundle), + "contributors": list(contributors), + }, + ) + return WorkflowDraft( + name=row.name, + workflow_id=workflow_id, + team_id=team_id, + etag=new_etag, + digest=digest, + bundle=bundle, + published_snapshot_id=row.published_snapshot_id, + ) + + async def publish( + self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str + ) -> WorkflowSnapshot | DefinitionError: + async with self._database.tx() as transaction: + row: Final = await _locked_definition(transaction, workflow_id, team_id, etag) + if isinstance(row, DefinitionError): + return row + if row.digest != digest: + return DefinitionError("stale_draft", "The reviewed digest does not match the draft", "digest") + if publisher in row.contributors: + return DefinitionError("forbidden", "Draft contributors cannot publish their own changes") + if row.published_etag == etag and row.published_snapshot_id is not None: + existing: Final = await transaction.litellm_workflowsnapshot.find_unique( + where={"snapshot_id": row.published_snapshot_id} + ) + return _SnapshotRow.model_validate(existing).snapshot() + snapshot: Final = await _snapshot(transaction, row, publisher) + await transaction.litellm_workflowdefinition.update( + where={"workflow_id": workflow_id}, + data={"published_snapshot_id": snapshot.snapshot_id, "published_etag": etag}, + ) + return snapshot + + async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: + async with self._database.tx() as transaction: + row: Final = await _locked_definition(transaction, workflow_id, team_id, etag) + if isinstance(row, DefinitionError): + return row + return await _snapshot(transaction, row, None) diff --git a/litellm/repositories/workflow_execution_repository.py b/litellm/repositories/workflow_execution_repository.py new file mode 100644 index 00000000000..257df3c1603 --- /dev/null +++ b/litellm/repositories/workflow_execution_repository.py @@ -0,0 +1,314 @@ +import base64 +import json +from contextlib import AbstractAsyncContextManager +from datetime import datetime +from itertools import chain +from typing import Final, Protocol +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError + +from litellm.models.workflow_execution import ( + EngineRun, + ExecutionError, + RunAdmission, + RunInput, + RunPage, + RunState, + RunView, + StepTrace, +) + + +class ExecutionQueries(Protocol): + async def query_raw(self, query: str, *args: object) -> object: ... + + +class ExecutionDatabase(ExecutionQueries, Protocol): + def tx(self) -> AbstractAsyncContextManager[ExecutionQueries]: ... + + +class _Cursor(BaseModel): + model_config = ConfigDict(frozen=True) + + created_at: datetime + run_id: str + + +class _TraceRow(BaseModel): + model_config = ConfigDict(frozen=True) + + trace: StepTrace + + +_RUNS: Final = TypeAdapter(tuple[EngineRun, ...]) +_VIEWS: Final = TypeAdapter(tuple[RunView, ...]) +_TRACES: Final = TypeAdapter(tuple[_TraceRow, ...]) +_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...]) + + +def _cursor(value: str | None) -> _Cursor | ExecutionError | None: + if value is None: + return None + try: + return _Cursor.model_validate_json(base64.urlsafe_b64decode(value.encode("ascii"))) + except (ValueError, UnicodeError, ValidationError): + return ExecutionError("invalid_cursor", "Invalid run cursor") + + +def _encoded_cursor(run: RunView) -> str: + return base64.urlsafe_b64encode( + _Cursor(created_at=run.created_at, run_id=run.run_id).model_dump_json().encode() + ).decode("ascii") + + +async def _first_run(database: ExecutionQueries, query: str, *args: object) -> EngineRun | None: + rows: Final = _RUNS.validate_python(await database.query_raw(query, *args)) + return rows[0] if rows else None + + +def _aggregate_attempts(attempts: tuple[StepTrace, ...]) -> StepTrace: + latest: Final = max(attempts, key=lambda step: step.attempts) + costs: Final = tuple(step.response_cost for step in attempts if step.response_cost is not None) + ids: Final = chain.from_iterable( + step.request_ids or ((step.request_id,) if step.request_id is not None else ()) for step in attempts + ) + return latest.model_copy( + update={ + "prompt_tokens": sum(step.prompt_tokens for step in attempts), + "completion_tokens": sum(step.completion_tokens for step in attempts), + "response_cost": sum(costs) if costs else None, + "request_ids": tuple(dict.fromkeys(ids)), + } + ) + + +class WorkflowExecutionRepository: + def __init__(self, database: ExecutionDatabase) -> None: + self._database: Final = database + + async def find_idempotent( + self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str + ) -> EngineRun | None: + return await _first_run( + self._database, + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND workflow_id = $1 ' + "AND team_id = $2 AND created_by = $3 AND idempotency_key = $4", + workflow_id, + team_id, + key_hash, + idempotency_key, + ) + + async def admit( + self, + *, + workflow_id: str, + snapshot_id: str, + team_id: str, + key_hash: str, + input: RunInput, + digest: str, + idempotency_key: str | None, + ) -> RunAdmission | ExecutionError: + async with self._database.tx() as transaction: + await transaction.query_raw( + "SELECT true AS acquired FROM pg_advisory_xact_lock(hashtextextended($1, 0))", team_id + ) + existing: Final = await _first_run( + transaction, + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 ' + "AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text", + key_hash, + team_id, + workflow_id, + idempotency_key, + ) + if existing is not None: + return ( + RunAdmission(existing, False) + if existing.request_digest == digest + else ExecutionError("idempotency_conflict", "Idempotency key was used with different input") + ) + capacity: Final = _ROWS.validate_python( + await transaction.query_raw( + 'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND team_id = $1 ' + "AND status IN ('pending', 'running') LIMIT 64", + team_id, + ) + ) + if len(capacity) >= 64: + return ExecutionError("capacity_exceeded", "The team's workflow execution queue is full") + created: Final = await _first_run( + transaction, + 'INSERT INTO "LiteLLM_WorkflowRun" ' + "(run_id, session_id, workflow_type, status, created_by, created_at, updated_at, " + "workflow_id, snapshot_id, team_id, engine_owned, input, request_digest, idempotency_key) " + "VALUES ($1, $2, 'liteflow', 'pending', $3, timezone('UTC', NOW()), timezone('UTC', NOW()), $4, $5, $6, true, $7::jsonb, $8, $9) " + "RETURNING *", + str(uuid4()), + str(uuid4()), + key_hash, + workflow_id, + snapshot_id, + team_id, + input.model_dump_json(), + digest, + idempotency_key, + ) + if created is None: + return ExecutionError("storage_error", "Workflow admission did not create a run") + return RunAdmission(created, True) + + async def get(self, run_id: str) -> EngineRun | None: + return await _first_run( + self._database, + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true', + run_id, + ) + + async def recover(self, limit: int) -> tuple[EngineRun, ...]: + return _RUNS.validate_python( + await self._database.query_raw( + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND ' + "(status = 'pending' OR (status = 'running' AND lease_expires_at < timezone('UTC', NOW()))) " + "ORDER BY created_at, run_id LIMIT $1", + limit, + ) + ) + + async def claim(self, run_id: str, token: str, lease_expires_at: datetime) -> EngineRun | None: + return await _first_run( + self._database, + "UPDATE \"LiteLLM_WorkflowRun\" SET status = 'running', lease_token = $2, " + "lease_expires_at = timezone('UTC', $3::text::timestamptz), updated_at = timezone('UTC', NOW()) " + "WHERE run_id = $1 AND engine_owned = true AND " + "(status = 'pending' OR (status = 'running' AND lease_expires_at < timezone('UTC', NOW()))) RETURNING *", + run_id, + token, + lease_expires_at.isoformat(), + ) + + async def renew(self, run_id: str, token: str, lease_expires_at: datetime) -> bool: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + "UPDATE \"LiteLLM_WorkflowRun\" SET lease_expires_at = timezone('UTC', $3::text::timestamptz) " + "WHERE run_id = $1 AND lease_token = $2 AND status = 'running' " + "AND lease_expires_at > timezone('UTC', NOW()) RETURNING run_id", + run_id, + token, + lease_expires_at.isoformat(), + ) + ) + return bool(rows) + + async def checkpoint(self, run_id: str, token: str, step: StepTrace) -> bool: + async with self._database.tx() as transaction: + owned: Final = _ROWS.validate_python( + await transaction.query_raw( + 'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND lease_token = $2 ' + "AND status = 'running' AND lease_expires_at > timezone('UTC', NOW()) FOR UPDATE", + run_id, + token, + ) + ) + if not owned: + return False + await transaction.query_raw( + 'INSERT INTO "LiteLLM_WorkflowEvent" ' + "(event_id, run_id, event_type, step_name, sequence_number, data, created_at) " + "SELECT $1, $2, 'liteflow.step', $3, COALESCE(MAX(sequence_number) + 1, 0), $4::jsonb, timezone('UTC', NOW()) " + 'FROM "LiteLLM_WorkflowEvent" WHERE run_id = $2 RETURNING event_id', + str(uuid4()), + run_id, + step.node_id, + json.dumps({"trace": step.model_dump(mode="json")}, allow_nan=False), + ) + return True + + async def steps(self, run_id: str) -> tuple[StepTrace, ...]: + rows: Final = _TRACES.validate_python( + await self._database.query_raw( + "SELECT trace FROM (" + "SELECT CASE WHEN receipts.ids IS NULL THEN latest.trace ELSE latest.trace || " + "jsonb_build_object('request_id', CASE WHEN cardinality(receipts.ids) = 1 THEN receipts.ids[1] END, " + "'request_ids', to_jsonb(receipts.ids), 'response_cost', receipts.cost, " + "'prompt_tokens', receipts.prompt_tokens, 'completion_tokens', receipts.completion_tokens) END AS trace, " + "latest.step_name FROM (" + "SELECT DISTINCT ON (step_name, data->'trace'->>'attempts') data->'trace' AS trace, step_name " + "FROM \"LiteLLM_WorkflowEvent\" WHERE run_id = $1 AND event_type = 'liteflow.step' " + "ORDER BY step_name, data->'trace'->>'attempts', sequence_number DESC" + ') latest JOIN "LiteLLM_WorkflowRun" run ON run.run_id = $1 AND run.engine_owned = true ' + "LEFT JOIN LATERAL (" + "SELECT array_agg(log.request_id ORDER BY log.request_id) AS ids, SUM(log.spend) AS cost, " + "SUM(log.prompt_tokens) AS prompt_tokens, SUM(log.completion_tokens) AS completion_tokens " + "FROM \"LiteLLM_SpendLogs\" log WHERE log.litellm_call_id = latest.trace->>'call_id' " + "AND log.api_key = run.created_by AND log.team_id = run.team_id AND log.session_id = run.session_id" + ") receipts ON true" + ") latest ORDER BY (trace->>'started_at')::timestamptz, step_name", + run_id, + ) + ) + node_ids: Final = dict.fromkeys(row.trace.node_id for row in rows) + return tuple( + _aggregate_attempts(tuple(row.trace for row in rows if row.trace.node_id == node_id)) + for node_id in node_ids + ) + + async def finish(self, run_id: str, token: str, state: RunState, output: JsonValue, error: str | None) -> bool: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + 'UPDATE "LiteLLM_WorkflowRun" SET status = $3, output = $4::jsonb, ' + "metadata = jsonb_build_object('error', $5::text), lease_token = NULL, lease_expires_at = NULL, " + "updated_at = timezone('UTC', NOW()) WHERE run_id = $1 AND lease_token = $2 AND status = 'running' " + "AND lease_expires_at > timezone('UTC', NOW()) RETURNING run_id", + run_id, + token, + state.value, + json.dumps(output, allow_nan=False), + error, + ) + ) + return bool(rows) + + async def cancel(self, run_id: str) -> None: + await self._database.query_raw( + "UPDATE \"LiteLLM_WorkflowRun\" SET cancel_requested = true, updated_at = timezone('UTC', NOW()) " + "WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id", + run_id, + ) + + async def list_runs( + self, + workflow_id: str, + team_id: str, + key_hash: str | None, + state: RunState | None, + limit: int, + cursor: str | None, + ) -> RunPage | ExecutionError: + position: Final = _cursor(cursor) + if isinstance(position, ExecutionError): + return position + rows: Final = _VIEWS.validate_python( + await self._database.query_raw( + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true ' + "AND workflow_id = $1 AND team_id = $2 " + "AND ($3::text IS NULL OR created_by = $3::text) AND ($4::text IS NULL OR status = $4::text) " + "AND ($5::text IS NULL OR (created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)) " + "ORDER BY created_at DESC, run_id DESC LIMIT $7", + workflow_id, + team_id, + key_hash, + state.value if state is not None else None, + position.created_at.isoformat() if position is not None else None, + position.run_id if position is not None else None, + limit + 1, + ) + ) + page: Final = rows[:limit] + return RunPage( + runs=page, + count=len(page), + next_cursor=_encoded_cursor(page[-1]) if len(rows) > limit else None, + ) diff --git a/pyproject.toml b/pyproject.toml index 47bed23c093..4a44a47433c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ Documentation = "https://docs.litellm.ai" # supply-chain-pinned to a vetted, aged release. [project.optional-dependencies] proxy = [ + "pydantic-monty==1.4.0", "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", "granian>=2.7.4,<3.0; python_version < '3.15'", @@ -257,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/schema.prisma b/schema.prisma index 818f2dc4858..c0dc2eb960b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1975,10 +1975,53 @@ model LiteLLM_ShadowEvalFunnel { // One instance of work being done. `status` is a materialized cache of the // latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowDefinition { + workflow_id String @id @default(uuid()) + team_id String + name String + etag String + digest String + bundle_json String + contributors String[] + published_snapshot_id String? + published_etag String? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + snapshots LiteLLM_WorkflowSnapshot[] + + @@unique([team_id, name]) + @@index([team_id, workflow_id]) +} + +model LiteLLM_WorkflowSnapshot { + snapshot_id String @id @default(uuid()) + workflow_id String + team_id String + name String + digest String + bundle_json String + contributors String[] + published_by String? + created_at DateTime @default(now()) + definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict) + + @@index([workflow_id, created_at, snapshot_id]) + @@index([team_id, snapshot_id]) +} + model LiteLLM_WorkflowRun { run_id String @id @default(uuid()) session_id String @unique @default(uuid()) workflow_type String + workflow_id String? + snapshot_id String? + team_id String? + engine_owned Boolean @default(false) + request_digest String? + idempotency_key String? + lease_token String? + lease_expires_at DateTime? + cancel_requested Boolean @default(false) status String @default("pending") created_by String? // user_id of the key that created this run; null = created by master key created_at DateTime @default(now()) @@ -1994,6 +2037,10 @@ model LiteLLM_WorkflowRun { @@index([session_id]) @@index([created_at]) @@index([created_by]) + @@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key") + @@index([workflow_id, team_id, created_at, run_id]) + @@index([engine_owned, status, lease_expires_at]) + @@index([team_id, engine_owned, status]) } // Append-only log of state transitions. Never mutate rows here. diff --git a/scripts/check_openapi_docs.py b/scripts/check_openapi_docs.py index 219ef912952..21e09d6bd72 100755 --- a/scripts/check_openapi_docs.py +++ b/scripts/check_openapi_docs.py @@ -7,18 +7,22 @@ count them per rule the same delta-vs-base way as the other lint gates. The rule OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix) OA002 no description (the handler has no docstring and the route no ``description=``) -OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``) +OA003 no 2xx success response has a schema (use ``response_model`` or a schema in ``responses``) """ from __future__ import annotations import json import sys -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from typing import Final, NamedTuple +from pydantic import TypeAdapter + GATED_METHODS: Final = frozenset({"post"}) SCHEMA_MEDIA_TYPES: Final = ("application/json", "application/sdp", "multipart/form-data") +SUCCESS_STATUS_CODES: Final = frozenset(str(code) for code in range(200, 300)) +_OBJECT: Final = TypeAdapter(Mapping[str, object]) class Violation(NamedTuple): @@ -28,7 +32,7 @@ class Violation(NamedTuple): def _as_mapping(node: object) -> Mapping[str, object]: - return node if isinstance(node, Mapping) else {} + return _OBJECT.validate_python(node) if isinstance(node, Mapping) else {} def _content_schema(container: object) -> object: @@ -45,7 +49,11 @@ def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]: description: Final = operation.get("description") if not isinstance(description, str) or not description.strip(): yield "OA002" - if not _content_schema(_as_mapping(operation.get("responses")).get("200")): + if not any( + _content_schema(response) + for status, response in _as_mapping(operation.get("responses")).items() + if status in SUCCESS_STATUS_CODES + ): yield "OA003" @@ -66,10 +74,10 @@ def proxy_spec() -> Mapping[str, object]: return _as_mapping(app.openapi()) -def main(argv: list[str] = sys.argv[1:]) -> int: +def main(argv: Sequence[str] = sys.argv[1:]) -> int: spec: Final = json.loads(open(argv[0]).read()) if argv else proxy_spec() for violation in violations(spec): - print(f"{violation.method} {violation.path} {violation.code}") + sys.stdout.write(f"{violation.method} {violation.path} {violation.code}\n") return 0 diff --git a/tests/integration/database/test_workflow_execution.py b/tests/integration/database/test_workflow_execution.py new file mode 100644 index 00000000000..0295a5659fc --- /dev/null +++ b/tests/integration/database/test_workflow_execution.py @@ -0,0 +1,149 @@ +import asyncio +import subprocess +import sys +from collections.abc import AsyncIterator, Iterator +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final, cast + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma +from prisma.types import DatasourceOverride +from psycopg import sql +from pydantic import TypeAdapter + +from litellm.models.workflow_execution import RunAdmission, RunInput, RunPage, RunState, StepTrace +from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] + + +@pytest.fixture(scope="module") +def schema_sql() -> bytes: + return subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), "--script"], + check=True, capture_output=True, cwd=_ROOT, + ).stdout + + +@pytest.fixture +def database(schema_sql: bytes) -> Iterator[str]: + with scratch_database("workflow_execution") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(schema_sql) + connection.execute( + sql.SQL("ALTER DATABASE {} SET timezone TO 'Asia/Singapore'").format(sql.Identifier(connection.info.dbname)) + ) + yield url + + +@pytest_asyncio.fixture(loop_scope="function") +async def ledger(database: str) -> AsyncIterator[Prisma]: + client: Final = Prisma(datasource=DatasourceOverride(url=database)) + await client.connect() + try: + assert TypeAdapter(tuple[dict[str, str], ...]).validate_python(await client.query_raw("SHOW TIME ZONE")) == ( + {"TimeZone": "Asia/Singapore"}, + ) + yield client + finally: + await client.disconnect() + + +def _repository(ledger: Prisma) -> WorkflowExecutionRepository: + return WorkflowExecutionRepository( + cast(ExecutionDatabase, ledger) # cast-ok: generated Prisma transaction delegates satisfy the protocol + ) + + +def _input(now: datetime) -> RunInput: + return RunInput(messages=({"role": "user", "content": "plan"},), deadline=now + timedelta(minutes=5), draft=True) + + +async def _admit( + repository: WorkflowExecutionRepository, data: RunInput, key: str | None, + *, snapshot: str = "snapshot", owner: str = "owner", +) -> RunAdmission: + result: Final = await repository.admit( + workflow_id="workflow", snapshot_id=snapshot, team_id="team", key_hash=owner, + input=data, digest="same-input", idempotency_key=key, + ) + assert isinstance(result, RunAdmission) + return result + + +@pytest.mark.asyncio(loop_scope="function") +async def test_concurrent_admission_pins_snapshot_and_execution_timestamps_are_utc(ledger: Prisma) -> None: + repository: Final = _repository(ledger) + before: Final = datetime.now(timezone.utc).replace(microsecond=0) + data: Final = _input(before) + first, second = await asyncio.gather( + _admit(repository, data, "idempotency-key", snapshot="first"), + _admit(repository, data, "idempotency-key", snapshot="edited"), + ) + assert sorted((first.created, second.created)) == [False, True] + assert first.run == second.run + assert await repository.find_idempotent("workflow", "team", "owner", "idempotency-key") == first.run + assert await ledger.litellm_workflowrun.count() == 1 + run_id: Final = first.run.run_id + claimed: Final = await repository.claim(run_id, "worker", data.deadline) + assert claimed is not None and claimed.status is RunState.RUNNING and claimed.lease_expires_at == data.deadline + assert await repository.claim(run_id, "other-worker", data.deadline) is None + extended: Final = data.deadline + timedelta(minutes=1) + assert await repository.renew(run_id, "worker", extended) is True + renewed: Final = await repository.get(run_id) + assert renewed is not None and renewed.lease_expires_at == extended + step: Final = StepTrace( + step_id="step", run_id=run_id, node_id="node", state=RunState.COMPLETED, input={"n": 1}, output={"n": 2}, + attempts=1, deadline=data.deadline, started_at=before, finished_at=before, + ) + assert await repository.checkpoint(run_id, "other-worker", step) is False + assert await repository.checkpoint(run_id, "worker", step) is True + assert await repository.steps(run_id) == (step,) + assert await repository.finish(run_id, "worker", RunState.COMPLETED, {"n": 2}, None) is True + finished: Final = await repository.get(run_id) + assert finished is not None and finished.status is RunState.COMPLETED + assert (finished.output, finished.metadata, finished.lease_token, finished.lease_expires_at) == ( + {"n": 2}, {"error": None}, None, None, + ) + stored: Final = await ledger.litellm_workflowrun.find_unique(where={"run_id": run_id}) + events: Final = await ledger.litellm_workflowevent.find_many(where={"run_id": run_id}) + assert stored is not None and stored.created_at == finished.created_at and stored.updated_at == finished.updated_at + assert len(events) == 1 + for timestamp in (first.run.created_at, claimed.updated_at, finished.updated_at, events[0].created_at): + assert before <= timestamp <= datetime.now(timezone.utc) + + +@pytest.mark.asyncio(loop_scope="function") +async def test_nullable_admission_and_history_filters_round_trip_through_prisma(ledger: Prisma) -> None: + repository: Final = _repository(ledger) + data: Final = _input(datetime.now(timezone.utc).replace(microsecond=0)) + first: Final = await _admit(repository, data, None) + foreign: Final = await _admit(repository, data, None, owner="other") + second: Final = await _admit(repository, data, None) + assert len({first.run.run_id, foreign.run.run_id, second.run.run_id}) == 3 + await ledger.litellm_workflowrun.update_many(where={}, data={"created_at": first.run.created_at}) + ordered: Final = tuple(sorted((first.run.run_id, foreign.run.run_id, second.run.run_id), reverse=True)) + owned: Final = tuple(run_id for run_id in ordered if run_id != foreign.run.run_id) + page: Final = await repository.list_runs("workflow", "team", None, None, 2, None) + assert isinstance(page, RunPage) and page.count == 2 and page.next_cursor is not None + assert tuple(run.run_id for run in page.runs) == ordered[:2] + tail: Final = await repository.list_runs("workflow", "team", None, None, 2, page.next_cursor) + assert isinstance(tail, RunPage) and tail.count == 1 and tail.next_cursor is None + assert tuple(run.run_id for run in tail.runs) == ordered[2:] + scoped: Final = await repository.list_runs("workflow", "team", "owner", RunState.PENDING, 1, None) + assert isinstance(scoped, RunPage) and scoped.next_cursor is not None + assert tuple(run.run_id for run in scoped.runs) == owned[:1] + continuation: Final = await repository.list_runs("workflow", "team", "owner", RunState.PENDING, 1, scoped.next_cursor) + assert isinstance(continuation, RunPage) and continuation.count == 1 and continuation.next_cursor is None + assert tuple(run.run_id for run in continuation.runs) == owned[1:] + assert await repository.claim(second.run.run_id, "worker", data.deadline) is not None + assert await repository.finish(second.run.run_id, "worker", RunState.FAILED, None, "execution_failed") is True + failed: Final = await repository.get(second.run.run_id) + assert failed is not None and failed.output is None and failed.metadata == {"error": "execution_failed"} + filtered: Final = await repository.list_runs("workflow", "team", "owner", RunState.FAILED, 2, None) + assert isinstance(filtered, RunPage) and filtered.count == 1 and filtered.next_cursor is None + assert tuple(run.run_id for run in filtered.runs) == (second.run.run_id,) diff --git a/tests/integration/security/test_monty_runtime.py b/tests/integration/security/test_monty_runtime.py new file mode 100644 index 00000000000..ac1b0ab5052 --- /dev/null +++ b/tests/integration/security/test_monty_runtime.py @@ -0,0 +1,624 @@ +import asyncio +import hashlib +import os +import sys +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._compat import timeout +from litellm.proxy.workflows.monty.compiler import compile_program +from litellm.proxy.workflows.monty.models import ( + CompiledProgram, + JsonSchema, + MontyErrorKind, + MontyFailure, + MontyLimits, + MontySuccess, + PythonProgram, +) +from litellm.proxy.workflows.monty.runtime import BinaryArtifact, MontyRuntime, installed_binary + +pytest.importorskip("pydantic_monty") + + +def _compile( + source: str, + output: JsonSchema | None = None, + argument: JsonSchema | None = None, + definitions: JsonSchema | None = None, +) -> CompiledProgram: + program: Final = PythonProgram( + contract_version="1", + source=source, + function_name="transform", + parameter_names=("value",), + input_schema={ + "type": "object", + "properties": {"value": argument if argument is not None else {}}, + "required": ["value"], + "additionalProperties": False, + "$defs": definitions or {}, + }, + output_schema={**(output or {}), "$defs": definitions or {}}, + source_sha256=hashlib.sha256(source.encode()).hexdigest(), + ) + result: Final = compile_program(program, MontyLimits()) + assert isinstance(result, CompiledProgram), result + return result + + +@pytest_asyncio.fixture(loop_scope="function") +async def runtime() -> AsyncIterator[MontyRuntime]: + instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1, max_queued_waiters=1)) + assert await instance.start() is None + try: + yield instance + finally: + await instance.close() + + +async def _run(runtime: MontyRuntime, compiled: CompiledProgram, value: object) -> MontySuccess | MontyFailure: + return await runtime.run_program(compiled, {"value": value}, deadline_monotonic=time.monotonic() + 10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value", + [ + {"__class__": "os.system"}, + {"type": "function", "callable": "os.system"}, + "Bearer sk-test-fake", + "${OPENAI_API_KEY}", + "https://example.invalid/action", + [], + None, + ], +) +async def test_model_output_remains_inert_json(runtime: MontyRuntime, value: object) -> None: + result: Final = await _run(runtime, _compile("def transform(value): return value"), value) + assert isinstance(result, MontySuccess), result + assert result.output == value + + +@pytest.mark.asyncio +async def test_typed_json_transform_and_escaped_stdout(runtime: MontyRuntime) -> None: + compiled: Final = _compile( + "def transform(value):\n print('\\x1b[31m')\n return {'n': value['n'] + 1}", + {"type": "object", "properties": {"n": {"type": "integer"}}, "required": ["n"]}, + ) + result: Final = await _run(runtime, compiled, {"n": 2}) + assert isinstance(result, MontySuccess), result + assert result.output == {"n": 3} + assert result.stdout == "\\x1b[31m\n" + + +_RECORDS: Final[JsonSchema] = { + "Expense": { + "type": "object", + "x-liteflow-record": True, + "additionalProperties": False, + "properties": { + "amount_cents": {"type": "integer", "minimum": 0}, + "currency": {"type": "string"}, + "receipt_present": {"type": "boolean"}, + }, + "required": ["amount_cents", "currency", "receipt_present"], + }, +} +_EXPENSE: Final[JsonSchema] = {"$ref": "#/$defs/Expense"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("receipt,amount,review", [(True, 100, False), (False, 100, True), (True, 10001, True)]) +async def test_expense_normalize_condition_and_both_branches( + runtime: MontyRuntime, receipt: bool, amount: int, review: bool +) -> None: + normalize: Final = _compile( + "def transform(value: Expense) -> Expense:\n return Expense(amount_cents=value.amount_cents, currency=value.currency.strip().upper(), receipt_present=value.receipt_present)", + _EXPENSE, + _EXPENSE, + _RECORDS, + ) + normalized: Final = await _run( + runtime, normalize, {"amount_cents": amount, "currency": " usd ", "receipt_present": receipt} + ) + assert isinstance(normalized, MontySuccess), normalized + assert normalized.output == {"amount_cents": amount, "currency": "USD", "receipt_present": receipt} + condition: Final = _compile( + "def transform(value: Expense) -> bool:\n return not value.receipt_present or value.currency != 'USD' or value.amount_cents > 10000", + {"type": "boolean"}, + _EXPENSE, + _RECORDS, + ) + assert ( + await runtime.evaluate_predicate( + condition, {"value": normalized.output}, deadline_monotonic=time.monotonic() + 10 + ) + is review + ) + status: Final = "review_required" if review else "ready_for_processing" + branch: Final = _compile( + f"def transform(value): return {{'status': {status!r}, 'amount_cents': value.amount_cents, 'currency': value.currency}}", + {"type": "object"}, + _EXPENSE, + _RECORDS, + ) + result: Final = await _run(runtime, branch, normalized.output) + assert isinstance(result, MontySuccess), result + assert result.output == {"status": status, "amount_cents": amount, "currency": "USD"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source", + [ + "def transform(value): return Expense(amount_cents=1, currency='USD', receipt_present=True)", + "def transform(value): return {'amount_cents':1, 'currency':'USD', 'receipt_present':True}", + ], +) +async def test_record_and_equal_dict_are_equivalent(runtime: MontyRuntime, source: str) -> None: + result: Final = await _run(runtime, _compile(source, _EXPENSE, definitions=_RECORDS), None) + assert isinstance(result, MontySuccess), result + assert result.output == {"amount_cents": 1, "currency": "USD", "receipt_present": True} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [None, {"amount_cents": 7, "currency": "USD", "receipt_present": True}]) +async def test_optional_record(runtime: MontyRuntime, value: object) -> None: + optional: Final[JsonSchema] = {"anyOf": [_EXPENSE, {"type": "null"}]} + result: Final = await _run( + runtime, + _compile( + "def transform(value): return None if value is None else value.amount_cents", + argument=optional, + definitions=_RECORDS, + ), + value, + ) + assert isinstance(result, MontySuccess), result + assert result.output == (None if value is None else 7) + + +@pytest.mark.asyncio +async def test_list_of_records(runtime: MontyRuntime) -> None: + result: Final = await _run( + runtime, + _compile( + "def transform(value): return [item.currency for item in value]", + argument={"type": "array", "items": _EXPENSE}, + definitions=_RECORDS, + ), + [{"amount_cents": 7, "currency": "USD", "receipt_present": True}], + ) + assert isinstance(result, MontySuccess), result + assert result.output == ["USD"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [1, "true", None]) +async def test_predicate_requires_exact_bool(runtime: MontyRuntime, value: object) -> None: + result: Final = await runtime.evaluate_predicate( + _compile("def transform(value): return value"), {"value": value}, deadline_monotonic=time.monotonic() + 10 + ) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.PREDICATE_NOT_BOOL + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source,schema,kind", + [ + ("def transform(value): return -1", {"type": "integer", "minimum": 0}, MontyErrorKind.OUTPUT_INVALID), + ("def transform(value): return 'wrong'", {"type": "integer"}, MontyErrorKind.OUTPUT_INVALID), + ("def helper(x): return x\ndef transform(value): return helper", {}, MontyErrorKind.OUTPUT_INVALID), + ("def transform(value): return 'x' * 1100000", {}, MontyErrorKind.OUTPUT_TOO_LARGE), + ("def transform(value):\n while True: pass", {}, MontyErrorKind.TIMEOUT), + ("def transform(value): return [0] * 1000000000", {}, MontyErrorKind.MEMORY_LIMIT), + ("def transform(value): return transform(value)", {}, MontyErrorKind.RECURSION_LIMIT), + ("def transform(value): print('x' * 100000)", {}, MontyErrorKind.STDOUT_LIMIT), + ("def transform(value): return 1 / 0", {}, MontyErrorKind.RUNTIME_ERROR), + ], +) +async def test_resource_and_output_failures( + runtime: MontyRuntime, source: str, schema: JsonSchema, kind: MontyErrorKind +) -> None: + result: Final = await _run(runtime, _compile(source, schema), None) + assert isinstance(result, MontyFailure), result + assert result.kind == kind, result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value,kind", + [ + ({1: "not JSON"}, MontyErrorKind.INPUT_INVALID), + (float("nan"), MontyErrorKind.INPUT_INVALID), + ("x" * 1100000, MontyErrorKind.INPUT_TOO_LARGE), + ], +) +async def test_input_failures(runtime: MontyRuntime, value: object, kind: MontyErrorKind) -> None: + result: Final = await _run(runtime, _compile("def transform(value): return value"), value) + assert isinstance(result, MontyFailure), result + assert result.kind == kind + + +@pytest.mark.asyncio +async def test_expired_deadline(runtime: MontyRuntime) -> None: + result: Final = await runtime.run_program( + _compile("def transform(value): return value"), {"value": 1}, deadline_monotonic=0 + ) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.TIMEOUT + + +@pytest.mark.asyncio +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)) + result: Final = await runtime.start() + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.RUNTIME_UNAVAILABLE + 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 = TypeAdapter(JsonSchema).validate_python(EmailAnalyses.model_json_schema()) + analysis_definitions: Final = TypeAdapter(dict[str, JsonSchema]).validate_python(analyses_schema["$defs"]) + message_contract: Final = TypeAdapter(JsonSchema).validate_python( + TypeAdapter(list[ChatCompletionMessageParam]).json_schema() + ) + message_definitions: Final = TypeAdapter(dict[str, JsonSchema]).validate_python(message_contract.get("$defs", {})) + messages_schema: Final = {key: value for key, value in message_contract.items() if key != "$defs"} + definitions: Final[JsonSchema] = { + **message_definitions, + "Finding": {**analysis_definitions["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 +@pytest.mark.parametrize("name", ["default", "pattern", "enum", "$ref", "$ref-like"]) +async def test_keyword_looking_properties_validate_as_data(runtime: MontyRuntime, name: str) -> None: + schema: Final[JsonSchema] = { + "type": "object", + "properties": {name: {"type": "string"}}, + "required": [name], + "additionalProperties": False, + } + compiled: Final = _compile("def transform(value): return value", schema, schema) + result: Final = await _run(runtime, compiled, {name: "data"}) + assert isinstance(result, MontySuccess), result + assert result.output == {name: "data"} + invalid: Final = await _run(runtime, compiled, {name: 1}) + assert isinstance(invalid, MontyFailure), invalid + assert invalid.kind == MontyErrorKind.INPUT_INVALID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("output", [False, True]) +async def test_bypassed_remote_schema_never_retrieves_from_host(runtime: MontyRuntime, output: bool) -> None: + requests: Final[asyncio.Queue[bytes]] = asyncio.Queue() + + async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + requests.put_nowait(await reader.readuntil(b"\r\n\r\n")) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}") + await writer.drain() + writer.close() + await writer.wait_closed() + + server: Final = await asyncio.start_server(respond, "127.0.0.1", 0) + async with server: + port: Final = TypeAdapter(int).validate_python(server.sockets[0].getsockname()[1], strict=True) + schema: Final[JsonSchema] = {"properties": {"default": {"$ref": f"http://127.0.0.1:{port}/schema"}}} + compiled: Final = _compile("def transform(value): return value") + bypassed: Final = ( + replace(compiled, output_schema=schema) + if output + else replace(compiled, input_schema={"properties": {"value": schema}}) + ) + result: Final = await _run(runtime, bypassed, {"default": 1}) + assert isinstance(result, MontyFailure), result + assert result.kind == (MontyErrorKind.OUTPUT_INVALID if output else MontyErrorKind.INPUT_INVALID) + assert requests.empty() + + +def _slow_schema(depth: int) -> JsonSchema: + definitions: Final[JsonSchema] = { + f"N{index}": {"anyOf": [{"$ref": f"#/$defs/N{index + 1}"}, {"$ref": f"#/$defs/N{index + 1}"}]} + for index in range(depth) + } + return {**definitions, f"N{depth}": {"type": "string"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("output", [False, True]) +async def test_expensive_allowed_validation_times_out_without_blocking_and_recovers( + runtime: MontyRuntime, output: bool +) -> None: + ticks: Final[asyncio.Queue[float]] = asyncio.Queue() + done: Final = asyncio.Event() + + async def heartbeat() -> None: + while not done.is_set(): + ticks.put_nowait(time.monotonic()) + await asyncio.sleep(0.005) + + schema: Final = _slow_schema(22) + ref: Final[JsonSchema] = {"$ref": "#/$defs/N0"} + compiled: Final = _compile( + "def transform(value): return value", ref if output else {}, {} if output else ref, schema + ) + pulse: Final = asyncio.create_task(heartbeat()) + await asyncio.sleep(0.01) + started: Final = time.monotonic() + try: + result: Final = await runtime.run_program(compiled, {"value": None}, deadline_monotonic=started + 0.15) + elapsed: Final = time.monotonic() - started + finally: + done.set() + await pulse + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.TIMEOUT + assert 0.1 <= elapsed < 0.5 # test-quality-ok: validation enforces the node deadline + observations: Final = tuple(ticks.get_nowait() for _ in range(ticks.qsize())) + (time.monotonic(),) + assert len(observations) > 5 + assert max(later - earlier for earlier, later in zip(observations, observations[1:])) < 0.05 + recovered: Final = await _run(runtime, _compile("def transform(value): return value"), "recovered") + assert isinstance(recovered, MontySuccess), recovered + assert recovered.output == "recovered" + + +def _validator_pids() -> frozenset[int]: + def is_validator(child: psutil.Process) -> bool: + try: + return any(argument.endswith("_validator_worker.py") for argument in child.cmdline()) + except (psutil.NoSuchProcess, psutil.ZombieProcess): + return False + + return frozenset(child.pid for child in psutil.Process().children() if is_validator(child)) + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform != "linux", reason="validator environment inspection requires Linux procfs") +async def test_validator_environment_is_empty_and_close_reaps_processes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITEFLOW_VALIDATOR_CANARY", "fake-value-no-credential") + previous: Final = _validator_pids() + instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1)) + assert await instance.start() is None + try: + created: Final = _validator_pids() - previous + assert len(created) == 1 + pid: Final = next(iter(created)) + assert b"LITEFLOW_VALIDATOR_CANARY" not in Path(f"/proc/{pid}/environ").read_bytes() + assert Path(f"/proc/{pid}/environ").read_bytes() == b"" + result: Final = await _run(instance, _compile("def transform(value): return value"), 1) + assert isinstance(result, MontySuccess), result + assert _validator_pids() - previous == created + finally: + await instance.close() + assert not (_validator_pids() - previous) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_validator_crash_and_cancellation_release_slot_and_recover(runtime: MontyRuntime, cancel: bool) -> None: + pid: Final = next(iter(_validator_pids())) + compiled: Final = _compile( + "def transform(value): return value", argument={"$ref": "#/$defs/N0"}, definitions=_slow_schema(22) + ) + task: Final = asyncio.create_task(_run(runtime, compiled, None)) + await asyncio.sleep(0.05) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + os.kill(pid, 9) # test-quality-ok: deliberately crash the owned validator to prove slot and subprocess recovery + result: Final = await task + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.WORKER_CRASHED + assert pid not in _validator_pids() + recovered: Final = await _run(runtime, _compile("def transform(value): return value"), "after failure") + assert isinstance(recovered, MontySuccess), recovered + assert recovered.output == "after failure" + + +@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]: + def is_worker(child: psutil.Process) -> bool: + try: + return child.name() == "monty" + except (psutil.NoSuchProcess, psutil.ZombieProcess): + return False + + return frozenset(child.pid for child in psutil.Process().children() if is_worker(child)) + + +async def _active_worker(previous: frozenset[int]) -> int: + async with timeout(5): + while True: + if current := _worker_pids() - previous: + return next(iter(current)) + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_cancellation_kills_worker_releases_slot_and_queue_is_bounded(runtime: MontyRuntime) -> None: + 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 + result: Final = await _run(runtime, _compile("def transform(value): return value"), 2) + assert isinstance(result, MontySuccess), result + assert result.output == 2 + assert not psutil.pid_exists(pid) + + +@pytest.mark.asyncio +async def test_killed_worker_fails_and_next_run_succeeds(runtime: MontyRuntime) -> None: + 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) + psutil.Process(pid).kill() # test-quality-ok: crash the owned Monty worker and prove recovery in the next run + result: Final = await task + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.WORKER_CRASHED + next_result: Final = await _run(runtime, _compile("def transform(value): return value"), 3) + assert isinstance(next_result, MontySuccess), next_result + assert next_result.output == 3 + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform != "linux", reason="worker environment and address-space limits require Linux procfs") +async def test_fresh_worker_has_no_parent_canary_and_has_address_space_limit( + runtime: MontyRuntime, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITEFLOW_ENV_CANARY", "fake-value-no-credential") + 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) + assert b"LITEFLOW_ENV_CANARY" not in Path(f"/proc/{pid}/environ").read_bytes() + async with timeout(5): + while str(runtime.limits.worker_address_space_bytes) not in Path(f"/proc/{pid}/limits").read_text(): + await asyncio.sleep(0.01) + assert str(runtime.limits.worker_address_space_bytes) in Path(f"/proc/{pid}/limits").read_text() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + second: Final = asyncio.create_task(_run(runtime, _compile("def transform(value):\n while True: pass"), None)) + next_pid: Final = await _active_worker(previous | {pid}) + assert next_pid != pid + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second diff --git a/tests/unit/models/test_workflow_definition.py b/tests/unit/models/test_workflow_definition.py new file mode 100644 index 00000000000..e26fc45be87 --- /dev/null +++ b/tests/unit/models/test_workflow_definition.py @@ -0,0 +1,41 @@ +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.models.workflow_definition import WorkflowBundle +from tests.unit.proxy.workflow_definition_rows import bundle + + +@pytest.mark.parametrize("node_type", ["parallel", "if", "shell", "http"]) +def test_v1_contract_rejects_nodes_outside_sequential_chat_and_python(node_type: str) -> None: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + unsupported: Final = {**original, "nodes": [{**original["nodes"][0], "type": node_type}]} + with pytest.raises(ValidationError) as rejected: + WorkflowBundle.model_validate(unsupported) + assert rejected.value.errors(include_input=False)[0]["type"] == "union_tag_invalid" + + +def test_binding_cannot_supply_both_literal_and_reference() -> None: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + ambiguous: Final = { + **original, + "nodes": [ + { + **original["nodes"][0], + "messages": {"from": "$input", "output": "messages", "value": [{"role": "user", "content": "hello"}]}, + } + ], + } + with pytest.raises(ValidationError) as rejected: + WorkflowBundle.model_validate(ambiguous) + assert { + error["loc"][-1] for error in rejected.value.errors(include_input=False) if "messages" in error["loc"] + } == {"from", "output", "value"} + + +def test_nonfinite_values_are_rejected_before_digest_serialization() -> None: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + with pytest.raises(ValidationError) as rejected: + WorkflowBundle.model_validate({**original, "input_schema": {"minimum": float("nan")}}) + assert rejected.value.errors(include_input=False)[0]["type"] == "finite_number" diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index c1011310af3..6f16b93d1a8 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -5130,3 +5130,69 @@ def test_member_can_reach_skill_history_visibility_checks(route: str) -> None: ) except HTTPException as exc: pytest.fail(f"History visibility could not run: {exc.detail}") + + +@pytest.mark.parametrize("route", [ + "/v1/workflows", + "/v1/workflows/definition-id", + "/v1/workflows/definition-id/validate", + "/v1/workflows/definition-id/publish", +]) +def test_definition_management_routes_reach_team_authorization(route: str) -> None: + token: Final = UserAPIKeyAuth( + user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER, + allowed_routes=["management_routes"], + ) + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={} + ) + assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True + + +@pytest.mark.parametrize("route", [ + "/v1/workflows/runs", + "/v1/workflows/runs/validate", + "/v1/workflows/runs/publish", + "/v1/workflows/definition-id/delete", + "/v1/workflows/definition-id/test", +]) +def test_definition_route_grants_do_not_grant_run_or_future_actions(route: str) -> None: + token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_definition_routes"]) + with pytest.raises(HTTPException) as rejected: + RouteChecks.is_virtual_key_allowed_to_call_route(route, token) + assert rejected.value.status_code == 403 + + +@pytest.mark.parametrize("route", [ + "/v1/workflows/definition-id/test", + "/v1/workflows/definition-id/runs", + "/v1/workflows/runs/run-id/steps", + "/v1/workflows/runs/run-id/definition", + "/v1/workflows/runs/run-id/cancel", +]) +def test_execution_routes_reach_their_scoped_authorization(route: str) -> None: + token: Final = UserAPIKeyAuth( + user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER, + allowed_routes=["management_routes"], + ) + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={} + ) + assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True + + +@pytest.mark.parametrize("route", [ + "/v1/workflows/runs", + "/v1/workflows/runs/run-id", + "/v1/workflows/runs/run-id/events", + "/v1/workflows/runs/run-id/messages", +]) +def test_execution_route_grants_preserve_legacy_tracking_admission(route: str) -> None: + token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_execution_routes"]) + with pytest.raises(HTTPException) as rejected: + RouteChecks.is_virtual_key_allowed_to_call_route(route, token) + assert rejected.value.status_code == 403 + assert RouteChecks.check_route_access(route, LiteLLMRoutes.self_managed_routes.value) is False + assert RouteChecks.check_route_access(route, LiteLLMRoutes.management_routes.value) is False diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 741669337cd..b4ea51a55f8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -75,7 +75,7 @@ from litellm.proxy.auth.auth_checks import ( jwt_key_mapping_cache_key, ) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model +from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model, safe_set_request_parsed_body from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, _check_team_member_model_budget_with_fallback, @@ -87,6 +87,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _reserve_budget_after_common_checks, _route_requires_auth_despite_public, _routing_selector_matches_claim, + authorize_internal_virtual_key, run_centralized_common_checks, _run_post_custom_auth_checks, user_api_key_auth_builder, @@ -107,6 +108,117 @@ class _RoutingRequest: self.state = SimpleNamespace() +async def _internal_key_cache( + monkeypatch: pytest.MonkeyPatch, *, models: tuple[str, ...] = (), max_budget: float | None = None, + key_hash: str = "b" * 64, +) -> UserApiKeyCache: + cache: Final = UserApiKeyCache() + await cache_key_object( + hashed_token=key_hash, + user_api_key_obj=UserAPIKeyAuth( + token=key_hash, team_id="original-team", models=list(models), spend=2.0, max_budget=max_budget, + ), + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + await cache.async_set_cache( + key="team_id:original-team", value=LiteLLM_TeamTableCachedObj(team_id="original-team") + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", object()) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {}) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + counters: Final = FakeSpendCounters(redis_default=2.0) + monkeypatch.setattr(litellm.proxy.proxy_server, "spend_counter_store", counters.store()) + monkeypatch.setattr(litellm, "max_budget", 0.0) + return cache + + +def _internal_management_request() -> Request: + return Request({ + "type": "http", "path": "/workflow/runs", "method": "POST", "headers": (), "query_string": b"", + }) + + +@pytest.mark.asyncio +async def test_internal_key_cannot_execute_for_a_different_team(monkeypatch: pytest.MonkeyPatch) -> None: + await _internal_key_cache(monkeypatch) + request: Final = _internal_management_request() + with pytest.raises(HTTPException, match="no longer belongs") as error: + await authorize_internal_virtual_key( + "b" * 64, request, {"model": "workflow-model"}, + route="/v1/chat/completions", expected_team_id="another-team", + ) + assert error.value.status_code == 403 + assert getattr(request.state, "budget_reservation", None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restriction", ["model", "budget"]) +async def test_internal_execution_route_enforces_inference_model_and_key_budget( + monkeypatch: pytest.MonkeyPatch, restriction: str, +) -> None: + await _internal_key_cache( + monkeypatch, + models=("allowed-model",) if restriction == "model" else (), + max_budget=1.0 if restriction == "budget" else None, + ) + with pytest.raises(litellm.BudgetExceededError if restriction == "budget" else ProxyException) as error: + await authorize_internal_virtual_key( + "b" * 64, _internal_management_request(), {"model": "workflow-model"}, + route="/v1/chat/completions", expected_team_id="original-team", + ) + if restriction == "budget": + assert isinstance(error.value, litellm.BudgetExceededError) + assert error.value.current_cost == 2.0 and error.value.max_budget == 1.0 + else: + assert isinstance(error.value, ProxyException) + assert error.value.type == ProxyErrorTypes.key_model_access_denied + + +@pytest.mark.asyncio +async def test_internal_execution_route_keeps_proxy_network_policy(monkeypatch: pytest.MonkeyPatch) -> None: + await _internal_key_cache(monkeypatch) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"allowed_ips": ["192.0.2.1"]}) + with pytest.raises(HTTPException, match="not allowed") as error: + await authorize_internal_virtual_key( + "b" * 64, _internal_management_request(), {"model": "workflow-model"}, + route="/v1/chat/completions", expected_team_id="original-team", + ) + assert error.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_workflow_wrapper_with_model_grant_passes_strict_pricing_but_unpriced_child_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + raw_key: Final = "sk-workflow-pricing-fixture" + key_hash: Final = proxy_utils.hash_token(raw_key) + await _internal_key_cache(monkeypatch, models=("workflow/example", "unpriced-child"), key_hash=key_hash) + router: Final = litellm.Router( + model_list=[{ + "model_name": "unpriced-child", + "litellm_params": {"model": "openai/workflow-unpriced-fixture", "api_key": "test-only"}, + }], + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + request: Final = Request({ + "type": "http", "path": "/v1/chat/completions", "method": "POST", + "headers": ((b"authorization", f"Bearer {raw_key}".encode()),), "query_string": b"", + }) + safe_set_request_parsed_body(request, {"model": "workflow/example", "messages": [{"role": "user", "content": "hello"}]}) + outer: Final = await user_api_key_auth(request=request, api_key=f"Bearer {raw_key}") + assert outer.team_id == "original-team" + with pytest.raises(ProxyException) as error: + await authorize_internal_virtual_key( + key_hash, _internal_management_request(), {"model": "unpriced-child"}, + route="/v1/chat/completions", expected_team_id="original-team", + ) + assert error.value.type == ProxyErrorTypes.model_cost_map_missing + + async def _real_model_budget_limiter( *, member_spend: Mapping[str, float] | None = None, diff --git a/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py new file mode 100644 index 00000000000..87a0cad7450 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py @@ -0,0 +1,222 @@ +from collections.abc import Mapping +from typing import Final + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from pydantic import TypeAdapter + +from litellm.models.workflow_definition import DefinitionPage, WorkflowDraft +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_workflow_definition_repository +from litellm.proxy.management_endpoints.workflow_definition_endpoints import router +from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository, bundle_digest +from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller + + +def _client(database: DefinitionRows | None, identity: UserAPIKeyAuth) -> TestClient: + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: identity + app.dependency_overrides[get_workflow_definition_repository] = lambda: ( + None if database is None else WorkflowDefinitionRepository(database.database) + ) + return TestClient(app) + + +def _body(name: str = "example", model: str = "gateway-model") -> Mapping[str, object]: + return {"bundle": bundle(name, model).model_dump(mode="json", by_alias=True, exclude_unset=True)} + + +def _register(client: TestClient, name: str = "example") -> WorkflowDraft: + response: Final = client.post("/v1/workflows", json=_body(name)) + assert response.status_code == 201, response.text + return WorkflowDraft.model_validate(response.json()) + + +def test_sdk_draft_contract_round_trip_has_exact_fields_and_etag() -> None: + database: Final = DefinitionRows() + client: Final = _client(database, caller()) + response: Final = client.post("/v1/workflows", json=_body()) + assert response.status_code == 201, response.text + draft: Final = WorkflowDraft.model_validate(response.json()) + assert response.json() == { + "name": "example", + "workflow_id": draft.workflow_id, + "team_id": "team", + "etag": response.headers["etag"], + "digest": bundle_digest(bundle()), + "bundle": _body()["bundle"], + "published_snapshot_id": None, + } + fetched: Final = client.get(f"/v1/workflows/{draft.workflow_id}") + assert fetched.json() == response.json() + assert fetched.headers["etag"] == draft.etag + + +def test_stale_put_and_publish_cannot_replace_or_approve_new_revision() -> None: + database: Final = DefinitionRows() + author: Final = _client(database, caller()) + reviewer: Final = _client(database, caller("reviewer")) + draft: Final = _register(author) + path: Final = f"/v1/workflows/{draft.workflow_id}" + changed: Final = author.put(path, json=_body(model="other-model"), headers={"If-Match": draft.etag}) + assert changed.status_code == 200, changed.text + assert changed.json()["etag"] != draft.etag + assert author.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 412 + published: Final = reviewer.post(f"{path}/publish", json={"digest": draft.digest}, headers={"If-Match": draft.etag}) + assert published.status_code == 412, published.text + assert author.get(path).json() == changed.json() + assert database.litellm_workflowsnapshot.rows == () + + +def test_publication_requires_noncontributor_administrator_and_matching_digest() -> None: + database: Final = DefinitionRows() + author: Final = _client(database, caller("reviewer")) + draft: Final = _register(author) + path: Final = f"/v1/workflows/{draft.workflow_id}/publish" + self_publish: Final = author.post(path, json={"digest": draft.digest}, headers={"If-Match": draft.etag}) + assert self_publish.status_code == 403 + ordinary: Final = _client(database, caller("other")) + assert ordinary.post(path, json={"digest": draft.digest}, headers={"If-Match": draft.etag}).status_code == 403 + admin: Final = _client( + database, UserAPIKeyAuth(user_id="independent-admin", user_role=LitellmUserRoles.PROXY_ADMIN, team_id="team") + ) + wrong_digest: Final = admin.post(path, json={"digest": "0" * 64}, headers={"If-Match": draft.etag}) + assert wrong_digest.status_code == 412 + published: Final = admin.post(path, json={"digest": draft.digest}, headers={"If-Match": draft.etag}) + assert published.status_code == 200, published.text + assert published.json() == { + "workflow_id": draft.workflow_id, + "snapshot_id": published.json()["snapshot_id"], + "model": "workflow/example", + "digest": draft.digest, + } + assert len(database.litellm_workflowsnapshot.rows) == 1 + + +def test_missing_etag_and_invalid_cursor_are_client_errors() -> None: + database: Final = DefinitionRows() + client: Final = _client(database, caller()) + draft: Final = _register(client) + path: Final = f"/v1/workflows/{draft.workflow_id}" + assert client.put(path, json=_body()).status_code == 412 + assert client.post(f"{path}/validate").status_code == 412 + assert client.get("/v1/workflows", params={"cursor": "malformed"}).status_code == 422 + valid: Final = client.post(f"{path}/validate", headers={"If-Match": draft.etag}) + assert valid.status_code == 200, valid.text + assert valid.json() == {"valid": True, "workflow_id": draft.workflow_id, "etag": draft.etag, "digest": draft.digest} + + +def test_openapi_describes_created_drafts_and_optional_empty_validation_body() -> None: + client: Final = _client(DefinitionRows(), caller()) + spec: Final = client.get("/openapi.json") + assert spec.status_code == 200 + assert "200" not in spec.json()["paths"]["/v1/workflows"]["post"]["responses"] + assert spec.json()["paths"]["/v1/workflows"]["post"]["responses"]["201"]["content"]["application/json"][ + "schema" + ] == {"$ref": "#/components/schemas/WorkflowDraft"} + request_body: Final = TypeAdapter(Mapping[str, object]).validate_python( + spec.json()["paths"]["/v1/workflows/{workflow_id}/validate"]["post"]["requestBody"] + ) + assert request_body.get("required", False) is False + assert spec.json()["components"]["schemas"]["WorkflowValidationRequest"] == { + "properties": {}, + "additionalProperties": False, + "type": "object", + "title": "WorkflowValidationRequest", + } + + +@pytest.mark.parametrize("body", [b"", b"{}", b"null"]) +def test_validation_accepts_the_documented_optional_empty_body(body: bytes) -> None: + client: Final = _client(DefinitionRows(), caller()) + draft: Final = _register(client) + response: Final = client.post( + f"/v1/workflows/{draft.workflow_id}/validate", + content=body, + headers={"If-Match": draft.etag, "Content-Type": "application/json"}, + ) + assert response.status_code == 200, response.text + assert response.json() == { + "valid": True, + "workflow_id": draft.workflow_id, + "etag": draft.etag, + "digest": draft.digest, + } + + +def test_validation_rejects_undocumented_action_fields() -> None: + client: Final = _client(DefinitionRows(), caller()) + draft: Final = _register(client) + response: Final = client.post( + f"/v1/workflows/{draft.workflow_id}/validate", + json={"execute": True}, + headers={"If-Match": draft.etag}, + ) + assert response.status_code == 422, response.text + + +def test_team_key_cannot_select_another_team_or_read_its_definition() -> None: + database: Final = DefinitionRows() + owner: Final = _client(database, caller()) + draft: Final = _register(owner) + outsider: Final = _client(database, caller("outsider", "other-team")) + assert outsider.get(f"/v1/workflows/{draft.workflow_id}").status_code == 404 + assert outsider.get(f"/v1/workflows/{draft.workflow_id}", params={"team_id": "team"}).status_code == 404 + assert outsider.post("/v1/workflows", json=_body(), params={"team_id": "team"}).status_code == 404 + assert outsider.get("/v1/workflows").json() == {"workflows": [], "next_cursor": None} + + +def test_same_team_key_requires_roster_membership_for_draft_access() -> None: + database: Final = DefinitionRows() + owner: Final = _client(database, caller()) + draft: Final = _register(owner) + outsider: Final = _client(database, caller("not-a-team-member")) + path: Final = f"/v1/workflows/{draft.workflow_id}" + assert outsider.get(path).status_code == 404 + assert outsider.get("/v1/workflows").status_code == 404 + assert outsider.post("/v1/workflows", json=_body("second")).status_code == 404 + assert outsider.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 404 + assert owner.get(path).json() == draft.model_dump(mode="json", by_alias=True, exclude_unset=True) + + +@pytest.mark.parametrize( + ("identity", "expected"), + [ + (caller(None), 403), + (UserAPIKeyAuth(user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY), 403), + (caller(), 500), + ], +) +def test_no_database_preserves_authorization_precedence(identity: UserAPIKeyAuth, expected: int) -> None: + response: Final = _client(None, identity).post("/v1/workflows", json=_body()) + assert response.status_code == expected, response.text + + +def test_duplicate_registration_is_conflict_and_list_cursor_never_repeats_rows() -> None: + database: Final = DefinitionRows() + client: Final = _client(database, caller()) + first: Final = _register(client, "first") + second: Final = _register(client, "second") + assert client.post("/v1/workflows", json=_body("first")).status_code == 409 + page: Final = DefinitionPage.model_validate(client.get("/v1/workflows", params={"limit": 1}).json()) + tail: Final = DefinitionPage.model_validate( + client.get("/v1/workflows", params={"limit": 1, "cursor": page.next_cursor}).json() + ) + assert {page.workflows[0].workflow_id, tail.workflows[0].workflow_id} == { + first.workflow_id, + second.workflow_id, + } + assert page.workflows[0].workflow_id != tail.workflows[0].workflow_id + assert tail.next_cursor is None + + +@pytest.mark.parametrize("field", ["api_key", "api_base", "callbacks", "headers", "router_settings"]) +def test_registration_rejects_untrusted_completion_capabilities(field: str) -> None: + client: Final = _client(DefinitionRows(), caller()) + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + unsafe: Final = {**original, "nodes": [{**original["nodes"][0], field: "untrusted"}]} + response: Final = client.post("/v1/workflows", json={"bundle": unsafe}) + assert response.status_code == 422, response.text diff --git a/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py new file mode 100644 index 00000000000..cb8459ede01 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py @@ -0,0 +1,332 @@ +from collections.abc import Mapping +from datetime import timedelta +from typing import Final + +import pytest +from fastapi import FastAPI, Request, Response +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient +from pydantic import JsonValue, TypeAdapter + +from litellm.models.workflow_definition import WorkflowDraft, WorkflowSnapshot +from litellm.models.workflow_execution import RunPage, RunState, StepTrace +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_workflow_definition_repository, get_workflow_execution_repository +from litellm.proxy.management_endpoints.workflow_definition_endpoints import router as definitions_router +from litellm.proxy.management_endpoints.workflow_execution_endpoints import invoke_workflow, router +from litellm.proxy.management_endpoints.workflow_management_endpoints import router as tracking_router +from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository +from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository +from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller +from tests.unit.proxy.workflows.test_runner import two_nodes +from tests.unit.repositories.test_workflow_execution_repository import NOW, ExecutionRows, pending_run + +_MESSAGES: Final = ({"role": "user", "content": "plan"},) +_ACCEPTED: Final = TypeAdapter(dict[str, str]) + + +def virtual_key(user: str = "author", team: str = "team", key: str = "owner-key") -> UserAPIKeyAuth: + return caller(user, team).model_copy(update={"token": key, "via_virtual_key": True}) + + +def client_for(definitions: DefinitionRows, executions: ExecutionRows, identity: UserAPIKeyAuth) -> TestClient: + app: Final = FastAPI() + app.include_router(tracking_router) + app.include_router(definitions_router) + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: identity + app.dependency_overrides[get_workflow_definition_repository] = lambda: WorkflowDefinitionRepository( + definitions.database + ) + app.dependency_overrides[get_workflow_execution_repository] = lambda: WorkflowExecutionRepository(executions) + return TestClient(app) + + +def register(client: TestClient) -> WorkflowDraft: + response: Final = client.post( + "/v1/workflows", json={"bundle": bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)} + ) + assert response.status_code == 201, response.text + return WorkflowDraft.model_validate(response.json()) + + +def test_paid_test_returns_202_and_persists_pending_admission_without_executing_inline() -> None: + definitions: Final = DefinitionRows() + executions: Final = ExecutionRows() + client: Final = client_for(definitions, executions, virtual_key()) + draft: Final = register(client) + response: Final = client.post( + f"/v1/workflows/{draft.workflow_id}/test", + json={"messages": _MESSAGES, "authorize_billing": True}, + headers={"If-Match": draft.etag}, + ) + assert response.status_code == 202, response.text + assert len(executions.runs) == 1 + admitted: Final = executions.runs[0] + assert response.json() == { + "run_id": admitted.run_id, + "snapshot_id": admitted.snapshot_id, + "trace_url": f"/v1/workflows/runs/{admitted.run_id}/steps", + } + assert response.headers["x-litellm-workflow-run-id"] == admitted.run_id + assert (admitted.status, admitted.output, admitted.lease_token) == (RunState.PENDING, None, None) + assert admitted.input.messages == _MESSAGES and admitted.input.draft + assert (admitted.created_by, admitted.team_id) == ("owner-key", "team") + assert executions.events == () + assert not any(query.sql.startswith("UPDATE") for query in executions.queries) + trace: Final = client.get(f"/v1/workflows/runs/{admitted.run_id}/steps") + assert trace.status_code == 200, trace.text + assert trace.json() == { + "run_id": admitted.run_id, + "snapshot_id": admitted.snapshot_id, + "state": "pending", + "cancel_requested": False, + "steps": [], + "error": None, + } + + +def test_idempotent_test_after_draft_edit_returns_original_snapshot_and_conflict_for_changed_input() -> None: + definitions: Final = DefinitionRows() + executions: Final = ExecutionRows() + client: Final = client_for(definitions, executions, virtual_key()) + draft: Final = register(client) + path: Final = f"/v1/workflows/{draft.workflow_id}" + headers: Final = {"If-Match": draft.etag, "Idempotency-Key": "retry-1"} + body: Final = {"messages": _MESSAGES, "authorize_billing": True} + admitted: Final = client.post(f"{path}/test", json=body, headers=headers) + assert admitted.status_code == 202, admitted.text + edited: Final = client.put( + path, + json={"bundle": bundle(model="changed-model").model_dump(mode="json", by_alias=True, exclude_unset=True)}, + headers={"If-Match": draft.etag}, + ) + assert edited.status_code == 200, edited.text + replay: Final = client.post(f"{path}/test", json=body, headers=headers) + assert replay.status_code == 202 and replay.json() == admitted.json(), replay.text + frozen: Final = client.get(f"/v1/workflows/runs/{admitted.json()['run_id']}/definition") + assert frozen.status_code == 200, frozen.text + assert WorkflowSnapshot.model_validate(frozen.json()).bundle == draft.bundle + assert len(definitions.litellm_workflowsnapshot.rows) == 1 and len(executions.runs) == 1 + conflict: Final = client.post( + f"{path}/test", json={**body, "messages": [{"role": "user", "content": "different"}]}, headers=headers + ) + assert conflict.status_code == 409, conflict.text + assert conflict.json()["detail"]["error"]["code"] == "idempotency_conflict" + assert len(executions.runs) == 1 and len(definitions.litellm_workflowsnapshot.rows) == 1 + + +@pytest.mark.parametrize("billing_authorized", [False, True]) +def test_billing_consent_and_virtual_key_are_required_before_creating_snapshot_or_run(billing_authorized: bool) -> None: + definitions: Final = DefinitionRows() + executions: Final = ExecutionRows() + author: Final = client_for(definitions, executions, virtual_key()) + draft: Final = register(author) + session: Final = virtual_key().model_copy(update={"is_session_token": True}) + client: Final = client_for(definitions, executions, session if billing_authorized else virtual_key()) + response: Final = client.post( + f"/v1/workflows/{draft.workflow_id}/test", + json={"messages": _MESSAGES, "authorize_billing": billing_authorized}, + headers={"If-Match": draft.etag}, + ) + assert response.status_code == (403 if billing_authorized else 400), response.text + assert response.json()["detail"]["error"]["code"] == ( + "forbidden" if billing_authorized else "billing_authorization_required" + ) + assert executions.runs == () and definitions.litellm_workflowsnapshot.rows == () + + +@pytest.mark.parametrize("identity", [virtual_key("other", key="other-key"), virtual_key("outsider", "other-team")]) +def test_foreign_run_denies_trace_definition_and_cancellation_without_side_effects(identity: UserAPIKeyAuth) -> None: + definitions: Final = DefinitionRows() + executions: Final = ExecutionRows() + owner: Final = client_for(definitions, executions, virtual_key()) + draft: Final = register(owner) + accepted: Final = owner.post( + f"/v1/workflows/{draft.workflow_id}/test", + json={"messages": _MESSAGES, "authorize_billing": True}, + headers={"If-Match": draft.etag}, + ) + assert accepted.status_code == 202, accepted.text + path: Final = f"/v1/workflows/runs/{accepted.json()['run_id']}" + foreign: Final = client_for(definitions, executions, identity) + assert foreign.get(f"{path}/steps").status_code == 404 + assert foreign.get(f"{path}/definition").status_code == 404 + denied: Final = foreign.post(f"{path}/cancel") + assert denied.status_code == 404, denied.text + assert not executions.runs[0].cancel_requested + assert not any(query.sql.startswith("UPDATE") for query in executions.queries) + assert owner.get(f"{path}/steps").status_code == 200 + assert owner.post(f"{path}/cancel").status_code == 202 + assert executions.runs[0].cancel_requested + + +def test_history_caps_page_size_and_scopes_ordinary_key_but_allows_team_admin() -> None: + definitions: Final = DefinitionRows() + executions: Final = ExecutionRows() + owner: Final = client_for(definitions, executions, virtual_key()) + draft: Final = register(owner) + executions.runs = ( + *(pending_run(f"run-{index:03}", workflow_id=draft.workflow_id) for index in range(101)), + pending_run("foreign-key", workflow_id=draft.workflow_id, key_hash="other-key"), + pending_run("foreign-team", workflow_id=draft.workflow_id, team_id="other-team"), + ) + response: Final = owner.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id, "limit": 250}) + assert response.status_code == 200, response.text + page: Final = RunPage.model_validate(response.json()) + assert page.count == 100 and len(page.runs) == 100 and page.next_cursor is not None + assert tuple(run.run_id for run in page.runs) == tuple(f"run-{index:03}" for index in range(100, 0, -1)) + assert executions.queries[-1].args == (draft.workflow_id, "team", "owner-key", None, None, None, 101) + tail_response: Final = owner.get( + f"/v1/workflows/{draft.workflow_id}/runs", params={"limit": 100, "after": page.next_cursor} + ) + tail: Final = RunPage.model_validate(tail_response.json()) + assert tuple(run.run_id for run in tail.runs) == ("run-000",) + assert tail.next_cursor is None + admin: Final = client_for(definitions, executions, virtual_key("reviewer", key="admin-key")) + team_response: Final = admin.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id, "limit": 250}) + assert team_response.status_code == 200, team_response.text + assert executions.queries[-1].args == (draft.workflow_id, "team", None, None, None, None, 101) + outsider: Final = client_for(definitions, executions, virtual_key("outsider", "other-team")) + denied: Final = outsider.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id, "team_id": "team"}) + assert denied.status_code == 404, denied.text + assert owner.get(f"/v1/workflows/{draft.workflow_id}/runs", params={"limit": 101}).status_code == 422 + + +@pytest.mark.asyncio +async def test_prefer_respond_async_returns_published_run_id_without_waiting_for_execution() -> None: + definitions: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await definitions.create(bundle(), "team", "author") + assert isinstance(draft, WorkflowDraft) + snapshot: Final = await definitions.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") + assert isinstance(snapshot, WorkflowSnapshot) + database: Final = ExecutionRows() + request: Final = Request( + {"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": [(b"prefer", b"respond-async")]} + ) + response: Final = await invoke_workflow( + request, + Response(), + virtual_key(), + {"model": "workflow/example", "messages": _MESSAGES}, + definitions, + WorkflowExecutionRepository(database), + ) + assert isinstance(response, JSONResponse) and response.status_code == 202 + assert len(database.runs) == 1 + run: Final = database.runs[0] + assert _ACCEPTED.validate_json(bytes(response.body)) == { + "run_id": run.run_id, + "snapshot_id": snapshot.snapshot_id, + "trace_url": f"/v1/workflows/runs/{run.run_id}/steps", + } + assert response.headers["x-litellm-workflow-run-id"] == run.run_id + assert run.status is RunState.PENDING and not run.input.draft and run.output is None + assert database.events == () + assert not any(query.sql.startswith("UPDATE") for query in database.queries) + + +@pytest.mark.asyncio +async def test_sdk_retry_after_republish_reuses_pinned_run_and_sums_node_usage() -> None: + definitions: Final = WorkflowDefinitionRepository(DefinitionRows().database) + original: Final = two_nodes() + draft: Final = await definitions.create(original, "team", "author") + assert isinstance(draft, WorkflowDraft) + snapshot: Final = await definitions.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") + assert isinstance(snapshot, WorkflowSnapshot) + database: Final = ExecutionRows() + executions: Final = WorkflowExecutionRepository(database) + body: Final = {"model": "workflow/example", "messages": _MESSAGES} + headers: Final = [(b"idempotency-key", b"published-retry")] + scope: Final = {"type": "http", "method": "POST", "path": "/v1/chat/completions"} + accepted: Final = await invoke_workflow( + Request({**scope, "headers": [*headers, (b"prefer", b"respond-async")]}), + Response(), + virtual_key(), + body, + definitions, + executions, + ) + assert isinstance(accepted, JSONResponse) and accepted.status_code == 202 + run: Final = database.runs[0] + assert await executions.claim(run.run_id, "worker", NOW + timedelta(seconds=15)) is not None + for node_id, prompt, completion, cost in (("node_1", 7, 2, 0.1), ("node_2", 11, 3, 0.2)): + assert await executions.checkpoint( + run.run_id, + "worker", + StepTrace( + step_id=node_id, + run_id=run.run_id, + node_id=node_id, + state=RunState.COMPLETED, + input=[dict(message) for message in run.input.messages], + output={"answer": "approved"}, + attempts=1, + deadline=run.input.deadline, + started_at=NOW, + finished_at=NOW, + prompt_tokens=prompt, + completion_tokens=completion, + response_cost=cost, + request_id=f"receipt-{node_id}", + ), + ) + assert await executions.finish(run.run_id, "worker", RunState.COMPLETED, {"answer": "approved"}, None) + changed: Final = await definitions.replace( + draft.workflow_id, + "team", + draft.etag, + original.model_copy(update={"timeout_seconds": 120}), + "author", + ) + assert isinstance(changed, WorkflowDraft) + published: Final = await definitions.publish(changed.workflow_id, "team", changed.etag, changed.digest, "reviewer") + assert isinstance(published, WorkflowSnapshot) and published.snapshot_id != snapshot.snapshot_id + response: Final = Response() + result: Final = await invoke_workflow( + Request({**scope, "headers": headers}), + response, + virtual_key(), + body, + definitions, + executions, + ) + assert result == { + "id": f"chatcmpl-workflow-{run.run_id}", + "object": "chat.completion", + "created": int(NOW.timestamp()), + "model": "workflow/example", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": '{"answer": "approved"}'}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 18, "completion_tokens": 5, "total_tokens": 23}, + } + assert response.headers["x-litellm-workflow-run-id"] == run.run_id + assert len(database.runs) == 1 and database.runs[0].snapshot_id == snapshot.snapshot_id + + +@pytest.mark.parametrize( + ("metadata", "error"), + [ + ({"error": "Pinned input schema rejected messages"}, "Pinned input schema rejected messages"), + ({"error": {"internal": "malformed metadata"}}, None), + ], +) +def test_trace_exposes_run_level_failure_without_steps_and_only_string_errors( + metadata: Mapping[str, JsonValue], + error: str | None, +) -> None: + run: Final = pending_run().model_copy(update={"status": RunState.FAILED, "metadata": metadata}) + response: Final = client_for(DefinitionRows(), ExecutionRows((run,)), virtual_key()).get( + f"/v1/workflows/runs/{run.run_id}/steps" + ) + assert response.status_code == 200, response.text + assert response.json() == { + "run_id": run.run_id, + "snapshot_id": run.snapshot_id, + "state": "failed", + "cancel_requested": False, + "steps": [], + "error": error, + } diff --git a/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py index 0c3d5107cb6..1b1119cb6ed 100644 --- a/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py @@ -4,8 +4,9 @@ Uses FastAPI TestClient with a mocked prisma_client. """ import asyncio +from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -113,8 +114,15 @@ def _make_prisma_client() -> MagicMock: def _make_app() -> FastAPI: - app = FastAPI() + from litellm.proxy.deps.repositories import ( + get_workflow_definition_repository, + get_workflow_execution_repository, + ) + + app: Final = FastAPI() app.include_router(router) + app.dependency_overrides[get_workflow_definition_repository] = lambda: None + app.dependency_overrides[get_workflow_execution_repository] = lambda: None return app @@ -734,3 +742,38 @@ class TestAdminViewerReadParity: with pytest.raises(HTTPException) as exc_info: asyncio.run(_require_run(prisma, "run-1", _override_auth_admin_viewer())) assert exc_info.value.status_code == 404 + + +@pytest.mark.parametrize( + ("method", "suffix", "body"), + [ + ("PATCH", "", {"status": "completed", "output": {"forged": True}}), + ("POST", "/events", {"event_type": "step.failed", "step_name": "node-1"}), + ("POST", "/messages", {"role": "assistant", "content": "forged result"}), + ], +) +def test_engine_owned_runs_reject_tracking_writes_even_for_proxy_admin( + method: str, suffix: str, body: Mapping[str, object], +) -> None: + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from tests.unit.repositories.test_workflow_execution_repository import NOW + + database: Final = _make_prisma_client() + run: Final = _make_run() + run.created_at = NOW + run.updated_at = NOW + run.engine_owned = True + database.db.litellm_workflowrun.find_unique = AsyncMock(return_value=run) + database.db.litellm_workflowrun.update = AsyncMock(return_value=run) + database.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + database.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) + database.db.litellm_workflowmessage.create = AsyncMock(return_value=_make_message()) + app: Final = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth_admin + with patch("litellm.proxy.proxy_server.prisma_client", database), TestClient(app) as client: + response: Final = client.request(method, f"/v1/workflows/runs/run-1{suffix}", json=body) + assert response.status_code == 403, response.text + assert response.json() == {"detail": "Gateway executions cannot be changed through tracking endpoints"} + database.db.litellm_workflowrun.update.assert_not_awaited() + database.db.litellm_workflowmessage.create.assert_not_awaited() + database.db.tx.assert_not_called() diff --git a/tests/unit/proxy/workflow_definition_rows.py b/tests/unit/proxy/workflow_definition_rows.py new file mode 100644 index 00000000000..d03e1d5bef7 --- /dev/null +++ b/tests/unit/proxy/workflow_definition_rows.py @@ -0,0 +1,137 @@ +import asyncio +from collections.abc import AsyncGenerator, Mapping, Sequence +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from typing import Final, cast # noqa: TID251 # the test database implements the delegates exercised by these routes + +from pydantic import TypeAdapter + +from litellm.models.workflow_definition import WorkflowBundle +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.repositories.workflow_definition_repository import DefinitionDatabase + +NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_COMPARISON: Final = TypeAdapter(Mapping[str, str]) + + +def bundle(name: str = "example", model: str = "gateway-model") -> WorkflowBundle: + return WorkflowBundle.model_validate( + { + "contract_version": "1", + "name": name, + "timeout_seconds": 300.0, + "input_schema": {"type": "array"}, + "output": {"from": "node_1", "output": "result"}, + "nodes": [ + { + "id": "node_1", + "type": "chat_completion", + "model": model, + "messages": {"from": "$input", "output": "messages"}, + "timeout_seconds": 60.0, + "retry": {"max_attempts": 1, "backoff_seconds": 0.0}, + } + ], + } + ) + + +def caller(user_id: str | None = "author", team_id: str = "team") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id=user_id, team_id=team_id, user_role=LitellmUserRoles.INTERNAL_USER) + + +def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool: + return all( + str(row.get(key)) > _COMPARISON.validate_python(value)["gt"] + if isinstance(value, Mapping) and "gt" in value + else row.get(key) == value + for key, value in where.items() + ) + + +class Rows: + def __init__(self, rows: tuple[Mapping[str, object], ...] = ()) -> None: + self.rows = rows + + async def find_first( + self, + skip: int | None = None, + where: Mapping[str, object] | None = None, + cursor: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + distinct: Sequence[str] | None = None, + ) -> object | None: + return next((row for row in self.rows if _matches(row, where or {})), None) + + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> object | None: + return await self.find_first(where=where) + + async def find_many( + self, + take: int | None = None, + skip: int | None = None, + where: Mapping[str, object] | None = None, + cursor: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + distinct: Sequence[str] | None = None, + ) -> Sequence[object]: + matches: Final = tuple(row for row in self.rows if _matches(row, where or {})) + return sorted(matches, key=lambda row: str(row.get("workflow_id")))[:take] + + async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> object: + if "etag" in data and any( + row.get("name") == data["name"] and row.get("team_id") == data["team_id"] for row in self.rows + ): + from prisma.errors import UniqueViolationError + + raise UniqueViolationError({"is_panic": False, "message": "duplicate", "error_code": "P2002"}) + row: Final = { + "created_at": NOW, + "published_etag": None, + "published_snapshot_id": None, + **data, + } + self.rows = (*self.rows, row) + return row + + async def update( + self, data: Mapping[str, object], where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> object | None: + self.rows = tuple({**row, **data} if _matches(row, where) else row for row in self.rows) + return await self.find_first(where=where) + + +class DefinitionRows: + def __init__(self) -> None: + self.litellm_workflowdefinition = Rows() + self.litellm_workflowsnapshot = Rows() + self.litellm_teamtable = Rows( + ( + { + "team_id": "team", + "members_with_roles": [ + {"user_id": "author", "role": "user"}, + {"user_id": "reviewer", "role": "admin"}, + {"user_id": "other", "role": "user"}, + ], + }, + {"team_id": "other-team", "members_with_roles": [{"user_id": "outsider", "role": "admin"}]}, + ) + ) + self.lock = asyncio.Lock() + + async def query_raw(self, query: str, *args: object) -> object: + return () + + @property + def database(self) -> DefinitionDatabase: + return cast(DefinitionDatabase, self) # cast-ok: isolated Prisma boundary implements the used table operations + + @asynccontextmanager + async def tx(self) -> AsyncGenerator["DefinitionRows"]: + async with self.lock: + yield self diff --git a/tests/unit/proxy/workflows/__init__.py b/tests/unit/proxy/workflows/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/workflows/monty/__init__.py b/tests/unit/proxy/workflows/monty/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/workflows/monty/test_compiler.py b/tests/unit/proxy/workflows/monty/test_compiler.py new file mode 100644 index 00000000000..2dfdff5a077 --- /dev/null +++ b/tests/unit/proxy/workflows/monty/test_compiler.py @@ -0,0 +1,395 @@ +import hashlib +from functools import reduce +from typing import Final + +import pytest + +from litellm.proxy.workflows.monty.compiler import compile_program, escape_diagnostics +from litellm.proxy.workflows.monty.models import ( + CompiledProgram, + JsonSchema, + MontyErrorKind, + MontyFailure, + MontyLimits, + PythonProgram, +) + + +def _program(source: str, *, output: JsonSchema | None = None) -> PythonProgram: + return PythonProgram( + contract_version="1", + source=source, + function_name="transform", + parameter_names=("value",), + input_schema={ + "type": "object", + "properties": {"value": {}}, + "required": ["value"], + "additionalProperties": False, + }, + output_schema=output if output is not None else {}, + source_sha256=hashlib.sha256(source.encode()).hexdigest(), + ) + + +@pytest.mark.parametrize( + "source", + [ + "import os\ndef transform(value): return value", + "from typing import TypedDict\ndef transform(value): return value", + "x=1\ndef transform(value): return value", + "print('side effect')\ndef transform(value): return value", + "async def transform(value): return value", + "class X: pass\ndef transform(value): return value", + "@print\ndef transform(value): return value", + "def transform(value):\n def closure(): return value\n return closure()", + "def transform(value): return (lambda: value)()", + "def transform(value):\n global x\n return value", + "def transform(value):\n nonlocal x\n return value", + "def transform(value): return value.__class__", + "def transform(value): return __name__", + "# /// script\ndef transform(value): return value", + "def transform(other): return other", + "def transform(value=1): return value", + "def transform(*value): return value", + "def transform(value, **kwargs): return value", + "def transform(*, value): return value", + "def transform(value, /): return value", + "def transform(value):\n import json\n return value", + "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()) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.INVALID_PROGRAM + assert "line" in result.message + + +@pytest.mark.parametrize( + "builtin", + [ + "eval", + "exec", + "compile", + "open", + "__import__", + "globals", + "locals", + "vars", + "getattr", + "setattr", + "delattr", + "breakpoint", + "input", + "help", + ], +) +def test_rejects_dangerous_builtin_references(builtin: str) -> None: + result: Final = compile_program(_program(f"def transform(value): return {builtin}"), MontyLimits()) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.INVALID_PROGRAM + + +@pytest.mark.parametrize( + "schema", + [ + {"$ref": "https://example.invalid/schema"}, + {"$ref": "file:///etc/passwd"}, + {"$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": {}}}}, + { + "$defs": { + "Bad-Name": { + "type": "object", + "x-liteflow-record": True, + "properties": {}, + "additionalProperties": False, + } + } + }, + ], +) +def test_rejects_unsupported_schema_without_remote_retrieval(schema: JsonSchema) -> None: + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + + +def test_hash_covers_exact_source_bytes() -> None: + program: Final = _program("def transform(value): return value") + result: Final = compile_program(program.model_copy(update={"source": program.source + "\n"}), MontyLimits()) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.INVALID_PROGRAM + assert "source_sha256" in result.message + + +def test_annotations_are_not_evaluated() -> None: + result: Final = compile_program( + _program("def transform(value: Unknown[Foo]) -> Missing: return value"), MontyLimits() + ) + assert isinstance(result, CompiledProgram) + + +def test_source_size_is_bounded() -> None: + result: Final = compile_program(_program("def transform(value): return value"), MontyLimits(max_source_bytes=8)) + assert isinstance(result, MontyFailure) + assert result.kind == MontyErrorKind.INVALID_PROGRAM + + +def test_diagnostics_escape_terminal_controls_but_preserve_layout() -> None: + assert escape_diagnostics("a\x1b[31m\x00\x7f\x9b\n\t") == "a\\x1b[31m\\x00\\x7f\\x9b\n\t" + + +def _array_schema(child: JsonSchema, index: int) -> JsonSchema: + return {"type": "array", "items": child} + + +@pytest.mark.parametrize("depth,accepted", [(4, True), (70, False)]) +def test_schema_nesting_is_bounded(depth: int, accepted: bool) -> None: + schema: Final = reduce(_array_schema, range(depth), {}) + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + if accepted: + assert isinstance(result, CompiledProgram), result + return + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert "depth limit" in result.message + + +@pytest.mark.parametrize("depth,accepted", [(4, True), (70, False)]) +def test_schema_reference_expansion_is_bounded(depth: int, accepted: bool) -> None: + schema: Final[JsonSchema] = { + "$defs": { + f"N{index}": {"$ref": f"#/$defs/N{index + 1}"} if index < depth else {"type": "integer"} + for index in range(depth + 1) + }, + "$ref": "#/$defs/N0", + } + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + if accepted: + assert isinstance(result, CompiledProgram), result + return + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert "depth limit" in result.message + + +def test_schema_reference_cycle_is_rejected() -> None: + schema: Final[JsonSchema] = {"$defs": {"X": {"$ref": "#/$defs/X"}}, "$ref": "#/$defs/X"} + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert "recursive schema" in result.message + + +_RECORD: Final[JsonSchema] = { + "type": "object", + "x-liteflow-record": True, + "properties": {"amount": {"type": "integer"}, "label": {"type": ["string", "null"]}}, + "required": ["amount"], + "additionalProperties": False, +} +_RECORD_REF: Final[JsonSchema] = {"$ref": "#/$defs/Record"} + + +@pytest.mark.parametrize( + "argument", + [ + _RECORD_REF, + {"type": "array", "items": _RECORD_REF}, + {"anyOf": [_RECORD_REF, {"type": "null"}]}, + {"type": "object", "properties": {"payload": _RECORD_REF}, "additionalProperties": True}, + {"type": "object", "properties": {}, "additionalProperties": _RECORD_REF}, + {"$ref": "#/$defs/Envelope"}, + ], +) +def test_registration_accepts_records_in_supported_input_shapes(argument: JsonSchema) -> None: + definitions: Final[JsonSchema] = { + "Record": _RECORD, + "Empty": {**_RECORD, "properties": {}, "required": []}, + "Envelope": {**_RECORD, "properties": {"child": _RECORD_REF}, "required": ["child"]}, + } + program: Final = _program("def transform(value): return value", output={"$defs": definitions}) + typed: Final = program.model_copy( + update={ + "input_schema": { + **program.input_schema, + "properties": {"value": argument}, + "$defs": definitions, + } + } + ) + result: Final = compile_program(typed, MontyLimits()) + assert isinstance(result, CompiledProgram), result + + +@pytest.mark.parametrize( + "schema", + [ + _RECORD, + {"$defs": {"Record": {**_RECORD, "type": "string"}}}, + {"$defs": {"Record": {**_RECORD, "properties": {"bad-field": {}}}}}, + {"$defs": {"str": _RECORD}}, + {"$defs": {"transform": _RECORD}}, + {"$defs": {"Record": _RECORD}, "oneOf": [_RECORD_REF, {"type": "null"}]}, + {"$defs": {"Record": _RECORD}, "anyOf": [_RECORD_REF, {"type": "string"}]}, + {"$defs": {"Record": _RECORD}, "anyOf": [_RECORD_REF, {"type": "null"}, {"type": "string"}]}, + {"$defs": {"Record": _RECORD}, "allOf": [_RECORD_REF]}, + ], +) +def test_registration_rejects_ambiguous_or_unconstructible_record_shapes(schema: JsonSchema) -> None: + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + + +@pytest.mark.parametrize( + "changes", + [ + {"type": "array"}, + {"properties": {"extra": {}}}, + {"required": []}, + {"required": ["value", "extra"]}, + {"additionalProperties": True}, + ], +) +def test_input_schema_must_require_exactly_the_entry_parameters(changes: JsonSchema) -> None: + program: Final = _program("def transform(value): return value") + invalid: Final = program.model_copy(update={"input_schema": {**program.input_schema, **changes}}) + result: Final = compile_program(invalid, MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + + +def test_registration_rejects_conflicting_record_definitions() -> None: + program: Final = _program("def transform(value): return value", output={"$defs": {"Record": _RECORD}}) + conflicting: Final = program.model_copy( + update={"input_schema": {**program.input_schema, "$defs": {"Record": {**_RECORD, "required": []}}}} + ) + result: Final = compile_program(conflicting, MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert "conflicting" in result.message + + +@pytest.mark.parametrize( + "source", + ["def transform(value):\n", "def transform(value): return value\ndef transform(value): return 1"], +) +def test_registration_rejects_syntax_errors_and_duplicate_entry_names(source: str) -> None: + result: Final = compile_program(_program(source), MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.INVALID_PROGRAM + + +def test_registration_accepts_local_json_pointer_escaping_and_annotations_import() -> None: + schema: Final[JsonSchema] = {"$defs": {"a/b~c": {"type": "integer"}}, "$ref": "#/$defs/a~1b~0c"} + source: Final = "from __future__ import annotations\ndef transform(value: Missing) -> Unknown: return value" + result: Final = compile_program(_program(source, output=schema), MontyLimits()) + assert isinstance(result, CompiledProgram), result + + +@pytest.mark.parametrize("name", ["default", "pattern", "enum", "$ref", "$ref-like"]) +def test_keyword_looking_property_names_are_data(name: str) -> None: + schema: Final[JsonSchema] = {"type": "object", "properties": {name: {"type": "string"}}} + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, CompiledProgram), result + + +@pytest.mark.parametrize( + "schema", + [ + {"properties": {"default": {"$ref": "https://example.invalid/schema"}}}, + {"$defs": {"default": {"items": {"$ref": "https://example.invalid/schema"}}}}, + {"anyOf": [{"properties": {"enum": {"$ref": "https://example.invalid/schema"}}}]}, + {"default": {"$ref": "https://example.invalid/schema"}, "$ref": "#/default"}, + ], +) +def test_remote_references_under_keyword_looking_names_are_rejected(schema: JsonSchema) -> None: + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert "local" in result.message + + +@pytest.mark.parametrize( + "keyword", + [ + "pattern", + "patternProperties", + "uniqueItems", + "$id", + "$anchor", + "$dynamicRef", + "$dynamicAnchor", + "$schema", + "if", + "then", + "else", + "dependentRequired", + "dependentSchemas", + "contains", + "minContains", + "maxContains", + "propertyNames", + "unevaluatedProperties", + "unevaluatedItems", + "discriminator", + "unknown", + ], +) +def test_unsafe_keywords_are_explicitly_refused(keyword: str) -> None: + result: Final = compile_program(_program("def transform(value): return value", output={keyword: {}}), MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert keyword in result.message + + +@pytest.mark.parametrize("keyword", ["const", "enum", "default", "examples"]) +def test_annotation_and_data_values_are_not_interpreted_as_schemas(keyword: str) -> None: + payload: Final[JsonSchema] = {"$ref": "https://example.invalid/data", "pattern": "inert"} + schema: Final[JsonSchema] = {keyword: [payload] if keyword in {"enum", "examples"} else payload} + result: Final = compile_program(_program("def transform(value): return value", output=schema), MontyLimits()) + assert isinstance(result, CompiledProgram), result + + +@pytest.mark.parametrize("depth", [8, 12, 14, 22]) +def test_repeated_references_compile_within_shared_work_budget(depth: int) -> None: + definitions: Final[JsonSchema] = { + f"N{index}": {"anyOf": [{"$ref": f"#/$defs/N{index + 1}"}, {"$ref": f"#/$defs/N{index + 1}"}]} + for index in range(depth) + } + schema: Final[JsonSchema] = {"$defs": {**definitions, f"N{depth}": {}}, "$ref": "#/$defs/N0"} + program: Final = _program("def transform(value): return value", output=schema) + result: Final = compile_program(program, MontyLimits()) + assert isinstance(result, CompiledProgram), result + + +def test_expanded_input_builder_fails_closed_at_shared_work_budget() -> None: + definitions: Final[JsonSchema] = { + f"N{index}": { + "type": "object", + "properties": {"left": {"$ref": f"#/$defs/N{index + 1}"}, "right": {"$ref": f"#/$defs/N{index + 1}"}}, + } + for index in range(14) + } + program: Final = _program("def transform(value): return value") + expanded: Final = program.model_copy( + update={ + "input_schema": { + **program.input_schema, + "$defs": {**definitions, "N14": {"type": "integer"}}, + "properties": {"value": {"$ref": "#/$defs/N0"}}, + } + } + ) + result: Final = compile_program(expanded, MontyLimits()) + assert isinstance(result, MontyFailure), result + assert result.kind == MontyErrorKind.UNSUPPORTED_SCHEMA + assert result.message == "schema expansion budget exceeded" 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/tests/unit/proxy/workflows/test_definitions.py b/tests/unit/proxy/workflows/test_definitions.py new file mode 100644 index 00000000000..ac5ccdb89e4 --- /dev/null +++ b/tests/unit/proxy/workflows/test_definitions.py @@ -0,0 +1,98 @@ +import hashlib +from typing import Final + +import pytest + +from litellm.models.workflow_definition import DefinitionError, WorkflowBundle, WorkflowDraft, WorkflowSnapshot +from litellm.proxy.workflows.definitions import freeze_draft, published_workflow, validate_bundle +from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository +from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller + + +def _python_bundle(source: str) -> WorkflowBundle: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + return WorkflowBundle.model_validate( + { + **original, + "output": {"from": "node_2", "output": "result"}, + "nodes": [ + *original["nodes"], + { + "id": "node_2", + "type": "python", + "timeout_seconds": 10.0, + "retry": {"max_attempts": 1, "backoff_seconds": 0.0}, + "inputs": {"text": {"from": "node_1", "output": "result"}}, + "program": { + "contract_version": "1", + "source": source, + "function_name": "transform", + "parameter_names": ["text"], + "input_schema": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + "additionalProperties": False, + }, + "output_schema": {"type": "string"}, + "source_sha256": hashlib.sha256(source.encode()).hexdigest(), + }, + }, + ], + } + ) + + +def test_validation_compiles_python_without_executing_its_body() -> None: + program: Final = _python_bundle("def transform(text: str) -> str:\n return str(1 // 0)\n") + assert validate_bundle(program) is None + + +@pytest.mark.parametrize( + "source", + [ + "import os\ndef transform(text: str) -> str:\n return text\n", + "def transform(text: str) -> str:\n return open(text).read()\n", + ], +) +def test_validation_rejects_imports_and_host_capabilities(source: str) -> None: + result: Final = validate_bundle(_python_bundle(source)) + assert isinstance(result, DefinitionError) + assert result.code == "invalid_bundle" + assert result.param == "nodes.node_2.program" + + +def test_validation_rejects_forward_references_and_external_schema_refs() -> None: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + forward: Final = WorkflowBundle.model_validate( + {**original, "nodes": [{**original["nodes"][0], "messages": {"from": "later", "output": "result"}}]} + ) + remote: Final = WorkflowBundle.model_validate({**original, "input_schema": {"$ref": "https://example.com/schema"}}) + assert validate_bundle(forward) == DefinitionError( + "invalid_bundle", "Bindings must reference input messages or an earlier node result", "nodes.node_1.messages" + ) + assert validate_bundle(remote) == DefinitionError( + "invalid_bundle", "Schemas must use local references and bounded nesting", "bundle.input_schema" + ) + + +@pytest.mark.parametrize("keyword", ["$id", "$dynamicRef", "$recursiveRef"]) +def test_validation_rejects_schema_reference_scope_changes(keyword: str) -> None: + original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True) + scoped: Final = WorkflowBundle.model_validate({**original, "input_schema": {keyword: "https://example.com/schema"}}) + assert validate_bundle(scoped) == DefinitionError( + "invalid_bundle", "Schemas must use local references and bounded nesting", "bundle.input_schema" + ) + + +@pytest.mark.asyncio +async def test_application_identity_can_resolve_published_workflow_but_cannot_freeze_draft() -> None: + store: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await store.create(bundle(), "team", "author") + assert isinstance(draft, WorkflowDraft) + snapshot: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") + assert isinstance(snapshot, WorkflowSnapshot) + assert await published_workflow(store, caller(None), "example") == snapshot + denied: Final = await freeze_draft(store, caller(None), draft.workflow_id, draft.etag) + assert isinstance(denied, DefinitionError) + assert denied.code == "not_found" diff --git a/tests/unit/proxy/workflows/test_gateway.py b/tests/unit/proxy/workflows/test_gateway.py new file mode 100644 index 00000000000..6c600702ee9 --- /dev/null +++ b/tests/unit/proxy/workflows/test_gateway.py @@ -0,0 +1,541 @@ +import asyncio +import json +from collections.abc import Iterator, Mapping +from contextvars import ContextVar +from dataclasses import dataclass, field, replace +from datetime import datetime, timedelta, timezone, tzinfo +from types import SimpleNamespace +from typing import Final, Literal + +import pytest +from fastapi import HTTPException +from pydantic import JsonValue + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.models.workflow_definition import ChatCompletionNode +from litellm.models.workflow_execution import ChildCallResult, ExecutionError, workflow_call_id +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import cache_key_object +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_cache_key, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, +) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import ProxyLogging +from litellm.proxy.workflows import gateway +from litellm.types.utils import ModelResponse +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + +_NOW: Final = datetime(2026, 11, 6, tzinfo=timezone.utc) +_PARENT_STATE: Final = ContextVar("workflow_gateway_test_parent", default="isolated") +_KEY: Final = "a" * 64 + + +class _Clock(datetime): + @classmethod + def now(cls, tz: tzinfo | None = None) -> datetime: + return _NOW + + +@pytest.fixture(autouse=True) +def fixed_clock(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(gateway, "datetime", _Clock) + token: Final = _PARENT_STATE.set("parent") + yield + _PARENT_STATE.reset(token) + + +def _call(*, structured: bool = False) -> gateway._NodeCall: + node: Final = ChatCompletionNode.model_validate( + { + "id": "answer", + "type": "chat_completion", + "model": "workflow-model", + "messages": {"from": "$input", "output": "messages"}, + "timeout_seconds": 60, + "retry": {"max_attempts": 2, "backoff_seconds": 0}, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "strict": True, + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + }, + } + if structured + else None, + } + ) + return gateway._NodeCall( + key_hash=_KEY, + team_id="team", + node=node, + messages=[{"role": "user", "content": "question"}], + workflow_id="workflow", + run_id="run", + session_id="session", + attempt=1, + deadline=_NOW + timedelta(minutes=2), + end_user_id="end-user", + ) + + +@dataclass(frozen=True, slots=True) +class _Gateway: + content: str | None = "answer" + failure: Exception | None = None + blocked_at: Literal["authorize", "process", "validate"] | None = None + response_body: Mapping[str, object] | None = None + receipt_ids: tuple[str, ...] = ("provider-id_cache_hit123",) + schema_error: ExecutionError | None = None + contexts: asyncio.Queue[gateway._RequestContext] = field(default_factory=asyncio.Queue) + schema_inputs: asyncio.Queue[JsonValue] = field(default_factory=asyncio.Queue) + started: asyncio.Event = field(default_factory=asyncio.Event) + stopped: asyncio.Event = field(default_factory=asyncio.Event) + processed: asyncio.Event = field(default_factory=asyncio.Event) + released: asyncio.Event = field(default_factory=asyncio.Event) + + async def _block(self, stage: str) -> None: + if self.blocked_at == stage: + self.started.set() + try: + await asyncio.Event().wait() + finally: + self.stopped.set() + + async def authorize(self, call: gateway._NodeCall, context: gateway._RequestContext) -> UserAPIKeyAuth: + assert _PARENT_STATE.get() == "isolated" + context.request.state.budget_reservation = {"reservation_id": context.request.headers["x-litellm-call-id"]} + self.contexts.put_nowait(context) + await self._block("authorize") + if self.failure is not None: + raise self.failure + return UserAPIKeyAuth(token=call.key_hash, team_id=call.team_id, end_user_id=call.end_user_id) + + async def process(self, context: gateway._RequestContext, auth: UserAPIKeyAuth) -> object: + self.processed.set() + await self._block("process") + context.response.headers["x-litellm-response-cost"] = "0.125" + return ( + self.response_body + if self.response_body is not None + else { + "id": "provider-id", + "choices": [{"message": {"role": "assistant", "content": self.content}}], + "usage": {"prompt_tokens": 7, "completion_tokens": 3}, + } + ) + + async def receipts(self, call: gateway._NodeCall, call_id: str) -> tuple[str, ...]: + return self.receipt_ids + + async def validate_output( + self, output: JsonValue, schema: Mapping[str, JsonValue], deadline: float + ) -> ExecutionError | None: + self.schema_inputs.put_nowait(output) + await self._block("validate") + return self.schema_error + + async def release(self, context: gateway._RequestContext) -> None: + self.released.set() + + +@pytest.mark.asyncio +async def test_attempts_keep_identity_and_unique_call_contexts_without_inheriting_parent() -> None: + port: Final = _Gateway() + call: Final = _call() + first: Final = await gateway._execute_isolated(call, port) + second: Final = await gateway._execute_isolated(replace(call, attempt=2), port) + assert first == second == ChildCallResult("answer", 7, 3, "provider-id_cache_hit123", 0.125) + assert _PARENT_STATE.get() == "parent" + first_context: Final = port.contexts.get_nowait() + second_context: Final = port.contexts.get_nowait() + assert first_context.request.state is not second_context.request.state + assert first_context.request.headers["x-litellm-call-id"] == workflow_call_id("run", "answer", 1) + assert second_context.request.headers["x-litellm-call-id"] == workflow_call_id("run", "answer", 2) + assert first_context.request.headers["x-litellm-call-id"] != second_context.request.headers["x-litellm-call-id"] + assert first_context.data["metadata"] == { + "workflow_id": "workflow", + "run_id": "run", + "node_id": "answer", + "attempt": 1, + "session_id": "session", + "spend_logs_metadata": { + "workflow_id": "workflow", + "run_id": "run", + "node_id": "answer", + "attempt": 1, + "session_id": "session", + }, + } + assert second_context.data["metadata"] == { + "workflow_id": "workflow", + "run_id": "run", + "node_id": "answer", + "attempt": 2, + "session_id": "session", + "spend_logs_metadata": { + "workflow_id": "workflow", + "run_id": "run", + "node_id": "answer", + "attempt": 2, + "session_id": "session", + }, + } + assert first_context.data["user"] == "end-user" + assert first_context.data["stream"] is False + assert first_context.data["messages"] == call.messages + assert call.messages == [{"role": "user", "content": "question"}] + assert port.released.is_set() + + +@pytest.mark.asyncio +async def test_recovered_attempt_uses_persisted_call_id_with_fresh_request_state() -> None: + port: Final = _Gateway(receipt_ids=()) + call: Final = _call() + first: Final = await gateway._execute_isolated(call, port) + recovered: Final = await gateway._execute_isolated(call, port) + assert first == recovered == ChildCallResult("answer", 7, 3, None, 0.125) + first_context: Final = port.contexts.get_nowait() + recovered_context: Final = port.contexts.get_nowait() + assert first_context.request.headers["x-litellm-call-id"] == recovered_context.request.headers["x-litellm-call-id"] + assert first_context.request.state is not recovered_context.request.state + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status,code", [(403, "authorization_failed"), (429, "rate_limit"), (503, "provider_error")]) +async def test_policy_denial_prevents_dispatch_and_redacts_failure(status: int, code: str) -> None: + port: Final = _Gateway(failure=HTTPException(status_code=status, detail="private provider configuration")) + result: Final = await gateway._execute_isolated(_call(), port) + assert isinstance(result, ExecutionError) + assert result.code == code + assert "private" not in result.message + assert not port.processed.is_set() + assert port.released.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["authorize", "process"]) +@pytest.mark.parametrize("expire", [False, True]) +async def test_cancellation_and_deadline_drain_auth_or_processing( + monkeypatch: pytest.MonkeyPatch, stage: Literal["authorize", "process"], expire: bool +) -> None: + port: Final = _Gateway(blocked_at=stage) + task: Final = asyncio.create_task(gateway._execute_isolated(_call(), port)) + await port.started.wait() + if expire: + loop: Final = asyncio.get_running_loop() + original_clock: Final = loop.time + monkeypatch.setattr(loop, "time", lambda: original_clock() + 3600) + result: Final = await task + assert result == ExecutionError("timeout", "The node deadline expired") + else: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert port.stopped.is_set() + assert port.released.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [(), ("first", "second")]) +async def test_missing_or_ambiguous_spend_receipt_never_substitutes_provider_id(ids: tuple[str, ...]) -> None: + result: Final = await gateway._execute_isolated(_call(), _Gateway(receipt_ids=ids)) + assert result == ChildCallResult("answer", 7, 3, None, 0.125) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid", [False, True]) +async def test_structured_result_is_parsed_and_schema_validation_is_required(invalid: bool) -> None: + error: Final = ExecutionError("output_invalid", "Schema rejected") if invalid else None + port: Final = _Gateway(content='{"answer":"structured"}', schema_error=error) + result: Final = await gateway._execute_isolated(_call(structured=True), port) + assert port.schema_inputs.get_nowait() == {"answer": "structured"} + assert result == ( + ExecutionError( + "output_invalid", + "Schema rejected", + receipt=ChildCallResult(None, 7, 3, "provider-id_cache_hit123", 0.125), + ) + if invalid + else ChildCallResult({"answer": "structured"}, 7, 3, "provider-id_cache_hit123", 0.125) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content", ['{"value":NaN}', '{"value":Infinity}', "not json"]) +@pytest.mark.parametrize("ids", [(), ("provider-id_cache_hit123",), ("first", "second")]) +async def test_invalid_json_retains_billing_without_reaching_schema_validator( + content: str, ids: tuple[str, ...] +) -> None: + port: Final = _Gateway(content=content, receipt_ids=ids) + result: Final = await gateway._execute_isolated(_call(structured=True), port) + assert result == ExecutionError( + "output_invalid", + "The model did not return valid JSON", + receipt=ChildCallResult(None, 7, 3, ids[0] if len(ids) == 1 else None, 0.125), + ) + assert port.schema_inputs.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content,code,message", + [ + (None, "invalid_response", "The model returned no assistant content"), + ("x" * (gateway._MAX_OUTPUT_BYTES + 1), "output_too_large", "The model output exceeded the workflow limit"), + ], + ids=["missing-content", "oversized-content"], +) +async def test_unusable_content_retains_billing_without_retaining_content( + content: str | None, code: str, message: str +) -> None: + result: Final = await gateway._execute_isolated(_call(), _Gateway(content=content)) + assert result == ExecutionError( + code, message, receipt=ChildCallResult(None, 7, 3, "provider-id_cache_hit123", 0.125) + ) + + +@pytest.mark.asyncio +async def test_malformed_chat_choices_retain_valid_reported_usage() -> None: + port: Final = _Gateway(response_body={"choices": [], "usage": {"prompt_tokens": 7, "completion_tokens": 3}}) + result: Final = await gateway._execute_isolated(_call(), port) + assert result == ExecutionError( + "invalid_response", + "The gateway returned an invalid chat response", + receipt=ChildCallResult(None, 7, 3, "provider-id_cache_hit123", 0.125), + ) + + +@pytest.mark.asyncio +async def test_schema_validation_deadline_retains_billing_for_delayed_spend_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + port: Final = _Gateway(content='{"answer":"structured"}', blocked_at="validate", receipt_ids=()) + task: Final = asyncio.create_task(gateway._execute_isolated(_call(structured=True), port)) + await port.started.wait() + loop: Final = asyncio.get_running_loop() + original_clock: Final = loop.time + monkeypatch.setattr(loop, "time", lambda: original_clock() + 3600) + result: Final = await task + assert result == ExecutionError( + "timeout", "The node deadline expired", receipt=ChildCallResult(None, 7, 3, None, 0.125) + ) + assert port.stopped.is_set() + assert port.released.is_set() + + +@pytest.mark.asyncio +async def test_expired_deadline_and_unpersisted_key_are_rejected_before_auth() -> None: + port: Final = _Gateway() + call: Final = _call() + assert await gateway._execute_isolated(replace(call, deadline=_NOW), port) == ExecutionError( + "timeout", "The node deadline expired" + ) + assert await gateway._execute_isolated(replace(call, key_hash="sk-unpersisted"), port) == ExecutionError( + "authorization_failed", "A persisted team virtual key is required" + ) + assert port.contexts.empty() + + +@pytest.mark.asyncio +async def test_internal_receive_yields_body_once_then_waits_without_signalling_disconnect() -> None: + context: Final = gateway._internal_request({"model": "workflow-model"}, "call", "session") + assert await context.request.body() == b'{"model": "workflow-model"}' + pending: Final = asyncio.create_task(context.request.receive()) + await asyncio.sleep(0) + assert not pending.done() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + +class _SpendSink(CustomLogger): + def __init__(self, outcome: str) -> None: + super().__init__() + self.outcome: Final = outcome + self.output_received: Final = asyncio.Event() + self.failure_received: Final = asyncio.Event() + self.admitted: Final[asyncio.Queue[UserAPIKeyAuth]] = asyncio.Queue() + self.rows: Final[asyncio.Queue[tuple[str, Mapping[str, object]]]] = asyncio.Queue() + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: Mapping[str, object], call_type: str + ) -> None: + self.admitted.put_nowait(user_api_key_dict) + + async def async_post_call_success_hook( + self, data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + self.output_received.set() + if self.outcome in {"cancelled", "timeout"}: + await asyncio.Event().wait() + if self.outcome in {"rejected", "failure_timeout"}: + raise HTTPException(status_code=422, detail="Private guardrail detail") + assert isinstance(response, ModelResponse) + return response.model_copy( + update={ + "choices": [litellm.Choices(message=litellm.Message(content="guarded answer", role="assistant"))], + } + ) + + async def async_post_call_failure_hook( + self, + request_data: Mapping[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.failure_received.set() + if self.outcome == "failure_timeout": + await asyncio.Event().wait() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.rows.put_nowait( + ( + "success", + get_logging_payload( + kwargs=dict(kwargs), + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ), + ) + ) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.rows.put_nowait( + ( + "failure", + get_logging_payload( + kwargs=dict(kwargs), + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ), + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "rejected", "cancelled", "timeout", "failure_timeout"]) +async def test_full_processor_reserves_child_budget_runs_hooks_and_emits_one_correlated_spend_payload( + monkeypatch: pytest.MonkeyPatch, + outcome: str, +) -> None: + cache: Final = UserApiKeyCache() + key: Final = UserAPIKeyAuth( + token=_KEY, team_id="team", user_id="owner", max_budget=100.0, models=["workflow-model"] + ) + await cache_key_object(hashed_token=_KEY, user_api_key_obj=key, user_api_key_cache=cache, proxy_logging_obj=None) + await cache.async_set_cache(key="team_id:team", value=LiteLLM_TeamTableCachedObj(team_id="team")) + await cache.async_set_cache(key="owner", value=LiteLLM_UserTable(user_id="owner")) + await cache.async_set_cache( + key=team_membership_auth_cache_key(team_id="team", user_id="owner"), + value=LiteLLM_TeamMembership(user_id="owner", team_id="team", spend=0), + ) + await cache.async_set_cache( + key=team_membership_reservation_cache_key(team_id="team", user_id="owner"), + value=LiteLLM_TeamMembership(user_id="owner", team_id="team", spend=0), + ) + await cache.async_set_cache( + key=end_user_cache_key("end-user"), value=LiteLLM_EndUserTable(user_id="end-user", blocked=False) + ) + sink: Final = _SpendSink(outcome) + counters: Final = FakeSpendCounters(redis_default=0.0) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "workflow-model", + "litellm_params": { + "model": "openai/workflow-fixture", + "api_key": "unit-only", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "mock_response": ModelResponse( + id="provider-fixture", + choices=[{"message": {"role": "assistant", "content": "unfiltered answer"}}], + usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, + ), + }, + } + ], + num_retries=0, + ) + monkeypatch.setattr(litellm, "callbacks", [sink]) + monkeypatch.setattr(litellm, "max_budget", 0.0) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "spend_counter_store", counters.store()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace()) + monkeypatch.setattr(proxy_server, "general_settings", {"cancel_on_disconnect": True}) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig()) + call: Final = _call() + task: Final = asyncio.create_task( + gateway._execute_isolated( + replace(call, node=call.node.model_copy(update={"max_tokens": 3})), gateway._ProxyGateway() + ) + ) + result: Final[ChildCallResult | ExecutionError | None] + if outcome == "cancelled": + await asyncio.wait_for(sink.output_received.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + result = None + elif outcome in {"timeout", "failure_timeout"}: + await asyncio.wait_for( + (sink.failure_received if outcome == "failure_timeout" else sink.output_received).wait(), timeout=5 + ) + loop: Final = asyncio.get_running_loop() + original_clock: Final = loop.time + monkeypatch.setattr(loop, "time", lambda: original_clock() + 3600) + result = await task + else: + result = await task + assert sink.admitted.qsize() == 1, result + admitted: Final = await asyncio.wait_for(sink.admitted.get(), timeout=5) + logged_status, row = await asyncio.wait_for(sink.rows.get(), timeout=5) + assert admitted.budget_reservation is not None + assert (admitted.user_id, admitted.team_id, admitted.end_user_id) == ("owner", "team", "end-user") + assert row["session_id"] == "session" and row["api_key"] == _KEY and row["team_id"] == "team" + assert row["end_user"] == "end-user" and row["user"] == "owner" + assert row["litellm_call_id"] == workflow_call_id("run", "answer", 1) + assert row["spend"] == pytest.approx(7 * 0.001 + 3 * 0.002) + assert isinstance(row["metadata"], str) + assert json.loads(row["metadata"])["spend_logs_metadata"] == { + "workflow_id": "workflow", + "run_id": "run", + "node_id": "answer", + "attempt": 1, + "session_id": "session", + } + assert sink.rows.empty() + if outcome != "success": + assert logged_status == "failure" + if outcome == "cancelled": + assert result is None + else: + assert isinstance(result, ExecutionError) + assert result.code == ("request_rejected" if outcome == "rejected" else "timeout") + assert result.receipt == ChildCallResult(None, 7, 3, None, row["spend"]) + else: + assert logged_status == "success" + assert isinstance(result, ChildCallResult) + assert (result.output, result.prompt_tokens, result.completion_tokens) == ("guarded answer", 7, 3) + assert result.response_cost == pytest.approx(row["spend"]) diff --git a/tests/unit/proxy/workflows/test_runner.py b/tests/unit/proxy/workflows/test_runner.py new file mode 100644 index 00000000000..12989251219 --- /dev/null +++ b/tests/unit/proxy/workflows/test_runner.py @@ -0,0 +1,347 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Final + +import pytest +from pydantic import JsonValue + +from litellm.models.workflow_definition import ( + ChatCompletionNode, + ReferenceBinding, + ValueBinding, + WorkflowBundle, + WorkflowDraft, + WorkflowRetry, + WorkflowSnapshot, +) +from litellm.models.workflow_execution import ( + ChildCallResult, + EngineRun, + ExecutionError, + RunState, + StepTrace, + workflow_call_id, +) +from litellm.proxy.workflows.runner import SchemaValidator, WorkflowRunner +from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository +from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository +from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle +from tests.unit.repositories.test_workflow_execution_repository import NOW, ExecutionRows, pending_run + + +class FixedClock: + def __init__(self) -> None: + self.value = NOW + + def __call__(self) -> datetime: + return self.value + + +class ValidationReply: + def __init__(self, error: ExecutionError | None = None) -> None: + self.error: Final = error + self.calls: tuple[tuple[JsonValue, Mapping[str, JsonValue]], ...] = () + + async def __call__(self, messages: JsonValue, schema: Mapping[str, JsonValue]) -> ExecutionError | None: + self.calls = (*self.calls, (messages, schema)) + return self.error + + +@dataclass(frozen=True, slots=True) +class ChatCall: + key_hash: str + team_id: str + workflow_id: str + node_id: str + messages: JsonValue + run_id: str + session_id: str + attempt: int + deadline: datetime + end_user_id: str | None + + +class ChatReplies: + def __init__( + self, + replies: tuple[ChildCallResult | ExecutionError | Exception, ...], + clock: FixedClock, + elapsed: timedelta = timedelta(0), + database: ExecutionRows | None = None, + ) -> None: + self.replies: Final = replies + self.clock: Final = clock + self.elapsed: Final = elapsed + self.database: Final = database + self.calls: tuple[ChatCall, ...] = () + self.before_dispatch: tuple[StepTrace, ...] = () + + async def __call__( + self, + *, + key_hash: str, + team_id: str, + workflow_id: str, + node: ChatCompletionNode, + messages: JsonValue, + run_id: str, + session_id: str, + attempt: int, + deadline: datetime, + end_user_id: str | None, + ) -> ChildCallResult | ExecutionError: + reply: Final = self.replies[len(self.calls)] + if self.database is not None: + self.before_dispatch = (*self.before_dispatch, self.database.events[-1]) + self.calls = ( + *self.calls, + ChatCall( + key_hash, team_id, workflow_id, node.id, messages, run_id, session_id, attempt, deadline, end_user_id + ), + ) + self.clock.value += self.elapsed + if isinstance(reply, Exception): + raise reply + return reply + + +def two_nodes(*, attempts: int = 1, timeout: float = 30) -> WorkflowBundle: + first: Final = ChatCompletionNode.model_validate(bundle().nodes[0].model_dump()).model_copy( + update={"retry": WorkflowRetry(max_attempts=attempts), "timeout_seconds": timeout} + ) + second: Final = first.model_copy( + update={"id": "node_2", "messages": ReferenceBinding.model_validate({"from": first.id, "output": "result"})} + ) + return bundle().model_copy( + update={ + "nodes": (first, second), + "output": ReferenceBinding.model_validate({"from": second.id, "output": "result"}), + } + ) + + +async def runner_inputs( + definition: WorkflowBundle, *, deadline: datetime = NOW + timedelta(seconds=90) +) -> tuple[WorkflowDefinitionRepository, ExecutionRows, EngineRun]: + definitions: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await definitions.create(definition, "team", "author") + assert isinstance(draft, WorkflowDraft) + snapshot: Final = await definitions.freeze_draft(draft.workflow_id, "team", draft.etag) + assert isinstance(snapshot, WorkflowSnapshot) + run: Final = pending_run(snapshot_id=snapshot.snapshot_id, workflow_id=snapshot.workflow_id, deadline=deadline) + return definitions, ExecutionRows((run,)), run + + +async def run_once( + repository: WorkflowExecutionRepository, + definitions: WorkflowDefinitionRepository, + chat: ChatReplies, + clock: FixedClock, + run_id: str, + *, + validate_input: SchemaValidator | None = None, +) -> None: + runner: Final = WorkflowRunner( + repository, + definitions, + chat, + clock=clock, + validate_input=ValidationReply() if validate_input is None else validate_input, + ) + try: + await runner._run(run_id) # pyright: ignore[reportPrivateUsage] # execute one claim without timed polling + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_sequential_nodes_preserve_json_values_and_record_usage_and_final_output() -> None: + definitions, database, run = await runner_inputs(two_nodes()) + clock: Final = FixedClock() + forwarded: Final[JsonValue] = [{"role": "user", "content": [{"type": "text", "text": "next"}]}] + output: Final[JsonValue] = {"count": 3, "enabled": False, "missing": None, "ratios": [0.5, 1]} + chat: Final = ChatReplies( + (ChildCallResult(forwarded, 7, 2, "request-1", 0.01), ChildCallResult(output, 11, 3, "request-2", 0.02)), + clock, + ) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + finished: Final = await repository.get(run.run_id) + assert finished is not None + assert (finished.status, finished.output, finished.lease_token) == (RunState.COMPLETED, output, None) + assert tuple((call.node_id, call.messages) for call in chat.calls) == ( + ("node_1", list(run.input.messages)), + ("node_2", forwarded), + ) + steps: Final = await repository.steps(run.run_id) + assert tuple( + (step.node_id, step.state, step.input, step.output, step.prompt_tokens, step.completion_tokens, step.request_id) + for step in steps + ) == ( + ("node_1", RunState.COMPLETED, list(run.input.messages), forwarded, 7, 2, "request-1"), + ("node_2", RunState.COMPLETED, forwarded, output, 11, 3, "request-2"), + ) + assert tuple(step.response_cost for step in steps) == (0.01, 0.02) + assert tuple((step.node_id, step.state) for step in database.events) == ( + ("node_1", RunState.RUNNING), + ("node_1", RunState.COMPLETED), + ("node_2", RunState.RUNNING), + ("node_2", RunState.COMPLETED), + ) + + +@pytest.mark.asyncio +async def test_first_nonretryable_error_stops_downstream_and_keeps_failure_trace() -> None: + original: Final = two_nodes(attempts=3) + independent: Final = original.nodes[1].model_copy( + update={"messages": ValueBinding(value=[{"role": "user", "content": "independent downstream"}])} + ) + definitions, database, run = await runner_inputs( + original.model_copy(update={"nodes": (original.nodes[0], independent)}) + ) + clock: Final = FixedClock() + chat: Final = ChatReplies( + (ExecutionError("invalid_input", "Expected messages"), ChildCallResult("must not execute", 0, 0, None, None)), + clock, + ) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + finished: Final = await repository.get(run.run_id) + assert finished is not None + assert (finished.status, finished.output, finished.metadata) == ( + RunState.FAILED, + None, + {"error": "Expected messages"}, + ) + assert tuple((call.node_id, call.attempt) for call in chat.calls) == (("node_1", 1),) + steps: Final = await repository.steps(run.run_id) + assert tuple((step.node_id, step.state, step.attempts, step.error) for step in steps) == ( + ("node_1", RunState.FAILED, 1, "Expected messages"), + ) + + +@pytest.mark.asyncio +async def test_retry_preserves_run_session_step_and_absolute_deadline() -> None: + definitions, database, run = await runner_inputs( + two_nodes(attempts=2, timeout=60), deadline=NOW + timedelta(seconds=20) + ) + clock: Final = FixedClock() + chat: Final = ChatReplies( + ( + ExecutionError( + "provider_error", "Retry later", receipt=ChildCallResult(None, 4, 2, "failed-receipt", 0.125) + ), + ChildCallResult([{"role": "user", "content": "continue"}], 2, 1, "retried", 0.25), + ChildCallResult("done", 3, 1, "last", None), + ), + clock, + elapsed=timedelta(seconds=3), + database=database, + ) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + assert chat.calls[:2] == tuple( + ChatCall( + run.created_by, + run.team_id, + run.workflow_id, + "node_1", + [dict(message) for message in run.input.messages], + run.run_id, + run.session_id, + attempt, + run.input.deadline, + run.input.end_user_id, + ) + for attempt in (1, 2) + ) + first_steps: Final = tuple(step for step in database.events if step.node_id == "node_1") + assert len(first_steps) == 4 + assert len({step.step_id for step in first_steps}) == 1 + assert tuple(step.attempts for step in first_steps) == (1, 1, 2, 2) + assert {step.deadline for step in first_steps} == {run.input.deadline} + assert {step.started_at for step in first_steps} == {NOW} + assert tuple((step.state, step.attempts, step.call_id) for step in chat.before_dispatch[:2]) == tuple( + (RunState.RUNNING, attempt, workflow_call_id(run.run_id, "node_1", attempt)) for attempt in (1, 2) + ) + assert len({step.call_id for step in first_steps}) == 2 + traces: Final = await repository.steps(run.run_id) + retried: Final = next(step for step in traces if step.node_id == "node_1") + assert (retried.prompt_tokens, retried.completion_tokens, retried.response_cost, retried.request_ids) == ( + 6, + 3, + 0.375, + ("failed-receipt", "retried"), + ) + finished: Final = await repository.get(run.run_id) + assert finished is not None and finished.status is RunState.COMPLETED and finished.output == "done" + + +@pytest.mark.asyncio +async def test_retry_budget_cannot_extend_node_deadline_or_dispatch_downstream() -> None: + definitions, database, run = await runner_inputs(two_nodes(attempts=3, timeout=5)) + clock: Final = FixedClock() + chat: Final = ChatReplies( + (ExecutionError("provider_error", "Unavailable"), ChildCallResult("must not execute", 0, 0, None, None)), + clock, + elapsed=timedelta(seconds=5), + ) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + assert tuple((call.node_id, call.attempt, call.deadline) for call in chat.calls) == ( + ("node_1", 1, NOW + timedelta(seconds=5)), + ) + finished: Final = await repository.get(run.run_id) + assert finished is not None and finished.status is RunState.FAILED and finished.output is None + + +@pytest.mark.asyncio +async def test_expired_workflow_never_dispatches_or_records_a_step() -> None: + definitions, database, run = await runner_inputs(two_nodes(), deadline=NOW) + clock: Final = FixedClock() + chat: Final = ChatReplies((ChildCallResult("must not execute", 0, 0, None, None),), clock) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + assert chat.calls == () and database.events == () + finished: Final = await repository.get(run.run_id) + assert finished is not None + assert (finished.status, finished.metadata) == (RunState.FAILED, {"error": "Workflow deadline expired"}) + + +@pytest.mark.asyncio +async def test_unexpected_executor_exception_records_failed_node_and_stops_downstream() -> None: + definitions, database, run = await runner_inputs(two_nodes()) + clock: Final = FixedClock() + failure: Final = RuntimeError("injected executor failure") + chat: Final = ChatReplies((failure,), clock) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id) + finished: Final = await repository.get(run.run_id) + assert finished is not None and finished.status is RunState.FAILED + assert tuple(call.node_id for call in chat.calls) == ("node_1",) + traces: Final = await repository.steps(run.run_id) + assert tuple((step.node_id, step.state, step.output, step.error) for step in traces) == ( + ("node_1", RunState.FAILED, None, "The node could not complete"), + ) + assert finished.metadata == {"error": traces[0].error} + + +@pytest.mark.asyncio +async def test_pinned_input_schema_rejection_fails_run_before_dispatch_or_checkpoint() -> None: + definition: Final = two_nodes().model_copy(update={"input_schema": {"type": "array", "maxItems": 0}}) + definitions, database, run = await runner_inputs(definition) + clock: Final = FixedClock() + chat: Final = ChatReplies((ChildCallResult("must not execute", 0, 0, None, None),), clock) + validation: Final = ValidationReply(ExecutionError("invalid_input", "Pinned input schema rejected messages")) + repository: Final = WorkflowExecutionRepository(database) + await run_once(repository, definitions, chat, clock, run.run_id, validate_input=validation) + assert validation.calls == ((list(run.input.messages), definition.input_schema),) + assert chat.calls == () and database.events == () + finished: Final = await repository.get(run.run_id) + assert finished is not None + assert (finished.status, finished.output, finished.metadata) == ( + RunState.FAILED, + None, + {"error": "Pinned input schema rejected messages"}, + ) diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 56843640888..1c410313bd6 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -3,6 +3,8 @@ Tests for gateway repository layer. """ import json +from collections.abc import Mapping +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, Dict, Final, List, Optional @@ -2242,6 +2244,39 @@ class TestTeamRepositoryArchiveData: class TestPrismaTableRepository: + @pytest.mark.asyncio + @pytest.mark.parametrize("changed_field", ["litellm_call_id", "api_key", "team_id", "session_id"]) + @pytest.mark.parametrize("row_status", ["success", "failure"]) + async def test_receipt_lookup_selects_only_the_exact_child_call( + self, changed_field: str, row_status: str + ) -> None: + from litellm.repositories.table_repositories import SpendLogsRepository + + @dataclass(frozen=True, slots=True) + class Receipt: + request_id: str + litellm_call_id: str = "child-call" + api_key: str = "key-hash" + team_id: str = "team" + session_id: str = "session" + status: str = "success" + + intended: Final = Receipt("response_cache_hit123", status=row_status) + unrelated: Final = replace(intended, request_id="unrelated", **{changed_field: "different"}) + + @dataclass(frozen=True, slots=True) + class Rows: + async def find_many(self, *, where: Mapping[str, object], take: int) -> tuple[Receipt, ...]: + return tuple( + row for row in (unrelated, intended) if all(getattr(row, key) == value for key, value in where.items()) + )[:take] + + client: Final = SimpleNamespace(writer_db=SimpleNamespace(litellm_spendlogs=Rows())) + result: Final = await SpendLogsRepository(client, use_writer=True).request_ids_for_call( + call_id="child-call", key_hash="key-hash", team_id="team", session_id="session", + ) + assert result == ("response_cache_hit123",) + def test_table_property_returns_named_delegate(self): from litellm.proxy.common_utils.config_sync_pubsub import ( _PublishOnWriteActions, diff --git a/tests/unit/repositories/test_workflow_definition_repository.py b/tests/unit/repositories/test_workflow_definition_repository.py new file mode 100644 index 00000000000..79606e12d13 --- /dev/null +++ b/tests/unit/repositories/test_workflow_definition_repository.py @@ -0,0 +1,71 @@ +import asyncio +from typing import Final + +import pytest + +from litellm.models.workflow_definition import DefinitionError, WorkflowDraft, WorkflowSnapshot +from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository +from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle + + +@pytest.mark.asyncio +async def test_edit_and_republish_leave_old_snapshot_and_frozen_draft_unchanged() -> None: + store: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await store.create(bundle(), "team", "author") + assert isinstance(draft, WorkflowDraft) + published: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") + assert isinstance(published, WorkflowSnapshot) + frozen: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag) + assert isinstance(frozen, WorkflowSnapshot) + assert frozen.published_by is None + updated: Final = await store.replace(draft.workflow_id, "team", draft.etag, bundle(model="second"), "other") + assert isinstance(updated, WorkflowDraft) + assert updated.published_snapshot_id == published.snapshot_id + assert await store.get_published("example", "team") == published + second: Final = await store.publish(updated.workflow_id, "team", updated.etag, updated.digest, "reviewer") + assert isinstance(second, WorkflowSnapshot) + assert second.snapshot_id != published.snapshot_id + assert second.contributors == ("other",) + assert await store.get_snapshot(published.snapshot_id) == published + assert await store.get_snapshot(frozen.snapshot_id) == frozen + assert await store.get_published("example", "team") == second + assert await store.get_snapshot(second.snapshot_id, "other-team") is None + + +@pytest.mark.asyncio +async def test_two_edits_of_same_revision_accept_only_one_and_preserve_contributors() -> None: + store: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await store.create(bundle(), "team", "author") + assert isinstance(draft, WorkflowDraft) + edits: Final = await asyncio.gather( + store.replace(draft.workflow_id, "team", draft.etag, bundle(model="one"), "editor-one"), + store.replace(draft.workflow_id, "team", draft.etag, bundle(model="two"), "editor-two"), + ) + accepted: Final = tuple(edit for edit in edits if isinstance(edit, WorkflowDraft)) + rejected: Final = tuple(edit for edit in edits if isinstance(edit, DefinitionError)) + assert len(accepted) == 1 + assert rejected == ( + DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match"), + ) + stored: Final = await store.get_draft(draft.workflow_id) + assert stored is not None + assert stored.draft == accepted[0] + assert "author" in stored.contributors + assert len(stored.contributors) == 2 + + +@pytest.mark.asyncio +async def test_publish_and_edit_race_never_snapshots_unreviewed_content() -> None: + store: Final = WorkflowDefinitionRepository(DefinitionRows().database) + draft: Final = await store.create(bundle(), "team", "author") + assert isinstance(draft, WorkflowDraft) + published, edited = await asyncio.gather( + store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer"), + store.replace(draft.workflow_id, "team", draft.etag, bundle(model="changed"), "other"), + ) + assert isinstance(published, WorkflowSnapshot) + assert isinstance(edited, WorkflowDraft) + assert published.bundle == draft.bundle + assert published.digest == draft.digest + assert edited.digest != published.digest + assert await store.get_published("example", "team") == published diff --git a/tests/unit/repositories/test_workflow_execution_repository.py b/tests/unit/repositories/test_workflow_execution_repository.py new file mode 100644 index 00000000000..4665913ebc3 --- /dev/null +++ b/tests/unit/repositories/test_workflow_execution_repository.py @@ -0,0 +1,428 @@ +import re +from collections.abc import AsyncGenerator, Mapping +from contextlib import asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.models.workflow_execution import ( + EngineRun, + ExecutionError, + RunAdmission, + RunInput, + RunPage, + RunState, + StepTrace, +) +from litellm.repositories.workflow_execution_repository import ExecutionQueries, WorkflowExecutionRepository + +NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_STRING: Final = TypeAdapter(str) +_INTEGER: Final = TypeAdapter(int) +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_TRACE: Final = TypeAdapter(dict[str, StepTrace]) + + +def pending_run( + run_id: str = "run-1", + *, + snapshot_id: str = "snapshot-1", + workflow_id: str = "workflow-1", + key_hash: str = "owner-key", + team_id: str = "team", + deadline: datetime = NOW + timedelta(seconds=90), +) -> EngineRun: + return EngineRun( + run_id=run_id, + session_id=f"session-{run_id}", + workflow_id=workflow_id, + snapshot_id=snapshot_id, + team_id=team_id, + created_by=key_hash, + status=RunState.PENDING, + created_at=NOW, + updated_at=NOW, + input=RunInput( + messages=({"role": "user", "content": "plan"},), + end_user_id="end-user", + deadline=deadline, + draft=True, + ), + request_digest="request-digest", + ) + + +@dataclass(frozen=True, slots=True) +class Query: + sql: str + args: tuple[object, ...] + in_transaction: bool + + +class ExecutionRows: + def __init__(self, runs: tuple[EngineRun, ...] = ()) -> None: + self.runs = runs + self.keys: tuple[tuple[str, str | None], ...] = () + self.events: tuple[StepTrace, ...] = () + self.queries: tuple[Query, ...] = () + self.in_transaction = False + + @asynccontextmanager + async def tx(self) -> AsyncGenerator[ExecutionQueries]: + self.in_transaction = True + try: + yield self + finally: + self.in_transaction = False + + def _rows(self, runs: tuple[EngineRun, ...]) -> tuple[Mapping[str, object], ...]: + return tuple(run.model_dump() for run in runs) + + def _replace(self, run: EngineRun) -> tuple[Mapping[str, object], ...]: + self.runs = tuple(run if current.run_id == run.run_id else current for current in self.runs) + return self._rows((run,)) + + async def query_raw(self, query: str, *args: object) -> object: + self.queries = (*self.queries, Query(query, args, self.in_transaction)) + parameters: Final = frozenset(int(match.group(1)) for match in re.finditer(r"\$(\d+)", query)) + assert parameters == frozenset(range(1, len(args) + 1)), (query, args) + if "pg_advisory_xact_lock" in query: + assert query.startswith("SELECT true AS acquired FROM "), "Prisma cannot decode a PostgreSQL void column" + return ({"acquired": True},) + if "idempotency_key =" in query: + owner, team, workflow, key = args if "created_by = $1" in query else (args[2], args[1], args[0], args[3]) + return self._rows( + tuple( + run + for run in self.runs + if key is not None + and (run.created_by, run.team_id, run.workflow_id) == (owner, team, workflow) + and (run.run_id, key) in self.keys + ) + ) + if query.startswith('INSERT INTO "LiteLLM_WorkflowRun"'): + run: Final = EngineRun.model_validate( + { + "run_id": args[0], + "session_id": args[1], + "created_by": args[2], + "workflow_id": args[3], + "snapshot_id": args[4], + "team_id": args[5], + "input": RunInput.model_validate_json(_STRING.validate_python(args[6])), + "request_digest": args[7], + "status": RunState.PENDING, + "created_at": NOW, + "updated_at": NOW, + "output": None, + "metadata": None, + "lease_token": None, + "lease_expires_at": None, + } + ) + self.runs = (*self.runs, run) + self.keys = (*self.keys, (run.run_id, None if args[8] is None else _STRING.validate_python(args[8]))) + return self._rows((run,)) + if "LIMIT 64" in query: + return tuple( + {"run_id": run.run_id} + for run in self.runs + if run.team_id == args[0] and run.status in {RunState.PENDING, RunState.RUNNING} + )[:64] + if "ORDER BY created_at DESC, run_id DESC" in query: + boundary: Final = None if args[4] is None else datetime.fromisoformat(_STRING.validate_python(args[4])) + matching: Final = tuple( + run + for run in self.runs + if run.workflow_id == args[0] + and run.team_id == args[1] + and (args[2] is None or run.created_by == args[2]) + and (args[3] is None or run.status.value == args[3]) + and (boundary is None or (run.created_at, run.run_id) < (boundary, _STRING.validate_python(args[5]))) + ) + return self._rows( + tuple(sorted(matching, key=lambda run: (run.created_at, run.run_id), reverse=True))[ + : _INTEGER.validate_python(args[6]) + ] + ) + if query.startswith("SELECT trace FROM"): + latest: Final = {(step.node_id, step.attempts): step for step in self.events if step.run_id == args[0]} + return tuple( + {"trace": step.model_dump(mode="json")} + for step in sorted(latest.values(), key=lambda step: (step.started_at, step.node_id)) + ) + if query.startswith('INSERT INTO "LiteLLM_WorkflowEvent"'): + payload: Final = _TRACE.validate_json(_STRING.validate_python(args[3])) + self.events = (*self.events, payload["trace"]) + return ({"event_id": args[0]},) + assert query.startswith( + ( + 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id', + 'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id', + 'UPDATE "LiteLLM_WorkflowRun"', + ) + ), query + current: Final = next((run for run in self.runs if run.run_id == args[0]), None) + if current is None: + return () + if query.startswith('SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id'): + return self._rows((current,)) + if "SET status = 'running'" in query: + if current.status is not RunState.PENDING: + return () + return self._replace( + current.model_copy( + update={ + "status": RunState.RUNNING, + "lease_token": args[1], + "lease_expires_at": datetime.fromisoformat(_STRING.validate_python(args[2])), + } + ) + ) + if "SET cancel_requested = true" in query: + return self._replace(current.model_copy(update={"cancel_requested": True})) + owned: Final = ( + current.status is RunState.RUNNING + and current.lease_token == args[1] + and current.lease_expires_at is not None + and current.lease_expires_at > NOW + ) + if "FOR UPDATE" in query: + return ({"run_id": current.run_id},) if owned else () + if "SET status = $3" in query: + if not owned: + return () + return self._replace( + current.model_copy( + update={ + "status": RunState(args[2]), + "output": _JSON.validate_json(_STRING.validate_python(args[3])), + "metadata": {"error": args[4]}, + "lease_token": None, + "lease_expires_at": None, + } + ) + ) + raise AssertionError(f"Unexpected database operation: {query}") + + +@pytest.mark.asyncio +async def test_idempotent_admission_pins_snapshot_and_rejects_different_input() -> None: + database: Final = ExecutionRows() + repository: Final = WorkflowExecutionRepository(database) + original: Final = pending_run() + key: Final = "key'; SELECT secret" + first: Final = await repository.admit( + workflow_id=original.workflow_id, + snapshot_id="snapshot-1", + team_id=original.team_id, + key_hash=original.created_by, + input=original.input, + digest=original.request_digest, + idempotency_key=key, + ) + replay: Final = await repository.admit( + workflow_id=original.workflow_id, + snapshot_id="snapshot-2", + team_id=original.team_id, + key_hash=original.created_by, + input=original.input, + digest=original.request_digest, + idempotency_key=key, + ) + conflict: Final = await repository.admit( + workflow_id=original.workflow_id, + snapshot_id="snapshot-2", + team_id=original.team_id, + key_hash=original.created_by, + input=original.input, + digest="different", + idempotency_key=key, + ) + assert isinstance(first, RunAdmission) + assert replay == RunAdmission(first.run, False) + assert first.created and first.run.snapshot_id == "snapshot-1" + assert conflict == ExecutionError("idempotency_conflict", "Idempotency key was used with different input") + assert len(database.runs) == 1 + assert all(query.in_transaction for query in database.queries) + assert database.queries[0].args == (original.team_id,) + assert database.queries[0].sql.startswith("SELECT true AS acquired FROM pg_advisory_xact_lock(") + assert all("key'; SELECT secret" not in query.sql for query in database.queries) + insert: Final = next(query for query in database.queries if query.sql.startswith("INSERT")) + assert "timezone('UTC', NOW()), timezone('UTC', NOW())" in insert.sql + assert insert.args[2:6] == (original.created_by, original.workflow_id, "snapshot-1", original.team_id) + assert RunInput.model_validate_json(_STRING.validate_python(insert.args[6])) == original.input + assert insert.args[7:] == (original.request_digest, key) + + +@pytest.mark.asyncio +async def test_null_idempotency_and_nullable_database_columns_round_trip_without_coalescing_runs() -> None: + database: Final = ExecutionRows() + repository: Final = WorkflowExecutionRepository(database) + original: Final = pending_run() + first: Final = await repository.admit( + workflow_id=original.workflow_id, + snapshot_id=original.snapshot_id, + team_id=original.team_id, + key_hash=original.created_by, + input=original.input, + digest=original.request_digest, + idempotency_key=None, + ) + second: Final = await repository.admit( + workflow_id=original.workflow_id, + snapshot_id=original.snapshot_id, + team_id=original.team_id, + key_hash=original.created_by, + input=original.input, + digest=original.request_digest, + idempotency_key=None, + ) + assert isinstance(first, RunAdmission) and isinstance(second, RunAdmission) + assert first.created and second.created and first.run.run_id != second.run.run_id + assert (first.run.output, first.run.metadata, first.run.lease_token, first.run.lease_expires_at) == (None,) * 4 + lookup: Final = next(query for query in database.queries if "idempotency_key =" in query.sql) + assert lookup.args == (original.created_by, original.team_id, original.workflow_id, None) + assert "idempotency_key = $4::text" in lookup.sql + assert await repository.claim(first.run.run_id, "lease", NOW + timedelta(seconds=15)) is not None + claim: Final = database.queries[-1] + assert claim.args == (first.run.run_id, "lease", (NOW + timedelta(seconds=15)).isoformat()) + assert "lease_expires_at = timezone('UTC', $3::text::timestamptz)" in claim.sql + assert await repository.finish(first.run.run_id, "lease", RunState.COMPLETED, {"value": None}, None) + finish: Final = database.queries[-1] + assert finish.args == (first.run.run_id, "lease", "completed", '{"value": null}', None) + assert "$4::jsonb" in finish.sql and "$5::text" in finish.sql + finished: Final = await repository.get(first.run.run_id) + assert finished is not None + assert (finished.output, finished.metadata, finished.lease_token) == ({"value": None}, {"error": None}, None) + + +@pytest.mark.asyncio +async def test_history_cursor_uses_timestamp_and_id_with_lookahead_and_scoped_parameters() -> None: + database: Final = ExecutionRows( + ( + pending_run("run-c"), + pending_run("run-a"), + pending_run("run-b"), + pending_run("foreign-team", team_id="other-team"), + pending_run("foreign-workflow", workflow_id="other-workflow"), + pending_run("foreign-key", key_hash="other-key"), + ) + ) + repository: Final = WorkflowExecutionRepository(database) + first: Final = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, None) + assert isinstance(first, RunPage) + assert tuple(run.run_id for run in first.runs) == ("run-c", "run-b") + assert first.count == 2 and first.next_cursor is not None + assert database.queries[-1].args == ("workflow-1", "team", "owner-key", "pending", None, None, 3) + database.runs = (*database.runs, pending_run("run-d")) + tail: Final = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, first.next_cursor) + assert isinstance(tail, RunPage) + assert tuple(run.run_id for run in tail.runs) == ("run-a",) + assert tail.count == 1 and tail.next_cursor is None + query: Final = database.queries[-1] + assert query.args == ("workflow-1", "team", "owner-key", "pending", NOW.isoformat(), "run-b", 3) + assert "(created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)" in query.sql + assert "ORDER BY created_at DESC, run_id DESC LIMIT $7" in query.sql + assert "workflow_id = $1 AND team_id = $2" in query.sql + assert "created_by = $3::text" in query.sql and "status = $4::text" in query.sql + assert "engine_owned = true" in query.sql + + +@pytest.mark.asyncio +async def test_invalid_cursor_never_queries_and_unfiltered_history_passes_sql_nulls() -> None: + database: Final = ExecutionRows() + repository: Final = WorkflowExecutionRepository(database) + assert await repository.list_runs("workflow-1", "team", None, None, 1, "not-a-cursor") == ExecutionError( + "invalid_cursor", "Invalid run cursor" + ) + assert database.queries == () + assert await repository.list_runs("workflow-1", "team", None, None, 1, None) == RunPage( + runs=(), count=0, next_cursor=None + ) + query: Final = database.queries[-1] + assert query.args == ("workflow-1", "team", None, None, None, None, 2) + assert "($3::text IS NULL" in query.sql and "($4::text IS NULL" in query.sql and "($5::text IS NULL" in query.sql + + +@pytest.mark.asyncio +async def test_lost_lease_cannot_write_trace_or_terminal_result() -> None: + database: Final = ExecutionRows((pending_run(),)) + repository: Final = WorkflowExecutionRepository(database) + assert await repository.claim("run-1", "owner-lease", NOW + timedelta(seconds=15)) is not None + step: Final = StepTrace( + step_id="step-1", + run_id="run-1", + node_id="node-1", + state=RunState.RUNNING, + input=None, + attempts=1, + deadline=NOW + timedelta(seconds=10), + started_at=NOW, + ) + assert not await repository.checkpoint("run-1", "stale-lease", step) + assert not await repository.finish("run-1", "stale-lease", RunState.COMPLETED, "forged", None) + assert database.events == () + current: Final = await repository.get("run-1") + assert current is not None and current.status is RunState.RUNNING and current.output is None + ownership: Final = database.queries[1] + assert ownership.in_transaction and ownership.args == ("run-1", "stale-lease") + assert "lease_token = $2" in ownership.sql + assert "lease_expires_at > timezone('UTC', NOW()) FOR UPDATE" in ownership.sql + + +@pytest.mark.asyncio +async def test_attempt_receipts_aggregate_without_counting_duplicate_checkpoints_twice() -> None: + database: Final = ExecutionRows((pending_run(),)) + repository: Final = WorkflowExecutionRepository(database) + assert await repository.claim("run-1", "lease", NOW + timedelta(seconds=15)) is not None + first: Final = StepTrace( + step_id="step-1", + run_id="run-1", + node_id="node-1", + state=RunState.FAILED, + input=None, + attempts=1, + deadline=NOW + timedelta(seconds=10), + started_at=NOW, + finished_at=NOW, + prompt_tokens=4, + completion_tokens=2, + response_cost=0.125, + request_id="failed-receipt", + call_id="call-1", + error="Provider failed after billing", + ) + second: Final = first.model_copy( + update={ + "state": RunState.COMPLETED, + "attempts": 2, + "output": {"done": True}, + "error": None, + "prompt_tokens": 6, + "completion_tokens": 3, + "response_cost": 0.25, + "request_id": "successful-receipt", + "call_id": "call-2", + } + ) + for checkpoint in (first, first, second, second): + assert await repository.checkpoint("run-1", "lease", checkpoint) + traces: Final = await repository.steps("run-1") + assert traces == ( + second.model_copy( + update={ + "prompt_tokens": 10, + "completion_tokens": 5, + "response_cost": 0.375, + "request_ids": ("failed-receipt", "successful-receipt"), + } + ), + ) + assert len(database.events) == 4 + query: Final = database.queries[-1] + assert query.args == ("run-1",) + assert "DISTINCT ON (step_name, data->'trace'->>'attempts')" in query.sql diff --git a/tests/unit/test_check_openapi_docs.py b/tests/unit/test_check_openapi_docs.py index c6ea0cd936e..46b89035e26 100644 --- a/tests/unit/test_check_openapi_docs.py +++ b/tests/unit/test_check_openapi_docs.py @@ -1,11 +1,11 @@ """Tests for scripts/check_openapi_docs.py, the scan behind scripts/openapi_docs_gate.py.""" from collections.abc import Mapping +from pathlib import Path from typing import Final -import pytest - import check_openapi_docs as checker +import pytest _BODY: Final = {"content": {"application/json": {"schema": {"type": "object"}}}} _RESPONSES: Final = {"200": {"content": {"application/json": {"schema": {"$ref": "#/components/schemas/Out"}}}}} @@ -20,15 +20,47 @@ def _post_with_media_types(request_media_type: str, response_media_type: str) -> } -@pytest.mark.parametrize("request_type, response_type", [ - ("application/json", "application/json"), ("application/sdp", "application/sdp"), - ("multipart/form-data", "application/json"), -]) +@pytest.mark.parametrize( + "request_type, response_type", + [ + ("application/json", "application/json"), + ("application/sdp", "application/sdp"), + ("multipart/form-data", "application/json"), + ], +) def test_a_fully_documented_post_has_no_findings(request_type: str, response_type: str) -> None: operation: Final = _post_with_media_types(request_type, response_type) assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == () +@pytest.mark.parametrize("status", ["200", "201", "202", "206", "299"]) +def test_any_documented_success_response_satisfies_the_response_rule(status: str) -> None: + operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}} + assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == () + + +@pytest.mark.parametrize("status", ["default", "100", "300", "400", "500", "2000"]) +def test_only_success_statuses_can_supply_the_response_schema(status: str) -> None: + operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}} + assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ( + checker.Violation("POST", "/v1/thing", "OA003"), + ) + + +@pytest.mark.parametrize("status", ["201", "202"]) +def test_success_without_a_schema_is_not_documented_by_an_error_response(status: str) -> None: + operation: Final[Mapping[str, object]] = { + **_DOCUMENTED, + "responses": { + status: {"content": {"application/json": {"schema": {}}}}, + "422": _RESPONSES["200"], + }, + } + assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ( + checker.Violation("POST", "/v1/thing", "OA003"), + ) + + def test_a_schema_under_another_media_type_still_counts_as_undocumented() -> None: operation: Final = _post_with_media_types("application/octet-stream", "audio/wav") assert {v.code for v in checker.violations({"paths": {"/v1/thing": {"post": operation}}})} == {"OA001", "OA003"} @@ -44,8 +76,11 @@ def test_each_missing_piece_is_its_own_rule() -> None: def test_only_the_missing_piece_is_reported(body: Mapping[str, object]) -> None: no_body: Final = {**_DOCUMENTED, "requestBody": body} blank_description: Final = {**_DOCUMENTED, "description": " "} - untyped_200: Final = {**_DOCUMENTED, "responses": {"200": {"content": {"application/json": {"schema": {}}}}}} - spec: Final = { + untyped_200: Final[Mapping[str, object]] = { + **_DOCUMENTED, + "responses": {"200": {"content": {"application/json": {"schema": {}}}}}, + } + spec: Final[Mapping[str, object]] = { "paths": { "/a": {"post": no_body}, "/b": {"post": blank_description}, @@ -65,11 +100,13 @@ def test_non_post_operations_are_not_gated() -> None: def test_findings_are_sorted_by_path_so_counts_are_stable() -> None: - spec: Final = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}} + spec: Final[Mapping[str, object]] = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}} assert [v.path for v in checker.violations(spec)] == ["/a", "/a", "/a", "/z", "/z", "/z"] -def test_main_reads_a_spec_file_and_prints_one_line_per_finding(tmp_path, capsys) -> None: +def test_main_reads_a_spec_file_and_prints_one_line_per_finding( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: spec_file: Final = tmp_path / "openapi.json" spec_file.write_text('{"paths": {"/v1/thing": {"post": {"description": "x", "requestBody": {}}}}}') assert checker.main([str(spec_file)]) == 0 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2272b4bea43..0555cc95b1f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24080,6 +24080,30 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/workflows": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Workflow Definitions + * @description List readable team drafts in immutable workflow ID order. + */ + get: operations["list_workflow_definitions_v1_workflows_get"]; + put?: never; + /** + * Register Workflow + * @description Create a team draft from a sequential LiteFlow v1 bundle. + */ + post: operations["register_workflow_v1_workflows_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/workflows/runs": { parameters: { query?: never; @@ -24133,6 +24157,46 @@ export interface paths { patch: operations["update_workflow_run_v1_workflows_runs__run_id__patch"]; trace?: never; }; + "/v1/workflows/runs/{run_id}/cancel": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Cancel Workflow Run + * @description Request cancellation and return the run ID to poll until execution stops. + */ + post: operations["cancel_workflow_run_v1_workflows_runs__run_id__cancel_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/runs/{run_id}/definition": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Run Definition + * @description Read the immutable workflow definition captured when this run was admitted. + */ + get: operations["get_run_definition_v1_workflows_runs__run_id__definition_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/workflows/runs/{run_id}/events": { parameters: { query?: never; @@ -24187,6 +24251,127 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/workflows/runs/{run_id}/steps": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Workflow Steps + * @description Read run progress, typed node inputs and outputs, and linked model-call receipts. + */ + get: operations["get_workflow_steps_v1_workflows_runs__run_id__steps_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/{workflow_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Workflow Definition + * @description Read the current draft and its publication pointer. + */ + get: operations["get_workflow_definition_v1_workflows__workflow_id__get"]; + /** + * Replace Workflow Definition + * @description Replace a draft only when If-Match identifies its current revision. + */ + put: operations["replace_workflow_definition_v1_workflows__workflow_id__put"]; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/{workflow_id}/publish": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Publish Workflow Definition + * @description Publish the reviewed revision as an immutable snapshot. + */ + post: operations["publish_workflow_definition_v1_workflows__workflow_id__publish_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/{workflow_id}/runs": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** List Definition Runs */ + get: operations["list_definition_runs_v1_workflows__workflow_id__runs_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/{workflow_id}/test": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Test Workflow + * @description Start a paid test of the exact draft revision and return its run ID immediately. + */ + post: operations["test_workflow_v1_workflows__workflow_id__test_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/{workflow_id}/validate": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Validate Workflow Definition + * @description Validate a draft without model calls or Python execution. + */ + post: operations["validate_workflow_definition_v1_workflows__workflow_id__validate_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1beta/agents": { parameters: { query?: never; @@ -26036,6 +26221,15 @@ export interface components { */ type: "apiKey"; }; + /** AcceptedRun */ + AcceptedRun: { + /** Run Id */ + run_id: string; + /** Snapshot Id */ + snapshot_id: string; + /** Trace Url */ + trace_url: string; + }; /** * AccessControl_UI_AccessMode * @description Model for Controlling UI Access Mode via SSO Groups @@ -29860,6 +30054,36 @@ export interface components { ChatCompletionMessageToolCall: { [key: string]: unknown; }; + /** ChatCompletionNode */ + ChatCompletionNode: { + /** Id */ + id: string; + /** Max Completion Tokens */ + max_completion_tokens?: number | null; + /** Max Tokens */ + max_tokens?: number | null; + /** Messages */ + messages: components["schemas"]["ValueBinding"] | components["schemas"]["ReferenceBinding"]; + /** Model */ + model: string; + response_format?: components["schemas"]["StructuredOutput"] | null; + retry: components["schemas"]["WorkflowRetry"]; + /** Seed */ + seed?: number | null; + /** Stop */ + stop?: string | string[] | null; + /** Temperature */ + temperature?: number | null; + /** Timeout Seconds */ + timeout_seconds: number; + /** Top P */ + top_p?: number | null; + /** + * @description discriminator enum property added by openapi-typescript + * @enum {string} + */ + type: "chat_completion"; + }; /** * ChatCompletionReasoningItem * @description Represents an OpenAI Responses API reasoning item for round-tripping in conversation history. @@ -32711,6 +32935,13 @@ export interface components { [key: string]: unknown; }; }; + /** DefinitionPage */ + DefinitionPage: { + /** Next Cursor */ + next_cursor: string | null; + /** Workflows */ + workflows: components["schemas"]["WorkflowDraft"][]; + }; /** DeleteAccessGroupBudgetResponse */ DeleteAccessGroupBudgetResponse: { /** Access Group */ @@ -33213,6 +33444,8 @@ export interface components { } & { [key: string]: unknown; }; + /** EmptyRunAction */ + EmptyRunAction: Record; /** EndpointProvider */ EndpointProvider: { /** Display Name */ @@ -43309,6 +43542,48 @@ export interface components { }; } | null; }; + /** PythonNode */ + PythonNode: { + /** Id */ + id: string; + /** Inputs */ + inputs: { + [key: string]: components["schemas"]["ValueBinding"] | components["schemas"]["ReferenceBinding"]; + }; + program: components["schemas"]["PythonProgram"]; + retry: components["schemas"]["WorkflowRetry"]; + /** Timeout Seconds */ + timeout_seconds: number; + /** + * @description discriminator enum property added by openapi-typescript + * @enum {string} + */ + type: "python"; + }; + /** PythonProgram */ + PythonProgram: { + /** + * Contract Version + * @constant + */ + contract_version: "1"; + /** Function Name */ + function_name: string; + /** Input Schema */ + input_schema: { + [key: string]: components["schemas"]["JsonValue"]; + }; + /** Output Schema */ + output_schema: { + [key: string]: components["schemas"]["JsonValue"]; + }; + /** Parameter Names */ + parameter_names: string[]; + /** Source */ + source: string; + /** Source Sha256 */ + source_sha256: string; + }; /** * RankingOptions * @description Ranking options for search. @@ -43397,6 +43672,16 @@ export interface components { /** Summary */ summary?: ("auto" | "concise" | "detailed") | null; }; + /** ReferenceBinding */ + ReferenceBinding: { + /** From */ + from: string; + /** + * Output + * @enum {string} + */ + output: "messages" | "result"; + }; /** RegenerateKeyRequest */ RegenerateKeyRequest: { /** Access Group Ids */ @@ -45972,6 +46257,15 @@ export interface components { /** Run Id */ run_id: string; }; + /** RunPage */ + RunPage: { + /** Count */ + count: number; + /** Next Cursor */ + next_cursor: string | null; + /** Runs */ + runs: components["schemas"]["RunView"][]; + }; /** RunSource */ RunSource: { /** Title */ @@ -45986,6 +46280,51 @@ export interface components { /** User */ user?: string; }; + /** + * RunState + * @enum {string} + */ + RunState: "pending" | "running" | "completed" | "failed" | "cancelled"; + /** RunSteps */ + RunSteps: { + /** Cancel Requested */ + cancel_requested: boolean; + /** Error */ + error?: string | null; + /** Run Id */ + run_id: string; + /** Snapshot Id */ + snapshot_id: string; + state: components["schemas"]["RunState"]; + /** Steps */ + steps: components["schemas"]["StepTrace"][]; + }; + /** RunView */ + RunView: { + /** Cancel Requested */ + cancel_requested: boolean; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Run Id */ + run_id: string; + /** Session Id */ + session_id: string; + /** Snapshot Id */ + snapshot_id: string; + status: components["schemas"]["RunState"]; + /** Team Id */ + team_id: string; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Workflow Id */ + workflow_id: string; + }; /** SCIMEnterpriseUser */ SCIMEnterpriseUser: { /** Costcenter */ @@ -47743,6 +48082,55 @@ export interface components { */ user_ids: string[]; }; + /** StepTrace */ + StepTrace: { + /** Attempts */ + attempts: number; + /** Call Id */ + call_id?: string | null; + /** + * Completion Tokens + * @default 0 + */ + completion_tokens: number; + /** + * Deadline + * Format: date-time + */ + deadline: string; + /** Error */ + error?: string | null; + /** Finished At */ + finished_at?: string | null; + input: components["schemas"]["JsonValue"]; + /** Node Id */ + node_id: string; + output?: components["schemas"]["JsonValue"]; + /** + * Prompt Tokens + * @default 0 + */ + prompt_tokens: number; + /** Request Id */ + request_id?: string | null; + /** + * Request Ids + * @default [] + */ + request_ids: string[]; + /** Response Cost */ + response_cost?: number | null; + /** Run Id */ + run_id: string; + /** + * Started At + * Format: date-time + */ + started_at: string; + state: components["schemas"]["RunState"]; + /** Step Id */ + step_id: string; + }; /** StoredReport */ StoredReport: { /** Id */ @@ -47753,6 +48141,29 @@ export interface components { /** Window Start */ window_start: number; }; + /** StructuredOutput */ + StructuredOutput: { + json_schema: components["schemas"]["StructuredSchema"]; + /** + * Type + * @constant + */ + type: "json_schema"; + }; + /** StructuredSchema */ + StructuredSchema: { + /** Name */ + name: string; + /** Schema */ + schema: { + [key: string]: components["schemas"]["JsonValue"]; + }; + /** + * Strict + * @constant + */ + strict: true; + }; /** SubscriptionAccountFeeRequest */ SubscriptionAccountFeeRequest: { /** Account Id */ @@ -48822,6 +49233,15 @@ export interface components { [key: string]: unknown; } | null; }; + /** TestRunRequest */ + TestRunRequest: { + /** Authorize Billing */ + authorize_billing: boolean; + /** Messages */ + messages: { + [key: string]: components["schemas"]["JsonValue"]; + }[]; + }; /** TestSearchToolConnectionRequest */ TestSearchToolConnectionRequest: { /** Litellm Params */ @@ -51773,6 +52193,10 @@ export interface components { /** Error Type */ type: string; }; + /** ValueBinding */ + ValueBinding: { + value: components["schemas"]["JsonValue"]; + }; /** * VantageDryRunRequest * @description Request model for Vantage dry-run operations (capped for preview) @@ -52085,6 +52509,41 @@ export interface components { /** Worker Id */ worker_id: string; }; + /** WorkflowBundle */ + WorkflowBundle: { + /** + * Contract Version + * @constant + */ + contract_version: "1"; + /** Input Schema */ + input_schema: { + [key: string]: components["schemas"]["JsonValue"]; + }; + /** Name */ + name: string; + /** Nodes */ + nodes: (components["schemas"]["ChatCompletionNode"] | components["schemas"]["PythonNode"])[]; + output: components["schemas"]["ReferenceBinding"]; + /** Timeout Seconds */ + timeout_seconds: number; + }; + /** WorkflowDraft */ + WorkflowDraft: { + bundle: components["schemas"]["WorkflowBundle"]; + /** Digest */ + digest: string; + /** Etag */ + etag: string; + /** Name */ + name: string; + /** Published Snapshot Id */ + published_snapshot_id: string | null; + /** Team Id */ + team_id: string; + /** Workflow Id */ + workflow_id: string; + }; /** WorkflowEventCreateRequest */ WorkflowEventCreateRequest: { /** Data */ @@ -52105,6 +52564,39 @@ export interface components { /** Session Id */ session_id?: string | null; }; + /** WorkflowPublication */ + WorkflowPublication: { + /** Digest */ + digest: string; + }; + /** WorkflowPublicationResult */ + WorkflowPublicationResult: { + /** Digest */ + digest: string; + /** Model */ + model: string; + /** Snapshot Id */ + snapshot_id: string; + /** Workflow Id */ + workflow_id: string; + }; + /** WorkflowRegistration */ + WorkflowRegistration: { + bundle: components["schemas"]["WorkflowBundle"]; + }; + /** WorkflowRetry */ + WorkflowRetry: { + /** + * Backoff Seconds + * @default 0 + */ + backoff_seconds: number; + /** + * Max Attempts + * @default 1 + */ + max_attempts: number; + }; /** WorkflowRunCreateRequest */ WorkflowRunCreateRequest: { /** Input */ @@ -52131,6 +52623,46 @@ export interface components { /** Status */ status?: ("pending" | "running" | "paused" | "completed" | "failed") | null; }; + /** WorkflowSnapshot */ + WorkflowSnapshot: { + bundle: components["schemas"]["WorkflowBundle"]; + /** Contributors */ + contributors: string[]; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Digest */ + digest: string; + /** Name */ + name: string; + /** Published By */ + published_by: string | null; + /** Snapshot Id */ + snapshot_id: string; + /** Team Id */ + team_id: string; + /** Workflow Id */ + workflow_id: string; + }; + /** WorkflowValidationRequest */ + WorkflowValidationRequest: Record; + /** WorkflowValidationResult */ + WorkflowValidationResult: { + /** Digest */ + digest: string; + /** Etag */ + etag: string; + /** + * Valid + * @default true + * @constant + */ + valid: true; + /** Workflow Id */ + workflow_id: string; + }; /** _A2AJsonRpcResponse */ _A2AJsonRpcResponse: { /** Error */ @@ -87094,12 +87626,83 @@ export interface operations { }; }; }; + list_workflow_definitions_v1_workflows_get: { + parameters: { + query?: { + team_id?: string | null; + limit?: number; + cursor?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DefinitionPage"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + register_workflow_v1_workflows_post: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["WorkflowRegistration"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowDraft"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_workflow_runs_v1_workflows_runs_get: { parameters: { query?: { workflow_type?: string | null; status?: string | null; limit?: number; + workflow_id?: string | null; + team_id?: string | null; + after?: string | null; }; header?: never; path?: never; @@ -87226,6 +87829,72 @@ export interface operations { }; }; }; + cancel_workflow_run_v1_workflows_runs__run_id__cancel_post: { + parameters: { + query?: never; + header?: never; + path: { + run_id: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": components["schemas"]["EmptyRunAction"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 202: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AcceptedRun"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_run_definition_v1_workflows_runs__run_id__definition_get: { + parameters: { + query?: never; + header?: never; + path: { + run_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowSnapshot"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_workflow_events_v1_workflows_runs__run_id__events_get: { parameters: { query?: { @@ -87362,6 +88031,261 @@ export interface operations { }; }; }; + get_workflow_steps_v1_workflows_runs__run_id__steps_get: { + parameters: { + query?: never; + header?: never; + path: { + run_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["RunSteps"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_workflow_definition_v1_workflows__workflow_id__get: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: never; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowDraft"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + replace_workflow_definition_v1_workflows__workflow_id__put: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: { + "If-Match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["WorkflowRegistration"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowDraft"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + publish_workflow_definition_v1_workflows__workflow_id__publish_post: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: { + "If-Match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["WorkflowPublication"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowPublicationResult"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + list_definition_runs_v1_workflows__workflow_id__runs_get: { + parameters: { + query?: { + team_id?: string | null; + status?: components["schemas"]["RunState"] | null; + limit?: number; + after?: string | null; + }; + header?: never; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["RunPage"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + test_workflow_v1_workflows__workflow_id__test_post: { + parameters: { + query?: never; + header?: { + "If-Match"?: string | null; + "Idempotency-Key"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TestRunRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 202: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AcceptedRun"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + validate_workflow_definition_v1_workflows__workflow_id__validate_post: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: { + "If-Match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": components["schemas"]["WorkflowValidationRequest"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowValidationResult"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_gemini_agents_v1beta_agents_get: { parameters: { query?: never; diff --git a/uv.lock b/uv.lock index fb591264fc3..3dc5ec031b1 100644 --- a/uv.lock +++ b/uv.lock @@ -5090,6 +5090,7 @@ proxy = [ { name = "polars" }, { name = "pydantic", version = "2.13.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" }, { name = "pydantic", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" }, + { name = "pydantic-monty" }, { name = "pyjwt" }, { name = "pynacl" }, { name = "pyroscope-io", marker = "sys_platform != 'win32'" }, @@ -5260,6 +5261,7 @@ proxy-dev = [ { name = "opentelemetry-sdk" }, { name = "prisma" }, { name = "prometheus-client" }, + { name = "pydantic-monty" }, { name = "semver" }, { name = "sentry-sdk" }, ] @@ -5349,6 +5351,7 @@ requires-dist = [ { name = "pydantic", marker = "python_full_version >= '3.15'", specifier = ">=2.14.0,<3.0.0" }, { name = "pydantic", marker = "extra == 'mcp'", specifier = ">=2.12.0,<3" }, { name = "pydantic", marker = "extra == 'proxy'", specifier = ">=2.12.0,<3" }, + { name = "pydantic-monty", marker = "extra == 'proxy'", specifier = "==1.4.0" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, @@ -5494,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" }, ] @@ -8785,6 +8789,147 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/24/f3/d15bc0b1fb0c4f1e07b3326ffb7489f7ae70da2b56a0884d22641b0eb47c/pydantic_core-2.50.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:7f476456ac2bb0d937f75191494a09c83a30765fea4f70f3b404942fe25f6cdf", size = 2336787, upload-time = "2026-10-11T18:35:42.558Z" }, ] +[[package]] +name = "pydantic-monty" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic-monty-client" }, + { name = "pydantic-monty-runtime" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/79/7c/58467370dd9f2db100ee2627791668f6f21a9fa150ccff9e60efc588ab8f/pydantic_monty-1.4.0.tar.gz", hash = "sha256:db1f1c779f509f7c547e580620a63ab8968cf167fd6a45a4f60b6625187a9848", size = 8025, upload-time = "2026-10-26T12:31:49.864Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/5a/5cb59f15ca5b89df8ea810397523877d0ffba33419a8975f5bde8795b262/pydantic_monty-1.4.0-py3-none-any.whl", hash = "sha256:e398217fdb701e4daaeeb9fd65ad9319627cb114334bc00d6cf2444845c43bca", size = 7552, upload-time = "2026-10-26T12:28:25.388Z" }, +] + +[[package]] +name = "pydantic-monty-client" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" }, + { name = "typing-extensions", version = "4.16.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/78/0f94fe1b8b45ba5e9c8bc622117f949820a71f7e96a67a7b691fdda96e3d/pydantic_monty_client-1.4.0.tar.gz", hash = "sha256:2252ca45e3cfd241a63950fde184b2ff06ddce5b3e1c260282c0733d70590acd", size = 2916981, upload-time = "2026-10-26T12:31:51.132Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/9f/c16b213349757ee032946bb29b59a50d9593820765a347cdeec66b2b4fff/pydantic_monty_client-1.4.0-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:6fc2b38d871932b43cf08d6ccd9cc38ca6f9fe386845658ea35f4f0de99e6c4a", size = 4288592, upload-time = "2026-10-26T12:28:26.929Z" }, + { url = "https://files.pythonhosted.org/packages/38/32/97f6b6bd430e32a69f3398e3528ac5c4140d79303b18a0d3d59ede7a02c9/pydantic_monty_client-1.4.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d4597ffa0664a0f2aac77d95f8d2c8786444b1d6aa7c5411e5538dde2d37ff96", size = 4060552, upload-time = "2026-10-26T12:28:28.661Z" }, + { url = "https://files.pythonhosted.org/packages/34/b1/327f06e728f8d9e0369ae6335775387fd2ee39eb0d6d528878590d316389/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:141422a3af8536a139f7e099e9f7ac3f78307d2381564ec91239b3a91bd13bf7", size = 4122496, upload-time = "2026-10-26T12:28:30.273Z" }, + { url = "https://files.pythonhosted.org/packages/17/a2/77563d71661923383addaae926f8bdb5799e54fe91434b1df41db814d544/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_armv7l.whl", hash = "sha256:32eaf5482170020b1bf85576c58d297a99dfdb126aecc451b266a6cf12b00bf2", size = 3807830, upload-time = "2026-10-26T12:28:31.926Z" }, + { url = "https://files.pythonhosted.org/packages/9b/4a/9d505fea5afebfcfcfee49374dcf618e0fc4768bfbd66c545e9e429663ff/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_i686.whl", hash = "sha256:0ec3c0205f02bc7d327fc8c53e985c68df5991607c8e18ec9b2d3562519a08ff", size = 4049655, upload-time = "2026-10-26T12:28:33.607Z" }, + { url = "https://files.pythonhosted.org/packages/0b/2a/f0b3300858e34a66d5a0270d7c2ef3d5c4bbca649b0ffcef548e9f122722/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_ppc64le.whl", hash = "sha256:33b9aa71bd8bf5cd0f65ab9dbf8241684f0cf3d7d70c99eb00bc45a6608b43ae", size = 4260607, upload-time = "2026-10-26T12:28:35.216Z" }, + { url = "https://files.pythonhosted.org/packages/71/d0/7bab7113a230bb21e28ccec28d2aa8cfd17426528c6a541baa01e8f18d90/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_s390x.whl", hash = "sha256:0a7e4a7c3d04651cf6c2f6503fec7694b979c02a1a286546949500d750f0210a", size = 4154742, upload-time = "2026-10-26T12:28:37.05Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d1/506e3450bddedc44f80ce40fda1840f65da26b15de4ff8f6c123cb3c8897/pydantic_monty_client-1.4.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:2204520a4f4c4360527aa957e9b79f0295f14dfd30ea1bc0b7db373e7d8958e5", size = 4411287, upload-time = "2026-10-26T12:28:38.494Z" }, + { url = "https://files.pythonhosted.org/packages/a0/91/59cc51af399c80490231207c54ca2557f36fd48024294a707c28d09392e9/pydantic_monty_client-1.4.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:309b2fa11943de2711ff9e453a2327cfd265192ac6b5ecad19fafde31671cdb5", size = 4295152, upload-time = "2026-10-26T12:28:39.921Z" }, + { url = "https://files.pythonhosted.org/packages/5a/17/ad54eab3c573c7e186ff270eef3389717c1c4d239b51f8901901daf5be6f/pydantic_monty_client-1.4.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:d45d0351d563dc652f8963339a098c0dd29b026df10a4dc92fdcae1c04f420af", size = 4605594, upload-time = "2026-10-26T12:28:41.769Z" }, + { url = "https://files.pythonhosted.org/packages/50/83/e9207c92d37f445b745cadcf0a5a1b52c3fdaac94b4aa9b04347093666dc/pydantic_monty_client-1.4.0-cp310-cp310-win32.whl", hash = "sha256:959ac1ec194f32d17b6e78fff5b7c7b31815e68b551dd17d0c6ab9b007c020e9", size = 3686561, upload-time = "2026-10-26T12:28:43.27Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e2/28ded8f46e31c635c0fea322fce526ec05df71c5f642373c0595ab96a1f4/pydantic_monty_client-1.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:959e5e03edd2f6ab719a306cd36f86e5cfb1bcd398f346ee439876dbce35108a", size = 4360024, upload-time = "2026-10-26T12:28:45.302Z" }, + { url = "https://files.pythonhosted.org/packages/91/7b/2798ebabdd90c1258274e4f85541130d31e9d543d330eb3ee98b98eaee2a/pydantic_monty_client-1.4.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:6e15c2a6f2833806ddb1f4334aa8272cdc9ab0a1c2301824af6354b88e2fb17c", size = 4287529, upload-time = "2026-10-26T12:28:46.93Z" }, + { url = "https://files.pythonhosted.org/packages/62/6c/f1618692da03841fb11a73731b52b6d11e0b861d47034059c877371c0950/pydantic_monty_client-1.4.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:485981ad29da00f295ad994299db809bb60aeb3b664beff7b6f5c31116c78b59", size = 4060688, upload-time = "2026-10-26T12:28:48.913Z" }, + { url = "https://files.pythonhosted.org/packages/82/f3/7c45f7f38b75d94a3f7164d79cb32f1bcd88ee657ed5646f4ea4c6bc2503/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:063534e65ae971099f2e64c674626d871aa3efcf6df983db5ba520d9085a60a6", size = 4121979, upload-time = "2026-10-26T12:28:50.783Z" }, + { url = "https://files.pythonhosted.org/packages/4f/94/e47425838420e172664868c93e99730cfc3ada642eb467be15deebbde4cc/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_armv7l.whl", hash = "sha256:fbd9ae34c2d68d0b14774468f87c43369041631b4de469e336ada7a34a57a30f", size = 3806350, upload-time = "2026-10-26T12:28:52.244Z" }, + { url = "https://files.pythonhosted.org/packages/6e/9f/4693140e8e27f2afa1c6d718f2d45ec17f201a57189f3353ed4dcad80577/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_i686.whl", hash = "sha256:814f551fe454c5db082192b53d75ec4f71c1b790ea944d68763528aba714e182", size = 4048270, upload-time = "2026-10-26T12:28:53.976Z" }, + { url = "https://files.pythonhosted.org/packages/b0/c6/69c7747f62b6a5a25d450563c630190cec254ef43e4dab0446329733bcb2/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_ppc64le.whl", hash = "sha256:69d7a87c60a8a2eea5b488f6c373b49f4c11a61917fbcd45c1247f7daf0a3a45", size = 4259483, upload-time = "2026-10-26T12:28:55.879Z" }, + { url = "https://files.pythonhosted.org/packages/e6/9c/b901f54c53ad9ba2cfb2a0319fb59211304b51512746b4c0be1516f3ca6a/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_s390x.whl", hash = "sha256:94cf89b243dc9dad970f36072849e7607fc4e7573d0b190be4613bd6167d1b9a", size = 4153767, upload-time = "2026-10-26T12:28:57.644Z" }, + { url = "https://files.pythonhosted.org/packages/f8/48/e342373a0164fe5ed71015707de82f66761e107cfe54bb77c2a2f67854fd/pydantic_monty_client-1.4.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:d5676ea4b2fead463c427f0a4536f59156feb978c56bd0862aef5ba24c00aff8", size = 4410295, upload-time = "2026-10-26T12:28:59.915Z" }, + { url = "https://files.pythonhosted.org/packages/fa/b0/d85b45cea1a542251ef229f21b31f5d9368cfc3f24ad7c88b7e7efbfc95a/pydantic_monty_client-1.4.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:12e35274ddac034f8c4772e1e0931a15d834ba8db84d347372e2d09f349f2168", size = 4293927, upload-time = "2026-10-26T12:29:01.998Z" }, + { url = "https://files.pythonhosted.org/packages/57/5d/4fbfae9af7c350b3672f2684253baf422299cea6e282d69a4648e127e26c/pydantic_monty_client-1.4.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:d28aa3d51cb522861cf49f7b21e3405afa256d0938606c330e72e2089a4a158e", size = 4604347, upload-time = "2026-10-26T12:29:03.604Z" }, + { url = "https://files.pythonhosted.org/packages/b5/de/a167924d35a052a74f94e14d9793cb9c0d3cb2fa159b0b2e48c7c1f4e180/pydantic_monty_client-1.4.0-cp311-cp311-win32.whl", hash = "sha256:ce66f43bea399ae8454c603664ff2bdb9bb3a58034a4e2855fdec85d412c4a08", size = 3685412, upload-time = "2026-10-26T12:29:05.397Z" }, + { url = "https://files.pythonhosted.org/packages/cd/7f/24317b4e747aca7aa85dd62d5540f976fdf5a7eaaa40f351e7586c1812b8/pydantic_monty_client-1.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:25440b0b3c5b7eebf010075729db1b4778cf90e9cdebd988a1fd5714df73f85b", size = 4359255, upload-time = "2026-10-26T12:29:06.95Z" }, + { url = "https://files.pythonhosted.org/packages/64/9f/6605612467001a29734947e3e66eb2f9c03172e524ad89a72ede2c00e2c5/pydantic_monty_client-1.4.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:8c771ded8705f24fdac9ea119547ffe3cd7126ac4d3f7a1514adf7f3f9961f4c", size = 4317732, upload-time = "2026-10-26T12:29:08.5Z" }, + { url = "https://files.pythonhosted.org/packages/a1/a0/bf17349667f3a3f76054fac14702840bd745a86c4de86cebaa4f50db742b/pydantic_monty_client-1.4.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0f0677e9a3be73159c772c4bc55292150356b5080cc4623196dd57407ddc5504", size = 4046984, upload-time = "2026-10-26T12:29:09.991Z" }, + { url = "https://files.pythonhosted.org/packages/b0/2e/5dca50a7086efb25fd7850595bd491c76692bc8c83e1331a0ca5dfd71331/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:287614a13d60c590475a39708171b5a037a6dbb6868e0125871cce365009846a", size = 4121512, upload-time = "2026-10-26T12:29:11.416Z" }, + { url = "https://files.pythonhosted.org/packages/0a/c4/b5dd4714b9cb204f1a992e6244066dbb1e880eae74303e526a77f2a1536f/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_armv7l.whl", hash = "sha256:8bfe5c08eea64423174712f641189dcc5e9fc2caad18c8b243af041bec8402a0", size = 3805501, upload-time = "2026-10-26T12:29:12.912Z" }, + { url = "https://files.pythonhosted.org/packages/12/1a/a54b7346c8b045318d65eebf97d09817dc3d04cf89dd2e678c6f3d2ddc6a/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_i686.whl", hash = "sha256:7471989402d267e181d7d5a528cc48a4674d3815fc56a67aeeb8caaeef324477", size = 4045097, upload-time = "2026-10-26T12:29:14.75Z" }, + { url = "https://files.pythonhosted.org/packages/b8/a4/94d8e46b0ed94c3f1f39c44269d7241363ca071d46351e4f0088064019cc/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_ppc64le.whl", hash = "sha256:d6f50a54f2a315897c047dc503881282d02bc601a5da04aa464af071ddc97fa0", size = 4257816, upload-time = "2026-10-26T12:29:16.563Z" }, + { url = "https://files.pythonhosted.org/packages/cd/23/80b53e023b05149ccd318f0873c4f00f1cddf53ad7bf4c72cda439ad3982/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_s390x.whl", hash = "sha256:35217394abf732e2df1033f7a5c3b28eb2b78ac4ce468886f88214d1d342e6ff", size = 4159259, upload-time = "2026-10-26T12:29:18.121Z" }, + { url = "https://files.pythonhosted.org/packages/23/3e/827db7c549b4b33cf3d3c66ec3e53920ec40ce86228b85ee3687053150e9/pydantic_monty_client-1.4.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5f3f32a9cc72aa831ab4c951b4c6e720a1c81cca0db8d8956b70f5c242381ba8", size = 4410442, upload-time = "2026-10-26T12:29:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/b1/94/ea72c9c9f05e859c1275191567274d63b0cfd84a40fd700270b91765ed04/pydantic_monty_client-1.4.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:4c4e54a0f670fe37fc8a6c065286c83061e5269bfb5bb3a9a0ee4f0369868e19", size = 4288907, upload-time = "2026-10-26T12:29:23.365Z" }, + { url = "https://files.pythonhosted.org/packages/08/86/3eec37a86a95418c9dfd36a34216b66aa0da08c94ff908f6745bdca49a0b/pydantic_monty_client-1.4.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1edfff6112e2757e2e1760d99a6b7d400f6a4f2fe7f77053e030f21e5d807255", size = 4605169, upload-time = "2026-10-26T12:29:25.005Z" }, + { url = "https://files.pythonhosted.org/packages/ac/39/c49055ba96c72eb365faa8abd88f0467f8c750c51a7fa01a0c19622973c5/pydantic_monty_client-1.4.0-cp312-cp312-win32.whl", hash = "sha256:f7ad83cb3ac4939c79ecc78b7cee6b0db8c5f69be5ca3396dc36d6311b3c2261", size = 3679130, upload-time = "2026-10-26T12:29:27.175Z" }, + { url = "https://files.pythonhosted.org/packages/30/bb/9bd7cb6d9975e10a62fb77836137fbc63111998639920c31a43d06e43b8d/pydantic_monty_client-1.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:be6f88686d3e93f73e006c7d88fee42350986afc7d307a08f4b0cd3f54bbcf49", size = 4346625, upload-time = "2026-10-26T12:29:28.716Z" }, + { url = "https://files.pythonhosted.org/packages/0e/90/7fb99bed2696f847f9a6cda709ae24f3e8eb188d3e9b869a9982a9a8ae76/pydantic_monty_client-1.4.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:f3104fab8df6995769c9b241f634e99d46782270e21d5b8b197bbe5874866b2f", size = 4324480, upload-time = "2026-10-26T12:29:31.232Z" }, + { url = "https://files.pythonhosted.org/packages/2f/03/47bf82469d8e7d451b1c5e5bd2fae427d4768346d6ea271d0a0f8fd4bb14/pydantic_monty_client-1.4.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c900bebb52812a58055a7cb00f8f4e678688aab2e66f8262e9d0733f1ff5874a", size = 4052142, upload-time = "2026-10-26T12:29:32.999Z" }, + { url = "https://files.pythonhosted.org/packages/2b/06/ac16fc6aded0eb73e18a27e4f7835f1810e6f9fa1e8f01fd2f8bcb81670f/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:3dc3e1d518f5732cdf254bdea8f56f54408ac3a5d84818ec992286f6859fd477", size = 4123662, upload-time = "2026-10-26T12:29:38.367Z" }, + { url = "https://files.pythonhosted.org/packages/39/8e/1ab5ef0e71fea32e29305fc7b2890107858008e8e5f6cc1e5b9a169092fb/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_armv7l.whl", hash = "sha256:6ad41d4809f51523698fb0fdf8dda2506abc120f7e4d409f75205c0bea40f158", size = 3806072, upload-time = "2026-10-26T12:29:39.849Z" }, + { url = "https://files.pythonhosted.org/packages/49/86/477ef5c02ec571b73853d20bf6927f81e04bd975525b168fcbdcdc3a98c9/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_i686.whl", hash = "sha256:aa0c3ecd210818c53fc9965ffd863ec7fa725373fda73e70627311fbbc8c12fc", size = 4044974, upload-time = "2026-10-26T12:29:41.35Z" }, + { url = "https://files.pythonhosted.org/packages/4b/92/03aee4671043f25ab2443979dd000486a3f6642ee860413c82febec340ca/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_ppc64le.whl", hash = "sha256:879b5c2351527d0d4ce7de73b0330bb934184fd62fe04104166849c233827e19", size = 4261162, upload-time = "2026-10-26T12:29:43.154Z" }, + { url = "https://files.pythonhosted.org/packages/b2/b3/ec5ab86e16987df0dd31c0dd3df2ddb0efa89367bc2106e1818ef5703715/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_s390x.whl", hash = "sha256:15e9949499463fdf707953689a210b7c8f7af216877db7474eadd35f25418d52", size = 4161701, upload-time = "2026-10-26T12:29:45.275Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/c6e3c5ce48dec3faae4a70690feeeaf4cd989b582f5aec1b305ae0f0fa13/pydantic_monty_client-1.4.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:15bd85bc75f1e6fafae2fc7f83f88348795d4376f25e7ede8d23b1cc504005a9", size = 4415744, upload-time = "2026-10-26T12:29:46.821Z" }, + { url = "https://files.pythonhosted.org/packages/e8/26/7db5c0400544bd63be3f6c7b4c6855ae4232610be7a81016f9c0cc2e1ba5/pydantic_monty_client-1.4.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:4079c47f62fce1655e1c3004f630676e18cf9fa2b652adff835fad03610b3762", size = 4292858, upload-time = "2026-10-26T12:29:48.405Z" }, + { url = "https://files.pythonhosted.org/packages/50/43/e560d6e2ae6f355e2306a8a1c7149ef49271ed4e349d1ecfcd26aba5a418/pydantic_monty_client-1.4.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:fc5b9e2d2ec0e2e6f38497803c61dff5c800e9b8f37f9b127e6ae339b7aa4f20", size = 4609688, upload-time = "2026-10-26T12:29:49.883Z" }, + { url = "https://files.pythonhosted.org/packages/12/7c/6af90b241f4adada9bfa7c5170b13ec6cfbb0aad319aea4fd6688080cfa1/pydantic_monty_client-1.4.0-cp313-cp313-win32.whl", hash = "sha256:272eb19ae41423448da0673a135bc0ecf3edf4f9d7860efb357b555efcab9563", size = 3679388, upload-time = "2026-10-26T12:29:51.321Z" }, + { url = "https://files.pythonhosted.org/packages/56/46/ffa8a123d9f6aa95984216faa90da223f3657f741f93573dee278e687ea8/pydantic_monty_client-1.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:ce9dbbceaa891fbef9edb2428839e636784c1fca10c4b4c72179ad8ae5f85338", size = 4350595, upload-time = "2026-10-26T12:29:52.9Z" }, + { url = "https://files.pythonhosted.org/packages/6f/4e/23788c86158a3ba16246716c6f0d0e331b0271a98ca2ee35a28772bcb09a/pydantic_monty_client-1.4.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:d68b922161f7ed637addd58ee9e25082406009832e114838683d6cae77dbe587", size = 4324862, upload-time = "2026-10-26T12:29:54.68Z" }, + { url = "https://files.pythonhosted.org/packages/d7/18/a1b951d4ebfa939f10c54200909f8d667fb58eef1489d07d94aea7de9c46/pydantic_monty_client-1.4.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f633e1786bc12acd2b31f2a0fd90f955982ef0a06c9a2ef1d2e3d5832987fb11", size = 4053512, upload-time = "2026-10-26T12:29:56.205Z" }, + { url = "https://files.pythonhosted.org/packages/2f/64/297f1580de9ef26b975b73f3020546145cd50f6000abb275a643474d0515/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:87260bf6a226f5bbcc03e4a8883ac3a04bebb1c8a3e3da747ee8c464443e6f37", size = 4123580, upload-time = "2026-10-26T12:29:57.975Z" }, + { url = "https://files.pythonhosted.org/packages/68/00/93bd1fa14eabf97c59b983ba020ef951391ff2497c06babb2d6c61635aad/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_armv7l.whl", hash = "sha256:ffa480dab03bbc0ae7e59d13e894c42c89909837a7a2ee3d8e41967d192213b4", size = 3807832, upload-time = "2026-10-26T12:29:59.388Z" }, + { url = "https://files.pythonhosted.org/packages/9c/3f/4f2fbb788c71f346a73fd08fdb172053aa09cb382c51091e7058868a8c86/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_i686.whl", hash = "sha256:731499cad98261d0e8f1d125349ef268c2db9081b2cb25d8bef68d99c62ecd80", size = 4046624, upload-time = "2026-10-26T12:30:01.16Z" }, + { url = "https://files.pythonhosted.org/packages/56/41/8055a6d78c10399602cdf605a94146245a13d2a6ad519be65903f8ad9406/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_ppc64le.whl", hash = "sha256:f50b807e1475e47cf813049c265070d3cdb76e8ef7e5b77b3e34908d15a8a4db", size = 4262907, upload-time = "2026-10-26T12:30:02.785Z" }, + { url = "https://files.pythonhosted.org/packages/fc/c6/e466e7a374fdc54718c8bb664a1bfbec063733315254490f4b21877f0459/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_s390x.whl", hash = "sha256:586d0f24faeabb6243584d2ce13bb2180a9607ac157a4e9e4edc78e40b2ac32f", size = 4163131, upload-time = "2026-10-26T12:30:04.228Z" }, + { url = "https://files.pythonhosted.org/packages/56/ac/3f868d04d3db7902ffcf0e4455498b5ad5baa73dbd1188fa978efc993750/pydantic_monty_client-1.4.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:41e7bd5b0b6ae8f1ac878f2c2a3c60656319fb3d51602bb638f72d5db0734711", size = 4416533, upload-time = "2026-10-26T12:30:05.78Z" }, + { url = "https://files.pythonhosted.org/packages/b7/27/913d7c73cdffb78ccd2fc516d05b0c392e48f61d42166910aff1fdf18d63/pydantic_monty_client-1.4.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:fa992fa81e1d753bd1c637d90f5713a92d7ce25d968c6195e88219fdf8e628c2", size = 4293701, upload-time = "2026-10-26T12:30:07.249Z" }, + { url = "https://files.pythonhosted.org/packages/85/ef/ec082ca4405889d784d0c667a63c17c3755cc1b0c775bf7e801509cd4fc9/pydantic_monty_client-1.4.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:868d269d20b356d513968aeecbd69161adf8f2a59bc506aac463fff30bd73ab1", size = 4610590, upload-time = "2026-10-26T12:30:09.211Z" }, + { url = "https://files.pythonhosted.org/packages/23/18/a8b360113cb09539ba67e012471661325549f314d3634162b7afbdd0b07e/pydantic_monty_client-1.4.0-cp314-cp314-win32.whl", hash = "sha256:878aede0148c7a9b51bf20940157bbcad089991f60ff5bf27ed34f402d31332f", size = 3680393, upload-time = "2026-10-26T12:30:11.07Z" }, + { url = "https://files.pythonhosted.org/packages/91/ac/fd15b7e8e2517b94cf5c2ac9124f1871367a9e94843e89eb3dcc829bf72a/pydantic_monty_client-1.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:59f038a326bb3ee8716236a62119f54d7fb27c00a3f9e56ac06ae7d468a48284", size = 4349750, upload-time = "2026-10-26T12:30:12.813Z" }, + { url = "https://files.pythonhosted.org/packages/6c/eb/1d513581f8e8d0e5363edc3cad334cce69777f8abfcbbd6883510e8a53c5/pydantic_monty_client-1.4.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:69aec529dd526655cb5861d6080b2df4bca2c5ba2fbdfd46c2b3ddea4050f31d", size = 4312356, upload-time = "2026-10-26T12:30:14.585Z" }, + { url = "https://files.pythonhosted.org/packages/26/6e/2e9187b663f3cb222beb8c6b5722250d36c12b8b65b71642825ad9c5013f/pydantic_monty_client-1.4.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f1cc47e6fd4a163f5bbb88934aa84fbbfa8ebfb5f76932896467d49dabfaf905", size = 4041090, upload-time = "2026-10-26T12:30:16.314Z" }, + { url = "https://files.pythonhosted.org/packages/f1/e3/f144d76a0956fa35487391f090af7cde4ffc48366f40945d2875b799bfd7/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:06290ce25c3b4a89941bdb9a08e1c42f2e288274ff98edcbd7fc2dbbbc2730c7", size = 4118219, upload-time = "2026-10-26T12:30:18.067Z" }, + { url = "https://files.pythonhosted.org/packages/c7/96/ae06f344a390a3cdb9aaff29cd89c6a28c680a8301f32b6d970873bac73e/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_armv7l.whl", hash = "sha256:4f7ed73540db9c17d6a64dcc1e823946444a44574883dac53c68cb93987dba97", size = 3795829, upload-time = "2026-10-26T12:30:19.521Z" }, + { url = "https://files.pythonhosted.org/packages/4c/54/023a9516f211a63f44881c5a4536f8827742541b5ccc50522c67712482b6/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_i686.whl", hash = "sha256:0493a65bc534445123834276bedf38baa4af964ca35b8fef66e440d6d23a1964", size = 4037589, upload-time = "2026-10-26T12:30:21.561Z" }, + { url = "https://files.pythonhosted.org/packages/09/d2/4aab45f23dad94fdea6705a236ff9eabf48577fcc6e4e80df031382cf231/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:41305b81fe9d61fcef7e9c6eabc7308053ebaa568a0ef9d21cc4715198392937", size = 4253700, upload-time = "2026-10-26T12:30:23.358Z" }, + { url = "https://files.pythonhosted.org/packages/09/40/fb1c9248efbd4b5e9d2219b63c3185ea3a0113f869415b56a489259ef3fe/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_s390x.whl", hash = "sha256:1943eb2505ac1b80e9e676abb92d015d5406a7413fb6a4f6d364d44905a521c9", size = 4147967, upload-time = "2026-10-26T12:30:25.107Z" }, + { url = "https://files.pythonhosted.org/packages/75/79/3b7b33118d03998ba1dee714919bf9f5d8099ae0771aa75495dcad1e9dfe/pydantic_monty_client-1.4.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1d1e9db032a2ff431d44de84fb569b0736c3f21dc58c84ef974985dbc722a34c", size = 4413011, upload-time = "2026-10-26T12:30:28.146Z" }, + { url = "https://files.pythonhosted.org/packages/45/76/40656a25a297112aedbe4d34f7f745e90aeb8f11a3377cbefd3a344cc445/pydantic_monty_client-1.4.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:076b256cc3b8b0301934f8c12a38f43e549fc2f07b945bf9a9251dfaa9d11aeb", size = 4285647, upload-time = "2026-10-26T12:30:29.77Z" }, + { url = "https://files.pythonhosted.org/packages/44/f7/10b71b443fdfd977e2e60913a2928f562cef97eb13119ecec91a6a719084/pydantic_monty_client-1.4.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3d466a34a6d111b07cff2464796d681a043279cb956a4286e7bfb8c730d22ab5", size = 4606819, upload-time = "2026-10-26T12:30:31.382Z" }, + { url = "https://files.pythonhosted.org/packages/63/c3/1a41dfed2d0fe15965a57919a0266e8fed06b20513f5fe773ae7b7aa26af/pydantic_monty_client-1.4.0-cp314-cp314t-win32.whl", hash = "sha256:fd7ccfa999bbcd05bd8b248ac2071c3d9b9a7e44b99e696596ee785e69aed91c", size = 3671748, upload-time = "2026-10-26T12:30:33.1Z" }, + { url = "https://files.pythonhosted.org/packages/01/b2/6a75a83dbc25e0244ac986ea33a822af7b6f0ca363b5c4ee852fd26fe396/pydantic_monty_client-1.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:8ad5d44fec1b0f60f68209eeff600fc8673dc6960d070bb05e307c6aba35c4fa", size = 4344965, upload-time = "2026-10-26T12:30:34.712Z" }, + { url = "https://files.pythonhosted.org/packages/24/86/61e5de22b9fb6963adba6d0b53b7aebd4a89f01a744f63596e9dbe56a5a0/pydantic_monty_client-1.4.0-cp315-cp315-macosx_10_12_x86_64.whl", hash = "sha256:1a8907baeca0f0bf0d257620bbe37f6867ac6a7ff2cd677591e6ad52b3ba8667", size = 4324609, upload-time = "2026-10-26T12:30:36.236Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5b/e766f4c65c47e18ed41dec5ceb3c795ec843c2c03f0d9cc0df4b00bdc52a/pydantic_monty_client-1.4.0-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:7e6423fc6026378c241e20d3cce48e549d1cbff744c1622f0bf625f85b2c3789", size = 4053763, upload-time = "2026-10-26T12:30:37.664Z" }, + { url = "https://files.pythonhosted.org/packages/bb/ed/419fbd3627764e67b9bacf8af72d1e28a214759b6efcbabcb03d3a373aae/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_aarch64.whl", hash = "sha256:f87a316222f38755b904902fe8c3c4e581165af25bcf872da741c354f004dd9d", size = 4123545, upload-time = "2026-10-26T12:30:39.322Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8d/40696d7bd4c5af8be86275943a98311beb011dd8aec6805b2a163924b4b4/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_armv7l.whl", hash = "sha256:7fdffd11c96710c81b85f7e35faa4641bf6edcee092724850344e2045c88ce6d", size = 3808033, upload-time = "2026-10-26T12:30:40.822Z" }, + { url = "https://files.pythonhosted.org/packages/b5/51/9242b4f151db9aa653e974dee222fb7a7f660b5474e17203e1668d3ff2b7/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_i686.whl", hash = "sha256:b24b44316bbb14dab9088fdcfa76f66ee92107b9e429186d2921aa2142f697e8", size = 4046971, upload-time = "2026-10-26T12:30:42.479Z" }, + { url = "https://files.pythonhosted.org/packages/35/7a/6619c87bb695b9cc0cc727ee427e1617812bec184dffa0fdc6441e0795f2/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_ppc64le.whl", hash = "sha256:ec085c9c5d66a540067148dabc4d894d160741277947a57ac4179c5607a5bc2c", size = 4264420, upload-time = "2026-10-26T12:30:44.554Z" }, + { url = "https://files.pythonhosted.org/packages/4a/90/b941236f60533f48c077c54b810d9f47d5c75b71f5057c13a1e4e1091905/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_s390x.whl", hash = "sha256:2d80aed4644d4ff907324582def94e76151ad862198776d544d760db8dc57219", size = 4163434, upload-time = "2026-10-26T12:30:46.765Z" }, + { url = "https://files.pythonhosted.org/packages/0a/a0/31b38b607c9a22dac273f860b21fb21e591b126a1ae251fb4a54bf110d43/pydantic_monty_client-1.4.0-cp315-cp315-manylinux_2_28_x86_64.whl", hash = "sha256:e459c3722c5db566b5f4295f5c5862e77eb0d09abb49b1feac9b0e39467c6fae", size = 4416755, upload-time = "2026-10-26T12:30:48.826Z" }, + { url = "https://files.pythonhosted.org/packages/50/c6/e07dccc12c598739fe442881b71d0d41293713f2eb2bd4f8f63270c0a854/pydantic_monty_client-1.4.0-cp315-cp315-musllinux_1_1_aarch64.whl", hash = "sha256:7130cf0b75557301764b57925da46cbaecae68440ea7b9401bc2d2b0e6bf9dd7", size = 4293680, upload-time = "2026-10-26T12:30:50.496Z" }, + { url = "https://files.pythonhosted.org/packages/5a/5f/be7b09bf2e159afcd91932dcd5f8e8878cdbdb2fec863b81824e8d370d3f/pydantic_monty_client-1.4.0-cp315-cp315-musllinux_1_1_x86_64.whl", hash = "sha256:29f76fb33adf72dda3277adeae13b11340178157b1ab0b544f8f13e9d032780c", size = 4610741, upload-time = "2026-10-26T12:30:52.095Z" }, + { url = "https://files.pythonhosted.org/packages/39/d5/bdefd85f079cbf6a6c8aea2b8b086f6a3ebd48d3291866aa45c6c2741bb9/pydantic_monty_client-1.4.0-cp315-cp315-win32.whl", hash = "sha256:7e073a55ac45eb001e2e702dca94b2bed882ee5fe65898642dd6f1e116589f93", size = 3680779, upload-time = "2026-10-26T12:30:53.66Z" }, + { url = "https://files.pythonhosted.org/packages/ab/4e/dc9403b65e914bba1b9507e33a52b243a6c1d83a8b85087306d71b412b7c/pydantic_monty_client-1.4.0-cp315-cp315-win_amd64.whl", hash = "sha256:f928cc867284e2f366c2b2677590baceccd1800ff34aab82221df9a7e3f58336", size = 4349931, upload-time = "2026-10-26T12:30:55.241Z" }, + { url = "https://files.pythonhosted.org/packages/12/4b/a032b0d12627c182b4b52cf9cd6e34a3c5fceaa635a2ce8b66ce3e753205/pydantic_monty_client-1.4.0-cp315-cp315t-macosx_10_12_x86_64.whl", hash = "sha256:5d0fbba5046db61eb82a41425c30c1317029694c68fe864cec8b05aff5fc8abc", size = 4312356, upload-time = "2026-10-26T12:30:56.765Z" }, + { url = "https://files.pythonhosted.org/packages/4c/90/4bc4258da5889fb7cd958f5e9fc79d254d51c5a9caba732abf48113e5332/pydantic_monty_client-1.4.0-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:639f4a7a8fb8f1978f9036734f7b8bc68b0f48074a91f9f2360fe51fb0023b78", size = 4041241, upload-time = "2026-10-26T12:30:58.258Z" }, + { url = "https://files.pythonhosted.org/packages/04/20/c8c385e7b756d284e1a6c77bf9cf2b1d7c41b5eeec033dc01c0bff38ab6f/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_aarch64.whl", hash = "sha256:cc66b78f433dd67ea4d9697fa50b4d5d24d7f6a96e99cff162ca1d46e89998bd", size = 4118063, upload-time = "2026-10-26T12:30:59.943Z" }, + { url = "https://files.pythonhosted.org/packages/61/0d/61bf503046e0954675e746ca468d1115c682125d96fb05e41a821b93f32d/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_armv7l.whl", hash = "sha256:74579bd3139c6c6cc707199b6956f97dbdf16e7c7ce4a802e1af64d48e10e68b", size = 3795982, upload-time = "2026-10-26T12:31:01.494Z" }, + { url = "https://files.pythonhosted.org/packages/e1/4a/11deac6408b5bfca8d9e2150802e366b59f80d5cba37b5d4a2bd99f9a545/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_i686.whl", hash = "sha256:527a10831d7bb56b728c878c8c80ee0c7cdcfbf1e4c3b82474e8b58d31beeba3", size = 4037807, upload-time = "2026-10-26T12:31:02.884Z" }, + { url = "https://files.pythonhosted.org/packages/bc/5e/16d5ed8a6d55d80847808248734ba0a7f68e5d4f077f46f4c5924344f8aa/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_ppc64le.whl", hash = "sha256:59bac93edeeb060b11b0b4d74bfb2afbd7548e8f67d70bc132d7526a4e4177fd", size = 4254709, upload-time = "2026-10-26T12:31:04.463Z" }, + { url = "https://files.pythonhosted.org/packages/9d/21/02c8a175b3fdc7098dcc2c25e6e0faff89c8346c86ea6b070ada9ecee2bd/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_s390x.whl", hash = "sha256:ee32f3242000c4020e5708abe1e4c36615c0ee4ac5b8d87318079f76a794e109", size = 4148932, upload-time = "2026-10-26T12:31:06.012Z" }, + { url = "https://files.pythonhosted.org/packages/54/32/1e97c4c2b4b35fbdd8c952f756038e6a92d354f9e30eeee2d5226daadad2/pydantic_monty_client-1.4.0-cp315-cp315t-manylinux_2_28_x86_64.whl", hash = "sha256:8cd489ea1fab085164185284cad43aa8780833d79e9e2f80071dd8494e5c4dd5", size = 4413415, upload-time = "2026-10-26T12:31:07.609Z" }, + { url = "https://files.pythonhosted.org/packages/a9/2f/11644f09adf9ebfe98a9e4a42f5d8583f4a465860754da1fd9926870f27b/pydantic_monty_client-1.4.0-cp315-cp315t-musllinux_1_1_aarch64.whl", hash = "sha256:02bdae4755e7c90ecc32b15ea68e13a527ea6f9d6574a54a92d0242ec2a2866a", size = 4285838, upload-time = "2026-10-26T12:31:09.182Z" }, + { url = "https://files.pythonhosted.org/packages/a7/32/528b03acf62fdcea145f7f7b2ea9b7a33a6364fd734874558f448f5c07c7/pydantic_monty_client-1.4.0-cp315-cp315t-musllinux_1_1_x86_64.whl", hash = "sha256:4412a0adc72e4c51bc3f23993500856523498676936f3eeab26d223bce27314b", size = 4607062, upload-time = "2026-10-26T12:31:11.578Z" }, + { url = "https://files.pythonhosted.org/packages/4b/62/72379b965360e953cb45c6aebc61605951f6b2e0e6d703571dcd780d6518/pydantic_monty_client-1.4.0-cp315-cp315t-win32.whl", hash = "sha256:29cdd6fc5c4ffd0e6580455185a138b63b7ffddbc4eda1f468e8ae96f611a4c3", size = 3671769, upload-time = "2026-10-26T12:31:13.115Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/ab795155a741ff3208ddb6f6c6475ccbb5cb65384329f4d0bad079443748/pydantic_monty_client-1.4.0-cp315-cp315t-win_amd64.whl", hash = "sha256:c7111b4239f745b63910212c3104c9f91a8f5ee1b7c64ebcc18286117d7cf1d6", size = 4344987, upload-time = "2026-10-26T12:31:14.729Z" }, +] + +[[package]] +name = "pydantic-monty-runtime" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/be/e7/1567009ab17fe091caba156ea9406487520ef382543e73ac7354a4eef88a/pydantic_monty_runtime-1.4.0.tar.gz", hash = "sha256:5275beb451201bd23669a253c8be93b90cad018b449fcdd0a1ad57daeed15187", size = 2581228, upload-time = "2026-10-26T12:31:53.171Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/6e/654bf4e3867e48564c7f4dabf03023660c860f53a0065915c4fa1e816825/pydantic_monty_runtime-1.4.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9ba7ad3163cbf601d7d36e5c9b6edb7b1625f27683a910885682d224edb8c05e", size = 12643236, upload-time = "2026-10-26T12:31:16.785Z" }, + { url = "https://files.pythonhosted.org/packages/cb/9b/0eedee19d738a168cb4fbcefb1be38cb7c729c750d0293379483e38fa71a/pydantic_monty_runtime-1.4.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:317edf6c62ea62c8a0f456ec13be889387cd16a917441fc087953cae042e6d75", size = 13763884, upload-time = "2026-10-26T12:31:19.94Z" }, + { url = "https://files.pythonhosted.org/packages/d6/91/613e52c0483bb6cbc92ad443450b6573b0d55924941e5235c72edfd08eef/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:fabe933bb157910aeb61b705c19500d8675b4033215bdb93bc7da1ee9a22a6c9", size = 14209091, upload-time = "2026-10-26T12:31:23.568Z" }, + { url = "https://files.pythonhosted.org/packages/35/cd/326b9d2333d3f18a1642ed0c6721f9454a1672a469b316c92942d764363e/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_armv7l.whl", hash = "sha256:b33e86f80adeb86f7928c5040602049d9a245d7cf08ff1398b2ee577fbf7d6fb", size = 13279002, upload-time = "2026-10-26T12:31:26.429Z" }, + { url = "https://files.pythonhosted.org/packages/66/d6/b9ab4adaa305b2aaa5624a9ac5b9a21fb65312c44f8dd658d336d2e20356/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_i686.whl", hash = "sha256:be875b3d248df50768ef9d2c420252e14458ec6f99a1ba2f10e34b0a2f7d88c1", size = 13509495, upload-time = "2026-10-26T12:31:28.996Z" }, + { url = "https://files.pythonhosted.org/packages/e9/e9/d812921782424db588b9772f57558df4ad8e58c65125d5101f511a08180d/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_ppc64le.whl", hash = "sha256:2be251cc1ee73337dbe3a7ea00d657e4a8385f2476523e74c5685d8f1fc2fea4", size = 13887394, upload-time = "2026-10-26T12:31:31.547Z" }, + { url = "https://files.pythonhosted.org/packages/20/60/2ebff28552652f00f4a8b7266570326f8e03383d44a9e06d492501cbaed5/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_s390x.whl", hash = "sha256:7e13192106dea260dc605a371ab44ffd248dc5afe0ab003f873779774c73a7e6", size = 14964165, upload-time = "2026-10-26T12:31:34.635Z" }, + { url = "https://files.pythonhosted.org/packages/d3/4a/fdf0ff01c3ee23f5af7a017e970967fc88d6d263a79db03d3dfce5c028a1/pydantic_monty_runtime-1.4.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:ff6301049a8e2f3407188dd174a749f2f2af6f82a2bc993c707dbebf44dcf651", size = 14726134, upload-time = "2026-10-26T12:31:37.607Z" }, + { url = "https://files.pythonhosted.org/packages/80/85/317a61daca57f3211d53f868977dbd1d491e0212082f0e79cd6df50137de/pydantic_monty_runtime-1.4.0-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:ae4e68d623b9038f8a41baa95e06e8d32402ad972120b068d14c336bb924eb03", size = 12315943, upload-time = "2026-10-26T12:31:40.22Z" }, + { url = "https://files.pythonhosted.org/packages/50/c4/fc238425268af5ca9642d5ae997cb1101c95b87e7876ceabf8275b42a1d3/pydantic_monty_runtime-1.4.0-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:af1f5d41c32207a2ecd9a9c4277290911c6e5f37299588097c7ff83240bb665c", size = 13021113, upload-time = "2026-10-26T12:31:42.899Z" }, + { url = "https://files.pythonhosted.org/packages/24/be/2d221328cd4017b3cfefb6c197740e56ab02689ac1300ad883450a15d565/pydantic_monty_runtime-1.4.0-py3-none-win32.whl", hash = "sha256:30c393985843e836b32e1ceb18c251661e54984bb1e992b7a88d5e897679af01", size = 12392040, upload-time = "2026-10-26T12:31:45.422Z" }, + { url = "https://files.pythonhosted.org/packages/5f/b2/5ac23283be20db8b127c42bb1072910d809792c98cb5b626a85e6ffdcc09/pydantic_monty_runtime-1.4.0-py3-none-win_amd64.whl", hash = "sha256:c2f5a7503288087669f7c5c6131dc9b68a26dfabf07d5fb36df33ff475afbe61", size = 13250757, upload-time = "2026-10-26T12:31:48.107Z" }, +] + [[package]] name = "pydantic-settings" version = "2.14.2"