feat(workflows): reuse LiteFlow SDK for sequential gateway execution

This commit is contained in:
Tin 2026-11-06 18:41:46 +08:00
parent 2a802eb8b9
commit 5fd7ee8e9e
49 changed files with 5634 additions and 68 deletions

View file

@ -546,6 +546,7 @@ jobs:
- shard: proxy-endpoints
test-path: |-
tests/unit/proxy/workflows
tests/unit/proxy/management_endpoints
tests/unit/proxy/management_helpers
tests/unit/proxy/list_api

View file

@ -0,0 +1,49 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE "LiteLLM_WorkflowDefinition" (
"workflow_id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"name" TEXT NOT NULL,
"etag" TEXT NOT NULL,
"digest" TEXT NOT NULL,
"bundle_json" TEXT NOT NULL,
"contributors" TEXT[],
"published_snapshot_id" TEXT,
"published_etag" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_WorkflowDefinition_pkey" PRIMARY KEY ("workflow_id")
);
CREATE TABLE "LiteLLM_WorkflowSnapshot" (
"snapshot_id" TEXT NOT NULL,
"workflow_id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"name" TEXT NOT NULL,
"digest" TEXT NOT NULL,
"bundle_json" TEXT NOT NULL,
"contributors" TEXT[],
"published_by" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_WorkflowSnapshot_pkey" PRIMARY KEY ("snapshot_id"),
CONSTRAINT "LiteLLM_WorkflowSnapshot_workflow_id_fkey" FOREIGN KEY ("workflow_id")
REFERENCES "LiteLLM_WorkflowDefinition"("workflow_id") ON DELETE RESTRICT ON UPDATE CASCADE
);
CREATE UNIQUE INDEX "LiteLLM_WorkflowDefinition_team_id_name_key" ON "LiteLLM_WorkflowDefinition"("team_id", "name");
CREATE INDEX "LiteLLM_WorkflowDefinition_team_id_workflow_id_idx" ON "LiteLLM_WorkflowDefinition"("team_id", "workflow_id");
CREATE INDEX "LiteLLM_WorkflowSnapshot_workflow_id_created_at_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("workflow_id", "created_at", "snapshot_id");
CREATE INDEX "LiteLLM_WorkflowSnapshot_team_id_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("team_id", "snapshot_id");
ALTER TABLE "LiteLLM_WorkflowRun"
ADD COLUMN "workflow_id" TEXT,
ADD COLUMN "snapshot_id" TEXT,
ADD COLUMN "team_id" TEXT,
ADD COLUMN "engine_owned" BOOLEAN NOT NULL DEFAULT false,
ADD COLUMN "request_digest" TEXT,
ADD COLUMN "idempotency_key" TEXT,
ADD COLUMN "cancel_requested" BOOLEAN NOT NULL DEFAULT false;
CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_principal_idempotency_key" ON "LiteLLM_WorkflowRun"("created_by", "team_id", "workflow_id", "idempotency_key");
CREATE INDEX "LiteLLM_WorkflowRun_workflow_id_team_id_created_at_run_id_idx" ON "LiteLLM_WorkflowRun"("workflow_id", "team_id", "created_at", "run_id");
CREATE INDEX "LiteLLM_WorkflowRun_team_id_engine_owned_status_idx" ON "LiteLLM_WorkflowRun"("team_id", "engine_owned", "status");

View file

@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
// One instance of work being done. `status` is a materialized cache of the
// latest event; the event log is the authoritative source of truth.
model LiteLLM_WorkflowDefinition {
workflow_id String @id @default(uuid())
team_id String
name String
etag String
digest String
bundle_json String
contributors String[]
published_snapshot_id String?
published_etag String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
snapshots LiteLLM_WorkflowSnapshot[]
@@unique([team_id, name])
@@index([team_id, workflow_id])
}
model LiteLLM_WorkflowSnapshot {
snapshot_id String @id @default(uuid())
workflow_id String
team_id String
name String
digest String
bundle_json String
contributors String[]
published_by String?
created_at DateTime @default(now())
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
@@index([workflow_id, created_at, snapshot_id])
@@index([team_id, snapshot_id])
}
model LiteLLM_WorkflowRun {
run_id String @id @default(uuid())
session_id String @unique @default(uuid())
workflow_type String
workflow_id String?
snapshot_id String?
team_id String?
engine_owned Boolean @default(false)
request_digest String?
idempotency_key String?
cancel_requested Boolean @default(false)
status String @default("pending")
created_by String? // user_id of the key that created this run; null = created by master key
created_at DateTime @default(now())
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
@@index([session_id])
@@index([created_at])
@@index([created_by])
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
@@index([workflow_id, team_id, created_at, run_id])
@@index([team_id, engine_owned, status])
}
// Append-only log of state transitions. Never mutate rows here.

View file

@ -355,6 +355,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_ShadowEvalJob",
"LiteLLM_ShadowEvalAttempt",
"LiteLLM_ShadowEvalFunnel",
"LiteLLM_WorkflowDefinition",
"LiteLLM_WorkflowSnapshot",
"LiteLLM_WorkflowRun",
"LiteLLM_WorkflowEvent",
"LiteLLM_WorkflowMessage",

View file

