mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat: add durable sequential workflow gateway execution
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
5085022580
commit
39c4fc5e20
28 changed files with 4163 additions and 68 deletions
|
|
@ -0,0 +1,87 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_WorkflowRun" ADD COLUMN "engine_owned" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowDefinition" (
|
||||
"workflow_id" TEXT NOT NULL,
|
||||
"team_id" TEXT NOT NULL,
|
||||
"name" TEXT NOT NULL,
|
||||
"etag" TEXT NOT NULL,
|
||||
"bundle_json" TEXT NOT NULL,
|
||||
"contributors" TEXT[],
|
||||
"published_snapshot_id" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowDefinition_pkey" PRIMARY KEY ("workflow_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowSnapshot" (
|
||||
"snapshot_id" TEXT NOT NULL,
|
||||
"workflow_id" TEXT NOT NULL,
|
||||
"team_id" TEXT NOT NULL,
|
||||
"name" TEXT NOT NULL,
|
||||
"bundle_json" TEXT NOT NULL,
|
||||
"published_by" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"origin" TEXT NOT NULL DEFAULT 'publication',
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowSnapshot_pkey" PRIMARY KEY ("snapshot_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowExecution" (
|
||||
"run_id" TEXT NOT NULL,
|
||||
"snapshot_id" TEXT NOT NULL,
|
||||
"team_id" TEXT NOT NULL,
|
||||
"name" TEXT NOT NULL,
|
||||
"principal" TEXT NOT NULL,
|
||||
"dispatch_key" TEXT NOT NULL,
|
||||
"request_digest" TEXT NOT NULL,
|
||||
"request_json" TEXT NOT NULL,
|
||||
"state" TEXT NOT NULL DEFAULT 'pending',
|
||||
"deadline" TIMESTAMP(3) NOT NULL,
|
||||
"lease_owner" TEXT,
|
||||
"lease_until" TIMESTAMP(3),
|
||||
"cancel_requested" BOOLEAN NOT NULL DEFAULT false,
|
||||
"checkpoints" JSONB NOT NULL DEFAULT '[]',
|
||||
"response_json" TEXT,
|
||||
"error" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowExecution_pkey" PRIMARY KEY ("run_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowStep" (
|
||||
"step_id" TEXT NOT NULL,
|
||||
"run_id" TEXT NOT NULL,
|
||||
"node_id" TEXT NOT NULL,
|
||||
"state" TEXT NOT NULL DEFAULT 'running',
|
||||
"input" JSONB NOT NULL,
|
||||
"output" JSONB,
|
||||
"attempts" INTEGER NOT NULL DEFAULT 0,
|
||||
"deadline" TIMESTAMP(3) NOT NULL,
|
||||
"started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"finished_at" TIMESTAMP(3),
|
||||
"error" TEXT,
|
||||
"prompt_tokens" INTEGER NOT NULL DEFAULT 0,
|
||||
"completion_tokens" INTEGER NOT NULL DEFAULT 0,
|
||||
"response_cost" DOUBLE PRECISION,
|
||||
"request_id" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowStep_pkey" PRIMARY KEY ("step_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowDefinition_team_id_name_key" ON "LiteLLM_WorkflowDefinition"("team_id", "name");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowSnapshot_workflow_id_idx" ON "LiteLLM_WorkflowSnapshot"("workflow_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowExecution_dispatch_key_key" ON "LiteLLM_WorkflowExecution"("dispatch_key");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowExecution_team_id_name_idx" ON "LiteLLM_WorkflowExecution"("team_id", "name");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowStep_run_id_node_id_key" ON "LiteLLM_WorkflowStep"("run_id", "node_id");
|
||||
|
|
@ -1986,6 +1986,7 @@ model LiteLLM_WorkflowRun {
|
|||
input Json?
|
||||
output Json?
|
||||
metadata Json?
|
||||
engine_owned Boolean @default(false)
|
||||
|
||||
events LiteLLM_WorkflowEvent[]
|
||||
messages LiteLLM_WorkflowMessage[]
|
||||
|
|
@ -2003,6 +2004,72 @@ model LiteLLM_WorkflowRun {
|
|||
// step.failed → run.status = failed
|
||||
// hook.waiting → run.status = paused
|
||||
// hook.received → run.status = running
|
||||
model LiteLLM_WorkflowDefinition {
|
||||
workflow_id String @id
|
||||
team_id String
|
||||
name String
|
||||
etag String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_snapshot_id String?
|
||||
|
||||
@@unique([team_id, name])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowSnapshot {
|
||||
snapshot_id String @id
|
||||
workflow_id String
|
||||
team_id String
|
||||
name String
|
||||
bundle_json String
|
||||
published_by String
|
||||
created_at DateTime @default(now())
|
||||
origin String @default("publication")
|
||||
|
||||
@@index([workflow_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowExecution {
|
||||
run_id String @id
|
||||
snapshot_id String
|
||||
team_id String
|
||||
name String
|
||||
principal String
|
||||
dispatch_key String @unique
|
||||
request_digest String
|
||||
request_json String
|
||||
state String @default("pending")
|
||||
deadline DateTime
|
||||
lease_owner String?
|
||||
lease_until DateTime?
|
||||
cancel_requested Boolean @default(false)
|
||||
checkpoints Json @default("[]")
|
||||
response_json String?
|
||||
error String?
|
||||
|
||||
@@index([team_id, name])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowStep {
|
||||
step_id String @id
|
||||
run_id String
|
||||
node_id String
|
||||
state String @default("running")
|
||||
input Json
|
||||
output Json?
|
||||
attempts Int @default(0)
|
||||
deadline DateTime
|
||||
started_at DateTime @default(now())
|
||||
finished_at DateTime?
|
||||
error String?
|
||||
prompt_tokens Int @default(0)
|
||||
completion_tokens Int @default(0)
|
||||
response_cost Float?
|
||||
request_id String?
|
||||
|
||||
@@unique([run_id, node_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowEvent {
|
||||
event_id String @id @default(uuid())
|
||||
run_id String
|
||||
|
|
|
|||
164
litellm/models/workflow.py
Normal file
164
litellm/models/workflow.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Protocol
|
||||
|
||||
from pydantic import ConfigDict, JsonValue
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
class WorkflowRunState(str, Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class WorkflowSnapshotOrigin(str, Enum):
|
||||
PUBLICATION = "publication"
|
||||
DRAFT_TEST = "draft_test"
|
||||
|
||||
|
||||
class WorkflowDraft(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
workflow_id: str
|
||||
team_id: str
|
||||
name: str
|
||||
etag: str
|
||||
bundle_json: str
|
||||
contributors: tuple[str, ...]
|
||||
published_snapshot_id: str | None
|
||||
|
||||
|
||||
class WorkflowSnapshot(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
snapshot_id: str
|
||||
workflow_id: str
|
||||
team_id: str
|
||||
name: str
|
||||
bundle_json: str
|
||||
published_by: str
|
||||
created_at: datetime
|
||||
origin: WorkflowSnapshotOrigin = WorkflowSnapshotOrigin.PUBLICATION
|
||||
|
||||
|
||||
class WorkflowCheckpoint(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
node_path: str
|
||||
input: JsonValue
|
||||
output: JsonValue
|
||||
attempts: int
|
||||
|
||||
|
||||
class WorkflowStep(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
step_id: str
|
||||
run_id: str
|
||||
node_id: str
|
||||
state: WorkflowRunState
|
||||
input: JsonValue
|
||||
output: JsonValue | None = None
|
||||
attempts: int
|
||||
deadline: datetime
|
||||
started_at: datetime
|
||||
finished_at: datetime | None = None
|
||||
error: str | None = None
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
response_cost: float | None = None
|
||||
request_id: str | None = None
|
||||
|
||||
|
||||
class WorkflowExecution(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
run_id: str
|
||||
snapshot_id: str
|
||||
team_id: str
|
||||
name: str
|
||||
principal: str
|
||||
dispatch_key: str
|
||||
request_digest: str
|
||||
request_json: str
|
||||
state: WorkflowRunState
|
||||
deadline: datetime
|
||||
lease_owner: str | None = None
|
||||
lease_until: datetime | None = None
|
||||
cancel_requested: bool
|
||||
checkpoints: tuple[WorkflowCheckpoint, ...] = ()
|
||||
response_json: str | None = None
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class WorkflowConflict(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
message: str
|
||||
status_code: int = 409
|
||||
|
||||
|
||||
class WorkflowStore(Protocol):
|
||||
async def create(
|
||||
self, team_id: str, name: str, bundle_json: str, author: str
|
||||
) -> WorkflowDraft | WorkflowConflict: ...
|
||||
|
||||
async def get(self, team_id: str, name: str) -> WorkflowDraft | None: ...
|
||||
|
||||
async def get_id(self, workflow_id: str) -> WorkflowDraft | None: ...
|
||||
|
||||
async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]: ...
|
||||
|
||||
async def replace(
|
||||
self, team_id: str, name: str, etag: str, bundle_json: str, author: str
|
||||
) -> WorkflowDraft | WorkflowConflict: ...
|
||||
|
||||
async def publish(
|
||||
self, team_id: str, name: str, etag: str, publisher: str
|
||||
) -> WorkflowSnapshot | WorkflowConflict: ...
|
||||
|
||||
async def snapshot(self, snapshot_id: str) -> WorkflowSnapshot | None: ...
|
||||
|
||||
async def pin_draft(
|
||||
self, workflow_id: str, etag: str, requested_by: str
|
||||
) -> WorkflowSnapshot | WorkflowConflict: ...
|
||||
|
||||
async def published(self, team_id: str, name: str) -> WorkflowSnapshot | None: ...
|
||||
|
||||
async def dispatch(self, execution: WorkflowExecution) -> WorkflowExecution: ...
|
||||
|
||||
async def by_dispatch_key(self, dispatch_key: str) -> WorkflowExecution | None: ...
|
||||
|
||||
async def execution(self, run_id: str) -> WorkflowExecution | None: ...
|
||||
|
||||
async def lease_next(self, owner: str) -> WorkflowExecution | None: ...
|
||||
|
||||
async def renew(self, run_id: str, owner: str) -> bool: ...
|
||||
|
||||
async def steps(self, run_id: str) -> tuple[WorkflowStep, ...]: ...
|
||||
|
||||
async def start_step(
|
||||
self, run_id: str, node_id: str, input: JsonValue, deadline: datetime, owner: str
|
||||
) -> WorkflowStep | WorkflowConflict: ...
|
||||
|
||||
async def attempt_step(self, step_id: str, owner: str) -> WorkflowStep | WorkflowConflict: ...
|
||||
|
||||
async def finish_step(self, step: WorkflowStep, owner: str) -> bool: ...
|
||||
|
||||
async def checkpoint(self, run_id: str, checkpoint: WorkflowCheckpoint, step: WorkflowStep, owner: str) -> bool: ...
|
||||
|
||||
async def finish(
|
||||
self,
|
||||
run_id: str,
|
||||
state: WorkflowRunState,
|
||||
response_json: str | None = None,
|
||||
error: str | None = None,
|
||||
*,
|
||||
owner: str,
|
||||
) -> bool: ...
|
||||
|
||||
async def cancel(self, run_id: str) -> None: ...
|
||||
157
litellm/models/workflow_bundle.py
Normal file
157
litellm/models/workflow_bundle.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Annotated, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import ConfigDict, Field, JsonValue
|
||||
|
||||
from litellm.proxy.workflows.monty.models import PythonProgram
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
class NodeType(str, Enum):
|
||||
CHAT_COMPLETION = "chat_completion"
|
||||
PYTHON = "python"
|
||||
|
||||
|
||||
class BindingOutput(str, Enum):
|
||||
MESSAGES = "messages"
|
||||
RESULT = "result"
|
||||
|
||||
|
||||
class Reference(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
|
||||
|
||||
source: str = Field(alias="from", min_length=1, max_length=128)
|
||||
output: BindingOutput
|
||||
path: str | None = Field(default=None, max_length=512)
|
||||
|
||||
|
||||
class Constant(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
value: JsonValue
|
||||
|
||||
|
||||
Binding: TypeAlias = Reference | Constant
|
||||
|
||||
|
||||
class RetryPolicy(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
max_attempts: int = Field(default=1, ge=1, le=3)
|
||||
backoff_seconds: float = Field(default=0, ge=0, le=10, allow_inf_nan=False)
|
||||
|
||||
|
||||
class ResponseSchema(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
|
||||
|
||||
name: str = Field(pattern=r"^[A-Za-z0-9_-]{1,64}$")
|
||||
strict: bool = True
|
||||
schema_: Mapping[str, JsonValue] = Field(alias="schema")
|
||||
|
||||
|
||||
class StructuredResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
type: Literal["json_schema"] = "json_schema"
|
||||
json_schema: ResponseSchema
|
||||
|
||||
|
||||
class ChatCompletionNode(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
id: str = Field(pattern=r"^[A-Za-z][A-Za-z0-9_-]{0,63}$")
|
||||
type: Literal[NodeType.CHAT_COMPLETION]
|
||||
model: str = Field(min_length=1, max_length=256)
|
||||
messages: Binding
|
||||
temperature: float | None = Field(default=None, ge=0, le=2, allow_inf_nan=False)
|
||||
top_p: float | None = Field(default=None, ge=0, le=1, allow_inf_nan=False)
|
||||
max_tokens: int | None = Field(default=None, gt=0)
|
||||
max_completion_tokens: int | None = Field(default=None, gt=0)
|
||||
seed: int | None = None
|
||||
stop: str | tuple[str, ...] | None = None
|
||||
response_format: StructuredResponse | None = None
|
||||
timeout_seconds: float = Field(default=60, gt=0, le=300, allow_inf_nan=False)
|
||||
retry: RetryPolicy = RetryPolicy()
|
||||
|
||||
|
||||
class PythonNode(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
id: str = Field(pattern=r"^[A-Za-z][A-Za-z0-9_-]{0,63}$")
|
||||
type: Literal[NodeType.PYTHON]
|
||||
program: PythonProgram
|
||||
inputs: Mapping[str, Binding]
|
||||
timeout_seconds: float = Field(default=60, gt=0, le=300, allow_inf_nan=False)
|
||||
retry: RetryPolicy = RetryPolicy()
|
||||
|
||||
|
||||
WorkflowNode: TypeAlias = Annotated[ChatCompletionNode | PythonNode, Field(discriminator="type")]
|
||||
|
||||
|
||||
class WorkflowBundle(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
contract_version: Literal["1"]
|
||||
name: str = Field(pattern=r"^[a-z][a-z0-9-]{0,62}$")
|
||||
nodes: tuple[WorkflowNode, ...] = Field(min_length=1, max_length=64)
|
||||
timeout_seconds: float = Field(default=300, gt=0, le=900, allow_inf_nan=False)
|
||||
input_schema: Mapping[str, JsonValue]
|
||||
output: Reference
|
||||
|
||||
|
||||
class WorkflowFailure(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
code: str
|
||||
message: str
|
||||
param: str | None = None
|
||||
status_code: int = 400
|
||||
|
||||
|
||||
def canonical_bundle(bundle: WorkflowBundle) -> str:
|
||||
return json.dumps(
|
||||
bundle.model_dump(mode="json", by_alias=True, exclude_none=True),
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
)
|
||||
|
||||
|
||||
def workflow_digest(bundle: WorkflowBundle) -> str:
|
||||
return hashlib.sha256(canonical_bundle(bundle).encode()).hexdigest()
|
||||
|
||||
|
||||
def reference_failure(reference: Reference, available: frozenset[str]) -> WorkflowFailure | None:
|
||||
if reference.source == "$input" and reference.output is not BindingOutput.MESSAGES:
|
||||
return WorkflowFailure(code="invalid_reference", message="Request input exposes only messages")
|
||||
if reference.source != "$input" and (
|
||||
reference.source not in available or reference.output is not BindingOutput.RESULT
|
||||
):
|
||||
return WorkflowFailure(code="invalid_reference", message="Reference must select a completed predecessor")
|
||||
if reference.path is not None and not reference.path.startswith("/") and reference.path != "":
|
||||
return WorkflowFailure(code="invalid_reference", message="path must be a JSON Pointer")
|
||||
return None
|
||||
|
||||
|
||||
def validate_bindings(bundle: WorkflowBundle) -> WorkflowFailure | None:
|
||||
if bundle.name == "runs":
|
||||
return WorkflowFailure(code="reserved_name", message="runs is reserved for run tracking")
|
||||
ids: Final = tuple(node.id for node in bundle.nodes)
|
||||
if len(set(ids)) != len(ids):
|
||||
return WorkflowFailure(code="duplicate_node", message="Node IDs must be unique")
|
||||
for index, node in enumerate(bundle.nodes):
|
||||
bindings = (node.messages,) if isinstance(node, ChatCompletionNode) else tuple(node.inputs.values())
|
||||
for binding in bindings:
|
||||
failure = reference_failure(binding, frozenset(ids[:index])) if isinstance(binding, Reference) else None
|
||||
if failure is not None:
|
||||
return failure.model_copy(update={"param": f"nodes.{node.id}"})
|
||||
if isinstance(node, ChatCompletionNode) and node.model.startswith("workflow/"):
|
||||
return WorkflowFailure(
|
||||
code="nested_workflow", message="Nested workflow calls are unsupported", param=node.id
|
||||
)
|
||||
if isinstance(node, PythonNode) and set(node.inputs) != set(node.program.parameter_names):
|
||||
return WorkflowFailure(code="invalid_inputs", message="Bind exactly the function parameters", param=node.id)
|
||||
return reference_failure(bundle.output, frozenset(ids))
|
||||
|
|
@ -2506,6 +2506,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
Documents all the fields supported by `general_settings` in config.yaml
|
||||
"""
|
||||
|
||||
enable_workflows: bool = Field(
|
||||
default=False, description="Enable durable typed workflows and the optional Monty worker"
|
||||
)
|
||||
|
||||
use_spend_logs_ledger: bool = Field(
|
||||
True,
|
||||
description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.",
|
||||
|
|
@ -2785,7 +2789,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
)
|
||||
alert_to_webhook_url: dict | None = Field(
|
||||
None,
|
||||
description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`",
|
||||
description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': '$ALERT_WEBHOOK_URL'}`",
|
||||
)
|
||||
alerting_args: dict | None = Field(None, description="Controllable params for slack alerting - e.g. ttl in cache.")
|
||||
alerting_threshold: int | None = Field(
|
||||
|
|
|
|||
|
|
@ -5008,6 +5008,9 @@ def _can_object_call_model(
|
|||
from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent
|
||||
|
||||
compaction_parent: Final = native_compaction_parent(model)
|
||||
from litellm.proxy.workflows.auth import workflow_parent
|
||||
|
||||
workflow_model: Final = workflow_parent(model)
|
||||
global_or_router_alias_target: Final = (
|
||||
litellm.model_alias_map[model]
|
||||
if model in litellm.model_alias_map
|
||||
|
|
@ -5031,6 +5034,7 @@ def _can_object_call_model(
|
|||
if key_alias_applied
|
||||
else (
|
||||
*((model, compaction_parent) if compaction_parent is not None else (model,)),
|
||||
*((workflow_model,) if workflow_model is not None else ()),
|
||||
*((global_or_router_alias_target,) if global_or_router_alias_target else ()),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -373,6 +373,8 @@ class RouteChecks:
|
|||
pass
|
||||
elif route == "/lens" or route.startswith("/lens/"):
|
||||
pass # Lens authorizes the delegated role, including its POST-based read endpoints.
|
||||
elif RouteChecks.is_workflow_execution_route(route, RouteChecks.get_request_method(request=request)):
|
||||
pass
|
||||
elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value:
|
||||
RouteChecks._check_proxy_admin_viewer_access(
|
||||
route=route,
|
||||
|
|
@ -753,6 +755,25 @@ class RouteChecks:
|
|||
|
||||
_get_request_method = get_request_method
|
||||
|
||||
@staticmethod
|
||||
def is_workflow_execution_route(route: str, method: str | None) -> bool:
|
||||
run_patterns: Final = (
|
||||
("GET", "/v1/workflows/runs/{run_id}/steps"),
|
||||
("GET", "/v1/workflows/runs/{run_id}/steps/{step_id}"),
|
||||
("POST", "/v1/workflows/runs/{run_id}/cancel"),
|
||||
)
|
||||
definition_patterns: Final = (
|
||||
("POST", "/v1/workflows"),
|
||||
("GET", "/v1/workflows"),
|
||||
("GET", "/v1/workflows/{workflow_id}"),
|
||||
("PUT", "/v1/workflows/{workflow_id}"),
|
||||
("POST", "/v1/workflows/{workflow_id}/validate"),
|
||||
("POST", "/v1/workflows/{workflow_id}/test"),
|
||||
("POST", "/v1/workflows/{workflow_id}/publish"),
|
||||
)
|
||||
patterns: Final = run_patterns if route.startswith("/v1/workflows/runs") else definition_patterns
|
||||
return any(method == verb and RouteChecks.check_route_access(route, (pattern,)) for verb, pattern in patterns)
|
||||
|
||||
@staticmethod
|
||||
def is_self_scoped_credential_route(route: str, request: Request | None) -> bool:
|
||||
method: Final = RouteChecks.get_request_method(request)
|
||||
|
|
|
|||
|
|
@ -2696,6 +2696,12 @@ async def user_api_key_auth_builder(
|
|||
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
|
||||
### CHECK IF ADMIN ###
|
||||
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
|
||||
from litellm.proxy.workflows.auth import workflow_child_identity
|
||||
|
||||
workflow_identity: Final = await workflow_child_identity(request, request_data)
|
||||
if workflow_identity is not None:
|
||||
valid_token = workflow_identity # rebind-ok: Existing auth state must use the verified child identity
|
||||
api_key = workflow_identity.token or "" # rebind-ok: Pass the trusted principal to existing auth checks
|
||||
if valid_token is None:
|
||||
## Check CACHE
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -7,12 +7,21 @@ from litellm.models.gateway_activity import GatewayActivityStore
|
|||
from litellm.models.plugin_upload import PluginUploadStore
|
||||
from litellm.models.skill_submission import SkillSubmissionStore
|
||||
from litellm.models.user_banner import UserBannerStore
|
||||
from litellm.models.workflow import WorkflowStore
|
||||
from litellm.proxy.deps.database import DatabaseDep
|
||||
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
|
||||
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
|
||||
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
|
||||
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
|
||||
from litellm.repositories.user_banner_repository import UserBannerRepository
|
||||
from litellm.repositories.workflow_repository import WorkflowRepository
|
||||
|
||||
|
||||
async def get_workflow_repository(database: DatabaseDep) -> WorkflowStore | None:
|
||||
return None if database is None else WorkflowRepository(database.prisma)
|
||||
|
||||
|
||||
WorkflowRepositoryDep: TypeAlias = Annotated[WorkflowStore | None, Depends(get_workflow_repository)]
|
||||
|
||||
|
||||
async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
|
||||
|
|
|
|||
|
|
@ -5034,7 +5034,7 @@ async def info_key_fn(
|
|||
default=None,
|
||||
description=(
|
||||
"Key to look up. Pass the key's sha256 hash so the raw key stays out of URLs and access "
|
||||
"logs. Example key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'"
|
||||
"logs. Example key='YOUR_VIRTUAL_KEY'"
|
||||
),
|
||||
),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
@ -5081,14 +5081,14 @@ async def info_key_fn(
|
|||
|
||||
Example Curl:
|
||||
```
|
||||
curl -X GET "http://0.0.0.0:4000/key/info?key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
|
||||
curl -X GET "http://0.0.0.0:4000/key/info?key=YOUR_VIRTUAL_KEY" \
|
||||
-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}"
|
||||
```
|
||||
|
||||
Example Curl - if no key is passed, it will use the Key Passed in Authorization Header
|
||||
```
|
||||
curl -X GET "http://0.0.0.0:4000/key/info" \
|
||||
-H "Authorization: Bearer sk-test-example-key-123"
|
||||
-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}"
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
|
|
@ -8076,7 +8076,7 @@ async def block_key(
|
|||
--header "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data '{
|
||||
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
|
||||
"key": "YOUR_VIRTUAL_KEY"
|
||||
}'
|
||||
```
|
||||
|
||||
|
|
@ -8190,7 +8190,7 @@ async def unblock_key(
|
|||
--header "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data '{
|
||||
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
|
||||
"key": "YOUR_VIRTUAL_KEY"
|
||||
}'
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -5399,7 +5399,7 @@ async def team_info(
|
|||
|
||||
```
|
||||
curl --location 'http://localhost:4000/team/info?team_id=your_team_id_here' \
|
||||
--header 'Authorization: Bearer your_api_key_here'
|
||||
-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}"
|
||||
```
|
||||
"""
|
||||
from litellm.proxy._types import TeamInfoResponseObjectTeamTable
|
||||
|
|
@ -5566,7 +5566,7 @@ async def team_member_me(
|
|||
|
||||
```
|
||||
curl --location 'http://localhost:4000/team/your_team_id/members/me' \
|
||||
--header 'Authorization: Bearer your_api_key_here'
|
||||
-H "Authorization: Bearer ${YOUR_VIRTUAL_KEY}"
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
|
|
|||
|
|
@ -15,9 +15,11 @@ GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final, Literal, Protocol, TypedDict
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from pydantic import ConfigDict, JsonValue
|
||||
|
||||
try:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
|
@ -102,10 +104,53 @@ class WorkflowMessageCreateRequest(LiteLLMBaseModel):
|
|||
session_id: str | None = None
|
||||
|
||||
|
||||
class WorkflowRunResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
run_id: str
|
||||
session_id: str
|
||||
workflow_type: str
|
||||
status: str
|
||||
created_by: str | None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
input: JsonValue
|
||||
output: JsonValue
|
||||
metadata: JsonValue
|
||||
engine_owned: bool
|
||||
|
||||
|
||||
class WorkflowEventResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
event_id: str
|
||||
run_id: str
|
||||
event_type: str
|
||||
step_name: str
|
||||
sequence_number: int
|
||||
data: JsonValue
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class WorkflowMessageResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
message_id: str
|
||||
run_id: str
|
||||
role: str
|
||||
content: str
|
||||
sequence_number: int
|
||||
session_id: str | None
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class _RunRow(Protocol):
|
||||
@property
|
||||
def created_by(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def engine_owned(self) -> bool: ...
|
||||
|
||||
|
||||
class _SeqRow(Protocol):
|
||||
@property
|
||||
|
|
@ -201,6 +246,8 @@ async def _require_run(
|
|||
prisma_client: object,
|
||||
run_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
*,
|
||||
mutation: bool = False,
|
||||
) -> _RunRow:
|
||||
"""Return the run or raise 404. For non-admin callers, also enforce key ownership."""
|
||||
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id})
|
||||
|
|
@ -210,6 +257,8 @@ async def _require_run(
|
|||
caller: Final = _caller_key(user_api_key_dict)
|
||||
if not caller or run.created_by != caller:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
if mutation and run.engine_owned:
|
||||
raise HTTPException(status_code=409, detail="Engine-owned runs cannot be modified through tracking APIs")
|
||||
return run
|
||||
|
||||
|
||||
|
|
@ -220,6 +269,7 @@ async def _require_run(
|
|||
|
||||
@router.post(
|
||||
"/v1/workflows/runs",
|
||||
responses={200: {"model": WorkflowRunResponse}},
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
|
|
@ -359,8 +409,7 @@ async def update_workflow_run(
|
|||
if not update:
|
||||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
|
||||
# Enforce ownership before writing.
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
try:
|
||||
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update(
|
||||
|
|
@ -379,6 +428,7 @@ async def update_workflow_run(
|
|||
|
||||
@router.post(
|
||||
"/v1/workflows/runs/{run_id}/events",
|
||||
responses={200: {"model": WorkflowEventResponse}},
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
|
|
@ -398,7 +448,7 @@ async def append_workflow_event(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
new_status: Final = _EVENT_STATUS_MAP.get(data.event_type)
|
||||
|
||||
|
|
@ -470,6 +520,7 @@ async def list_workflow_events(
|
|||
|
||||
@router.post(
|
||||
"/v1/workflows/runs/{run_id}/messages",
|
||||
responses={200: {"model": WorkflowMessageResponse}},
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
|
|
@ -487,7 +538,7 @@ async def append_workflow_message(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
for attempt in range(_MAX_SEQUENCE_RETRIES):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -683,6 +683,7 @@ from litellm.proxy.db.request_error_tracking import (
|
|||
request_error_accumulator,
|
||||
)
|
||||
from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed
|
||||
from litellm.proxy.deps.repositories import WorkflowRepositoryDep
|
||||
from litellm.proxy.discovery_endpoints import (
|
||||
agent_skills_discovery_router,
|
||||
ui_discovery_endpoints_router,
|
||||
|
|
@ -887,6 +888,8 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware
|
|||
from litellm.proxy.telemetry.runtime import TelemetryRuntime
|
||||
from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors
|
||||
from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings
|
||||
from litellm.proxy.workflows.management import router as workflow_definition_router
|
||||
from litellm.proxy.workflows.runs import router as workflow_execution_router
|
||||
from litellm.router import router_arguments
|
||||
from litellm.router_config import build_router_config
|
||||
|
||||
|
|
@ -1975,7 +1978,12 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
|
|||
await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app))
|
||||
gc.freeze()
|
||||
verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count())
|
||||
yield state
|
||||
from litellm.proxy.workflows.dispatcher import workflow_dispatch_lifespan
|
||||
|
||||
async with workflow_dispatch_lifespan(
|
||||
state, app, general_settings_view().get("enable_workflows") is True
|
||||
):
|
||||
yield state
|
||||
finally:
|
||||
if model_info_scheduler is not None and model_info_scheduler.running:
|
||||
model_info_scheduler.remove_job("refresh_model_info")
|
||||
|
|
@ -3137,6 +3145,10 @@ def general_settings_view() -> Mapping[str, object]:
|
|||
return _GENERAL_SETTINGS_VIEW.validate_python(general_settings)
|
||||
|
||||
|
||||
def llm_model_list_view() -> tuple[object, ...] | None:
|
||||
return TypeAdapter[tuple[object, ...] | None](tuple[object, ...] | None).validate_python(llm_model_list)
|
||||
|
||||
|
||||
config_passthrough_endpoints: list[dict[str, object]] | None = None
|
||||
log_file: Final = "api_log.json"
|
||||
worker_config: Final = None
|
||||
|
|
@ -13129,6 +13141,7 @@ async def chat_completion(
|
|||
fastapi_response: Response,
|
||||
model: str | None = None,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
workflows: WorkflowRepositoryDep = None,
|
||||
):
|
||||
"""
|
||||
|
||||
|
|
@ -13156,6 +13169,12 @@ async def chat_completion(
|
|||
global general_settings, user_debug, proxy_logging_obj, llm_model_list
|
||||
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
||||
data: Final = await read_request_body(request=request)
|
||||
workflow_data: Final = _GENERAL_SETTINGS_VIEW.validate_python(data)
|
||||
workflow_model: Final = workflow_data.get("model")
|
||||
if isinstance(workflow_model, str) and workflow_model.startswith("workflow/"):
|
||||
from litellm.proxy.workflows.invocation import invoke_workflow
|
||||
|
||||
return await invoke_workflow(request, user_api_key_dict, workflow_data, workflows)
|
||||
if user_api_key_dict is not None:
|
||||
if not isinstance(data.get("metadata"), dict):
|
||||
# Covers both missing and JSON-string metadata (multipart /
|
||||
|
|
@ -22432,6 +22451,8 @@ app.include_router(model_access_group_management_router)
|
|||
app.include_router(auto_router_management_router)
|
||||
app.include_router(tag_management_router)
|
||||
app.include_router(workflow_management_router)
|
||||
app.include_router(workflow_execution_router)
|
||||
app.include_router(workflow_definition_router)
|
||||
app.include_router(memory_router)
|
||||
app.include_router(lens_router)
|
||||
app.include_router(plugin_router)
|
||||
|
|
|
|||
|
|
@ -1986,6 +1986,7 @@ model LiteLLM_WorkflowRun {
|
|||
input Json?
|
||||
output Json?
|
||||
metadata Json?
|
||||
engine_owned Boolean @default(false)
|
||||
|
||||
events LiteLLM_WorkflowEvent[]
|
||||
messages LiteLLM_WorkflowMessage[]
|
||||
|
|
@ -1996,6 +1997,72 @@ model LiteLLM_WorkflowRun {
|
|||
@@index([created_by])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowDefinition {
|
||||
workflow_id String @id
|
||||
team_id String
|
||||
name String
|
||||
etag String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_snapshot_id String?
|
||||
|
||||
@@unique([team_id, name])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowSnapshot {
|
||||
snapshot_id String @id
|
||||
workflow_id String
|
||||
team_id String
|
||||
name String
|
||||
bundle_json String
|
||||
published_by String
|
||||
created_at DateTime @default(now())
|
||||
origin String @default("publication")
|
||||
|
||||
@@index([workflow_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowExecution {
|
||||
run_id String @id
|
||||
snapshot_id String
|
||||
team_id String
|
||||
name String
|
||||
principal String
|
||||
dispatch_key String @unique
|
||||
request_digest String
|
||||
request_json String
|
||||
state String @default("pending")
|
||||
deadline DateTime
|
||||
lease_owner String?
|
||||
lease_until DateTime?
|
||||
cancel_requested Boolean @default(false)
|
||||
checkpoints Json @default("[]")
|
||||
response_json String?
|
||||
error String?
|
||||
|
||||
@@index([team_id, name])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowStep {
|
||||
step_id String @id
|
||||
run_id String
|
||||
node_id String
|
||||
state String @default("running")
|
||||
input Json
|
||||
output Json?
|
||||
attempts Int @default(0)
|
||||
deadline DateTime
|
||||
started_at DateTime @default(now())
|
||||
finished_at DateTime?
|
||||
error String?
|
||||
prompt_tokens Int @default(0)
|
||||
completion_tokens Int @default(0)
|
||||
response_cost Float?
|
||||
request_id String?
|
||||
|
||||
@@unique([run_id, node_id])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
// `step_name` and `event_type` are caller-defined strings — no hardcoded enums.
|
||||
// Status auto-update rules (applied by the append endpoint):
|
||||
|
|
|
|||
|
|
@ -1470,7 +1470,7 @@ async def get_global_spend_report(
|
|||
default=None,
|
||||
description=(
|
||||
"View spend for a specific api_key. Pass the key's sha256 hash so the raw key stays "
|
||||
"out of URLs and access logs. Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'"
|
||||
"out of URLs and access logs. Example api_key='YOUR_VIRTUAL_KEY'"
|
||||
),
|
||||
),
|
||||
internal_user_id: str | None = fastapi.Query(
|
||||
|
|
@ -1945,7 +1945,7 @@ async def get_key_spend_report(
|
|||
description=(
|
||||
"View spend for a specific api_key. Proxy admin only; other callers are scoped to their "
|
||||
"own key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. "
|
||||
"Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'"
|
||||
"Example api_key='YOUR_VIRTUAL_KEY'"
|
||||
)
|
||||
),
|
||||
] = None,
|
||||
|
|
@ -3447,7 +3447,7 @@ async def view_spend_logs(
|
|||
|
||||
Example Request for specific api_key
|
||||
```
|
||||
curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \
|
||||
curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=YOUR_VIRTUAL_KEY" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
|
||||
```
|
||||
|
||||
|
|
|
|||
64
litellm/proxy/workflows/auth.py
Normal file
64
litellm/proxy/workflows/auth.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
from collections.abc import Mapping
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkflowChildGrant:
|
||||
principal: str
|
||||
team_id: str
|
||||
snapshot_id: str
|
||||
workflow_model: str
|
||||
node_id: str
|
||||
model: str
|
||||
run_id: str
|
||||
|
||||
|
||||
workflow_child_grant: Final[ContextVar[WorkflowChildGrant | None]] = ContextVar("workflow_child_grant", default=None)
|
||||
|
||||
|
||||
def workflow_parent(model: str) -> str | None:
|
||||
grant: Final = workflow_child_grant.get()
|
||||
return grant.workflow_model if grant is not None and grant.model == model else None
|
||||
|
||||
|
||||
async def workflow_child_identity(request: Request, request_data: Mapping[str, object]) -> UserAPIKeyAuth | None:
|
||||
grant: Final = workflow_child_grant.get()
|
||||
if grant is None:
|
||||
return None
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if (
|
||||
request.method != "POST"
|
||||
or not request.url.path.endswith("/v1/chat/completions")
|
||||
or request_data.get("model") != grant.model
|
||||
):
|
||||
from fastapi import HTTPException
|
||||
|
||||
raise HTTPException(status_code=403, detail="Invalid internal workflow child scope")
|
||||
identity: Final = IdentityStore.key_from_principal(
|
||||
await IdentityStore(
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
).resolve(hashed_token=grant.principal)
|
||||
)
|
||||
if (
|
||||
identity.team_id != grant.team_id
|
||||
or identity.is_session_token
|
||||
or identity.user_role
|
||||
in {
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
}
|
||||
):
|
||||
from fastapi import HTTPException
|
||||
|
||||
raise HTTPException(status_code=403, detail="Workflow requires its original team virtual key")
|
||||
return identity.model_copy(update={"via_virtual_key": True})
|
||||
168
litellm/proxy/workflows/child.py
Normal file
168
litellm/proxy/workflows/child.py
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
from typing import Annotated, Final
|
||||
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from litellm.llms.custom_httpx.asgi_handler import get_async_asgi_client
|
||||
from litellm.models.workflow import WorkflowExecution
|
||||
from litellm.models.workflow_bundle import ChatCompletionNode, WorkflowFailure
|
||||
from litellm.proxy._types import ProxyLifespanState
|
||||
from litellm.proxy.workflows.auth import WorkflowChildGrant, workflow_child_grant
|
||||
from litellm.proxy.workflows.runner import ChildResult
|
||||
from litellm.proxy.workflows.validation import validate_value
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
class _Message(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
content: str | None = None
|
||||
refusal: str | None = None
|
||||
|
||||
|
||||
class _Choice(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
finish_reason: str
|
||||
message: _Message
|
||||
|
||||
|
||||
class _Usage(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
|
||||
|
||||
class _Completion(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
choices: tuple[_Choice, ...]
|
||||
usage: _Usage = _Usage()
|
||||
|
||||
|
||||
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
|
||||
_JSON: Final = TypeAdapter[JsonValue](JsonValue)
|
||||
_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_COST: Final = TypeAdapter[float | None](Annotated[float | None, Field(ge=0, allow_inf_nan=False)])
|
||||
|
||||
|
||||
class GatewayChildCaller:
|
||||
def __init__(self, app: ASGIApp, state: ProxyLifespanState) -> None:
|
||||
self._app: Final = app
|
||||
self._state: Final = state
|
||||
|
||||
async def dispatch(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
await self._app({**scope, "state": {**self._state}}, receive, send)
|
||||
|
||||
async def call(
|
||||
self, execution: WorkflowExecution, node: ChatCompletionNode, messages: JsonValue, remaining: float
|
||||
) -> ChildResult | WorkflowFailure:
|
||||
try:
|
||||
validated_messages: Final = _MESSAGES.validate_python(messages)
|
||||
except ValidationError:
|
||||
return WorkflowFailure(
|
||||
code="invalid_messages",
|
||||
message="Resolved messages must be standard Chat Completions messages",
|
||||
param=node.id,
|
||||
)
|
||||
original_request: Final = _OBJECT.validate_json(execution.request_json)
|
||||
payload: Final = {
|
||||
**node.model_dump(
|
||||
mode="json",
|
||||
by_alias=True,
|
||||
exclude_none=True,
|
||||
exclude={"id", "type", "messages", "retry", "timeout_seconds"},
|
||||
),
|
||||
"messages": _MESSAGES.dump_python(validated_messages, mode="json"),
|
||||
"stream": False,
|
||||
"n": 1,
|
||||
"timeout": remaining,
|
||||
**({"user": original_request["user"]} if "user" in original_request else {}),
|
||||
"metadata": {
|
||||
"workflow_run_id": execution.run_id,
|
||||
"workflow_node_id": node.id,
|
||||
"workflow_snapshot_id": execution.snapshot_id,
|
||||
},
|
||||
}
|
||||
token: Final = workflow_child_grant.set(
|
||||
WorkflowChildGrant(
|
||||
principal=execution.principal,
|
||||
team_id=execution.team_id,
|
||||
snapshot_id=execution.snapshot_id,
|
||||
workflow_model=f"workflow/{execution.name}",
|
||||
node_id=node.id,
|
||||
model=node.model,
|
||||
run_id=execution.run_id,
|
||||
)
|
||||
)
|
||||
try:
|
||||
with get_async_asgi_client(app=self.dispatch) as client:
|
||||
response: Final = await client.post(
|
||||
"http://litellm.internal/v1/chat/completions",
|
||||
json=payload,
|
||||
headers={"Authorization": "Bearer sk-workflow-internal", "x-litellm-session-id": execution.run_id},
|
||||
timeout=remaining,
|
||||
)
|
||||
if not response.is_success:
|
||||
return WorkflowFailure(
|
||||
code="child_call_failed",
|
||||
message=f"Gateway child returned HTTP {response.status_code}",
|
||||
param=node.id,
|
||||
status_code=response.status_code,
|
||||
)
|
||||
try:
|
||||
completion: Final = _Completion.model_validate_json(response.content)
|
||||
response_cost: Final = _COST.validate_python(response.headers.get("x-litellm-response-cost"))
|
||||
except ValidationError:
|
||||
return WorkflowFailure(
|
||||
code="invalid_completion",
|
||||
message="Gateway child returned an invalid completion",
|
||||
param=node.id,
|
||||
status_code=502,
|
||||
)
|
||||
if len(completion.choices) != 1:
|
||||
return WorkflowFailure(
|
||||
code="invalid_completion",
|
||||
message="Workflow nodes require one completion choice",
|
||||
param=node.id,
|
||||
status_code=502,
|
||||
)
|
||||
choice: Final = completion.choices[0]
|
||||
if choice.finish_reason != "stop" or choice.message.refusal is not None or choice.message.content is None:
|
||||
return WorkflowFailure(
|
||||
code="invalid_completion",
|
||||
message="Completion refused, truncated, or produced no text result",
|
||||
param=node.id,
|
||||
status_code=502,
|
||||
)
|
||||
if node.response_format is None:
|
||||
return ChildResult(
|
||||
output=choice.message.content,
|
||||
prompt_tokens=completion.usage.prompt_tokens,
|
||||
completion_tokens=completion.usage.completion_tokens,
|
||||
request_id=completion.id,
|
||||
response_cost=response_cost,
|
||||
)
|
||||
try:
|
||||
output: Final = _JSON.validate_json(choice.message.content)
|
||||
except ValidationError:
|
||||
return WorkflowFailure(
|
||||
code="invalid_json", message="Completion content is not valid JSON", param=node.id, status_code=502
|
||||
)
|
||||
failure: Final = validate_value(output, node.response_format.json_schema.schema_, node.id)
|
||||
return (
|
||||
failure
|
||||
if failure is not None
|
||||
else ChildResult(
|
||||
output=output,
|
||||
prompt_tokens=completion.usage.prompt_tokens,
|
||||
completion_tokens=completion.usage.completion_tokens,
|
||||
request_id=completion.id,
|
||||
response_cost=response_cost,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
workflow_child_grant.reset(token)
|
||||
158
litellm/proxy/workflows/dispatcher.py
Normal file
158
litellm/proxy/workflows/dispatcher.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
from starlette.types import ASGIApp
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.models.workflow import WorkflowExecution, WorkflowRunState, WorkflowStore
|
||||
from litellm.models.workflow_bundle import WorkflowFailure
|
||||
from litellm.proxy._types import ProxyLifespanState
|
||||
from litellm.proxy.workflows.child import GatewayChildCaller
|
||||
from litellm.proxy.workflows.monty.runtime import MontyRuntime
|
||||
from litellm.proxy.workflows.runner import ChildCaller, execute_run
|
||||
from litellm.proxy.workflows.validation import DEFAULT_MONTY_LIMITS
|
||||
from litellm.repositories.workflow_repository import WorkflowRepository
|
||||
|
||||
|
||||
async def maintain_lease(
|
||||
store: WorkflowStore, execution: WorkflowExecution, work: asyncio.Task[JsonValue | WorkflowFailure]
|
||||
) -> None:
|
||||
try:
|
||||
while not work.done():
|
||||
await asyncio.sleep(2)
|
||||
current = await store.execution(execution.run_id)
|
||||
if (
|
||||
current is None
|
||||
or current.cancel_requested
|
||||
or not await store.renew(execution.run_id, execution.lease_owner or "")
|
||||
):
|
||||
work.cancel()
|
||||
return
|
||||
except Exception: # noqa: BLE001 # Any persistence failure must stop the worker before its lease expires
|
||||
work.cancel()
|
||||
verbose_proxy_logger.exception("Workflow lease renewal failed; stopping dispatch")
|
||||
|
||||
|
||||
async def format_response(store: WorkflowStore, execution: WorkflowExecution, result: JsonValue) -> str:
|
||||
steps: Final = await store.steps(execution.run_id)
|
||||
prompt_tokens: Final = sum(step.prompt_tokens for step in steps)
|
||||
completion_tokens: Final = sum(step.completion_tokens for step in steps)
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{execution.run_id}",
|
||||
"object": "chat.completion",
|
||||
"created": int(datetime.now(timezone.utc).timestamp()),
|
||||
"model": f"workflow/{execution.name}",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": result if isinstance(result, str) else json.dumps(result, allow_nan=False),
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
},
|
||||
allow_nan=False,
|
||||
)
|
||||
|
||||
|
||||
async def dispatch_one(
|
||||
store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, execution: WorkflowExecution
|
||||
) -> None:
|
||||
work: Final = asyncio.create_task(execute_run(store, caller, runtime, execution))
|
||||
heartbeat: Final = asyncio.create_task(maintain_lease(store, execution, work))
|
||||
try:
|
||||
result: Final = await work
|
||||
if isinstance(result, WorkflowFailure):
|
||||
await store.finish(
|
||||
execution.run_id,
|
||||
WorkflowRunState.CANCELLED if result.code == "cancelled" else WorkflowRunState.FAILED,
|
||||
error=result.model_dump_json(),
|
||||
owner=execution.lease_owner or "",
|
||||
)
|
||||
return
|
||||
await store.finish(
|
||||
execution.run_id,
|
||||
WorkflowRunState.COMPLETED,
|
||||
response_json=await format_response(store, execution, result),
|
||||
owner=execution.lease_owner or "",
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
current: Final = await store.execution(execution.run_id)
|
||||
if current is not None and current.cancel_requested and current.lease_owner == execution.lease_owner:
|
||||
await store.finish(execution.run_id, WorkflowRunState.CANCELLED, owner=execution.lease_owner or "")
|
||||
elif (task := asyncio.current_task()) is not None and task.cancelling():
|
||||
raise
|
||||
finally:
|
||||
work.cancel()
|
||||
heartbeat.cancel()
|
||||
await asyncio.gather(work, heartbeat, return_exceptions=True)
|
||||
|
||||
|
||||
async def worker(store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, owner: str) -> None:
|
||||
while True:
|
||||
try:
|
||||
execution = await store.lease_next(f"{owner}/{uuid.uuid4()}")
|
||||
if execution is None:
|
||||
await asyncio.sleep(0.25)
|
||||
continue
|
||||
await dispatch_one(store, caller, runtime, execution)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception: # noqa: BLE001 # Durable unfinished runs must remain recoverable after unexpected failures
|
||||
verbose_proxy_logger.exception("Workflow dispatcher failed; unfinished work remains lease-recoverable")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkflowDispatcher:
|
||||
tasks: tuple[asyncio.Task[None], ...]
|
||||
|
||||
async def aclose(self) -> None:
|
||||
for task in self.tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*self.tasks, return_exceptions=True)
|
||||
|
||||
|
||||
def start_dispatcher(
|
||||
store: WorkflowStore, app: ASGIApp, state: ProxyLifespanState, runtime: MontyRuntime
|
||||
) -> WorkflowDispatcher:
|
||||
caller: Final = GatewayChildCaller(app, state)
|
||||
owner: Final = str(uuid.uuid4())
|
||||
return WorkflowDispatcher(
|
||||
tasks=tuple(asyncio.create_task(worker(store, caller, runtime, f"{owner}/{index}")) for index in range(4))
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def workflow_dispatch_lifespan(
|
||||
state: ProxyLifespanState, app: ASGIApp, enabled: bool
|
||||
) -> AsyncGenerator[None, None]:
|
||||
database: Final = state["database"]
|
||||
if not enabled or database is None:
|
||||
yield
|
||||
return
|
||||
runtime: Final = MontyRuntime(DEFAULT_MONTY_LIMITS)
|
||||
failure: Final = await runtime.start()
|
||||
if failure is not None:
|
||||
raise RuntimeError(f"Cannot enable workflows, sandbox unavailable ({failure.kind.value})")
|
||||
dispatcher: Final = start_dispatcher(WorkflowRepository(database.prisma), app, state, runtime)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await dispatcher.aclose()
|
||||
await runtime.close()
|
||||
219
litellm/proxy/workflows/invocation.py
Normal file
219
litellm/proxy/workflows/invocation.py
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final, Literal
|
||||
|
||||
from fastapi import Request
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.models.workflow import WorkflowExecution, WorkflowRunState, WorkflowSnapshot, WorkflowStore
|
||||
from litellm.models.workflow_bundle import WorkflowFailure
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import hash_token
|
||||
from litellm.proxy.workflows.management import raise_public, require_store
|
||||
from litellm.proxy.workflows.validation import canonical_request
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
class InvocationMetadata(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
tags: tuple[str, ...] = Field(default=(), max_length=20)
|
||||
|
||||
|
||||
class WorkflowInvocation(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
model: str = Field(pattern=r"^workflow/[a-z][a-z0-9-]{0,62}$")
|
||||
messages: tuple[ChatCompletionMessageParam, ...] = Field(min_length=1, max_length=128)
|
||||
user: str | None = Field(default=None, max_length=256)
|
||||
metadata: InvocationMetadata | None = None
|
||||
stream: Literal[False] = False
|
||||
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_RUN_HEADER: Final = "x-litellm-workflow-run-id"
|
||||
|
||||
|
||||
def principal_hash(auth: UserAPIKeyAuth) -> str:
|
||||
from litellm.proxy.proxy_server import user_custom_auth
|
||||
|
||||
if (
|
||||
not auth.token
|
||||
or not auth.team_id
|
||||
or not auth.via_virtual_key
|
||||
or auth.is_session_token
|
||||
or user_custom_auth is not None
|
||||
):
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="virtual_key_required", message="Workflow invocation requires a team virtual key", status_code=403
|
||||
)
|
||||
)
|
||||
return hash_token(auth.token) if auth.token.startswith("sk-") else auth.token
|
||||
|
||||
|
||||
async def admit(
|
||||
store: WorkflowStore,
|
||||
auth: UserAPIKeyAuth,
|
||||
invocation: WorkflowInvocation,
|
||||
idempotency_key: str | None,
|
||||
*,
|
||||
snapshot: WorkflowSnapshot | None = None,
|
||||
operation: str = "invoke",
|
||||
) -> WorkflowExecution:
|
||||
principal: Final = principal_hash(auth)
|
||||
if snapshot is not None and snapshot.team_id != auth.team_id:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="workflow_team_mismatch",
|
||||
message="Workflow execution requires a virtual key for its owning team",
|
||||
status_code=403,
|
||||
)
|
||||
)
|
||||
if idempotency_key is not None and not 1 <= len(idempotency_key) <= 256:
|
||||
raise_public(
|
||||
WorkflowFailure(code="invalid_idempotency_key", message="Idempotency-Key must contain 1 to 256 characters")
|
||||
)
|
||||
request_json: Final = canonical_request(
|
||||
_JSON_OBJECT.validate_python(invocation.model_dump(mode="json", exclude_none=True))
|
||||
)
|
||||
digest: Final = hashlib.sha256(request_json.encode()).hexdigest()
|
||||
name: Final = invocation.model.removeprefix("workflow/")
|
||||
dispatch_key: Final = hashlib.sha256(
|
||||
canonical_request(
|
||||
{
|
||||
"principal": principal,
|
||||
"team": auth.team_id,
|
||||
"operation": operation,
|
||||
"workflow": name,
|
||||
"idempotency_key": idempotency_key if idempotency_key is not None else str(uuid.uuid4()),
|
||||
}
|
||||
).encode()
|
||||
).hexdigest()
|
||||
existing: Final = await store.by_dispatch_key(dispatch_key)
|
||||
if existing is not None:
|
||||
if existing.request_digest != digest:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="idempotency_conflict",
|
||||
message="Idempotency-Key was accepted with a different request",
|
||||
status_code=409,
|
||||
)
|
||||
)
|
||||
return existing
|
||||
pinned: Final = snapshot if snapshot is not None else await store.published(auth.team_id or "", name)
|
||||
if pinned is None:
|
||||
raise_public(WorkflowFailure(code="not_found", message="Published workflow not found", status_code=404))
|
||||
from litellm.models.workflow_bundle import WorkflowBundle
|
||||
|
||||
bundle: Final = WorkflowBundle.model_validate_json(pinned.bundle_json)
|
||||
accepted: Final = await store.dispatch(
|
||||
WorkflowExecution(
|
||||
run_id=str(uuid.uuid4()),
|
||||
snapshot_id=pinned.snapshot_id,
|
||||
team_id=pinned.team_id,
|
||||
name=name,
|
||||
principal=principal,
|
||||
dispatch_key=dispatch_key,
|
||||
request_digest=digest,
|
||||
request_json=request_json,
|
||||
state=WorkflowRunState.PENDING,
|
||||
deadline=datetime.now(timezone.utc) + timedelta(seconds=bundle.timeout_seconds),
|
||||
cancel_requested=False,
|
||||
)
|
||||
)
|
||||
if accepted.request_digest != digest:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="idempotency_conflict",
|
||||
message="Idempotency-Key was accepted with a different request",
|
||||
status_code=409,
|
||||
)
|
||||
)
|
||||
return accepted
|
||||
|
||||
|
||||
def error_response(failure: WorkflowFailure, run_id: str) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=failure.status_code,
|
||||
content={
|
||||
"error": {
|
||||
"type": "workflow_error",
|
||||
"code": failure.code,
|
||||
"message": failure.message,
|
||||
"param": failure.param,
|
||||
}
|
||||
},
|
||||
headers={_RUN_HEADER: run_id},
|
||||
)
|
||||
|
||||
|
||||
async def wait_for_run(store: WorkflowStore, execution: WorkflowExecution, wait_seconds: float = 60.0) -> JSONResponse:
|
||||
try:
|
||||
async with asyncio.timeout(wait_seconds):
|
||||
while True:
|
||||
current = await store.execution(execution.run_id)
|
||||
if current is None:
|
||||
return error_response(
|
||||
WorkflowFailure(
|
||||
code="run_unavailable", message="Accepted run is temporarily unavailable", status_code=503
|
||||
),
|
||||
execution.run_id,
|
||||
)
|
||||
if current.state is WorkflowRunState.COMPLETED and current.response_json is not None:
|
||||
return JSONResponse(
|
||||
content=_JSON_OBJECT.validate_json(current.response_json), headers={_RUN_HEADER: current.run_id}
|
||||
)
|
||||
if current.state is WorkflowRunState.FAILED:
|
||||
failure = (
|
||||
WorkflowFailure.model_validate_json(current.error)
|
||||
if current.error is not None
|
||||
else WorkflowFailure(
|
||||
code="execution_failed",
|
||||
message="Workflow execution failed",
|
||||
status_code=502,
|
||||
)
|
||||
)
|
||||
return error_response(failure, current.run_id)
|
||||
if current.state is WorkflowRunState.CANCELLED:
|
||||
return error_response(
|
||||
WorkflowFailure(code="cancelled", message="Run was cancelled", status_code=409), current.run_id
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
except TimeoutError:
|
||||
return error_response(
|
||||
WorkflowFailure(
|
||||
code="wait_timeout",
|
||||
message="Run continues; retry with the same Idempotency-Key or inspect it",
|
||||
status_code=504,
|
||||
),
|
||||
execution.run_id,
|
||||
)
|
||||
|
||||
|
||||
async def invoke_workflow(
|
||||
request: Request, auth: UserAPIKeyAuth, data: Mapping[str, object], workflows: WorkflowStore | None
|
||||
) -> JSONResponse:
|
||||
from litellm.proxy.proxy_server import general_settings_view
|
||||
|
||||
if general_settings_view().get("enable_workflows") is not True:
|
||||
raise_public(
|
||||
WorkflowFailure(code="workflow_disabled", message="Workflow execution is disabled", status_code=503)
|
||||
)
|
||||
try:
|
||||
invocation: Final = WorkflowInvocation.model_validate(data)
|
||||
except ValidationError:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="unsupported_workflow_request",
|
||||
message="Workflow requests accept messages, user, metadata.tags and non-streaming invocation only",
|
||||
)
|
||||
)
|
||||
store: Final = require_store(workflows)
|
||||
accepted: Final = await admit(store, auth, invocation, request.headers.get("Idempotency-Key"))
|
||||
return await wait_for_run(store, accepted)
|
||||
273
litellm/proxy/workflows/management.py
Normal file
273
litellm/proxy/workflows/management.py
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
from typing import Annotated, Final, NoReturn, TypeAlias
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Response
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from litellm.models.workflow import WorkflowConflict, WorkflowDraft, WorkflowStore
|
||||
from litellm.models.workflow_bundle import (
|
||||
ChatCompletionNode,
|
||||
WorkflowBundle,
|
||||
WorkflowFailure,
|
||||
canonical_bundle,
|
||||
workflow_digest,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, get_team_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.deps.repositories import WorkflowRepositoryDep
|
||||
from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, is_team_member
|
||||
from litellm.proxy.management.teams.dependencies import get_team_access
|
||||
from litellm.proxy.workflows.validation import validate_bundle
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
router: Final = APIRouter(tags=["Workflows"])
|
||||
Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
|
||||
|
||||
|
||||
class EmptyWorkflowAction(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
class RegisterWorkflow(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
bundle: WorkflowBundle
|
||||
|
||||
|
||||
class PublishWorkflow(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
digest: str
|
||||
|
||||
|
||||
class DraftView(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
workflow_id: str
|
||||
name: str
|
||||
team_id: str
|
||||
etag: str
|
||||
digest: str
|
||||
bundle: WorkflowBundle
|
||||
published_snapshot_id: str | None
|
||||
|
||||
|
||||
class DraftList(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
workflows: tuple[DraftView, ...]
|
||||
next_cursor: str | None
|
||||
|
||||
|
||||
class ValidationView(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
valid: bool
|
||||
digest: str
|
||||
|
||||
|
||||
class PublicationView(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
snapshot_id: str
|
||||
model: str
|
||||
digest: str
|
||||
|
||||
|
||||
def raise_public(failure: WorkflowFailure) -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=failure.status_code,
|
||||
detail={
|
||||
"type": "workflow_error",
|
||||
"code": failure.code,
|
||||
"message": failure.message,
|
||||
"param": failure.param,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def require_store(store: WorkflowStore | None) -> WorkflowStore:
|
||||
if store is None:
|
||||
raise_public(
|
||||
WorkflowFailure(code="database_required", message="Workflow execution requires a database", status_code=503)
|
||||
)
|
||||
return store
|
||||
|
||||
|
||||
async def require_team(auth: UserAPIKeyAuth, team_id: str | None, *, admin: bool = False, write: bool = False) -> str:
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
resolved_id: Final = team_id if team_id is not None else auth.team_id
|
||||
if write and auth.user_role in {LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY}:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="write_denied", message="View-only users cannot change or execute drafts", status_code=403
|
||||
)
|
||||
)
|
||||
if resolved_id is None or not auth.user_id:
|
||||
raise_public(
|
||||
WorkflowFailure(
|
||||
code="team_identity_required", message="A team and a bound user are required", status_code=403
|
||||
)
|
||||
)
|
||||
team: Final = await get_team_object(
|
||||
resolved_id,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
check_db_only=True,
|
||||
)
|
||||
allowed_admin: Final = await get_team_access().allows(auth, team, TEAM_OR_ORG_ADMIN)
|
||||
if not allowed_admin and (admin or not is_team_member(auth, team)):
|
||||
raise_public(WorkflowFailure(code="not_found", message="Workflow not found", status_code=404))
|
||||
return resolved_id
|
||||
|
||||
|
||||
async def require_draft(
|
||||
store: WorkflowStore, workflow_id: str, auth: UserAPIKeyAuth, *, admin: bool = False, write: bool = False
|
||||
) -> WorkflowDraft:
|
||||
draft: Final = await store.get_id(workflow_id)
|
||||
if draft is None:
|
||||
raise_public(WorkflowFailure(code="not_found", message="Workflow not found", status_code=404))
|
||||
await require_team(auth, draft.team_id, admin=admin, write=write)
|
||||
return draft
|
||||
|
||||
|
||||
def draft_view(draft: WorkflowDraft) -> DraftView:
|
||||
bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json)
|
||||
return DraftView(
|
||||
workflow_id=draft.workflow_id,
|
||||
name=draft.name,
|
||||
team_id=draft.team_id,
|
||||
etag=f'"{draft.etag}"',
|
||||
digest=workflow_digest(bundle),
|
||||
bundle=bundle,
|
||||
published_snapshot_id=draft.published_snapshot_id,
|
||||
)
|
||||
|
||||
|
||||
def require_etag(draft: WorkflowDraft, if_match: str | None) -> None:
|
||||
if if_match != f'"{draft.etag}"':
|
||||
raise_public(
|
||||
WorkflowFailure(code="stale_draft", message="If-Match must match the reviewed draft ETag", status_code=412)
|
||||
)
|
||||
|
||||
|
||||
async def check_bundle(bundle: WorkflowBundle, auth: UserAPIKeyAuth) -> None:
|
||||
from litellm.proxy.proxy_server import llm_model_list_view, llm_router
|
||||
|
||||
failure: Final = validate_bundle(bundle)
|
||||
if failure is not None:
|
||||
raise_public(failure)
|
||||
for node in bundle.nodes:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
if llm_router is None or node.model not in llm_router.registry().names(team_id=auth.team_id):
|
||||
raise_public(
|
||||
WorkflowFailure(code="unknown_model", message="Use a configured gateway model alias", param=node.id)
|
||||
)
|
||||
await can_key_call_resolved_model(node.model, llm_model_list_view(), auth, llm_router)
|
||||
|
||||
|
||||
@router.post("/v1/workflows", status_code=201, response_model=DraftView)
|
||||
async def create_workflow(
|
||||
data: RegisterWorkflow, response: Response, auth: Auth, workflows: WorkflowRepositoryDep, team_id: str | None = None
|
||||
) -> DraftView:
|
||||
"""Register a team-scoped draft without publishing or executing it."""
|
||||
team: Final = await require_team(auth, team_id, write=True)
|
||||
await check_bundle(data.bundle, auth)
|
||||
result: Final = await require_store(workflows).create(
|
||||
team, data.bundle.name, canonical_bundle(data.bundle), auth.user_id or ""
|
||||
)
|
||||
if isinstance(result, WorkflowConflict):
|
||||
raise_public(WorkflowFailure(code="workflow_conflict", message=result.message, status_code=result.status_code))
|
||||
view: Final = draft_view(result)
|
||||
response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers
|
||||
return view
|
||||
|
||||
|
||||
@router.get("/v1/workflows", response_model=DraftList)
|
||||
async def list_workflows(
|
||||
auth: Auth,
|
||||
workflows: WorkflowRepositoryDep,
|
||||
team_id: str | None = None,
|
||||
after: str | None = None,
|
||||
limit: int = Query(default=50, ge=1, le=100),
|
||||
) -> DraftList:
|
||||
team: Final = await require_team(auth, team_id)
|
||||
rows: Final = await require_store(workflows).list(team, after, limit + 1)
|
||||
return DraftList(
|
||||
workflows=tuple(draft_view(row) for row in rows[:limit]),
|
||||
next_cursor=rows[limit - 1].workflow_id if len(rows) > limit else None,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/v1/workflows/{workflow_id}", response_model=DraftView)
|
||||
async def get_workflow(workflow_id: str, response: Response, auth: Auth, workflows: WorkflowRepositoryDep) -> DraftView:
|
||||
view: Final = draft_view(await require_draft(require_store(workflows), workflow_id, auth))
|
||||
response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers
|
||||
return view
|
||||
|
||||
|
||||
@router.put("/v1/workflows/{workflow_id}", response_model=DraftView)
|
||||
async def replace_workflow(
|
||||
workflow_id: str,
|
||||
data: RegisterWorkflow,
|
||||
response: Response,
|
||||
auth: Auth,
|
||||
workflows: WorkflowRepositoryDep,
|
||||
if_match: Annotated[str | None, Header()] = None,
|
||||
) -> DraftView:
|
||||
store: Final = require_store(workflows)
|
||||
draft: Final = await require_draft(store, workflow_id, auth, write=True)
|
||||
require_etag(draft, if_match)
|
||||
if data.bundle.name != draft.name:
|
||||
raise_public(WorkflowFailure(code="immutable_name", message="Workflow names cannot be changed"))
|
||||
await check_bundle(data.bundle, auth)
|
||||
result: Final = await store.replace(
|
||||
draft.team_id, draft.name, draft.etag, canonical_bundle(data.bundle), auth.user_id or ""
|
||||
)
|
||||
if isinstance(result, WorkflowConflict):
|
||||
raise_public(WorkflowFailure(code="stale_draft", message=result.message, status_code=result.status_code))
|
||||
view: Final = draft_view(result)
|
||||
response.headers["ETag"] = view.etag # rebind-ok: FastAPI injects the response to set protocol headers
|
||||
return view
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/validate", response_model=ValidationView)
|
||||
async def validate_workflow(
|
||||
workflow_id: str,
|
||||
auth: Auth,
|
||||
workflows: WorkflowRepositoryDep,
|
||||
if_match: Annotated[str | None, Header()] = None,
|
||||
action: EmptyWorkflowAction | None = None,
|
||||
) -> ValidationView:
|
||||
"""Validate the reviewed draft without executing Python or calling models."""
|
||||
draft: Final = await require_draft(require_store(workflows), workflow_id, auth)
|
||||
require_etag(draft, if_match)
|
||||
bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json)
|
||||
await check_bundle(bundle, auth)
|
||||
return ValidationView(valid=True, digest=workflow_digest(bundle))
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/publish", response_model=PublicationView)
|
||||
async def publish_workflow(
|
||||
workflow_id: str,
|
||||
data: PublishWorkflow,
|
||||
auth: Auth,
|
||||
workflows: WorkflowRepositoryDep,
|
||||
if_match: Annotated[str | None, Header()] = None,
|
||||
) -> PublicationView:
|
||||
"""Publish an immutable snapshot after independent approval of the current ETag."""
|
||||
store: Final = require_store(workflows)
|
||||
draft: Final = await require_draft(store, workflow_id, auth, admin=True, write=True)
|
||||
require_etag(draft, if_match)
|
||||
bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json)
|
||||
if data.digest != workflow_digest(bundle):
|
||||
raise_public(
|
||||
WorkflowFailure(code="stale_draft", message="Reviewed digest does not match the draft", status_code=412)
|
||||
)
|
||||
await check_bundle(bundle, auth)
|
||||
result: Final = await store.publish(draft.team_id, draft.name, draft.etag, auth.user_id or "")
|
||||
if isinstance(result, WorkflowConflict):
|
||||
raise_public(WorkflowFailure(code="publication_denied", message=result.message, status_code=result.status_code))
|
||||
return PublicationView(snapshot_id=result.snapshot_id, model=f"workflow/{draft.name}", digest=data.digest)
|
||||
232
litellm/proxy/workflows/runner.py
Normal file
232
litellm/proxy/workflows/runner.py
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.models.workflow import (
|
||||
WorkflowCheckpoint,
|
||||
WorkflowConflict,
|
||||
WorkflowExecution,
|
||||
WorkflowRunState,
|
||||
WorkflowStep,
|
||||
WorkflowStore,
|
||||
)
|
||||
from litellm.models.workflow_bundle import (
|
||||
Binding,
|
||||
ChatCompletionNode,
|
||||
Constant,
|
||||
WorkflowBundle,
|
||||
WorkflowFailure,
|
||||
WorkflowNode,
|
||||
)
|
||||
from litellm.proxy.workflows.monty.compiler import compile_program
|
||||
from litellm.proxy.workflows.monty.models import MontyFailure
|
||||
from litellm.proxy.workflows.monty.runtime import MontyRuntime
|
||||
from litellm.proxy.workflows.validation import validate_value
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
class ChildResult(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
output: JsonValue
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
response_cost: float | None = None
|
||||
request_id: str | None = None
|
||||
|
||||
|
||||
class ChildCaller(Protocol):
|
||||
async def call(
|
||||
self, execution: WorkflowExecution, node: ChatCompletionNode, messages: JsonValue, remaining: float
|
||||
) -> ChildResult | WorkflowFailure: ...
|
||||
|
||||
|
||||
_JSON: Final = TypeAdapter[JsonValue](JsonValue)
|
||||
_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def resolve(binding: Binding, messages: JsonValue, outputs: Mapping[str, JsonValue]) -> JsonValue | WorkflowFailure:
|
||||
if isinstance(binding, Constant):
|
||||
return binding.value
|
||||
if binding.source != "$input" and binding.source not in outputs:
|
||||
return WorkflowFailure(
|
||||
code="missing_result", message="Referenced predecessor has no result", param=binding.source
|
||||
)
|
||||
value: Final = messages if binding.source == "$input" else outputs[binding.source]
|
||||
if binding.path in (None, ""):
|
||||
return value
|
||||
return project(value, tuple(part.replace("~1", "/").replace("~0", "~") for part in binding.path[1:].split("/")))
|
||||
|
||||
|
||||
def project(value: JsonValue, parts: tuple[str, ...]) -> JsonValue | WorkflowFailure:
|
||||
if not parts:
|
||||
return value
|
||||
key: Final = parts[0]
|
||||
if isinstance(value, dict) and key in value:
|
||||
return project(value[key], parts[1:])
|
||||
if isinstance(value, list) and key.isdigit() and str(int(key)) == key and int(key) < len(value):
|
||||
return project(value[int(key)], parts[1:])
|
||||
return WorkflowFailure(code="missing_path", message="Referenced JSON Pointer has no value")
|
||||
|
||||
|
||||
def node_input(
|
||||
node: WorkflowNode, messages: JsonValue, outputs: Mapping[str, JsonValue]
|
||||
) -> JsonValue | WorkflowFailure:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
return resolve(node.messages, messages, outputs)
|
||||
arguments: Final = {name: resolve(binding, messages, outputs) for name, binding in node.inputs.items()}
|
||||
failures: Final = tuple(value for value in arguments.values() if isinstance(value, WorkflowFailure))
|
||||
if failures:
|
||||
return failures[0]
|
||||
return _JSON.validate_python(arguments)
|
||||
|
||||
|
||||
async def execute_attempt(
|
||||
caller: ChildCaller,
|
||||
runtime: MontyRuntime,
|
||||
execution: WorkflowExecution,
|
||||
node: WorkflowNode,
|
||||
input: JsonValue,
|
||||
remaining: float,
|
||||
) -> ChildResult | WorkflowFailure:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
return await caller.call(execution, node, input, remaining)
|
||||
compiled: Final = compile_program(node.program, runtime.limits)
|
||||
if isinstance(compiled, MontyFailure):
|
||||
return WorkflowFailure(code="invalid_program", message=compiled.message, param=node.id)
|
||||
result: Final = await runtime.run_program(compiled, input, deadline_monotonic=time.monotonic() + remaining)
|
||||
if isinstance(result, MontyFailure):
|
||||
return WorkflowFailure(code="python_execution_failed", message=result.message, param=node.id)
|
||||
return ChildResult(output=result.output)
|
||||
|
||||
|
||||
async def execute_step(
|
||||
store: WorkflowStore,
|
||||
caller: ChildCaller,
|
||||
runtime: MontyRuntime,
|
||||
execution: WorkflowExecution,
|
||||
node: WorkflowNode,
|
||||
step: WorkflowStep,
|
||||
) -> ChildResult | WorkflowFailure:
|
||||
remaining: Final = (min(step.deadline, execution.deadline) - datetime.now(timezone.utc)).total_seconds()
|
||||
if remaining <= 0:
|
||||
return WorkflowFailure(
|
||||
code="node_deadline", message="Node total deadline exceeded", param=node.id, status_code=504
|
||||
)
|
||||
if step.attempts >= node.retry.max_attempts:
|
||||
return WorkflowFailure(
|
||||
code="attempts_exhausted", message="Node attempt policy exhausted", param=node.id, status_code=502
|
||||
)
|
||||
attempt: Final = await store.attempt_step(step.step_id, execution.lease_owner or "")
|
||||
if isinstance(attempt, WorkflowConflict):
|
||||
return WorkflowFailure(code="lease_lost", message=attempt.message, status_code=503)
|
||||
attempt_remaining: Final = (min(step.deadline, execution.deadline) - datetime.now(timezone.utc)).total_seconds()
|
||||
try:
|
||||
async with asyncio.timeout(max(attempt_remaining, 0)):
|
||||
result: Final = await execute_attempt(caller, runtime, execution, node, step.input, attempt_remaining)
|
||||
except TimeoutError:
|
||||
return WorkflowFailure(
|
||||
code="node_deadline", message="Node total deadline exceeded", param=node.id, status_code=504
|
||||
)
|
||||
if not isinstance(result, WorkflowFailure) or attempt.attempts >= node.retry.max_attempts:
|
||||
return result
|
||||
if result.status_code not in {429, 500, 502, 503, 504}:
|
||||
return result
|
||||
backoff: Final = min(
|
||||
node.retry.backoff_seconds * 2.0 ** (attempt.attempts - 1),
|
||||
max((step.deadline - datetime.now(timezone.utc)).total_seconds(), 0),
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
return await execute_step(store, caller, runtime, execution, node, attempt)
|
||||
|
||||
|
||||
async def execute_sequence(
|
||||
store: WorkflowStore,
|
||||
caller: ChildCaller,
|
||||
runtime: MontyRuntime,
|
||||
execution: WorkflowExecution,
|
||||
bundle: WorkflowBundle,
|
||||
messages: JsonValue,
|
||||
outputs: Mapping[str, JsonValue],
|
||||
index: int = 0,
|
||||
) -> JsonValue | WorkflowFailure:
|
||||
if index == len(bundle.nodes):
|
||||
return resolve(bundle.output, messages, outputs)
|
||||
current: Final = await store.execution(execution.run_id)
|
||||
if current is None or current.cancel_requested:
|
||||
return WorkflowFailure(code="cancelled", message="Run cancellation requested", status_code=409)
|
||||
if current.lease_owner != execution.lease_owner:
|
||||
return WorkflowFailure(code="lease_lost", message="Run lease lost", status_code=503)
|
||||
node: Final = bundle.nodes[index]
|
||||
input: Final = node_input(node, messages, outputs)
|
||||
if isinstance(input, WorkflowFailure):
|
||||
return input
|
||||
deadline: Final = min(execution.deadline, datetime.now(timezone.utc) + timedelta(seconds=node.timeout_seconds))
|
||||
step: Final = await store.start_step(execution.run_id, node.id, input, deadline, execution.lease_owner or "")
|
||||
if isinstance(step, WorkflowConflict):
|
||||
return WorkflowFailure(code="lease_lost", message=step.message, status_code=503)
|
||||
if step.state is WorkflowRunState.COMPLETED:
|
||||
return await execute_sequence(
|
||||
store, caller, runtime, execution, bundle, messages, {**outputs, node.id: step.output}, index + 1
|
||||
)
|
||||
result: Final = await execute_step(store, caller, runtime, execution, node, step)
|
||||
steps: Final = await store.steps(execution.run_id)
|
||||
attempted: Final = next(item for item in steps if item.step_id == step.step_id)
|
||||
if isinstance(result, WorkflowFailure):
|
||||
await store.finish_step(
|
||||
attempted.model_copy(
|
||||
update={
|
||||
"state": WorkflowRunState.FAILED,
|
||||
"error": result.code,
|
||||
"finished_at": datetime.now(timezone.utc),
|
||||
}
|
||||
),
|
||||
execution.lease_owner or "",
|
||||
)
|
||||
return result
|
||||
completed: Final = attempted.model_copy(
|
||||
update={
|
||||
"state": WorkflowRunState.COMPLETED,
|
||||
"output": result.output,
|
||||
"finished_at": datetime.now(timezone.utc),
|
||||
"prompt_tokens": result.prompt_tokens,
|
||||
"completion_tokens": result.completion_tokens,
|
||||
"response_cost": result.response_cost,
|
||||
"request_id": result.request_id,
|
||||
}
|
||||
)
|
||||
checkpointed: Final = await store.checkpoint(
|
||||
execution.run_id,
|
||||
WorkflowCheckpoint(
|
||||
node_path=node.id,
|
||||
input=step.input,
|
||||
output=result.output,
|
||||
attempts=attempted.attempts,
|
||||
),
|
||||
completed,
|
||||
execution.lease_owner or "",
|
||||
)
|
||||
if not checkpointed:
|
||||
return WorkflowFailure(code="lease_lost", message="Unable to commit node checkpoint", status_code=503)
|
||||
return await execute_sequence(
|
||||
store, caller, runtime, execution, bundle, messages, {**outputs, node.id: result.output}, index + 1
|
||||
)
|
||||
|
||||
|
||||
async def execute_run(
|
||||
store: WorkflowStore, caller: ChildCaller, runtime: MontyRuntime, execution: WorkflowExecution
|
||||
) -> JsonValue | WorkflowFailure:
|
||||
snapshot: Final = await store.snapshot(execution.snapshot_id)
|
||||
if snapshot is None:
|
||||
return WorkflowFailure(code="snapshot_missing", message="Accepted snapshot is unavailable", status_code=503)
|
||||
bundle: Final = WorkflowBundle.model_validate_json(snapshot.bundle_json)
|
||||
request: Final = _OBJECT.validate_json(execution.request_json)
|
||||
messages: Final = request.get("messages")
|
||||
failure: Final = validate_value(messages, bundle.input_schema, "messages")
|
||||
if failure is not None:
|
||||
return failure
|
||||
return await execute_sequence(store, caller, runtime, execution, bundle, messages, {})
|
||||
147
litellm/proxy/workflows/runs.py
Normal file
147
litellm/proxy/workflows/runs.py
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
from typing import Annotated, Final, Literal
|
||||
|
||||
from fastapi import APIRouter, Header, Request
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from litellm.models.workflow import WorkflowConflict, WorkflowExecution, WorkflowRunState, WorkflowStep, WorkflowStore
|
||||
from litellm.models.workflow_bundle import WorkflowBundle, WorkflowFailure
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.deps.repositories import WorkflowRepositoryDep
|
||||
from litellm.proxy.workflows.invocation import InvocationMetadata, WorkflowInvocation, admit, principal_hash
|
||||
from litellm.proxy.workflows.management import (
|
||||
Auth,
|
||||
EmptyWorkflowAction,
|
||||
check_bundle,
|
||||
raise_public,
|
||||
require_draft,
|
||||
require_etag,
|
||||
require_store,
|
||||
require_team,
|
||||
)
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
router: Final = APIRouter(tags=["Workflows"])
|
||||
|
||||
|
||||
class DraftTest(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
messages: tuple[ChatCompletionMessageParam, ...]
|
||||
user: str | None = None
|
||||
metadata: InvocationMetadata | None = None
|
||||
authorize_billing: Literal[True]
|
||||
|
||||
|
||||
class RunAccepted(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
run_id: str
|
||||
snapshot_id: str
|
||||
trace_url: str
|
||||
|
||||
|
||||
class RunSteps(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
run_id: str
|
||||
snapshot_id: str
|
||||
state: WorkflowRunState
|
||||
cancel_requested: bool
|
||||
steps: tuple[WorkflowStep, ...]
|
||||
|
||||
|
||||
class CancellationView(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
run_id: str
|
||||
cancel_requested: bool
|
||||
|
||||
|
||||
async def require_execution(
|
||||
store: WorkflowStore, run_id: str, auth: UserAPIKeyAuth, *, write: bool = False
|
||||
) -> WorkflowExecution:
|
||||
execution: Final = await store.execution(run_id)
|
||||
if execution is None:
|
||||
raise_public(WorkflowFailure(code="not_found", message="Workflow run not found", status_code=404))
|
||||
if principal_hash(auth) != execution.principal:
|
||||
await require_team(auth, execution.team_id, admin=True, write=write)
|
||||
elif auth.team_id != execution.team_id:
|
||||
raise_public(WorkflowFailure(code="not_found", message="Workflow run not found", status_code=404))
|
||||
return execution
|
||||
|
||||
|
||||
@router.get("/v1/workflows/runs/{run_id}/steps", response_model=RunSteps)
|
||||
async def read_steps(run_id: str, auth: Auth, workflows: WorkflowRepositoryDep) -> RunSteps:
|
||||
store: Final = require_store(workflows)
|
||||
execution: Final = await require_execution(store, run_id, auth)
|
||||
return RunSteps(
|
||||
run_id=run_id,
|
||||
snapshot_id=execution.snapshot_id,
|
||||
state=execution.state,
|
||||
cancel_requested=execution.cancel_requested,
|
||||
steps=await store.steps(run_id),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/v1/workflows/runs/{run_id}/steps/{step_id}", response_model=WorkflowStep)
|
||||
async def read_step(run_id: str, step_id: str, auth: Auth, workflows: WorkflowRepositoryDep) -> WorkflowStep:
|
||||
store: Final = require_store(workflows)
|
||||
await require_execution(store, run_id, auth)
|
||||
steps: Final = await store.steps(run_id)
|
||||
selected: Final = next((step for step in steps if step.step_id == step_id), None)
|
||||
if selected is None:
|
||||
raise_public(WorkflowFailure(code="not_found", message="Workflow step not found", status_code=404))
|
||||
return selected
|
||||
|
||||
|
||||
@router.post("/v1/workflows/runs/{run_id}/cancel", status_code=202, response_model=CancellationView)
|
||||
async def cancel_run(
|
||||
run_id: str, auth: Auth, workflows: WorkflowRepositoryDep, action: EmptyWorkflowAction | None = None
|
||||
) -> CancellationView:
|
||||
"""Request cancellation and prevent further node dispatch for this run."""
|
||||
store: Final = require_store(workflows)
|
||||
await require_execution(store, run_id, auth, write=True)
|
||||
await store.cancel(run_id)
|
||||
return CancellationView(run_id=run_id, cancel_requested=True)
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/test", status_code=202, response_model=RunAccepted)
|
||||
async def test_draft(
|
||||
workflow_id: str,
|
||||
data: DraftTest,
|
||||
request: Request,
|
||||
auth: Auth,
|
||||
workflows: WorkflowRepositoryDep,
|
||||
if_match: Annotated[str | None, Header()] = None,
|
||||
) -> RunAccepted:
|
||||
"""Execute the reviewed draft with explicit billing authorization."""
|
||||
from litellm.proxy.proxy_server import general_settings_view
|
||||
|
||||
if general_settings_view().get("enable_workflows") is not True:
|
||||
raise_public(
|
||||
WorkflowFailure(code="workflow_disabled", message="Workflow execution is disabled", status_code=503)
|
||||
)
|
||||
store: Final = require_store(workflows)
|
||||
draft: Final = await require_draft(store, workflow_id, auth, write=True)
|
||||
require_etag(draft, if_match)
|
||||
bundle: Final = WorkflowBundle.model_validate_json(draft.bundle_json)
|
||||
await check_bundle(bundle, auth)
|
||||
snapshot: Final = await store.pin_draft(draft.workflow_id, draft.etag, auth.user_id or "")
|
||||
if isinstance(snapshot, WorkflowConflict):
|
||||
raise_public(WorkflowFailure(code="stale_draft", message=snapshot.message, status_code=snapshot.status_code))
|
||||
accepted: Final = await admit(
|
||||
store,
|
||||
auth,
|
||||
WorkflowInvocation(
|
||||
model=f"workflow/{draft.name}", messages=data.messages, user=data.user, metadata=data.metadata
|
||||
),
|
||||
request.headers.get("Idempotency-Key"),
|
||||
snapshot=snapshot,
|
||||
operation="draft_test",
|
||||
)
|
||||
return RunAccepted(
|
||||
run_id=accepted.run_id,
|
||||
snapshot_id=accepted.snapshot_id,
|
||||
trace_url=f"/v1/workflows/runs/{accepted.run_id}/steps",
|
||||
)
|
||||
163
litellm/proxy/workflows/validation.py
Normal file
163
litellm/proxy/workflows/validation.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
from jsonschema import Draft202012Validator, SchemaError
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.models.workflow_bundle import (
|
||||
Binding,
|
||||
ChatCompletionNode,
|
||||
Constant,
|
||||
PythonNode,
|
||||
Reference,
|
||||
WorkflowBundle,
|
||||
WorkflowFailure,
|
||||
canonical_bundle,
|
||||
validate_bindings,
|
||||
)
|
||||
from litellm.proxy.workflows.monty.compiler import compile_program
|
||||
from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits
|
||||
|
||||
DEFAULT_MONTY_LIMITS: Final = MontyLimits()
|
||||
|
||||
|
||||
def schema_is_local(value: JsonValue, depth: int = 0) -> bool:
|
||||
if depth > 48:
|
||||
return False
|
||||
if isinstance(value, list):
|
||||
return all(schema_is_local(item, depth + 1) for item in value)
|
||||
if not isinstance(value, dict):
|
||||
return True
|
||||
return all(
|
||||
(key not in {"$ref", "$dynamicRef"} or isinstance(child, str) and child.startswith("#/$defs/"))
|
||||
and key not in {"$id", "$dynamicAnchor"}
|
||||
and schema_is_local(child, depth + 1)
|
||||
for key, child in value.items()
|
||||
)
|
||||
|
||||
|
||||
def schema_failure(schema: Mapping[str, JsonValue], param: str) -> WorkflowFailure | None:
|
||||
if not schema_is_local(dict(schema)):
|
||||
return WorkflowFailure(
|
||||
code="unsupported_schema", message="Only local $defs references are supported", param=param
|
||||
)
|
||||
try:
|
||||
Draft202012Validator.check_schema(schema)
|
||||
except SchemaError:
|
||||
return WorkflowFailure(code="invalid_schema", message="Invalid JSON Schema", param=param)
|
||||
return None
|
||||
|
||||
|
||||
def validate_value(value: JsonValue, schema: Mapping[str, JsonValue], param: str) -> WorkflowFailure | None:
|
||||
if next(Draft202012Validator(schema).iter_errors(value), None) is not None: # pyright: ignore[reportUnknownMemberType] # jsonschema stubs leave variadic arguments untyped
|
||||
return WorkflowFailure(
|
||||
code="schema_mismatch", message="Value does not match its declared JSON Schema", param=param
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def validate_bundle(bundle: WorkflowBundle) -> WorkflowFailure | None:
|
||||
if len(canonical_bundle(bundle).encode()) > 1_048_576:
|
||||
return WorkflowFailure(code="bundle_too_large", message="Workflow bundle exceeds 1 MiB")
|
||||
bindings_failure: Final = validate_bindings(bundle)
|
||||
if bindings_failure is not None:
|
||||
return bindings_failure
|
||||
input_failure: Final = schema_failure(bundle.input_schema, "input_schema")
|
||||
if input_failure is not None:
|
||||
return input_failure
|
||||
for node in bundle.nodes:
|
||||
if isinstance(node, ChatCompletionNode) and node.response_format is not None:
|
||||
response_failure = schema_failure(
|
||||
node.response_format.json_schema.schema_, f"nodes.{node.id}.response_format"
|
||||
)
|
||||
if response_failure is not None:
|
||||
return response_failure
|
||||
if isinstance(node, PythonNode):
|
||||
compiled = compile_program(node.program, DEFAULT_MONTY_LIMITS)
|
||||
if isinstance(compiled, MontyFailure):
|
||||
return WorkflowFailure(code="invalid_program", message=compiled.message, param=f"nodes.{node.id}")
|
||||
return compatibility_failure(bundle)
|
||||
|
||||
|
||||
def _types(
|
||||
schema: JsonValue | Mapping[str, JsonValue], root: Mapping[str, JsonValue], seen: frozenset[str] = frozenset()
|
||||
) -> frozenset[str]:
|
||||
if not isinstance(schema, Mapping):
|
||||
return frozenset()
|
||||
reference: Final = schema.get("$ref")
|
||||
if isinstance(reference, str) and reference.startswith("#/$defs/") and reference not in seen:
|
||||
definitions: Final = root.get("$defs")
|
||||
return (
|
||||
_types(definitions.get(reference[8:]), root, seen | {reference})
|
||||
if isinstance(definitions, dict)
|
||||
else frozenset()
|
||||
)
|
||||
type_: Final = schema.get("type")
|
||||
if isinstance(type_, str):
|
||||
return frozenset({type_})
|
||||
alternatives: Final = schema.get("anyOf", schema.get("oneOf"))
|
||||
if isinstance(alternatives, list):
|
||||
return frozenset(chain.from_iterable(_types(alternative, root, seen) for alternative in alternatives))
|
||||
return frozenset()
|
||||
|
||||
|
||||
def _result_schema(bundle: WorkflowBundle, source: str) -> Mapping[str, JsonValue]:
|
||||
if source == "$input":
|
||||
return bundle.input_schema
|
||||
node: Final = next(node for node in bundle.nodes if node.id == source)
|
||||
if isinstance(node, PythonNode):
|
||||
return node.program.output_schema
|
||||
return node.response_format.json_schema.schema_ if node.response_format is not None else {"type": "string"}
|
||||
|
||||
|
||||
def _binding_compatibility(
|
||||
bundle: WorkflowBundle, binding: Binding, expected: JsonValue, expected_root: Mapping[str, JsonValue], param: str
|
||||
) -> WorkflowFailure | None:
|
||||
if isinstance(binding, Constant) and isinstance(expected, dict):
|
||||
return validate_value(binding.value, {**expected, "$defs": expected_root.get("$defs", {})}, param)
|
||||
if not isinstance(binding, Reference) or binding.path not in (None, ""):
|
||||
return None
|
||||
actual: Final = _result_schema(bundle, binding.source)
|
||||
actual_types: Final = _types(actual, actual)
|
||||
expected_types: Final = _types(expected, expected_root)
|
||||
if (
|
||||
actual_types
|
||||
and expected_types
|
||||
and not actual_types & expected_types
|
||||
and not ("integer" in actual_types and "number" in expected_types)
|
||||
):
|
||||
return WorkflowFailure(
|
||||
code="incompatible_binding", message="Binding has an incompatible declared result type", param=param
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def compatibility_failure(bundle: WorkflowBundle) -> WorkflowFailure | None:
|
||||
for node in bundle.nodes:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
failure = _binding_compatibility(bundle, node.messages, {"type": "array"}, {}, f"nodes.{node.id}.messages")
|
||||
if failure is not None:
|
||||
return failure
|
||||
continue
|
||||
properties = node.program.input_schema.get("properties")
|
||||
if not isinstance(properties, dict):
|
||||
return WorkflowFailure(
|
||||
code="invalid_inputs", message="Function argument schema must declare properties", param=node.id
|
||||
)
|
||||
for name, binding in node.inputs.items():
|
||||
input_failure = _binding_compatibility(
|
||||
bundle,
|
||||
binding,
|
||||
properties.get(name),
|
||||
node.program.input_schema,
|
||||
f"nodes.{node.id}.inputs.{name}",
|
||||
)
|
||||
if input_failure is not None:
|
||||
return input_failure
|
||||
return None
|
||||
|
||||
|
||||
def canonical_request(value: Mapping[str, JsonValue]) -> str:
|
||||
return json.dumps(dict(value), sort_keys=True, separators=(",", ":"), allow_nan=False)
|
||||
438
litellm/repositories/workflow_repository.py
Normal file
438
litellm/repositories/workflow_repository.py
Normal file
|
|
@ -0,0 +1,438 @@
|
|||
import json
|
||||
from collections.abc import Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # PrismaWrapper forwards transaction operations dynamically
|
||||
)
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.models.workflow import (
|
||||
WorkflowCheckpoint,
|
||||
WorkflowConflict,
|
||||
WorkflowDraft,
|
||||
WorkflowExecution,
|
||||
WorkflowRunState,
|
||||
WorkflowSnapshot,
|
||||
WorkflowStep,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class WorkflowTransaction(Protocol):
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int: ...
|
||||
|
||||
|
||||
class WorkflowDatabase(WorkflowTransaction, Protocol):
|
||||
def tx(self) -> AbstractAsyncContextManager[WorkflowTransaction]: ...
|
||||
|
||||
|
||||
_DRAFTS: Final = TypeAdapter(tuple[WorkflowDraft, ...])
|
||||
_SNAPSHOTS: Final = TypeAdapter(tuple[WorkflowSnapshot, ...])
|
||||
_EXECUTIONS: Final = TypeAdapter(tuple[WorkflowExecution, ...])
|
||||
_CHECKPOINTS: Final = TypeAdapter(tuple[WorkflowCheckpoint, ...])
|
||||
_STEPS: Final = TypeAdapter(tuple[WorkflowStep, ...])
|
||||
_DRAFT_COLUMNS: Final = "workflow_id, team_id, name, etag, bundle_json, contributors, published_snapshot_id"
|
||||
_SNAPSHOT_COLUMNS: Final = "snapshot_id, workflow_id, team_id, name, bundle_json, published_by, created_at, origin"
|
||||
_EXECUTION_COLUMNS: Final = (
|
||||
"run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, state, "
|
||||
"deadline, lease_owner, lease_until, cancel_requested, checkpoints, response_json, error"
|
||||
)
|
||||
_STEP_COLUMNS: Final = (
|
||||
"step_id, run_id, node_id, state, input, output, attempts, deadline, started_at, finished_at, error, "
|
||||
"prompt_tokens, completion_tokens, response_cost, request_id"
|
||||
)
|
||||
|
||||
|
||||
_RowT: Final = TypeVar("_RowT")
|
||||
|
||||
|
||||
def _first(rows: Sequence[_RowT]) -> _RowT | None:
|
||||
return rows[0] if rows else None
|
||||
|
||||
|
||||
class WorkflowRepository:
|
||||
def __init__(self, prisma_client: "PrismaClient") -> None:
|
||||
self._database: Final = cast( # cast-ok: Prisma exposes raw SQL and transaction methods dynamically
|
||||
WorkflowDatabase, prisma_client.writer_db
|
||||
)
|
||||
|
||||
async def create(self, team_id: str, name: str, bundle_json: str, author: str) -> WorkflowDraft | WorkflowConflict:
|
||||
rows: Final = _DRAFTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowDefinition" '
|
||||
"(workflow_id, team_id, name, etag, bundle_json, contributors) "
|
||||
"VALUES ($1, $2, $3, $4, $5, ARRAY[$6]::text[]) ON CONFLICT (team_id, name) DO NOTHING "
|
||||
f"RETURNING {_DRAFT_COLUMNS}",
|
||||
str(uuid.uuid4()),
|
||||
team_id,
|
||||
name,
|
||||
str(uuid.uuid4()),
|
||||
bundle_json,
|
||||
author,
|
||||
)
|
||||
)
|
||||
return rows[0] if rows else WorkflowConflict(message="Workflow already exists")
|
||||
|
||||
async def get(self, team_id: str, name: str) -> WorkflowDraft | None:
|
||||
return _first(
|
||||
_DRAFTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2',
|
||||
team_id,
|
||||
name,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def replace(
|
||||
self, team_id: str, name: str, etag: str, bundle_json: str, author: str
|
||||
) -> WorkflowDraft | WorkflowConflict:
|
||||
rows: Final = _DRAFTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'UPDATE "LiteLLM_WorkflowDefinition" SET etag = $4, bundle_json = $5, '
|
||||
"contributors = ARRAY(SELECT DISTINCT unnest(contributors || ARRAY[$6]::text[])) "
|
||||
f"WHERE team_id = $1 AND name = $2 AND etag = $3 RETURNING {_DRAFT_COLUMNS}",
|
||||
team_id,
|
||||
name,
|
||||
etag,
|
||||
str(uuid.uuid4()),
|
||||
bundle_json,
|
||||
author,
|
||||
)
|
||||
)
|
||||
return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412)
|
||||
|
||||
async def get_id(self, workflow_id: str) -> WorkflowDraft | None:
|
||||
return _first(
|
||||
_DRAFTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1',
|
||||
workflow_id,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]:
|
||||
return _DRAFTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" '
|
||||
"WHERE team_id = $1 AND workflow_id > $2 ORDER BY workflow_id LIMIT $3",
|
||||
team_id,
|
||||
after or "",
|
||||
limit,
|
||||
)
|
||||
)
|
||||
|
||||
async def publish(self, team_id: str, name: str, etag: str, publisher: str) -> WorkflowSnapshot | WorkflowConflict:
|
||||
async with self._database.tx() as tx:
|
||||
drafts: Final = _DRAFTS.validate_python(
|
||||
await tx.query_raw(
|
||||
f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" '
|
||||
"WHERE team_id = $1 AND name = $2 FOR UPDATE",
|
||||
team_id,
|
||||
name,
|
||||
)
|
||||
)
|
||||
if not drafts or drafts[0].etag != etag:
|
||||
return WorkflowConflict(message="Draft changed or does not exist", status_code=412)
|
||||
if publisher in drafts[0].contributors:
|
||||
return WorkflowConflict(message="Draft contributors cannot publish their own draft", status_code=403)
|
||||
rows: Final = _SNAPSHOTS.validate_python(
|
||||
await tx.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowSnapshot" '
|
||||
"(snapshot_id, workflow_id, team_id, name, bundle_json, published_by) "
|
||||
f"VALUES ($1, $2, $3, $4, $5, $6) RETURNING {_SNAPSHOT_COLUMNS}",
|
||||
str(uuid.uuid4()),
|
||||
drafts[0].workflow_id,
|
||||
team_id,
|
||||
name,
|
||||
drafts[0].bundle_json,
|
||||
publisher,
|
||||
)
|
||||
)
|
||||
await tx.execute_raw(
|
||||
'UPDATE "LiteLLM_WorkflowDefinition" SET published_snapshot_id = $3, etag = $4 '
|
||||
"WHERE team_id = $1 AND name = $2",
|
||||
team_id,
|
||||
name,
|
||||
rows[0].snapshot_id,
|
||||
str(uuid.uuid4()),
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
async def snapshot(self, snapshot_id: str) -> WorkflowSnapshot | None:
|
||||
return _first(
|
||||
_SNAPSHOTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = $1',
|
||||
snapshot_id,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def published(self, team_id: str, name: str) -> WorkflowSnapshot | None:
|
||||
return _first(
|
||||
_SNAPSHOTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = '
|
||||
'(SELECT published_snapshot_id FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2)',
|
||||
team_id,
|
||||
name,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def pin_draft(self, workflow_id: str, etag: str, requested_by: str) -> WorkflowSnapshot | WorkflowConflict:
|
||||
rows: Final = _SNAPSHOTS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowSnapshot" '
|
||||
"(snapshot_id, workflow_id, team_id, name, bundle_json, published_by, origin) "
|
||||
"SELECT $3, workflow_id, team_id, name, bundle_json, $4, 'draft_test' "
|
||||
'FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND etag = $2 '
|
||||
f"RETURNING {_SNAPSHOT_COLUMNS}",
|
||||
workflow_id,
|
||||
etag,
|
||||
str(uuid.uuid4()),
|
||||
requested_by,
|
||||
)
|
||||
)
|
||||
return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412)
|
||||
|
||||
async def by_dispatch_key(self, dispatch_key: str) -> WorkflowExecution | None:
|
||||
return _first(
|
||||
_EXECUTIONS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1',
|
||||
dispatch_key,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def execution(self, run_id: str) -> WorkflowExecution | None:
|
||||
return _first(
|
||||
_EXECUTIONS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE run_id = $1',
|
||||
run_id,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def dispatch(self, execution: WorkflowExecution) -> WorkflowExecution:
|
||||
async with self._database.tx() as tx:
|
||||
await tx.query_raw(
|
||||
"SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended($1, 0))", execution.dispatch_key
|
||||
)
|
||||
existing: Final = _EXECUTIONS.validate_python(
|
||||
await tx.query_raw(
|
||||
f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1',
|
||||
execution.dispatch_key,
|
||||
)
|
||||
)
|
||||
if existing:
|
||||
return existing[0]
|
||||
await tx.execute_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowRun" '
|
||||
"(run_id, session_id, workflow_type, created_by, input, metadata, engine_owned, updated_at) "
|
||||
"VALUES ($1, $1, 'liteflow', $2, $3::jsonb, $4::jsonb, true, NOW())",
|
||||
execution.run_id,
|
||||
execution.principal,
|
||||
execution.request_json,
|
||||
json.dumps(
|
||||
{
|
||||
"snapshot_id": execution.snapshot_id,
|
||||
"team_id": execution.team_id,
|
||||
"workflow_name": execution.name,
|
||||
}
|
||||
),
|
||||
)
|
||||
rows: Final = _EXECUTIONS.validate_python(
|
||||
await tx.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowExecution" '
|
||||
"(run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, deadline) "
|
||||
f"VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9::timestamptz AT TIME ZONE 'UTC') RETURNING {_EXECUTION_COLUMNS}",
|
||||
execution.run_id,
|
||||
execution.snapshot_id,
|
||||
execution.team_id,
|
||||
execution.name,
|
||||
execution.principal,
|
||||
execution.dispatch_key,
|
||||
execution.request_digest,
|
||||
execution.request_json,
|
||||
execution.deadline,
|
||||
)
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
async def lease_next(self, owner: str) -> WorkflowExecution | None:
|
||||
return _first(
|
||||
_EXECUTIONS.validate_python(
|
||||
await self._database.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowExecution\" SET lease_owner = $1, lease_until = NOW() + INTERVAL '30 seconds', "
|
||||
"state = 'running' WHERE run_id = (SELECT run_id FROM \"LiteLLM_WorkflowExecution\" "
|
||||
"WHERE state IN ('pending', 'running') AND (lease_until IS NULL OR lease_until < NOW()) "
|
||||
f"ORDER BY deadline FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING {_EXECUTION_COLUMNS}",
|
||||
owner,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def renew(self, run_id: str, owner: str) -> bool:
|
||||
return bool(
|
||||
await self._database.execute_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowExecution\" SET lease_until = NOW() + INTERVAL '30 seconds' "
|
||||
"WHERE run_id = $1 AND lease_owner = $2 AND lease_until > NOW() AND state = 'running'",
|
||||
run_id,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
|
||||
async def steps(self, run_id: str) -> tuple[WorkflowStep, ...]:
|
||||
return _STEPS.validate_python(
|
||||
await self._database.query_raw(
|
||||
f'SELECT {_STEP_COLUMNS} FROM "LiteLLM_WorkflowStep" WHERE run_id = $1 ORDER BY started_at, step_id',
|
||||
run_id,
|
||||
)
|
||||
)
|
||||
|
||||
async def start_step(
|
||||
self, run_id: str, node_id: str, input: JsonValue, deadline: datetime, owner: str
|
||||
) -> WorkflowStep | WorkflowConflict:
|
||||
rows: Final = _STEPS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowStep" (step_id, run_id, node_id, input, deadline) '
|
||||
"SELECT $1,$2,$3,$4::jsonb,$5::timestamptz AT TIME ZONE 'UTC' WHERE EXISTS (SELECT 1 FROM \"LiteLLM_WorkflowExecution\" "
|
||||
"WHERE run_id=$2 AND lease_owner=$6 AND lease_until>NOW() AND NOT cancel_requested) "
|
||||
"ON CONFLICT (run_id, node_id) DO UPDATE SET node_id = EXCLUDED.node_id "
|
||||
f"RETURNING {_STEP_COLUMNS}",
|
||||
str(uuid.uuid4()),
|
||||
run_id,
|
||||
node_id,
|
||||
json.dumps(input, allow_nan=False),
|
||||
deadline,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
return rows[0] if rows else WorkflowConflict(message="Execution lease lost")
|
||||
|
||||
async def attempt_step(self, step_id: str, owner: str) -> WorkflowStep | WorkflowConflict:
|
||||
rows: Final = _STEPS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'UPDATE "LiteLLM_WorkflowStep" SET attempts = attempts + 1 '
|
||||
'WHERE step_id = $1 AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e '
|
||||
'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$2 '
|
||||
f"AND e.lease_until>NOW() AND NOT e.cancel_requested) RETURNING {_STEP_COLUMNS}",
|
||||
step_id,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
return rows[0] if rows else WorkflowConflict(message="Execution lease lost")
|
||||
|
||||
async def _finish_step(self, database: WorkflowTransaction, step: WorkflowStep, owner: str) -> bool:
|
||||
return bool(
|
||||
await database.execute_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowStep\" SET state=$2, output=$3::jsonb, finished_at=$4::timestamptz AT TIME ZONE 'UTC', error=$5, "
|
||||
"prompt_tokens=$6, completion_tokens=$7, response_cost=$8, request_id=$9 WHERE step_id=$1 "
|
||||
'AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e '
|
||||
'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$10 '
|
||||
"AND e.lease_until>NOW() AND NOT e.cancel_requested)",
|
||||
step.step_id,
|
||||
step.state.value,
|
||||
json.dumps(step.output, allow_nan=False),
|
||||
step.finished_at,
|
||||
step.error,
|
||||
step.prompt_tokens,
|
||||
step.completion_tokens,
|
||||
step.response_cost,
|
||||
step.request_id,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
|
||||
async def finish_step(self, step: WorkflowStep, owner: str) -> bool:
|
||||
return await self._finish_step(self._database, step, owner)
|
||||
|
||||
async def checkpoint(self, run_id: str, checkpoint: WorkflowCheckpoint, step: WorkflowStep, owner: str) -> bool:
|
||||
async with self._database.tx() as tx:
|
||||
rows: Final = _EXECUTIONS.validate_python(
|
||||
await tx.query_raw(
|
||||
f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" '
|
||||
"WHERE run_id = $1 AND lease_owner=$2 AND lease_until>NOW() AND NOT cancel_requested FOR UPDATE",
|
||||
run_id,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
if not rows or not await self._finish_step(tx, step, owner):
|
||||
return False
|
||||
if any(item.node_path == checkpoint.node_path for item in rows[0].checkpoints):
|
||||
return True
|
||||
checkpoints: Final = (*rows[0].checkpoints, checkpoint)
|
||||
await tx.execute_raw(
|
||||
'UPDATE "LiteLLM_WorkflowExecution" SET checkpoints = $2::jsonb WHERE run_id = $1',
|
||||
run_id,
|
||||
_CHECKPOINTS.dump_json(checkpoints).decode(),
|
||||
)
|
||||
await tx.execute_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowEvent" '
|
||||
"(event_id, run_id, event_type, step_name, sequence_number, data) "
|
||||
"VALUES ($1,$2,'node.completed',$3,$4,$5::jsonb)",
|
||||
str(uuid.uuid4()),
|
||||
run_id,
|
||||
checkpoint.node_path,
|
||||
len(checkpoints) - 1,
|
||||
checkpoint.model_dump_json(),
|
||||
)
|
||||
return True
|
||||
|
||||
async def finish(
|
||||
self,
|
||||
run_id: str,
|
||||
state: WorkflowRunState,
|
||||
response_json: str | None = None,
|
||||
error: str | None = None,
|
||||
*,
|
||||
owner: str,
|
||||
) -> bool:
|
||||
async with self._database.tx() as tx:
|
||||
updated: Final = _EXECUTIONS.validate_python(
|
||||
await tx.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowExecution\" SET state = CASE WHEN cancel_requested THEN 'cancelled' ELSE $2 END, "
|
||||
"response_json = CASE WHEN cancel_requested THEN NULL ELSE $3 END, error = $4 "
|
||||
f"WHERE run_id = $1 AND state IN ($5,$6) AND lease_owner=$7 AND lease_until>NOW() RETURNING {_EXECUTION_COLUMNS}",
|
||||
run_id,
|
||||
state.value,
|
||||
response_json,
|
||||
error,
|
||||
WorkflowRunState.RUNNING.value,
|
||||
WorkflowRunState.PENDING.value,
|
||||
owner,
|
||||
)
|
||||
)
|
||||
if not updated:
|
||||
return False
|
||||
await tx.execute_raw(
|
||||
'UPDATE "LiteLLM_WorkflowRun" SET status = $2, output = $3::jsonb, updated_at = NOW() WHERE run_id = $1',
|
||||
run_id,
|
||||
updated[0].state.value,
|
||||
json.dumps({"response": json.loads(updated[0].response_json) if updated[0].response_json else None}),
|
||||
)
|
||||
return True
|
||||
|
||||
async def cancel(self, run_id: str) -> None:
|
||||
await self._database.execute_raw(
|
||||
'UPDATE "LiteLLM_WorkflowExecution" SET cancel_requested = true WHERE run_id = $1 AND state IN ($2,$3)',
|
||||
run_id,
|
||||
WorkflowRunState.PENDING.value,
|
||||
WorkflowRunState.RUNNING.value,
|
||||
)
|
||||
405
tests/integration/management/test_workflow_execution.py
Normal file
405
tests/integration/management/test_workflow_execution.py
Normal file
|
|
@ -0,0 +1,405 @@
|
|||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
_SOURCE: Final = (
|
||||
"def prepare_followup(text: str) -> list[dict[str, str]]:\n"
|
||||
' return [{"role": "user", "content": "Review this draft: " + text}]\n'
|
||||
)
|
||||
_MESSAGES: Final[Mapping[str, JsonValue]] = {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {"role": {"type": "string", "enum": ["user"]}, "content": {"type": "string"}},
|
||||
"required": ["role", "content"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class Draft(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
workflow_id: str
|
||||
etag: str
|
||||
digest: str
|
||||
|
||||
|
||||
class Step(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
step_id: str
|
||||
state: str
|
||||
input: JsonValue
|
||||
output: JsonValue
|
||||
attempts: int
|
||||
response_cost: float | None
|
||||
|
||||
|
||||
class Trace(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
run_id: str
|
||||
snapshot_id: str
|
||||
state: str
|
||||
steps: tuple[Step, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Actors:
|
||||
author: str
|
||||
publisher: str
|
||||
caller: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def workflow_gateway(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
with owned_proxy(
|
||||
session_gateway,
|
||||
tmp_path_factory.mktemp("workflows"),
|
||||
{},
|
||||
workers=2,
|
||||
config=Path(__file__).with_name("workflow_proxy_config.yaml"),
|
||||
) as gateway:
|
||||
yield gateway
|
||||
|
||||
|
||||
def _actors(scenario: Scenario, model: str) -> Actors:
|
||||
author: Final = scenario.user(user_role="internal_user")
|
||||
publisher: Final = scenario.user(user_role="internal_user")
|
||||
team: Final = scenario.team(
|
||||
models=[model, "workflow/email-reply"],
|
||||
members_with_roles=[{"role": "admin", "user_id": author}, {"role": "admin", "user_id": publisher}],
|
||||
)
|
||||
return Actors(
|
||||
scenario.key(team_id=team, user_id=author, models=[model]),
|
||||
scenario.key(team_id=team, user_id=publisher, models=[model]),
|
||||
scenario.key(team_id=team, user_id=author, models=["workflow/email-reply"]),
|
||||
)
|
||||
|
||||
|
||||
def _bundle(model: str, *, structured: bool = False, timeout: float = 60) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"contract_version": "1",
|
||||
"name": "email-reply",
|
||||
"timeout_seconds": timeout,
|
||||
"input_schema": dict(_MESSAGES),
|
||||
"nodes": [
|
||||
{
|
||||
"id": "draft",
|
||||
"type": "chat_completion",
|
||||
"model": model,
|
||||
"messages": {"from": "$input", "output": "messages"},
|
||||
"timeout_seconds": timeout,
|
||||
},
|
||||
{
|
||||
"id": "prepare",
|
||||
"type": "python",
|
||||
"program": {
|
||||
"contract_version": "1",
|
||||
"source": _SOURCE,
|
||||
"function_name": "prepare_followup",
|
||||
"parameter_names": ["text"],
|
||||
"source_sha256": hashlib.sha256(_SOURCE.encode()).hexdigest(),
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"output_schema": dict(_MESSAGES),
|
||||
},
|
||||
"inputs": {"text": {"from": "draft", "output": "result"}},
|
||||
},
|
||||
{
|
||||
"id": "review",
|
||||
"type": "chat_completion",
|
||||
"model": model,
|
||||
"messages": {"from": "prepare", "output": "result"},
|
||||
**(
|
||||
{
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "Review",
|
||||
"strict": True,
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"approved": {"type": "boolean"}},
|
||||
"required": ["approved"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
if structured
|
||||
else {}
|
||||
),
|
||||
},
|
||||
],
|
||||
"output": {"from": "review", "output": "result"},
|
||||
}
|
||||
|
||||
|
||||
def _completion(request: Request) -> Reply:
|
||||
body: Final = _OBJECT.validate_json(request.body)
|
||||
assert request.target == "/v1/chat/completions", request.target
|
||||
assert body["messages"] in (
|
||||
[{"role": "user", "content": "Help me update my account email."}],
|
||||
[{"role": "user", "content": "Review this draft: synthetic draft"}],
|
||||
), body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-controlled-" + uuid.uuid4().hex,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": body["model"],
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": '{"approved":true}' if "response_format" in body else "synthetic draft",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _register(gateway: Gateway, actors: Actors, model: str, *, structured: bool = False, timeout: float = 60) -> Draft:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/workflows",
|
||||
{"bundle": dict(_bundle(model, structured=structured, timeout=timeout))},
|
||||
key=actors.author,
|
||||
)
|
||||
assert response.status_code == 201, response.text
|
||||
return Draft.model_validate_json(response.content)
|
||||
|
||||
|
||||
def _publish(gateway: Gateway, actors: Actors, draft: Draft) -> None:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1/workflows/{draft.workflow_id}/publish",
|
||||
{"digest": draft.digest},
|
||||
key=actors.publisher,
|
||||
headers={"If-Match": draft.etag},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _request() -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"model": "workflow/email-reply",
|
||||
"messages": [{"role": "user", "content": "Help me update my account email."}],
|
||||
}
|
||||
|
||||
|
||||
def _trace(gateway: Gateway, actors: Actors, run_id: str) -> Trace:
|
||||
response: Final = gateway.request("GET", f"/v1/workflows/runs/{run_id}/steps", key=actors.caller)
|
||||
assert response.status_code == 200, response.text
|
||||
return Trace.model_validate_json(response.content)
|
||||
|
||||
|
||||
def test_draft_execution_requires_its_owning_team_key_even_for_multiteam_members(workflow_gateway: Gateway) -> None:
|
||||
with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire:
|
||||
model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1")
|
||||
author: Final = scenario.user(user_role="internal_user")
|
||||
members: Final = [{"role": "admin", "user_id": author}]
|
||||
owning_team: Final = scenario.team(models=[model], members_with_roles=members)
|
||||
other_team: Final = scenario.team(models=[model], members_with_roles=members)
|
||||
owning_key: Final = scenario.key(team_id=owning_team, user_id=author, models=[model])
|
||||
other_key: Final = scenario.key(team_id=other_team, user_id=author, models=[model])
|
||||
actors: Final = Actors(owning_key, owning_key, owning_key)
|
||||
draft: Final = _register(workflow_gateway, actors, model)
|
||||
refused: Final = workflow_gateway.request(
|
||||
"POST",
|
||||
f"/v1/workflows/{draft.workflow_id}/test",
|
||||
{"messages": _request()["messages"], "authorize_billing": True},
|
||||
key=other_key,
|
||||
headers={"If-Match": draft.etag},
|
||||
)
|
||||
assert refused.status_code == 403, refused.text
|
||||
assert (
|
||||
workflow_gateway.rows(
|
||||
'SELECT run_id FROM "LiteLLM_WorkflowExecution" WHERE principal=%s',
|
||||
(hashlib.sha256(other_key.encode()).hexdigest(),),
|
||||
)
|
||||
== []
|
||||
)
|
||||
assert wire.drain() == ()
|
||||
accepted: Final = workflow_gateway.request(
|
||||
"POST",
|
||||
f"/v1/workflows/{draft.workflow_id}/test",
|
||||
{"messages": _request()["messages"], "authorize_billing": True},
|
||||
key=owning_key,
|
||||
headers={"If-Match": draft.etag},
|
||||
)
|
||||
assert accepted.status_code == 202, accepted.text
|
||||
run_id: Final = TypeAdapter(str).validate_python(_OBJECT.validate_json(accepted.content)["run_id"])
|
||||
trace: Final = eventually(
|
||||
lambda: _trace(workflow_gateway, actors, run_id),
|
||||
lambda result: result.state == "completed",
|
||||
seconds=70,
|
||||
)
|
||||
assert tuple(step.state for step in trace.steps) == ("completed",) * 3
|
||||
|
||||
|
||||
def test_etags_and_independent_publication_do_not_execute_drafts(workflow_gateway: Gateway) -> None:
|
||||
with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire:
|
||||
model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1")
|
||||
actors: Final = _actors(scenario, model)
|
||||
draft: Final = _register(workflow_gateway, actors, model)
|
||||
validation: Final = workflow_gateway.request(
|
||||
"POST", f"/v1/workflows/{draft.workflow_id}/validate", key=actors.author, headers={"If-Match": draft.etag}
|
||||
)
|
||||
assert validation.status_code == 200, validation.text
|
||||
self_publication: Final = workflow_gateway.request(
|
||||
"POST",
|
||||
f"/v1/workflows/{draft.workflow_id}/publish",
|
||||
{"digest": draft.digest},
|
||||
key=actors.author,
|
||||
headers={"If-Match": draft.etag},
|
||||
)
|
||||
assert self_publication.status_code == 403, self_publication.text
|
||||
for method, suffix, body in (
|
||||
("PUT", "", {"bundle": dict(_bundle(model))}),
|
||||
("POST", "/publish", {"digest": draft.digest}),
|
||||
):
|
||||
stale = workflow_gateway.request(
|
||||
method,
|
||||
f"/v1/workflows/{draft.workflow_id}{suffix}",
|
||||
body,
|
||||
key=actors.publisher,
|
||||
headers={"If-Match": '"stale"'},
|
||||
)
|
||||
assert stale.status_code == 412, stale.text
|
||||
listed: Final = workflow_gateway.request("GET", "/v1/workflows", key=actors.author)
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert draft.workflow_id in listed.text
|
||||
assert wire.drain() == ()
|
||||
_publish(workflow_gateway, actors, draft)
|
||||
invoked: Final = workflow_gateway.request("POST", "/v1/chat/completions", _request(), key=actors.caller)
|
||||
assert invoked.status_code == 200, invoked.text
|
||||
assert len(wire.drain()) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("structured", [False, True])
|
||||
def test_typed_checkpoints_idempotency_bounded_delegation_and_single_child_billing(
|
||||
workflow_gateway: Gateway,
|
||||
structured: bool,
|
||||
) -> None:
|
||||
with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire:
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-6-sol", api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
|
||||
)
|
||||
actors: Final = _actors(scenario, model)
|
||||
draft: Final = _register(workflow_gateway, actors, model, structured=structured)
|
||||
_publish(workflow_gateway, actors, draft)
|
||||
denied: Final = workflow_gateway.request(
|
||||
"POST", "/v1/chat/completions", {**_request(), "model": model}, key=actors.caller
|
||||
)
|
||||
assert denied.status_code in (401, 403), denied.text
|
||||
response: Final = workflow_gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
_request(),
|
||||
key=actors.caller,
|
||||
headers={"Idempotency-Key": "controlled-retry"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
run_id: Final = response.headers["x-litellm-workflow-run-id"]
|
||||
retried: Final = workflow_gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
_request(),
|
||||
key=actors.caller,
|
||||
headers={"Idempotency-Key": "controlled-retry"},
|
||||
)
|
||||
assert retried.status_code == 200, retried.text
|
||||
assert retried.headers["x-litellm-workflow-run-id"] == run_id
|
||||
assert retried.json() == response.json()
|
||||
trace: Final = _trace(workflow_gateway, actors, run_id)
|
||||
assert trace.state == "completed"
|
||||
assert tuple(step.state for step in trace.steps) == ("completed",) * 3
|
||||
assert trace.steps[1].input == {"text": "synthetic draft"}
|
||||
assert trace.steps[1].output == [{"role": "user", "content": "Review this draft: synthetic draft"}]
|
||||
assert trace.steps[2].output == ({"approved": True} if structured else "synthetic draft")
|
||||
assert tuple(step.attempts for step in trace.steps) == (1, 1, 1)
|
||||
assert trace.steps[0].response_cost == pytest.approx(2 * 0.001 + 3 * 0.002)
|
||||
checkpoints: Final = workflow_gateway.rows(
|
||||
'SELECT checkpoints FROM "LiteLLM_WorkflowExecution" WHERE run_id=%s', (run_id,)
|
||||
)
|
||||
assert len(TypeAdapter(tuple[object, ...]).validate_python(checkpoints[0]["checkpoints"])) == 3
|
||||
assert len(wire.drain()) == 2
|
||||
charges: Final = eventually(
|
||||
lambda: workflow_gateway.rows(
|
||||
'SELECT spend, model FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (run_id,)
|
||||
),
|
||||
lambda rows: len(rows) == 2,
|
||||
seconds=70,
|
||||
)
|
||||
assert sum(TypeAdapter(float).validate_python(row["spend"]) for row in charges) == pytest.approx(
|
||||
2 * (2 * 0.001 + 3 * 0.002)
|
||||
)
|
||||
assert all(row["model"] != "workflow/email-reply" for row in charges)
|
||||
mutation: Final = workflow_gateway.request("PATCH", f"/v1/workflows/runs/{run_id}", {"status": "failed"})
|
||||
assert mutation.status_code == 409, mutation.text
|
||||
|
||||
|
||||
def test_cancellation_stops_the_next_node_without_cancelling_disconnected_work(workflow_gateway: Gateway) -> None:
|
||||
entered: Final = threading.Event()
|
||||
release: Final = threading.Event()
|
||||
|
||||
def blocked(request: Request) -> Reply:
|
||||
entered.set()
|
||||
assert release.wait(30), "test did not release the controlled upstream"
|
||||
return _completion(request)
|
||||
|
||||
with workflow_gateway.scenario() as scenario, wire_server(blocked) as wire:
|
||||
model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1")
|
||||
actors: Final = _actors(scenario, model)
|
||||
draft: Final = _register(workflow_gateway, actors, model)
|
||||
_publish(workflow_gateway, actors, draft)
|
||||
try:
|
||||
with pytest.raises(httpx.ReadTimeout):
|
||||
workflow_gateway.client.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
json=_request(),
|
||||
headers={"Idempotency-Key": "detached-run", "Authorization": f"Bearer {actors.caller}"},
|
||||
timeout=0.5,
|
||||
)
|
||||
assert entered.wait(15), "detached invocation never reached the controlled upstream"
|
||||
executions: Final = workflow_gateway.rows(
|
||||
'SELECT run_id, state FROM "LiteLLM_WorkflowExecution" WHERE name=%s AND state=%s',
|
||||
("email-reply", "running"),
|
||||
)
|
||||
assert len(executions) == 1, executions
|
||||
run_id: Final = str(executions[0]["run_id"])
|
||||
cancelled: Final = workflow_gateway.request(
|
||||
"POST", f"/v1/workflows/runs/{run_id}/cancel", key=actors.caller
|
||||
)
|
||||
assert cancelled.status_code == 202, cancelled.text
|
||||
final: Final = eventually(
|
||||
lambda: _trace(workflow_gateway, actors, run_id), lambda trace: trace.state == "cancelled", seconds=15
|
||||
)
|
||||
assert len(final.steps) == 1
|
||||
finally:
|
||||
release.set()
|
||||
assert len(wire.drain()) == 1
|
||||
18
tests/integration/management/workflow_proxy_config.yaml
Normal file
18
tests/integration/management/workflow_proxy_config.yaml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
model_list: []
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
database_url: os.environ/DATABASE_URL
|
||||
store_model_in_db: true
|
||||
disable_model_info_refresh: true
|
||||
disable_spend_logs: false
|
||||
proxy_batch_write_at: 1
|
||||
proxy_batch_polling_interval: 1
|
||||
enable_workflows: true
|
||||
coordination_redis:
|
||||
host: os.environ/REDIS_HOST
|
||||
port: os.environ/REDIS_PORT
|
||||
router_settings:
|
||||
redis_host: os.environ/REDIS_HOST
|
||||
redis_port: os.environ/REDIS_PORT
|
||||
num_retries: 0
|
||||
disable_cooldowns: true
|
||||
|
|
@ -21,6 +21,23 @@ from litellm.proxy.auth.login_utils import PASSWORD_RESET_ALLOWED_ROUTES
|
|||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("method", "path", "allowed"), (
|
||||
("POST", "/v1/workflows", True),
|
||||
("GET", "/v1/workflows", True),
|
||||
("PUT", "/v1/workflows/draft-id", True),
|
||||
("POST", "/v1/workflows/draft-id/publish", True),
|
||||
("GET", "/v1/workflows/runs/run-id/steps", True),
|
||||
("POST", "/v1/workflows/runs/run-id/cancel", True),
|
||||
("POST", "/v1/workflows/runs", False),
|
||||
("GET", "/v1/workflows/runs", False),
|
||||
("PATCH", "/v1/workflows/runs/run-id", False),
|
||||
("POST", "/v1/workflows/runs/run-id/events", False),
|
||||
("DELETE", "/v1/workflows/draft-id", False),
|
||||
))
|
||||
def test_workflow_execution_rbac_does_not_grant_legacy_tracking_access(method: str, path: str, allowed: bool) -> None:
|
||||
assert RouteChecks.is_workflow_execution_route(path, method) is allowed
|
||||
|
||||
DAILY_ACTIVITY_ROUTE_PAIRS: Final[tuple[tuple[str, str], ...]] = (
|
||||
("/user/daily/activity", "/user/daily/activity/aggregated"),
|
||||
("/user/daily/activity", "/user/daily/activity/aggregated/keys"),
|
||||
|
|
@ -2904,7 +2921,7 @@ def test_internal_user_can_access_key_reset_spend_route():
|
|||
request = MagicMock(spec=Request)
|
||||
request.query_params = {}
|
||||
|
||||
key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae"
|
||||
key_hash = "a" * 64
|
||||
route = f"/key/{key_hash}/reset_spend"
|
||||
|
||||
# Should not raise — the route-level check must pass for team admins
|
||||
|
|
@ -2936,7 +2953,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re
|
|||
request = MagicMock(spec=Request)
|
||||
request.query_params = {}
|
||||
|
||||
key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae"
|
||||
key_hash = "a" * 64
|
||||
|
||||
# /key/{hash}/reset_spend passes the route check for internal_user
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
|
|
|
|||
1233
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
1233
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue