diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_workflow_execution/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_workflow_execution/migration.sql new file mode 100644 index 00000000000..f04c059d82e --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_workflow_execution/migration.sql @@ -0,0 +1,87 @@ +-- AlterTable +ALTER TABLE "LiteLLM_WorkflowRun" ADD COLUMN "engine_owned" BOOLEAN NOT NULL DEFAULT false; + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowDefinition" ( + "workflow_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "etag" TEXT NOT NULL, + "bundle_json" TEXT NOT NULL, + "contributors" TEXT[], + "published_snapshot_id" TEXT, + + CONSTRAINT "LiteLLM_WorkflowDefinition_pkey" PRIMARY KEY ("workflow_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowSnapshot" ( + "snapshot_id" TEXT NOT NULL, + "workflow_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "bundle_json" TEXT NOT NULL, + "published_by" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "origin" TEXT NOT NULL DEFAULT 'publication', + + CONSTRAINT "LiteLLM_WorkflowSnapshot_pkey" PRIMARY KEY ("snapshot_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowExecution" ( + "run_id" TEXT NOT NULL, + "snapshot_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "principal" TEXT NOT NULL, + "dispatch_key" TEXT NOT NULL, + "request_digest" TEXT NOT NULL, + "request_json" TEXT NOT NULL, + "state" TEXT NOT NULL DEFAULT 'pending', + "deadline" TIMESTAMP(3) NOT NULL, + "lease_owner" TEXT, + "lease_until" TIMESTAMP(3), + "cancel_requested" BOOLEAN NOT NULL DEFAULT false, + "checkpoints" JSONB NOT NULL DEFAULT '[]', + "response_json" TEXT, + "error" TEXT, + + CONSTRAINT "LiteLLM_WorkflowExecution_pkey" PRIMARY KEY ("run_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowStep" ( + "step_id" TEXT NOT NULL, + "run_id" TEXT NOT NULL, + "node_id" TEXT NOT NULL, + "state" TEXT NOT NULL DEFAULT 'running', + "input" JSONB NOT NULL, + "output" JSONB, + "attempts" INTEGER NOT NULL DEFAULT 0, + "deadline" TIMESTAMP(3) NOT NULL, + "started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "finished_at" TIMESTAMP(3), + "error" TEXT, + "prompt_tokens" INTEGER NOT NULL DEFAULT 0, + "completion_tokens" INTEGER NOT NULL DEFAULT 0, + "response_cost" DOUBLE PRECISION, + "request_id" TEXT, + + CONSTRAINT "LiteLLM_WorkflowStep_pkey" PRIMARY KEY ("step_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowDefinition_team_id_name_key" ON "LiteLLM_WorkflowDefinition"("team_id", "name"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowSnapshot_workflow_id_idx" ON "LiteLLM_WorkflowSnapshot"("workflow_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowExecution_dispatch_key_key" ON "LiteLLM_WorkflowExecution"("dispatch_key"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowExecution_team_id_name_idx" ON "LiteLLM_WorkflowExecution"("team_id", "name"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowStep_run_id_node_id_key" ON "LiteLLM_WorkflowStep"("run_id", "node_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 818f2dc4858..dd8a4039b35 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1986,6 +1986,7 @@ model LiteLLM_WorkflowRun { input Json? output Json? metadata Json? + engine_owned Boolean @default(false) events LiteLLM_WorkflowEvent[] messages LiteLLM_WorkflowMessage[] @@ -2003,6 +2004,72 @@ model LiteLLM_WorkflowRun { // step.failed → run.status = failed // hook.waiting → run.status = paused // hook.received → run.status = running +model LiteLLM_WorkflowDefinition { + workflow_id String @id + team_id String + name String + etag String + bundle_json String + contributors String[] + published_snapshot_id String? + + @@unique([team_id, name]) +} + +model LiteLLM_WorkflowSnapshot { + snapshot_id String @id + workflow_id String + team_id String + name String + bundle_json String + published_by String + created_at DateTime @default(now()) + origin String @default("publication") + + @@index([workflow_id]) +} + +model LiteLLM_WorkflowExecution { + run_id String @id + snapshot_id String + team_id String + name String + principal String + dispatch_key String @unique + request_digest String + request_json String + state String @default("pending") + deadline DateTime + lease_owner String? + lease_until DateTime? + cancel_requested Boolean @default(false) + checkpoints Json @default("[]") + response_json String? + error String? + + @@index([team_id, name]) +} + +model LiteLLM_WorkflowStep { + step_id String @id + run_id String + node_id String + state String @default("running") + input Json + output Json? + attempts Int @default(0) + deadline DateTime + started_at DateTime @default(now()) + finished_at DateTime? + error String? + prompt_tokens Int @default(0) + completion_tokens Int @default(0) + response_cost Float? + request_id String? + + @@unique([run_id, node_id]) +} + model LiteLLM_WorkflowEvent { event_id String @id @default(uuid()) run_id String diff --git a/litellm/models/workflow.py b/litellm/models/workflow.py new file mode 100644 index 00000000000..a9cdca83624 --- /dev/null +++ b/litellm/models/workflow.py @@ -0,0 +1,164 @@ +from datetime import datetime +from enum import Enum +from typing import Protocol + +from pydantic import ConfigDict, JsonValue + +from litellm.types.llms.base import LiteLLMBaseModel + + +class WorkflowRunState(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class WorkflowSnapshotOrigin(str, Enum): + PUBLICATION = "publication" + DRAFT_TEST = "draft_test" + + +class WorkflowDraft(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + workflow_id: str + team_id: str + name: str + etag: str + bundle_json: str + contributors: tuple[str, ...] + published_snapshot_id: str | None + + +class WorkflowSnapshot(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + snapshot_id: str + workflow_id: str + team_id: str + name: str + bundle_json: str + published_by: str + created_at: datetime + origin: WorkflowSnapshotOrigin = WorkflowSnapshotOrigin.PUBLICATION + + +class WorkflowCheckpoint(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + node_path: str + input: JsonValue + output: JsonValue + attempts: int + + +class WorkflowStep(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + step_id: str + run_id: str + node_id: str + state: WorkflowRunState + input: JsonValue + output: JsonValue | None = None + attempts: int + 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 + + +class WorkflowExecution(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + run_id: str + snapshot_id: str + team_id: str + name: str + principal: str + dispatch_key: str + request_digest: str + request_json: str + state: WorkflowRunState + deadline: datetime + lease_owner: str | None = None + lease_until: datetime | None = None + cancel_requested: bool + checkpoints: tuple[WorkflowCheckpoint, ...] = () + response_json: str | None = None + error: str | None = None + + +class WorkflowConflict(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + message: str + status_code: int = 409 + + +class WorkflowStore(Protocol): + async def create( + self, team_id: str, name: str, bundle_json: str, author: str + ) -> WorkflowDraft | WorkflowConflict: ... + + async def get(self, team_id: str, name: str) -> WorkflowDraft | None: ... + + async def get_id(self, workflow_id: str) -> WorkflowDraft | None: ... + + async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]: ... + + async def replace( + self, team_id: str, name: str, etag: str, bundle_json: str, author: str + ) -> WorkflowDraft | WorkflowConflict: ... + + async def publish( + self, team_id: str, name: str, etag: str, publisher: str + ) -> WorkflowSnapshot | WorkflowConflict: ... + + async def snapshot(self, snapshot_id: str) -> WorkflowSnapshot | None: ... + + async def pin_draft( + self, workflow_id: str, etag: str, requested_by: str + ) -> WorkflowSnapshot | WorkflowConflict: ... + + async def published(self, team_id: str, name: str) -> WorkflowSnapshot | None: ... + + async def dispatch(self, execution: WorkflowExecution) -> WorkflowExecution: ... + + async def by_dispatch_key(self, dispatch_key: str) -> WorkflowExecution | None: ... + + async def execution(self, run_id: str) -> WorkflowExecution | None: ... + + async def lease_next(self, owner: str) -> WorkflowExecution | None: ... + + async def renew(self, run_id: str, owner: str) -> bool: ... + + async def steps(self, run_id: str) -> tuple[WorkflowStep, ...]: ... + + async def start_step( + self, run_id: str, node_id: str, input: JsonValue, deadline: datetime, owner: str + ) -> WorkflowStep | WorkflowConflict: ... + + async def attempt_step(self, step_id: str, owner: str) -> WorkflowStep | WorkflowConflict: ... + + async def finish_step(self, step: WorkflowStep, owner: str) -> bool: ... + + async def checkpoint(self, run_id: str, checkpoint: WorkflowCheckpoint, step: WorkflowStep, owner: str) -> bool: ... + + async def finish( + self, + run_id: str, + state: WorkflowRunState, + response_json: str | None = None, + error: str | None = None, + *, + owner: str, + ) -> bool: ... + + async def cancel(self, run_id: str) -> None: ... diff --git a/litellm/models/workflow_bundle.py b/litellm/models/workflow_bundle.py new file mode 100644 index 00000000000..6987164b756 --- /dev/null +++ b/litellm/models/workflow_bundle.py @@ -0,0 +1,157 @@ +import hashlib +import json +from collections.abc import Mapping +from enum import Enum +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import ConfigDict, Field, JsonValue + +from litellm.proxy.workflows.monty.models import PythonProgram +from litellm.types.llms.base import LiteLLMBaseModel + + +class NodeType(str, Enum): + CHAT_COMPLETION = "chat_completion" + PYTHON = "python" + + +class BindingOutput(str, Enum): + MESSAGES = "messages" + RESULT = "result" + + +class Reference(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + source: str = Field(alias="from", min_length=1, max_length=128) + output: BindingOutput + path: str | None = Field(default=None, max_length=512) + + +class Constant(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + value: JsonValue + + +Binding: TypeAlias = Reference | Constant + + +class RetryPolicy(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + max_attempts: int = Field(default=1, ge=1, le=3) + backoff_seconds: float = Field(default=0, ge=0, le=10, allow_inf_nan=False) + + +class ResponseSchema(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + name: str = Field(pattern=r"^[A-Za-z0-9_-]{1,64}$") + strict: bool = True + schema_: Mapping[str, JsonValue] = Field(alias="schema") + + +class StructuredResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["json_schema"] = "json_schema" + json_schema: ResponseSchema + + +class ChatCompletionNode(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + id: str = Field(pattern=r"^[A-Za-z][A-Za-z0-9_-]{0,63}$") + type: Literal[NodeType.CHAT_COMPLETION] + model: str = Field(min_length=1, max_length=256) + messages: Binding + temperature: float | None = Field(default=None, ge=0, le=2, allow_inf_nan=False) + top_p: float | None = Field(default=None, ge=0, le=1, allow_inf_nan=False) + max_tokens: int | None = Field(default=None, gt=0) + max_completion_tokens: int | None = Field(default=None, gt=0) + seed: int | None = None + stop: str | tuple[str, ...] | None = None + response_format: StructuredResponse | None = None + timeout_seconds: float = Field(default=60, gt=0, le=300, allow_inf_nan=False) + retry: RetryPolicy = RetryPolicy() + + +class PythonNode(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + id: str = Field(pattern=r"^[A-Za-z][A-Za-z0-9_-]{0,63}$") + type: Literal[NodeType.PYTHON] + program: PythonProgram + inputs: Mapping[str, Binding] + timeout_seconds: float = Field(default=60, gt=0, le=300, allow_inf_nan=False) + retry: RetryPolicy = RetryPolicy() + + +WorkflowNode: TypeAlias = Annotated[ChatCompletionNode | PythonNode, Field(discriminator="type")] + + +class WorkflowBundle(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + contract_version: Literal["1"] + name: str = Field(pattern=r"^[a-z][a-z0-9-]{0,62}$") + nodes: tuple[WorkflowNode, ...] = Field(min_length=1, max_length=64) + timeout_seconds: float = Field(default=300, gt=0, le=900, allow_inf_nan=False) + input_schema: Mapping[str, JsonValue] + output: Reference + + +class WorkflowFailure(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + code: str + message: str + param: str | None = None + status_code: int = 400 + + +def canonical_bundle(bundle: WorkflowBundle) -> str: + return json.dumps( + bundle.model_dump(mode="json", by_alias=True, exclude_none=True), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + +def workflow_digest(bundle: WorkflowBundle) -> str: + return hashlib.sha256(canonical_bundle(bundle).encode()).hexdigest() + + +def reference_failure(reference: Reference, available: frozenset[str]) -> WorkflowFailure | None: + if reference.source == "$input" and reference.output is not BindingOutput.MESSAGES: + return WorkflowFailure(code="invalid_reference", message="Request input exposes only messages") + if reference.source != "$input" and ( + reference.source not in available or reference.output is not BindingOutput.RESULT + ): + return WorkflowFailure(code="invalid_reference", message="Reference must select a completed predecessor") + if reference.path is not None and not reference.path.startswith("/") and reference.path != "": + return WorkflowFailure(code="invalid_reference", message="path must be a JSON Pointer") + return None + + +def validate_bindings(bundle: WorkflowBundle) -> WorkflowFailure | None: + if bundle.name == "runs": + return WorkflowFailure(code="reserved_name", message="runs is reserved for run tracking") + ids: Final = tuple(node.id for node in bundle.nodes) + if len(set(ids)) != len(ids): + return WorkflowFailure(code="duplicate_node", message="Node IDs must be unique") + for index, node in enumerate(bundle.nodes): + bindings = (node.messages,) if isinstance(node, ChatCompletionNode) else tuple(node.inputs.values()) + for binding in bindings: + failure = reference_failure(binding, frozenset(ids[:index])) if isinstance(binding, Reference) else None + if failure is not None: + return failure.model_copy(update={"param": f"nodes.{node.id}"}) + if isinstance(node, ChatCompletionNode) and node.model.startswith("workflow/"): + return WorkflowFailure( + code="nested_workflow", message="Nested workflow calls are unsupported", param=node.id + ) + if isinstance(node, PythonNode) and set(node.inputs) != set(node.program.parameter_names): + return WorkflowFailure(code="invalid_inputs", message="Bind exactly the function parameters", param=node.id) + return reference_failure(bundle.output, frozenset(ids)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d3ac264518e..f82efb7dd78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2506,6 +2506,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + enable_workflows: bool = Field( + default=False, description="Enable durable typed workflows and the optional Monty worker" + ) + use_spend_logs_ledger: bool = Field( True, description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", @@ -2785,7 +2789,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) alert_to_webhook_url: dict | None = Field( None, - description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`", + description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': '$ALERT_WEBHOOK_URL'}`", ) alerting_args: dict | None = Field(None, description="Controllable params for slack alerting - e.g. ttl in cache.") alerting_threshold: int | None = Field( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af60a315c60..f9ba399cb1c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5008,6 +5008,9 @@ def _can_object_call_model( from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent compaction_parent: Final = native_compaction_parent(model) + from litellm.proxy.workflows.auth import workflow_parent + + workflow_model: Final = workflow_parent(model) global_or_router_alias_target: Final = ( litellm.model_alias_map[model] if model in litellm.model_alias_map @@ -5031,6 +5034,7 @@ def _can_object_call_model( if key_alias_applied else ( *((model, compaction_parent) if compaction_parent is not None else (model,)), + *((workflow_model,) if workflow_model is not None else ()), *((global_or_router_alias_target,) if global_or_router_alias_target else ()), ) ) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 125bfab4edf..84437512ed4 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -373,6 +373,8 @@ class RouteChecks: pass elif route == "/lens" or route.startswith("/lens/"): pass # Lens authorizes the delegated role, including its POST-based read endpoints. + elif RouteChecks.is_workflow_execution_route(route, RouteChecks.get_request_method(request=request)): + pass elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: RouteChecks._check_proxy_admin_viewer_access( route=route, @@ -753,6 +755,25 @@ class RouteChecks: _get_request_method = get_request_method + @staticmethod + def is_workflow_execution_route(route: str, method: str | None) -> bool: + run_patterns: Final = ( + ("GET", "/v1/workflows/runs/{run_id}/steps"), + ("GET", "/v1/workflows/runs/{run_id}/steps/{step_id}"), + ("POST", "/v1/workflows/runs/{run_id}/cancel"), + ) + definition_patterns: Final = ( + ("POST", "/v1/workflows"), + ("GET", "/v1/workflows"), + ("GET", "/v1/workflows/{workflow_id}"), + ("PUT", "/v1/workflows/{workflow_id}"), + ("POST", "/v1/workflows/{workflow_id}/validate"), + ("POST", "/v1/workflows/{workflow_id}/test"), + ("POST", "/v1/workflows/{workflow_id}/publish"), + ) + patterns: Final = run_patterns if route.startswith("/v1/workflows/runs") else definition_patterns + return any(method == verb and RouteChecks.check_route_access(route, (pattern,)) for verb, pattern in patterns) + @staticmethod def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: method: Final = RouteChecks.get_request_method(request) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..5eb37001282 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2696,6 +2696,12 @@ async def user_api_key_auth_builder( # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead ### CHECK IF ADMIN ### # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead + from litellm.proxy.workflows.auth import workflow_child_identity + + workflow_identity: Final = await workflow_child_identity(request, request_data) + if workflow_identity is not None: + valid_token = workflow_identity # rebind-ok: Existing auth state must use the verified child identity + api_key = workflow_identity.token or "" # rebind-ok: Pass the trusted principal to existing auth checks if valid_token is None: ## Check CACHE try: diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 323cf57dfec..2642980d851 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -7,12 +7,21 @@ 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 import WorkflowStore 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_repository import WorkflowRepository + + +async def get_workflow_repository(database: DatabaseDep) -> WorkflowStore | None: + return None if database is None else WorkflowRepository(database.prisma) + + +WorkflowRepositoryDep: TypeAlias = Annotated[WorkflowStore | None, Depends(get_workflow_repository)] async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d28401524ed..db58ef6ad89 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5034,7 +5034,7 @@ async def info_key_fn( default=None, description=( "Key to look up. Pass the key's sha256 hash so the raw key stays out of URLs and access " - "logs. Example key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + "logs. Example key='YOUR_VIRTUAL_KEY'" ), ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -5081,14 +5081,14 @@ async def info_key_fn( Example Curl: ``` - curl -X GET "http://0.0.0.0:4000/key/info?key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \ --H "Authorization: Bearer $LITELLM_MASTER_KEY" + curl -X GET "http://0.0.0.0:4000/key/info?key=YOUR_VIRTUAL_KEY" \ +-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" ``` Example Curl - if no key is passed, it will use the Key Passed in Authorization Header ``` curl -X GET "http://0.0.0.0:4000/key/info" \ --H "Authorization: Bearer sk-test-example-key-123" +-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" ``` """ from litellm.proxy.proxy_server import ( @@ -8076,7 +8076,7 @@ async def block_key( --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ --header 'Content-Type: application/json' \ --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + "key": "YOUR_VIRTUAL_KEY" }' ``` @@ -8190,7 +8190,7 @@ async def unblock_key( --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ --header 'Content-Type: application/json' \ --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + "key": "YOUR_VIRTUAL_KEY" }' ``` diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..0eb277346f9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5399,7 +5399,7 @@ async def team_info( ``` curl --location 'http://localhost:4000/team/info?team_id=your_team_id_here' \ - --header 'Authorization: Bearer your_api_key_here' + -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" ``` """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable @@ -5566,7 +5566,7 @@ async def team_member_me( ``` curl --location 'http://localhost:4000/team/your_team_id/members/me' \ - --header 'Authorization: Bearer your_api_key_here' + -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" ``` """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 649941bc233..6d03650eeb6 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -15,9 +15,11 @@ GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history import json from collections.abc import Mapping, Sequence +from datetime import datetime from typing import Final, Literal, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import ConfigDict, JsonValue try: from prisma.errors import UniqueViolationError @@ -102,10 +104,53 @@ class WorkflowMessageCreateRequest(LiteLLMBaseModel): session_id: str | None = None +class WorkflowRunResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + session_id: str + workflow_type: str + status: str + created_by: str | None + created_at: datetime + updated_at: datetime + input: JsonValue + output: JsonValue + metadata: JsonValue + engine_owned: bool + + +class WorkflowEventResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + event_id: str + run_id: str + event_type: str + step_name: str + sequence_number: int + data: JsonValue + created_at: datetime + + +class WorkflowMessageResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + message_id: str + run_id: str + role: str + content: str + sequence_number: int + session_id: str | None + created_at: datetime + + class _RunRow(Protocol): @property def created_by(self) -> str | None: ... + @property + def engine_owned(self) -> bool: ... + class _SeqRow(Protocol): @property @@ -201,6 +246,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 +257,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 run.engine_owned: + raise HTTPException(status_code=409, detail="Engine-owned runs cannot be modified through tracking APIs") return run @@ -220,6 +269,7 @@ async def _require_run( @router.post( "/v1/workflows/runs", + responses={200: {"model": WorkflowRunResponse}}, tags=["workflow management"], dependencies=[Depends(user_api_key_auth)], ) @@ -359,8 +409,7 @@ async def update_workflow_run( if not update: 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( @@ -379,6 +428,7 @@ async def update_workflow_run( @router.post( "/v1/workflows/runs/{run_id}/events", + responses={200: {"model": WorkflowEventResponse}}, tags=["workflow management"], dependencies=[Depends(user_api_key_auth)], ) @@ -398,7 +448,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) @@ -470,6 +520,7 @@ async def list_workflow_events( @router.post( "/v1/workflows/runs/{run_id}/messages", + responses={200: {"model": WorkflowMessageResponse}}, tags=["workflow management"], dependencies=[Depends(user_api_key_auth)], ) @@ -487,7 +538,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 c1af5d6f980..c69780f26f8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -683,6 +683,7 @@ from litellm.proxy.db.request_error_tracking import ( request_error_accumulator, ) from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed +from litellm.proxy.deps.repositories import WorkflowRepositoryDep from litellm.proxy.discovery_endpoints import ( agent_skills_discovery_router, ui_discovery_endpoints_router, @@ -887,6 +888,8 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware from litellm.proxy.telemetry.runtime import TelemetryRuntime from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings +from litellm.proxy.workflows.management import router as workflow_definition_router +from litellm.proxy.workflows.runs import router as workflow_execution_router from litellm.router import router_arguments from litellm.router_config import build_router_config @@ -1975,7 +1978,12 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) gc.freeze() verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) - yield state + from litellm.proxy.workflows.dispatcher import workflow_dispatch_lifespan + + async with workflow_dispatch_lifespan( + state, app, general_settings_view().get("enable_workflows") is True + ): + yield state finally: if model_info_scheduler is not None and model_info_scheduler.running: model_info_scheduler.remove_job("refresh_model_info") @@ -3137,6 +3145,10 @@ def general_settings_view() -> Mapping[str, object]: return _GENERAL_SETTINGS_VIEW.validate_python(general_settings) +def llm_model_list_view() -> tuple[object, ...] | None: + return TypeAdapter[tuple[object, ...] | None](tuple[object, ...] | None).validate_python(llm_model_list) + + config_passthrough_endpoints: list[dict[str, object]] | None = None log_file: Final = "api_log.json" worker_config: Final = None @@ -13129,6 +13141,7 @@ async def chat_completion( fastapi_response: Response, model: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + workflows: WorkflowRepositoryDep = None, ): """ @@ -13156,6 +13169,12 @@ 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 = _GENERAL_SETTINGS_VIEW.validate_python(data) + workflow_model: Final = workflow_data.get("model") + if isinstance(workflow_model, str) and workflow_model.startswith("workflow/"): + from litellm.proxy.workflows.invocation import invoke_workflow + + return await invoke_workflow(request, user_api_key_dict, workflow_data, workflows) if user_api_key_dict is not None: if not isinstance(data.get("metadata"), dict): # Covers both missing and JSON-string metadata (multipart / @@ -22432,6 +22451,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..76aeef7c0fa 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1986,6 +1986,7 @@ model LiteLLM_WorkflowRun { input Json? output Json? metadata Json? + engine_owned Boolean @default(false) events LiteLLM_WorkflowEvent[] messages LiteLLM_WorkflowMessage[] @@ -1996,6 +1997,72 @@ model LiteLLM_WorkflowRun { @@index([created_by]) } +model LiteLLM_WorkflowDefinition { + workflow_id String @id + team_id String + name String + etag String + bundle_json String + contributors String[] + published_snapshot_id String? + + @@unique([team_id, name]) +} + +model LiteLLM_WorkflowSnapshot { + snapshot_id String @id + workflow_id String + team_id String + name String + bundle_json String + published_by String + created_at DateTime @default(now()) + origin String @default("publication") + + @@index([workflow_id]) +} + +model LiteLLM_WorkflowExecution { + run_id String @id + snapshot_id String + team_id String + name String + principal String + dispatch_key String @unique + request_digest String + request_json String + state String @default("pending") + deadline DateTime + lease_owner String? + lease_until DateTime? + cancel_requested Boolean @default(false) + checkpoints Json @default("[]") + response_json String? + error String? + + @@index([team_id, name]) +} + +model LiteLLM_WorkflowStep { + step_id String @id + run_id String + node_id String + state String @default("running") + input Json + output Json? + attempts Int @default(0) + deadline DateTime + started_at DateTime @default(now()) + finished_at DateTime? + error String? + prompt_tokens Int @default(0) + completion_tokens Int @default(0) + response_cost Float? + request_id String? + + @@unique([run_id, node_id]) +} + // Append-only log of state transitions. Never mutate rows here. // `step_name` and `event_type` are caller-defined strings — no hardcoded enums. // Status auto-update rules (applied by the append endpoint): diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0348d8e1f82..9f4db00e5ec 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1470,7 +1470,7 @@ async def get_global_spend_report( default=None, description=( "View spend for a specific api_key. Pass the key's sha256 hash so the raw key stays " - "out of URLs and access logs. Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + "out of URLs and access logs. Example api_key='YOUR_VIRTUAL_KEY'" ), ), internal_user_id: str | None = fastapi.Query( @@ -1945,7 +1945,7 @@ async def get_key_spend_report( description=( "View spend for a specific api_key. Proxy admin only; other callers are scoped to their " "own key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. " - "Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + "Example api_key='YOUR_VIRTUAL_KEY'" ) ), ] = None, @@ -3447,7 +3447,7 @@ async def view_spend_logs( Example Request for specific api_key ``` - curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \ + curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=YOUR_VIRTUAL_KEY" \ -H "Authorization: Bearer $LITELLM_MASTER_KEY" ``` diff --git a/litellm/proxy/workflows/auth.py b/litellm/proxy/workflows/auth.py new file mode 100644 index 00000000000..18b7ee49690 --- /dev/null +++ b/litellm/proxy/workflows/auth.py @@ -0,0 +1,64 @@ +from collections.abc import Mapping +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Final + +from fastapi import Request + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.resolvers.store import IdentityStore + + +@dataclass(frozen=True, slots=True) +class WorkflowChildGrant: + principal: str + team_id: str + snapshot_id: str + workflow_model: str + node_id: str + model: str + run_id: str + + +workflow_child_grant: Final[ContextVar[WorkflowChildGrant | None]] = ContextVar("workflow_child_grant", default=None) + + +def workflow_parent(model: str) -> str | None: + grant: Final = workflow_child_grant.get() + return grant.workflow_model if grant is not None and grant.model == model else None + + +async def workflow_child_identity(request: Request, request_data: Mapping[str, object]) -> UserAPIKeyAuth | None: + grant: Final = workflow_child_grant.get() + if grant is None: + return None + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if ( + request.method != "POST" + or not request.url.path.endswith("/v1/chat/completions") + or request_data.get("model") != grant.model + ): + from fastapi import HTTPException + + raise HTTPException(status_code=403, detail="Invalid internal workflow child scope") + identity: Final = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=grant.principal) + ) + if ( + identity.team_id != grant.team_id + or identity.is_session_token + or identity.user_role + in { + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + } + ): + from fastapi import HTTPException + + raise HTTPException(status_code=403, detail="Workflow requires its original team virtual key") + return identity.model_copy(update={"via_virtual_key": True}) diff --git a/litellm/proxy/workflows/child.py b/litellm/proxy/workflows/child.py new file mode 100644 index 00000000000..3b9a0f3ac90 --- /dev/null +++ b/litellm/proxy/workflows/child.py @@ -0,0 +1,168 @@ +from typing import Annotated, Final + +from openai.types.chat import ChatCompletionMessageParam +from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from starlette.types import ASGIApp, Receive, Scope, Send + +from litellm.llms.custom_httpx.asgi_handler import get_async_asgi_client +from litellm.models.workflow import WorkflowExecution +from litellm.models.workflow_bundle import ChatCompletionNode, WorkflowFailure +from litellm.proxy._types import ProxyLifespanState +from litellm.proxy.workflows.auth import WorkflowChildGrant, workflow_child_grant +from litellm.proxy.workflows.runner import ChildResult +from litellm.proxy.workflows.validation import validate_value +from litellm.types.llms.base import LiteLLMBaseModel + + +class _Message(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + content: str | None = None + refusal: str | None = None + + +class _Choice(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + finish_reason: str + message: _Message + + +class _Usage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + prompt_tokens: int = 0 + completion_tokens: int = 0 + + +class _Completion(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + choices: tuple[_Choice, ...] + usage: _Usage = _Usage() + + +_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam]) +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST: Final = TypeAdapter[float | None](Annotated[float | None, Field(ge=0, allow_inf_nan=False)]) + + +class GatewayChildCaller: + def __init__(self, app: ASGIApp, state: ProxyLifespanState) -> None: + self._app: Final = app + self._state: Final = state + + async def dispatch(self, scope: Scope, receive: Receive, send: Send) -> None: + await self._app({**scope, "state": {**self._state}}, receive, send) + + async def call( + self, execution: WorkflowExecution, node: ChatCompletionNode, messages: JsonValue, remaining: float + ) -> ChildResult | WorkflowFailure: + try: + validated_messages: Final = _MESSAGES.validate_python(messages) + except ValidationError: + return WorkflowFailure( + code="invalid_messages", + message="Resolved messages must be standard Chat Completions messages", + param=node.id, + ) + original_request: Final = _OBJECT.validate_json(execution.request_json) + payload: Final = { + **node.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + exclude={"id", "type", "messages", "retry", "timeout_seconds"}, + ), + "messages": _MESSAGES.dump_python(validated_messages, mode="json"), + "stream": False, + "n": 1, + "timeout": remaining, + **({"user": original_request["user"]} if "user" in original_request else {}), + "metadata": { + "workflow_run_id": execution.run_id, + "workflow_node_id": node.id, + "workflow_snapshot_id": execution.snapshot_id, + }, + } + token: Final = workflow_child_grant.set( + WorkflowChildGrant( + principal=execution.principal, + team_id=execution.team_id, + snapshot_id=execution.snapshot_id, + workflow_model=f"workflow/{execution.name}", + node_id=node.id, + model=node.model, + run_id=execution.run_id, + ) + ) + try: + with get_async_asgi_client(app=self.dispatch) as client: + response: Final = await client.post( + "http://litellm.internal/v1/chat/completions", + json=payload, + headers={"Authorization": "Bearer sk-workflow-internal", "x-litellm-session-id": execution.run_id}, + timeout=remaining, + ) + if not response.is_success: + return WorkflowFailure( + code="child_call_failed", + message=f"Gateway child returned HTTP {response.status_code}", + param=node.id, + status_code=response.status_code, + ) + try: + completion: Final = _Completion.model_validate_json(response.content) + response_cost: Final = _COST.validate_python(response.headers.get("x-litellm-response-cost")) + except ValidationError: + return WorkflowFailure( + code="invalid_completion", + message="Gateway child returned an invalid completion", + param=node.id, + status_code=502, + ) + if len(completion.choices) != 1: + return WorkflowFailure( + code="invalid_completion", + message="Workflow nodes require one completion choice", + param=node.id, + status_code=502, + ) + choice: Final = completion.choices[0] + if choice.finish_reason != "stop" or choice.message.refusal is not None or choice.message.content is None: + return WorkflowFailure( + code="invalid_completion", + message="Completion refused, truncated, or produced no text result", + param=node.id, + status_code=502, + ) + if node.response_format is None: + return ChildResult( + output=choice.message.content, + prompt_tokens=completion.usage.prompt_tokens, + completion_tokens=completion.usage.completion_tokens, + request_id=completion.id, + response_cost=response_cost, + ) + try: + output: Final = _JSON.validate_json(choice.message.content) + except ValidationError: + return WorkflowFailure( + code="invalid_json", message="Completion content is not valid JSON", param=node.id, status_code=502 + ) + failure: Final = validate_value(output, node.response_format.json_schema.schema_, node.id) + return ( + failure + if failure is not None + else ChildResult( + output=output, + prompt_tokens=completion.usage.prompt_tokens, + completion_tokens=completion.usage.completion_tokens, + request_id=completion.id, + response_cost=response_cost, + ) + ) + finally: + workflow_child_grant.reset(token) diff --git a/litellm/proxy/workflows/dispatcher.py b/litellm/proxy/workflows/dispatcher.py new file mode 100644 index 00000000000..55a9954f70e --- /dev/null +++ b/litellm/proxy/workflows/dispatcher.py @@ -0,0 +1,158 @@ +import asyncio +import json +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Final + +from pydantic import JsonValue +from starlette.types import ASGIApp + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.models.workflow import WorkflowExecution, WorkflowRunState, WorkflowStore +from litellm.models.workflow_bundle import WorkflowFailure +from litellm.proxy._types import ProxyLifespanState +from litellm.proxy.workflows.child import GatewayChildCaller +from litellm.proxy.workflows.monty.runtime import MontyRuntime +from litellm.proxy.workflows.runner import ChildCaller, execute_run +from litellm.proxy.workflows.validation import DEFAULT_MONTY_LIMITS +from litellm.repositories.workflow_repository import WorkflowRepository + + +async def maintain_lease( + store: WorkflowStore, execution: WorkflowExecution, work: asyncio.Task[JsonValue | WorkflowFailure] +) -> None: + try: + while not work.done(): + await asyncio.sleep(2) + current = await store.execution(execution.run_id) + if ( + current is None + or current.cancel_requested + or not await store.renew(execution.run_id, execution.lease_owner or "") + ): + work.cancel() + return + except Exception: # noqa: BLE001 # Any persistence failure must stop the worker before its lease expires + work.cancel() + verbose_proxy_logger.exception("Workflow lease renewal failed; stopping dispatch") + + +async def format_response(store: WorkflowStore, execution: WorkflowExecution, result: JsonValue) -> str: + steps: Final = await store.steps(execution.run_id) + prompt_tokens: Final = sum(step.prompt_tokens for step in steps) + completion_tokens: Final = sum(step.completion_tokens for step in steps) + return json.dumps( + { + "id": f"chatcmpl-{execution.run_id}", + "object": "chat.completion", + "created": int(datetime.now(timezone.utc).timestamp()), + "model": f"workflow/{execution.name}", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": result if isinstance(result, str) else json.dumps(result, allow_nan=False), + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + }, + allow_nan=False, + ) + + +async def dispatch_one( + store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, execution: WorkflowExecution +) -> None: + work: Final = asyncio.create_task(execute_run(store, caller, runtime, execution)) + heartbeat: Final = asyncio.create_task(maintain_lease(store, execution, work)) + try: + result: Final = await work + if isinstance(result, WorkflowFailure): + await store.finish( + execution.run_id, + WorkflowRunState.CANCELLED if result.code == "cancelled" else WorkflowRunState.FAILED, + error=result.model_dump_json(), + owner=execution.lease_owner or "", + ) + return + await store.finish( + execution.run_id, + WorkflowRunState.COMPLETED, + response_json=await format_response(store, execution, result), + owner=execution.lease_owner or "", + ) + except asyncio.CancelledError: + current: Final = await store.execution(execution.run_id) + if current is not None and current.cancel_requested and current.lease_owner == execution.lease_owner: + await store.finish(execution.run_id, WorkflowRunState.CANCELLED, owner=execution.lease_owner or "") + elif (task := asyncio.current_task()) is not None and task.cancelling(): + raise + finally: + work.cancel() + heartbeat.cancel() + await asyncio.gather(work, heartbeat, return_exceptions=True) + + +async def worker(store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, owner: str) -> None: + while True: + try: + execution = await store.lease_next(f"{owner}/{uuid.uuid4()}") + if execution is None: + await asyncio.sleep(0.25) + continue + await dispatch_one(store, caller, runtime, execution) + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 # Durable unfinished runs must remain recoverable after unexpected failures + verbose_proxy_logger.exception("Workflow dispatcher failed; unfinished work remains lease-recoverable") + await asyncio.sleep(1) + + +@dataclass(frozen=True, slots=True) +class WorkflowDispatcher: + tasks: tuple[asyncio.Task[None], ...] + + async def aclose(self) -> None: + for task in self.tasks: + task.cancel() + await asyncio.gather(*self.tasks, return_exceptions=True) + + +def start_dispatcher( + store: WorkflowStore, app: ASGIApp, state: ProxyLifespanState, runtime: MontyRuntime +) -> WorkflowDispatcher: + caller: Final = GatewayChildCaller(app, state) + owner: Final = str(uuid.uuid4()) + return WorkflowDispatcher( + tasks=tuple(asyncio.create_task(worker(store, caller, runtime, f"{owner}/{index}")) for index in range(4)) + ) + + +@asynccontextmanager +async def workflow_dispatch_lifespan( + state: ProxyLifespanState, app: ASGIApp, enabled: bool +) -> AsyncGenerator[None, None]: + database: Final = state["database"] + if not enabled or database is None: + yield + return + runtime: Final = MontyRuntime(DEFAULT_MONTY_LIMITS) + failure: Final = await runtime.start() + if failure is not None: + raise RuntimeError(f"Cannot enable workflows, sandbox unavailable ({failure.kind.value})") + dispatcher: Final = start_dispatcher(WorkflowRepository(database.prisma), app, state, runtime) + try: + yield + finally: + await dispatcher.aclose() + await runtime.close() diff --git a/litellm/proxy/workflows/invocation.py b/litellm/proxy/workflows/invocation.py new file mode 100644 index 00000000000..71b48e7516f --- /dev/null +++ b/litellm/proxy/workflows/invocation.py @@ -0,0 +1,219 @@ +import asyncio +import hashlib +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final, Literal + +from fastapi import Request +from openai.types.chat import ChatCompletionMessageParam +from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from starlette.responses import JSONResponse + +from litellm._uuid import uuid +from litellm.models.workflow import WorkflowExecution, WorkflowRunState, WorkflowSnapshot, WorkflowStore +from litellm.models.workflow_bundle import WorkflowFailure +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import hash_token +from litellm.proxy.workflows.management import raise_public, require_store +from litellm.proxy.workflows.validation import canonical_request +from litellm.types.llms.base import LiteLLMBaseModel + + +class InvocationMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + tags: tuple[str, ...] = Field(default=(), max_length=20) + + +class WorkflowInvocation(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + model: str = Field(pattern=r"^workflow/[a-z][a-z0-9-]{0,62}$") + messages: tuple[ChatCompletionMessageParam, ...] = Field(min_length=1, max_length=128) + user: str | None = Field(default=None, max_length=256) + metadata: InvocationMetadata | None = None + stream: Literal[False] = False + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_RUN_HEADER: Final = "x-litellm-workflow-run-id" + + +def principal_hash(auth: UserAPIKeyAuth) -> str: + from litellm.proxy.proxy_server import user_custom_auth + + if ( + not auth.token + or not auth.team_id + or not auth.via_virtual_key + or auth.is_session_token + or user_custom_auth is not None + ): + raise_public( + WorkflowFailure( + code="virtual_key_required", message="Workflow invocation requires a team virtual key", status_code=403 + ) + ) + return hash_token(auth.token) if auth.token.startswith("sk-") else auth.token + + +async def admit( + store: WorkflowStore, + auth: UserAPIKeyAuth, + invocation: WorkflowInvocation, + idempotency_key: str | None, + *, + snapshot: WorkflowSnapshot | None = None, + operation: str = "invoke", +) -> WorkflowExecution: + principal: Final = principal_hash(auth) + if snapshot is not None and snapshot.team_id != auth.team_id: + raise_public( + WorkflowFailure( + code="workflow_team_mismatch", + message="Workflow execution requires a virtual key for its owning team", + status_code=403, + ) + ) + if idempotency_key is not None and not 1 <= len(idempotency_key) <= 256: + raise_public( + WorkflowFailure(code="invalid_idempotency_key", message="Idempotency-Key must contain 1 to 256 characters") + ) + request_json: Final = canonical_request( + _JSON_OBJECT.validate_python(invocation.model_dump(mode="json", exclude_none=True)) + ) + digest: Final = hashlib.sha256(request_json.encode()).hexdigest() + name: Final = invocation.model.removeprefix("workflow/") + dispatch_key: Final = hashlib.sha256( + canonical_request( + { + "principal": principal, + "team": auth.team_id, + "operation": operation, + "workflow": name, + "idempotency_key": idempotency_key if idempotency_key is not None else str(uuid.uuid4()), + } + ).encode() + ).hexdigest() + existing: Final = await store.by_dispatch_key(dispatch_key) + if existing is not None: + if existing.request_digest != digest: + raise_public( + WorkflowFailure( + code="idempotency_conflict", + message="Idempotency-Key was accepted with a different request", + status_code=409, + ) + ) + return existing + pinned: Final = snapshot if snapshot is not None else await store.published(auth.team_id or "", name) + if pinned is None: + raise_public(WorkflowFailure(code="not_found", message="Published workflow not found", status_code=404)) + from litellm.models.workflow_bundle import WorkflowBundle + + bundle: Final = WorkflowBundle.model_validate_json(pinned.bundle_json) + accepted: Final = await store.dispatch( + WorkflowExecution( + run_id=str(uuid.uuid4()), + snapshot_id=pinned.snapshot_id, + team_id=pinned.team_id, + name=name, + principal=principal, + dispatch_key=dispatch_key, + request_digest=digest, + request_json=request_json, + state=WorkflowRunState.PENDING, + deadline=datetime.now(timezone.utc) + timedelta(seconds=bundle.timeout_seconds), + cancel_requested=False, + ) + ) + if accepted.request_digest != digest: + raise_public( + WorkflowFailure( + code="idempotency_conflict", + message="Idempotency-Key was accepted with a different request", + status_code=409, + ) + ) + return accepted + + +def error_response(failure: WorkflowFailure, run_id: str) -> JSONResponse: + return JSONResponse( + status_code=failure.status_code, + content={ + "error": { + "type": "workflow_error", + "code": failure.code, + "message": failure.message, + "param": failure.param, + } + }, + headers={_RUN_HEADER: run_id}, + ) + + +async def wait_for_run(store: WorkflowStore, execution: WorkflowExecution, wait_seconds: float = 60.0) -> JSONResponse: + try: + async with asyncio.timeout(wait_seconds): + while True: + current = await store.execution(execution.run_id) + if current is None: + return error_response( + WorkflowFailure( + code="run_unavailable", message="Accepted run is temporarily unavailable", status_code=503 + ), + execution.run_id, + ) + if current.state is WorkflowRunState.COMPLETED and current.response_json is not None: + return JSONResponse( + content=_JSON_OBJECT.validate_json(current.response_json), headers={_RUN_HEADER: current.run_id} + ) + if current.state is WorkflowRunState.FAILED: + failure = ( + WorkflowFailure.model_validate_json(current.error) + if current.error is not None + else WorkflowFailure( + code="execution_failed", + message="Workflow execution failed", + status_code=502, + ) + ) + return error_response(failure, current.run_id) + if current.state is WorkflowRunState.CANCELLED: + return error_response( + WorkflowFailure(code="cancelled", message="Run was cancelled", status_code=409), current.run_id + ) + await asyncio.sleep(0.05) + except TimeoutError: + return error_response( + WorkflowFailure( + code="wait_timeout", + message="Run continues; retry with the same Idempotency-Key or inspect it", + status_code=504, + ), + execution.run_id, + ) + + +async def invoke_workflow( + request: Request, auth: UserAPIKeyAuth, data: Mapping[str, object], workflows: WorkflowStore | None +) -> JSONResponse: + from litellm.proxy.proxy_server import general_settings_view + + if general_settings_view().get("enable_workflows") is not True: + raise_public( + WorkflowFailure(code="workflow_disabled", message="Workflow execution is disabled", status_code=503) + ) + try: + invocation: Final = WorkflowInvocation.model_validate(data) + except ValidationError: + raise_public( + WorkflowFailure( + code="unsupported_workflow_request", + message="Workflow requests accept messages, user, metadata.tags and non-streaming invocation only", + ) + ) + store: Final = require_store(workflows) + accepted: Final = await admit(store, auth, invocation, request.headers.get("Idempotency-Key")) + return await wait_for_run(store, accepted) diff --git a/litellm/proxy/workflows/management.py b/litellm/proxy/workflows/management.py new file mode 100644 index 00000000000..a2679706ba0 --- /dev/null +++ b/litellm/proxy/workflows/management.py @@ -0,0 +1,273 @@ +from typing import Annotated, Final, NoReturn, TypeAlias + +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Response +from pydantic import ConfigDict + +from litellm.models.workflow import WorkflowConflict, WorkflowDraft, WorkflowStore +from litellm.models.workflow_bundle import ( + ChatCompletionNode, + WorkflowBundle, + WorkflowFailure, + canonical_bundle, + workflow_digest, +) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, get_team_object +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import WorkflowRepositoryDep +from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, is_team_member +from litellm.proxy.management.teams.dependencies import get_team_access +from litellm.proxy.workflows.validation import validate_bundle +from litellm.types.llms.base import LiteLLMBaseModel + +router: Final = APIRouter(tags=["Workflows"]) +Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] + + +class EmptyWorkflowAction(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class RegisterWorkflow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + bundle: WorkflowBundle + + +class PublishWorkflow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + digest: str + + +class DraftView(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + workflow_id: str + name: str + team_id: str + etag: str + digest: str + bundle: WorkflowBundle + published_snapshot_id: str | None + + +class DraftList(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + workflows: tuple[DraftView, ...] + next_cursor: str | None + + +class ValidationView(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + valid: bool + digest: str + + +class PublicationView(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + snapshot_id: str + model: str + digest: str + + +def raise_public(failure: WorkflowFailure) -> NoReturn: + raise HTTPException( + status_code=failure.status_code, + detail={ + "type": "workflow_error", + "code": failure.code, + "message": failure.message, + "param": failure.param, + }, + ) + + +def require_store(store: WorkflowStore | None) -> WorkflowStore: + if store is None: + raise_public( + WorkflowFailure(code="database_required", message="Workflow execution requires a database", status_code=503) + ) + return store + + +async def require_team(auth: UserAPIKeyAuth, team_id: str | None, *, admin: bool = False, write: bool = False) -> str: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + resolved_id: Final = team_id if team_id is not None else auth.team_id + if write and auth.user_role in {LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY}: + raise_public( + WorkflowFailure( + code="write_denied", message="View-only users cannot change or execute drafts", status_code=403 + ) + ) + if resolved_id is None or not auth.user_id: + raise_public( + WorkflowFailure( + code="team_identity_required", message="A team and a bound user are required", status_code=403 + ) + ) + team: Final = await get_team_object( + resolved_id, + prisma_client, + user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + allowed_admin: Final = await get_team_access().allows(auth, team, TEAM_OR_ORG_ADMIN) + if not allowed_admin and (admin or not is_team_member(auth, team)): + raise_public(WorkflowFailure(code="not_found", message="Workflow not found", status_code=404)) + return resolved_id + + +async def require_draft( + store: WorkflowStore, workflow_id: str, auth: UserAPIKeyAuth, *, admin: bool = False, write: bool = False +) -> WorkflowDraft: + draft: Final = await store.get_id(workflow_id) + if draft is None: + raise_public(WorkflowFailure(code="not_found", message="Workflow not found", status_code=404)) + await require_team(auth, draft.team_id, admin=admin, write=write) + return draft + + +def draft_view(draft: WorkflowDraft) -> DraftView: + bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json) + return DraftView( + workflow_id=draft.workflow_id, + name=draft.name, + team_id=draft.team_id, + etag=f'"{draft.etag}"', + digest=workflow_digest(bundle), + bundle=bundle, + published_snapshot_id=draft.published_snapshot_id, + ) + + +def require_etag(draft: WorkflowDraft, if_match: str | None) -> None: + if if_match != f'"{draft.etag}"': + raise_public( + WorkflowFailure(code="stale_draft", message="If-Match must match the reviewed draft ETag", status_code=412) + ) + + +async def check_bundle(bundle: WorkflowBundle, auth: UserAPIKeyAuth) -> None: + from litellm.proxy.proxy_server import llm_model_list_view, llm_router + + failure: Final = validate_bundle(bundle) + if failure is not None: + raise_public(failure) + for node in bundle.nodes: + if isinstance(node, ChatCompletionNode): + if llm_router is None or node.model not in llm_router.registry().names(team_id=auth.team_id): + raise_public( + WorkflowFailure(code="unknown_model", message="Use a configured gateway model alias", param=node.id) + ) + await can_key_call_resolved_model(node.model, llm_model_list_view(), auth, llm_router) + + +@router.post("/v1/workflows", status_code=201, response_model=DraftView) +async def create_workflow( + data: RegisterWorkflow, response: Response, auth: Auth, workflows: WorkflowRepositoryDep, team_id: str | None = None +) -> DraftView: + """Register a team-scoped draft without publishing or executing it.""" + team: Final = await require_team(auth, team_id, write=True) + await check_bundle(data.bundle, auth) + result: Final = await require_store(workflows).create( + team, data.bundle.name, canonical_bundle(data.bundle), auth.user_id or "" + ) + if isinstance(result, WorkflowConflict): + raise_public(WorkflowFailure(code="workflow_conflict", message=result.message, status_code=result.status_code)) + view: Final = draft_view(result) + response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers + return view + + +@router.get("/v1/workflows", response_model=DraftList) +async def list_workflows( + auth: Auth, + workflows: WorkflowRepositoryDep, + team_id: str | None = None, + after: str | None = None, + limit: int = Query(default=50, ge=1, le=100), +) -> DraftList: + team: Final = await require_team(auth, team_id) + rows: Final = await require_store(workflows).list(team, after, limit + 1) + return DraftList( + workflows=tuple(draft_view(row) for row in rows[:limit]), + next_cursor=rows[limit - 1].workflow_id if len(rows) > limit else None, + ) + + +@router.get("/v1/workflows/{workflow_id}", response_model=DraftView) +async def get_workflow(workflow_id: str, response: Response, auth: Auth, workflows: WorkflowRepositoryDep) -> DraftView: + view: Final = draft_view(await require_draft(require_store(workflows), workflow_id, auth)) + response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers + return view + + +@router.put("/v1/workflows/{workflow_id}", response_model=DraftView) +async def replace_workflow( + workflow_id: str, + data: RegisterWorkflow, + response: Response, + auth: Auth, + workflows: WorkflowRepositoryDep, + if_match: Annotated[str | None, Header()] = None, +) -> DraftView: + store: Final = require_store(workflows) + draft: Final = await require_draft(store, workflow_id, auth, write=True) + require_etag(draft, if_match) + if data.bundle.name != draft.name: + raise_public(WorkflowFailure(code="immutable_name", message="Workflow names cannot be changed")) + await check_bundle(data.bundle, auth) + result: Final = await store.replace( + draft.team_id, draft.name, draft.etag, canonical_bundle(data.bundle), auth.user_id or "" + ) + if isinstance(result, WorkflowConflict): + raise_public(WorkflowFailure(code="stale_draft", message=result.message, status_code=result.status_code)) + view: Final = draft_view(result) + response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers + return view + + +@router.post("/v1/workflows/{workflow_id}/validate", response_model=ValidationView) +async def validate_workflow( + workflow_id: str, + auth: Auth, + workflows: WorkflowRepositoryDep, + if_match: Annotated[str | None, Header()] = None, + action: EmptyWorkflowAction | None = None, +) -> ValidationView: + """Validate the reviewed draft without executing Python or calling models.""" + draft: Final = await require_draft(require_store(workflows), workflow_id, auth) + require_etag(draft, if_match) + bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json) + await check_bundle(bundle, auth) + return ValidationView(valid=True, digest=workflow_digest(bundle)) + + +@router.post("/v1/workflows/{workflow_id}/publish", response_model=PublicationView) +async def publish_workflow( + workflow_id: str, + data: PublishWorkflow, + auth: Auth, + workflows: WorkflowRepositoryDep, + if_match: Annotated[str | None, Header()] = None, +) -> PublicationView: + """Publish an immutable snapshot after independent approval of the current ETag.""" + store: Final = require_store(workflows) + draft: Final = await require_draft(store, workflow_id, auth, admin=True, write=True) + require_etag(draft, if_match) + bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json) + if data.digest != workflow_digest(bundle): + raise_public( + WorkflowFailure(code="stale_draft", message="Reviewed digest does not match the draft", status_code=412) + ) + await check_bundle(bundle, auth) + result: Final = await store.publish(draft.team_id, draft.name, draft.etag, auth.user_id or "") + if isinstance(result, WorkflowConflict): + raise_public(WorkflowFailure(code="publication_denied", message=result.message, status_code=result.status_code)) + return PublicationView(snapshot_id=result.snapshot_id, model=f"workflow/{draft.name}", digest=data.digest) diff --git a/litellm/proxy/workflows/runner.py b/litellm/proxy/workflows/runner.py new file mode 100644 index 00000000000..af48409b620 --- /dev/null +++ b/litellm/proxy/workflows/runner.py @@ -0,0 +1,232 @@ +import asyncio +import time +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final, Protocol + +from pydantic import ConfigDict, JsonValue, TypeAdapter + +from litellm.models.workflow import ( + WorkflowCheckpoint, + WorkflowConflict, + WorkflowExecution, + WorkflowRunState, + WorkflowStep, + WorkflowStore, +) +from litellm.models.workflow_bundle import ( + Binding, + ChatCompletionNode, + Constant, + WorkflowBundle, + WorkflowFailure, + WorkflowNode, +) +from litellm.proxy.workflows.monty.compiler import compile_program +from litellm.proxy.workflows.monty.models import MontyFailure +from litellm.proxy.workflows.monty.runtime import MontyRuntime +from litellm.proxy.workflows.validation import validate_value +from litellm.types.llms.base import LiteLLMBaseModel + + +class ChildResult(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + output: JsonValue + prompt_tokens: int = 0 + completion_tokens: int = 0 + response_cost: float | None = None + request_id: str | None = None + + +class ChildCaller(Protocol): + async def call( + self, execution: WorkflowExecution, node: ChatCompletionNode, messages: JsonValue, remaining: float + ) -> ChildResult | WorkflowFailure: ... + + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def resolve(binding: Binding, messages: JsonValue, outputs: Mapping[str, JsonValue]) -> JsonValue | WorkflowFailure: + if isinstance(binding, Constant): + return binding.value + if binding.source != "$input" and binding.source not in outputs: + return WorkflowFailure( + code="missing_result", message="Referenced predecessor has no result", param=binding.source + ) + value: Final = messages if binding.source == "$input" else outputs[binding.source] + if binding.path in (None, ""): + return value + return project(value, tuple(part.replace("~1", "/").replace("~0", "~") for part in binding.path[1:].split("/"))) + + +def project(value: JsonValue, parts: tuple[str, ...]) -> JsonValue | WorkflowFailure: + if not parts: + return value + key: Final = parts[0] + if isinstance(value, dict) and key in value: + return project(value[key], parts[1:]) + if isinstance(value, list) and key.isdigit() and str(int(key)) == key and int(key) < len(value): + return project(value[int(key)], parts[1:]) + return WorkflowFailure(code="missing_path", message="Referenced JSON Pointer has no value") + + +def node_input( + node: WorkflowNode, messages: JsonValue, outputs: Mapping[str, JsonValue] +) -> JsonValue | WorkflowFailure: + if isinstance(node, ChatCompletionNode): + return resolve(node.messages, messages, outputs) + arguments: Final = {name: resolve(binding, messages, outputs) for name, binding in node.inputs.items()} + failures: Final = tuple(value for value in arguments.values() if isinstance(value, WorkflowFailure)) + if failures: + return failures[0] + return _JSON.validate_python(arguments) + + +async def execute_attempt( + caller: ChildCaller, + runtime: MontyRuntime, + execution: WorkflowExecution, + node: WorkflowNode, + input: JsonValue, + remaining: float, +) -> ChildResult | WorkflowFailure: + if isinstance(node, ChatCompletionNode): + return await caller.call(execution, node, input, remaining) + compiled: Final = compile_program(node.program, runtime.limits) + if isinstance(compiled, MontyFailure): + return WorkflowFailure(code="invalid_program", message=compiled.message, param=node.id) + result: Final = await runtime.run_program(compiled, input, deadline_monotonic=time.monotonic() + remaining) + if isinstance(result, MontyFailure): + return WorkflowFailure(code="python_execution_failed", message=result.message, param=node.id) + return ChildResult(output=result.output) + + +async def execute_step( + store: WorkflowStore, + caller: ChildCaller, + runtime: MontyRuntime, + execution: WorkflowExecution, + node: WorkflowNode, + step: WorkflowStep, +) -> ChildResult | WorkflowFailure: + remaining: Final = (min(step.deadline, execution.deadline) - datetime.now(timezone.utc)).total_seconds() + if remaining <= 0: + return WorkflowFailure( + code="node_deadline", message="Node total deadline exceeded", param=node.id, status_code=504 + ) + if step.attempts >= node.retry.max_attempts: + return WorkflowFailure( + code="attempts_exhausted", message="Node attempt policy exhausted", param=node.id, status_code=502 + ) + attempt: Final = await store.attempt_step(step.step_id, execution.lease_owner or "") + if isinstance(attempt, WorkflowConflict): + return WorkflowFailure(code="lease_lost", message=attempt.message, status_code=503) + attempt_remaining: Final = (min(step.deadline, execution.deadline) - datetime.now(timezone.utc)).total_seconds() + try: + async with asyncio.timeout(max(attempt_remaining, 0)): + result: Final = await execute_attempt(caller, runtime, execution, node, step.input, attempt_remaining) + except TimeoutError: + return WorkflowFailure( + code="node_deadline", message="Node total deadline exceeded", param=node.id, status_code=504 + ) + if not isinstance(result, WorkflowFailure) or attempt.attempts >= node.retry.max_attempts: + return result + if result.status_code not in {429, 500, 502, 503, 504}: + return result + backoff: Final = min( + node.retry.backoff_seconds * 2.0 ** (attempt.attempts - 1), + max((step.deadline - datetime.now(timezone.utc)).total_seconds(), 0), + ) + await asyncio.sleep(backoff) + return await execute_step(store, caller, runtime, execution, node, attempt) + + +async def execute_sequence( + store: WorkflowStore, + caller: ChildCaller, + runtime: MontyRuntime, + execution: WorkflowExecution, + bundle: WorkflowBundle, + messages: JsonValue, + outputs: Mapping[str, JsonValue], + index: int = 0, +) -> JsonValue | WorkflowFailure: + if index == len(bundle.nodes): + return resolve(bundle.output, messages, outputs) + current: Final = await store.execution(execution.run_id) + if current is None or current.cancel_requested: + return WorkflowFailure(code="cancelled", message="Run cancellation requested", status_code=409) + if current.lease_owner != execution.lease_owner: + return WorkflowFailure(code="lease_lost", message="Run lease lost", status_code=503) + node: Final = bundle.nodes[index] + input: Final = node_input(node, messages, outputs) + if isinstance(input, WorkflowFailure): + return input + deadline: Final = min(execution.deadline, datetime.now(timezone.utc) + timedelta(seconds=node.timeout_seconds)) + step: Final = await store.start_step(execution.run_id, node.id, input, deadline, execution.lease_owner or "") + if isinstance(step, WorkflowConflict): + return WorkflowFailure(code="lease_lost", message=step.message, status_code=503) + if step.state is WorkflowRunState.COMPLETED: + return await execute_sequence( + store, caller, runtime, execution, bundle, messages, {**outputs, node.id: step.output}, index + 1 + ) + result: Final = await execute_step(store, caller, runtime, execution, node, step) + steps: Final = await store.steps(execution.run_id) + attempted: Final = next(item for item in steps if item.step_id == step.step_id) + if isinstance(result, WorkflowFailure): + await store.finish_step( + attempted.model_copy( + update={ + "state": WorkflowRunState.FAILED, + "error": result.code, + "finished_at": datetime.now(timezone.utc), + } + ), + execution.lease_owner or "", + ) + return result + completed: Final = attempted.model_copy( + update={ + "state": WorkflowRunState.COMPLETED, + "output": result.output, + "finished_at": datetime.now(timezone.utc), + "prompt_tokens": result.prompt_tokens, + "completion_tokens": result.completion_tokens, + "response_cost": result.response_cost, + "request_id": result.request_id, + } + ) + checkpointed: Final = await store.checkpoint( + execution.run_id, + WorkflowCheckpoint( + node_path=node.id, + input=step.input, + output=result.output, + attempts=attempted.attempts, + ), + completed, + execution.lease_owner or "", + ) + if not checkpointed: + return WorkflowFailure(code="lease_lost", message="Unable to commit node checkpoint", status_code=503) + return await execute_sequence( + store, caller, runtime, execution, bundle, messages, {**outputs, node.id: result.output}, index + 1 + ) + + +async def execute_run( + store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, execution: WorkflowExecution +) -> JsonValue | WorkflowFailure: + snapshot: Final = await store.snapshot(execution.snapshot_id) + if snapshot is None: + return WorkflowFailure(code="snapshot_missing", message="Accepted snapshot is unavailable", status_code=503) + bundle: Final = WorkflowBundle.model_validate_json(snapshot.bundle_json) + request: Final = _OBJECT.validate_json(execution.request_json) + messages: Final = request.get("messages") + failure: Final = validate_value(messages, bundle.input_schema, "messages") + if failure is not None: + return failure + return await execute_sequence(store, caller, runtime, execution, bundle, messages, {}) diff --git a/litellm/proxy/workflows/runs.py b/litellm/proxy/workflows/runs.py new file mode 100644 index 00000000000..8be1f395153 --- /dev/null +++ b/litellm/proxy/workflows/runs.py @@ -0,0 +1,147 @@ +from typing import Annotated, Final, Literal + +from fastapi import APIRouter, Header, Request +from openai.types.chat import ChatCompletionMessageParam +from pydantic import ConfigDict + +from litellm.models.workflow import WorkflowConflict, WorkflowExecution, WorkflowRunState, WorkflowStep, WorkflowStore +from litellm.models.workflow_bundle import WorkflowBundle, WorkflowFailure +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.deps.repositories import WorkflowRepositoryDep +from litellm.proxy.workflows.invocation import InvocationMetadata, WorkflowInvocation, admit, principal_hash +from litellm.proxy.workflows.management import ( + Auth, + EmptyWorkflowAction, + check_bundle, + raise_public, + require_draft, + require_etag, + require_store, + require_team, +) +from litellm.types.llms.base import LiteLLMBaseModel + +router: Final = APIRouter(tags=["Workflows"]) + + +class DraftTest(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + messages: tuple[ChatCompletionMessageParam, ...] + user: str | None = None + metadata: InvocationMetadata | None = None + authorize_billing: Literal[True] + + +class RunAccepted(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + snapshot_id: str + trace_url: str + + +class RunSteps(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + snapshot_id: str + state: WorkflowRunState + cancel_requested: bool + steps: tuple[WorkflowStep, ...] + + +class CancellationView(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + run_id: str + cancel_requested: bool + + +async def require_execution( + store: WorkflowStore, run_id: str, auth: UserAPIKeyAuth, *, write: bool = False +) -> WorkflowExecution: + execution: Final = await store.execution(run_id) + if execution is None: + raise_public(WorkflowFailure(code="not_found", message="Workflow run not found", status_code=404)) + if principal_hash(auth) != execution.principal: + await require_team(auth, execution.team_id, admin=True, write=write) + elif auth.team_id != execution.team_id: + raise_public(WorkflowFailure(code="not_found", message="Workflow run not found", status_code=404)) + return execution + + +@router.get("/v1/workflows/runs/{run_id}/steps", response_model=RunSteps) +async def read_steps(run_id: str, auth: Auth, workflows: WorkflowRepositoryDep) -> RunSteps: + store: Final = require_store(workflows) + execution: Final = await require_execution(store, run_id, auth) + return RunSteps( + run_id=run_id, + snapshot_id=execution.snapshot_id, + state=execution.state, + cancel_requested=execution.cancel_requested, + steps=await store.steps(run_id), + ) + + +@router.get("/v1/workflows/runs/{run_id}/steps/{step_id}", response_model=WorkflowStep) +async def read_step(run_id: str, step_id: str, auth: Auth, workflows: WorkflowRepositoryDep) -> WorkflowStep: + store: Final = require_store(workflows) + await require_execution(store, run_id, auth) + steps: Final = await store.steps(run_id) + selected: Final = next((step for step in steps if step.step_id == step_id), None) + if selected is None: + raise_public(WorkflowFailure(code="not_found", message="Workflow step not found", status_code=404)) + return selected + + +@router.post("/v1/workflows/runs/{run_id}/cancel", status_code=202, response_model=CancellationView) +async def cancel_run( + run_id: str, auth: Auth, workflows: WorkflowRepositoryDep, action: EmptyWorkflowAction | None = None +) -> CancellationView: + """Request cancellation and prevent further node dispatch for this run.""" + store: Final = require_store(workflows) + await require_execution(store, run_id, auth, write=True) + await store.cancel(run_id) + return CancellationView(run_id=run_id, cancel_requested=True) + + +@router.post("/v1/workflows/{workflow_id}/test", status_code=202, response_model=RunAccepted) +async def test_draft( + workflow_id: str, + data: DraftTest, + request: Request, + auth: Auth, + workflows: WorkflowRepositoryDep, + if_match: Annotated[str | None, Header()] = None, +) -> RunAccepted: + """Execute the reviewed draft with explicit billing authorization.""" + from litellm.proxy.proxy_server import general_settings_view + + if general_settings_view().get("enable_workflows") is not True: + raise_public( + WorkflowFailure(code="workflow_disabled", message="Workflow execution is disabled", status_code=503) + ) + store: Final = require_store(workflows) + draft: Final = await require_draft(store, workflow_id, auth, write=True) + require_etag(draft, if_match) + bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json) + await check_bundle(bundle, auth) + snapshot: Final = await store.pin_draft(draft.workflow_id, draft.etag, auth.user_id or "") + if isinstance(snapshot, WorkflowConflict): + raise_public(WorkflowFailure(code="stale_draft", message=snapshot.message, status_code=snapshot.status_code)) + accepted: Final = await admit( + store, + auth, + WorkflowInvocation( + model=f"workflow/{draft.name}", messages=data.messages, user=data.user, metadata=data.metadata + ), + request.headers.get("Idempotency-Key"), + snapshot=snapshot, + operation="draft_test", + ) + return RunAccepted( + run_id=accepted.run_id, + snapshot_id=accepted.snapshot_id, + trace_url=f"/v1/workflows/runs/{accepted.run_id}/steps", + ) diff --git a/litellm/proxy/workflows/validation.py b/litellm/proxy/workflows/validation.py new file mode 100644 index 00000000000..01cac0fcb43 --- /dev/null +++ b/litellm/proxy/workflows/validation.py @@ -0,0 +1,163 @@ +import json +from collections.abc import Mapping +from itertools import chain +from typing import Final + +from jsonschema import Draft202012Validator, SchemaError +from pydantic import JsonValue + +from litellm.models.workflow_bundle import ( + Binding, + ChatCompletionNode, + Constant, + PythonNode, + Reference, + WorkflowBundle, + WorkflowFailure, + canonical_bundle, + validate_bindings, +) +from litellm.proxy.workflows.monty.compiler import compile_program +from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits + +DEFAULT_MONTY_LIMITS: Final = MontyLimits() + + +def schema_is_local(value: JsonValue, depth: int = 0) -> bool: + if depth > 48: + return False + if isinstance(value, list): + return all(schema_is_local(item, depth + 1) for item in value) + if not isinstance(value, dict): + return True + return all( + (key not in {"$ref", "$dynamicRef"} or isinstance(child, str) and child.startswith("#/$defs/")) + and key not in {"$id", "$dynamicAnchor"} + and schema_is_local(child, depth + 1) + for key, child in value.items() + ) + + +def schema_failure(schema: Mapping[str, JsonValue], param: str) -> WorkflowFailure | None: + if not schema_is_local(dict(schema)): + return WorkflowFailure( + code="unsupported_schema", message="Only local $defs references are supported", param=param + ) + try: + Draft202012Validator.check_schema(schema) + except SchemaError: + return WorkflowFailure(code="invalid_schema", message="Invalid JSON Schema", param=param) + return None + + +def validate_value(value: JsonValue, schema: Mapping[str, JsonValue], param: str) -> WorkflowFailure | None: + if next(Draft202012Validator(schema).iter_errors(value), None) is not None: # pyright: ignore[reportUnknownMemberType] # jsonschema stubs leave variadic arguments untyped + return WorkflowFailure( + code="schema_mismatch", message="Value does not match its declared JSON Schema", param=param + ) + return None + + +def validate_bundle(bundle: WorkflowBundle) -> WorkflowFailure | None: + if len(canonical_bundle(bundle).encode()) > 1_048_576: + return WorkflowFailure(code="bundle_too_large", message="Workflow bundle exceeds 1 MiB") + bindings_failure: Final = validate_bindings(bundle) + if bindings_failure is not None: + return bindings_failure + input_failure: Final = schema_failure(bundle.input_schema, "input_schema") + if input_failure is not None: + return input_failure + for node in bundle.nodes: + if isinstance(node, ChatCompletionNode) and node.response_format is not None: + response_failure = schema_failure( + node.response_format.json_schema.schema_, f"nodes.{node.id}.response_format" + ) + if response_failure is not None: + return response_failure + if isinstance(node, PythonNode): + compiled = compile_program(node.program, DEFAULT_MONTY_LIMITS) + if isinstance(compiled, MontyFailure): + return WorkflowFailure(code="invalid_program", message=compiled.message, param=f"nodes.{node.id}") + return compatibility_failure(bundle) + + +def _types( + schema: JsonValue | Mapping[str, JsonValue], root: Mapping[str, JsonValue], seen: frozenset[str] = frozenset() +) -> frozenset[str]: + if not isinstance(schema, Mapping): + return frozenset() + reference: Final = schema.get("$ref") + if isinstance(reference, str) and reference.startswith("#/$defs/") and reference not in seen: + definitions: Final = root.get("$defs") + return ( + _types(definitions.get(reference[8:]), root, seen | {reference}) + if isinstance(definitions, dict) + else frozenset() + ) + type_: Final = schema.get("type") + if isinstance(type_, str): + return frozenset({type_}) + alternatives: Final = schema.get("anyOf", schema.get("oneOf")) + if isinstance(alternatives, list): + return frozenset(chain.from_iterable(_types(alternative, root, seen) for alternative in alternatives)) + return frozenset() + + +def _result_schema(bundle: WorkflowBundle, source: str) -> Mapping[str, JsonValue]: + if source == "$input": + return bundle.input_schema + node: Final = next(node for node in bundle.nodes if node.id == source) + if isinstance(node, PythonNode): + return node.program.output_schema + return node.response_format.json_schema.schema_ if node.response_format is not None else {"type": "string"} + + +def _binding_compatibility( + bundle: WorkflowBundle, binding: Binding, expected: JsonValue, expected_root: Mapping[str, JsonValue], param: str +) -> WorkflowFailure | None: + if isinstance(binding, Constant) and isinstance(expected, dict): + return validate_value(binding.value, {**expected, "$defs": expected_root.get("$defs", {})}, param) + if not isinstance(binding, Reference) or binding.path not in (None, ""): + return None + actual: Final = _result_schema(bundle, binding.source) + actual_types: Final = _types(actual, actual) + expected_types: Final = _types(expected, expected_root) + if ( + actual_types + and expected_types + and not actual_types & expected_types + and not ("integer" in actual_types and "number" in expected_types) + ): + return WorkflowFailure( + code="incompatible_binding", message="Binding has an incompatible declared result type", param=param + ) + return None + + +def compatibility_failure(bundle: WorkflowBundle) -> WorkflowFailure | None: + for node in bundle.nodes: + if isinstance(node, ChatCompletionNode): + failure = _binding_compatibility(bundle, node.messages, {"type": "array"}, {}, f"nodes.{node.id}.messages") + if failure is not None: + return failure + continue + properties = node.program.input_schema.get("properties") + if not isinstance(properties, dict): + return WorkflowFailure( + code="invalid_inputs", message="Function argument schema must declare properties", param=node.id + ) + for name, binding in node.inputs.items(): + input_failure = _binding_compatibility( + bundle, + binding, + properties.get(name), + node.program.input_schema, + f"nodes.{node.id}.inputs.{name}", + ) + if input_failure is not None: + return input_failure + return None + + +def canonical_request(value: Mapping[str, JsonValue]) -> str: + return json.dumps(dict(value), sort_keys=True, separators=(",", ":"), allow_nan=False) diff --git a/litellm/repositories/workflow_repository.py b/litellm/repositories/workflow_repository.py new file mode 100644 index 00000000000..c6e94022987 --- /dev/null +++ b/litellm/repositories/workflow_repository.py @@ -0,0 +1,438 @@ +import json +from collections.abc import Sequence +from contextlib import AbstractAsyncContextManager +from datetime import datetime +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + TypeVar, + cast, # noqa: TID251 # PrismaWrapper forwards transaction operations dynamically +) + +from pydantic import JsonValue, TypeAdapter + +from litellm._uuid import uuid +from litellm.models.workflow import ( + WorkflowCheckpoint, + WorkflowConflict, + WorkflowDraft, + WorkflowExecution, + WorkflowRunState, + WorkflowSnapshot, + WorkflowStep, +) + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +class WorkflowTransaction(Protocol): + async def query_raw(self, query: str, *args: object) -> object: ... + + async def execute_raw(self, query: str, *args: object) -> int: ... + + +class WorkflowDatabase(WorkflowTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[WorkflowTransaction]: ... + + +_DRAFTS: Final = TypeAdapter(tuple[WorkflowDraft, ...]) +_SNAPSHOTS: Final = TypeAdapter(tuple[WorkflowSnapshot, ...]) +_EXECUTIONS: Final = TypeAdapter(tuple[WorkflowExecution, ...]) +_CHECKPOINTS: Final = TypeAdapter(tuple[WorkflowCheckpoint, ...]) +_STEPS: Final = TypeAdapter(tuple[WorkflowStep, ...]) +_DRAFT_COLUMNS: Final = "workflow_id, team_id, name, etag, bundle_json, contributors, published_snapshot_id" +_SNAPSHOT_COLUMNS: Final = "snapshot_id, workflow_id, team_id, name, bundle_json, published_by, created_at, origin" +_EXECUTION_COLUMNS: Final = ( + "run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, state, " + "deadline, lease_owner, lease_until, cancel_requested, checkpoints, response_json, error" +) +_STEP_COLUMNS: Final = ( + "step_id, run_id, node_id, state, input, output, attempts, deadline, started_at, finished_at, error, " + "prompt_tokens, completion_tokens, response_cost, request_id" +) + + +_RowT: Final = TypeVar("_RowT") + + +def _first(rows: Sequence[_RowT]) -> _RowT | None: + return rows[0] if rows else None + + +class WorkflowRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._database: Final = cast( # cast-ok: Prisma exposes raw SQL and transaction methods dynamically + WorkflowDatabase, prisma_client.writer_db + ) + + async def create(self, team_id: str, name: str, bundle_json: str, author: str) -> WorkflowDraft | WorkflowConflict: + rows: Final = _DRAFTS.validate_python( + await self._database.query_raw( + 'INSERT INTO "LiteLLM_WorkflowDefinition" ' + "(workflow_id, team_id, name, etag, bundle_json, contributors) " + "VALUES ($1, $2, $3, $4, $5, ARRAY[$6]::text[]) ON CONFLICT (team_id, name) DO NOTHING " + f"RETURNING {_DRAFT_COLUMNS}", + str(uuid.uuid4()), + team_id, + name, + str(uuid.uuid4()), + bundle_json, + author, + ) + ) + return rows[0] if rows else WorkflowConflict(message="Workflow already exists") + + async def get(self, team_id: str, name: str) -> WorkflowDraft | None: + return _first( + _DRAFTS.validate_python( + await self._database.query_raw( + f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2', + team_id, + name, + ) + ) + ) + + async def replace( + self, team_id: str, name: str, etag: str, bundle_json: str, author: str + ) -> WorkflowDraft | WorkflowConflict: + rows: Final = _DRAFTS.validate_python( + await self._database.query_raw( + 'UPDATE "LiteLLM_WorkflowDefinition" SET etag = $4, bundle_json = $5, ' + "contributors = ARRAY(SELECT DISTINCT unnest(contributors || ARRAY[$6]::text[])) " + f"WHERE team_id = $1 AND name = $2 AND etag = $3 RETURNING {_DRAFT_COLUMNS}", + team_id, + name, + etag, + str(uuid.uuid4()), + bundle_json, + author, + ) + ) + return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412) + + async def get_id(self, workflow_id: str) -> WorkflowDraft | None: + return _first( + _DRAFTS.validate_python( + await self._database.query_raw( + f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1', + workflow_id, + ) + ) + ) + + async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]: + return _DRAFTS.validate_python( + await self._database.query_raw( + f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" ' + "WHERE team_id = $1 AND workflow_id > $2 ORDER BY workflow_id LIMIT $3", + team_id, + after or "", + limit, + ) + ) + + async def publish(self, team_id: str, name: str, etag: str, publisher: str) -> WorkflowSnapshot | WorkflowConflict: + async with self._database.tx() as tx: + drafts: Final = _DRAFTS.validate_python( + await tx.query_raw( + f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" ' + "WHERE team_id = $1 AND name = $2 FOR UPDATE", + team_id, + name, + ) + ) + if not drafts or drafts[0].etag != etag: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + if publisher in drafts[0].contributors: + return WorkflowConflict(message="Draft contributors cannot publish their own draft", status_code=403) + rows: Final = _SNAPSHOTS.validate_python( + await tx.query_raw( + 'INSERT INTO "LiteLLM_WorkflowSnapshot" ' + "(snapshot_id, workflow_id, team_id, name, bundle_json, published_by) " + f"VALUES ($1, $2, $3, $4, $5, $6) RETURNING {_SNAPSHOT_COLUMNS}", + str(uuid.uuid4()), + drafts[0].workflow_id, + team_id, + name, + drafts[0].bundle_json, + publisher, + ) + ) + await tx.execute_raw( + 'UPDATE "LiteLLM_WorkflowDefinition" SET published_snapshot_id = $3, etag = $4 ' + "WHERE team_id = $1 AND name = $2", + team_id, + name, + rows[0].snapshot_id, + str(uuid.uuid4()), + ) + return rows[0] + + async def snapshot(self, snapshot_id: str) -> WorkflowSnapshot | None: + return _first( + _SNAPSHOTS.validate_python( + await self._database.query_raw( + f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = $1', + snapshot_id, + ) + ) + ) + + async def published(self, team_id: str, name: str) -> WorkflowSnapshot | None: + return _first( + _SNAPSHOTS.validate_python( + await self._database.query_raw( + f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = ' + '(SELECT published_snapshot_id FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2)', + team_id, + name, + ) + ) + ) + + async def pin_draft(self, workflow_id: str, etag: str, requested_by: str) -> WorkflowSnapshot | WorkflowConflict: + rows: Final = _SNAPSHOTS.validate_python( + await self._database.query_raw( + 'INSERT INTO "LiteLLM_WorkflowSnapshot" ' + "(snapshot_id, workflow_id, team_id, name, bundle_json, published_by, origin) " + "SELECT $3, workflow_id, team_id, name, bundle_json, $4, 'draft_test' " + 'FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND etag = $2 ' + f"RETURNING {_SNAPSHOT_COLUMNS}", + workflow_id, + etag, + str(uuid.uuid4()), + requested_by, + ) + ) + return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412) + + async def by_dispatch_key(self, dispatch_key: str) -> WorkflowExecution | None: + return _first( + _EXECUTIONS.validate_python( + await self._database.query_raw( + f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1', + dispatch_key, + ) + ) + ) + + async def execution(self, run_id: str) -> WorkflowExecution | None: + return _first( + _EXECUTIONS.validate_python( + await self._database.query_raw( + f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE run_id = $1', + run_id, + ) + ) + ) + + async def dispatch(self, execution: WorkflowExecution) -> WorkflowExecution: + async with self._database.tx() as tx: + await tx.query_raw( + "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended($1, 0))", execution.dispatch_key + ) + existing: Final = _EXECUTIONS.validate_python( + await tx.query_raw( + f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1', + execution.dispatch_key, + ) + ) + if existing: + return existing[0] + await tx.execute_raw( + 'INSERT INTO "LiteLLM_WorkflowRun" ' + "(run_id, session_id, workflow_type, created_by, input, metadata, engine_owned, updated_at) " + "VALUES ($1, $1, 'liteflow', $2, $3::jsonb, $4::jsonb, true, NOW())", + execution.run_id, + execution.principal, + execution.request_json, + json.dumps( + { + "snapshot_id": execution.snapshot_id, + "team_id": execution.team_id, + "workflow_name": execution.name, + } + ), + ) + rows: Final = _EXECUTIONS.validate_python( + await tx.query_raw( + 'INSERT INTO "LiteLLM_WorkflowExecution" ' + "(run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, deadline) " + f"VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9::timestamptz AT TIME ZONE 'UTC') RETURNING {_EXECUTION_COLUMNS}", + execution.run_id, + execution.snapshot_id, + execution.team_id, + execution.name, + execution.principal, + execution.dispatch_key, + execution.request_digest, + execution.request_json, + execution.deadline, + ) + ) + return rows[0] + + async def lease_next(self, owner: str) -> WorkflowExecution | None: + return _first( + _EXECUTIONS.validate_python( + await self._database.query_raw( + "UPDATE \"LiteLLM_WorkflowExecution\" SET lease_owner = $1, lease_until = NOW() + INTERVAL '30 seconds', " + "state = 'running' WHERE run_id = (SELECT run_id FROM \"LiteLLM_WorkflowExecution\" " + "WHERE state IN ('pending', 'running') AND (lease_until IS NULL OR lease_until < NOW()) " + f"ORDER BY deadline FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING {_EXECUTION_COLUMNS}", + owner, + ) + ) + ) + + async def renew(self, run_id: str, owner: str) -> bool: + return bool( + await self._database.execute_raw( + "UPDATE \"LiteLLM_WorkflowExecution\" SET lease_until = NOW() + INTERVAL '30 seconds' " + "WHERE run_id = $1 AND lease_owner = $2 AND lease_until > NOW() AND state = 'running'", + run_id, + owner, + ) + ) + + async def steps(self, run_id: str) -> tuple[WorkflowStep, ...]: + return _STEPS.validate_python( + await self._database.query_raw( + f'SELECT {_STEP_COLUMNS} FROM "LiteLLM_WorkflowStep" WHERE run_id = $1 ORDER BY started_at, step_id', + run_id, + ) + ) + + async def start_step( + self, run_id: str, node_id: str, input: JsonValue, deadline: datetime, owner: str + ) -> WorkflowStep | WorkflowConflict: + rows: Final = _STEPS.validate_python( + await self._database.query_raw( + 'INSERT INTO "LiteLLM_WorkflowStep" (step_id, run_id, node_id, input, deadline) ' + "SELECT $1,$2,$3,$4::jsonb,$5::timestamptz AT TIME ZONE 'UTC' WHERE EXISTS (SELECT 1 FROM \"LiteLLM_WorkflowExecution\" " + "WHERE run_id=$2 AND lease_owner=$6 AND lease_until>NOW() AND NOT cancel_requested) " + "ON CONFLICT (run_id, node_id) DO UPDATE SET node_id = EXCLUDED.node_id " + f"RETURNING {_STEP_COLUMNS}", + str(uuid.uuid4()), + run_id, + node_id, + json.dumps(input, allow_nan=False), + deadline, + owner, + ) + ) + return rows[0] if rows else WorkflowConflict(message="Execution lease lost") + + async def attempt_step(self, step_id: str, owner: str) -> WorkflowStep | WorkflowConflict: + rows: Final = _STEPS.validate_python( + await self._database.query_raw( + 'UPDATE "LiteLLM_WorkflowStep" SET attempts = attempts + 1 ' + 'WHERE step_id = $1 AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e ' + 'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$2 ' + f"AND e.lease_until>NOW() AND NOT e.cancel_requested) RETURNING {_STEP_COLUMNS}", + step_id, + owner, + ) + ) + return rows[0] if rows else WorkflowConflict(message="Execution lease lost") + + async def _finish_step(self, database: WorkflowTransaction, step: WorkflowStep, owner: str) -> bool: + return bool( + await database.execute_raw( + "UPDATE \"LiteLLM_WorkflowStep\" SET state=$2, output=$3::jsonb, finished_at=$4::timestamptz AT TIME ZONE 'UTC', error=$5, " + "prompt_tokens=$6, completion_tokens=$7, response_cost=$8, request_id=$9 WHERE step_id=$1 " + 'AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e ' + 'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$10 ' + "AND e.lease_until>NOW() AND NOT e.cancel_requested)", + step.step_id, + step.state.value, + json.dumps(step.output, allow_nan=False), + step.finished_at, + step.error, + step.prompt_tokens, + step.completion_tokens, + step.response_cost, + step.request_id, + owner, + ) + ) + + async def finish_step(self, step: WorkflowStep, owner: str) -> bool: + return await self._finish_step(self._database, step, owner) + + async def checkpoint(self, run_id: str, checkpoint: WorkflowCheckpoint, step: WorkflowStep, owner: str) -> bool: + async with self._database.tx() as tx: + rows: Final = _EXECUTIONS.validate_python( + await tx.query_raw( + f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" ' + "WHERE run_id = $1 AND lease_owner=$2 AND lease_until>NOW() AND NOT cancel_requested FOR UPDATE", + run_id, + owner, + ) + ) + if not rows or not await self._finish_step(tx, step, owner): + return False + if any(item.node_path == checkpoint.node_path for item in rows[0].checkpoints): + return True + checkpoints: Final = (*rows[0].checkpoints, checkpoint) + await tx.execute_raw( + 'UPDATE "LiteLLM_WorkflowExecution" SET checkpoints = $2::jsonb WHERE run_id = $1', + run_id, + _CHECKPOINTS.dump_json(checkpoints).decode(), + ) + await tx.execute_raw( + 'INSERT INTO "LiteLLM_WorkflowEvent" ' + "(event_id, run_id, event_type, step_name, sequence_number, data) " + "VALUES ($1,$2,'node.completed',$3,$4,$5::jsonb)", + str(uuid.uuid4()), + run_id, + checkpoint.node_path, + len(checkpoints) - 1, + checkpoint.model_dump_json(), + ) + return True + + async def finish( + self, + run_id: str, + state: WorkflowRunState, + response_json: str | None = None, + error: str | None = None, + *, + owner: str, + ) -> bool: + async with self._database.tx() as tx: + updated: Final = _EXECUTIONS.validate_python( + await tx.query_raw( + "UPDATE \"LiteLLM_WorkflowExecution\" SET state = CASE WHEN cancel_requested THEN 'cancelled' ELSE $2 END, " + "response_json = CASE WHEN cancel_requested THEN NULL ELSE $3 END, error = $4 " + f"WHERE run_id = $1 AND state IN ($5,$6) AND lease_owner=$7 AND lease_until>NOW() RETURNING {_EXECUTION_COLUMNS}", + run_id, + state.value, + response_json, + error, + WorkflowRunState.RUNNING.value, + WorkflowRunState.PENDING.value, + owner, + ) + ) + if not updated: + return False + await tx.execute_raw( + 'UPDATE "LiteLLM_WorkflowRun" SET status = $2, output = $3::jsonb, updated_at = NOW() WHERE run_id = $1', + run_id, + updated[0].state.value, + json.dumps({"response": json.loads(updated[0].response_json) if updated[0].response_json else None}), + ) + return True + + async def cancel(self, run_id: str) -> None: + await self._database.execute_raw( + 'UPDATE "LiteLLM_WorkflowExecution" SET cancel_requested = true WHERE run_id = $1 AND state IN ($2,$3)', + run_id, + WorkflowRunState.PENDING.value, + WorkflowRunState.RUNNING.value, + ) diff --git a/tests/integration/management/test_workflow_execution.py b/tests/integration/management/test_workflow_execution.py new file mode 100644 index 00000000000..d68271bdcbb --- /dev/null +++ b/tests/integration/management/test_workflow_execution.py @@ -0,0 +1,405 @@ +import hashlib +import json +import threading +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_SOURCE: Final = ( + "def prepare_followup(text: str) -> list[dict[str, str]]:\n" + ' return [{"role": "user", "content": "Review this draft: " + text}]\n' +) +_MESSAGES: Final[Mapping[str, JsonValue]] = { + "type": "array", + "items": { + "type": "object", + "properties": {"role": {"type": "string", "enum": ["user"]}, "content": {"type": "string"}}, + "required": ["role", "content"], + "additionalProperties": False, + }, +} +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +class Draft(BaseModel): + model_config = ConfigDict(frozen=True) + workflow_id: str + etag: str + digest: str + + +class Step(BaseModel): + model_config = ConfigDict(frozen=True) + step_id: str + state: str + input: JsonValue + output: JsonValue + attempts: int + response_cost: float | None + + +class Trace(BaseModel): + model_config = ConfigDict(frozen=True) + run_id: str + snapshot_id: str + state: str + steps: tuple[Step, ...] + + +@dataclass(frozen=True, slots=True) +class Actors: + author: str + publisher: str + caller: str + + +@pytest.fixture(scope="module") +def workflow_gateway(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with owned_proxy( + session_gateway, + tmp_path_factory.mktemp("workflows"), + {}, + workers=2, + config=Path(__file__).with_name("workflow_proxy_config.yaml"), + ) as gateway: + yield gateway + + +def _actors(scenario: Scenario, model: str) -> Actors: + author: Final = scenario.user(user_role="internal_user") + publisher: Final = scenario.user(user_role="internal_user") + team: Final = scenario.team( + models=[model, "workflow/email-reply"], + members_with_roles=[{"role": "admin", "user_id": author}, {"role": "admin", "user_id": publisher}], + ) + return Actors( + scenario.key(team_id=team, user_id=author, models=[model]), + scenario.key(team_id=team, user_id=publisher, models=[model]), + scenario.key(team_id=team, user_id=author, models=["workflow/email-reply"]), + ) + + +def _bundle(model: str, *, structured: bool = False, timeout: float = 60) -> Mapping[str, JsonValue]: + return { + "contract_version": "1", + "name": "email-reply", + "timeout_seconds": timeout, + "input_schema": dict(_MESSAGES), + "nodes": [ + { + "id": "draft", + "type": "chat_completion", + "model": model, + "messages": {"from": "$input", "output": "messages"}, + "timeout_seconds": timeout, + }, + { + "id": "prepare", + "type": "python", + "program": { + "contract_version": "1", + "source": _SOURCE, + "function_name": "prepare_followup", + "parameter_names": ["text"], + "source_sha256": hashlib.sha256(_SOURCE.encode()).hexdigest(), + "input_schema": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + "additionalProperties": False, + }, + "output_schema": dict(_MESSAGES), + }, + "inputs": {"text": {"from": "draft", "output": "result"}}, + }, + { + "id": "review", + "type": "chat_completion", + "model": model, + "messages": {"from": "prepare", "output": "result"}, + **( + { + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "Review", + "strict": True, + "schema": { + "type": "object", + "properties": {"approved": {"type": "boolean"}}, + "required": ["approved"], + "additionalProperties": False, + }, + }, + } + } + if structured + else {} + ), + }, + ], + "output": {"from": "review", "output": "result"}, + } + + +def _completion(request: Request) -> Reply: + body: Final = _OBJECT.validate_json(request.body) + assert request.target == "/v1/chat/completions", request.target + assert body["messages"] in ( + [{"role": "user", "content": "Help me update my account email."}], + [{"role": "user", "content": "Review this draft: synthetic draft"}], + ), body + return Reply( + body=json.dumps( + { + "id": "chatcmpl-controlled-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": '{"approved":true}' if "response_format" in body else "synthetic draft", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}, + } + ).encode() + ) + + +def _register(gateway: Gateway, actors: Actors, model: str, *, structured: bool = False, timeout: float = 60) -> Draft: + response: Final = gateway.request( + "POST", + "/v1/workflows", + {"bundle": dict(_bundle(model, structured=structured, timeout=timeout))}, + key=actors.author, + ) + assert response.status_code == 201, response.text + return Draft.model_validate_json(response.content) + + +def _publish(gateway: Gateway, actors: Actors, draft: Draft) -> None: + response: Final = gateway.request( + "POST", + f"/v1/workflows/{draft.workflow_id}/publish", + {"digest": draft.digest}, + key=actors.publisher, + headers={"If-Match": draft.etag}, + ) + assert response.status_code == 200, response.text + + +def _request() -> Mapping[str, JsonValue]: + return { + "model": "workflow/email-reply", + "messages": [{"role": "user", "content": "Help me update my account email."}], + } + + +def _trace(gateway: Gateway, actors: Actors, run_id: str) -> Trace: + response: Final = gateway.request("GET", f"/v1/workflows/runs/{run_id}/steps", key=actors.caller) + assert response.status_code == 200, response.text + return Trace.model_validate_json(response.content) + + +def test_draft_execution_requires_its_owning_team_key_even_for_multiteam_members(workflow_gateway: Gateway) -> None: + with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: + model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") + author: Final = scenario.user(user_role="internal_user") + members: Final = [{"role": "admin", "user_id": author}] + owning_team: Final = scenario.team(models=[model], members_with_roles=members) + other_team: Final = scenario.team(models=[model], members_with_roles=members) + owning_key: Final = scenario.key(team_id=owning_team, user_id=author, models=[model]) + other_key: Final = scenario.key(team_id=other_team, user_id=author, models=[model]) + actors: Final = Actors(owning_key, owning_key, owning_key) + draft: Final = _register(workflow_gateway, actors, model) + refused: Final = workflow_gateway.request( + "POST", + f"/v1/workflows/{draft.workflow_id}/test", + {"messages": _request()["messages"], "authorize_billing": True}, + key=other_key, + headers={"If-Match": draft.etag}, + ) + assert refused.status_code == 403, refused.text + assert ( + workflow_gateway.rows( + 'SELECT run_id FROM "LiteLLM_WorkflowExecution" WHERE principal=%s', + (hashlib.sha256(other_key.encode()).hexdigest(),), + ) + == [] + ) + assert wire.drain() == () + accepted: Final = workflow_gateway.request( + "POST", + f"/v1/workflows/{draft.workflow_id}/test", + {"messages": _request()["messages"], "authorize_billing": True}, + key=owning_key, + headers={"If-Match": draft.etag}, + ) + assert accepted.status_code == 202, accepted.text + run_id: Final = TypeAdapter(str).validate_python(_OBJECT.validate_json(accepted.content)["run_id"]) + trace: Final = eventually( + lambda: _trace(workflow_gateway, actors, run_id), + lambda result: result.state == "completed", + seconds=70, + ) + assert tuple(step.state for step in trace.steps) == ("completed",) * 3 + + +def test_etags_and_independent_publication_do_not_execute_drafts(workflow_gateway: Gateway) -> None: + with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: + model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") + actors: Final = _actors(scenario, model) + draft: Final = _register(workflow_gateway, actors, model) + validation: Final = workflow_gateway.request( + "POST", f"/v1/workflows/{draft.workflow_id}/validate", key=actors.author, headers={"If-Match": draft.etag} + ) + assert validation.status_code == 200, validation.text + self_publication: Final = workflow_gateway.request( + "POST", + f"/v1/workflows/{draft.workflow_id}/publish", + {"digest": draft.digest}, + key=actors.author, + headers={"If-Match": draft.etag}, + ) + assert self_publication.status_code == 403, self_publication.text + for method, suffix, body in ( + ("PUT", "", {"bundle": dict(_bundle(model))}), + ("POST", "/publish", {"digest": draft.digest}), + ): + stale = workflow_gateway.request( + method, + f"/v1/workflows/{draft.workflow_id}{suffix}", + body, + key=actors.publisher, + headers={"If-Match": '"stale"'}, + ) + assert stale.status_code == 412, stale.text + listed: Final = workflow_gateway.request("GET", "/v1/workflows", key=actors.author) + assert listed.status_code == 200, listed.text + assert draft.workflow_id in listed.text + assert wire.drain() == () + _publish(workflow_gateway, actors, draft) + invoked: Final = workflow_gateway.request("POST", "/v1/chat/completions", _request(), key=actors.caller) + assert invoked.status_code == 200, invoked.text + assert len(wire.drain()) == 2 + + +@pytest.mark.parametrize("structured", [False, True]) +def test_typed_checkpoints_idempotency_bounded_delegation_and_single_child_billing( + workflow_gateway: Gateway, + structured: bool, +) -> None: + with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: + model: Final = scenario.model( + model="openai/gpt-6-sol", api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002 + ) + actors: Final = _actors(scenario, model) + draft: Final = _register(workflow_gateway, actors, model, structured=structured) + _publish(workflow_gateway, actors, draft) + denied: Final = workflow_gateway.request( + "POST", "/v1/chat/completions", {**_request(), "model": model}, key=actors.caller + ) + assert denied.status_code in (401, 403), denied.text + response: Final = workflow_gateway.request( + "POST", + "/v1/chat/completions", + _request(), + key=actors.caller, + headers={"Idempotency-Key": "controlled-retry"}, + ) + assert response.status_code == 200, response.text + run_id: Final = response.headers["x-litellm-workflow-run-id"] + retried: Final = workflow_gateway.request( + "POST", + "/v1/chat/completions", + _request(), + key=actors.caller, + headers={"Idempotency-Key": "controlled-retry"}, + ) + assert retried.status_code == 200, retried.text + assert retried.headers["x-litellm-workflow-run-id"] == run_id + assert retried.json() == response.json() + trace: Final = _trace(workflow_gateway, actors, run_id) + assert trace.state == "completed" + assert tuple(step.state for step in trace.steps) == ("completed",) * 3 + assert trace.steps[1].input == {"text": "synthetic draft"} + assert trace.steps[1].output == [{"role": "user", "content": "Review this draft: synthetic draft"}] + assert trace.steps[2].output == ({"approved": True} if structured else "synthetic draft") + assert tuple(step.attempts for step in trace.steps) == (1, 1, 1) + assert trace.steps[0].response_cost == pytest.approx(2 * 0.001 + 3 * 0.002) + checkpoints: Final = workflow_gateway.rows( + 'SELECT checkpoints FROM "LiteLLM_WorkflowExecution" WHERE run_id=%s', (run_id,) + ) + assert len(TypeAdapter(tuple[object, ...]).validate_python(checkpoints[0]["checkpoints"])) == 3 + assert len(wire.drain()) == 2 + charges: Final = eventually( + lambda: workflow_gateway.rows( + 'SELECT spend, model FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (run_id,) + ), + lambda rows: len(rows) == 2, + seconds=70, + ) + assert sum(TypeAdapter(float).validate_python(row["spend"]) for row in charges) == pytest.approx( + 2 * (2 * 0.001 + 3 * 0.002) + ) + assert all(row["model"] != "workflow/email-reply" for row in charges) + mutation: Final = workflow_gateway.request("PATCH", f"/v1/workflows/runs/{run_id}", {"status": "failed"}) + assert mutation.status_code == 409, mutation.text + + +def test_cancellation_stops_the_next_node_without_cancelling_disconnected_work(workflow_gateway: Gateway) -> None: + entered: Final = threading.Event() + release: Final = threading.Event() + + def blocked(request: Request) -> Reply: + entered.set() + assert release.wait(30), "test did not release the controlled upstream" + return _completion(request) + + with workflow_gateway.scenario() as scenario, wire_server(blocked) as wire: + model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") + actors: Final = _actors(scenario, model) + draft: Final = _register(workflow_gateway, actors, model) + _publish(workflow_gateway, actors, draft) + try: + with pytest.raises(httpx.ReadTimeout): + workflow_gateway.client.request( + "POST", + "/v1/chat/completions", + json=_request(), + headers={"Idempotency-Key": "detached-run", "Authorization": f"Bearer {actors.caller}"}, + timeout=0.5, + ) + assert entered.wait(15), "detached invocation never reached the controlled upstream" + executions: Final = workflow_gateway.rows( + 'SELECT run_id, state FROM "LiteLLM_WorkflowExecution" WHERE name=%s AND state=%s', + ("email-reply", "running"), + ) + assert len(executions) == 1, executions + run_id: Final = str(executions[0]["run_id"]) + cancelled: Final = workflow_gateway.request( + "POST", f"/v1/workflows/runs/{run_id}/cancel", key=actors.caller + ) + assert cancelled.status_code == 202, cancelled.text + final: Final = eventually( + lambda: _trace(workflow_gateway, actors, run_id), lambda trace: trace.state == "cancelled", seconds=15 + ) + assert len(final.steps) == 1 + finally: + release.set() + assert len(wire.drain()) == 1 diff --git a/tests/integration/management/workflow_proxy_config.yaml b/tests/integration/management/workflow_proxy_config.yaml new file mode 100644 index 00000000000..69a17874eed --- /dev/null +++ b/tests/integration/management/workflow_proxy_config.yaml @@ -0,0 +1,18 @@ +model_list: [] +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + disable_model_info_refresh: true + disable_spend_logs: false + proxy_batch_write_at: 1 + proxy_batch_polling_interval: 1 + enable_workflows: true + coordination_redis: + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT +router_settings: + redis_host: os.environ/REDIS_HOST + redis_port: os.environ/REDIS_PORT + num_retries: 0 + disable_cooldowns: true diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index c1011310af3..ec3993ea6bd 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -21,6 +21,23 @@ from litellm.proxy.auth.login_utils import PASSWORD_RESET_ALLOWED_ROUTES from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router + +@pytest.mark.parametrize(("method", "path", "allowed"), ( + ("POST", "/v1/workflows", True), + ("GET", "/v1/workflows", True), + ("PUT", "/v1/workflows/draft-id", True), + ("POST", "/v1/workflows/draft-id/publish", True), + ("GET", "/v1/workflows/runs/run-id/steps", True), + ("POST", "/v1/workflows/runs/run-id/cancel", True), + ("POST", "/v1/workflows/runs", False), + ("GET", "/v1/workflows/runs", False), + ("PATCH", "/v1/workflows/runs/run-id", False), + ("POST", "/v1/workflows/runs/run-id/events", False), + ("DELETE", "/v1/workflows/draft-id", False), +)) +def test_workflow_execution_rbac_does_not_grant_legacy_tracking_access(method: str, path: str, allowed: bool) -> None: + assert RouteChecks.is_workflow_execution_route(path, method) is allowed + DAILY_ACTIVITY_ROUTE_PAIRS: Final[tuple[tuple[str, str], ...]] = ( ("/user/daily/activity", "/user/daily/activity/aggregated"), ("/user/daily/activity", "/user/daily/activity/aggregated/keys"), @@ -2904,7 +2921,7 @@ def test_internal_user_can_access_key_reset_spend_route(): request = MagicMock(spec=Request) request.query_params = {} - key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae" + key_hash = "a" * 64 route = f"/key/{key_hash}/reset_spend" # Should not raise — the route-level check must pass for team admins @@ -2936,7 +2953,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re request = MagicMock(spec=Request) request.query_params = {} - key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae" + key_hash = "a" * 64 # /key/{hash}/reset_spend passes the route check for internal_user RouteChecks.non_proxy_admin_allowed_routes_check( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f9dae53bd03..04adf9ed69a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8503,7 +8503,7 @@ export interface paths { * Example: * ```bash * curl --location 'http://0.0.0.0:4000/key/block' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ - * "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + * "key": "YOUR_VIRTUAL_KEY" * }' * ``` * @@ -8819,12 +8819,12 @@ export interface paths { * * Example Curl: * ``` - * curl -X GET "http://0.0.0.0:4000/key/info?key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" -H "Authorization: Bearer $LITELLM_MASTER_KEY" + * curl -X GET "http://0.0.0.0:4000/key/info?key=YOUR_VIRTUAL_KEY" -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" * ``` * * Example Curl - if no key is passed, it will use the Key Passed in Authorization Header * ``` - * curl -X GET "http://0.0.0.0:4000/key/info" -H "Authorization: Bearer sk-test-example-key-123" + * curl -X GET "http://0.0.0.0:4000/key/info" -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" * ``` */ get: operations["info_key_fn_key_info_get"]; @@ -9062,7 +9062,7 @@ export interface paths { * Example: * ```bash * curl --location 'http://0.0.0.0:4000/key/unblock' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ - * "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + * "key": "YOUR_VIRTUAL_KEY" * }' * ``` * @@ -16246,7 +16246,7 @@ export interface paths { * * Example Request for specific api_key * ``` - * curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" -H "Authorization: Bearer $LITELLM_MASTER_KEY" + * curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=YOUR_VIRTUAL_KEY" -H "Authorization: Bearer $LITELLM_MASTER_KEY" * ``` * * Example Request for specific user_id @@ -17557,7 +17557,7 @@ export interface paths { * - team_id: str - Required. The unique identifier of the team to get info on. * * ``` - * curl --location 'http://localhost:4000/team/info?team_id=your_team_id_here' --header 'Authorization: Bearer your_api_key_here' + * curl --location 'http://localhost:4000/team/info?team_id=your_team_id_here' -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" * ``` */ get: operations["team_info_team_info_get"]; @@ -18381,7 +18381,7 @@ export interface paths { * Returns 404 if the caller is not a member of the team. * * ``` - * curl --location 'http://localhost:4000/team/your_team_id/members/me' --header 'Authorization: Bearer your_api_key_here' + * curl --location 'http://localhost:4000/team/your_team_id/members/me' -H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}" * ``` */ get: operations["team_member_me_team__team_id__members_me_get"]; @@ -24080,6 +24080,27 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/workflows": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** List Workflows */ + get: operations["list_workflows_v1_workflows_get"]; + put?: never; + /** + * Create Workflow + * @description Register a team-scoped draft without publishing or executing it. + */ + post: operations["create_workflow_v1_workflows_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/workflows/runs": { parameters: { query?: never; @@ -24133,6 +24154,26 @@ 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 Run + * @description Request cancellation and prevent further node dispatch for this run. + */ + post: operations["cancel_run_v1_workflows_runs__run_id__cancel_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/workflows/runs/{run_id}/events": { parameters: { query?: never; @@ -24187,6 +24228,118 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/workflows/runs/{run_id}/steps": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Read Steps */ + get: operations["read_steps_v1_workflows_runs__run_id__steps_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/workflows/runs/{run_id}/steps/{step_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Read Step */ + get: operations["read_step_v1_workflows_runs__run_id__steps__step_id__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 */ + get: operations["get_workflow_v1_workflows__workflow_id__get"]; + /** Replace Workflow */ + put: operations["replace_workflow_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 + * @description Publish an immutable snapshot after independent approval of the current ETag. + */ + post: operations["publish_workflow_v1_workflows__workflow_id__publish_post"]; + 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 Draft + * @description Execute the reviewed draft with explicit billing authorization. + */ + post: operations["test_draft_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 + * @description Validate the reviewed draft without executing Python or calling models. + */ + post: operations["validate_workflow_v1_workflows__workflow_id__validate_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1beta/agents": { parameters: { query?: never; @@ -27298,6 +27451,15 @@ export interface components { */ unnamed_teams_count: number; }; + /** + * Audio + * @description Data about a previous audio response from the model. + * [Learn more](https://platform.openai.com/docs/guides/audio). + */ + Audio: { + /** Id */ + id: string; + }; /** * AuditLogResponse * @description Response model for a single audit log entry @@ -28262,6 +28424,11 @@ export interface components { */ start: string; }; + /** + * BindingOutput + * @enum {string} + */ + BindingOutput: "messages" | "result"; /** BlockKeyRequest */ BlockKeyRequest: { /** Key */ @@ -29541,6 +29708,13 @@ export interface components { */ status: "cancelled"; }; + /** CancellationView */ + CancellationView: { + /** Cancel Requested */ + cancel_requested: boolean; + /** Run Id */ + run_id: string; + }; /** CapabilityCalibrationConfig */ CapabilityCalibrationConfig: { /** Intercept */ @@ -29687,6 +29861,27 @@ export interface components { /** Tool Calls */ tool_calls?: components["schemas"]["ChatCompletionAssistantToolCall"][] | null; }; + /** + * ChatCompletionAssistantMessageParam + * @description Messages sent by the model in response to user messages. + */ + ChatCompletionAssistantMessageParam: { + audio?: components["schemas"]["Audio"] | null; + /** Content */ + content?: string | (components["schemas"]["ChatCompletionContentPartTextParam"] | components["schemas"]["ChatCompletionContentPartRefusalParam"])[] | null; + function_call?: components["schemas"]["openai__types__chat__chat_completion_assistant_message_param__FunctionCall"] | null; + /** Name */ + name?: string; + /** Refusal */ + refusal?: string | null; + /** + * Role + * @constant + */ + role: "assistant"; + /** Tool Calls */ + tool_calls?: (components["schemas"]["ChatCompletionMessageFunctionToolCallParam"] | components["schemas"]["ChatCompletionMessageCustomToolCallParam"])[]; + }; /** ChatCompletionAssistantToolCall */ ChatCompletionAssistantToolCall: { function: components["schemas"]["ChatCompletionToolCallFunctionChunk"]; @@ -29736,6 +29931,53 @@ export interface components { } & { [key: string]: unknown; }; + /** + * ChatCompletionContentPartImageParam + * @description Learn about [image inputs](https://platform.openai.com/docs/guides/vision). + */ + ChatCompletionContentPartImageParam: { + image_url: components["schemas"]["ImageURL"]; + /** + * Type + * @constant + */ + type: "image_url"; + }; + /** + * ChatCompletionContentPartInputAudioParam + * @description Learn about [audio inputs](https://platform.openai.com/docs/guides/audio). + */ + ChatCompletionContentPartInputAudioParam: { + input_audio: components["schemas"]["InputAudio"]; + /** + * Type + * @constant + */ + type: "input_audio"; + }; + /** ChatCompletionContentPartRefusalParam */ + ChatCompletionContentPartRefusalParam: { + /** Refusal */ + refusal: string; + /** + * Type + * @constant + */ + type: "refusal"; + }; + /** + * ChatCompletionContentPartTextParam + * @description Learn about [text inputs](https://platform.openai.com/docs/guides/text-generation). + */ + ChatCompletionContentPartTextParam: { + /** Text */ + text: string; + /** + * Type + * @constant + */ + type: "text"; + }; /** ChatCompletionCustomToolCallPayload */ ChatCompletionCustomToolCallPayload: { /** Input */ @@ -29758,6 +30000,23 @@ export interface components { */ role: "developer"; }; + /** + * ChatCompletionDeveloperMessageParam + * @description Developer-provided instructions that the model should follow, regardless of + * messages sent by the user. With o1 models and newer, `developer` messages + * replace the previous `system` messages. + */ + ChatCompletionDeveloperMessageParam: { + /** Content */ + content: string | components["schemas"]["ChatCompletionContentPartTextParam"][]; + /** Name */ + name?: string; + /** + * Role + * @constant + */ + role: "developer"; + }; /** ChatCompletionDocumentObject */ ChatCompletionDocumentObject: { citations: components["schemas"]["CitationsObject"] | null; @@ -29812,6 +30071,18 @@ export interface components { /** Tool Call Id */ tool_call_id: string | null; }; + /** ChatCompletionFunctionMessageParam */ + ChatCompletionFunctionMessageParam: { + /** Content */ + content: string | null; + /** Name */ + name: string; + /** + * Role + * @constant + */ + role: "function"; + }; /** ChatCompletionImageObject */ ChatCompletionImageObject: { /** Image Url */ @@ -29856,10 +30127,77 @@ export interface components { } & { [key: string]: unknown; }; + /** + * ChatCompletionMessageCustomToolCallParam + * @description A call to a custom tool created by the model. + */ + ChatCompletionMessageCustomToolCallParam: { + custom: components["schemas"]["Custom"]; + /** Id */ + id: string; + /** + * Type + * @constant + */ + type: "custom"; + }; + /** + * ChatCompletionMessageFunctionToolCallParam + * @description A call to a function tool created by the model. + */ + ChatCompletionMessageFunctionToolCallParam: { + function: components["schemas"]["Function"]; + /** Id */ + id: string; + /** + * Type + * @constant + */ + type: "function"; + }; /** ChatCompletionMessageToolCall */ 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"]["Reference"] | components["schemas"]["Constant"]; + /** Model */ + model: string; + response_format?: components["schemas"]["StructuredResponse"] | null; + /** + * @default { + * "backoff_seconds": 0, + * "max_attempts": 1 + * } + */ + retry: components["schemas"]["litellm__models__workflow_bundle__RetryPolicy"]; + /** Seed */ + seed?: number | null; + /** Stop */ + stop?: string | string[] | null; + /** Temperature */ + temperature?: number | null; + /** + * Timeout Seconds + * @default 60 + */ + 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. @@ -29922,6 +30260,23 @@ export interface components { */ role: "system"; }; + /** + * ChatCompletionSystemMessageParam + * @description Developer-provided instructions that the model should follow, regardless of + * messages sent by the user. With o1 models and newer, use `developer` messages + * for this purpose instead. + */ + ChatCompletionSystemMessageParam: { + /** Content */ + content: string | components["schemas"]["ChatCompletionContentPartTextParam"][]; + /** Name */ + name?: string; + /** + * Role + * @constant + */ + role: "system"; + }; /** ChatCompletionTextObject */ ChatCompletionTextObject: { cache_control?: components["schemas"]["ChatCompletionCachedContent"]; @@ -30002,6 +30357,18 @@ export interface components { /** Tool Call Id */ tool_call_id: string; }; + /** ChatCompletionToolMessageParam */ + ChatCompletionToolMessageParam: { + /** Content */ + content: string | components["schemas"]["ChatCompletionContentPartTextParam"][]; + /** + * Role + * @constant + */ + role: "tool"; + /** Tool Call Id */ + tool_call_id: string; + }; /** ChatCompletionToolParam */ ChatCompletionToolParam: { /** Allowed Callers */ @@ -30054,6 +30421,22 @@ export interface components { */ role: "user"; }; + /** + * ChatCompletionUserMessageParam + * @description Messages sent by an end user, containing prompts or additional context + * information. + */ + ChatCompletionUserMessageParam: { + /** Content */ + content: string | (components["schemas"]["ChatCompletionContentPartTextParam"] | components["schemas"]["ChatCompletionContentPartImageParam"] | components["schemas"]["ChatCompletionContentPartInputAudioParam"] | components["schemas"]["File"])[]; + /** Name */ + name?: string; + /** + * Role + * @constant + */ + role: "user"; + }; /** ChatCompletionVideoObject */ ChatCompletionVideoObject: { /** @@ -30754,7 +31137,7 @@ export interface components { agent_access_default_deny: boolean; /** * Alert To Webhook Url - * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}` + * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': '$ALERT_WEBHOOK_URL'}` */ alert_to_webhook_url?: { [key: string]: unknown; @@ -30964,6 +31347,12 @@ export interface components { * @default false */ enable_public_model_hub: boolean; + /** + * Enable Workflows + * @description Enable durable typed workflows and the optional Monty worker + * @default false + */ + enable_workflows: boolean; /** * Enforce Fallback Model Access * @description If True, router fallbacks configured in router_settings are only attempted when the calling key (and its team and project) is allowed to call the fallback model; unauthorized fallback targets are skipped and the primary model's error is returned. Default is True. Set False to disable fallback model access enforcement. @@ -31500,6 +31889,10 @@ export interface components { /** Api Base */ api_base: string; }; + /** Constant */ + Constant: { + value: components["schemas"]["JsonValue"]; + }; /** ContainerAuto */ ContainerAuto: { /** File Ids */ @@ -32205,6 +32598,16 @@ export interface components { user_id: string; user_role: components["schemas"]["LitellmUserRoles"] | null; }; + /** + * Custom + * @description The custom tool that the model called. + */ + Custom: { + /** Input */ + input: string; + /** Name */ + name: string; + }; /** CustomDimension */ CustomDimension: { /** @@ -32994,6 +33397,42 @@ export interface components { } & { [key: string]: unknown; }; + /** DraftList */ + DraftList: { + /** Next Cursor */ + next_cursor: string | null; + /** Workflows */ + workflows: components["schemas"]["DraftView"][]; + }; + /** DraftTest */ + DraftTest: { + /** + * Authorize Billing + * @constant + */ + authorize_billing: true; + /** Messages */ + messages: (components["schemas"]["ChatCompletionDeveloperMessageParam"] | components["schemas"]["ChatCompletionSystemMessageParam"] | components["schemas"]["ChatCompletionUserMessageParam"] | components["schemas"]["ChatCompletionAssistantMessageParam"] | components["schemas"]["ChatCompletionToolMessageParam"] | components["schemas"]["ChatCompletionFunctionMessageParam"])[]; + metadata?: components["schemas"]["InvocationMetadata"] | null; + /** User */ + user?: string | null; + }; + /** DraftView */ + DraftView: { + 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; + }; /** * Drag * @description A drag action. @@ -33193,6 +33632,8 @@ export interface components { } & { [key: string]: unknown; }; + /** EmptyWorkflowAction */ + EmptyWorkflowAction: Record; /** EndpointProvider */ EndpointProvider: { /** Display Name */ @@ -33562,6 +34003,27 @@ export interface components { /** Stored In Db */ stored_in_db: boolean | null; }; + /** + * File + * @description Learn about [file inputs](https://platform.openai.com/docs/guides/text) for text generation. + */ + File: { + file: components["schemas"]["FileFile"]; + /** + * Type + * @constant + */ + type: "file"; + }; + /** FileFile */ + FileFile: { + /** File Data */ + file_data?: string; + /** File Id */ + file_id?: string; + /** Filename */ + filename?: string; + }; /** * FileSearchTool * @description A tool that searches for relevant content from uploaded files. @@ -33614,14 +34076,15 @@ export interface components { } & { [key: string]: unknown; }; - /** FunctionCall */ - FunctionCall: { + /** + * Function + * @description The function that the model called. + */ + Function: { /** Arguments */ arguments: string; /** Name */ - name?: string | null; - } & { - [key: string]: unknown; + name: string; }; /** * FunctionCallOutput @@ -34618,6 +35081,16 @@ export interface components { } & { [key: string]: unknown; }; + /** ImageURL */ + ImageURL: { + /** + * Detail + * @enum {string} + */ + detail?: "auto" | "low" | "high"; + /** Url */ + url: string; + }; /** ImageURLListItem */ ImageURLListItem: { image_url: components["schemas"]["ImageURLObject"]; @@ -34773,6 +35246,16 @@ export interface components { */ type: "base64"; }; + /** InputAudio */ + InputAudio: { + /** Data */ + data: string; + /** + * Format + * @enum {string} + */ + format: "wav" | "mp3"; + }; /** InputTokensDetails */ InputTokensDetails: { /** Audio Tokens */ @@ -34865,6 +35348,14 @@ export interface components { /** Is Accepted */ is_accepted: boolean; }; + /** InvocationMetadata */ + InvocationMetadata: { + /** + * Tags + * @default [] + */ + tags: string[]; + }; /** * ItemReference * @description An internal identifier for an item to reference. @@ -39504,7 +39995,7 @@ export interface components { audio?: components["schemas"]["ChatCompletionAudioResponse"] | null; /** Content */ content: string | null; - function_call: components["schemas"]["FunctionCall"] | null; + function_call: components["schemas"]["litellm__types__utils__FunctionCall"] | null; /** Images */ images?: components["schemas"]["ImageURLListItem"][] | null; /** Provider Specific Fields */ @@ -43287,6 +43778,71 @@ export interface components { }; } | null; }; + /** PublicationView */ + PublicationView: { + /** Digest */ + digest: string; + /** Model */ + model: string; + /** Snapshot Id */ + snapshot_id: string; + }; + /** PublishWorkflow */ + PublishWorkflow: { + /** Digest */ + digest: string; + }; + /** PythonNode */ + PythonNode: { + /** Id */ + id: string; + /** Inputs */ + inputs: { + [key: string]: components["schemas"]["Reference"] | components["schemas"]["Constant"]; + }; + program: components["schemas"]["PythonProgram"]; + /** + * @default { + * "backoff_seconds": 0, + * "max_attempts": 1 + * } + */ + retry: components["schemas"]["litellm__models__workflow_bundle__RetryPolicy"]; + /** + * Timeout Seconds + * @default 60 + */ + 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. @@ -43375,6 +43931,14 @@ export interface components { /** Summary */ summary?: ("auto" | "concise" | "detailed") | null; }; + /** Reference */ + Reference: { + /** From */ + from: string; + output: components["schemas"]["BindingOutput"]; + /** Path */ + path?: string | null; + }; /** RegenerateKeyRequest */ RegenerateKeyRequest: { /** Access Group Ids */ @@ -43660,6 +44224,10 @@ export interface components { */ status: string; }; + /** RegisterWorkflow */ + RegisterWorkflow: { + bundle: components["schemas"]["WorkflowBundle"]; + }; /** RejectMCPServerRequest */ RejectMCPServerRequest: { /** Review Notes */ @@ -45360,6 +45928,20 @@ export interface components { */ type: "reasoning"; }; + /** ResponseSchema */ + ResponseSchema: { + /** Name */ + name: string; + /** Schema */ + schema: { + [key: string]: components["schemas"]["JsonValue"]; + }; + /** + * Strict + * @default true + */ + strict: boolean; + }; /** * ResponseTextConfigParam * @description Configuration options for a text response from the model. @@ -45727,33 +46309,6 @@ export interface components { /** Total */ total: number; }; - /** - * RetryPolicy - * @description Use this to set a custom number of retries per exception type - * If RateLimitErrorRetries = 3, then 3 retries will be made for RateLimitError - * Mapping of Exception type to number of retries - * https://docs.litellm.ai/docs/exception_mapping - */ - RetryPolicy: { - /** Authenticationerrorretries */ - AuthenticationErrorRetries?: number | null; - /** Badrequesterrorretries */ - BadRequestErrorRetries?: number | null; - /** Contentpolicyviolationerrorretries */ - ContentPolicyViolationErrorRetries?: number | null; - /** Defaultretries */ - DefaultRetries?: number | null; - /** Internalservererrorretries */ - InternalServerErrorRetries?: number | null; - /** Notfounderrorretries */ - NotFoundErrorRetries?: number | null; - /** Ratelimiterrorretries */ - RateLimitErrorRetries?: number | null; - /** Serviceunavailableerrorretries */ - ServiceUnavailableErrorRetries?: number | null; - /** Timeouterrorretries */ - TimeoutErrorRetries?: number | null; - }; /** * RoleMappings * @description Configuration for mapping SSO groups to LiteLLM roles. @@ -45932,6 +46487,15 @@ export interface components { */ status: "queued" | "running" | "completed" | "failed" | "cancelled"; }; + /** RunAccepted */ + RunAccepted: { + /** Run Id */ + run_id: string; + /** Snapshot Id */ + snapshot_id: string; + /** Trace Url */ + trace_url: string; + }; /** * RunDeleteResponse * @description Response from deleting a run @@ -45964,6 +46528,18 @@ export interface components { /** User */ user?: string; }; + /** RunSteps */ + RunSteps: { + /** Cancel Requested */ + cancel_requested: boolean; + /** Run Id */ + run_id: string; + /** Snapshot Id */ + snapshot_id: string; + state: components["schemas"]["WorkflowRunState"]; + /** Steps */ + steps: components["schemas"]["WorkflowStep"][]; + }; /** SCIMEnterpriseUser */ SCIMEnterpriseUser: { /** Costcenter */ @@ -47731,6 +48307,16 @@ export interface components { /** Window Start */ window_start: number; }; + /** StructuredResponse */ + StructuredResponse: { + json_schema: components["schemas"]["ResponseSchema"]; + /** + * Type + * @default json_schema + * @constant + */ + type: "json_schema"; + }; /** SubscriptionAccountFeeRequest */ SubscriptionAccountFeeRequest: { /** Account Id */ @@ -50485,7 +51071,7 @@ export interface components { } | null; /** Model Group Retry Policy */ model_group_retry_policy?: { - [key: string]: components["schemas"]["RetryPolicy"]; + [key: string]: components["schemas"]["litellm__types__router__RetryPolicy"]; } | null; /** Num Retries */ num_retries?: number | null; @@ -50493,7 +51079,7 @@ export interface components { optional_pre_call_checks?: ("prompt_caching" | "router_budget_limiting" | "responses_api_deployment_check" | "deployment_affinity" | "session_affinity" | "forward_client_headers_by_model_group" | "enforce_model_rate_limits" | "encrypted_content_affinity")[] | null; /** Retry After */ retry_after?: number | null; - retry_policy?: components["schemas"]["RetryPolicy"] | null; + retry_policy?: components["schemas"]["litellm__types__router__RetryPolicy"] | null; /** Routing Groups */ routing_groups?: components["schemas"]["RoutingGroup"][] | null; /** Routing Strategy */ @@ -51751,6 +52337,13 @@ export interface components { /** Error Type */ type: string; }; + /** ValidationView */ + ValidationView: { + /** Digest */ + digest: string; + /** Valid */ + valid: boolean; + }; /** * VantageDryRunRequest * @description Request model for Vantage dry-run operations (capped for preview) @@ -52063,6 +52656,28 @@ 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"]["Reference"]; + /** + * Timeout Seconds + * @default 300 + */ + timeout_seconds: number; + }; /** WorkflowEventCreateRequest */ WorkflowEventCreateRequest: { /** Data */ @@ -52074,6 +52689,25 @@ export interface components { /** Step Name */ step_name: string; }; + /** WorkflowEventResponse */ + WorkflowEventResponse: { + /** + * Created At + * Format: date-time + */ + created_at: string; + data: components["schemas"]["JsonValue"]; + /** Event Id */ + event_id: string; + /** Event Type */ + event_type: string; + /** Run Id */ + run_id: string; + /** Sequence Number */ + sequence_number: number; + /** Step Name */ + step_name: string; + }; /** WorkflowMessageCreateRequest */ WorkflowMessageCreateRequest: { /** Content */ @@ -52083,6 +52717,26 @@ export interface components { /** Session Id */ session_id?: string | null; }; + /** WorkflowMessageResponse */ + WorkflowMessageResponse: { + /** Content */ + content: string; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Message Id */ + message_id: string; + /** Role */ + role: string; + /** Run Id */ + run_id: string; + /** Sequence Number */ + sequence_number: number; + /** Session Id */ + session_id: string | null; + }; /** WorkflowRunCreateRequest */ WorkflowRunCreateRequest: { /** Input */ @@ -52096,6 +52750,39 @@ export interface components { /** Workflow Type */ workflow_type: string; }; + /** WorkflowRunResponse */ + WorkflowRunResponse: { + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string | null; + /** Engine Owned */ + engine_owned: boolean; + input: components["schemas"]["JsonValue"]; + metadata: components["schemas"]["JsonValue"]; + output: components["schemas"]["JsonValue"]; + /** Run Id */ + run_id: string; + /** Session Id */ + session_id: string; + /** Status */ + status: string; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Workflow Type */ + workflow_type: string; + }; + /** + * WorkflowRunState + * @enum {string} + */ + WorkflowRunState: "pending" | "running" | "completed" | "failed" | "cancelled"; /** WorkflowRunUpdateRequest */ WorkflowRunUpdateRequest: { /** Metadata */ @@ -52109,6 +52796,48 @@ export interface components { /** Status */ status?: ("pending" | "running" | "paused" | "completed" | "failed") | null; }; + /** WorkflowStep */ + WorkflowStep: { + /** Attempts */ + attempts: number; + /** + * 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"] | null; + /** + * Prompt Tokens + * @default 0 + */ + prompt_tokens: number; + /** Request Id */ + request_id?: string | null; + /** Response Cost */ + response_cost?: number | null; + /** Run Id */ + run_id: string; + /** + * Started At + * Format: date-time + */ + started_at: string; + state: components["schemas"]["WorkflowRunState"]; + /** Step Id */ + step_id: string; + }; /** _A2AJsonRpcResponse */ _A2AJsonRpcResponse: { /** Error */ @@ -52127,6 +52856,19 @@ export interface components { } & { [key: string]: unknown; }; + /** RetryPolicy */ + litellm__models__workflow_bundle__RetryPolicy: { + /** + * Backoff Seconds + * @default 0 + */ + backoff_seconds: number; + /** + * Max Attempts + * @default 1 + */ + max_attempts: number; + }; /** ModelInfo */ litellm__proxy___types__ModelInfo: { /** Base Model */ @@ -52269,6 +53011,54 @@ export interface components { } & { [key: string]: unknown; }; + /** + * RetryPolicy + * @description Use this to set a custom number of retries per exception type + * If RateLimitErrorRetries = 3, then 3 retries will be made for RateLimitError + * Mapping of Exception type to number of retries + * https://docs.litellm.ai/docs/exception_mapping + */ + litellm__types__router__RetryPolicy: { + /** Authenticationerrorretries */ + AuthenticationErrorRetries?: number | null; + /** Badrequesterrorretries */ + BadRequestErrorRetries?: number | null; + /** Contentpolicyviolationerrorretries */ + ContentPolicyViolationErrorRetries?: number | null; + /** Defaultretries */ + DefaultRetries?: number | null; + /** Internalservererrorretries */ + InternalServerErrorRetries?: number | null; + /** Notfounderrorretries */ + NotFoundErrorRetries?: number | null; + /** Ratelimiterrorretries */ + RateLimitErrorRetries?: number | null; + /** Serviceunavailableerrorretries */ + ServiceUnavailableErrorRetries?: number | null; + /** Timeouterrorretries */ + TimeoutErrorRetries?: number | null; + }; + /** FunctionCall */ + litellm__types__utils__FunctionCall: { + /** Arguments */ + arguments: string; + /** Name */ + name?: string | null; + } & { + [key: string]: unknown; + }; + /** + * FunctionCall + * @description Deprecated and replaced by `tool_calls`. + * + * The name and arguments of a function that should be called, as generated by the model. + */ + openai__types__chat__chat_completion_assistant_message_param__FunctionCall: { + /** Arguments */ + arguments: string; + /** Name */ + name: string; + }; /** * ComputerToolParam * @description A tool that controls a virtual computer. @@ -63583,7 +64373,7 @@ export interface operations { end_date?: string | null; /** @description Group spend by internal team or customer or api_key */ group_by?: ("team" | "customer" | "api_key") | null; - /** @description View spend for a specific api_key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa' */ + /** @description View spend for a specific api_key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example api_key='YOUR_VIRTUAL_KEY' */ api_key?: string | null; /** @description View spend for a specific internal_user_id. Example internal_user_id='1234 */ internal_user_id?: string | null; @@ -65513,7 +66303,7 @@ export interface operations { info_key_fn_key_info_get: { parameters: { query?: { - /** @description Key to look up. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa' */ + /** @description Key to look up. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example key='YOUR_VIRTUAL_KEY' */ key?: string | null; }; header?: never; @@ -65693,7 +66483,7 @@ export interface operations { start_date?: string | null; /** @description Time till which to view spend (YYYY-MM-DD) */ end_date?: string | null; - /** @description View spend for a specific api_key. Proxy admin only; other callers are scoped to their own key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa' */ + /** @description View spend for a specific api_key. Proxy admin only; other callers are scoped to their own key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. Example api_key='YOUR_VIRTUAL_KEY' */ api_key?: string | null; }; header?: never; @@ -87070,6 +87860,74 @@ export interface operations { }; }; }; + list_workflows_v1_workflows_get: { + parameters: { + query?: { + team_id?: string | null; + after?: string | null; + limit?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DraftList"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + create_workflow_v1_workflows_post: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["RegisterWorkflow"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DraftView"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_workflow_runs_v1_workflows_runs_get: { parameters: { query?: { @@ -87122,7 +87980,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["WorkflowRunResponse"]; }; }; /** @description Validation Error */ @@ -87202,6 +88060,41 @@ export interface operations { }; }; }; + cancel_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"]["EmptyWorkflowAction"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 202: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CancellationView"]; + }; + }; + /** @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?: { @@ -87256,7 +88149,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["WorkflowEventResponse"]; }; }; /** @description Validation Error */ @@ -87324,7 +88217,249 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["WorkflowMessageResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + read_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"]; + }; + }; + }; + }; + read_step_v1_workflows_runs__run_id__steps__step_id__get: { + parameters: { + query?: never; + header?: never; + path: { + run_id: string; + step_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkflowStep"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_workflow_v1_workflows__workflow_id__get: { + parameters: { + query?: never; + header?: never; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DraftView"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + replace_workflow_v1_workflows__workflow_id__put: { + parameters: { + query?: never; + header?: { + "if-match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["RegisterWorkflow"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DraftView"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + publish_workflow_v1_workflows__workflow_id__publish_post: { + parameters: { + query?: never; + header?: { + "if-match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["PublishWorkflow"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["PublicationView"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + test_draft_v1_workflows__workflow_id__test_post: { + parameters: { + query?: never; + header?: { + "if-match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["DraftTest"]; + }; + }; + responses: { + /** @description Successful Response */ + 202: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["RunAccepted"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + validate_workflow_v1_workflows__workflow_id__validate_post: { + parameters: { + query?: never; + header?: { + "if-match"?: string | null; + }; + path: { + workflow_id: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": components["schemas"]["EmptyWorkflowAction"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ValidationView"]; }; }; /** @description Validation Error */