@ -0,0 +1,165 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Annotated, Literal, Protocol, TypeAlias
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
Draft as WorkflowDraft,
)
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StringConstraints
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy.workflows.monty.models import PythonProgram
WorkflowName: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$")]
Digest: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")]
class ValueBinding(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
value: JsonValue
class ReferenceBinding(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
source: str = Field(alias="from", min_length=1, max_length=128)
output: Literal["messages", "result"]
Binding: TypeAlias = ValueBinding | ReferenceBinding
class WorkflowRetry(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
max_attempts: int = Field(default=1, ge=1, le=3, strict=True)
backoff_seconds: float = Field(default=0, ge=0, le=30)
class StructuredSchema(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
name: str = Field(min_length=1, max_length=128)
strict: Literal[True]
schema_: Mapping[str, JsonValue] = Field(alias="schema")
class StructuredOutput(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
type: Literal["json_schema"]
json_schema: StructuredSchema
class ChatCompletionNode(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
id: WorkflowName
type: Literal["chat_completion"]
model: str = Field(min_length=1, max_length=256)
messages: Binding
timeout_seconds: float = Field(gt=0, le=300)
retry: WorkflowRetry
temperature: float | None = Field(default=None, ge=0, le=2)
top_p: float | None = Field(default=None, ge=0, le=1)
max_tokens: int | None = Field(default=None, gt=0, strict=True)
max_completion_tokens: int | None = Field(default=None, gt=0, strict=True)
seed: int | None = Field(default=None, strict=True)
stop: str | tuple[str, ...] | None = None
response_format: StructuredOutput | None = None
class PythonNode(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
id: WorkflowName
type: Literal["python"]
program: PythonProgram
inputs: Mapping[str, Binding]
timeout_seconds: float = Field(gt=0, le=300)
retry: WorkflowRetry
WorkflowNode: TypeAlias = Annotated[ChatCompletionNode | PythonNode, Field(discriminator="type")]
class WorkflowBundle(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
contract_version: Literal["1"]
name: WorkflowName
timeout_seconds: float = Field(gt=0, le=1800)
input_schema: Mapping[str, JsonValue]
output: ReferenceBinding
nodes: tuple[WorkflowNode, ...] = Field(min_length=1, max_length=64)
def parsed_bundle(bundle: Mapping[str, JsonValue]) -> WorkflowBundle:
return WorkflowBundle.model_validate(bundle)
@dataclass(frozen=True, slots=True)
class WorkflowDefinition:
draft: WorkflowDraft
contributors: tuple[str, ...]
class WorkflowSnapshot(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
snapshot_id: str
workflow_id: str
team_id: str
name: str
digest: Digest
bundle: WorkflowBundle
contributors: tuple[str, ...]
published_by: str | None
created_at: datetime
class DefinitionPage(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
workflows: tuple[WorkflowDraft, ...]
next_cursor: str | None
DefinitionErrorCode: TypeAlias = Literal[
"not_found", "forbidden", "team_required", "invalid_bundle", "conflict", "stale_draft", "precondition_required"
]
@dataclass(frozen=True, slots=True)
class DefinitionError:
code: DefinitionErrorCode
message: str
param: str | None = None
class WorkflowDefinitionStore(Protocol):
async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None: ...
async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None: ...
async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage: ...
async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None: ...
async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None: ...
async def create(
self, bundle: WorkflowBundle, team_id: str, contributor: str
) -> WorkflowDraft | DefinitionError: ...
async def replace(
self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str
) -> WorkflowDraft | DefinitionError: ...
async def publish(
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
) -> WorkflowSnapshot | DefinitionError: ...
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: ...

View file

@ -0,0 +1,151 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Final, Protocol
from uuid import NAMESPACE_URL, uuid5
from liteflow import RunState # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
RunSteps as SDKRunSteps,
)
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
StepTrace as SDKStepTrace,
)
from pydantic import BaseModel, ConfigDict, Field, JsonValue
__all__: Final = ("RunState",)
class RunInput(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
deadline: datetime
draft: bool
end_user_id: str | None = None
class EngineRun(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
run_id: str
session_id: str
workflow_id: str
snapshot_id: str
team_id: str
created_by: str
status: RunState
created_at: datetime
updated_at: datetime
input: RunInput
request_digest: str
output: JsonValue = None
metadata: Mapping[str, JsonValue] | None = None
cancel_requested: bool = False
class StepTrace(SDKStepTrace):
model_config = ConfigDict(frozen=True, extra="forbid")
output: JsonValue = None
finished_at: datetime | None = None
error: str | None = None
prompt_tokens: int = 0
completion_tokens: int = 0
response_cost: float | None = None
request_id: str | None = None
call_id: str | None = None
request_ids: tuple[str, ...] = ()
class RunSteps(SDKRunSteps):
model_config = ConfigDict(frozen=True)
steps: tuple[StepTrace, ...] # pyright: ignore[reportIncompatibleVariableOverride] # both SDK and extension models are frozen
error: str | None = None
class RunView(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
run_id: str
session_id: str
workflow_id: str
snapshot_id: str
team_id: str
status: RunState
created_at: datetime
updated_at: datetime
cancel_requested: bool
class RunPage(BaseModel):
model_config = ConfigDict(frozen=True)
runs: tuple[RunView, ...]
count: int
next_cursor: str | None
@dataclass(frozen=True, slots=True)
class ChildCallResult:
output: JsonValue
prompt_tokens: int
completion_tokens: int
request_id: str | None
response_cost: float | None
@dataclass(frozen=True, slots=True)
class ExecutionError:
code: str
message: str
receipt: ChildCallResult | None = None
def workflow_call_id(run_id: str, node_id: str, attempt: int) -> str:
return str(uuid5(NAMESPACE_URL, f"litellm/workflows/{run_id}/{node_id}/{attempt}"))
@dataclass(frozen=True, slots=True)
class RunAdmission:
run: EngineRun
created: bool
class WorkflowExecutionStore(Protocol):
async def find_idempotent(
self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str
) -> EngineRun | None: ...
async def admit(
self,
*,
workflow_id: str,
snapshot_id: str,
team_id: str,
key_hash: str,
input: RunInput,
digest: str,
idempotency_key: str | None,
) -> RunAdmission | ExecutionError: ...
async def get(self, run_id: str) -> EngineRun | None: ...
async def checkpoint(self, run_id: str, step: StepTrace) -> bool: ...
async def steps(self, run_id: str) -> tuple[StepTrace, ...]: ...
async def finish(self, run_id: str, state: RunState, output: JsonValue, error: str | None) -> None: ...
async def cancel(self, run_id: str) -> None: ...
async def list_runs(
self,
workflow_id: str,
team_id: str,
key_hash: str | None,
state: RunState | None,
limit: int,
cursor: str | None,
) -> RunPage | ExecutionError: ...

View file

@ -126,6 +126,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy.workflows.runner import WorkflowRunner
from litellm.repositories.database import Database
from litellm.tracing import TraceReceiver
@ -135,6 +136,7 @@ else:
class ProxyLifespanState(TypedDict):
workflow_runner: NotRequired[ReadOnly["WorkflowRunner | None"]]
tracing_receiver: ReadOnly["TraceReceiver | None"]
database: ReadOnly["Database | None"]
@ -660,6 +662,21 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_UPDATE.value,
)
workflow_definition_routes = (
"/v1/workflows",
"/v1/workflows/{workflow_id}",
"/v1/workflows/{workflow_id}/validate",
"/v1/workflows/{workflow_id}/publish",
)
workflow_execution_routes = (
"/v1/workflows/{workflow_id}/test",
"/v1/workflows/{workflow_id}/runs",
"/v1/workflows/runs/{run_id}/steps",
"/v1/workflows/runs/{run_id}/definition",
"/v1/workflows/runs/{run_id}/cancel",
)
management_routes = (
[
# user
@ -723,6 +740,8 @@ class LiteLLMRoutes(enum.Enum):
+ mcp_management_routes
+ list(agent_management_routes)
+ list(skill_hub_routes)
+ list(workflow_definition_routes)
+ list(workflow_execution_routes)
)
spend_tracking_routes = [
@ -966,6 +985,8 @@ class LiteLLMRoutes(enum.Enum):
# inside agent_endpoints/endpoints.py
*agent_management_routes,
*skill_hub_routes,
*workflow_definition_routes,
*workflow_execution_routes,
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##

View file

@ -603,8 +603,10 @@ class RouteChecks:
if not isinstance(route, str):
return False
def _placeholder_to_regex(match: re.Match) -> str:
def _placeholder_to_regex(match: re.Match[str]) -> str:
placeholder: Final = match.group(0).strip("{}")
if placeholder == "workflow_id" and match.string.startswith("/v1/workflows/"):
return r"(?!runs(?:/|$))[^/]+"
if not placeholder.endswith(":path"):
return r"[^/]+"
# A ":path" placeholder takes whatever the router's own path

View file

@ -4849,7 +4849,12 @@ async def _run_post_custom_auth_checks(
async def authorize_internal_virtual_key(
key_hash: str, request: Request, request_data: dict[str, object]
key_hash: str,
request: Request,
request_data: dict[str, object],
*,
route: str | None = None,
expected_team_id: str | None = None,
) -> UserAPIKeyAuth:
"""Authorize a server-owned job against its persisted virtual-key assignment, never a client-supplied bearer hash."""
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
@ -4858,15 +4863,17 @@ async def authorize_internal_virtual_key(
await IdentityStore(prisma_client, user_api_key_cache, proxy_logging_obj=proxy_logging_obj).resolve(
hashed_token=key_hash
)
)
route: Final = get_request_route(request=request)
await pre_db_read_auth_checks(request_data=request_data, request=request, route=route)
).model_copy(deep=True)
if expected_team_id is not None and identity.team_id != expected_team_id:
raise HTTPException(status_code=403, detail="The execution key no longer belongs to its original team")
execution_route: Final = route if route is not None else get_request_route(request=request)
await pre_db_read_auth_checks(request_data=request_data, request=request, route=execution_route)
auth: Final = await validate_resolved_virtual_key(
request=request,
request_data=request_data,
valid_token=identity,
api_key=key_hash,
route=route,
route=execution_route,
start_time=datetime.now(timezone.utc),
parent_otel_span=None,
end_user_id=None,
@ -4878,13 +4885,13 @@ async def authorize_internal_virtual_key(
user_api_key_auth_obj=auth,
request=request,
request_data=request_data,
route=route,
route=execution_route,
api_key=key_hash,
force_virtual_key_checks=True,
)
if recovered is not None:
return recovered
_seed_request_destinations(auth, request)
auth.request_route = route
auth.request_route = execution_route
request.state.principal = _resolve_request_principal(request, auth)
return auth

View file

@ -4,6 +4,7 @@ import json
import logging
import math
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
from contextvars import Context
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
@ -29,7 +30,7 @@ import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.types import Receive, Scope, Send
from typing_extensions import Never
@ -94,6 +95,7 @@ from litellm.litellm_core_utils.streaming_handler import (
backfill_missing_cache_usage_fields,
)
from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields
from litellm.models.workflow_execution import ExecutionError
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload, with_openai_overload_status
from litellm.proxy.auth.auth_checks import (
@ -725,6 +727,8 @@ def serialize_http_exception_detail(
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True))
_WORKFLOW_FIELDS: Final = TypeAdapter(dict[str, object])
_WORKFLOW_JSON: Final = TypeAdapter[JsonValue](JsonValue)
_UNSEEN_DEPLOYMENT: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset())
@ -2190,10 +2194,63 @@ def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]:
)
async def process_workflow_request(
*,
request: Request,
workflow_model: str,
caller: UserAPIKeyAuth,
messages: JsonValue,
execute: Callable[[JsonValue], Awaitable[ModelResponse | ExecutionError]],
) -> ModelResponse | ExecutionError:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.budget_reservation import release_unbound_budget_reservation
await release_unbound_budget_reservation(caller.budget_reservation)
auth: Final = caller.model_copy(update={"budget_reservation": None})
async def process() -> ModelResponse | ExecutionError:
try:
prepared: Final = _WORKFLOW_FIELDS.validate_python(
await add_litellm_data_to_request(
data={"model": workflow_model, "messages": messages, "stream": False},
request=request,
general_settings=_WORKFLOW_FIELDS.validate_python(proxy_server.general_settings), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
user_api_key_dict=auth,
version=proxy_server.version,
proxy_config=proxy_server.proxy_config,
)
)
data: Final = _WORKFLOW_FIELDS.validate_python(
check_and_merge_model_level_guardrails(
data=prepared, llm_router=proxy_server.llm_router, trust_client_model_info=False
)
)
guarded: Final = _WORKFLOW_FIELDS.validate_python(
await proxy_server.proxy_logging_obj.pre_call_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook body is validated at this boundary
user_api_key_dict=auth, data=data, call_type="acompletion", guardrails_only=True
)
)
await proxy_server.proxy_logging_obj.during_call_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook accepts the validated body
data=guarded, user_api_key_dict=auth, call_type="acompletion", guardrails_only=True
)
result: Final = await execute(_WORKFLOW_JSON.validate_python(guarded["messages"]))
if isinstance(result, ExecutionError):
return result
response: Final = await proxy_server.proxy_logging_obj.post_call_success_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook accepts the validated body
data=guarded, user_api_key_dict=auth, response=result, guardrails_only=True
)
return response if isinstance(response, ModelResponse) else ModelResponse.model_validate(response)
except Exception: # noqa: BLE001 # workflow policy failures are values and must not expose guarded content
return ExecutionError("request_rejected", "The gateway could not complete the workflow")
return await Context().run(asyncio.create_task, process())
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict, served_model: str | None = None):
self.data = data
self.served_model = served_model
self.logging_obj: LiteLLMLoggingObj | None = None
self._tags_before_guardrails: frozenset[str] | None = None
self.route_type: str | None = None
@ -2636,6 +2693,7 @@ class ProxyBaseLLMRequestProcessing:
)
self.data["litellm_logging_obj"] = logging_obj
self.logging_obj = logging_obj
# Merge model-level guardrails before pre_call_hook so DB/UI-configured
# guardrails actually execute on pre_call. Without this, guardrails set

View file

@ -1,4 +1,4 @@
from typing import Annotated, TypeAlias
from typing import Annotated, TypeAlias, cast # noqa: TID251 # Prisma writer delegates are generated dynamically
from fastapi import Depends
@ -7,12 +7,16 @@ from litellm.models.gateway_activity import GatewayActivityStore
from litellm.models.plugin_upload import PluginUploadStore
from litellm.models.skill_submission import SkillSubmissionStore
from litellm.models.user_banner import UserBannerStore
from litellm.models.workflow_definition import WorkflowDefinitionStore
from litellm.models.workflow_execution import WorkflowExecutionStore
from litellm.proxy.deps.database import DatabaseDep
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
@ -54,3 +58,33 @@ async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmiss
SkillSubmissionRepositoryDep: TypeAlias = Annotated[
SkillSubmissionStore | None, Depends(get_skill_submission_repository)
]
async def get_workflow_definition_repository(database: DatabaseDep) -> WorkflowDefinitionStore | None:
return (
None
if database is None
else WorkflowDefinitionRepository(
cast(DefinitionDatabase, database.prisma.writer_db) # cast-ok: Prisma delegates implement the protocol
)
)
WorkflowDefinitionRepositoryDep: TypeAlias = Annotated[
WorkflowDefinitionStore | None, Depends(get_workflow_definition_repository)
]
async def get_workflow_execution_repository(database: DatabaseDep) -> WorkflowExecutionStore | None:
return (
None
if database is None
else WorkflowExecutionRepository(
cast(ExecutionDatabase, database.prisma.writer_db) # cast-ok: generated Prisma transaction delegates
)
)
WorkflowExecutionRepositoryDep: TypeAlias = Annotated[
WorkflowExecutionStore | None, Depends(get_workflow_execution_repository)
]

View file

@ -0,0 +1,223 @@
from types import MappingProxyType
from typing import Annotated, Final, Literal, NoReturn, TypeAlias, TypeVar
from fastapi import APIRouter, Depends, Header, HTTPException, Query
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict
from litellm.models.workflow_definition import (
DefinitionError,
DefinitionPage,
Digest,
WorkflowBundle,
WorkflowDefinitionStore,
WorkflowDraft,
parsed_bundle,
)
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep
from litellm.proxy.workflows.definitions import (
author_identity,
authorize_team,
authorized_draft,
check_etag,
requested_team,
validate_bundle,
)
router: Final = APIRouter(tags=["workflow management"])
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
_Team: TypeAlias = Annotated[str | None, Query(min_length=1, max_length=256)]
_ETag: TypeAlias = Annotated[str | None, Header(alias="If-Match")]
_Result: Final = TypeVar("_Result")
_STATUS: Final = MappingProxyType(
{
"not_found": 404,
"forbidden": 403,
"team_required": 400,
"invalid_bundle": 400,
"conflict": 409,
"stale_draft": 412,
"precondition_required": 412,
}
)
class WorkflowRegistration(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
bundle: WorkflowBundle
class WorkflowPublication(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
digest: Digest
class WorkflowValidationRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
class WorkflowValidationResult(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
valid: Literal[True] = True
workflow_id: str
etag: str
digest: str
class WorkflowPublicationResult(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
workflow_id: str
snapshot_id: str
model: str
digest: str
def raise_public(error: DefinitionError) -> NoReturn:
raise HTTPException(
status_code=_STATUS[error.code],
detail={
"error": {
"type": "workflow_definition_error",
"code": error.code,
"message": error.message,
"param": error.param,
}
},
)
def _result(result: _Result | DefinitionError) -> _Result:
if isinstance(result, DefinitionError):
raise_public(result)
return result
def _store(store: WorkflowDefinitionStore | None) -> WorkflowDefinitionStore:
if store is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
return store
def _author(caller: UserAPIKeyAuth) -> str:
return _result(author_identity(caller))
def _draft_response(draft: WorkflowDraft, status_code: int = 200) -> JSONResponse:
return JSONResponse(
content=draft.model_dump(mode="json", by_alias=True, exclude_unset=True),
status_code=status_code,
headers={"ETag": draft.etag},
)
@router.post("/v1/workflows", status_code=201, response_model=WorkflowDraft, response_model_exclude_unset=True)
async def register_workflow(
data: WorkflowRegistration,
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
) -> JSONResponse:
"""Create a team draft from a sequential LiteFlow v1 bundle."""
contributor: Final = _author(caller)
store: Final = _store(repository)
scope: Final = _result(requested_team(caller, team_id))
_result(await authorize_team(store, caller, scope, "write"))
_result(validate_bundle(data.bundle))
draft: Final = _result(await store.create(data.bundle, scope, contributor))
return _draft_response(draft, 201)
@router.get("/v1/workflows", response_model=DefinitionPage, response_model_exclude_unset=True)
async def list_workflow_definitions(
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
limit: Annotated[int, Query(ge=1, le=100)] = 50,
cursor: Annotated[
str | None, Query(pattern=r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")
] = None,
) -> DefinitionPage:
"""List readable team drafts in immutable workflow ID order."""
store: Final = _store(repository)
scope: Final = _result(requested_team(caller, team_id))
_result(await authorize_team(store, caller, scope, "read"))
return await store.list_drafts(scope, limit, cursor)
@router.get("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True)
async def get_workflow_definition(
workflow_id: str,
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
) -> JSONResponse:
"""Read the current draft and its publication pointer."""
draft: Final = _result(await authorized_draft(_store(repository), caller, workflow_id, team_id))
return _draft_response(draft)
@router.put("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True)
async def replace_workflow_definition(
workflow_id: str,
data: WorkflowRegistration,
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
if_match: _ETag = None,
) -> JSONResponse:
"""Replace a draft only when If-Match identifies its current revision."""
contributor: Final = _author(caller)
store: Final = _store(repository)
current: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="write"))
_result(check_etag(current, if_match))
_result(validate_bundle(data.bundle))
draft: Final = _result(await store.replace(workflow_id, current.team_id, current.etag, data.bundle, contributor))
return _draft_response(draft)
@router.post("/v1/workflows/{workflow_id}/validate", response_model=WorkflowValidationResult)
async def validate_workflow_definition(
workflow_id: str,
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
if_match: _ETag = None,
data: WorkflowValidationRequest | None = None,
) -> WorkflowValidationResult:
"""Validate a draft without model calls or Python execution."""
_author(caller)
store: Final = _store(repository)
draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="write"))
_result(check_etag(draft, if_match))
_result(validate_bundle(parsed_bundle(draft.bundle)))
return WorkflowValidationResult(workflow_id=workflow_id, etag=draft.etag, digest=draft.digest)
@router.post("/v1/workflows/{workflow_id}/publish", response_model=WorkflowPublicationResult)
async def publish_workflow_definition(
workflow_id: str,
data: WorkflowPublication,
repository: WorkflowDefinitionRepositoryDep,
caller: _Caller,
team_id: _Team = None,
if_match: _ETag = None,
) -> WorkflowPublicationResult:
"""Publish the reviewed revision as an immutable snapshot."""
publisher: Final = _author(caller)
store: Final = _store(repository)
draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="publish"))
_result(check_etag(draft, if_match))
_result(validate_bundle(parsed_bundle(draft.bundle)))
snapshot: Final = _result(await store.publish(workflow_id, draft.team_id, draft.etag, data.digest, publisher))
return WorkflowPublicationResult(
workflow_id=workflow_id,
snapshot_id=snapshot.snapshot_id,
model=f"workflow/{snapshot.name}",
digest=snapshot.digest,
)

View file

@ -0,0 +1,413 @@
import asyncio
import hashlib
import json
from collections.abc import Mapping
from datetime import timedelta
from typing import Annotated, Final, Literal, NoReturn, TypeAlias
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, Response
from fastapi.responses import JSONResponse
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError
from litellm.models.workflow_definition import DefinitionError, WorkflowDefinitionStore, WorkflowDraft, WorkflowSnapshot
from litellm.models.workflow_execution import (
EngineRun,
ExecutionError,
RunInput,
RunPage,
RunState,
RunSteps,
WorkflowExecutionStore,
)
from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep
from litellm.proxy.management.teams.authz import is_team_admin
from litellm.proxy.management_endpoints.workflow_definition_endpoints import raise_public
from litellm.proxy.workflows.definitions import (
authorize_team,
authorized_draft,
freeze_draft,
published_workflow,
requested_team,
)
from litellm.proxy.workflows.lifespan import get_workflow_runner
from litellm.proxy.workflows.monty.models import MontyFailure
from litellm.proxy.workflows.monty.runtime import validate_json
from litellm.proxy.workflows.runner import utc_now
router: Final = APIRouter(tags=["workflow management"])
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
_Idempotency: TypeAlias = Annotated[str | None, Header(alias="Idempotency-Key", min_length=1, max_length=256)]
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
class TestRunRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
authorize_billing: StrictBool
class AcceptedRun(BaseModel):
model_config = ConfigDict(frozen=True)
run_id: str
snapshot_id: str
trace_url: str
class EmptyRunAction(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
class WorkflowInvocation(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
model: str
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
stream: Literal[False] = False
user: str | None = Field(default=None, max_length=256)
metadata: Mapping[str, JsonValue] | None = None
def _error(error: ExecutionError, run_id: str | None = None) -> NoReturn:
code: Final = {
"not_found": 404,
"forbidden": 403,
"idempotency_conflict": 409,
"capacity_exceeded": 429,
"execution_failed": 422,
"database_unavailable": 503,
}.get(error.code, 400)
raise HTTPException(
status_code=code,
detail={"error": {"type": "workflow_error", "code": error.code, "message": error.message}},
headers={"x-litellm-workflow-run-id": run_id} if run_id is not None else None,
)
def _stores(
definitions: WorkflowDefinitionStore | None, runs: WorkflowExecutionStore | None
) -> tuple[WorkflowDefinitionStore, WorkflowExecutionStore]:
if definitions is None or runs is None:
_error(ExecutionError("database_unavailable", "Workflow execution requires a database"))
return definitions, runs
def _execution_key(caller: UserAPIKeyAuth, team_id: str) -> str:
if not caller.token or not caller.via_virtual_key or caller.is_session_token or caller.team_id != team_id:
_error(ExecutionError("forbidden", "Execution requires a virtual key assigned to the workflow team"))
return caller.token
def _idempotency_key(operation: str, key: str | None) -> str | None:
return None if key is None else hashlib.sha256(f"{operation}:{key}".encode()).hexdigest()
def _digest(messages: tuple[Mapping[str, JsonValue], ...], user: str | None, draft_etag: str | None) -> str:
return hashlib.sha256(
json.dumps(
{"messages": tuple(dict(message) for message in messages), "user": user, "etag": draft_etag},
sort_keys=True,
separators=(",", ":"),
allow_nan=False,
).encode()
).hexdigest()
def _validate_messages(messages: tuple[Mapping[str, JsonValue], ...]) -> None:
bounded: Final = validate_json([dict(message) for message in messages], max_bytes=1024 * 1024, max_depth=64)
if isinstance(bounded, MontyFailure):
_error(ExecutionError("invalid_input", "Workflow messages exceed the JSON size or depth limit"))
try:
_MESSAGES.validate_python(bounded, strict=True)
except ValidationError:
_error(ExecutionError("invalid_input", "Expected standard Chat Completions messages"))
async def _admit(
request: Request,
runs: WorkflowExecutionStore,
snapshot: WorkflowSnapshot,
caller: UserAPIKeyAuth,
messages: tuple[Mapping[str, JsonValue], ...],
user: str | None,
key: str | None,
digest: str,
*,
draft: bool,
) -> EngineRun:
_validate_messages(messages)
runner: Final = get_workflow_runner(request)
if runner is None:
_error(ExecutionError("database_unavailable", "Workflow execution is unavailable"))
admitted: Final = await runs.admit(
workflow_id=snapshot.workflow_id,
snapshot_id=snapshot.snapshot_id,
team_id=snapshot.team_id,
key_hash=_execution_key(caller, snapshot.team_id),
input=RunInput(
messages=messages,
end_user_id=user,
deadline=utc_now() + timedelta(seconds=snapshot.bundle.timeout_seconds),
draft=draft,
),
digest=digest,
idempotency_key=key,
)
if isinstance(admitted, ExecutionError):
_error(admitted)
if admitted.created and (failure := await runner.submit(admitted.run, request, caller)) is not None:
_error(failure, admitted.run.run_id)
return admitted.run
def _accepted(run: EngineRun) -> JSONResponse:
return JSONResponse(
status_code=202,
content=AcceptedRun(
run_id=run.run_id,
snapshot_id=run.snapshot_id,
trace_url=f"/v1/workflows/runs/{run.run_id}/steps",
).model_dump(),
headers={"x-litellm-workflow-run-id": run.run_id},
)
async def _replay_test(
runs: WorkflowExecutionStore, draft: WorkflowDraft, key_hash: str, key: str | None, digest: str
) -> JSONResponse | None:
if key is None:
return None
existing: Final = await runs.find_idempotent(draft.workflow_id, draft.team_id, key_hash, key)
if existing is None:
return None
if existing.request_digest != digest:
_error(ExecutionError("idempotency_conflict", "Idempotency key was used with different input"))
return _accepted(existing)
@router.post("/v1/workflows/{workflow_id}/test", status_code=202, response_model=AcceptedRun)
async def test_workflow(
request: Request,
workflow_id: str,
data: TestRunRequest,
caller: _Caller,
definitions: WorkflowDefinitionRepositoryDep,
executions: WorkflowExecutionRepositoryDep,
if_match: Annotated[str | None, Header(alias="If-Match")] = None,
idempotency_key: _Idempotency = None,
) -> JSONResponse:
"""Start a paid test of the exact draft revision and return its run ID immediately."""
store, runs = _stores(definitions, executions)
if not data.authorize_billing:
_error(ExecutionError("billing_authorization_required", "Set authorize_billing to true for a paid test run"))
draft: Final = await authorized_draft(store, caller, workflow_id, action="write")
if isinstance(draft, DefinitionError):
raise_public(draft)
if if_match is None:
raise_public(DefinitionError("precondition_required", "If-Match is required", "If-Match"))
key_hash: Final = _execution_key(caller, draft.team_id)
digest: Final = _digest(data.messages, caller.end_user_id, if_match)
key: Final = _idempotency_key("test", idempotency_key)
if (replay := await _replay_test(runs, draft, key_hash, key, digest)) is not None:
return replay
snapshot: Final = await freeze_draft(store, caller, workflow_id, if_match)
if isinstance(snapshot, DefinitionError):
if (
snapshot.code == "stale_draft"
and (accepted := await _replay_test(runs, draft, key_hash, key, digest)) is not None
):
return accepted
raise_public(snapshot)
return _accepted(
await _admit(request, runs, snapshot, caller, data.messages, caller.end_user_id, key, digest, draft=True)
)
async def _read_run(
store: WorkflowDefinitionStore, runs: WorkflowExecutionStore, caller: UserAPIKeyAuth, run_id: str
) -> EngineRun:
run: Final = await runs.get(run_id)
if run is None:
_error(ExecutionError("not_found", "Run not found"))
if isinstance(requested_team(caller, run.team_id), DefinitionError):
_error(ExecutionError("not_found", "Run not found"))
team: Final = await store.get_team(run.team_id)
if team is None or team.blocked:
_error(ExecutionError("not_found", "Run not found"))
owns_run: Final = bool(caller.token) and caller.token == run.created_by and caller.team_id == run.team_id
if not owns_run and not user_api_key_has_admin_view(caller) and not is_team_admin(caller, team):
_error(ExecutionError("not_found", "Run not found"))
return run
@router.get("/v1/workflows/runs/{run_id}/steps")
async def get_workflow_steps(
run_id: str,
caller: _Caller,
definitions: WorkflowDefinitionRepositoryDep,
executions: WorkflowExecutionRepositoryDep,
) -> RunSteps:
"""Read run progress, typed node inputs and outputs, and linked model-call receipts."""
store, runs = _stores(definitions, executions)
run: Final = await _read_run(store, runs, caller, run_id)
error: Final = (run.metadata or {}).get("error")
return RunSteps(
run_id=run_id,
snapshot_id=run.snapshot_id,
state=run.status,
cancel_requested=run.cancel_requested,
steps=await runs.steps(run_id),
error=error if isinstance(error, str) else None,
)
@router.get("/v1/workflows/runs/{run_id}/definition")
async def get_run_definition(
run_id: str,
caller: _Caller,
definitions: WorkflowDefinitionRepositoryDep,
executions: WorkflowExecutionRepositoryDep,
) -> WorkflowSnapshot:
"""Read the immutable workflow definition captured when this run was admitted."""
store, runs = _stores(definitions, executions)
run: Final = await _read_run(store, runs, caller, run_id)
snapshot: Final = await store.get_snapshot(run.snapshot_id, run.team_id)
if snapshot is None or snapshot.workflow_id != run.workflow_id:
_error(ExecutionError("not_found", "Pinned workflow definition not found"))
return snapshot
@router.post("/v1/workflows/runs/{run_id}/cancel", status_code=202, response_model=AcceptedRun)
async def cancel_workflow_run(
run_id: str,
caller: _Caller,
definitions: WorkflowDefinitionRepositoryDep,
executions: WorkflowExecutionRepositoryDep,
data: EmptyRunAction | None = None,
) -> JSONResponse:
"""Request cancellation and return the run ID to poll until execution stops."""
store, runs = _stores(definitions, executions)
run: Final = await _read_run(store, runs, caller, run_id)
if caller.token != run.created_by:
permission: Final = await authorize_team(store, caller, run.team_id, "write")
if isinstance(permission, DefinitionError):
raise_public(permission)
await runs.cancel(run_id)
return _accepted(run)
async def list_engine_runs(
*,
workflow_id: str,
team_id: str | None,
status: str | None,
limit: int,
after: str | None,
caller: UserAPIKeyAuth,
definitions: WorkflowDefinitionStore | None,
executions: WorkflowExecutionStore | None,
) -> RunPage:
store, runs = _stores(definitions, executions)
definition: Final = await store.get_draft(workflow_id, team_id)
if definition is None:
_error(ExecutionError("not_found", "Workflow not found"))
permission: Final = await authorize_team(store, caller, definition.draft.team_id, "invoke")
if isinstance(permission, DefinitionError):
raise_public(permission)
may_read_team: Final = user_api_key_has_admin_view(caller) or is_team_admin(caller, permission)
if not may_read_team and not caller.token:
_error(ExecutionError("forbidden", "A run reader identity is required"))
try:
state: Final = RunState(status) if status is not None else None
except ValueError:
_error(ExecutionError("invalid_status", "Unknown workflow run status"))
result: Final = await runs.list_runs(
workflow_id,
definition.draft.team_id,
None if may_read_team else caller.token,
state,
min(limit, 100),
after,
)
if isinstance(result, ExecutionError):
_error(result)
return result
@router.get("/v1/workflows/{workflow_id}/runs")
async def list_definition_runs(
workflow_id: str,
caller: _Caller,
definitions: WorkflowDefinitionRepositoryDep,
executions: WorkflowExecutionRepositoryDep,
team_id: str | None = None,
status: RunState | None = None,
limit: Annotated[int, Query(ge=1, le=100)] = 50,
after: Annotated[str | None, Query(max_length=512)] = None,
) -> RunPage:
return await list_engine_runs(
workflow_id=workflow_id,
team_id=team_id,
status=status.value if status is not None else None,
limit=limit,
after=after,
caller=caller,
definitions=definitions,
executions=executions,
)
async def _wait_for_run(runs: WorkflowExecutionStore, run_id: str) -> EngineRun:
while True:
current: EngineRun | None = await runs.get(run_id)
if current is None:
_error(ExecutionError("not_found", "Run not found"), run_id)
if current.status in {RunState.completed, RunState.failed, RunState.cancelled}:
return current
await asyncio.sleep(0.2)
async def invoke_workflow(
request: Request,
response: Response,
caller: UserAPIKeyAuth,
raw: object,
definitions: WorkflowDefinitionStore | None,
executions: WorkflowExecutionStore | None,
) -> Mapping[str, JsonValue] | JSONResponse:
store, runs = _stores(definitions, executions)
try:
data: Final = WorkflowInvocation.model_validate(raw)
except ValidationError:
_error(ExecutionError("unsupported_request", "Use model, messages, user, metadata and stream=false"))
snapshot: Final = await published_workflow(store, caller, data.model.removeprefix("workflow/"))
if isinstance(snapshot, DefinitionError):
raise_public(snapshot)
key: Final = request.headers.get("Idempotency-Key")
if key is not None and not 1 <= len(key) <= 256:
_error(ExecutionError("invalid_idempotency_key", "Idempotency-Key must contain 1 to 256 characters"))
admitted: Final = await _admit(
request,
runs,
snapshot,
caller,
data.messages,
data.user or caller.end_user_id,
_idempotency_key("invoke", key),
_digest(data.messages, data.user or caller.end_user_id, None),
draft=False,
)
response.headers["x-litellm-workflow-run-id"] = admitted.run_id # rebind-ok: FastAPI response header
if request.headers.get("Prefer") == "respond-async":
return _accepted(admitted)
completed: Final = await _wait_for_run(runs, admitted.run_id)
if completed.status is not RunState.completed:
_error(
ExecutionError("execution_failed", "Workflow execution did not complete. Inspect the run trace"),
admitted.run_id,
)
return TypeAdapter(Mapping[str, JsonValue]).validate_python(completed.output)

View file

@ -18,6 +18,7 @@ from collections.abc import Mapping, Sequence
from typing import Final, Literal, Protocol, TypedDict
from fastapi import APIRouter, Depends, HTTPException, Query
from typing_extensions import ReadOnly
try:
from prisma.errors import UniqueViolationError
@ -31,6 +32,7 @@ from litellm.proxy._types import (
user_api_key_has_admin_view,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep
from litellm.repositories.table_repositories import (
WorkflowEventRepository,
WorkflowMessageRepository,
@ -120,6 +122,7 @@ class _RunCreateData(TypedDict, total=False):
class _RunWhere(TypedDict, total=False):
engine_owned: ReadOnly[bool]
workflow_type: str
status: str | Mapping[str, Sequence[str]]
created_by: str
@ -201,6 +204,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 +215,12 @@ 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 getattr(run, "engine_owned", False) is True:
if mutation:
raise HTTPException(
status_code=403, detail="Gateway executions cannot be changed through tracking endpoints"
)
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
return run
@ -263,17 +274,36 @@ async def list_workflow_runs(
status: str | None = Query(None),
limit: int = Query(50, ge=1, le=250),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
workflow_id: str | None = None,
team_id: str | None = None,
after: str | None = None,
definitions: WorkflowDefinitionRepositoryDep = None,
executions: WorkflowExecutionRepositoryDep = None,
):
"""List workflow runs. Filter by workflow_type and/or status.
Non-admin callers only see runs created by their own API key.
"""
if workflow_id is not None:
from litellm.proxy.management_endpoints.workflow_execution_endpoints import list_engine_runs
return await list_engine_runs(
workflow_id=workflow_id,
team_id=team_id,
status=status,
limit=limit,
after=after,
caller=user_api_key_dict,
definitions=definitions,
executions=executions,
)
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
where: Final[_RunWhere] = {}
where: Final[_RunWhere] = {"engine_owned": False}
if workflow_type:
where["workflow_type"] = workflow_type
if status:
@ -318,7 +348,7 @@ async def get_workflow_run(
where={"run_id": run_id},
include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}},
)
if run is None:
if run is None or getattr(run, "engine_owned", False) is True:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
if not user_api_key_has_admin_view(user_api_key_dict):
caller: Final = _caller_key(user_api_key_dict)
@ -360,7 +390,7 @@ async def update_workflow_run(
raise HTTPException(status_code=400, detail="No fields to update")
# Enforce ownership before writing.
await _require_run(prisma_client, run_id, user_api_key_dict)
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
try:
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update(
@ -398,7 +428,7 @@ async def append_workflow_event(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
await _require_run(prisma_client, run_id, user_api_key_dict)
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
new_status: Final = _EVENT_STATUS_MAP.get(data.event_type)
@ -487,7 +517,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:

View file

@ -847,6 +847,12 @@ from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router
from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import (
router as user_agent_analytics_router,
)
from litellm.proxy.management_endpoints.workflow_definition_endpoints import (
router as workflow_definition_router,
)
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
router as workflow_execution_router,
)
from litellm.proxy.management_endpoints.workflow_management_endpoints import (
router as workflow_management_router,
)
@ -1973,9 +1979,12 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
async with AsyncExitStack() as admin_mcp_stack:
try:
await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app))
from litellm.proxy.workflows.lifespan import workflow_lifespan
workflow_runner: Final = await admin_mcp_stack.enter_async_context(workflow_lifespan(database))
gc.freeze()
verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count())
yield state
yield {**state, "workflow_runner": workflow_runner}
finally:
if model_info_scheduler is not None and model_info_scheduler.running:
model_info_scheduler.remove_job("refresh_model_info")
@ -13156,6 +13165,26 @@ async def chat_completion(
global general_settings, user_debug, proxy_logging_obj, llm_model_list
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
data: Final = await read_request_body(request=request)
workflow_data: Final = as_str_mapping(_UNTYPED_VALUE.validate_python(data))
workflow_model: Final = workflow_data.get("model") if workflow_data is not None else None
if isinstance(workflow_model, str) and workflow_model.startswith("workflow/"):
from litellm.proxy.deps.database import get_database
from litellm.proxy.deps.repositories import (
get_workflow_definition_repository,
get_workflow_execution_repository,
)
from litellm.proxy.management_endpoints.workflow_execution_endpoints import invoke_workflow
workflow_database: Final = await get_database(request)
await release_unbound_budget_reservation(user_api_key_dict.budget_reservation)
return await invoke_workflow(
request,
fastapi_response,
user_api_key_dict,
workflow_data,
await get_workflow_definition_repository(workflow_database),
await get_workflow_execution_repository(workflow_database),
)
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 +22461,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)

View file

@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
// One instance of work being done. `status` is a materialized cache of the
// latest event; the event log is the authoritative source of truth.
model LiteLLM_WorkflowDefinition {
workflow_id String @id @default(uuid())
team_id String
name String
etag String
digest String
bundle_json String
contributors String[]
published_snapshot_id String?
published_etag String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
snapshots LiteLLM_WorkflowSnapshot[]
@@unique([team_id, name])
@@index([team_id, workflow_id])
}
model LiteLLM_WorkflowSnapshot {
snapshot_id String @id @default(uuid())
workflow_id String
team_id String
name String
digest String
bundle_json String
contributors String[]
published_by String?
created_at DateTime @default(now())
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
@@index([workflow_id, created_at, snapshot_id])
@@index([team_id, snapshot_id])
}
model LiteLLM_WorkflowRun {
run_id String @id @default(uuid())
session_id String @unique @default(uuid())
workflow_type String
workflow_id String?
snapshot_id String?
team_id String?
engine_owned Boolean @default(false)
request_digest String?
idempotency_key String?
cancel_requested Boolean @default(false)
status String @default("pending")
created_by String? // user_id of the key that created this run; null = created by master key
created_at DateTime @default(now())
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
@@index([session_id])
@@index([created_at])
@@index([created_by])
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
@@index([workflow_id, team_id, created_at, run_id])
@@index([team_id, engine_owned, status])
}
// Append-only log of state transitions. Never mutate rows here.

View file

@ -458,7 +458,7 @@ async def release_or_invalidate_budget_reservation(
budget_reservation["finalized"] = True
async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None:
async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object] | None) -> None:
"""Release a reservation no logging callback took ownership of, once the request ended.
A handler whose litellm call never builds a logging object (batch cancel, file

View file

@ -4367,6 +4367,8 @@ class ProxyLogging:
data: dict,
response: LLMResponseTypes,
user_api_key_dict: UserAPIKeyAuth,
*,
guardrails_only: bool = False,
):
"""
Allow user to modify outgoing data
@ -4462,6 +4464,8 @@ class ProxyLogging:
#################################################################
for callback in other_callbacks:
if guardrails_only and not callback.enforces_request_content:
continue
callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, data=data, response=response
)

View file

@ -0,0 +1,91 @@
# LiteFlow gateway workflows
The gateway accepts LiteFlow compiled contract version `1` and executes sequential `chat_completion` and `python` nodes. Python nodes use the Monty adapter. Conditional, parallel and streaming execution are not supported
Workflow definitions belong to a team. Authors register and edit drafts. A team administrator who did not contribute to the draft publishes an immutable snapshot. Invocation requires a virtual key assigned to that team, with access to the workflow model and each child model
## SDK example
The proxy imports the actual LiteFlow SDK models and digest helper, pinned to SDK revision `27bcbce4`. Its compiled JSON is validated at the server boundary because the SDK does not export a server bundle parser or executor. No SDK implementation is copied
This contract matches LiteFlow PRs [3](https://github.com/BerriAI/liteflow/pull/3) and [5](https://github.com/BerriAI/liteflow/pull/5). Configure `gateway-model` on the proxy and provide two different team members' keys
```python
import os
import time
from liteflow import ChatCompletion, GatewayClient, Input, Workflow
workflow = Workflow(
name="welcome",
nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)],
)
messages = [{"role": "user", "content": "Write one short welcome sentence."}]
with GatewayClient(
base_url=os.environ["PROXY_BASE_URL"],
api_key=os.environ["WORKFLOW_AUTHOR_KEY"],
team_id=os.environ["WORKFLOW_TEAM_ID"],
) as author, GatewayClient(
base_url=os.environ["PROXY_BASE_URL"],
api_key=os.environ["WORKFLOW_PUBLISHER_KEY"],
team_id=os.environ["WORKFLOW_TEAM_ID"],
) as publisher:
draft = author.register(workflow)
author.validate(draft)
publisher.publish(draft)
accepted = author.test(draft, messages=messages, authorize_billing=True)
for _ in range(300):
trace = author.inspect(accepted["run_id"])
print(trace)
if trace.state.value in ("completed", "failed", "cancelled"):
break
time.sleep(1)
result = author.run("welcome", messages=messages)
print(result.run_id, result.output)
```
## API contract
| Operation | Endpoint | Contract |
| --- | --- | --- |
| Register | `POST /v1/workflows` | `{"bundle": <compiled bundle>}`, optional `team_id` query, returns `201` |
| List definitions | `GET /v1/workflows` | `team_id`, `limit`, `cursor`; returns `workflows`, `next_cursor` |
| Retrieve draft | `GET /v1/workflows/{workflow_id}` | Returns the draft and `ETag` header |
| Replace draft | `PUT /v1/workflows/{workflow_id}` | Compiled bundle and exact `If-Match` |
| Validate | `POST /v1/workflows/{workflow_id}/validate` | Exact `If-Match`, no body required |
| Publish | `POST /v1/workflows/{workflow_id}/publish` | Exact `If-Match` and `{"digest": "<draft digest>"}` |
| Test draft | `POST /v1/workflows/{workflow_id}/test` | Exact `If-Match`, `messages`, `authorize_billing: true`; returns `202` |
| Invoke published | `POST /v1/chat/completions` | `model: "workflow/<name>"`, `messages`, `stream: false` |
| List runs | `GET /v1/workflows/{workflow_id}/runs` | `limit`, `status`, `after`; returns `runs`, `count`, `next_cursor` |
| Inspect | `GET /v1/workflows/runs/{run_id}/steps` | Run state and typed step traces |
| Pinned definition | `GET /v1/workflows/runs/{run_id}/definition` | Exact snapshot used by that run |
| Cancel | `POST /v1/workflows/runs/{run_id}/cancel` | Returns `202`; poll until terminal state |
Draft responses contain `name`, `workflow_id`, `team_id`, `etag`, `digest`, `bundle` and `published_snapshot_id`. Missing or stale `If-Match` returns `412`. Validation checks the compiled graph and Python program without executing model or Python nodes
Test runs and chat invocations accept `Idempotency-Key`. Retrying the same operation with the same key and input returns the original run. Reusing that key with different input returns `409`. A test run pins the exact draft revision even if the draft changes later
For live UI progress, send `Prefer: respond-async` on a published chat invocation. The gateway returns `202` with `run_id`, `snapshot_id`, `trace_url` and the `x-litellm-workflow-run-id` header. The ordinary SDK call waits for a Chat Completions response carrying the same run-ID header. It waits until execution finishes or the workflow deadline expires; clients with shorter HTTP timeouts should use asynchronous invocation and an idempotency key
## Trace and spend data
Draw the pinned bundle's `nodes` and resolve each binding's `from` reference. Match `step.node_id` to `node.id`. Run and node states are `pending`, `running`, `completed`, `failed` and `cancelled`
Each step exposes typed `input` and `output`, `attempts`, `started_at`, `finished_at`, `deadline`, sanitized `error`, token counts and `response_cost`. A run-level `error` explains failures before any node starts
Model steps carry a deterministic `call_id` before dispatch. `request_id` and `request_ids` link to actual `LiteLLM_SpendLogs.request_id` values. Spend writes are asynchronous, so receipt IDs can be absent on the first read and appear on later polls. The trace reader sums persisted attempts once and does not duplicate earlier checkpoints
Each child uses the standard gateway authorization, hooks, guardrails, limits and billing path. The workflow wrapper emits no additional bill. Python nodes have no model spend receipt. Known costs remain visible when output validation fails
Definitions and snapshots use separate tables. Execution reuses `LiteLLM_WorkflowRun` and `LiteLLM_WorkflowEvent`. Engine-owned runs are exposed only through the scoped execution endpoints; existing tracking endpoints continue to serve externally tracked runs
## Execution limits and interruption
The initial runner allows four active runs per worker and 64 pending or running runs per team. The request that creates the durable run owns its execution task; idempotent retries observe that same run. Workers drain for up to 30 seconds on graceful shutdown, then fail unfinished work. After a crash, unfinished runs fail at their original deadline when read or during subsequent admission. Stored traces survive; executions are never replayed automatically
Chat retries are limited to rate limits confirmed before provider dispatch. Provider failures, timeouts and attempts with receipts are terminal because replaying an uncertain request could charge twice
Cancellation is cooperative. It stops downstream nodes but cannot undo a model request already sent. Trace receipts can still arrive after cancellation
Ordinary virtual keys can read their own runs. Team administrators and proxy administrators can inspect team runs within their credential scope. Definition access does not grant another caller access to private run inputs or outputs

View file

@ -0,0 +1,218 @@
from collections.abc import Mapping
from typing import Final, Literal
from jsonschema import Draft202012Validator
from jsonschema.exceptions import SchemaError
from openai.types.chat import ChatCompletionMessageParam
from pydantic import JsonValue, TypeAdapter, ValidationError
from referencing.jsonschema import DRAFT202012
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.workflow_definition import (
Binding,
ChatCompletionNode,
DefinitionError,
PythonNode,
ValueBinding,
WorkflowBundle,
WorkflowDefinitionStore,
WorkflowDraft,
WorkflowSnapshot,
parsed_bundle,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view
from litellm.proxy.management.teams.authz import is_team_admin, is_team_member
from litellm.proxy.workflows.monty.compiler import compile_program
from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
_SCHEMA: Final = TypeAdapter[Mapping[str, JsonValue] | bool](Mapping[str, JsonValue] | bool)
_LIMITS: Final = MontyLimits()
_WRITER_ROLES: Final = frozenset(
{LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.ORG_ADMIN, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM}
)
def requested_team(caller: UserAPIKeyAuth, team_id: str | None) -> str | DefinitionError:
selected: Final = team_id if team_id is not None else caller.team_id
if not selected:
return DefinitionError("team_required", "Specify team_id for this workflow", "team_id")
if caller.team_id and caller.team_id != selected and not caller.is_session_token:
return DefinitionError("not_found", "Workflow not found")
return selected
def author_identity(caller: UserAPIKeyAuth) -> str | DefinitionError:
if not caller.user_id or caller.user_role not in _WRITER_ROLES:
return DefinitionError("forbidden", "Workflow authoring requires a team member identity")
return caller.user_id
async def authorize_team(
store: WorkflowDefinitionStore,
caller: UserAPIKeyAuth,
team_id: str,
action: Literal["read", "write", "publish", "invoke"],
) -> LiteLLM_TeamTable | DefinitionError:
scope: Final = requested_team(caller, team_id)
if isinstance(scope, DefinitionError):
return scope
team: Final = await store.get_team(scope)
if team is None or team.blocked:
return DefinitionError("not_found", "Workflow not found")
if action == "invoke" and (caller.team_id == scope or user_api_key_has_admin_view(caller)):
return team
if not user_api_key_has_admin_view(caller) and not is_team_member(caller, team):
return DefinitionError("not_found", "Workflow not found")
if action == "read":
return team
if isinstance(identity := author_identity(caller), DefinitionError):
return identity
if action == "publish" and caller.user_role != LitellmUserRoles.PROXY_ADMIN and not is_team_admin(caller, team):
return DefinitionError("forbidden", "Publishing requires a team administrator")
return team
async def authorized_draft(
store: WorkflowDefinitionStore,
caller: UserAPIKeyAuth,
workflow_id: str,
team_id: str | None = None,
*,
action: Literal["read", "write", "publish"] = "read",
) -> WorkflowDraft | DefinitionError:
key_team: Final = None if caller.is_session_token else caller.team_id
definition: Final = await store.get_draft(workflow_id, team_id if team_id is not None else key_team)
if definition is None:
return DefinitionError("not_found", "Workflow not found")
access: Final = await authorize_team(store, caller, definition.draft.team_id, action)
return access if isinstance(access, DefinitionError) else definition.draft
async def published_workflow(
store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, name: str, team_id: str | None = None
) -> WorkflowSnapshot | DefinitionError:
scope: Final = requested_team(caller, team_id)
if isinstance(scope, DefinitionError):
return scope
access: Final = await authorize_team(store, caller, scope, "invoke")
if isinstance(access, DefinitionError):
return access
snapshot: Final = await store.get_published(name, scope)
return snapshot if snapshot is not None else DefinitionError("not_found", "Published workflow not found")
async def snapshot_by_id(
store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, snapshot_id: str
) -> WorkflowSnapshot | DefinitionError:
key_team: Final = None if caller.is_session_token else caller.team_id
snapshot: Final = await store.get_snapshot(snapshot_id, key_team)
if snapshot is None:
return DefinitionError("not_found", "Workflow snapshot not found")
access: Final = await authorize_team(store, caller, snapshot.team_id, "read")
return access if isinstance(access, DefinitionError) else snapshot
def check_etag(draft: WorkflowDraft, etag: str | None) -> DefinitionError | None:
if etag is None:
return DefinitionError("precondition_required", "If-Match is required", "If-Match")
if draft.etag != etag:
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
return None
async def freeze_draft(
store: WorkflowDefinitionStore,
caller: UserAPIKeyAuth,
workflow_id: str,
etag: str | None,
team_id: str | None = None,
) -> WorkflowSnapshot | DefinitionError:
draft: Final = await authorized_draft(store, caller, workflow_id, team_id, action="write")
if isinstance(draft, DefinitionError):
return draft
if (failure := check_etag(draft, etag)) is not None:
return failure
if (invalid := validate_bundle(parsed_bundle(draft.bundle))) is not None:
return invalid
return await store.freeze_draft(workflow_id, draft.team_id, draft.etag)
def _local_references(value: Mapping[str, JsonValue] | bool, depth: int = 0) -> bool:
if depth > 64:
return False
if isinstance(value, Mapping):
if any(keyword in value for keyword in ("$id", "$dynamicRef", "$recursiveRef", "$dynamicAnchor")):
return False
reference: Final = value.get("$ref")
if reference is not None and (not isinstance(reference, str) or not reference.startswith("#/")):
return False
return all(
_local_references(_SCHEMA.validate_python(child), depth + 1) for child in DRAFT202012.subresources_of(value)
)
return True
def _schema_error(schema: Mapping[str, JsonValue], param: str) -> DefinitionError | None:
try:
Draft202012Validator.check_schema(schema)
if not _local_references(schema):
return DefinitionError("invalid_bundle", "Schemas must use local references and bounded nesting", param)
except (SchemaError, RecursionError):
return DefinitionError("invalid_bundle", "Invalid JSON schema", param)
return None
def _binding_error(binding: Binding, previous: frozenset[str], param: str) -> DefinitionError | None:
if isinstance(binding, ValueBinding):
return None
if binding.source == "$input" and binding.output == "messages":
return None
if binding.source in previous and binding.output == "result":
return None
return DefinitionError("invalid_bundle", "Bindings must reference input messages or an earlier node result", param)
def _node_error(node: ChatCompletionNode | PythonNode, previous: frozenset[str]) -> DefinitionError | None:
if node.id in previous or node.id == "runs":
return DefinitionError("invalid_bundle", "Node IDs must be unique and cannot be reserved", f"nodes.{node.id}")
if isinstance(node, ChatCompletionNode):
if node.model.startswith("workflow/") or "://" in node.model:
return DefinitionError("invalid_bundle", "Use a configured gateway model group", f"nodes.{node.id}.model")
if (failure := _binding_error(node.messages, previous, f"nodes.{node.id}.messages")) is not None:
return failure
if isinstance(node.messages, ValueBinding):
try:
_MESSAGES.validate_python(node.messages.value, strict=True)
except ValidationError:
return DefinitionError("invalid_bundle", "Invalid chat messages", f"nodes.{node.id}.messages")
if node.response_format is not None:
return _schema_error(node.response_format.json_schema.schema_, f"nodes.{node.id}.response_format")
return None
if set(node.inputs) != set(node.program.parameter_names):
return DefinitionError(
"invalid_bundle", "Bind exactly the Python program parameters", f"nodes.{node.id}.inputs"
)
for name, binding in node.inputs.items():
if (failure := _binding_error(binding, previous, f"nodes.{node.id}.inputs.{name}")) is not None:
return failure
compiled: Final = compile_program(node.program, _LIMITS)
if isinstance(compiled, MontyFailure):
return DefinitionError("invalid_bundle", compiled.message[:1024], f"nodes.{node.id}.program")
return None
def validate_bundle(bundle: WorkflowBundle) -> DefinitionError | None:
if bundle.name == "runs":
return DefinitionError("invalid_bundle", "The workflow name is reserved", "bundle.name")
if len(bundle.model_dump_json(by_alias=True).encode()) > 1024 * 1024:
return DefinitionError("invalid_bundle", "Workflow bundles cannot exceed 1 MiB", "bundle")
if (failure := _schema_error(bundle.input_schema, "bundle.input_schema")) is not None:
return failure
for index, node in enumerate(bundle.nodes):
previous: frozenset[str] = frozenset(prior.id for prior in bundle.nodes[:index])
if (node_failure := _node_error(node, previous)) is not None:
return node_failure
if bundle.output.source != bundle.nodes[-1].id or bundle.output.output != "result":
return DefinitionError("invalid_bundle", "The workflow output must be the final node result", "bundle.output")
return None

View file

@ -0,0 +1,287 @@
import asyncio
import json
import math
import re
import sys
import time
from collections.abc import Mapping
from contextvars import Context
from dataclasses import replace
from datetime import datetime, timezone
from typing import Annotated, Final
from fastapi import Request, Response
from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from pydantic_core import from_json
from starlette.types import Message
from litellm.exceptions import BudgetExceededError
from litellm.models.workflow_definition import ChatCompletionNode
from litellm.models.workflow_execution import ChildCallResult, EngineRun, ExecutionError, StepTrace
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import authorize_internal_virtual_key
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import (
safe_set_request_parsed_body, # pyright: ignore[reportUnknownVariableType] # legacy parsed-body API
)
from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers, redact_credential_headers
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation_on_cancel, # pyright: ignore[reportUnknownVariableType] # legacy reservation dictionary
release_unbound_budget_reservation,
)
from litellm.proxy.workflows.monty.models import MontyErrorKind
from litellm.proxy.workflows.monty.validation import ValidatorProcess
from litellm.types.utils import Usage
if sys.version_info >= (3, 11):
from asyncio import timeout
else:
from async_timeout import timeout
_ROUTE: Final = "/v1/chat/completions"
_MESSAGES: Final = TypeAdapter[list[ChatCompletionMessageParam]](
Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
)
_JSON: Final = TypeAdapter[JsonValue](JsonValue, config=ConfigDict(allow_inf_nan=False))
_FIELDS: Final = TypeAdapter(dict[str, object])
_STRING: Final = TypeAdapter[str | None](str | None)
_NUMBER: Final = TypeAdapter[float | None](float | None)
def _child_request(
parent: Request,
data: dict[str, object], # mutable-ok: shared request body mutated by the existing proxy pipeline
call_id: str,
session_id: str,
) -> Request:
encoded: Final = json.dumps(data, allow_nan=False).encode()
incoming: Final[asyncio.Queue[Message]] = asyncio.Queue(maxsize=1)
incoming.put_nowait({"type": "http.request", "body": encoded, "more_body": False})
async def receive() -> Message:
return await incoming.get()
session_headers: Final = (
()
if get_chain_id_from_headers(dict(parent.headers)) or parent.headers.get("traceparent")
else ((b"x-litellm-session-id", session_id.encode()),)
)
original_headers: Final = dict(parent.headers)
headers: Final = tuple(
(key.encode("latin-1"), value.encode("latin-1"))
for key, value in redact_credential_headers(original_headers).items()
if key.lower() not in {"content-type", "content-length", "x-litellm-call-id"} and value == original_headers[key]
)
shared_state: Final = _FIELDS.validate_python(parent.scope.get("state", {}))
request: Final = Request(
{
**parent.scope,
"method": "POST",
"path": _ROUTE,
"raw_path": _ROUTE.encode(),
"query_string": b"",
"path_params": {},
"state": {key: value for key, value in shared_state.items() if key == "database"},
"headers": (
*headers,
*session_headers,
(b"content-type", b"application/json"),
(b"content-length", str(len(encoded)).encode()),
(b"x-litellm-call-id", call_id.encode()),
),
},
receive,
)
safe_set_request_parsed_body(request, data)
return request
def _execution_error(exc: Exception, receipt: ChildCallResult | None = None) -> ExecutionError:
if isinstance(exc, BudgetExceededError) or (
isinstance(exc, ProxyException) and exc.type == ProxyErrorTypes.budget_exceeded
):
return ExecutionError("budget_exceeded", "The execution identity has insufficient budget", receipt)
status: Final = str(exc.code if isinstance(exc, ProxyException) else getattr(exc, "status_code", None))
if status in {"401", "403"}:
return ExecutionError("authorization_failed", "The execution identity is not authorized for this node", receipt)
if status == "429":
return ExecutionError("rate_limit", "The gateway could not admit this node", receipt)
if status in {"400", "422"}:
return ExecutionError("request_rejected", "The gateway rejected this node", receipt)
if status == "408" or status.startswith("5"):
return ExecutionError("provider_error", "The provider could not complete this node", receipt)
return ExecutionError("runtime_error", "The gateway could not complete this node", receipt)
async def _authorize_child(
request: Request,
run: EngineRun,
data: dict[str, object], # mutable-ok: authentication annotates the shared proxy request body
) -> UserAPIKeyAuth | ExecutionError:
try:
return await authorize_internal_virtual_key(
key_hash=run.created_by, request=request, request_data=data, route=_ROUTE, expected_team_id=run.team_id
)
except Exception as exc: # noqa: BLE001 # authentication failures are known to precede dispatch
failure: Final = _execution_error(exc)
return replace(failure, code="admission_rate_limit") if failure.code == "rate_limit" else failure
def _receipt(
processor: ProxyBaseLLMRequestProcessing, response: Response, raw: object = None
) -> ChildCallResult | None:
logging: Final = processor.logging_obj
prepared: Final = logging.prepared_cost if logging is not None else None
finalized: Final = logging.finalized_cost if logging is not None else None
fields: Final = _FIELDS.validate_python(raw.model_dump() if isinstance(raw, BaseModel) else raw or {})
reported: Final = fields.get("usage")
usage: Final = (
prepared.usage if prepared is not None else Usage.model_validate(reported) if reported is not None else None
)
if usage is None or any(
type(value) is not int or value < 0 for value in (usage.prompt_tokens, usage.completion_tokens)
):
return None
try:
cost: Final = finalized.amount if finalized is not None else float(response.headers["x-litellm-response-cost"])
except (KeyError, ValueError):
return ChildCallResult(None, usage.prompt_tokens, usage.completion_tokens, None, None)
return ChildCallResult(
None,
usage.prompt_tokens,
usage.completion_tokens,
None,
cost if math.isfinite(cost) and cost >= 0 else None,
)
async def _output(raw: object, node: ChatCompletionNode, deadline: float) -> JsonValue | ExecutionError:
try:
result: Final = ChatCompletion.model_validate(raw.model_dump() if isinstance(raw, BaseModel) else raw)
except ValueError:
return ExecutionError("invalid_response", "The gateway returned an invalid chat response")
if len(result.choices) != 1:
return ExecutionError("invalid_response", "The gateway returned an invalid chat response")
choice: Final = result.choices[0]
if choice.message.refusal or choice.finish_reason == "content_filter":
return ExecutionError("content_rejected", "The gateway rejected the node output")
content: Final = choice.message.content
if content is None:
return ExecutionError("invalid_response", "The model returned no assistant content")
if len(content.encode()) > 1024 * 1024:
return ExecutionError("output_too_large", "The model output exceeded the workflow limit")
if node.response_format is None:
return content
try:
output: Final = _JSON.validate_python(from_json(content, allow_inf_nan=False))
except ValueError:
return ExecutionError("output_invalid", "The model did not return valid JSON")
validator: Final = ValidatorProcess()
try:
failure: Final = await validator.validate(
output, dict(node.response_format.json_schema.schema_), deadline=deadline, output=True
)
if failure is None:
return output
if failure.kind is MontyErrorKind.TIMEOUT:
return ExecutionError("timeout", "The node deadline expired")
return ExecutionError("output_invalid", "The model output did not satisfy the node schema")
finally:
await validator.close()
async def execute_chat_node(
request: Request, run: EngineRun, node: ChatCompletionNode, step: StepTrace
) -> ChildCallResult | ExecutionError:
from litellm.proxy import proxy_server
if step.deadline.tzinfo is None or step.deadline.utcoffset() is None:
return ExecutionError("invalid_deadline", "The node deadline must include a timezone")
if re.fullmatch(r"[0-9a-f]{64}", run.created_by) is None or not run.team_id or not step.call_id:
return ExecutionError("authorization_failed", "A persisted team virtual key and child call ID are required")
remaining: Final = min(node.timeout_seconds, (step.deadline - datetime.now(timezone.utc)).total_seconds())
if remaining <= 0:
return ExecutionError("timeout", "The node deadline expired")
deadline: Final = time.monotonic() + remaining
try:
messages: Final = _MESSAGES.validate_python(step.input, strict=True)
trace: Final = {
"workflow_id": run.workflow_id,
"run_id": run.run_id,
"node_id": node.id,
"attempt": step.attempts,
"session_id": run.session_id,
}
data: Final[dict[str, object]] = { # mutable-ok: proxy processor and parsed-body cache share this object
**node.model_dump(
mode="json",
by_alias=True,
exclude_none=True,
exclude={"id", "type", "messages", "timeout_seconds", "retry"},
),
"messages": messages,
"stream": False,
"num_retries": 0,
"timeout": remaining,
"metadata": {**trace, "spend_logs_metadata": trace},
**({"user": run.input.end_user_id} if run.input.end_user_id is not None else {}),
}
child: Final = _child_request(request, data, step.call_id, run.session_id)
except (TypeError, ValueError):
return ExecutionError("invalid_messages", "The node input must contain valid chat messages")
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
response: Final = Response()
async def execute() -> ChildCallResult | ExecutionError:
try:
async with timeout(remaining):
auth: Final = await _authorize_child(child, run, data)
if isinstance(auth, ExecutionError):
return auth
try:
try:
raw: Final = await processor.base_process_llm_request(
request=child,
fastapi_response=response,
user_api_key_dict=auth,
route_type="acompletion",
proxy_logging_obj=proxy_server.proxy_logging_obj,
general_settings=_FIELDS.validate_python(proxy_server.general_settings), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
proxy_config=proxy_server.proxy_config,
llm_router=proxy_server.llm_router,
user_model=_STRING.validate_python(proxy_server.user_model), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
user_temperature=_NUMBER.validate_python(proxy_server.user_temperature), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
user_request_timeout=proxy_server.user_request_timeout,
user_max_tokens=proxy_server.user_max_tokens,
user_api_base=_STRING.validate_python(proxy_server.user_api_base), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
version=proxy_server.version,
)
except Exception as exc: # noqa: BLE001 # provider and hook failures must run the proxy's accounting
try:
await processor.handle_llm_api_exception(
e=exc, user_api_key_dict=auth, proxy_logging_obj=proxy_server.proxy_logging_obj
)
except Exception as failure: # noqa: BLE001 # failures become sanitized workflow values
return _execution_error(failure, _receipt(processor, response))
except asyncio.CancelledError:
await release_budget_reservation_on_cancel(auth.budget_reservation)
await proxy_server.proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth)
raise
receipt: Final = _receipt(processor, response, raw)
output: Final = await _output(raw, node, deadline)
if isinstance(output, ExecutionError):
return replace(output, receipt=receipt)
if receipt is None:
return ExecutionError("invalid_response", "The gateway returned no valid usage")
return replace(receipt, output=output)
except asyncio.TimeoutError:
return ExecutionError("timeout", "The node deadline expired", _receipt(processor, response))
except Exception as exc: # noqa: BLE001 # auth and validator failures must not expose credentials or content
return _execution_error(exc, _receipt(processor, response))
finally:
reservation: Final[object] = getattr(child.state, "budget_reservation", None)
if isinstance(reservation, Mapping):
await release_unbound_budget_reservation(reservation) # pyright: ignore[reportUnknownArgumentType] # preserve callback-shared reservation identity
return await Context().run(asyncio.create_task, execute())

View file

@ -0,0 +1,31 @@
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Final
from fastapi import Request
from litellm.proxy.deps.repositories import get_workflow_definition_repository, get_workflow_execution_repository
from litellm.proxy.workflows.gateway import execute_chat_node
from litellm.proxy.workflows.runner import WorkflowRunner
from litellm.repositories.database import Database
@asynccontextmanager
async def workflow_lifespan(database: Database | None) -> AsyncGenerator[WorkflowRunner | None]:
definitions: Final = await get_workflow_definition_repository(database)
executions: Final = await get_workflow_execution_repository(database)
runner: Final = (
WorkflowRunner(executions, definitions, execute_chat_node)
if definitions is not None and executions is not None
else None
)
try:
yield runner
finally:
if runner is not None:
await runner.close()
def get_workflow_runner(request: Request) -> WorkflowRunner | None:
runner: Final[object] = getattr(request.state, "workflow_runner", None)
return runner if isinstance(runner, WorkflowRunner) else None

View file

@ -0,0 +1,9 @@
import sys
from typing import Final
__all__: Final = ("timeout",)
if sys.version_info >= (3, 11):
from asyncio import timeout
else:
from async_timeout import timeout

View file

@ -1,6 +1,6 @@
import keyword
from dataclasses import dataclass
from enum import StrEnum
from enum import Enum
from typing import Annotated, Literal, TypeAlias
from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, StringConstraints
@ -29,7 +29,7 @@ class PythonProgram(BaseModel):
source_sha256: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")]
class MontyErrorKind(StrEnum):
class MontyErrorKind(str, Enum):
INVALID_PROGRAM = "INVALID_PROGRAM"
UNSUPPORTED_SCHEMA = "UNSUPPORTED_SCHEMA"
INPUT_INVALID = "INPUT_INVALID"
@ -47,6 +47,9 @@ class MontyErrorKind(StrEnum):
CANCELLED = "CANCELLED"
RUNTIME_UNAVAILABLE = "RUNTIME_UNAVAILABLE"
def __str__(self) -> str:
return self.value
@dataclass(frozen=True, slots=True)
class MontyFailure:

View file

@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Final
from pydantic import JsonValue, TypeAdapter
from pydantic import ValidationError as PydanticValidationError
from ._compat import timeout
from .compiler import escape_diagnostics
from .models import CompiledProgram, JsonSchema, MontyErrorKind, MontyFailure, MontyLimits, MontySuccess
from .validation import ValidatorProcess
@ -147,7 +148,9 @@ class MontyRuntime:
self._pool: AsyncMonty | None = None
self._slots = asyncio.Semaphore(limits.max_concurrent_workers)
self._admitted = 0
self._validators = tuple(ValidatorProcess() for _ in range(limits.max_concurrent_workers))
self._validators = tuple(
ValidatorProcess(limits.worker_address_space_bytes) for _ in range(limits.max_concurrent_workers)
)
async def start(self) -> MontyFailure | None:
if self._pool is not None:
@ -200,23 +203,20 @@ class MontyRuntime:
return MontyFailure(MontyErrorKind.QUEUE_FULL, "Monty execution queue is full")
self._admitted += 1
try:
async with asyncio.timeout_at(asyncio.get_running_loop().time() + deadline_monotonic - time.monotonic()):
async with self._slots:
worker: Final = next(worker for worker in self._validators if not worker.busy)
worker.busy = True
try:
validated: Final = await self._validate(
input, compiled.input_schema, worker, deadline_monotonic
)
if isinstance(validated, MontyFailure):
return validated
return await self._execute(compiled, validated, deadline_monotonic, started, worker)
except asyncio.CancelledError:
await worker.close()
raise
finally:
worker.busy = False
except TimeoutError:
async with timeout(max(0, deadline_monotonic - time.monotonic())), self._slots:
worker: Final = next(worker for worker in self._validators if not worker.busy)
worker.busy = True
try:
validated: Final = await self._validate(input, compiled.input_schema, worker, deadline_monotonic)
if isinstance(validated, MontyFailure):
return validated
return await self._execute(compiled, validated, deadline_monotonic, started, worker)
except asyncio.CancelledError:
await worker.close()
raise
finally:
worker.busy = False
except asyncio.TimeoutError:
return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired")
finally:
self._admitted -= 1

View file

@ -9,9 +9,10 @@ from typing import Final
from pydantic import BaseModel, ConfigDict, JsonValue, StrictBool, StrictStr
from ._compat import timeout
from ._validator_worker import MAX_MESSAGE_BYTES
from .compiler import escape_diagnostics
from .models import JsonSchema, MontyErrorKind, MontyFailure
from .models import JsonSchema, MontyErrorKind, MontyFailure, MontyLimits
class _Response(BaseModel):
@ -21,7 +22,8 @@ class _Response(BaseModel):
class ValidatorProcess:
def __init__(self) -> None:
def __init__(self, address_space_bytes: int = MontyLimits().worker_address_space_bytes) -> None:
self._address_space_bytes = address_space_bytes
self.busy = False
self.process: asyncio.subprocess.Process | None = None
@ -29,7 +31,7 @@ class ValidatorProcess:
if self.process is not None and self.process.returncode is None:
return None
try:
async with asyncio.timeout(5):
async with timeout(5):
self.process = await asyncio.create_subprocess_exec(
sys.executable,
"-I",
@ -43,10 +45,16 @@ class ValidatorProcess:
if self.process.stdout is None or await self.process.stdout.readline() != b"ready\n":
await self.close()
return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "validator startup failed")
if sys.platform == "linux":
import resource
resource.prlimit(
self.process.pid, resource.RLIMIT_AS, (self._address_space_bytes, self._address_space_bytes)
)
except asyncio.CancelledError:
await self.close()
raise
except (OSError, ValueError, TimeoutError) as exc:
except (OSError, ValueError, asyncio.TimeoutError) as exc:
await self.close()
return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc)))
return None
@ -66,7 +74,7 @@ class ValidatorProcess:
self, value: JsonValue, schema: JsonSchema, *, deadline: float, output: bool = False
) -> MontyFailure | None:
try:
async with asyncio.timeout(max(0, deadline - time.monotonic())):
async with timeout(max(0, deadline - time.monotonic())):
if (failure := await self.start()) is not None:
return failure
process: Final = self.process
@ -89,7 +97,7 @@ class ValidatorProcess:
except asyncio.CancelledError:
await self.close()
raise
except TimeoutError:
except asyncio.TimeoutError:
await self.close()
return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired during schema validation")
except (OSError, ValueError, RecursionError) as exc:

View file

@ -0,0 +1,345 @@
import asyncio
import json
import time
from collections.abc import Callable
from datetime import datetime, timedelta, timezone
from typing import Final, Protocol
from uuid import uuid4
from fastapi import Request
from pydantic import JsonValue, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.models.workflow_definition import (
Binding,
ChatCompletionNode,
ValueBinding,
WorkflowBundle,
WorkflowDefinitionStore,
WorkflowNode,
)
from litellm.models.workflow_execution import (
ChildCallResult,
EngineRun,
ExecutionError,
RunInput,
RunState,
StepTrace,
WorkflowExecutionStore,
workflow_call_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.workflows.monty.compiler import compile_program
from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits
from litellm.proxy.workflows.monty.runtime import MontyRuntime
from litellm.proxy.workflows.monty.validation import ValidatorProcess
from litellm.types.utils import ModelResponse
class ChatExecutor(Protocol):
async def __call__(
self,
request: Request,
run: EngineRun,
node: ChatCompletionNode,
step: StepTrace,
) -> ChildCallResult | ExecutionError: ...
def utc_now() -> datetime:
return datetime.now(timezone.utc)
def resolve_binding(binding: Binding, run: EngineRun, steps: tuple[StepTrace, ...]) -> JsonValue | ExecutionError:
if isinstance(binding, ValueBinding):
return binding.value
if binding.source == "$input":
return [dict(message) for message in run.input.messages]
source: Final = next((step for step in steps if step.node_id == binding.source), None)
return source.output if source is not None else ExecutionError("missing_input", "Node output is unavailable")
def _inputs(node: WorkflowNode, run: EngineRun, steps: tuple[StepTrace, ...]) -> JsonValue | ExecutionError:
if isinstance(node, ChatCompletionNode):
return resolve_binding(node.messages, run, steps)
resolved: Final = {name: resolve_binding(binding, run, steps) for name, binding in node.inputs.items()}
failure: Final = next((value for value in resolved.values() if isinstance(value, ExecutionError)), None)
return failure if failure is not None else {k: v for k, v in resolved.items() if not isinstance(v, ExecutionError)}
def _settled(task: asyncio.Task[None]) -> None:
if not task.cancelled() and task.exception() is not None:
verbose_proxy_logger.error("Workflow execution could not persist its final status")
class WorkflowRunner:
def __init__(
self,
runs: WorkflowExecutionStore,
definitions: WorkflowDefinitionStore,
execute_chat: ChatExecutor,
*,
clock: Callable[[], datetime] = utc_now,
) -> None:
self.runs: Final = runs
self.definitions: Final = definitions
self._execute_chat: Final = execute_chat
self._clock: Final = clock
self._slots: Final = asyncio.Semaphore(4)
self._monty: Final = MontyRuntime(MontyLimits())
self._monty_start: Final = asyncio.Lock()
self._active: tuple[asyncio.Task[None], ...] = ()
async def submit(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> ExecutionError | None:
self._active = tuple(task for task in self._active if not task.done())
if len(self._active) >= 64:
await self.runs.finish(run.run_id, RunState.failed, None, "Workflow capacity reached")
return ExecutionError("capacity_exceeded", "Workflow capacity reached")
task: Final = asyncio.create_task(self._run(run, request, caller))
task.add_done_callback(_settled)
self._active = (*self._active, task)
return None
async def close(self) -> None:
if self._active:
_, pending = await asyncio.wait(self._active, timeout=30)
for task in pending:
task.cancel()
await asyncio.gather(*self._active, return_exceptions=True)
await self._monty.close()
async def _watch(self, run_id: str) -> ExecutionError:
while True:
await asyncio.sleep(0.5)
current: EngineRun | None = await self.runs.get(run_id)
if current is None or current.status in {RunState.failed, RunState.cancelled}:
return ExecutionError("interrupted", "Execution stopped or its deadline expired")
if current.cancel_requested:
return ExecutionError("cancelled", "Cancellation requested")
async def _run(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> None:
execution: Final = asyncio.create_task(self._execute(run, request, caller))
watch: Final = asyncio.create_task(self._watch(run.run_id))
try:
done, _ = await asyncio.wait((execution, watch), return_when=asyncio.FIRST_COMPLETED)
result: Final = execution.result() if execution in done else watch.result()
state: Final = (
RunState.cancelled
if isinstance(result, ExecutionError) and result.code == "cancelled"
else RunState.failed
if isinstance(result, ExecutionError)
else RunState.completed
)
await self.runs.finish(
run.run_id,
state,
None if isinstance(result, ExecutionError) else result,
result.message if isinstance(result, ExecutionError) else None,
)
except asyncio.CancelledError:
await self.runs.finish(run.run_id, RunState.failed, None, "Execution worker stopped")
raise
except Exception: # noqa: BLE001 # settle background failures without exposing inputs or provider credentials
await self.runs.finish(run.run_id, RunState.failed, None, "Execution could not complete")
verbose_proxy_logger.error("Workflow execution failed")
finally:
execution.cancel()
watch.cancel()
await asyncio.gather(execution, watch, return_exceptions=True)
async def _execute(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> JsonValue | ExecutionError:
from litellm.proxy.common_request_processing import process_workflow_request
async with self._slots:
snapshot: Final = await self.definitions.get_snapshot(run.snapshot_id, run.team_id)
if snapshot is None or snapshot.workflow_id != run.workflow_id:
return ExecutionError("snapshot_unavailable", "Pinned definition unavailable")
async def execute(messages: JsonValue) -> ModelResponse | ExecutionError:
guarded: Final = run.model_copy(
update={
"input": RunInput.model_validate({**run.input.model_dump(), "messages": messages}),
}
)
validator: Final = ValidatorProcess()
try:
invalid: Final = await validator.validate(
messages,
dict(snapshot.bundle.input_schema),
deadline=time.monotonic()
+ min(5, max(0, (run.input.deadline - self._clock()).total_seconds())),
)
finally:
await validator.close()
if invalid is not None:
return ExecutionError("invalid_input", "Messages did not satisfy the workflow input schema")
output: Final = await self._sequence(guarded, request, snapshot.bundle, ())
if isinstance(output, ExecutionError):
return output
steps: Final = await self.runs.steps(run.run_id)
return ModelResponse(
id=f"chatcmpl-workflow-{run.run_id}",
model=f"workflow/{snapshot.name}",
created=int(run.created_at.timestamp()),
choices=[
{
"index": 0,
"message": {
"role": "assistant",
"content": output
if isinstance(output, str)
else json.dumps(output, ensure_ascii=False, allow_nan=False),
},
"finish_reason": "stop",
}
],
usage={
"prompt_tokens": sum(step.prompt_tokens for step in steps),
"completion_tokens": sum(step.completion_tokens for step in steps),
"total_tokens": sum(step.prompt_tokens + step.completion_tokens for step in steps),
},
)
response: Final = await process_workflow_request(
request=request,
caller=caller,
workflow_model=f"workflow/{snapshot.name}",
messages=[dict(message) for message in run.input.messages],
execute=execute,
)
return (
response
if isinstance(response, ExecutionError)
else TypeAdapter[JsonValue](JsonValue).validate_json(response.model_dump_json())
)
async def _sequence(
self,
run: EngineRun,
request: Request,
bundle: WorkflowBundle,
steps: tuple[StepTrace, ...],
) -> JsonValue | ExecutionError:
if len(steps) == len(bundle.nodes):
return resolve_binding(bundle.output, run, steps)
node: Final = bundle.nodes[len(steps)]
value: Final = _inputs(node, run, steps)
if isinstance(value, ExecutionError):
return value
started: Final = self._clock()
step: Final = StepTrace(
step_id=str(uuid4()),
run_id=run.run_id,
node_id=node.id,
state=RunState.running,
input=value,
attempts=1,
started_at=started,
deadline=min(run.input.deadline, started + timedelta(seconds=node.timeout_seconds)),
)
result: Final = await self._attempt(run, request, node, step)
return (
result
if isinstance(result, ExecutionError)
else await self._sequence(run, request, bundle, (*steps, result))
)
async def _attempt(
self,
run: EngineRun,
request: Request,
node: WorkflowNode,
step: StepTrace,
) -> StepTrace | ExecutionError:
attempt: Final = step.model_copy(
update={
"call_id": workflow_call_id(run.run_id, node.id, step.attempts)
if isinstance(node, ChatCompletionNode)
else None,
}
)
if not await self.runs.checkpoint(run.run_id, attempt):
return ExecutionError("cancelled", "Execution stopped before dispatch")
try:
result: Final = await self._safe_node(run, request, node, attempt)
except asyncio.CancelledError:
await self.runs.checkpoint(
run.run_id,
attempt.model_copy(
update={
"state": RunState.cancelled,
"finished_at": self._clock(),
"error": "Node execution interrupted",
}
),
)
raise
receipt: Final = result.receipt if isinstance(result, ExecutionError) else result
finished: Final = attempt.model_copy(
update={
"state": RunState.failed if isinstance(result, ExecutionError) else RunState.completed,
"finished_at": self._clock(),
"error": result.message if isinstance(result, ExecutionError) else None,
"output": None if isinstance(result, ExecutionError) else result.output,
"prompt_tokens": receipt.prompt_tokens if receipt else 0,
"completion_tokens": receipt.completion_tokens if receipt else 0,
"response_cost": receipt.response_cost if receipt else None,
"request_id": receipt.request_id if receipt else None,
}
)
await self.runs.checkpoint(run.run_id, finished)
if not isinstance(result, ExecutionError):
return finished
retryable: Final = (
result.code == "admission_rate_limit"
if isinstance(node, ChatCompletionNode)
else result.code in {"runtime_error", "worker_crashed"}
)
if (
not retryable
or result.receipt is not None
or step.attempts >= node.retry.max_attempts
or (step.deadline - self._clock()).total_seconds() <= node.retry.backoff_seconds
):
return result
await asyncio.sleep(node.retry.backoff_seconds)
return await self._attempt(run, request, node, step.model_copy(update={"attempts": step.attempts + 1}))
async def _safe_node(
self,
run: EngineRun,
request: Request,
node: WorkflowNode,
step: StepTrace,
) -> ChildCallResult | ExecutionError:
try:
return await self._node(run, request, node, step)
except Exception: # noqa: BLE001 # persist a sanitized terminal step for unexpected runtime failures
return ExecutionError("execution_failed", "The node could not complete")
async def _node(
self,
run: EngineRun,
request: Request,
node: WorkflowNode,
step: StepTrace,
) -> ChildCallResult | ExecutionError:
if isinstance(node, ChatCompletionNode):
return await self._execute_chat(request, run, node, step)
compiled: Final = compile_program(node.program, self._monty.limits)
if isinstance(compiled, MontyFailure):
return ExecutionError("invalid_program", compiled.message)
async with self._monty_start:
availability: Final = await self._monty.start()
if availability is not None:
return ExecutionError("runtime_unavailable", "The Monty runtime is unavailable")
result: Final = await self._monty.run_program(
compiled,
step.input,
deadline_monotonic=time.monotonic() + max(0, (step.deadline - self._clock()).total_seconds()),
)
return (
ExecutionError(result.kind.value.lower(), result.message)
if isinstance(result, MontyFailure)
else ChildCallResult(result.output, 0, 0, None, None)
)

View file

@ -0,0 +1,213 @@
from collections.abc import Mapping
from contextlib import AbstractAsyncContextManager
from typing import Final, Protocol, TypeVar
from uuid import uuid4
from liteflow import bundle_digest # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.workflow_definition import (
DefinitionError,
DefinitionPage,
WorkflowBundle,
WorkflowDefinition,
WorkflowDraft,
WorkflowSnapshot,
)
from litellm.repositories.base_repository import DbRecord, record_to_dict
from litellm.repositories.prisma_protocols import TableActions
_Model: Final = TypeVar("_Model", bound=BaseModel)
_JSON: Final = TypeAdapter(dict[str, JsonValue])
_STRING: Final = TypeAdapter(str)
class DefinitionTransaction(Protocol):
@property
def litellm_workflowdefinition(self) -> TableActions[DbRecord]: ...
@property
def litellm_workflowsnapshot(self) -> TableActions[DbRecord]: ...
@property
def litellm_teamtable(self) -> TableActions[DbRecord]: ...
async def query_raw(self, query: str, *args: object) -> object: ...
class DefinitionDatabase(DefinitionTransaction, Protocol):
def tx(self) -> AbstractAsyncContextManager[DefinitionTransaction]: ...
class _DefinitionRow(BaseModel):
model_config = ConfigDict(frozen=True)
draft: WorkflowDraft
contributors: tuple[str, ...]
published_etag: str | None
def _decoded(row: DbRecord, model: type[_Model]) -> _Model:
values: Final = record_to_dict(row)
bundle: Final = _JSON.validate_json(_STRING.validate_python(values["bundle_json"]))
return model.model_validate({field: bundle if field == "bundle" else values[field] for field in model.model_fields})
def _definition(row: DbRecord) -> _DefinitionRow:
return _DefinitionRow.model_validate({**record_to_dict(row), "draft": _decoded(row, WorkflowDraft)})
def _content(bundle: WorkflowBundle) -> Mapping[str, object]:
serialized: Final = bundle.model_dump_json(by_alias=True, exclude_unset=True)
return {
"etag": f'"{uuid4()}"',
"digest": bundle_digest(_JSON.validate_json(serialized)),
"bundle_json": serialized,
}
async def _locked_definition(
transaction: DefinitionTransaction, workflow_id: str, team_id: str, etag: str
) -> _DefinitionRow | DefinitionError:
await transaction.query_raw(
'SELECT workflow_id FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND team_id = $2 FOR UPDATE',
workflow_id,
team_id,
)
stored: Final = await transaction.litellm_workflowdefinition.find_first(
where={"workflow_id": workflow_id, "team_id": team_id}
)
if stored is None:
return DefinitionError("not_found", "Workflow not found")
row: Final = _definition(stored)
if row.draft.etag != etag:
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
return row
async def _snapshot(
transaction: DefinitionTransaction, row: _DefinitionRow, published_by: str | None
) -> WorkflowSnapshot:
stored: Final = await transaction.litellm_workflowsnapshot.create(
data={
"snapshot_id": str(uuid4()),
"workflow_id": row.draft.workflow_id,
"team_id": row.draft.team_id,
"name": row.draft.name,
"digest": row.draft.digest,
"bundle_json": _JSON.dump_json(dict(row.draft.bundle)).decode(),
"contributors": list(row.contributors),
"published_by": published_by,
}
)
return _decoded(stored, WorkflowSnapshot)
class WorkflowDefinitionRepository:
def __init__(self, database: DefinitionDatabase) -> None:
self._database: Final = database
async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None:
row: Final = await self._database.litellm_teamtable.find_unique(where={"team_id": team_id})
return None if row is None else LiteLLM_TeamTable.model_validate(record_to_dict(row))
async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None:
stored: Final = await self._database.litellm_workflowdefinition.find_first(
where={"workflow_id": workflow_id, **({"team_id": team_id} if team_id is not None else {})}
)
if stored is None:
return None
row: Final = _definition(stored)
return WorkflowDefinition(draft=row.draft, contributors=row.contributors)
async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage:
rows: Final = await self._database.litellm_workflowdefinition.find_many(
where={"team_id": team_id, **({"workflow_id": {"gt": cursor}} if cursor is not None else {})},
order={"workflow_id": "asc"},
take=limit + 1,
)
drafts: Final = tuple(_decoded(row, WorkflowDraft) for row in rows[:limit])
return DefinitionPage(workflows=drafts, next_cursor=drafts[-1].workflow_id if len(rows) > limit else None)
async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None:
row: Final = await self._database.litellm_workflowsnapshot.find_first(
where={"snapshot_id": snapshot_id, **({"team_id": team_id} if team_id is not None else {})}
)
return None if row is None else _decoded(row, WorkflowSnapshot)
async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None:
row: Final = await self._database.litellm_workflowdefinition.find_first(
where={"name": name, "team_id": team_id}
)
if row is None or (snapshot_id := _decoded(row, WorkflowDraft).published_snapshot_id) is None:
return None
return await self.get_snapshot(snapshot_id, team_id)
async def create(self, bundle: WorkflowBundle, team_id: str, contributor: str) -> WorkflowDraft | DefinitionError:
from prisma.errors import UniqueViolationError
try:
stored: Final = await self._database.litellm_workflowdefinition.create(
data={
"workflow_id": str(uuid4()),
"name": bundle.name,
"team_id": team_id,
"contributors": [contributor],
**_content(bundle),
}
)
except UniqueViolationError:
return DefinitionError("conflict", "A workflow with this name already exists in the team", "bundle.name")
return _decoded(stored, WorkflowDraft)
async def replace(
self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str
) -> WorkflowDraft | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
if row.draft.name != bundle.name:
return DefinitionError("conflict", "A workflow name cannot change", "bundle.name")
contributors: Final = (
(contributor,)
if etag == row.published_etag
else tuple(sorted(frozenset(row.contributors) | {contributor}))
)
stored: Final = await transaction.litellm_workflowdefinition.update(
where={"workflow_id": workflow_id}, data={**_content(bundle), "contributors": list(contributors)}
)
assert stored is not None
return _decoded(stored, WorkflowDraft)
async def publish(
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
) -> WorkflowSnapshot | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
if row.draft.digest != digest:
return DefinitionError("stale_draft", "The reviewed digest does not match the draft", "digest")
if publisher in row.contributors:
return DefinitionError("forbidden", "Draft contributors cannot publish their own changes")
if row.published_etag == etag and row.draft.published_snapshot_id is not None:
existing: Final = await transaction.litellm_workflowsnapshot.find_unique(
where={"snapshot_id": row.draft.published_snapshot_id}
)
assert existing is not None
return _decoded(existing, WorkflowSnapshot)
snapshot: Final = await _snapshot(transaction, row, publisher)
await transaction.litellm_workflowdefinition.update(
where={"workflow_id": workflow_id},
data={"published_snapshot_id": snapshot.snapshot_id, "published_etag": etag},
)
return snapshot
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
return await _snapshot(transaction, row, None)

View file

@ -0,0 +1,297 @@
import base64
import json
from contextlib import AbstractAsyncContextManager
from datetime import datetime, timezone
from itertools import chain
from typing import Final, Protocol
from uuid import uuid4
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from litellm.models.workflow_execution import (
EngineRun,
ExecutionError,
RunAdmission,
RunInput,
RunPage,
RunState,
RunView,
StepTrace,
)
class ExecutionQueries(Protocol):
async def query_raw(self, query: str, *args: object) -> object: ...
class ExecutionDatabase(ExecutionQueries, Protocol):
def tx(self) -> AbstractAsyncContextManager[ExecutionQueries]: ...
class _Cursor(BaseModel):
model_config = ConfigDict(frozen=True)
created_at: datetime
run_id: str
class _TraceRow(BaseModel):
model_config = ConfigDict(frozen=True)
trace: StepTrace
_RUNS: Final = TypeAdapter(tuple[EngineRun, ...])
_VIEWS: Final = TypeAdapter(tuple[RunView, ...])
_TRACES: Final = TypeAdapter(tuple[_TraceRow, ...])
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
def _cursor(value: str | None) -> _Cursor | ExecutionError | None:
if value is None:
return None
try:
return _Cursor.model_validate_json(base64.urlsafe_b64decode(value.encode("ascii")))
except (ValueError, UnicodeError, ValidationError):
return ExecutionError("invalid_cursor", "Invalid run cursor")
def _encoded_cursor(run: RunView) -> str:
return base64.urlsafe_b64encode(
_Cursor(created_at=run.created_at, run_id=run.run_id).model_dump_json().encode()
).decode("ascii")
async def _first_run(database: ExecutionQueries, query: str, *args: object) -> EngineRun | None:
rows: Final = _RUNS.validate_python(await database.query_raw(query, *args))
return rows[0] if rows else None
def _aggregate_attempts(attempts: tuple[StepTrace, ...]) -> StepTrace:
latest: Final = max(attempts, key=lambda step: step.attempts)
costs: Final = tuple(step.response_cost for step in attempts if step.response_cost is not None)
ids: Final = chain.from_iterable(
step.request_ids or ((step.request_id,) if step.request_id is not None else ()) for step in attempts
)
return latest.model_copy(
update={
"prompt_tokens": sum(step.prompt_tokens for step in attempts),
"completion_tokens": sum(step.completion_tokens for step in attempts),
"response_cost": sum(costs) if costs else None,
"request_ids": tuple(dict.fromkeys(ids)),
}
)
class WorkflowExecutionRepository:
def __init__(self, database: ExecutionDatabase) -> None:
self._database: Final = database
async def find_idempotent(
self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str
) -> EngineRun | None:
return await _first_run(
self._database,
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 '
"AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text",
key_hash,
team_id,
workflow_id,
idempotency_key,
)
async def admit(
self,
*,
workflow_id: str,
snapshot_id: str,
team_id: str,
key_hash: str,
input: RunInput,
digest: str,
idempotency_key: str | None,
) -> RunAdmission | ExecutionError:
await self._expire_team(team_id)
async with self._database.tx() as transaction:
await transaction.query_raw(
"SELECT true AS acquired FROM pg_advisory_xact_lock(hashtextextended($1, 0))", team_id
)
existing: Final = await _first_run(
transaction,
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 '
"AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text",
key_hash,
team_id,
workflow_id,
idempotency_key,
)
if existing is not None:
return (
RunAdmission(existing, False)
if existing.request_digest == digest
else ExecutionError("idempotency_conflict", "Idempotency key was used with different input")
)
capacity: Final = _ROWS.validate_python(
await transaction.query_raw(
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND team_id = $1 '
"AND status IN ('pending', 'running') LIMIT 64",
team_id,
)
)
if len(capacity) >= 64:
return ExecutionError("capacity_exceeded", "The team's workflow execution queue is full")
created: Final = await _first_run(
transaction,
'INSERT INTO "LiteLLM_WorkflowRun" '
"(run_id, session_id, workflow_type, status, created_by, created_at, updated_at, "
"workflow_id, snapshot_id, team_id, engine_owned, input, request_digest, idempotency_key) "
"VALUES ($1, $2, 'liteflow', 'pending', $3, timezone('UTC', NOW()), timezone('UTC', NOW()), $4, $5, $6, true, $7::jsonb, $8, $9) "
"RETURNING *",
str(uuid4()),
str(uuid4()),
key_hash,
workflow_id,
snapshot_id,
team_id,
input.model_dump_json(),
digest,
idempotency_key,
)
if created is None:
return ExecutionError("storage_error", "Workflow admission did not create a run")
return RunAdmission(created, True)
async def _expire_team(self, team_id: str) -> None:
await self._database.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET status = CASE WHEN cancel_requested THEN 'cancelled' ELSE 'failed' END, "
"metadata = jsonb_build_object('error', 'Execution deadline expired'), updated_at = timezone('UTC', NOW()) "
"WHERE team_id = $1 AND engine_owned = true AND status IN ('pending', 'running') "
"AND (input->>'deadline')::timestamptz <= NOW() RETURNING run_id",
team_id,
)
async def get(self, run_id: str) -> EngineRun | None:
query: Final = 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true'
run: Final = await _first_run(self._database, query, run_id)
if (
run is not None
and run.status in {RunState.pending, RunState.running}
and run.input.deadline <= datetime.now(timezone.utc)
):
await self._expire_team(run.team_id)
return await _first_run(self._database, query, run_id)
return run
async def checkpoint(self, run_id: str, step: StepTrace) -> bool:
async with self._database.tx() as transaction:
owned: Final = _ROWS.validate_python(
await transaction.query_raw(
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true '
"AND ($2::boolean = false OR (status IN ('pending', 'running') AND NOT cancel_requested "
"AND (input->>'deadline')::timestamptz > NOW())) FOR UPDATE",
run_id,
step.state is RunState.running,
)
)
if not owned:
return False
await transaction.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET status = 'running', updated_at = timezone('UTC', NOW()) "
"WHERE run_id = $1 AND status = 'pending' RETURNING run_id",
run_id,
)
await transaction.query_raw(
'INSERT INTO "LiteLLM_WorkflowEvent" '
"(event_id, run_id, event_type, step_name, sequence_number, data, created_at) "
"SELECT $1, $2, 'liteflow.step', $3, COALESCE(MAX(sequence_number) + 1, 0), $4::jsonb, timezone('UTC', NOW()) "
'FROM "LiteLLM_WorkflowEvent" WHERE run_id = $2 RETURNING event_id',
str(uuid4()),
run_id,
step.node_id,
json.dumps({"trace": step.model_dump(mode="json")}, allow_nan=False),
)
return True
async def steps(self, run_id: str) -> tuple[StepTrace, ...]:
rows: Final = _TRACES.validate_python(
await self._database.query_raw(
"SELECT trace FROM ("
"SELECT CASE WHEN receipts.ids IS NULL THEN latest.trace ELSE latest.trace || "
"jsonb_build_object('request_id', CASE WHEN cardinality(receipts.ids) = 1 THEN receipts.ids[1] END, "
"'request_ids', to_jsonb(receipts.ids), 'response_cost', receipts.cost, "
"'prompt_tokens', receipts.prompt_tokens, 'completion_tokens', receipts.completion_tokens) END AS trace, "
"latest.step_name FROM ("
"SELECT DISTINCT ON (step_name, data->'trace'->>'attempts') data->'trace' AS trace, step_name "
"FROM \"LiteLLM_WorkflowEvent\" WHERE run_id = $1 AND event_type = 'liteflow.step' "
"ORDER BY step_name, data->'trace'->>'attempts', sequence_number DESC"
') latest JOIN "LiteLLM_WorkflowRun" run ON run.run_id = $1 AND run.engine_owned = true '
"LEFT JOIN LATERAL ("
"SELECT array_agg(log.request_id ORDER BY log.request_id) AS ids, SUM(log.spend) AS cost, "
"SUM(log.prompt_tokens) AS prompt_tokens, SUM(log.completion_tokens) AS completion_tokens "
"FROM \"LiteLLM_SpendLogs\" log WHERE log.litellm_call_id = latest.trace->>'call_id' "
"AND log.api_key = run.created_by AND log.team_id = run.team_id"
") receipts ON true"
") latest ORDER BY (trace->>'started_at')::timestamptz, step_name",
run_id,
)
)
node_ids: Final = dict.fromkeys(row.trace.node_id for row in rows)
return tuple(
_aggregate_attempts(tuple(row.trace for row in rows if row.trace.node_id == node_id))
for node_id in node_ids
)
async def finish(self, run_id: str, state: RunState, output: JsonValue, error: str | None) -> None:
await self._database.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET status = CASE WHEN cancel_requested THEN 'cancelled' ELSE $2 END, "
"output = CASE WHEN cancel_requested THEN NULL ELSE $3::jsonb END, "
"metadata = jsonb_build_object('error', $4::text), updated_at = timezone('UTC', NOW()) "
"WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id",
run_id,
state.value,
json.dumps(output, allow_nan=False),
error,
)
async def cancel(self, run_id: str) -> None:
await self._database.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET cancel_requested = true, updated_at = timezone('UTC', NOW()) "
"WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id",
run_id,
)
async def list_runs(
self,
workflow_id: str,
team_id: str,
key_hash: str | None,
state: RunState | None,
limit: int,
cursor: str | None,
) -> RunPage | ExecutionError:
await self._expire_team(team_id)
position: Final = _cursor(cursor)
if isinstance(position, ExecutionError):
return position
rows: Final = _VIEWS.validate_python(
await self._database.query_raw(
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true '
"AND workflow_id = $1 AND team_id = $2 "
"AND ($3::text IS NULL OR created_by = $3::text) AND ($4::text IS NULL OR status = $4::text) "
"AND ($5::text IS NULL OR (created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)) "
"ORDER BY created_at DESC, run_id DESC LIMIT $7",
workflow_id,
team_id,
key_hash,
state.value if state is not None else None,
position.created_at.isoformat() if position is not None else None,
position.run_id if position is not None else None,
limit + 1,
)
)
page: Final = rows[:limit]
return RunPage(
runs=page,
count=len(page),
next_cursor=_encoded_cursor(page[-1]) if len(rows) > limit else None,
)

View file

@ -55,6 +55,7 @@ Documentation = "https://docs.litellm.ai"
# supply-chain-pinned to a vetted, aged release.
[project.optional-dependencies]
proxy = [
"liteflow @ https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz#sha256=c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072",
"pydantic-monty==1.4.0",
"gunicorn>=23.0.0,<24.0",
"uvicorn>=0.33.0,<1.0",
@ -258,6 +259,7 @@ admin-mcp = [
"litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/31fb096904ae05684c0779110a93b135b5601182.tar.gz#sha256=5803a20a00a18177f4591706956f31be311879543049b8b38143ad90bb1cb0a6 ; python_version >= '3.12'",
]
proxy-dev = [
"liteflow @ https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz#sha256=c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072",
"pydantic-monty==1.4.0",
{ include-group = "admin-mcp" },
"prisma==0.11.0",

View file

@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
// One instance of work being done. `status` is a materialized cache of the
// latest event; the event log is the authoritative source of truth.
model LiteLLM_WorkflowDefinition {
workflow_id String @id @default(uuid())
team_id String
name String
etag String
digest String
bundle_json String
contributors String[]
published_snapshot_id String?
published_etag String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
snapshots LiteLLM_WorkflowSnapshot[]
@@unique([team_id, name])
@@index([team_id, workflow_id])
}
model LiteLLM_WorkflowSnapshot {
snapshot_id String @id @default(uuid())
workflow_id String
team_id String
name String
digest String
bundle_json String
contributors String[]
published_by String?
created_at DateTime @default(now())
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
@@index([workflow_id, created_at, snapshot_id])
@@index([team_id, snapshot_id])
}
model LiteLLM_WorkflowRun {
run_id String @id @default(uuid())
session_id String @unique @default(uuid())
workflow_type String
workflow_id String?
snapshot_id String?
team_id String?
engine_owned Boolean @default(false)
request_digest String?
idempotency_key String?
cancel_requested Boolean @default(false)
status String @default("pending")
created_by String? // user_id of the key that created this run; null = created by master key
created_at DateTime @default(now())
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
@@index([session_id])
@@index([created_at])
@@index([created_by])
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
@@index([workflow_id, team_id, created_at, run_id])
@@index([team_id, engine_owned, status])
}
// Append-only log of state transitions. Never mutate rows here.

View file

@ -7,18 +7,22 @@ count them per rule the same delta-vs-base way as the other lint gates. The rule
OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI
what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix)
OA002 no description (the handler has no docstring and the route no ``description=``)
OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``)
OA003 no 2xx success response has a schema (use ``response_model`` or a schema in ``responses``)
"""
from __future__ import annotations
import json
import sys
from collections.abc import Iterator, Mapping
from collections.abc import Iterator, Mapping, Sequence
from typing import Final, NamedTuple
from pydantic import TypeAdapter
GATED_METHODS: Final = frozenset({"post"})
SCHEMA_MEDIA_TYPES: Final = ("application/json", "application/sdp", "multipart/form-data")
SUCCESS_STATUS_CODES: Final = frozenset(str(code) for code in range(200, 300))
_OBJECT: Final = TypeAdapter(Mapping[str, object])
class Violation(NamedTuple):
@ -28,7 +32,7 @@ class Violation(NamedTuple):
def _as_mapping(node: object) -> Mapping[str, object]:
return node if isinstance(node, Mapping) else {}
return _OBJECT.validate_python(node) if isinstance(node, Mapping) else {}
def _content_schema(container: object) -> object:
@ -45,7 +49,11 @@ def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]:
description: Final = operation.get("description")
if not isinstance(description, str) or not description.strip():
yield "OA002"
if not _content_schema(_as_mapping(operation.get("responses")).get("200")):
if not any(
_content_schema(response)
for status, response in _as_mapping(operation.get("responses")).items()
if status in SUCCESS_STATUS_CODES
):
yield "OA003"
@ -66,10 +74,10 @@ def proxy_spec() -> Mapping[str, object]:
return _as_mapping(app.openapi())
def main(argv: list[str] = sys.argv[1:]) -> int:
def main(argv: Sequence[str] = sys.argv[1:]) -> int:
spec: Final = json.loads(open(argv[0]).read()) if argv else proxy_spec()
for violation in violations(spec):
print(f"{violation.method} {violation.path} {violation.code}")
sys.stdout.write(f"{violation.method} {violation.path} {violation.code}\n")
return 0

View file

@ -0,0 +1,427 @@
import asyncio
import json
import subprocess
import sys
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final, cast
import psycopg
import pytest
import pytest_asyncio
from fastapi import HTTPException, Request
from openai.types.chat import ChatCompletion
from prisma import Json, Prisma
from prisma.types import DatasourceOverride
from psycopg import sql
from litellm.models.workflow_definition import ChatCompletionNode, WorkflowDraft, WorkflowSnapshot
from litellm.models.workflow_execution import (
ChildCallResult,
EngineRun,
ExecutionError,
RunAdmission,
RunInput,
RunPage,
RunState,
StepTrace,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
TestRunRequest as PaidRunRequest,
)
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
cancel_workflow_run,
get_run_definition,
get_workflow_steps,
list_definition_runs,
)
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
test_workflow as run_draft_test,
)
from litellm.proxy.workflows.runner import WorkflowRunner
from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
from tests.integration._support.database import scratch_database
from tests.unit.proxy.workflows.test_runner import two_nodes
pytestmark: Final = pytest.mark.asyncio(loop_scope="function")
_ROOT: Final = Path(__file__).resolve().parents[3]
_MESSAGES: Final = ({"role": "user", "content": "plan"},)
_CALLER: Final = UserAPIKeyAuth(token="owner", team_id="team").model_copy(update={"via_virtual_key": True})
_REQUEST: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []})
@dataclass(frozen=True, slots=True)
class Storage:
db: Prisma
runs: WorkflowExecutionRepository
definitions: WorkflowDefinitionRepository
async def admit(self, *, key: str | None = None, digest: str = "input", owner: str = "owner",
team: str = "team", snapshot: WorkflowSnapshot | None = None,
deadline: datetime | None = None) -> RunAdmission | ExecutionError:
return await self.runs.admit(
workflow_id=snapshot.workflow_id if snapshot else "workflow",
snapshot_id=snapshot.snapshot_id if snapshot else "snapshot",
team_id=team, key_hash=owner, input=RunInput(messages=_MESSAGES, draft=False,
deadline=deadline or datetime.now(timezone.utc) + timedelta(minutes=5)), digest=digest, idempotency_key=key,
)
async def snapshot(self, team: str = "team") -> WorkflowSnapshot:
draft: Final = await self.definitions.create(two_nodes(), team, "author")
assert isinstance(draft, WorkflowDraft)
snapshot: Final = await self.definitions.publish(draft.workflow_id, team, draft.etag, draft.digest, "admin")
assert isinstance(snapshot, WorkflowSnapshot)
return snapshot
@pytest.fixture(scope="module")
def schema_sql() -> bytes:
return subprocess.run(
[sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel",
str(_ROOT / "schema.prisma"), "--script"], check=True, capture_output=True, cwd=_ROOT,
).stdout
@pytest_asyncio.fixture(loop_scope="function")
async def store(schema_sql: bytes) -> AsyncIterator[Storage]:
with scratch_database("workflow_execution") as url, psycopg.connect(url, autocommit=True) as connection:
connection.execute(schema_sql)
connection.execute(sql.SQL("ALTER DATABASE {} SET timezone TO 'Asia/Singapore'").format(
sql.Identifier(connection.info.dbname)
))
async with Prisma(datasource=DatasourceOverride(url=url)) as client:
assert await client.query_raw("SHOW TIME ZONE") == [{"TimeZone": "Asia/Singapore"}]
for team in ("team", "other"):
await client.litellm_teamtable.create(data={"team_id": team, "members_with_roles": Json([
{"user_id": "admin", "role": "admin"}, {"user_id": "member", "role": "user"},
])})
yield Storage(client, WorkflowExecutionRepository(cast(ExecutionDatabase, client)),
WorkflowDefinitionRepository(cast(DefinitionDatabase, client)))
class ChatReplies:
def __init__(self, store: Storage, replies: tuple[ChildCallResult | ExecutionError | Exception, ...],
cancel_after: int | None = None) -> None:
self.store, self.replies, self.cancel_after = store, replies, cancel_after
self.calls: tuple[StepTrace, ...] = ()
self.models: tuple[str, ...] = ()
self.persisted: tuple[StepTrace, ...] = ()
async def __call__(self, request: Request, run: EngineRun, node: ChatCompletionNode,
step: StepTrace) -> ChildCallResult | ExecutionError:
self.calls = (*self.calls, step)
self.models = (*self.models, node.model)
self.persisted = (*self.persisted, (await self.store.runs.steps(run.run_id))[-1])
if len(self.calls) == self.cancel_after:
await self.store.runs.cancel(run.run_id)
reply: Final = self.replies[len(self.calls) - 1]
if isinstance(reply, Exception):
raise reply
return reply
async def test_concurrent_idempotency_pins_first_snapshot_and_conflicts_on_input(store: Storage) -> None:
original: Final = await store.snapshot()
different: Final = original.model_copy(update={"snapshot_id": "later-snapshot"})
results: Final = await asyncio.gather(*(
store.admit(key="same", snapshot=snapshot) for snapshot in (original, different) * 4
))
assert all(isinstance(result, RunAdmission) for result in results)
admitted: Final = tuple(result for result in results if isinstance(result, RunAdmission))
assert sum(result.created for result in admitted) == 1
assert len({result.run.run_id for result in admitted}) == 1
assert len({result.run.snapshot_id for result in admitted}) == 1
assert all(result.run == admitted[0].run for result in admitted)
assert await store.admit(key="same", snapshot=different, digest="changed") == ExecutionError(
"idempotency_conflict", "Idempotency key was used with different input"
)
assert isinstance(await store.admit(key="same", snapshot=original, owner="another"), RunAdmission)
assert isinstance(await store.admit(key="same", snapshot=original, team="other"), RunAdmission)
assert await store.db.litellm_workflowrun.count() == 3
async def test_draft_test_retry_after_edit_recovers_only_the_same_scoped_request(store: Storage) -> None:
draft: Final = await store.definitions.create(two_nodes(), "team", "member")
assert isinstance(draft, WorkflowDraft)
caller: Final = _CALLER.model_copy(update={"user_id": "member", "user_role": LitellmUserRoles.INTERNAL_USER})
data: Final = PaidRunRequest(messages=_MESSAGES, authorize_billing=True)
chat: Final = ChatReplies(store, (
ChildCallResult(list(_MESSAGES), 1, 1, "first", 0.125), ChildCallResult("done", 1, 1, "last", 0.125),
))
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
request: Final = Request({**_REQUEST.scope, "state": {"workflow_runner": runner}})
try:
accepted: Final = await run_draft_test(
request, draft.workflow_id, data, caller, store.definitions, store.runs,
if_match=draft.etag, idempotency_key="retry",
)
finally:
await runner.close()
changed: Final = await store.definitions.replace(
draft.workflow_id, "team", draft.etag, two_nodes(), "member",
)
assert isinstance(changed, WorkflowDraft) and changed.etag != draft.etag
replay: Final = await run_draft_test(
_REQUEST, draft.workflow_id, data, caller, store.definitions, store.runs,
if_match=draft.etag, idempotency_key="retry",
)
assert replay.status_code == accepted.status_code == 202
assert replay.body == accepted.body and replay.headers == accepted.headers
for body, identity, etag, key, status in (
(data.model_copy(update={"messages": ({"role": "user", "content": "changed"},)}), caller, draft.etag, "retry", 409),
(data, caller, changed.etag, "retry", 409),
(data, caller, None, "retry", 412),
(data, caller, draft.etag, "new-key", 412),
(data, caller.model_copy(update={"token": "peer"}), draft.etag, "retry", 412),
(data, caller.model_copy(update={"team_id": "other"}), draft.etag, "retry", 404),
(data, caller.model_copy(update={"user_id": "removed-member"}), draft.etag, "retry", 404),
(data, caller.model_copy(update={"via_virtual_key": False}), draft.etag, "retry", 403),
(data.model_copy(update={"authorize_billing": False}), caller, draft.etag, "retry", 400),
):
with pytest.raises(HTTPException) as rejected:
await run_draft_test(
_REQUEST, draft.workflow_id, body, identity, store.definitions, store.runs,
if_match=etag, idempotency_key=key,
)
assert rejected.value.status_code == status
assert await store.db.litellm_workflowrun.count() == await store.db.litellm_workflowsnapshot.count() == 1
assert len(chat.calls) == 2
async def test_nullable_history_filters_cursor_ties_and_utc_in_non_utc_database(store: Storage) -> None:
before: Final = datetime.now(timezone.utc).replace(microsecond=0)
admissions: Final = tuple([await store.admit(owner=owner) for owner in ("owner", "foreign", "owner")])
assert all(isinstance(item, RunAdmission) for item in admissions)
runs: Final = tuple(item.run for item in admissions if isinstance(item, RunAdmission))
await store.db.litellm_workflowrun.update_many(where={}, data={"created_at": runs[0].created_at})
ordered: Final = tuple(sorted((run.run_id for run in runs), reverse=True))
for owner, state in ((None, None), ("owner", None), (None, RunState.pending), ("owner", RunState.pending)):
page: Final = await store.runs.list_runs("workflow", "team", owner, state, 1, None)
assert isinstance(page, RunPage) and page.next_cursor is not None
tail: Final = await store.runs.list_runs("workflow", "team", owner, state, 10, page.next_cursor)
assert isinstance(tail, RunPage) and tail.next_cursor is None
expected: Final = tuple(run_id for run_id in ordered if owner is None or run_id != runs[1].run_id)
assert tuple(run.run_id for run in (*page.runs, *tail.runs)) == expected
await store.runs.finish(runs[0].run_id, RunState.failed, None, "failed")
finished: Final = await store.runs.get(runs[0].run_id)
assert finished is not None and finished.output is None and finished.metadata == {"error": "failed"}
filtered: Final = await store.runs.list_runs("workflow", "team", "owner", RunState.failed, 10, None)
assert isinstance(filtered, RunPage) and tuple(run.run_id for run in filtered.runs) == (finished.run_id,)
for timestamp in (runs[0].created_at, finished.updated_at):
assert before <= timestamp <= datetime.now(timezone.utc)
assert isinstance(await store.runs.list_runs("workflow", "team", None, None, 10, "bad-cursor"), ExecutionError)
async def test_attempt_receipts_deduplicate_checkpoints_and_exclude_foreign_spend(store: Storage) -> None:
admitted: Final = await store.admit()
assert isinstance(admitted, RunAdmission)
run: Final = admitted.run
first: Final = StepTrace(step_id="step", run_id=run.run_id, node_id="node", state=RunState.failed, input=None,
attempts=1, started_at=run.created_at, deadline=run.input.deadline, call_id="first",
prompt_tokens=4, completion_tokens=2, response_cost=0.125, request_id="failed")
second: Final = first.model_copy(update={
"attempts": 2, "state": RunState.completed, "call_id": "second", "output": False,
})
for step in (first, first, second, second):
assert await store.runs.checkpoint(run.run_id, step)
for request_id, key, team, call_id, session in (
("a", "owner", "team", "second", run.session_id),
("b", "owner", "team", "second", "parent-session"),
("wrong-key", "foreign", "team", "second", run.session_id),
("wrong-team", "owner", "other", "second", run.session_id),
("wrong-call", "owner", "team", "unrelated", run.session_id),
):
await store.db.litellm_spendlogs.create(data={
"request_id": request_id, "call_type": "acompletion", "api_key": key, "team_id": team,
"session_id": session, "litellm_call_id": call_id, "startTime": run.created_at, "endTime": run.created_at,
"prompt_tokens": 3, "completion_tokens": 1, "spend": 0.25,
})
assert await store.runs.steps(run.run_id) == (second.model_copy(update={
"prompt_tokens": 10, "completion_tokens": 4, "response_cost": 0.625,
"request_id": None, "request_ids": ("failed", "a", "b"),
}),)
events: Final = await store.db.litellm_workflowevent.find_many(where={"run_id": run.run_id})
assert len(events) == 4
assert all(run.created_at <= event.created_at <= datetime.now(timezone.utc) for event in events)
@pytest.mark.parametrize("phase", ["before", "retry", "downstream"])
async def test_cancellation_fences_every_dispatch_but_keeps_completed_receipts(store: Storage, phase: str) -> None:
snapshot: Final = await store.snapshot()
admitted: Final = await store.admit(snapshot=snapshot)
assert isinstance(admitted, RunAdmission)
run: Final = admitted.run
receipt: Final = ChildCallResult(list(_MESSAGES), 4, 2, "receipt", 0.125)
reply: Final = ExecutionError("admission_rate_limit", "retry") if phase == "retry" else receipt
chat: Final = ChatReplies(store, (reply,), cancel_after=1)
if phase == "before":
await store.runs.cancel(run.run_id)
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
try:
result: Final = await runner._sequence(run, _REQUEST, snapshot.bundle, ())
assert isinstance(result, ExecutionError) and result.code == "cancelled"
assert len(chat.calls) == (0 if phase == "before" else 1)
traces: Final = await store.runs.steps(run.run_id)
expected_ids: Final = () if phase == "before" else ((),) if phase == "retry" else (("receipt",),)
assert tuple(trace.request_ids for trace in traces) == expected_ids
await store.runs.finish(run.run_id, RunState.completed, {"must": "not win cancellation"}, None)
settled: Final = await store.runs.get(run.run_id)
assert settled is not None and settled.status is RunState.cancelled and settled.output is None
finally:
await runner.close()
@pytest.mark.parametrize("error", [
ExecutionError("invalid_input", "rejected"), ExecutionError("provider_error", "unavailable"),
ExecutionError("rate_limit", "upstream 429"), ExecutionError("runtime_error", "runtime failed"),
ExecutionError("admission_rate_limit", "retry budget exhausted"),
ExecutionError("provider_error", "billed failure", ChildCallResult(None, 4, 2, "billed", 0.125)),
ExecutionError("admission_rate_limit", "known receipt", ChildCallResult(None, 4, 2, "billed", 0.125)),
RuntimeError("private details"),
])
async def test_failure_stops_downstream_and_persists_sanitized_terminal_trace(
store: Storage, error: ExecutionError | Exception,
) -> None:
snapshot: Final = await store.snapshot()
admitted: Final = await store.admit(snapshot=snapshot)
assert isinstance(admitted, RunAdmission)
chat: Final = ChatReplies(store, (error, error))
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
try:
result: Final = await runner._sequence(admitted.run, _REQUEST, snapshot.bundle, ())
expected: Final = error if isinstance(error, ExecutionError) else ExecutionError(
"execution_failed", "The node could not complete"
)
assert result == expected
attempts: Final = 2 if expected.code == "admission_rate_limit" and expected.receipt is None else 1
traces: Final = await store.runs.steps(admitted.run.run_id)
assert tuple((step.node_id, step.state, step.attempts, step.error) for step in traces) == (
("node_1", RunState.failed, attempts, expected.message),
)
assert len(chat.calls) == attempts
assert (traces[0].prompt_tokens, traces[0].completion_tokens, traces[0].response_cost, traces[0].request_ids) == (
(4, 2, 0.125, ("billed",)) if expected.receipt is not None else (0, 0, None, ())
)
finally:
await runner.close()
@pytest.mark.parametrize("trigger", ["get", "admit", "list"])
async def test_deadline_expiry_is_terminal_and_fences_dispatch(store: Storage, trigger: str) -> None:
admitted: Final = await store.admit(deadline=datetime.now(timezone.utc) - timedelta(seconds=1))
assert isinstance(admitted, RunAdmission)
if trigger == "admit":
await store.admit()
elif trigger == "list":
await store.runs.list_runs("workflow", "team", None, None, 10, None)
else:
await store.runs.get(admitted.run.run_id)
stored: Final = await store.db.litellm_workflowrun.find_unique(where={"run_id": admitted.run.run_id})
assert stored is not None and stored.status == "failed"
step: Final = StepTrace(
step_id="step", run_id=stored.run_id, node_id="node", state=RunState.running,
input=None, attempts=1, started_at=admitted.run.created_at, deadline=admitted.run.input.deadline,
)
assert not await store.runs.checkpoint(stored.run_id, step)
await store.runs.finish(stored.run_id, RunState.completed, "late output", None)
expired: Final = await store.runs.get(stored.run_id)
assert expired is not None and expired.status is RunState.failed and expired.output is None
assert await store.runs.steps(stored.run_id) == ()
@pytest.mark.parametrize("role,user", [
(LitellmUserRoles.PROXY_ADMIN, "admin"), (LitellmUserRoles.INTERNAL_USER, "admin"),
(LitellmUserRoles.INTERNAL_USER, "member"), (None, None),
])
async def test_key_scope_intersects_user_scope_and_own_application_key_can_cancel(
store: Storage, role: LitellmUserRoles | None, user: str | None,
) -> None:
local, foreign = await store.snapshot(), await store.snapshot("other")
own: Final = await store.admit(snapshot=local)
peer: Final = await store.admit(snapshot=local, owner="peer")
other: Final = await store.admit(snapshot=foreign, team="other")
assert isinstance(own, RunAdmission) and isinstance(peer, RunAdmission) and isinstance(other, RunAdmission)
caller: Final = _CALLER.model_copy(update={"user_role": role, "user_id": user})
for endpoint in (get_workflow_steps, get_run_definition, cancel_workflow_run):
with pytest.raises(HTTPException) as denied:
await endpoint(other.run.run_id, caller, store.definitions, store.runs)
assert denied.value.status_code == 404
with pytest.raises(HTTPException) as denied_history:
await list_definition_runs(
foreign.workflow_id, caller, store.definitions, store.runs, team_id="other", limit=10, after=None,
)
assert denied_history.value.status_code == 404
history: Final = await list_definition_runs(
local.workflow_id, caller, store.definitions, store.runs, limit=10, after=None,
)
expected: Final = {own.run.run_id, peer.run.run_id} if user == "admin" else {own.run.run_id}
assert {run.run_id for run in history.runs} == expected
for endpoint in (get_workflow_steps, get_run_definition, cancel_workflow_run):
if user == "admin":
assert await endpoint(peer.run.run_id, caller, store.definitions, store.runs) is not None
else:
with pytest.raises(HTTPException) as denied_peer:
await endpoint(peer.run.run_id, caller, store.definitions, store.runs)
assert denied_peer.value.status_code == 404
pinned: Final = await get_run_definition(own.run.run_id, caller, store.definitions, store.runs)
assert pinned.snapshot_id == local.snapshot_id
assert (await cancel_workflow_run(own.run.run_id, caller, store.definitions, store.runs)).status_code == 202
assert (await get_workflow_steps(own.run.run_id, caller, store.definitions, store.runs)).cancel_requested
foreign_trace: Final = await get_workflow_steps(
other.run.run_id, _CALLER.model_copy(update={"team_id": "other"}), store.definitions, store.runs,
)
assert not foreign_trace.cancel_requested
async def test_submit_drains_pinned_snapshot_serializes_chat_and_never_cold_replays(store: Storage) -> None:
snapshot: Final = await store.snapshot()
admitted: Final = await store.admit(snapshot=snapshot, key="original")
orphan: Final = await store.admit(snapshot=snapshot)
assert isinstance(admitted, RunAdmission) and isinstance(orphan, RunAdmission)
draft: Final = await store.definitions.get_draft(snapshot.workflow_id, "team")
assert draft is not None
changed: Final = snapshot.bundle.model_copy(update={
"nodes": tuple(node.model_copy(update={"model": "edited"}) for node in snapshot.bundle.nodes),
})
replaced: Final = await store.definitions.replace(snapshot.workflow_id, "team", draft.draft.etag, changed, "author")
assert isinstance(replaced, WorkflowDraft)
republished: Final = await store.definitions.publish(
snapshot.workflow_id, "team", replaced.etag, replaced.digest, "admin",
)
assert isinstance(republished, WorkflowSnapshot)
chat: Final = ChatReplies(store, (
ExecutionError("admission_rate_limit", "retry"),
ChildCallResult(list(_MESSAGES), 4, 2, "first", 0.125),
ChildCallResult({"answer": False}, 6, 3, "last", 0.25),
))
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
try:
assert await runner.submit(admitted.run, _REQUEST, _CALLER) is None
finally:
await runner.close()
completed: Final = await store.runs.get(admitted.run.run_id)
assert completed is not None and completed.status is RunState.completed, completed
completion: Final = ChatCompletion.model_validate(completed.output)
assert json.loads(completion.choices[0].message.content or "null") == {"answer": False}
assert completion.usage is not None
usage: Final = completion.usage
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert chat.models == ("gateway-model",) * 3
assert tuple((step.node_id, step.attempts) for step in chat.calls) == (("node_1", 1), ("node_1", 2), ("node_2", 1))
assert chat.calls[0].step_id == chat.calls[1].step_id and chat.calls[0].deadline == chat.calls[1].deadline
assert len({step.call_id for step in chat.calls}) == 3
assert tuple((step.state, step.call_id) for step in chat.persisted) == tuple(
(RunState.running, step.call_id) for step in chat.calls
)
assert chat.calls[1].input == chat.calls[0].input == chat.calls[2].input == list(_MESSAGES)
pinned: Final = await get_run_definition(completed.run_id, _CALLER, store.definitions, store.runs)
assert pinned.bundle == snapshot.bundle
restarted: Final = WorkflowRunner(store.runs, store.definitions, chat)
await restarted.close()
pending: Final = await store.runs.get(orphan.run.run_id)
assert pending is not None and pending.status is RunState.pending and len(chat.calls) == 3
assert await store.runs.steps(orphan.run.run_id) == ()

View file

@ -1,6 +1,7 @@
import asyncio
import hashlib
import os
import sys
import time
from collections.abc import AsyncIterator
from dataclasses import replace
@ -13,6 +14,7 @@ import pytest_asyncio
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel, ConfigDict, TypeAdapter
from litellm.proxy.workflows.monty._compat import timeout
from litellm.proxy.workflows.monty.compiler import compile_program
from litellm.proxy.workflows.monty.models import (
CompiledProgram,
@ -24,6 +26,7 @@ from litellm.proxy.workflows.monty.models import (
PythonProgram,
)
from litellm.proxy.workflows.monty.runtime import BinaryArtifact, MontyRuntime, installed_binary
from litellm.proxy.workflows.monty.validation import ValidatorProcess
from tests.integration._support.process import stop_processes
pytest.importorskip("pydantic_monty")
@ -426,6 +429,45 @@ def _slow_schema(depth: int) -> JsonSchema:
return {**definitions, f"N{depth}": {"type": "string"}}
@pytest.mark.skipif(sys.platform != "linux", reason="address-space limits require Linux prlimit")
@pytest.mark.asyncio
async def test_standalone_validator_bounds_child_memory_without_limiting_gateway() -> None:
import resource
parent_limit: Final = resource.getrlimit(resource.RLIMIT_AS)
validator: Final = ValidatorProcess()
try:
assert await validator.start() is None
assert validator.process is not None
soft, hard = resource.prlimit(validator.process.pid, resource.RLIMIT_AS)
assert 0 < soft <= hard <= MontyLimits().worker_address_space_bytes
assert resource.getrlimit(resource.RLIMIT_AS) == parent_limit
assert await validator.validate("valid", {"type": "string"}, deadline=time.monotonic() + 5) is None
finally:
await validator.close()
@pytest.mark.skipif(sys.platform != "linux", reason="address-space limits require Linux prlimit")
@pytest.mark.asyncio
@pytest.mark.parametrize("output", [False, True])
async def test_validator_memory_exhaustion_fails_closed_and_recovers(output: bool) -> None:
instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1, worker_address_space_bytes=64 * 1024 * 1024))
assert await instance.start() is None
ref: Final[JsonSchema] = {"$ref": "#/$defs/N0"}
compiled: Final = _compile(
"def transform(value): return value", ref if output else {}, {} if output else ref, _slow_schema(22)
)
try:
result: Final = await instance.run_program(compiled, {"value": None}, deadline_monotonic=time.monotonic() + 15)
assert isinstance(result, MontyFailure), result
assert result.kind == MontyErrorKind.WORKER_CRASHED, result
recovered: Final = await _run(instance, _compile("def transform(value): return value"), "recovered")
assert isinstance(recovered, MontySuccess), recovered
assert recovered.output == "recovered"
finally:
await instance.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("output", [False, True])
async def test_expensive_allowed_validation_times_out_without_blocking_and_recovers(
@ -546,7 +588,7 @@ def _worker_pids() -> frozenset[int]:
async def _active_worker(previous: frozenset[int]) -> int:
async with asyncio.timeout(5):
async with timeout(5):
while True:
current: Final = _worker_pids() - previous
if current:
@ -599,7 +641,7 @@ async def test_fresh_worker_has_no_parent_canary_and_has_address_space_limit(
task: Final = asyncio.create_task(_run(runtime, _compile("def transform(value):\n while True: pass"), None))
pid: Final = await _active_worker(previous)
assert b"LITEFLOW_ENV_CANARY" not in Path(f"/proc/{pid}/environ").read_bytes()
async with asyncio.timeout(5):
async with timeout(5):
while str(runtime.limits.worker_address_space_bytes) not in Path(f"/proc/{pid}/limits").read_text():
await asyncio.sleep(0.01)
assert str(runtime.limits.worker_address_space_bytes) in Path(f"/proc/{pid}/limits").read_text()

View file

@ -0,0 +1,51 @@
from typing import Final
import pytest
from pydantic import ValidationError
from litellm.models.workflow_definition import WorkflowBundle
from tests.unit.proxy.workflow_definition_rows import bundle
@pytest.mark.parametrize("node_type", ["parallel", "if", "shell", "http"])
def test_v1_contract_rejects_nodes_outside_sequential_chat_and_python(node_type: str) -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
unsupported: Final = {**original, "nodes": [{**original["nodes"][0], "type": node_type}]}
with pytest.raises(ValidationError) as rejected:
WorkflowBundle.model_validate(unsupported)
assert rejected.value.errors(include_input=False)[0]["type"] == "union_tag_invalid"
def test_binding_cannot_supply_both_literal_and_reference() -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
ambiguous: Final = {
**original,
"nodes": [
{
**original["nodes"][0],
"messages": {"from": "$input", "output": "messages", "value": [{"role": "user", "content": "hello"}]},
}
],
}
with pytest.raises(ValidationError) as rejected:
WorkflowBundle.model_validate(ambiguous)
assert {error["loc"][-1] for error in rejected.value.errors(include_input=False) if "messages" in error["loc"]} == {
"from",
"output",
"value",
}
def test_nonfinite_values_are_rejected_before_digest_serialization() -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
with pytest.raises(ValidationError) as rejected:
WorkflowBundle.model_validate({**original, "input_schema": {"minimum": float("nan")}})
assert rejected.value.errors(include_input=False)[0]["type"] == "finite_number"
@pytest.mark.parametrize("field", ["api_key", "api_base", "base_url", "headers", "callbacks", "user_config"])
def test_definition_cannot_supply_provider_credentials_or_runtime_overrides(field: str) -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
with pytest.raises(ValidationError) as rejected:
WorkflowBundle.model_validate({**original, "nodes": [{**original["nodes"][0], field: "untrusted"}]})
assert rejected.value.errors(include_input=False)[0]["type"] == "extra_forbidden"

View file

@ -5130,3 +5130,69 @@ def test_member_can_reach_skill_history_visibility_checks(route: str) -> None:
)
except HTTPException as exc:
pytest.fail(f"History visibility could not run: {exc.detail}")
@pytest.mark.parametrize("route", [
"/v1/workflows",
"/v1/workflows/definition-id",
"/v1/workflows/definition-id/validate",
"/v1/workflows/definition-id/publish",
])
def test_definition_management_routes_reach_team_authorization(route: str) -> None:
token: Final = UserAPIKeyAuth(
user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER,
allowed_routes=["management_routes"],
)
request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""})
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={}
)
assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True
@pytest.mark.parametrize("route", [
"/v1/workflows/runs",
"/v1/workflows/runs/validate",
"/v1/workflows/runs/publish",
"/v1/workflows/definition-id/delete",
"/v1/workflows/definition-id/test",
])
def test_definition_route_grants_do_not_grant_run_or_future_actions(route: str) -> None:
token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_definition_routes"])
with pytest.raises(HTTPException) as rejected:
RouteChecks.is_virtual_key_allowed_to_call_route(route, token)
assert rejected.value.status_code == 403
@pytest.mark.parametrize("route", [
"/v1/workflows/definition-id/test",
"/v1/workflows/definition-id/runs",
"/v1/workflows/runs/run-id/steps",
"/v1/workflows/runs/run-id/definition",
"/v1/workflows/runs/run-id/cancel",
])
def test_execution_routes_reach_their_scoped_authorization(route: str) -> None:
token: Final = UserAPIKeyAuth(
user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER,
allowed_routes=["management_routes"],
)
request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""})
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={}
)
assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True
@pytest.mark.parametrize("route", [
"/v1/workflows/runs",
"/v1/workflows/runs/run-id",
"/v1/workflows/runs/run-id/events",
"/v1/workflows/runs/run-id/messages",
])
def test_execution_route_grants_preserve_legacy_tracking_admission(route: str) -> None:
token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_execution_routes"])
with pytest.raises(HTTPException) as rejected:
RouteChecks.is_virtual_key_allowed_to_call_route(route, token)
assert rejected.value.status_code == 403
assert RouteChecks.check_route_access(route, LiteLLMRoutes.self_managed_routes.value) is False
assert RouteChecks.check_route_access(route, LiteLLMRoutes.management_routes.value) is False

View file

@ -7,7 +7,7 @@ import sys
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager, nullcontext
from datetime import datetime, timedelta, timezone
from datetime import datetime, timedelta, timezone, tzinfo
from functools import partial
from io import BytesIO
from pathlib import Path
@ -69,7 +69,7 @@ from litellm.proxy.auth.auth_checks import (
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model, safe_set_request_parsed_body
from litellm.proxy.auth.user_api_key_auth import (
_check_key_model_budget_with_fallback,
_check_team_member_model_budget_with_fallback,
@ -81,6 +81,7 @@ from litellm.proxy.auth.user_api_key_auth import (
_reserve_budget_after_common_checks,
_route_requires_auth_despite_public,
_routing_selector_matches_claim,
authorize_internal_virtual_key,
run_centralized_common_checks,
_run_post_custom_auth_checks,
user_api_key_auth_builder,
@ -101,6 +102,118 @@ class _RoutingRequest:
self.state = SimpleNamespace()
async def _internal_key_cache(
monkeypatch: pytest.MonkeyPatch, *, models: tuple[str, ...] = (), max_budget: float | None = None,
key_hash: str = "b" * 64,
) -> UserApiKeyCache:
cache: Final = UserApiKeyCache()
await cache_key_object(
hashed_token=key_hash,
user_api_key_obj=UserAPIKeyAuth(
token=key_hash, team_id="original-team", models=list(models), spend=2.0, max_budget=max_budget,
),
user_api_key_cache=cache,
proxy_logging_obj=None,
)
await cache.async_set_cache(
key="team_id:original-team", value=LiteLLM_TeamTableCachedObj(team_id="original-team")
)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", object())
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {})
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
counters: Final = DualCache()
counters.in_memory_cache.set_cache(key=f"spend:key:{key_hash}", value=2.0)
monkeypatch.setattr(litellm.proxy.proxy_server, "spend_counter_cache", counters)
monkeypatch.setattr(litellm, "max_budget", 0.0)
return cache
def _internal_management_request() -> Request:
return Request({
"type": "http", "path": "/workflow/runs", "method": "POST", "headers": (), "query_string": b"",
})
@pytest.mark.asyncio
async def test_internal_key_cannot_execute_for_a_different_team(monkeypatch: pytest.MonkeyPatch) -> None:
await _internal_key_cache(monkeypatch)
request: Final = _internal_management_request()
with pytest.raises(HTTPException, match="no longer belongs") as error:
await authorize_internal_virtual_key(
"b" * 64, request, {"model": "workflow-model"},
route="/v1/chat/completions", expected_team_id="another-team",
)
assert error.value.status_code == 403
assert getattr(request.state, "budget_reservation", None) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("restriction", ["model", "budget"])
async def test_internal_execution_route_enforces_inference_model_and_key_budget(
monkeypatch: pytest.MonkeyPatch, restriction: str,
) -> None:
await _internal_key_cache(
monkeypatch,
models=("allowed-model",) if restriction == "model" else (),
max_budget=1.0 if restriction == "budget" else None,
)
with pytest.raises(litellm.BudgetExceededError if restriction == "budget" else ProxyException) as error:
await authorize_internal_virtual_key(
"b" * 64, _internal_management_request(), {"model": "workflow-model"},
route="/v1/chat/completions", expected_team_id="original-team",
)
if restriction == "budget":
assert isinstance(error.value, litellm.BudgetExceededError)
assert error.value.current_cost == 2.0 and error.value.max_budget == 1.0
else:
assert isinstance(error.value, ProxyException)
assert error.value.type == ProxyErrorTypes.key_model_access_denied
@pytest.mark.asyncio
async def test_internal_execution_route_keeps_proxy_network_policy(monkeypatch: pytest.MonkeyPatch) -> None:
await _internal_key_cache(monkeypatch)
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"allowed_ips": ["192.0.2.1"]})
with pytest.raises(HTTPException, match="not allowed") as error:
await authorize_internal_virtual_key(
"b" * 64, _internal_management_request(), {"model": "workflow-model"},
route="/v1/chat/completions", expected_team_id="original-team",
)
assert error.value.status_code == 403
@pytest.mark.asyncio
async def test_workflow_wrapper_with_model_grant_passes_strict_pricing_but_unpriced_child_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
raw_key: Final = "sk-workflow-pricing-fixture"
key_hash: Final = proxy_utils.hash_token(raw_key)
await _internal_key_cache(monkeypatch, models=("workflow/example", "unpriced-child"), key_hash=key_hash)
router: Final = litellm.Router(
model_list=[{
"model_name": "unpriced-child",
"litellm_params": {"model": "openai/workflow-unpriced-fixture", "api_key": "test-only"},
}],
)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True)
request: Final = Request({
"type": "http", "path": "/v1/chat/completions", "method": "POST",
"headers": ((b"authorization", f"Bearer {raw_key}".encode()),), "query_string": b"",
})
safe_set_request_parsed_body(request, {"model": "workflow/example", "messages": [{"role": "user", "content": "hello"}]})
outer: Final = await user_api_key_auth(request=request, api_key=f"Bearer {raw_key}")
assert outer.team_id == "original-team"
with pytest.raises(ProxyException) as error:
await authorize_internal_virtual_key(
key_hash, _internal_management_request(), {"model": "unpriced-child"},
route="/v1/chat/completions", expected_team_id="original-team",
)
assert error.value.type == ProxyErrorTypes.model_cost_map_missing
async def _real_model_budget_limiter(
*,
member_spend: Mapping[str, float] | None = None,
@ -466,10 +579,17 @@ async def test_custom_auth_does_not_enforce_key_model_access_by_default():
@pytest.mark.asyncio
async def test_post_custom_auth_expired_key_returns_unauthorized():
expired_token = UserAPIKeyAuth(
async def test_post_custom_auth_expired_key_returns_unauthorized(monkeypatch: pytest.MonkeyPatch) -> None:
class Clock(datetime):
@classmethod
def now(cls, tz: tzinfo | None = None) -> datetime:
current: Final = datetime(2026, 1, 1, 12, tzinfo=timezone.utc)
return current.astimezone(tz) if tz is not None else current.replace(tzinfo=None)
monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.datetime", Clock)
expired_token: Final = UserAPIKeyAuth(
token="test_token",
expires=datetime.now() - timedelta(minutes=1),
expires=Clock(2026, 1, 1, 11, 59, tzinfo=timezone.utc),
)
with pytest.raises(ProxyException) as exc_info:

View file

@ -0,0 +1,157 @@
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from liteflow import ChatCompletion, Draft, GatewayClient, Input, Workflow, bundle_digest, compile_workflow
from litellm.models.workflow_definition import DefinitionPage
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import get_workflow_definition_repository
from litellm.proxy.management_endpoints.workflow_definition_endpoints import router
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
def _client(database: DefinitionRows | None, identity: UserAPIKeyAuth) -> TestClient:
app: Final = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: identity
app.dependency_overrides[get_workflow_definition_repository] = lambda: (
None if database is None else WorkflowDefinitionRepository(database.database)
)
return TestClient(app)
def _sdk(client: TestClient) -> GatewayClient:
def send(request: httpx.Request) -> httpx.Response:
response: Final = client.request(
request.method, str(request.url), content=request.content, headers=dict(request.headers)
)
return httpx.Response(response.status_code, headers=dict(response.headers), content=response.content)
return GatewayClient(
base_url="http://testserver",
api_key="synthetic",
team_id="team",
transport=httpx.MockTransport(send),
)
def _body(name: str = "example", model: str = "gateway-model") -> Mapping[str, object]:
return {"bundle": bundle(name, model).model_dump(mode="json", by_alias=True, exclude_unset=True)}
def _register(client: TestClient, name: str = "example") -> Draft:
response: Final = client.post("/v1/workflows", json=_body(name))
assert response.status_code == 201, response.text
return Draft.model_validate(response.json())
def test_real_sdk_registers_validates_replaces_and_publishes_with_an_independent_admin() -> None:
database: Final = DefinitionRows()
author_http: Final = _client(database, caller("reviewer"))
publisher_http: Final = _client(
database, UserAPIKeyAuth(user_id="publisher", user_role=LitellmUserRoles.PROXY_ADMIN, team_id="team")
)
workflow: Final = Workflow(name="example", nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)])
with _sdk(author_http) as author, _sdk(publisher_http) as publisher:
draft: Final = author.register(workflow)
assert draft.bundle == compile_workflow(workflow) and draft.digest == bundle_digest(draft.bundle)
assert author.get(draft.workflow_id) == draft
assert author_http.get(f"/v1/workflows/{draft.workflow_id}").headers["etag"] == draft.etag
assert author.validate(draft) == {
"valid": True,
"workflow_id": draft.workflow_id,
"etag": draft.etag,
"digest": draft.digest,
}
with pytest.raises(httpx.HTTPStatusError) as denied:
author.publish(draft)
assert denied.value.response.status_code == 403
replacement: Final = author.register(workflow, draft=draft)
assert replacement.etag != draft.etag
assert publisher.publish(replacement) == {
"workflow_id": draft.workflow_id,
"snapshot_id": author.get(draft.workflow_id).published_snapshot_id,
"model": "workflow/example",
"digest": replacement.digest,
}
def test_stale_or_missing_revision_and_wrong_digest_cannot_publish_or_replace() -> None:
database: Final = DefinitionRows()
client: Final = _client(database, caller("reviewer"))
draft: Final = _register(client)
path: Final = f"/v1/workflows/{draft.workflow_id}"
changed: Final = client.put(path, json=_body(model="other-model"), headers={"If-Match": draft.etag})
assert changed.status_code == 200, changed.text
assert client.put(path, json=_body()).status_code == 412
assert client.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 412
assert client.post(f"{path}/validate").status_code == 412
assert (
client.post(f"{path}/publish", json={"digest": draft.digest}, headers={"If-Match": draft.etag}).status_code
== 412
)
assert (
client.post(
f"{path}/publish", json={"digest": "0" * 64}, headers={"If-Match": changed.json()["etag"]}
).status_code
== 412
)
assert client.get(path).json() == changed.json()
@pytest.mark.parametrize("body, status", [(b"", 200), (b"{}", 200), (b"null", 200), (b'{"execute":true}', 422)])
def test_validate_accepts_only_an_optional_empty_action(body: bytes, status: int) -> None:
client: Final = _client(DefinitionRows(), caller())
draft: Final = _register(client)
response: Final = client.post(
f"/v1/workflows/{draft.workflow_id}/validate",
content=body,
headers={"If-Match": draft.etag, "Content-Type": "application/json"},
)
assert response.status_code == status, response.text
@pytest.mark.parametrize("user, team", [("outsider", "other-team"), ("not-a-team-member", "team")])
def test_inaccessible_draft_cannot_be_read_edited_or_published(user: str, team: str) -> None:
database: Final = DefinitionRows()
draft: Final = _register(_client(database, caller()))
client: Final = _client(database, caller(user, team))
path: Final = f"/v1/workflows/{draft.workflow_id}"
assert client.get(path).status_code == 404
assert client.get(path, params={"team_id": "team"}).status_code == 404
assert client.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 404
assert (
client.post(f"{path}/publish", json={"digest": draft.digest}, headers={"If-Match": draft.etag}).status_code
== 404
)
@pytest.mark.parametrize(
"identity, expected",
[
(caller(None), 403),
(UserAPIKeyAuth(user_id="author", user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY), 403),
(caller(), 500),
],
)
def test_no_database_preserves_authorization_precedence(identity: UserAPIKeyAuth, expected: int) -> None:
response: Final = _client(None, identity).post("/v1/workflows", json=_body())
assert response.status_code == expected, response.text
def test_duplicate_registration_conflicts_and_cursor_pagination_returns_each_draft_once() -> None:
client: Final = _client(DefinitionRows(), caller())
first: Final = _register(client, "first")
second: Final = _register(client, "second")
assert client.post("/v1/workflows", json=_body("first")).status_code == 409
assert client.get("/v1/workflows", params={"cursor": "malformed"}).status_code == 422
page: Final = DefinitionPage.model_validate(client.get("/v1/workflows", params={"limit": 1}).json())
tail: Final = DefinitionPage.model_validate(client.get("/v1/workflows", params={"cursor": page.next_cursor}).json())
assert {page.workflows[0].workflow_id, tail.workflows[0].workflow_id} == {first.workflow_id, second.workflow_id}
assert page.workflows[0].workflow_id != tail.workflows[0].workflow_id and tail.next_cursor is None

View file

@ -0,0 +1,35 @@
from typing import Final
import pytest
from fastapi import HTTPException
from pydantic import ValidationError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
TestRunRequest as PaidRunRequest, WorkflowInvocation, _execution_key,
)
@pytest.mark.parametrize("overrides", [
{"is_session_token": True}, {"via_virtual_key": False}, {"team_id": "other"}, {"token": None},
])
def test_execution_requires_a_team_bound_application_key(overrides: dict[str, object]) -> None:
owner: Final = UserAPIKeyAuth(token="owner", team_id="team").model_copy(update={"via_virtual_key": True})
assert _execution_key(owner, "team") == "owner"
with pytest.raises(HTTPException) as denied:
_execution_key(owner.model_copy(update=overrides), "team")
assert denied.value.status_code == 403
@pytest.mark.parametrize("extra", [{"stream": True}, {"tools": []}, {"temperature": 0}])
def test_workflow_invocation_rejects_unsupported_chat_parameters(extra: dict[str, object]) -> None:
with pytest.raises(ValidationError):
WorkflowInvocation.model_validate({
"model": "workflow/example", "messages": [{"role": "user", "content": "go"}], **extra,
})
@pytest.mark.parametrize("consent", ["true", 1, None])
def test_paid_test_requires_explicit_boolean_billing_consent(consent: object) -> None:
with pytest.raises(ValidationError):
PaidRunRequest.model_validate({"messages": [{"role": "user", "content": "go"}], "authorize_billing": consent})

View file

@ -4,8 +4,9 @@ Uses FastAPI TestClient with a mocked prisma_client.
"""
import asyncio
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -113,8 +114,15 @@ def _make_prisma_client() -> MagicMock:
def _make_app() -> FastAPI:
app = FastAPI()
from litellm.proxy.deps.repositories import (
get_workflow_definition_repository,
get_workflow_execution_repository,
)
app: Final = FastAPI()
app.include_router(router)
app.dependency_overrides[get_workflow_definition_repository] = lambda: None
app.dependency_overrides[get_workflow_execution_repository] = lambda: None
return app
@ -223,6 +231,7 @@ class TestListWorkflowRuns:
assert resp.status_code == 200
data = resp.json()
assert data["count"] == 1
assert self._prisma.db.litellm_workflowrun.find_many.call_args.kwargs["where"]["engine_owned"] is False
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_filters_by_status(self, mock_pc):
@ -734,3 +743,54 @@ class TestAdminViewerReadParity:
with pytest.raises(HTTPException) as exc_info:
asyncio.run(_require_run(prisma, "run-1", _override_auth_admin_viewer()))
assert exc_info.value.status_code == 404
@pytest.mark.parametrize(
("method", "suffix", "body"),
[
("PATCH", "", {"status": "completed", "output": {"forged": True}}),
("POST", "/events", {"event_type": "step.failed", "step_name": "node-1"}),
("POST", "/messages", {"role": "assistant", "content": "forged result"}),
],
)
def test_engine_owned_runs_reject_tracking_writes_even_for_proxy_admin(
method: str, suffix: str, body: Mapping[str, object],
) -> None:
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
database: Final = _make_prisma_client()
run: Final = _make_run()
run.engine_owned = True
database.db.litellm_workflowrun.find_unique = AsyncMock(return_value=run)
database.db.litellm_workflowrun.update = AsyncMock(return_value=run)
database.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
database.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
database.db.litellm_workflowmessage.create = AsyncMock(return_value=_make_message())
app: Final = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth_admin
with patch("litellm.proxy.proxy_server.prisma_client", database), TestClient(app) as client:
response: Final = client.request(method, f"/v1/workflows/runs/run-1{suffix}", json=body)
assert response.status_code == 403, response.text
assert response.json() == {"detail": "Gateway executions cannot be changed through tracking endpoints"}
database.db.litellm_workflowrun.update.assert_not_awaited()
database.db.litellm_workflowmessage.create.assert_not_awaited()
database.db.tx.assert_not_called()
@pytest.mark.parametrize("suffix", ["", "/events", "/messages"])
def test_engine_run_content_is_only_exposed_by_scoped_execution_routes(suffix: str) -> None:
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
database: Final = _make_prisma_client()
run: Final = _make_run()
run.engine_owned = True
database.db.litellm_workflowrun.find_unique = AsyncMock(return_value=run)
database.db.litellm_workflowevent.find_many = AsyncMock()
database.db.litellm_workflowmessage.find_many = AsyncMock()
app: Final = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth_admin
with patch("litellm.proxy.proxy_server.prisma_client", database), TestClient(app) as client:
response: Final = client.get(f"/v1/workflows/runs/run-1{suffix}")
assert response.status_code == 404
database.db.litellm_workflowevent.find_many.assert_not_awaited()
database.db.litellm_workflowmessage.find_many.assert_not_awaited()

View file

@ -0,0 +1,122 @@
import asyncio
from collections.abc import AsyncGenerator, Mapping, Sequence
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Final, cast # noqa: TID251 # the test database implements the delegates exercised by these routes
from liteflow import ChatCompletion, Input, Workflow, compile_workflow
from pydantic import TypeAdapter
from litellm.models.workflow_definition import WorkflowBundle
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.repositories.workflow_definition_repository import DefinitionDatabase
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_COMPARISON: Final = TypeAdapter(Mapping[str, str])
def bundle(name: str = "example", model: str = "gateway-model") -> WorkflowBundle:
return WorkflowBundle.model_validate(
compile_workflow(Workflow(name=name, nodes=[ChatCompletion(model=model, messages=Input.messages)]))
)
def caller(user_id: str | None = "author", team_id: str = "team") -> UserAPIKeyAuth:
return UserAPIKeyAuth(user_id=user_id, team_id=team_id, user_role=LitellmUserRoles.INTERNAL_USER)
def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool:
return all(
str(row.get(key)) > _COMPARISON.validate_python(value)["gt"]
if isinstance(value, Mapping) and "gt" in value
else row.get(key) == value
for key, value in where.items()
)
class Rows:
def __init__(self, rows: tuple[Mapping[str, object], ...] = ()) -> None:
self.rows = rows
async def find_first(
self,
skip: int | None = None,
where: Mapping[str, object] | None = None,
cursor: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> object | None:
return next((row for row in self.rows if _matches(row, where or {})), None)
async def find_unique(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> object | None:
return await self.find_first(where=where)
async def find_many(
self,
take: int | None = None,
skip: int | None = None,
where: Mapping[str, object] | None = None,
cursor: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> Sequence[object]:
matches: Final = tuple(row for row in self.rows if _matches(row, where or {}))
return sorted(matches, key=lambda row: str(row.get("workflow_id")))[:take]
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> object:
if "etag" in data and any(
row.get("name") == data["name"] and row.get("team_id") == data["team_id"] for row in self.rows
):
from prisma.errors import UniqueViolationError
raise UniqueViolationError({"is_panic": False, "message": "duplicate", "error_code": "P2002"})
row: Final = {
"created_at": NOW,
"published_etag": None,
"published_snapshot_id": None,
**data,
}
self.rows = (*self.rows, row)
return row
async def update(
self, data: Mapping[str, object], where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> object | None:
self.rows = tuple({**row, **data} if _matches(row, where) else row for row in self.rows)
return await self.find_first(where=where)
class DefinitionRows:
def __init__(self) -> None:
self.litellm_workflowdefinition = Rows()
self.litellm_workflowsnapshot = Rows()
self.litellm_teamtable = Rows(
(
{
"team_id": "team",
"members_with_roles": [
{"user_id": "author", "role": "user"},
{"user_id": "reviewer", "role": "admin"},
{"user_id": "other", "role": "user"},
],
},
{"team_id": "other-team", "members_with_roles": [{"user_id": "outsider", "role": "admin"}]},
)
)
self.lock = asyncio.Lock()
async def query_raw(self, query: str, *args: object) -> object:
return ()
@property
def database(self) -> DefinitionDatabase:
return cast(DefinitionDatabase, self) # cast-ok: isolated Prisma boundary implements the used table operations
@asynccontextmanager
async def tx(self) -> AsyncGenerator["DefinitionRows"]:
async with self.lock:
yield self

View file

@ -0,0 +1,118 @@
import hashlib
from collections.abc import Mapping
from typing import Final
import pytest
from liteflow import ChatCompletion, Input, Python, Workflow, compile_workflow
from pydantic import JsonValue
from litellm.models.workflow_definition import (
DefinitionError,
PythonNode,
WorkflowBundle,
WorkflowDraft,
WorkflowSnapshot,
)
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.workflows.definitions import authorize_team, freeze_draft, published_workflow, validate_bundle
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
def transform(text: str) -> str:
return str(1 // 0)
def _python_bundle(source: str) -> WorkflowBundle:
workflow: Final = Workflow(
name="example",
nodes=[
ChatCompletion(model="gateway-model", messages=Input.messages),
Python(fn=transform),
],
)
original: Final = WorkflowBundle.model_validate(compile_workflow(workflow))
node: Final = original.nodes[-1]
assert isinstance(node, PythonNode)
program: Final = node.program.model_copy(
update={"source": source, "source_sha256": hashlib.sha256(source.encode()).hexdigest()}
)
return original.model_copy(update={"nodes": (*original.nodes[:-1], node.model_copy(update={"program": program}))})
def test_validation_compiles_python_without_executing_its_body() -> None:
assert validate_bundle(_python_bundle("def transform(text: str) -> str:\n return str(1 // 0)\n")) is None
@pytest.mark.parametrize(
"source",
[
"import os\ndef transform(text: str) -> str:\n return text\n",
"def transform(text: str) -> str:\n return open(text).read()\n",
],
)
def test_validation_rejects_imports_and_host_capabilities(source: str) -> None:
result: Final = validate_bundle(_python_bundle(source))
assert isinstance(result, DefinitionError) and result.code == "invalid_bundle"
assert result.param == "nodes.node_2.program"
@pytest.mark.parametrize(
"schema",
[
{"type": "object", "properties": {"$ref": {"type": "string"}, "$id": {"type": "number"}}},
{"const": {"$ref": "literal"}},
{"enum": [{"$dynamicRef": "literal"}]},
],
)
def test_schema_keywords_inside_property_names_and_literals_are_data(schema: Mapping[str, JsonValue]) -> None:
assert validate_bundle(bundle().model_copy(update={"input_schema": schema})) is None
@pytest.mark.parametrize(
"schema",
[
{"properties": {"safe": {"$ref": "https://example.com/schema"}}},
{"patternProperties": {"^a": {"$id": "https://example.com/schema"}}},
{"$dynamicRef": "https://example.com/schema"},
{"$recursiveRef": "https://example.com/schema"},
],
)
def test_reference_scope_is_checked_only_at_schema_positions(schema: Mapping[str, JsonValue]) -> None:
assert validate_bundle(bundle().model_copy(update={"input_schema": schema})) == DefinitionError(
"invalid_bundle", "Schemas must use local references and bounded nesting", "bundle.input_schema"
)
def test_forward_reference_cannot_reach_a_provider() -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
forward: Final = WorkflowBundle.model_validate(
{**original, "nodes": [{**original["nodes"][0], "messages": {"from": "later", "output": "result"}}]}
)
assert validate_bundle(forward) == DefinitionError(
"invalid_bundle", "Bindings must reference input messages or an earlier node result", "nodes.node_1.messages"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
async def test_key_team_restriction_precedes_global_user_role(role: LitellmUserRoles) -> None:
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
identity: Final = caller("reviewer").model_copy(update={"user_role": role, "via_virtual_key": True})
assert await authorize_team(store, identity, "other-team", "read") == DefinitionError(
"not_found", "Workflow not found"
)
session: Final = identity.model_copy(update={"is_session_token": True})
assert not isinstance(await authorize_team(store, session, "other-team", "read"), DefinitionError)
@pytest.mark.asyncio
async def test_application_identity_can_invoke_but_cannot_freeze_a_draft() -> None:
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
draft: Final = await store.create(bundle(), "team", "author")
assert isinstance(draft, WorkflowDraft)
snapshot: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
assert isinstance(snapshot, WorkflowSnapshot)
assert await published_workflow(store, caller(None), "example") == snapshot
denied: Final = await freeze_draft(store, caller(None), draft.workflow_id, draft.etag)
assert isinstance(denied, DefinitionError) and denied.code == "not_found"

View file

@ -0,0 +1,443 @@
import asyncio
from collections.abc import AsyncIterator, Mapping
from datetime import datetime, timedelta, timezone, tzinfo
from types import SimpleNamespace
from typing import Final
import pytest
from fastapi import HTTPException, Request
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.models.workflow_definition import ChatCompletionNode
from litellm.models.workflow_execution import (
ChildCallResult,
EngineRun,
ExecutionError,
RunInput,
RunState,
StepTrace,
workflow_call_id,
)
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import cache_key_object
from litellm.proxy.auth.user_api_key_auth import authorize_internal_virtual_key
from litellm.proxy.common_request_processing import process_workflow_request
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.deps.database import get_database
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3, get_request_stash
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.workflows import gateway
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails
from litellm.types.utils import ModelResponse
_NOW: Final = datetime(2026, 11, 6, tzinfo=timezone.utc)
_KEY: Final = "a" * 64
_FIELDS: Final = TypeAdapter(dict[str, object])
_MESSAGES: Final = [{"role": "user", "content": "private input"}]
class _Clock(datetime):
@classmethod
def now(cls, tz: tzinfo | None = None) -> datetime:
return _NOW
def _response(content: str) -> ModelResponse:
return ModelResponse(
model="workflow-model",
choices=[{"message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
)
def _request() -> Request:
return Request(
{
"type": "http",
"app": proxy_server.app,
"http_version": "1.1",
"method": "POST",
"scheme": "https",
"path": "/v1/chat/completions",
"query_string": b"",
"server": ("gateway.example", 443),
"client": ("192.0.2.10", 1234),
"headers": [
(b"x-litellm-trace-id", b"trace-root"),
(b"x-litellm-session-id", b"caller-session"),
(b"x-forwarded-for", b"198.51.100.8"),
(b"content-type", b"application/json"),
(b"authorization", b"Bearer caller-secret"),
],
"state": {"parent": "untouched", "database": None},
}
)
def _run() -> EngineRun:
return EngineRun(
run_id="run",
session_id="trace-root",
workflow_id="review",
snapshot_id="snapshot",
team_id="team",
created_by=_KEY,
status=RunState.running,
created_at=_NOW,
updated_at=_NOW,
request_digest="digest",
input=RunInput(messages=tuple(_MESSAGES), deadline=_NOW + timedelta(minutes=2), draft=False),
)
def _node() -> ChatCompletionNode:
return ChatCompletionNode.model_validate(
{
"id": "answer",
"type": "chat_completion",
"model": "workflow-model",
"messages": {"from": "$input", "output": "messages"},
"timeout_seconds": 60,
"max_tokens": 3,
"retry": {"max_attempts": 1, "backoff_seconds": 0},
}
)
def _step(attempt: int = 1) -> StepTrace:
return StepTrace(
step_id=f"step-{attempt}",
run_id="run",
node_id="answer",
state=RunState.running,
input=_MESSAGES,
attempts=attempt,
deadline=_NOW + timedelta(minutes=1),
started_at=_NOW,
call_id=workflow_call_id("run", "answer", attempt),
)
class _Observer(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.outcome = "success"
self.failure_status = 500
self.entered: Final = asyncio.Event()
self.failure_entered: Final = asyncio.Event()
self.completed: Final[asyncio.Queue[str]] = asyncio.Queue()
self.admissions: Final[asyncio.Queue[UserAPIKeyAuth]] = asyncio.Queue()
self.rows: Final[asyncio.Queue[Mapping[str, object]]] = asyncio.Queue()
self.requests: Final[asyncio.Queue[Mapping[str, object]]] = asyncio.Queue()
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: Mapping[str, object], call_type: str
) -> Mapping[str, object] | None:
self.admissions.put_nowait(user_api_key_dict)
self.requests.put_nowait(_FIELDS.validate_python(data["proxy_server_request"]))
if self.outcome == "pre":
raise HTTPException(self.failure_status, "private processor failure")
async def async_post_call_success_hook(
self, data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, response: object
) -> object:
self.completed.put_nowait(str(data["model"]))
if data["model"] == "workflow-model":
self.entered.set()
if self.outcome == "post":
raise HTTPException(self.failure_status, "private post-call failure")
if self.outcome in {"cancel", "timeout"}:
await asyncio.Event().wait()
if self.outcome in {"reject", "failure_timeout"}:
raise HTTPException(422, "private guardrail detail")
return response
async def async_post_call_failure_hook(
self,
request_data: Mapping[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> None:
self.failure_entered.set()
if self.outcome == "failure_timeout":
await asyncio.Event().wait()
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.rows.put_nowait(
get_logging_payload(
kwargs=dict(kwargs),
response_obj=response_obj,
start_time=start_time,
end_time=end_time,
)
)
async def async_log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
await self.async_log_success_event(kwargs, response_obj, start_time, end_time)
class _WorkflowPolicy(CustomGuardrail):
def __init__(self, reject: bool = False) -> None:
super().__init__(
guardrail_name="workflow-policy",
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
)
self.reject: Final = reject
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: Mapping[str, object], call_type: str
) -> Mapping[str, object] | None:
return {**data, "messages": [{"role": "user", "content": "guarded input"}]}
async def async_post_call_success_hook(
self, data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, response: object
) -> object:
if self.reject:
raise HTTPException(422, "private Python output")
assert isinstance(response, ModelResponse)
return response.model_copy(update={"choices": _response("guarded Python output").choices})
@pytest.fixture
async def proxy(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[tuple[_Observer, UserAPIKeyAuth]]:
cache: Final = UserApiKeyCache()
key: Final = UserAPIKeyAuth(
token=_KEY,
team_id="team",
max_budget=100.0,
max_parallel_requests=1,
models=["workflow-model", "workflow/review"],
)
await cache_key_object(hashed_token=_KEY, user_api_key_obj=key, user_api_key_cache=cache, proxy_logging_obj=None)
await cache.async_set_cache(
key="team_id:team",
value=LiteLLM_TeamTableCachedObj(
team_id="team",
metadata={"require_trace_id": True},
),
)
observer: Final = _Observer()
logging: Final = ProxyLogging(user_api_key_cache=cache)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(logging.internal_usage_cache)
counters: Final = DualCache()
router: Final = litellm.Router(
model_list=[
{
"model_name": "workflow-model",
"litellm_params": {
"model": "openai/workflow-fixture",
"api_key": "unit-only",
"mock_response": _response("child answer"),
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
},
}
],
num_retries=0,
)
monkeypatch.setattr(gateway, "datetime", _Clock)
monkeypatch.setattr(litellm, "callbacks", [observer, limiter])
monkeypatch.setattr(litellm, "max_budget", 0.0)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "spend_counter_cache", counters)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace())
monkeypatch.setattr(proxy_server, "general_settings", {"allowed_ips": ["192.0.2.10"], "cancel_on_disconnect": True})
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig())
yield observer, key
get_policy_registry().clear()
get_attachment_registry().clear()
router.reset()
@pytest.mark.asyncio
@pytest.mark.parametrize("forwarded", [False, True])
async def test_child_preserves_trusted_context_and_rechecks_model_access(
proxy: tuple[_Observer, UserAPIKeyAuth], forwarded: bool
) -> None:
observer, key = proxy
proxy_server.general_settings.update(
{
"use_x_forwarded_for": forwarded,
"allowed_ips": ["198.51.100.8" if forwarded else "192.0.2.10"],
}
)
request: Final = _request()
child: Final = gateway._child_request(request, {}, workflow_call_id("run", "answer", 1), "trace-root")
assert await get_database(child) is None
assert "authorization" not in child.headers
result: Final = await gateway.execute_chat_node(request, _run(), _node(), _step())
assert isinstance(result, ChildCallResult), result
row: Final = await asyncio.wait_for(observer.rows.get(), 5)
seen: Final = observer.requests.get_nowait()
assert _FIELDS.validate_python(seen["headers"])["x-litellm-trace-id"] == "trace-root"
assert (row["session_id"], row["litellm_call_id"], row["api_key"]) == ("trace-root", _step().call_id, _KEY)
assert result.request_id is None
assert result.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
assert request.scope["state"] == {"parent": "untouched", "database": None}
assert "caller-secret" not in str(seen)
await cache_key_object(
hashed_token=_KEY,
user_api_key_obj=key.model_copy(update={"models": ["another-model"]}),
user_api_key_cache=proxy_server.user_api_key_cache,
proxy_logging_obj=None,
)
denied: Final = await gateway.execute_chat_node(request, _run(), _node(), _step(2))
assert isinstance(denied, ExecutionError), denied
assert denied.receipt is None and observer.requests.empty()
assert "private" not in denied.message
@pytest.mark.asyncio
@pytest.mark.parametrize(
("phase", "status", "code"),
[
("auth", 429, "admission_rate_limit"),
("auth", 500, "provider_error"),
("pre", 429, "rate_limit"),
("pre", 500, "provider_error"),
("post", 429, "rate_limit"),
("post", 500, "provider_error"),
],
)
async def test_only_auth_rate_limits_are_safe_to_retry(
proxy: tuple[_Observer, UserAPIKeyAuth], phase: str, status: int, code: str, monkeypatch: pytest.MonkeyPatch
) -> None:
observer, _ = proxy
observer.outcome = phase
observer.failure_status = status
async def reject_auth(
key_hash: str,
request: Request,
request_data: Mapping[str, object],
*,
route: str | None = None,
expected_team_id: str | None = None,
) -> UserAPIKeyAuth:
raise HTTPException(status, "private authentication failure")
if phase == "auth":
monkeypatch.setattr(gateway, "authorize_internal_virtual_key", reject_auth)
result: Final = await gateway.execute_chat_node(_request(), _run(), _node(), _step())
assert isinstance(result, ExecutionError), result
assert result.code == code and "private" not in result.message
assert observer.admissions.qsize() == (0 if phase == "auth" else 1)
if phase == "post":
assert result.receipt is not None
assert result.receipt.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
await asyncio.wait_for(observer.rows.get(), 5)
else:
assert result.receipt is None
if phase == "auth":
assert observer.requests.empty() and observer.rows.empty()
@pytest.mark.asyncio
@pytest.mark.parametrize("outcome", ["reject", "cancel", "timeout", "failure_timeout"])
async def test_billed_failure_and_cancellation_release_child_slot(
proxy: tuple[_Observer, UserAPIKeyAuth], outcome: str, monkeypatch: pytest.MonkeyPatch
) -> None:
observer, _ = proxy
observer.outcome = outcome
task: Final = asyncio.create_task(gateway.execute_chat_node(_request(), _run(), _node(), _step()))
await asyncio.wait_for((observer.failure_entered if outcome == "failure_timeout" else observer.entered).wait(), 5)
if outcome == "cancel":
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
else:
if outcome in {"timeout", "failure_timeout"}:
loop: Final = asyncio.get_running_loop()
original_time: Final = loop.time
monkeypatch.setattr(loop, "time", lambda: original_time() + 120)
failure: Final = await task
assert isinstance(failure, ExecutionError) and failure.receipt is not None, failure
assert failure.receipt.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
assert failure.message == (
"The node deadline expired"
if outcome in {"timeout", "failure_timeout"}
else "The gateway rejected this node"
)
await asyncio.wait_for(observer.rows.get(), 5)
admitted: Final = observer.admissions.get_nowait()
if outcome in {"cancel", "timeout", "failure_timeout"}:
assert admitted.budget_reservation is not None and admitted.budget_reservation["finalized"] is True
observer.outcome = "success"
retried: Final = await gateway.execute_chat_node(_request(), _run(), _node(), _step(2))
assert isinstance(retried, ChildCallResult), retried
@pytest.mark.asyncio
@pytest.mark.parametrize("detached", [False, True])
@pytest.mark.parametrize("reject", [False, True])
async def test_workflow_policy_guards_modified_input_and_final_python_output_without_wrapper_charge(
proxy: tuple[_Observer, UserAPIKeyAuth],
detached: bool,
reject: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
observer, key = proxy
policy: Final = _WorkflowPolicy(reject)
monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, policy])
get_policy_registry().add_policy("workflow-dlp", Policy(guardrails=PolicyGuardrails(add=["workflow-policy"])))
get_attachment_registry().add_attachment(PolicyAttachment(policy="workflow-dlp", models=["workflow/review"]))
parent_request: Final = _request()
parent_auth: Final = await authorize_internal_virtual_key(
key_hash=_KEY,
request=parent_request,
request_data={"model": "workflow-model", "messages": _MESSAGES, "max_tokens": 3},
expected_team_id="team",
)
assert parent_auth.budget_reservation is not None
inputs: Final[asyncio.Queue[JsonValue]] = asyncio.Queue()
async def execute(messages: JsonValue) -> ModelResponse | ExecutionError:
inputs.put_nowait(messages)
assert get_request_stash() is None
assert parent_auth.budget_reservation["finalized"] is True
assert proxy_server.spend_counter_cache.in_memory_cache.get_cache(key=f"spend:key:{_KEY}") == 0
child: Final = await gateway.execute_chat_node(
_request(),
_run(),
_node(),
_step().model_copy(update={"input": messages}),
)
assert isinstance(child, ChildCallResult), child
return _response("private Python output")
invocation: Final = process_workflow_request(
request=parent_request,
workflow_model="workflow/review",
caller=parent_auth,
messages=_MESSAGES,
execute=execute,
)
result: Final = await (asyncio.create_task(invocation) if detached else invocation)
assert inputs.get_nowait() == [{"role": "user", "content": "guarded input"}]
if reject:
assert isinstance(result, ExecutionError), result
assert "private" not in result.message
else:
assert isinstance(result, ModelResponse), result
assert result.choices[0].message.content == "guarded Python output"
row: Final = await asyncio.wait_for(observer.rows.get(), 5)
assert row["litellm_call_id"] == _step().call_id
assert observer.rows.empty()
assert observer.admissions.qsize() == 1
assert observer.completed.get_nowait() == "workflow-model" and observer.completed.empty()

View file

@ -0,0 +1,33 @@
from typing import Final
import pytest
from pydantic import JsonValue
from litellm.models.workflow_definition import ReferenceBinding, ValueBinding, WorkflowBundle, WorkflowRetry
from litellm.models.workflow_execution import ExecutionError, RunState, StepTrace
from litellm.proxy.workflows.runner import resolve_binding
from tests.unit.proxy.workflow_definition_rows import bundle
from tests.unit.repositories.test_workflow_execution_repository import NOW, pending_run
def two_nodes() -> WorkflowBundle:
first: Final = bundle().nodes[0].model_copy(update={"retry": WorkflowRetry(max_attempts=2)})
second: Final = first.model_copy(update={
"id": "node_2", "messages": ReferenceBinding(source=first.id, output="result"),
})
return bundle().model_copy(update={
"nodes": (first, second), "output": ReferenceBinding(source=second.id, output="result"),
})
@pytest.mark.parametrize("value", [None, False, 0, [], {}, [{"role": "user", "content": "next"}]])
def test_bindings_preserve_falsy_and_structured_values(value: JsonValue) -> None:
run: Final = pending_run()
step: Final = StepTrace(step_id="step", run_id=run.run_id, node_id="node", state=RunState.completed,
input=None, output=value, attempts=1, started_at=NOW, deadline=NOW)
assert resolve_binding(ValueBinding(value=value), run, ()) == value
assert resolve_binding(ReferenceBinding(source="node", output="result"), run, (step,)) == value
assert resolve_binding(ReferenceBinding(source="$input", output="messages"), run, ()) == list(run.input.messages)
assert resolve_binding(ReferenceBinding(source="missing", output="result"), run, (step,)) == ExecutionError(
"missing_input", "Node output is unavailable"
)

View file

@ -0,0 +1,33 @@
from typing import Final
import pytest
from liteflow import Draft, bundle_digest
from litellm.models.workflow_definition import DefinitionError, WorkflowSnapshot, parsed_bundle
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle
@pytest.mark.asyncio
async def test_publication_and_draft_tests_keep_immutable_snapshots_across_edits() -> None:
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
draft: Final = await store.create(bundle(), "team", "author")
assert isinstance(draft, Draft) and draft.digest == bundle_digest(draft.bundle)
published: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
frozen: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(published, WorkflowSnapshot) and isinstance(frozen, WorkflowSnapshot)
assert published.bundle == frozen.bundle == parsed_bundle(draft.bundle)
assert frozen.published_by is None
assert await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") == published
updated: Final = await store.replace(draft.workflow_id, "team", draft.etag, bundle(model="second"), "other")
assert isinstance(updated, Draft) and updated.published_snapshot_id == published.snapshot_id
assert await store.get_published("example", "team") == published
stale: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(stale, DefinitionError) and stale.code == "stale_draft"
second: Final = await store.publish(updated.workflow_id, "team", updated.etag, updated.digest, "reviewer")
assert isinstance(second, WorkflowSnapshot) and second.contributors == ("other",)
assert second.snapshot_id != published.snapshot_id and second.digest != published.digest
assert await store.get_snapshot(published.snapshot_id) == published
assert await store.get_snapshot(frozen.snapshot_id) == frozen
assert await store.get_snapshot(second.snapshot_id, "other-team") is None
assert await store.get_published("example", "team") == second

View file

@ -0,0 +1,34 @@
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from litellm.models.workflow_execution import EngineRun, RunInput, RunState, StepTrace
from litellm.repositories.workflow_execution_repository import _aggregate_attempts
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
def pending_run() -> EngineRun:
return EngineRun(
run_id="run", session_id="session", workflow_id="workflow", snapshot_id="snapshot", team_id="team",
created_by="owner", status=RunState.pending, created_at=NOW, updated_at=NOW, request_digest="digest",
input=RunInput(
messages=({"role": "user", "content": "plan"},), deadline=NOW + timedelta(minutes=5), draft=True,
),
)
@pytest.mark.parametrize("cost", [None, 0.0, 0.125])
def test_attempt_aggregation_keeps_unknown_cost_and_unique_receipts(cost: float | None) -> None:
first: Final = StepTrace(
step_id="step", run_id="run", node_id="node", state=RunState.failed, input=None, attempts=1,
deadline=NOW, started_at=NOW, prompt_tokens=4, completion_tokens=2, response_cost=cost,
request_ids=("fallback-a", "fallback-b"),
)
latest: Final = first.model_copy(update={"attempts": 2, "state": RunState.completed, "output": False,
"request_ids": ("fallback-b", "retry"), "response_cost": None})
assert _aggregate_attempts((latest, first)) == latest.model_copy(update={
"prompt_tokens": 8, "completion_tokens": 4, "response_cost": cost,
"request_ids": ("fallback-b", "retry", "fallback-a"),
})

View file

@ -1,11 +1,11 @@
"""Tests for scripts/check_openapi_docs.py, the scan behind scripts/openapi_docs_gate.py."""
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import pytest
import check_openapi_docs as checker
import pytest
_BODY: Final = {"content": {"application/json": {"schema": {"type": "object"}}}}
_RESPONSES: Final = {"200": {"content": {"application/json": {"schema": {"$ref": "#/components/schemas/Out"}}}}}
@ -20,15 +20,47 @@ def _post_with_media_types(request_media_type: str, response_media_type: str) ->
}
@pytest.mark.parametrize("request_type, response_type", [
("application/json", "application/json"), ("application/sdp", "application/sdp"),
("multipart/form-data", "application/json"),
])
@pytest.mark.parametrize(
"request_type, response_type",
[
("application/json", "application/json"),
("application/sdp", "application/sdp"),
("multipart/form-data", "application/json"),
],
)
def test_a_fully_documented_post_has_no_findings(request_type: str, response_type: str) -> None:
operation: Final = _post_with_media_types(request_type, response_type)
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
@pytest.mark.parametrize("status", ["200", "201", "202", "206", "299"])
def test_any_documented_success_response_satisfies_the_response_rule(status: str) -> None:
operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}}
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
@pytest.mark.parametrize("status", ["default", "100", "300", "400", "500", "2000"])
def test_only_success_statuses_can_supply_the_response_schema(status: str) -> None:
operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}}
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == (
checker.Violation("POST", "/v1/thing", "OA003"),
)
@pytest.mark.parametrize("status", ["201", "202"])
def test_success_without_a_schema_is_not_documented_by_an_error_response(status: str) -> None:
operation: Final[Mapping[str, object]] = {
**_DOCUMENTED,
"responses": {
status: {"content": {"application/json": {"schema": {}}}},
"422": _RESPONSES["200"],
},
}
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == (
checker.Violation("POST", "/v1/thing", "OA003"),
)
def test_a_schema_under_another_media_type_still_counts_as_undocumented() -> None:
operation: Final = _post_with_media_types("application/octet-stream", "audio/wav")
assert {v.code for v in checker.violations({"paths": {"/v1/thing": {"post": operation}}})} == {"OA001", "OA003"}
@ -44,8 +76,11 @@ def test_each_missing_piece_is_its_own_rule() -> None:
def test_only_the_missing_piece_is_reported(body: Mapping[str, object]) -> None:
no_body: Final = {**_DOCUMENTED, "requestBody": body}
blank_description: Final = {**_DOCUMENTED, "description": " "}
untyped_200: Final = {**_DOCUMENTED, "responses": {"200": {"content": {"application/json": {"schema": {}}}}}}
spec: Final = {
untyped_200: Final[Mapping[str, object]] = {
**_DOCUMENTED,
"responses": {"200": {"content": {"application/json": {"schema": {}}}}},
}
spec: Final[Mapping[str, object]] = {
"paths": {
"/a": {"post": no_body},
"/b": {"post": blank_description},
@ -65,11 +100,13 @@ def test_non_post_operations_are_not_gated() -> None:
def test_findings_are_sorted_by_path_so_counts_are_stable() -> None:
spec: Final = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}}
spec: Final[Mapping[str, object]] = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}}
assert [v.path for v in checker.violations(spec)] == ["/a", "/a", "/a", "/z", "/z", "/z"]
def test_main_reads_a_spec_file_and_prints_one_line_per_finding(tmp_path, capsys) -> None:
def test_main_reads_a_spec_file_and_prints_one_line_per_finding(
tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
spec_file: Final = tmp_path / "openapi.json"
spec_file.write_text('{"paths": {"/v1/thing": {"post": {"description": "x", "requestBody": {}}}}}')
assert checker.main([str(spec_file)]) == 0

File diff suppressed because it is too large Load diff

26
uv.lock generated
View file

@ -4996,6 +4996,28 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/38/a6/800800bfed7b1fb10fc3f3d557785c3854e80d3f7a9800d784b176a1fc2d/librt-0.15.0-cp315-cp315t-win_arm64.whl", hash = "sha256:84d244b00604d17df3fc7736c327892d6bba66181254aa4087be807b6c342bdc", size = 110700, upload-time = "2026-08-07T10:49:15.499Z" },
]
[[package]]
name = "liteflow"
version = "0.1.0a1"
source = { url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" }
dependencies = [
{ name = "httpx" },
{ name = "openai" },
{ name = "pydantic", version = "2.13.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" },
{ name = "pydantic", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" },
]
sdist = { hash = "sha256:c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072" }
[package.metadata]
requires-dist = [
{ name = "httpx", specifier = ">=0.27,<1" },
{ name = "openai", specifier = ">=1.99,<3" },
{ name = "pydantic", specifier = ">=2.8,<3" },
{ name = "pytest", marker = "extra == 'test'", specifier = ">=8,<9" },
{ name = "rich", marker = "extra == 'cli'", specifier = ">=13,<15" },
]
provides-extras = ["test", "cli"]
[[package]]
name = "litellm"
version = "1.109.0"
@ -5083,6 +5105,7 @@ proxy = [
{ name = "hiredis", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" },
{ name = "httpx2" },
{ name = "inquirerpy" },
{ name = "liteflow" },
{ name = "litellm-enterprise" },
{ name = "litellm-proxy-extras" },
{ name = "mcp" },
@ -5254,6 +5277,7 @@ proxy-dev = [
{ name = "backports-zstd", marker = "python_full_version < '3.14'" },
{ name = "httptools" },
{ name = "hypercorn" },
{ name = "liteflow" },
{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'" },
{ name = "opentelemetry-api" },
{ name = "opentelemetry-exporter-otlp" },
@ -5323,6 +5347,7 @@ requires-dist = [
{ name = "jsonschema", specifier = ">=4.0.0,<5.0" },
{ name = "keyring", marker = "extra == 'cli'", specifier = ">=25.6.0,<26.0" },
{ name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=4.7,<5.0" },
{ name = "liteflow", marker = "extra == 'proxy'", url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" },
{ name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" },
{ name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" },
{ name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" },
@ -5490,6 +5515,7 @@ proxy-dev = [
{ name = "backports-zstd", marker = "python_full_version < '3.14'", specifier = "==1.8.0" },
{ name = "httptools", specifier = ">=0.7.1,<0.8" },
{ name = "hypercorn", specifier = "==0.17.3" },
{ name = "liteflow", url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" },
{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'", url = "https://github.com/BerriAI/liteadmin-mcp/archive/31fb096904ae05684c0779110a93b135b5601182.tar.gz" },
{ name = "opentelemetry-api", specifier = "==1.45.1" },
{ name = "opentelemetry-exporter-otlp", specifier = "==1.45.1" },