mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(workflows): reuse LiteFlow SDK for sequential gateway execution
This commit is contained in:
parent
2a802eb8b9
commit
5fd7ee8e9e
49 changed files with 5634 additions and 68 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -546,6 +546,7 @@ jobs:
|
|||
|
||||
- shard: proxy-endpoints
|
||||
test-path: |-
|
||||
tests/unit/proxy/workflows
|
||||
tests/unit/proxy/management_endpoints
|
||||
tests/unit/proxy/management_helpers
|
||||
tests/unit/proxy/list_api
|
||||
|
|
|
|||
|
|
@ -0,0 +1,49 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TABLE "LiteLLM_WorkflowDefinition" (
|
||||
"workflow_id" TEXT NOT NULL,
|
||||
"team_id" TEXT NOT NULL,
|
||||
"name" TEXT NOT NULL,
|
||||
"etag" TEXT NOT NULL,
|
||||
"digest" TEXT NOT NULL,
|
||||
"bundle_json" TEXT NOT NULL,
|
||||
"contributors" TEXT[],
|
||||
"published_snapshot_id" TEXT,
|
||||
"published_etag" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
CONSTRAINT "LiteLLM_WorkflowDefinition_pkey" PRIMARY KEY ("workflow_id")
|
||||
);
|
||||
|
||||
CREATE TABLE "LiteLLM_WorkflowSnapshot" (
|
||||
"snapshot_id" TEXT NOT NULL,
|
||||
"workflow_id" TEXT NOT NULL,
|
||||
"team_id" TEXT NOT NULL,
|
||||
"name" TEXT NOT NULL,
|
||||
"digest" TEXT NOT NULL,
|
||||
"bundle_json" TEXT NOT NULL,
|
||||
"contributors" TEXT[],
|
||||
"published_by" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
CONSTRAINT "LiteLLM_WorkflowSnapshot_pkey" PRIMARY KEY ("snapshot_id"),
|
||||
CONSTRAINT "LiteLLM_WorkflowSnapshot_workflow_id_fkey" FOREIGN KEY ("workflow_id")
|
||||
REFERENCES "LiteLLM_WorkflowDefinition"("workflow_id") ON DELETE RESTRICT ON UPDATE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowDefinition_team_id_name_key" ON "LiteLLM_WorkflowDefinition"("team_id", "name");
|
||||
CREATE INDEX "LiteLLM_WorkflowDefinition_team_id_workflow_id_idx" ON "LiteLLM_WorkflowDefinition"("team_id", "workflow_id");
|
||||
CREATE INDEX "LiteLLM_WorkflowSnapshot_workflow_id_created_at_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("workflow_id", "created_at", "snapshot_id");
|
||||
CREATE INDEX "LiteLLM_WorkflowSnapshot_team_id_snapshot_id_idx" ON "LiteLLM_WorkflowSnapshot"("team_id", "snapshot_id");
|
||||
|
||||
ALTER TABLE "LiteLLM_WorkflowRun"
|
||||
ADD COLUMN "workflow_id" TEXT,
|
||||
ADD COLUMN "snapshot_id" TEXT,
|
||||
ADD COLUMN "team_id" TEXT,
|
||||
ADD COLUMN "engine_owned" BOOLEAN NOT NULL DEFAULT false,
|
||||
ADD COLUMN "request_digest" TEXT,
|
||||
ADD COLUMN "idempotency_key" TEXT,
|
||||
ADD COLUMN "cancel_requested" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_principal_idempotency_key" ON "LiteLLM_WorkflowRun"("created_by", "team_id", "workflow_id", "idempotency_key");
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_workflow_id_team_id_created_at_run_id_idx" ON "LiteLLM_WorkflowRun"("workflow_id", "team_id", "created_at", "run_id");
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_team_id_engine_owned_status_idx" ON "LiteLLM_WorkflowRun"("team_id", "engine_owned", "status");
|
||||
|
|
@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
|
|||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowDefinition {
|
||||
workflow_id String @id @default(uuid())
|
||||
team_id String
|
||||
name String
|
||||
etag String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_snapshot_id String?
|
||||
published_etag String?
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
snapshots LiteLLM_WorkflowSnapshot[]
|
||||
|
||||
@@unique([team_id, name])
|
||||
@@index([team_id, workflow_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowSnapshot {
|
||||
snapshot_id String @id @default(uuid())
|
||||
workflow_id String
|
||||
team_id String
|
||||
name String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_by String?
|
||||
created_at DateTime @default(now())
|
||||
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
|
||||
|
||||
@@index([workflow_id, created_at, snapshot_id])
|
||||
@@index([team_id, snapshot_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
workflow_id String?
|
||||
snapshot_id String?
|
||||
team_id String?
|
||||
engine_owned Boolean @default(false)
|
||||
request_digest String?
|
||||
idempotency_key String?
|
||||
cancel_requested Boolean @default(false)
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
|
|
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
|
|||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
|
||||
@@index([workflow_id, team_id, created_at, run_id])
|
||||
@@index([team_id, engine_owned, status])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
|
|
|
|||
|
|
@ -355,6 +355,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_ShadowEvalJob",
|
||||
"LiteLLM_ShadowEvalAttempt",
|
||||
"LiteLLM_ShadowEvalFunnel",
|
||||
"LiteLLM_WorkflowDefinition",
|
||||
"LiteLLM_WorkflowSnapshot",
|
||||
"LiteLLM_WorkflowRun",
|
||||
"LiteLLM_WorkflowEvent",
|
||||
"LiteLLM_WorkflowMessage",
|
||||
|
|
|
|||
165
litellm/models/workflow_definition.py
Normal file
165
litellm/models/workflow_definition.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Literal, Protocol, TypeAlias
|
||||
|
||||
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
|
||||
Draft as WorkflowDraft,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StringConstraints
|
||||
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.proxy.workflows.monty.models import PythonProgram
|
||||
|
||||
WorkflowName: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$")]
|
||||
Digest: TypeAlias = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")]
|
||||
|
||||
|
||||
class ValueBinding(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
|
||||
|
||||
value: JsonValue
|
||||
|
||||
|
||||
class ReferenceBinding(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
|
||||
|
||||
source: str = Field(alias="from", min_length=1, max_length=128)
|
||||
output: Literal["messages", "result"]
|
||||
|
||||
|
||||
Binding: TypeAlias = ValueBinding | ReferenceBinding
|
||||
|
||||
|
||||
class WorkflowRetry(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
|
||||
|
||||
max_attempts: int = Field(default=1, ge=1, le=3, strict=True)
|
||||
backoff_seconds: float = Field(default=0, ge=0, le=30)
|
||||
|
||||
|
||||
class StructuredSchema(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
|
||||
|
||||
name: str = Field(min_length=1, max_length=128)
|
||||
strict: Literal[True]
|
||||
schema_: Mapping[str, JsonValue] = Field(alias="schema")
|
||||
|
||||
|
||||
class StructuredOutput(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
type: Literal["json_schema"]
|
||||
json_schema: StructuredSchema
|
||||
|
||||
|
||||
class ChatCompletionNode(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
|
||||
|
||||
id: WorkflowName
|
||||
type: Literal["chat_completion"]
|
||||
model: str = Field(min_length=1, max_length=256)
|
||||
messages: Binding
|
||||
timeout_seconds: float = Field(gt=0, le=300)
|
||||
retry: WorkflowRetry
|
||||
temperature: float | None = Field(default=None, ge=0, le=2)
|
||||
top_p: float | None = Field(default=None, ge=0, le=1)
|
||||
max_tokens: int | None = Field(default=None, gt=0, strict=True)
|
||||
max_completion_tokens: int | None = Field(default=None, gt=0, strict=True)
|
||||
seed: int | None = Field(default=None, strict=True)
|
||||
stop: str | tuple[str, ...] | None = None
|
||||
response_format: StructuredOutput | None = None
|
||||
|
||||
|
||||
class PythonNode(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
|
||||
|
||||
id: WorkflowName
|
||||
type: Literal["python"]
|
||||
program: PythonProgram
|
||||
inputs: Mapping[str, Binding]
|
||||
timeout_seconds: float = Field(gt=0, le=300)
|
||||
retry: WorkflowRetry
|
||||
|
||||
|
||||
WorkflowNode: TypeAlias = Annotated[ChatCompletionNode | PythonNode, Field(discriminator="type")]
|
||||
|
||||
|
||||
class WorkflowBundle(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)
|
||||
|
||||
contract_version: Literal["1"]
|
||||
name: WorkflowName
|
||||
timeout_seconds: float = Field(gt=0, le=1800)
|
||||
input_schema: Mapping[str, JsonValue]
|
||||
output: ReferenceBinding
|
||||
nodes: tuple[WorkflowNode, ...] = Field(min_length=1, max_length=64)
|
||||
|
||||
|
||||
def parsed_bundle(bundle: Mapping[str, JsonValue]) -> WorkflowBundle:
|
||||
return WorkflowBundle.model_validate(bundle)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkflowDefinition:
|
||||
draft: WorkflowDraft
|
||||
contributors: tuple[str, ...]
|
||||
|
||||
|
||||
class WorkflowSnapshot(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
snapshot_id: str
|
||||
workflow_id: str
|
||||
team_id: str
|
||||
name: str
|
||||
digest: Digest
|
||||
bundle: WorkflowBundle
|
||||
contributors: tuple[str, ...]
|
||||
published_by: str | None
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class DefinitionPage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
workflows: tuple[WorkflowDraft, ...]
|
||||
next_cursor: str | None
|
||||
|
||||
|
||||
DefinitionErrorCode: TypeAlias = Literal[
|
||||
"not_found", "forbidden", "team_required", "invalid_bundle", "conflict", "stale_draft", "precondition_required"
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DefinitionError:
|
||||
code: DefinitionErrorCode
|
||||
message: str
|
||||
param: str | None = None
|
||||
|
||||
|
||||
class WorkflowDefinitionStore(Protocol):
|
||||
async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None: ...
|
||||
|
||||
async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None: ...
|
||||
|
||||
async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage: ...
|
||||
|
||||
async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None: ...
|
||||
|
||||
async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None: ...
|
||||
|
||||
async def create(
|
||||
self, bundle: WorkflowBundle, team_id: str, contributor: str
|
||||
) -> WorkflowDraft | DefinitionError: ...
|
||||
|
||||
async def replace(
|
||||
self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str
|
||||
) -> WorkflowDraft | DefinitionError: ...
|
||||
|
||||
async def publish(
|
||||
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
|
||||
) -> WorkflowSnapshot | DefinitionError: ...
|
||||
|
||||
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: ...
|
||||
151
litellm/models/workflow_execution.py
Normal file
151
litellm/models/workflow_execution.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Final, Protocol
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
from liteflow import RunState # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
|
||||
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
|
||||
RunSteps as SDKRunSteps,
|
||||
)
|
||||
from liteflow import ( # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
|
||||
StepTrace as SDKStepTrace,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue
|
||||
|
||||
__all__: Final = ("RunState",)
|
||||
|
||||
|
||||
class RunInput(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
|
||||
deadline: datetime
|
||||
draft: bool
|
||||
end_user_id: str | None = None
|
||||
|
||||
|
||||
class EngineRun(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, from_attributes=True)
|
||||
|
||||
run_id: str
|
||||
session_id: str
|
||||
workflow_id: str
|
||||
snapshot_id: str
|
||||
team_id: str
|
||||
created_by: str
|
||||
status: RunState
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
input: RunInput
|
||||
request_digest: str
|
||||
output: JsonValue = None
|
||||
metadata: Mapping[str, JsonValue] | None = None
|
||||
cancel_requested: bool = False
|
||||
|
||||
|
||||
class StepTrace(SDKStepTrace):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
output: JsonValue = None
|
||||
finished_at: datetime | None = None
|
||||
error: str | None = None
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
response_cost: float | None = None
|
||||
request_id: str | None = None
|
||||
call_id: str | None = None
|
||||
request_ids: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class RunSteps(SDKRunSteps):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
steps: tuple[StepTrace, ...] # pyright: ignore[reportIncompatibleVariableOverride] # both SDK and extension models are frozen
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class RunView(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, from_attributes=True)
|
||||
|
||||
run_id: str
|
||||
session_id: str
|
||||
workflow_id: str
|
||||
snapshot_id: str
|
||||
team_id: str
|
||||
status: RunState
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
cancel_requested: bool
|
||||
|
||||
|
||||
class RunPage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
runs: tuple[RunView, ...]
|
||||
count: int
|
||||
next_cursor: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChildCallResult:
|
||||
output: JsonValue
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
request_id: str | None
|
||||
response_cost: float | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ExecutionError:
|
||||
code: str
|
||||
message: str
|
||||
receipt: ChildCallResult | None = None
|
||||
|
||||
|
||||
def workflow_call_id(run_id: str, node_id: str, attempt: int) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, f"litellm/workflows/{run_id}/{node_id}/{attempt}"))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RunAdmission:
|
||||
run: EngineRun
|
||||
created: bool
|
||||
|
||||
|
||||
class WorkflowExecutionStore(Protocol):
|
||||
async def find_idempotent(
|
||||
self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str
|
||||
) -> EngineRun | None: ...
|
||||
|
||||
async def admit(
|
||||
self,
|
||||
*,
|
||||
workflow_id: str,
|
||||
snapshot_id: str,
|
||||
team_id: str,
|
||||
key_hash: str,
|
||||
input: RunInput,
|
||||
digest: str,
|
||||
idempotency_key: str | None,
|
||||
) -> RunAdmission | ExecutionError: ...
|
||||
|
||||
async def get(self, run_id: str) -> EngineRun | None: ...
|
||||
|
||||
async def checkpoint(self, run_id: str, step: StepTrace) -> bool: ...
|
||||
|
||||
async def steps(self, run_id: str) -> tuple[StepTrace, ...]: ...
|
||||
|
||||
async def finish(self, run_id: str, state: RunState, output: JsonValue, error: str | None) -> None: ...
|
||||
|
||||
async def cancel(self, run_id: str) -> None: ...
|
||||
|
||||
async def list_runs(
|
||||
self,
|
||||
workflow_id: str,
|
||||
team_id: str,
|
||||
key_hash: str | None,
|
||||
state: RunState | None,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> RunPage | ExecutionError: ...
|
||||
|
|
@ -126,6 +126,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_
|
|||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.proxy.workflows.runner import WorkflowRunner
|
||||
from litellm.repositories.database import Database
|
||||
from litellm.tracing import TraceReceiver
|
||||
|
||||
|
|
@ -135,6 +136,7 @@ else:
|
|||
|
||||
|
||||
class ProxyLifespanState(TypedDict):
|
||||
workflow_runner: NotRequired[ReadOnly["WorkflowRunner | None"]]
|
||||
tracing_receiver: ReadOnly["TraceReceiver | None"]
|
||||
database: ReadOnly["Database | None"]
|
||||
|
||||
|
|
@ -660,6 +662,21 @@ class LiteLLMRoutes(enum.Enum):
|
|||
KeyManagementRoutes.KEY_UPDATE.value,
|
||||
)
|
||||
|
||||
workflow_definition_routes = (
|
||||
"/v1/workflows",
|
||||
"/v1/workflows/{workflow_id}",
|
||||
"/v1/workflows/{workflow_id}/validate",
|
||||
"/v1/workflows/{workflow_id}/publish",
|
||||
)
|
||||
|
||||
workflow_execution_routes = (
|
||||
"/v1/workflows/{workflow_id}/test",
|
||||
"/v1/workflows/{workflow_id}/runs",
|
||||
"/v1/workflows/runs/{run_id}/steps",
|
||||
"/v1/workflows/runs/{run_id}/definition",
|
||||
"/v1/workflows/runs/{run_id}/cancel",
|
||||
)
|
||||
|
||||
management_routes = (
|
||||
[
|
||||
# user
|
||||
|
|
@ -723,6 +740,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
+ mcp_management_routes
|
||||
+ list(agent_management_routes)
|
||||
+ list(skill_hub_routes)
|
||||
+ list(workflow_definition_routes)
|
||||
+ list(workflow_execution_routes)
|
||||
)
|
||||
|
||||
spend_tracking_routes = [
|
||||
|
|
@ -966,6 +985,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
# inside agent_endpoints/endpoints.py
|
||||
*agent_management_routes,
|
||||
*skill_hub_routes,
|
||||
*workflow_definition_routes,
|
||||
*workflow_execution_routes,
|
||||
] # routes that manage their own allowed/disallowed logic
|
||||
|
||||
## Org Admin Routes ##
|
||||
|
|
|
|||
|
|
@ -603,8 +603,10 @@ class RouteChecks:
|
|||
if not isinstance(route, str):
|
||||
return False
|
||||
|
||||
def _placeholder_to_regex(match: re.Match) -> str:
|
||||
def _placeholder_to_regex(match: re.Match[str]) -> str:
|
||||
placeholder: Final = match.group(0).strip("{}")
|
||||
if placeholder == "workflow_id" and match.string.startswith("/v1/workflows/"):
|
||||
return r"(?!runs(?:/|$))[^/]+"
|
||||
if not placeholder.endswith(":path"):
|
||||
return r"[^/]+"
|
||||
# A ":path" placeholder takes whatever the router's own path
|
||||
|
|
|
|||
|
|
@ -4849,7 +4849,12 @@ async def _run_post_custom_auth_checks(
|
|||
|
||||
|
||||
async def authorize_internal_virtual_key(
|
||||
key_hash: str, request: Request, request_data: dict[str, object]
|
||||
key_hash: str,
|
||||
request: Request,
|
||||
request_data: dict[str, object],
|
||||
*,
|
||||
route: str | None = None,
|
||||
expected_team_id: str | None = None,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""Authorize a server-owned job against its persisted virtual-key assignment, never a client-supplied bearer hash."""
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
|
@ -4858,15 +4863,17 @@ async def authorize_internal_virtual_key(
|
|||
await IdentityStore(prisma_client, user_api_key_cache, proxy_logging_obj=proxy_logging_obj).resolve(
|
||||
hashed_token=key_hash
|
||||
)
|
||||
)
|
||||
route: Final = get_request_route(request=request)
|
||||
await pre_db_read_auth_checks(request_data=request_data, request=request, route=route)
|
||||
).model_copy(deep=True)
|
||||
if expected_team_id is not None and identity.team_id != expected_team_id:
|
||||
raise HTTPException(status_code=403, detail="The execution key no longer belongs to its original team")
|
||||
execution_route: Final = route if route is not None else get_request_route(request=request)
|
||||
await pre_db_read_auth_checks(request_data=request_data, request=request, route=execution_route)
|
||||
auth: Final = await validate_resolved_virtual_key(
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
valid_token=identity,
|
||||
api_key=key_hash,
|
||||
route=route,
|
||||
route=execution_route,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
parent_otel_span=None,
|
||||
end_user_id=None,
|
||||
|
|
@ -4878,13 +4885,13 @@ async def authorize_internal_virtual_key(
|
|||
user_api_key_auth_obj=auth,
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
route=execution_route,
|
||||
api_key=key_hash,
|
||||
force_virtual_key_checks=True,
|
||||
)
|
||||
if recovered is not None:
|
||||
return recovered
|
||||
_seed_request_destinations(auth, request)
|
||||
auth.request_route = route
|
||||
auth.request_route = execution_route
|
||||
request.state.principal = _resolve_request_principal(request, auth)
|
||||
return auth
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
import logging
|
||||
import math
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
|
||||
from contextvars import Context
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import lru_cache
|
||||
|
|
@ -29,7 +30,7 @@ import httpx
|
|||
import orjson
|
||||
from fastapi import HTTPException, Request, status
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from starlette.types import Receive, Scope, Send
|
||||
from typing_extensions import Never
|
||||
|
|
@ -94,6 +95,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
backfill_missing_cache_usage_fields,
|
||||
)
|
||||
from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields
|
||||
from litellm.models.workflow_execution import ExecutionError
|
||||
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload, with_openai_overload_status
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
|
|
@ -725,6 +727,8 @@ def serialize_http_exception_detail(
|
|||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True))
|
||||
_WORKFLOW_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_WORKFLOW_JSON: Final = TypeAdapter[JsonValue](JsonValue)
|
||||
_UNSEEN_DEPLOYMENT: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset())
|
||||
|
||||
|
||||
|
|
@ -2190,10 +2194,63 @@ def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]:
|
|||
)
|
||||
|
||||
|
||||
async def process_workflow_request(
|
||||
*,
|
||||
request: Request,
|
||||
workflow_model: str,
|
||||
caller: UserAPIKeyAuth,
|
||||
messages: JsonValue,
|
||||
execute: Callable[[JsonValue], Awaitable[ModelResponse | ExecutionError]],
|
||||
) -> ModelResponse | ExecutionError:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_unbound_budget_reservation
|
||||
|
||||
await release_unbound_budget_reservation(caller.budget_reservation)
|
||||
auth: Final = caller.model_copy(update={"budget_reservation": None})
|
||||
|
||||
async def process() -> ModelResponse | ExecutionError:
|
||||
try:
|
||||
prepared: Final = _WORKFLOW_FIELDS.validate_python(
|
||||
await add_litellm_data_to_request(
|
||||
data={"model": workflow_model, "messages": messages, "stream": False},
|
||||
request=request,
|
||||
general_settings=_WORKFLOW_FIELDS.validate_python(proxy_server.general_settings), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
|
||||
user_api_key_dict=auth,
|
||||
version=proxy_server.version,
|
||||
proxy_config=proxy_server.proxy_config,
|
||||
)
|
||||
)
|
||||
data: Final = _WORKFLOW_FIELDS.validate_python(
|
||||
check_and_merge_model_level_guardrails(
|
||||
data=prepared, llm_router=proxy_server.llm_router, trust_client_model_info=False
|
||||
)
|
||||
)
|
||||
guarded: Final = _WORKFLOW_FIELDS.validate_python(
|
||||
await proxy_server.proxy_logging_obj.pre_call_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook body is validated at this boundary
|
||||
user_api_key_dict=auth, data=data, call_type="acompletion", guardrails_only=True
|
||||
)
|
||||
)
|
||||
await proxy_server.proxy_logging_obj.during_call_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook accepts the validated body
|
||||
data=guarded, user_api_key_dict=auth, call_type="acompletion", guardrails_only=True
|
||||
)
|
||||
result: Final = await execute(_WORKFLOW_JSON.validate_python(guarded["messages"]))
|
||||
if isinstance(result, ExecutionError):
|
||||
return result
|
||||
response: Final = await proxy_server.proxy_logging_obj.post_call_success_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook accepts the validated body
|
||||
data=guarded, user_api_key_dict=auth, response=result, guardrails_only=True
|
||||
)
|
||||
return response if isinstance(response, ModelResponse) else ModelResponse.model_validate(response)
|
||||
except Exception: # noqa: BLE001 # workflow policy failures are values and must not expose guarded content
|
||||
return ExecutionError("request_rejected", "The gateway could not complete the workflow")
|
||||
|
||||
return await Context().run(asyncio.create_task, process())
|
||||
|
||||
|
||||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict, served_model: str | None = None):
|
||||
self.data = data
|
||||
self.served_model = served_model
|
||||
self.logging_obj: LiteLLMLoggingObj | None = None
|
||||
self._tags_before_guardrails: frozenset[str] | None = None
|
||||
self.route_type: str | None = None
|
||||
|
||||
|
|
@ -2636,6 +2693,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
)
|
||||
|
||||
self.data["litellm_logging_obj"] = logging_obj
|
||||
self.logging_obj = logging_obj
|
||||
|
||||
# Merge model-level guardrails before pre_call_hook so DB/UI-configured
|
||||
# guardrails actually execute on pre_call. Without this, guardrails set
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Annotated, TypeAlias
|
||||
from typing import Annotated, TypeAlias, cast # noqa: TID251 # Prisma writer delegates are generated dynamically
|
||||
|
||||
from fastapi import Depends
|
||||
|
||||
|
|
@ -7,12 +7,16 @@ from litellm.models.gateway_activity import GatewayActivityStore
|
|||
from litellm.models.plugin_upload import PluginUploadStore
|
||||
from litellm.models.skill_submission import SkillSubmissionStore
|
||||
from litellm.models.user_banner import UserBannerStore
|
||||
from litellm.models.workflow_definition import WorkflowDefinitionStore
|
||||
from litellm.models.workflow_execution import WorkflowExecutionStore
|
||||
from litellm.proxy.deps.database import DatabaseDep
|
||||
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
|
||||
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
|
||||
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
|
||||
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
|
||||
from litellm.repositories.user_banner_repository import UserBannerRepository
|
||||
from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository
|
||||
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
|
||||
|
||||
|
||||
async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
|
||||
|
|
@ -54,3 +58,33 @@ async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmiss
|
|||
SkillSubmissionRepositoryDep: TypeAlias = Annotated[
|
||||
SkillSubmissionStore | None, Depends(get_skill_submission_repository)
|
||||
]
|
||||
|
||||
|
||||
async def get_workflow_definition_repository(database: DatabaseDep) -> WorkflowDefinitionStore | None:
|
||||
return (
|
||||
None
|
||||
if database is None
|
||||
else WorkflowDefinitionRepository(
|
||||
cast(DefinitionDatabase, database.prisma.writer_db) # cast-ok: Prisma delegates implement the protocol
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
WorkflowDefinitionRepositoryDep: TypeAlias = Annotated[
|
||||
WorkflowDefinitionStore | None, Depends(get_workflow_definition_repository)
|
||||
]
|
||||
|
||||
|
||||
async def get_workflow_execution_repository(database: DatabaseDep) -> WorkflowExecutionStore | None:
|
||||
return (
|
||||
None
|
||||
if database is None
|
||||
else WorkflowExecutionRepository(
|
||||
cast(ExecutionDatabase, database.prisma.writer_db) # cast-ok: generated Prisma transaction delegates
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
WorkflowExecutionRepositoryDep: TypeAlias = Annotated[
|
||||
WorkflowExecutionStore | None, Depends(get_workflow_execution_repository)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,223 @@
|
|||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, NoReturn, TypeAlias, TypeVar
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm.models.workflow_definition import (
|
||||
DefinitionError,
|
||||
DefinitionPage,
|
||||
Digest,
|
||||
WorkflowBundle,
|
||||
WorkflowDefinitionStore,
|
||||
WorkflowDraft,
|
||||
parsed_bundle,
|
||||
)
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep
|
||||
from litellm.proxy.workflows.definitions import (
|
||||
author_identity,
|
||||
authorize_team,
|
||||
authorized_draft,
|
||||
check_etag,
|
||||
requested_team,
|
||||
validate_bundle,
|
||||
)
|
||||
|
||||
router: Final = APIRouter(tags=["workflow management"])
|
||||
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
|
||||
_Team: TypeAlias = Annotated[str | None, Query(min_length=1, max_length=256)]
|
||||
_ETag: TypeAlias = Annotated[str | None, Header(alias="If-Match")]
|
||||
_Result: Final = TypeVar("_Result")
|
||||
_STATUS: Final = MappingProxyType(
|
||||
{
|
||||
"not_found": 404,
|
||||
"forbidden": 403,
|
||||
"team_required": 400,
|
||||
"invalid_bundle": 400,
|
||||
"conflict": 409,
|
||||
"stale_draft": 412,
|
||||
"precondition_required": 412,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class WorkflowRegistration(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
bundle: WorkflowBundle
|
||||
|
||||
|
||||
class WorkflowPublication(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
digest: Digest
|
||||
|
||||
|
||||
class WorkflowValidationRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
class WorkflowValidationResult(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
valid: Literal[True] = True
|
||||
workflow_id: str
|
||||
etag: str
|
||||
digest: str
|
||||
|
||||
|
||||
class WorkflowPublicationResult(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
workflow_id: str
|
||||
snapshot_id: str
|
||||
model: str
|
||||
digest: str
|
||||
|
||||
|
||||
def raise_public(error: DefinitionError) -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=_STATUS[error.code],
|
||||
detail={
|
||||
"error": {
|
||||
"type": "workflow_definition_error",
|
||||
"code": error.code,
|
||||
"message": error.message,
|
||||
"param": error.param,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _result(result: _Result | DefinitionError) -> _Result:
|
||||
if isinstance(result, DefinitionError):
|
||||
raise_public(result)
|
||||
return result
|
||||
|
||||
|
||||
def _store(store: WorkflowDefinitionStore | None) -> WorkflowDefinitionStore:
|
||||
if store is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
return store
|
||||
|
||||
|
||||
def _author(caller: UserAPIKeyAuth) -> str:
|
||||
return _result(author_identity(caller))
|
||||
|
||||
|
||||
def _draft_response(draft: WorkflowDraft, status_code: int = 200) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
content=draft.model_dump(mode="json", by_alias=True, exclude_unset=True),
|
||||
status_code=status_code,
|
||||
headers={"ETag": draft.etag},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/v1/workflows", status_code=201, response_model=WorkflowDraft, response_model_exclude_unset=True)
|
||||
async def register_workflow(
|
||||
data: WorkflowRegistration,
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
) -> JSONResponse:
|
||||
"""Create a team draft from a sequential LiteFlow v1 bundle."""
|
||||
contributor: Final = _author(caller)
|
||||
store: Final = _store(repository)
|
||||
scope: Final = _result(requested_team(caller, team_id))
|
||||
_result(await authorize_team(store, caller, scope, "write"))
|
||||
_result(validate_bundle(data.bundle))
|
||||
draft: Final = _result(await store.create(data.bundle, scope, contributor))
|
||||
return _draft_response(draft, 201)
|
||||
|
||||
|
||||
@router.get("/v1/workflows", response_model=DefinitionPage, response_model_exclude_unset=True)
|
||||
async def list_workflow_definitions(
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 50,
|
||||
cursor: Annotated[
|
||||
str | None, Query(pattern=r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")
|
||||
] = None,
|
||||
) -> DefinitionPage:
|
||||
"""List readable team drafts in immutable workflow ID order."""
|
||||
store: Final = _store(repository)
|
||||
scope: Final = _result(requested_team(caller, team_id))
|
||||
_result(await authorize_team(store, caller, scope, "read"))
|
||||
return await store.list_drafts(scope, limit, cursor)
|
||||
|
||||
|
||||
@router.get("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True)
|
||||
async def get_workflow_definition(
|
||||
workflow_id: str,
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
) -> JSONResponse:
|
||||
"""Read the current draft and its publication pointer."""
|
||||
draft: Final = _result(await authorized_draft(_store(repository), caller, workflow_id, team_id))
|
||||
return _draft_response(draft)
|
||||
|
||||
|
||||
@router.put("/v1/workflows/{workflow_id}", response_model=WorkflowDraft, response_model_exclude_unset=True)
|
||||
async def replace_workflow_definition(
|
||||
workflow_id: str,
|
||||
data: WorkflowRegistration,
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
if_match: _ETag = None,
|
||||
) -> JSONResponse:
|
||||
"""Replace a draft only when If-Match identifies its current revision."""
|
||||
contributor: Final = _author(caller)
|
||||
store: Final = _store(repository)
|
||||
current: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="write"))
|
||||
_result(check_etag(current, if_match))
|
||||
_result(validate_bundle(data.bundle))
|
||||
draft: Final = _result(await store.replace(workflow_id, current.team_id, current.etag, data.bundle, contributor))
|
||||
return _draft_response(draft)
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/validate", response_model=WorkflowValidationResult)
|
||||
async def validate_workflow_definition(
|
||||
workflow_id: str,
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
if_match: _ETag = None,
|
||||
data: WorkflowValidationRequest | None = None,
|
||||
) -> WorkflowValidationResult:
|
||||
"""Validate a draft without model calls or Python execution."""
|
||||
_author(caller)
|
||||
store: Final = _store(repository)
|
||||
draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="write"))
|
||||
_result(check_etag(draft, if_match))
|
||||
_result(validate_bundle(parsed_bundle(draft.bundle)))
|
||||
return WorkflowValidationResult(workflow_id=workflow_id, etag=draft.etag, digest=draft.digest)
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/publish", response_model=WorkflowPublicationResult)
|
||||
async def publish_workflow_definition(
|
||||
workflow_id: str,
|
||||
data: WorkflowPublication,
|
||||
repository: WorkflowDefinitionRepositoryDep,
|
||||
caller: _Caller,
|
||||
team_id: _Team = None,
|
||||
if_match: _ETag = None,
|
||||
) -> WorkflowPublicationResult:
|
||||
"""Publish the reviewed revision as an immutable snapshot."""
|
||||
publisher: Final = _author(caller)
|
||||
store: Final = _store(repository)
|
||||
draft: Final = _result(await authorized_draft(store, caller, workflow_id, team_id, action="publish"))
|
||||
_result(check_etag(draft, if_match))
|
||||
_result(validate_bundle(parsed_bundle(draft.bundle)))
|
||||
snapshot: Final = _result(await store.publish(workflow_id, draft.team_id, draft.etag, data.digest, publisher))
|
||||
return WorkflowPublicationResult(
|
||||
workflow_id=workflow_id,
|
||||
snapshot_id=snapshot.snapshot_id,
|
||||
model=f"workflow/{snapshot.name}",
|
||||
digest=snapshot.digest,
|
||||
)
|
||||
|
|
@ -0,0 +1,413 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import timedelta
|
||||
from typing import Annotated, Final, Literal, NoReturn, TypeAlias
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.models.workflow_definition import DefinitionError, WorkflowDefinitionStore, WorkflowDraft, WorkflowSnapshot
|
||||
from litellm.models.workflow_execution import (
|
||||
EngineRun,
|
||||
ExecutionError,
|
||||
RunInput,
|
||||
RunPage,
|
||||
RunState,
|
||||
RunSteps,
|
||||
WorkflowExecutionStore,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep
|
||||
from litellm.proxy.management.teams.authz import is_team_admin
|
||||
from litellm.proxy.management_endpoints.workflow_definition_endpoints import raise_public
|
||||
from litellm.proxy.workflows.definitions import (
|
||||
authorize_team,
|
||||
authorized_draft,
|
||||
freeze_draft,
|
||||
published_workflow,
|
||||
requested_team,
|
||||
)
|
||||
from litellm.proxy.workflows.lifespan import get_workflow_runner
|
||||
from litellm.proxy.workflows.monty.models import MontyFailure
|
||||
from litellm.proxy.workflows.monty.runtime import validate_json
|
||||
from litellm.proxy.workflows.runner import utc_now
|
||||
|
||||
router: Final = APIRouter(tags=["workflow management"])
|
||||
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
|
||||
_Idempotency: TypeAlias = Annotated[str | None, Header(alias="Idempotency-Key", min_length=1, max_length=256)]
|
||||
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
|
||||
|
||||
|
||||
class TestRunRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
|
||||
authorize_billing: StrictBool
|
||||
|
||||
|
||||
class AcceptedRun(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
run_id: str
|
||||
snapshot_id: str
|
||||
trace_url: str
|
||||
|
||||
|
||||
class EmptyRunAction(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
class WorkflowInvocation(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
model: str
|
||||
messages: tuple[Mapping[str, JsonValue], ...] = Field(min_length=1, max_length=256)
|
||||
stream: Literal[False] = False
|
||||
user: str | None = Field(default=None, max_length=256)
|
||||
metadata: Mapping[str, JsonValue] | None = None
|
||||
|
||||
|
||||
def _error(error: ExecutionError, run_id: str | None = None) -> NoReturn:
|
||||
code: Final = {
|
||||
"not_found": 404,
|
||||
"forbidden": 403,
|
||||
"idempotency_conflict": 409,
|
||||
"capacity_exceeded": 429,
|
||||
"execution_failed": 422,
|
||||
"database_unavailable": 503,
|
||||
}.get(error.code, 400)
|
||||
raise HTTPException(
|
||||
status_code=code,
|
||||
detail={"error": {"type": "workflow_error", "code": error.code, "message": error.message}},
|
||||
headers={"x-litellm-workflow-run-id": run_id} if run_id is not None else None,
|
||||
)
|
||||
|
||||
|
||||
def _stores(
|
||||
definitions: WorkflowDefinitionStore | None, runs: WorkflowExecutionStore | None
|
||||
) -> tuple[WorkflowDefinitionStore, WorkflowExecutionStore]:
|
||||
if definitions is None or runs is None:
|
||||
_error(ExecutionError("database_unavailable", "Workflow execution requires a database"))
|
||||
return definitions, runs
|
||||
|
||||
|
||||
def _execution_key(caller: UserAPIKeyAuth, team_id: str) -> str:
|
||||
if not caller.token or not caller.via_virtual_key or caller.is_session_token or caller.team_id != team_id:
|
||||
_error(ExecutionError("forbidden", "Execution requires a virtual key assigned to the workflow team"))
|
||||
return caller.token
|
||||
|
||||
|
||||
def _idempotency_key(operation: str, key: str | None) -> str | None:
|
||||
return None if key is None else hashlib.sha256(f"{operation}:{key}".encode()).hexdigest()
|
||||
|
||||
|
||||
def _digest(messages: tuple[Mapping[str, JsonValue], ...], user: str | None, draft_etag: str | None) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(
|
||||
{"messages": tuple(dict(message) for message in messages), "user": user, "etag": draft_etag},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _validate_messages(messages: tuple[Mapping[str, JsonValue], ...]) -> None:
|
||||
bounded: Final = validate_json([dict(message) for message in messages], max_bytes=1024 * 1024, max_depth=64)
|
||||
if isinstance(bounded, MontyFailure):
|
||||
_error(ExecutionError("invalid_input", "Workflow messages exceed the JSON size or depth limit"))
|
||||
try:
|
||||
_MESSAGES.validate_python(bounded, strict=True)
|
||||
except ValidationError:
|
||||
_error(ExecutionError("invalid_input", "Expected standard Chat Completions messages"))
|
||||
|
||||
|
||||
async def _admit(
|
||||
request: Request,
|
||||
runs: WorkflowExecutionStore,
|
||||
snapshot: WorkflowSnapshot,
|
||||
caller: UserAPIKeyAuth,
|
||||
messages: tuple[Mapping[str, JsonValue], ...],
|
||||
user: str | None,
|
||||
key: str | None,
|
||||
digest: str,
|
||||
*,
|
||||
draft: bool,
|
||||
) -> EngineRun:
|
||||
_validate_messages(messages)
|
||||
runner: Final = get_workflow_runner(request)
|
||||
if runner is None:
|
||||
_error(ExecutionError("database_unavailable", "Workflow execution is unavailable"))
|
||||
admitted: Final = await runs.admit(
|
||||
workflow_id=snapshot.workflow_id,
|
||||
snapshot_id=snapshot.snapshot_id,
|
||||
team_id=snapshot.team_id,
|
||||
key_hash=_execution_key(caller, snapshot.team_id),
|
||||
input=RunInput(
|
||||
messages=messages,
|
||||
end_user_id=user,
|
||||
deadline=utc_now() + timedelta(seconds=snapshot.bundle.timeout_seconds),
|
||||
draft=draft,
|
||||
),
|
||||
digest=digest,
|
||||
idempotency_key=key,
|
||||
)
|
||||
if isinstance(admitted, ExecutionError):
|
||||
_error(admitted)
|
||||
if admitted.created and (failure := await runner.submit(admitted.run, request, caller)) is not None:
|
||||
_error(failure, admitted.run.run_id)
|
||||
return admitted.run
|
||||
|
||||
|
||||
def _accepted(run: EngineRun) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=202,
|
||||
content=AcceptedRun(
|
||||
run_id=run.run_id,
|
||||
snapshot_id=run.snapshot_id,
|
||||
trace_url=f"/v1/workflows/runs/{run.run_id}/steps",
|
||||
).model_dump(),
|
||||
headers={"x-litellm-workflow-run-id": run.run_id},
|
||||
)
|
||||
|
||||
|
||||
async def _replay_test(
|
||||
runs: WorkflowExecutionStore, draft: WorkflowDraft, key_hash: str, key: str | None, digest: str
|
||||
) -> JSONResponse | None:
|
||||
if key is None:
|
||||
return None
|
||||
existing: Final = await runs.find_idempotent(draft.workflow_id, draft.team_id, key_hash, key)
|
||||
if existing is None:
|
||||
return None
|
||||
if existing.request_digest != digest:
|
||||
_error(ExecutionError("idempotency_conflict", "Idempotency key was used with different input"))
|
||||
return _accepted(existing)
|
||||
|
||||
|
||||
@router.post("/v1/workflows/{workflow_id}/test", status_code=202, response_model=AcceptedRun)
|
||||
async def test_workflow(
|
||||
request: Request,
|
||||
workflow_id: str,
|
||||
data: TestRunRequest,
|
||||
caller: _Caller,
|
||||
definitions: WorkflowDefinitionRepositoryDep,
|
||||
executions: WorkflowExecutionRepositoryDep,
|
||||
if_match: Annotated[str | None, Header(alias="If-Match")] = None,
|
||||
idempotency_key: _Idempotency = None,
|
||||
) -> JSONResponse:
|
||||
"""Start a paid test of the exact draft revision and return its run ID immediately."""
|
||||
store, runs = _stores(definitions, executions)
|
||||
if not data.authorize_billing:
|
||||
_error(ExecutionError("billing_authorization_required", "Set authorize_billing to true for a paid test run"))
|
||||
draft: Final = await authorized_draft(store, caller, workflow_id, action="write")
|
||||
if isinstance(draft, DefinitionError):
|
||||
raise_public(draft)
|
||||
if if_match is None:
|
||||
raise_public(DefinitionError("precondition_required", "If-Match is required", "If-Match"))
|
||||
key_hash: Final = _execution_key(caller, draft.team_id)
|
||||
digest: Final = _digest(data.messages, caller.end_user_id, if_match)
|
||||
key: Final = _idempotency_key("test", idempotency_key)
|
||||
if (replay := await _replay_test(runs, draft, key_hash, key, digest)) is not None:
|
||||
return replay
|
||||
snapshot: Final = await freeze_draft(store, caller, workflow_id, if_match)
|
||||
if isinstance(snapshot, DefinitionError):
|
||||
if (
|
||||
snapshot.code == "stale_draft"
|
||||
and (accepted := await _replay_test(runs, draft, key_hash, key, digest)) is not None
|
||||
):
|
||||
return accepted
|
||||
raise_public(snapshot)
|
||||
return _accepted(
|
||||
await _admit(request, runs, snapshot, caller, data.messages, caller.end_user_id, key, digest, draft=True)
|
||||
)
|
||||
|
||||
|
||||
async def _read_run(
|
||||
store: WorkflowDefinitionStore, runs: WorkflowExecutionStore, caller: UserAPIKeyAuth, run_id: str
|
||||
) -> EngineRun:
|
||||
run: Final = await runs.get(run_id)
|
||||
if run is None:
|
||||
_error(ExecutionError("not_found", "Run not found"))
|
||||
if isinstance(requested_team(caller, run.team_id), DefinitionError):
|
||||
_error(ExecutionError("not_found", "Run not found"))
|
||||
team: Final = await store.get_team(run.team_id)
|
||||
if team is None or team.blocked:
|
||||
_error(ExecutionError("not_found", "Run not found"))
|
||||
owns_run: Final = bool(caller.token) and caller.token == run.created_by and caller.team_id == run.team_id
|
||||
if not owns_run and not user_api_key_has_admin_view(caller) and not is_team_admin(caller, team):
|
||||
_error(ExecutionError("not_found", "Run not found"))
|
||||
return run
|
||||
|
||||
|
||||
@router.get("/v1/workflows/runs/{run_id}/steps")
|
||||
async def get_workflow_steps(
|
||||
run_id: str,
|
||||
caller: _Caller,
|
||||
definitions: WorkflowDefinitionRepositoryDep,
|
||||
executions: WorkflowExecutionRepositoryDep,
|
||||
) -> RunSteps:
|
||||
"""Read run progress, typed node inputs and outputs, and linked model-call receipts."""
|
||||
store, runs = _stores(definitions, executions)
|
||||
run: Final = await _read_run(store, runs, caller, run_id)
|
||||
error: Final = (run.metadata or {}).get("error")
|
||||
return RunSteps(
|
||||
run_id=run_id,
|
||||
snapshot_id=run.snapshot_id,
|
||||
state=run.status,
|
||||
cancel_requested=run.cancel_requested,
|
||||
steps=await runs.steps(run_id),
|
||||
error=error if isinstance(error, str) else None,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/v1/workflows/runs/{run_id}/definition")
|
||||
async def get_run_definition(
|
||||
run_id: str,
|
||||
caller: _Caller,
|
||||
definitions: WorkflowDefinitionRepositoryDep,
|
||||
executions: WorkflowExecutionRepositoryDep,
|
||||
) -> WorkflowSnapshot:
|
||||
"""Read the immutable workflow definition captured when this run was admitted."""
|
||||
store, runs = _stores(definitions, executions)
|
||||
run: Final = await _read_run(store, runs, caller, run_id)
|
||||
snapshot: Final = await store.get_snapshot(run.snapshot_id, run.team_id)
|
||||
if snapshot is None or snapshot.workflow_id != run.workflow_id:
|
||||
_error(ExecutionError("not_found", "Pinned workflow definition not found"))
|
||||
return snapshot
|
||||
|
||||
|
||||
@router.post("/v1/workflows/runs/{run_id}/cancel", status_code=202, response_model=AcceptedRun)
|
||||
async def cancel_workflow_run(
|
||||
run_id: str,
|
||||
caller: _Caller,
|
||||
definitions: WorkflowDefinitionRepositoryDep,
|
||||
executions: WorkflowExecutionRepositoryDep,
|
||||
data: EmptyRunAction | None = None,
|
||||
) -> JSONResponse:
|
||||
"""Request cancellation and return the run ID to poll until execution stops."""
|
||||
store, runs = _stores(definitions, executions)
|
||||
run: Final = await _read_run(store, runs, caller, run_id)
|
||||
if caller.token != run.created_by:
|
||||
permission: Final = await authorize_team(store, caller, run.team_id, "write")
|
||||
if isinstance(permission, DefinitionError):
|
||||
raise_public(permission)
|
||||
await runs.cancel(run_id)
|
||||
return _accepted(run)
|
||||
|
||||
|
||||
async def list_engine_runs(
|
||||
*,
|
||||
workflow_id: str,
|
||||
team_id: str | None,
|
||||
status: str | None,
|
||||
limit: int,
|
||||
after: str | None,
|
||||
caller: UserAPIKeyAuth,
|
||||
definitions: WorkflowDefinitionStore | None,
|
||||
executions: WorkflowExecutionStore | None,
|
||||
) -> RunPage:
|
||||
store, runs = _stores(definitions, executions)
|
||||
definition: Final = await store.get_draft(workflow_id, team_id)
|
||||
if definition is None:
|
||||
_error(ExecutionError("not_found", "Workflow not found"))
|
||||
permission: Final = await authorize_team(store, caller, definition.draft.team_id, "invoke")
|
||||
if isinstance(permission, DefinitionError):
|
||||
raise_public(permission)
|
||||
may_read_team: Final = user_api_key_has_admin_view(caller) or is_team_admin(caller, permission)
|
||||
if not may_read_team and not caller.token:
|
||||
_error(ExecutionError("forbidden", "A run reader identity is required"))
|
||||
try:
|
||||
state: Final = RunState(status) if status is not None else None
|
||||
except ValueError:
|
||||
_error(ExecutionError("invalid_status", "Unknown workflow run status"))
|
||||
result: Final = await runs.list_runs(
|
||||
workflow_id,
|
||||
definition.draft.team_id,
|
||||
None if may_read_team else caller.token,
|
||||
state,
|
||||
min(limit, 100),
|
||||
after,
|
||||
)
|
||||
if isinstance(result, ExecutionError):
|
||||
_error(result)
|
||||
return result
|
||||
|
||||
|
||||
@router.get("/v1/workflows/{workflow_id}/runs")
|
||||
async def list_definition_runs(
|
||||
workflow_id: str,
|
||||
caller: _Caller,
|
||||
definitions: WorkflowDefinitionRepositoryDep,
|
||||
executions: WorkflowExecutionRepositoryDep,
|
||||
team_id: str | None = None,
|
||||
status: RunState | None = None,
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 50,
|
||||
after: Annotated[str | None, Query(max_length=512)] = None,
|
||||
) -> RunPage:
|
||||
return await list_engine_runs(
|
||||
workflow_id=workflow_id,
|
||||
team_id=team_id,
|
||||
status=status.value if status is not None else None,
|
||||
limit=limit,
|
||||
after=after,
|
||||
caller=caller,
|
||||
definitions=definitions,
|
||||
executions=executions,
|
||||
)
|
||||
|
||||
|
||||
async def _wait_for_run(runs: WorkflowExecutionStore, run_id: str) -> EngineRun:
|
||||
while True:
|
||||
current: EngineRun | None = await runs.get(run_id)
|
||||
if current is None:
|
||||
_error(ExecutionError("not_found", "Run not found"), run_id)
|
||||
if current.status in {RunState.completed, RunState.failed, RunState.cancelled}:
|
||||
return current
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
|
||||
async def invoke_workflow(
|
||||
request: Request,
|
||||
response: Response,
|
||||
caller: UserAPIKeyAuth,
|
||||
raw: object,
|
||||
definitions: WorkflowDefinitionStore | None,
|
||||
executions: WorkflowExecutionStore | None,
|
||||
) -> Mapping[str, JsonValue] | JSONResponse:
|
||||
store, runs = _stores(definitions, executions)
|
||||
try:
|
||||
data: Final = WorkflowInvocation.model_validate(raw)
|
||||
except ValidationError:
|
||||
_error(ExecutionError("unsupported_request", "Use model, messages, user, metadata and stream=false"))
|
||||
snapshot: Final = await published_workflow(store, caller, data.model.removeprefix("workflow/"))
|
||||
if isinstance(snapshot, DefinitionError):
|
||||
raise_public(snapshot)
|
||||
key: Final = request.headers.get("Idempotency-Key")
|
||||
if key is not None and not 1 <= len(key) <= 256:
|
||||
_error(ExecutionError("invalid_idempotency_key", "Idempotency-Key must contain 1 to 256 characters"))
|
||||
admitted: Final = await _admit(
|
||||
request,
|
||||
runs,
|
||||
snapshot,
|
||||
caller,
|
||||
data.messages,
|
||||
data.user or caller.end_user_id,
|
||||
_idempotency_key("invoke", key),
|
||||
_digest(data.messages, data.user or caller.end_user_id, None),
|
||||
draft=False,
|
||||
)
|
||||
response.headers["x-litellm-workflow-run-id"] = admitted.run_id # rebind-ok: FastAPI response header
|
||||
if request.headers.get("Prefer") == "respond-async":
|
||||
return _accepted(admitted)
|
||||
completed: Final = await _wait_for_run(runs, admitted.run_id)
|
||||
if completed.status is not RunState.completed:
|
||||
_error(
|
||||
ExecutionError("execution_failed", "Workflow execution did not complete. Inspect the run trace"),
|
||||
admitted.run_id,
|
||||
)
|
||||
return TypeAdapter(Mapping[str, JsonValue]).validate_python(completed.output)
|
||||
|
|
@ -18,6 +18,7 @@ from collections.abc import Mapping, Sequence
|
|||
from typing import Final, Literal, Protocol, TypedDict
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
try:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
|
@ -31,6 +32,7 @@ from litellm.proxy._types import (
|
|||
user_api_key_has_admin_view,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.deps.repositories import WorkflowDefinitionRepositoryDep, WorkflowExecutionRepositoryDep
|
||||
from litellm.repositories.table_repositories import (
|
||||
WorkflowEventRepository,
|
||||
WorkflowMessageRepository,
|
||||
|
|
@ -120,6 +122,7 @@ class _RunCreateData(TypedDict, total=False):
|
|||
|
||||
|
||||
class _RunWhere(TypedDict, total=False):
|
||||
engine_owned: ReadOnly[bool]
|
||||
workflow_type: str
|
||||
status: str | Mapping[str, Sequence[str]]
|
||||
created_by: str
|
||||
|
|
@ -201,6 +204,8 @@ async def _require_run(
|
|||
prisma_client: object,
|
||||
run_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
*,
|
||||
mutation: bool = False,
|
||||
) -> _RunRow:
|
||||
"""Return the run or raise 404. For non-admin callers, also enforce key ownership."""
|
||||
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id})
|
||||
|
|
@ -210,6 +215,12 @@ async def _require_run(
|
|||
caller: Final = _caller_key(user_api_key_dict)
|
||||
if not caller or run.created_by != caller:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
if getattr(run, "engine_owned", False) is True:
|
||||
if mutation:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="Gateway executions cannot be changed through tracking endpoints"
|
||||
)
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
return run
|
||||
|
||||
|
||||
|
|
@ -263,17 +274,36 @@ async def list_workflow_runs(
|
|||
status: str | None = Query(None),
|
||||
limit: int = Query(50, ge=1, le=250),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
workflow_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
after: str | None = None,
|
||||
definitions: WorkflowDefinitionRepositoryDep = None,
|
||||
executions: WorkflowExecutionRepositoryDep = None,
|
||||
):
|
||||
"""List workflow runs. Filter by workflow_type and/or status.
|
||||
|
||||
Non-admin callers only see runs created by their own API key.
|
||||
"""
|
||||
if workflow_id is not None:
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import list_engine_runs
|
||||
|
||||
return await list_engine_runs(
|
||||
workflow_id=workflow_id,
|
||||
team_id=team_id,
|
||||
status=status,
|
||||
limit=limit,
|
||||
after=after,
|
||||
caller=user_api_key_dict,
|
||||
definitions=definitions,
|
||||
executions=executions,
|
||||
)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
where: Final[_RunWhere] = {}
|
||||
where: Final[_RunWhere] = {"engine_owned": False}
|
||||
if workflow_type:
|
||||
where["workflow_type"] = workflow_type
|
||||
if status:
|
||||
|
|
@ -318,7 +348,7 @@ async def get_workflow_run(
|
|||
where={"run_id": run_id},
|
||||
include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}},
|
||||
)
|
||||
if run is None:
|
||||
if run is None or getattr(run, "engine_owned", False) is True:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
if not user_api_key_has_admin_view(user_api_key_dict):
|
||||
caller: Final = _caller_key(user_api_key_dict)
|
||||
|
|
@ -360,7 +390,7 @@ async def update_workflow_run(
|
|||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
|
||||
# Enforce ownership before writing.
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
try:
|
||||
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update(
|
||||
|
|
@ -398,7 +428,7 @@ async def append_workflow_event(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
new_status: Final = _EVENT_STATUS_MAP.get(data.event_type)
|
||||
|
||||
|
|
@ -487,7 +517,7 @@ async def append_workflow_message(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict, mutation=True)
|
||||
|
||||
for attempt in range(_MAX_SEQUENCE_RETRIES):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -847,6 +847,12 @@ from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router
|
|||
from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import (
|
||||
router as user_agent_analytics_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_definition_endpoints import (
|
||||
router as workflow_definition_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
|
||||
router as workflow_execution_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_management_endpoints import (
|
||||
router as workflow_management_router,
|
||||
)
|
||||
|
|
@ -1973,9 +1979,12 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
|
|||
async with AsyncExitStack() as admin_mcp_stack:
|
||||
try:
|
||||
await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app))
|
||||
from litellm.proxy.workflows.lifespan import workflow_lifespan
|
||||
|
||||
workflow_runner: Final = await admin_mcp_stack.enter_async_context(workflow_lifespan(database))
|
||||
gc.freeze()
|
||||
verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count())
|
||||
yield state
|
||||
yield {**state, "workflow_runner": workflow_runner}
|
||||
finally:
|
||||
if model_info_scheduler is not None and model_info_scheduler.running:
|
||||
model_info_scheduler.remove_job("refresh_model_info")
|
||||
|
|
@ -13156,6 +13165,26 @@ async def chat_completion(
|
|||
global general_settings, user_debug, proxy_logging_obj, llm_model_list
|
||||
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
||||
data: Final = await read_request_body(request=request)
|
||||
workflow_data: Final = as_str_mapping(_UNTYPED_VALUE.validate_python(data))
|
||||
workflow_model: Final = workflow_data.get("model") if workflow_data is not None else None
|
||||
if isinstance(workflow_model, str) and workflow_model.startswith("workflow/"):
|
||||
from litellm.proxy.deps.database import get_database
|
||||
from litellm.proxy.deps.repositories import (
|
||||
get_workflow_definition_repository,
|
||||
get_workflow_execution_repository,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import invoke_workflow
|
||||
|
||||
workflow_database: Final = await get_database(request)
|
||||
await release_unbound_budget_reservation(user_api_key_dict.budget_reservation)
|
||||
return await invoke_workflow(
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
workflow_data,
|
||||
await get_workflow_definition_repository(workflow_database),
|
||||
await get_workflow_execution_repository(workflow_database),
|
||||
)
|
||||
if user_api_key_dict is not None:
|
||||
if not isinstance(data.get("metadata"), dict):
|
||||
# Covers both missing and JSON-string metadata (multipart /
|
||||
|
|
@ -22432,6 +22461,8 @@ app.include_router(model_access_group_management_router)
|
|||
app.include_router(auto_router_management_router)
|
||||
app.include_router(tag_management_router)
|
||||
app.include_router(workflow_management_router)
|
||||
app.include_router(workflow_execution_router)
|
||||
app.include_router(workflow_definition_router)
|
||||
app.include_router(memory_router)
|
||||
app.include_router(lens_router)
|
||||
app.include_router(plugin_router)
|
||||
|
|
|
|||
|
|
@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
|
|||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowDefinition {
|
||||
workflow_id String @id @default(uuid())
|
||||
team_id String
|
||||
name String
|
||||
etag String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_snapshot_id String?
|
||||
published_etag String?
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
snapshots LiteLLM_WorkflowSnapshot[]
|
||||
|
||||
@@unique([team_id, name])
|
||||
@@index([team_id, workflow_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowSnapshot {
|
||||
snapshot_id String @id @default(uuid())
|
||||
workflow_id String
|
||||
team_id String
|
||||
name String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_by String?
|
||||
created_at DateTime @default(now())
|
||||
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
|
||||
|
||||
@@index([workflow_id, created_at, snapshot_id])
|
||||
@@index([team_id, snapshot_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
workflow_id String?
|
||||
snapshot_id String?
|
||||
team_id String?
|
||||
engine_owned Boolean @default(false)
|
||||
request_digest String?
|
||||
idempotency_key String?
|
||||
cancel_requested Boolean @default(false)
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
|
|
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
|
|||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
|
||||
@@index([workflow_id, team_id, created_at, run_id])
|
||||
@@index([team_id, engine_owned, status])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
|
|
|
|||
|
|
@ -458,7 +458,7 @@ async def release_or_invalidate_budget_reservation(
|
|||
budget_reservation["finalized"] = True
|
||||
|
||||
|
||||
async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None:
|
||||
async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object] | None) -> None:
|
||||
"""Release a reservation no logging callback took ownership of, once the request ended.
|
||||
|
||||
A handler whose litellm call never builds a logging object (batch cancel, file
|
||||
|
|
|
|||
|
|
@ -4367,6 +4367,8 @@ class ProxyLogging:
|
|||
data: dict,
|
||||
response: LLMResponseTypes,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
*,
|
||||
guardrails_only: bool = False,
|
||||
):
|
||||
"""
|
||||
Allow user to modify outgoing data
|
||||
|
|
@ -4462,6 +4464,8 @@ class ProxyLogging:
|
|||
#################################################################
|
||||
|
||||
for callback in other_callbacks:
|
||||
if guardrails_only and not callback.enforces_request_content:
|
||||
continue
|
||||
callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, response=response
|
||||
)
|
||||
|
|
|
|||
91
litellm/proxy/workflows/LITEFLOW.md
Normal file
91
litellm/proxy/workflows/LITEFLOW.md
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
# LiteFlow gateway workflows
|
||||
|
||||
The gateway accepts LiteFlow compiled contract version `1` and executes sequential `chat_completion` and `python` nodes. Python nodes use the Monty adapter. Conditional, parallel and streaming execution are not supported
|
||||
|
||||
Workflow definitions belong to a team. Authors register and edit drafts. A team administrator who did not contribute to the draft publishes an immutable snapshot. Invocation requires a virtual key assigned to that team, with access to the workflow model and each child model
|
||||
|
||||
## SDK example
|
||||
|
||||
The proxy imports the actual LiteFlow SDK models and digest helper, pinned to SDK revision `27bcbce4`. Its compiled JSON is validated at the server boundary because the SDK does not export a server bundle parser or executor. No SDK implementation is copied
|
||||
|
||||
This contract matches LiteFlow PRs [3](https://github.com/BerriAI/liteflow/pull/3) and [5](https://github.com/BerriAI/liteflow/pull/5). Configure `gateway-model` on the proxy and provide two different team members' keys
|
||||
|
||||
```python
|
||||
import os
|
||||
import time
|
||||
|
||||
from liteflow import ChatCompletion, GatewayClient, Input, Workflow
|
||||
|
||||
workflow = Workflow(
|
||||
name="welcome",
|
||||
nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)],
|
||||
)
|
||||
messages = [{"role": "user", "content": "Write one short welcome sentence."}]
|
||||
|
||||
with GatewayClient(
|
||||
base_url=os.environ["PROXY_BASE_URL"],
|
||||
api_key=os.environ["WORKFLOW_AUTHOR_KEY"],
|
||||
team_id=os.environ["WORKFLOW_TEAM_ID"],
|
||||
) as author, GatewayClient(
|
||||
base_url=os.environ["PROXY_BASE_URL"],
|
||||
api_key=os.environ["WORKFLOW_PUBLISHER_KEY"],
|
||||
team_id=os.environ["WORKFLOW_TEAM_ID"],
|
||||
) as publisher:
|
||||
draft = author.register(workflow)
|
||||
author.validate(draft)
|
||||
publisher.publish(draft)
|
||||
accepted = author.test(draft, messages=messages, authorize_billing=True)
|
||||
for _ in range(300):
|
||||
trace = author.inspect(accepted["run_id"])
|
||||
print(trace)
|
||||
if trace.state.value in ("completed", "failed", "cancelled"):
|
||||
break
|
||||
time.sleep(1)
|
||||
result = author.run("welcome", messages=messages)
|
||||
print(result.run_id, result.output)
|
||||
```
|
||||
|
||||
## API contract
|
||||
|
||||
| Operation | Endpoint | Contract |
|
||||
| --- | --- | --- |
|
||||
| Register | `POST /v1/workflows` | `{"bundle": <compiled bundle>}`, optional `team_id` query, returns `201` |
|
||||
| List definitions | `GET /v1/workflows` | `team_id`, `limit`, `cursor`; returns `workflows`, `next_cursor` |
|
||||
| Retrieve draft | `GET /v1/workflows/{workflow_id}` | Returns the draft and `ETag` header |
|
||||
| Replace draft | `PUT /v1/workflows/{workflow_id}` | Compiled bundle and exact `If-Match` |
|
||||
| Validate | `POST /v1/workflows/{workflow_id}/validate` | Exact `If-Match`, no body required |
|
||||
| Publish | `POST /v1/workflows/{workflow_id}/publish` | Exact `If-Match` and `{"digest": "<draft digest>"}` |
|
||||
| Test draft | `POST /v1/workflows/{workflow_id}/test` | Exact `If-Match`, `messages`, `authorize_billing: true`; returns `202` |
|
||||
| Invoke published | `POST /v1/chat/completions` | `model: "workflow/<name>"`, `messages`, `stream: false` |
|
||||
| List runs | `GET /v1/workflows/{workflow_id}/runs` | `limit`, `status`, `after`; returns `runs`, `count`, `next_cursor` |
|
||||
| Inspect | `GET /v1/workflows/runs/{run_id}/steps` | Run state and typed step traces |
|
||||
| Pinned definition | `GET /v1/workflows/runs/{run_id}/definition` | Exact snapshot used by that run |
|
||||
| Cancel | `POST /v1/workflows/runs/{run_id}/cancel` | Returns `202`; poll until terminal state |
|
||||
|
||||
Draft responses contain `name`, `workflow_id`, `team_id`, `etag`, `digest`, `bundle` and `published_snapshot_id`. Missing or stale `If-Match` returns `412`. Validation checks the compiled graph and Python program without executing model or Python nodes
|
||||
|
||||
Test runs and chat invocations accept `Idempotency-Key`. Retrying the same operation with the same key and input returns the original run. Reusing that key with different input returns `409`. A test run pins the exact draft revision even if the draft changes later
|
||||
|
||||
For live UI progress, send `Prefer: respond-async` on a published chat invocation. The gateway returns `202` with `run_id`, `snapshot_id`, `trace_url` and the `x-litellm-workflow-run-id` header. The ordinary SDK call waits for a Chat Completions response carrying the same run-ID header. It waits until execution finishes or the workflow deadline expires; clients with shorter HTTP timeouts should use asynchronous invocation and an idempotency key
|
||||
|
||||
## Trace and spend data
|
||||
|
||||
Draw the pinned bundle's `nodes` and resolve each binding's `from` reference. Match `step.node_id` to `node.id`. Run and node states are `pending`, `running`, `completed`, `failed` and `cancelled`
|
||||
|
||||
Each step exposes typed `input` and `output`, `attempts`, `started_at`, `finished_at`, `deadline`, sanitized `error`, token counts and `response_cost`. A run-level `error` explains failures before any node starts
|
||||
|
||||
Model steps carry a deterministic `call_id` before dispatch. `request_id` and `request_ids` link to actual `LiteLLM_SpendLogs.request_id` values. Spend writes are asynchronous, so receipt IDs can be absent on the first read and appear on later polls. The trace reader sums persisted attempts once and does not duplicate earlier checkpoints
|
||||
|
||||
Each child uses the standard gateway authorization, hooks, guardrails, limits and billing path. The workflow wrapper emits no additional bill. Python nodes have no model spend receipt. Known costs remain visible when output validation fails
|
||||
|
||||
Definitions and snapshots use separate tables. Execution reuses `LiteLLM_WorkflowRun` and `LiteLLM_WorkflowEvent`. Engine-owned runs are exposed only through the scoped execution endpoints; existing tracking endpoints continue to serve externally tracked runs
|
||||
|
||||
## Execution limits and interruption
|
||||
|
||||
The initial runner allows four active runs per worker and 64 pending or running runs per team. The request that creates the durable run owns its execution task; idempotent retries observe that same run. Workers drain for up to 30 seconds on graceful shutdown, then fail unfinished work. After a crash, unfinished runs fail at their original deadline when read or during subsequent admission. Stored traces survive; executions are never replayed automatically
|
||||
|
||||
Chat retries are limited to rate limits confirmed before provider dispatch. Provider failures, timeouts and attempts with receipts are terminal because replaying an uncertain request could charge twice
|
||||
|
||||
Cancellation is cooperative. It stops downstream nodes but cannot undo a model request already sent. Trace receipts can still arrive after cancellation
|
||||
|
||||
Ordinary virtual keys can read their own runs. Team administrators and proxy administrators can inspect team runs within their credential scope. Definition access does not grant another caller access to private run inputs or outputs
|
||||
218
litellm/proxy/workflows/definitions.py
Normal file
218
litellm/proxy/workflows/definitions.py
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
|
||||
from jsonschema import Draft202012Validator
|
||||
from jsonschema.exceptions import SchemaError
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from referencing.jsonschema import DRAFT202012
|
||||
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.models.workflow_definition import (
|
||||
Binding,
|
||||
ChatCompletionNode,
|
||||
DefinitionError,
|
||||
PythonNode,
|
||||
ValueBinding,
|
||||
WorkflowBundle,
|
||||
WorkflowDefinitionStore,
|
||||
WorkflowDraft,
|
||||
WorkflowSnapshot,
|
||||
parsed_bundle,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.management.teams.authz import is_team_admin, is_team_member
|
||||
from litellm.proxy.workflows.monty.compiler import compile_program
|
||||
from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits
|
||||
|
||||
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
|
||||
_SCHEMA: Final = TypeAdapter[Mapping[str, JsonValue] | bool](Mapping[str, JsonValue] | bool)
|
||||
_LIMITS: Final = MontyLimits()
|
||||
_WRITER_ROLES: Final = frozenset(
|
||||
{LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.ORG_ADMIN, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM}
|
||||
)
|
||||
|
||||
|
||||
def requested_team(caller: UserAPIKeyAuth, team_id: str | None) -> str | DefinitionError:
|
||||
selected: Final = team_id if team_id is not None else caller.team_id
|
||||
if not selected:
|
||||
return DefinitionError("team_required", "Specify team_id for this workflow", "team_id")
|
||||
if caller.team_id and caller.team_id != selected and not caller.is_session_token:
|
||||
return DefinitionError("not_found", "Workflow not found")
|
||||
return selected
|
||||
|
||||
|
||||
def author_identity(caller: UserAPIKeyAuth) -> str | DefinitionError:
|
||||
if not caller.user_id or caller.user_role not in _WRITER_ROLES:
|
||||
return DefinitionError("forbidden", "Workflow authoring requires a team member identity")
|
||||
return caller.user_id
|
||||
|
||||
|
||||
async def authorize_team(
|
||||
store: WorkflowDefinitionStore,
|
||||
caller: UserAPIKeyAuth,
|
||||
team_id: str,
|
||||
action: Literal["read", "write", "publish", "invoke"],
|
||||
) -> LiteLLM_TeamTable | DefinitionError:
|
||||
scope: Final = requested_team(caller, team_id)
|
||||
if isinstance(scope, DefinitionError):
|
||||
return scope
|
||||
team: Final = await store.get_team(scope)
|
||||
if team is None or team.blocked:
|
||||
return DefinitionError("not_found", "Workflow not found")
|
||||
if action == "invoke" and (caller.team_id == scope or user_api_key_has_admin_view(caller)):
|
||||
return team
|
||||
if not user_api_key_has_admin_view(caller) and not is_team_member(caller, team):
|
||||
return DefinitionError("not_found", "Workflow not found")
|
||||
if action == "read":
|
||||
return team
|
||||
if isinstance(identity := author_identity(caller), DefinitionError):
|
||||
return identity
|
||||
if action == "publish" and caller.user_role != LitellmUserRoles.PROXY_ADMIN and not is_team_admin(caller, team):
|
||||
return DefinitionError("forbidden", "Publishing requires a team administrator")
|
||||
return team
|
||||
|
||||
|
||||
async def authorized_draft(
|
||||
store: WorkflowDefinitionStore,
|
||||
caller: UserAPIKeyAuth,
|
||||
workflow_id: str,
|
||||
team_id: str | None = None,
|
||||
*,
|
||||
action: Literal["read", "write", "publish"] = "read",
|
||||
) -> WorkflowDraft | DefinitionError:
|
||||
key_team: Final = None if caller.is_session_token else caller.team_id
|
||||
definition: Final = await store.get_draft(workflow_id, team_id if team_id is not None else key_team)
|
||||
if definition is None:
|
||||
return DefinitionError("not_found", "Workflow not found")
|
||||
access: Final = await authorize_team(store, caller, definition.draft.team_id, action)
|
||||
return access if isinstance(access, DefinitionError) else definition.draft
|
||||
|
||||
|
||||
async def published_workflow(
|
||||
store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, name: str, team_id: str | None = None
|
||||
) -> WorkflowSnapshot | DefinitionError:
|
||||
scope: Final = requested_team(caller, team_id)
|
||||
if isinstance(scope, DefinitionError):
|
||||
return scope
|
||||
access: Final = await authorize_team(store, caller, scope, "invoke")
|
||||
if isinstance(access, DefinitionError):
|
||||
return access
|
||||
snapshot: Final = await store.get_published(name, scope)
|
||||
return snapshot if snapshot is not None else DefinitionError("not_found", "Published workflow not found")
|
||||
|
||||
|
||||
async def snapshot_by_id(
|
||||
store: WorkflowDefinitionStore, caller: UserAPIKeyAuth, snapshot_id: str
|
||||
) -> WorkflowSnapshot | DefinitionError:
|
||||
key_team: Final = None if caller.is_session_token else caller.team_id
|
||||
snapshot: Final = await store.get_snapshot(snapshot_id, key_team)
|
||||
if snapshot is None:
|
||||
return DefinitionError("not_found", "Workflow snapshot not found")
|
||||
access: Final = await authorize_team(store, caller, snapshot.team_id, "read")
|
||||
return access if isinstance(access, DefinitionError) else snapshot
|
||||
|
||||
|
||||
def check_etag(draft: WorkflowDraft, etag: str | None) -> DefinitionError | None:
|
||||
if etag is None:
|
||||
return DefinitionError("precondition_required", "If-Match is required", "If-Match")
|
||||
if draft.etag != etag:
|
||||
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
|
||||
return None
|
||||
|
||||
|
||||
async def freeze_draft(
|
||||
store: WorkflowDefinitionStore,
|
||||
caller: UserAPIKeyAuth,
|
||||
workflow_id: str,
|
||||
etag: str | None,
|
||||
team_id: str | None = None,
|
||||
) -> WorkflowSnapshot | DefinitionError:
|
||||
draft: Final = await authorized_draft(store, caller, workflow_id, team_id, action="write")
|
||||
if isinstance(draft, DefinitionError):
|
||||
return draft
|
||||
if (failure := check_etag(draft, etag)) is not None:
|
||||
return failure
|
||||
if (invalid := validate_bundle(parsed_bundle(draft.bundle))) is not None:
|
||||
return invalid
|
||||
return await store.freeze_draft(workflow_id, draft.team_id, draft.etag)
|
||||
|
||||
|
||||
def _local_references(value: Mapping[str, JsonValue] | bool, depth: int = 0) -> bool:
|
||||
if depth > 64:
|
||||
return False
|
||||
if isinstance(value, Mapping):
|
||||
if any(keyword in value for keyword in ("$id", "$dynamicRef", "$recursiveRef", "$dynamicAnchor")):
|
||||
return False
|
||||
reference: Final = value.get("$ref")
|
||||
if reference is not None and (not isinstance(reference, str) or not reference.startswith("#/")):
|
||||
return False
|
||||
return all(
|
||||
_local_references(_SCHEMA.validate_python(child), depth + 1) for child in DRAFT202012.subresources_of(value)
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _schema_error(schema: Mapping[str, JsonValue], param: str) -> DefinitionError | None:
|
||||
try:
|
||||
Draft202012Validator.check_schema(schema)
|
||||
if not _local_references(schema):
|
||||
return DefinitionError("invalid_bundle", "Schemas must use local references and bounded nesting", param)
|
||||
except (SchemaError, RecursionError):
|
||||
return DefinitionError("invalid_bundle", "Invalid JSON schema", param)
|
||||
return None
|
||||
|
||||
|
||||
def _binding_error(binding: Binding, previous: frozenset[str], param: str) -> DefinitionError | None:
|
||||
if isinstance(binding, ValueBinding):
|
||||
return None
|
||||
if binding.source == "$input" and binding.output == "messages":
|
||||
return None
|
||||
if binding.source in previous and binding.output == "result":
|
||||
return None
|
||||
return DefinitionError("invalid_bundle", "Bindings must reference input messages or an earlier node result", param)
|
||||
|
||||
|
||||
def _node_error(node: ChatCompletionNode | PythonNode, previous: frozenset[str]) -> DefinitionError | None:
|
||||
if node.id in previous or node.id == "runs":
|
||||
return DefinitionError("invalid_bundle", "Node IDs must be unique and cannot be reserved", f"nodes.{node.id}")
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
if node.model.startswith("workflow/") or "://" in node.model:
|
||||
return DefinitionError("invalid_bundle", "Use a configured gateway model group", f"nodes.{node.id}.model")
|
||||
if (failure := _binding_error(node.messages, previous, f"nodes.{node.id}.messages")) is not None:
|
||||
return failure
|
||||
if isinstance(node.messages, ValueBinding):
|
||||
try:
|
||||
_MESSAGES.validate_python(node.messages.value, strict=True)
|
||||
except ValidationError:
|
||||
return DefinitionError("invalid_bundle", "Invalid chat messages", f"nodes.{node.id}.messages")
|
||||
if node.response_format is not None:
|
||||
return _schema_error(node.response_format.json_schema.schema_, f"nodes.{node.id}.response_format")
|
||||
return None
|
||||
if set(node.inputs) != set(node.program.parameter_names):
|
||||
return DefinitionError(
|
||||
"invalid_bundle", "Bind exactly the Python program parameters", f"nodes.{node.id}.inputs"
|
||||
)
|
||||
for name, binding in node.inputs.items():
|
||||
if (failure := _binding_error(binding, previous, f"nodes.{node.id}.inputs.{name}")) is not None:
|
||||
return failure
|
||||
compiled: Final = compile_program(node.program, _LIMITS)
|
||||
if isinstance(compiled, MontyFailure):
|
||||
return DefinitionError("invalid_bundle", compiled.message[:1024], f"nodes.{node.id}.program")
|
||||
return None
|
||||
|
||||
|
||||
def validate_bundle(bundle: WorkflowBundle) -> DefinitionError | None:
|
||||
if bundle.name == "runs":
|
||||
return DefinitionError("invalid_bundle", "The workflow name is reserved", "bundle.name")
|
||||
if len(bundle.model_dump_json(by_alias=True).encode()) > 1024 * 1024:
|
||||
return DefinitionError("invalid_bundle", "Workflow bundles cannot exceed 1 MiB", "bundle")
|
||||
if (failure := _schema_error(bundle.input_schema, "bundle.input_schema")) is not None:
|
||||
return failure
|
||||
for index, node in enumerate(bundle.nodes):
|
||||
previous: frozenset[str] = frozenset(prior.id for prior in bundle.nodes[:index])
|
||||
if (node_failure := _node_error(node, previous)) is not None:
|
||||
return node_failure
|
||||
if bundle.output.source != bundle.nodes[-1].id or bundle.output.output != "result":
|
||||
return DefinitionError("invalid_bundle", "The workflow output must be the final node result", "bundle.output")
|
||||
return None
|
||||
287
litellm/proxy/workflows/gateway.py
Normal file
287
litellm/proxy/workflows/gateway.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from contextvars import Context
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timezone
|
||||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import Request, Response
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
|
||||
from pydantic_core import from_json
|
||||
from starlette.types import Message
|
||||
|
||||
from litellm.exceptions import BudgetExceededError
|
||||
from litellm.models.workflow_definition import ChatCompletionNode
|
||||
from litellm.models.workflow_execution import ChildCallResult, EngineRun, ExecutionError, StepTrace
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import authorize_internal_virtual_key
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
safe_set_request_parsed_body, # pyright: ignore[reportUnknownVariableType] # legacy parsed-body API
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers, redact_credential_headers
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_budget_reservation_on_cancel, # pyright: ignore[reportUnknownVariableType] # legacy reservation dictionary
|
||||
release_unbound_budget_reservation,
|
||||
)
|
||||
from litellm.proxy.workflows.monty.models import MontyErrorKind
|
||||
from litellm.proxy.workflows.monty.validation import ValidatorProcess
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from asyncio import timeout
|
||||
else:
|
||||
from async_timeout import timeout
|
||||
|
||||
_ROUTE: Final = "/v1/chat/completions"
|
||||
_MESSAGES: Final = TypeAdapter[list[ChatCompletionMessageParam]](
|
||||
Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
|
||||
)
|
||||
_JSON: Final = TypeAdapter[JsonValue](JsonValue, config=ConfigDict(allow_inf_nan=False))
|
||||
_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_STRING: Final = TypeAdapter[str | None](str | None)
|
||||
_NUMBER: Final = TypeAdapter[float | None](float | None)
|
||||
|
||||
|
||||
def _child_request(
|
||||
parent: Request,
|
||||
data: dict[str, object], # mutable-ok: shared request body mutated by the existing proxy pipeline
|
||||
call_id: str,
|
||||
session_id: str,
|
||||
) -> Request:
|
||||
encoded: Final = json.dumps(data, allow_nan=False).encode()
|
||||
incoming: Final[asyncio.Queue[Message]] = asyncio.Queue(maxsize=1)
|
||||
incoming.put_nowait({"type": "http.request", "body": encoded, "more_body": False})
|
||||
|
||||
async def receive() -> Message:
|
||||
return await incoming.get()
|
||||
|
||||
session_headers: Final = (
|
||||
()
|
||||
if get_chain_id_from_headers(dict(parent.headers)) or parent.headers.get("traceparent")
|
||||
else ((b"x-litellm-session-id", session_id.encode()),)
|
||||
)
|
||||
original_headers: Final = dict(parent.headers)
|
||||
headers: Final = tuple(
|
||||
(key.encode("latin-1"), value.encode("latin-1"))
|
||||
for key, value in redact_credential_headers(original_headers).items()
|
||||
if key.lower() not in {"content-type", "content-length", "x-litellm-call-id"} and value == original_headers[key]
|
||||
)
|
||||
shared_state: Final = _FIELDS.validate_python(parent.scope.get("state", {}))
|
||||
request: Final = Request(
|
||||
{
|
||||
**parent.scope,
|
||||
"method": "POST",
|
||||
"path": _ROUTE,
|
||||
"raw_path": _ROUTE.encode(),
|
||||
"query_string": b"",
|
||||
"path_params": {},
|
||||
"state": {key: value for key, value in shared_state.items() if key == "database"},
|
||||
"headers": (
|
||||
*headers,
|
||||
*session_headers,
|
||||
(b"content-type", b"application/json"),
|
||||
(b"content-length", str(len(encoded)).encode()),
|
||||
(b"x-litellm-call-id", call_id.encode()),
|
||||
),
|
||||
},
|
||||
receive,
|
||||
)
|
||||
safe_set_request_parsed_body(request, data)
|
||||
return request
|
||||
|
||||
|
||||
def _execution_error(exc: Exception, receipt: ChildCallResult | None = None) -> ExecutionError:
|
||||
if isinstance(exc, BudgetExceededError) or (
|
||||
isinstance(exc, ProxyException) and exc.type == ProxyErrorTypes.budget_exceeded
|
||||
):
|
||||
return ExecutionError("budget_exceeded", "The execution identity has insufficient budget", receipt)
|
||||
status: Final = str(exc.code if isinstance(exc, ProxyException) else getattr(exc, "status_code", None))
|
||||
if status in {"401", "403"}:
|
||||
return ExecutionError("authorization_failed", "The execution identity is not authorized for this node", receipt)
|
||||
if status == "429":
|
||||
return ExecutionError("rate_limit", "The gateway could not admit this node", receipt)
|
||||
if status in {"400", "422"}:
|
||||
return ExecutionError("request_rejected", "The gateway rejected this node", receipt)
|
||||
if status == "408" or status.startswith("5"):
|
||||
return ExecutionError("provider_error", "The provider could not complete this node", receipt)
|
||||
return ExecutionError("runtime_error", "The gateway could not complete this node", receipt)
|
||||
|
||||
|
||||
async def _authorize_child(
|
||||
request: Request,
|
||||
run: EngineRun,
|
||||
data: dict[str, object], # mutable-ok: authentication annotates the shared proxy request body
|
||||
) -> UserAPIKeyAuth | ExecutionError:
|
||||
try:
|
||||
return await authorize_internal_virtual_key(
|
||||
key_hash=run.created_by, request=request, request_data=data, route=_ROUTE, expected_team_id=run.team_id
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # authentication failures are known to precede dispatch
|
||||
failure: Final = _execution_error(exc)
|
||||
return replace(failure, code="admission_rate_limit") if failure.code == "rate_limit" else failure
|
||||
|
||||
|
||||
def _receipt(
|
||||
processor: ProxyBaseLLMRequestProcessing, response: Response, raw: object = None
|
||||
) -> ChildCallResult | None:
|
||||
logging: Final = processor.logging_obj
|
||||
prepared: Final = logging.prepared_cost if logging is not None else None
|
||||
finalized: Final = logging.finalized_cost if logging is not None else None
|
||||
fields: Final = _FIELDS.validate_python(raw.model_dump() if isinstance(raw, BaseModel) else raw or {})
|
||||
reported: Final = fields.get("usage")
|
||||
usage: Final = (
|
||||
prepared.usage if prepared is not None else Usage.model_validate(reported) if reported is not None else None
|
||||
)
|
||||
if usage is None or any(
|
||||
type(value) is not int or value < 0 for value in (usage.prompt_tokens, usage.completion_tokens)
|
||||
):
|
||||
return None
|
||||
try:
|
||||
cost: Final = finalized.amount if finalized is not None else float(response.headers["x-litellm-response-cost"])
|
||||
except (KeyError, ValueError):
|
||||
return ChildCallResult(None, usage.prompt_tokens, usage.completion_tokens, None, None)
|
||||
return ChildCallResult(
|
||||
None,
|
||||
usage.prompt_tokens,
|
||||
usage.completion_tokens,
|
||||
None,
|
||||
cost if math.isfinite(cost) and cost >= 0 else None,
|
||||
)
|
||||
|
||||
|
||||
async def _output(raw: object, node: ChatCompletionNode, deadline: float) -> JsonValue | ExecutionError:
|
||||
try:
|
||||
result: Final = ChatCompletion.model_validate(raw.model_dump() if isinstance(raw, BaseModel) else raw)
|
||||
except ValueError:
|
||||
return ExecutionError("invalid_response", "The gateway returned an invalid chat response")
|
||||
if len(result.choices) != 1:
|
||||
return ExecutionError("invalid_response", "The gateway returned an invalid chat response")
|
||||
choice: Final = result.choices[0]
|
||||
if choice.message.refusal or choice.finish_reason == "content_filter":
|
||||
return ExecutionError("content_rejected", "The gateway rejected the node output")
|
||||
content: Final = choice.message.content
|
||||
if content is None:
|
||||
return ExecutionError("invalid_response", "The model returned no assistant content")
|
||||
if len(content.encode()) > 1024 * 1024:
|
||||
return ExecutionError("output_too_large", "The model output exceeded the workflow limit")
|
||||
if node.response_format is None:
|
||||
return content
|
||||
try:
|
||||
output: Final = _JSON.validate_python(from_json(content, allow_inf_nan=False))
|
||||
except ValueError:
|
||||
return ExecutionError("output_invalid", "The model did not return valid JSON")
|
||||
validator: Final = ValidatorProcess()
|
||||
try:
|
||||
failure: Final = await validator.validate(
|
||||
output, dict(node.response_format.json_schema.schema_), deadline=deadline, output=True
|
||||
)
|
||||
if failure is None:
|
||||
return output
|
||||
if failure.kind is MontyErrorKind.TIMEOUT:
|
||||
return ExecutionError("timeout", "The node deadline expired")
|
||||
return ExecutionError("output_invalid", "The model output did not satisfy the node schema")
|
||||
finally:
|
||||
await validator.close()
|
||||
|
||||
|
||||
async def execute_chat_node(
|
||||
request: Request, run: EngineRun, node: ChatCompletionNode, step: StepTrace
|
||||
) -> ChildCallResult | ExecutionError:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
if step.deadline.tzinfo is None or step.deadline.utcoffset() is None:
|
||||
return ExecutionError("invalid_deadline", "The node deadline must include a timezone")
|
||||
if re.fullmatch(r"[0-9a-f]{64}", run.created_by) is None or not run.team_id or not step.call_id:
|
||||
return ExecutionError("authorization_failed", "A persisted team virtual key and child call ID are required")
|
||||
remaining: Final = min(node.timeout_seconds, (step.deadline - datetime.now(timezone.utc)).total_seconds())
|
||||
if remaining <= 0:
|
||||
return ExecutionError("timeout", "The node deadline expired")
|
||||
deadline: Final = time.monotonic() + remaining
|
||||
try:
|
||||
messages: Final = _MESSAGES.validate_python(step.input, strict=True)
|
||||
trace: Final = {
|
||||
"workflow_id": run.workflow_id,
|
||||
"run_id": run.run_id,
|
||||
"node_id": node.id,
|
||||
"attempt": step.attempts,
|
||||
"session_id": run.session_id,
|
||||
}
|
||||
data: Final[dict[str, object]] = { # mutable-ok: proxy processor and parsed-body cache share this object
|
||||
**node.model_dump(
|
||||
mode="json",
|
||||
by_alias=True,
|
||||
exclude_none=True,
|
||||
exclude={"id", "type", "messages", "timeout_seconds", "retry"},
|
||||
),
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"num_retries": 0,
|
||||
"timeout": remaining,
|
||||
"metadata": {**trace, "spend_logs_metadata": trace},
|
||||
**({"user": run.input.end_user_id} if run.input.end_user_id is not None else {}),
|
||||
}
|
||||
child: Final = _child_request(request, data, step.call_id, run.session_id)
|
||||
except (TypeError, ValueError):
|
||||
return ExecutionError("invalid_messages", "The node input must contain valid chat messages")
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
response: Final = Response()
|
||||
|
||||
async def execute() -> ChildCallResult | ExecutionError:
|
||||
try:
|
||||
async with timeout(remaining):
|
||||
auth: Final = await _authorize_child(child, run, data)
|
||||
if isinstance(auth, ExecutionError):
|
||||
return auth
|
||||
try:
|
||||
try:
|
||||
raw: Final = await processor.base_process_llm_request(
|
||||
request=child,
|
||||
fastapi_response=response,
|
||||
user_api_key_dict=auth,
|
||||
route_type="acompletion",
|
||||
proxy_logging_obj=proxy_server.proxy_logging_obj,
|
||||
general_settings=_FIELDS.validate_python(proxy_server.general_settings), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
|
||||
proxy_config=proxy_server.proxy_config,
|
||||
llm_router=proxy_server.llm_router,
|
||||
user_model=_STRING.validate_python(proxy_server.user_model), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
|
||||
user_temperature=_NUMBER.validate_python(proxy_server.user_temperature), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
|
||||
user_request_timeout=proxy_server.user_request_timeout,
|
||||
user_max_tokens=proxy_server.user_max_tokens,
|
||||
user_api_base=_STRING.validate_python(proxy_server.user_api_base), # pyright: ignore[reportUnknownMemberType] # legacy proxy global validated at this boundary
|
||||
version=proxy_server.version,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # provider and hook failures must run the proxy's accounting
|
||||
try:
|
||||
await processor.handle_llm_api_exception(
|
||||
e=exc, user_api_key_dict=auth, proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
)
|
||||
except Exception as failure: # noqa: BLE001 # failures become sanitized workflow values
|
||||
return _execution_error(failure, _receipt(processor, response))
|
||||
except asyncio.CancelledError:
|
||||
await release_budget_reservation_on_cancel(auth.budget_reservation)
|
||||
await proxy_server.proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth)
|
||||
raise
|
||||
receipt: Final = _receipt(processor, response, raw)
|
||||
output: Final = await _output(raw, node, deadline)
|
||||
if isinstance(output, ExecutionError):
|
||||
return replace(output, receipt=receipt)
|
||||
if receipt is None:
|
||||
return ExecutionError("invalid_response", "The gateway returned no valid usage")
|
||||
return replace(receipt, output=output)
|
||||
except asyncio.TimeoutError:
|
||||
return ExecutionError("timeout", "The node deadline expired", _receipt(processor, response))
|
||||
except Exception as exc: # noqa: BLE001 # auth and validator failures must not expose credentials or content
|
||||
return _execution_error(exc, _receipt(processor, response))
|
||||
finally:
|
||||
reservation: Final[object] = getattr(child.state, "budget_reservation", None)
|
||||
if isinstance(reservation, Mapping):
|
||||
await release_unbound_budget_reservation(reservation) # pyright: ignore[reportUnknownArgumentType] # preserve callback-shared reservation identity
|
||||
|
||||
return await Context().run(asyncio.create_task, execute())
|
||||
31
litellm/proxy/workflows/lifespan.py
Normal file
31
litellm/proxy/workflows/lifespan.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Final
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.deps.repositories import get_workflow_definition_repository, get_workflow_execution_repository
|
||||
from litellm.proxy.workflows.gateway import execute_chat_node
|
||||
from litellm.proxy.workflows.runner import WorkflowRunner
|
||||
from litellm.repositories.database import Database
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def workflow_lifespan(database: Database | None) -> AsyncGenerator[WorkflowRunner | None]:
|
||||
definitions: Final = await get_workflow_definition_repository(database)
|
||||
executions: Final = await get_workflow_execution_repository(database)
|
||||
runner: Final = (
|
||||
WorkflowRunner(executions, definitions, execute_chat_node)
|
||||
if definitions is not None and executions is not None
|
||||
else None
|
||||
)
|
||||
try:
|
||||
yield runner
|
||||
finally:
|
||||
if runner is not None:
|
||||
await runner.close()
|
||||
|
||||
|
||||
def get_workflow_runner(request: Request) -> WorkflowRunner | None:
|
||||
runner: Final[object] = getattr(request.state, "workflow_runner", None)
|
||||
return runner if isinstance(runner, WorkflowRunner) else None
|
||||
9
litellm/proxy/workflows/monty/_compat.py
Normal file
9
litellm/proxy/workflows/monty/_compat.py
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
import sys
|
||||
from typing import Final
|
||||
|
||||
__all__: Final = ("timeout",)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from asyncio import timeout
|
||||
else:
|
||||
from async_timeout import timeout
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
import keyword
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from enum import Enum
|
||||
from typing import Annotated, Literal, TypeAlias
|
||||
|
||||
from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, StringConstraints
|
||||
|
|
@ -29,7 +29,7 @@ class PythonProgram(BaseModel):
|
|||
source_sha256: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")]
|
||||
|
||||
|
||||
class MontyErrorKind(StrEnum):
|
||||
class MontyErrorKind(str, Enum):
|
||||
INVALID_PROGRAM = "INVALID_PROGRAM"
|
||||
UNSUPPORTED_SCHEMA = "UNSUPPORTED_SCHEMA"
|
||||
INPUT_INVALID = "INPUT_INVALID"
|
||||
|
|
@ -47,6 +47,9 @@ class MontyErrorKind(StrEnum):
|
|||
CANCELLED = "CANCELLED"
|
||||
RUNTIME_UNAVAILABLE = "RUNTIME_UNAVAILABLE"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MontyFailure:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Final
|
|||
from pydantic import JsonValue, TypeAdapter
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
|
||||
from ._compat import timeout
|
||||
from .compiler import escape_diagnostics
|
||||
from .models import CompiledProgram, JsonSchema, MontyErrorKind, MontyFailure, MontyLimits, MontySuccess
|
||||
from .validation import ValidatorProcess
|
||||
|
|
@ -147,7 +148,9 @@ class MontyRuntime:
|
|||
self._pool: AsyncMonty | None = None
|
||||
self._slots = asyncio.Semaphore(limits.max_concurrent_workers)
|
||||
self._admitted = 0
|
||||
self._validators = tuple(ValidatorProcess() for _ in range(limits.max_concurrent_workers))
|
||||
self._validators = tuple(
|
||||
ValidatorProcess(limits.worker_address_space_bytes) for _ in range(limits.max_concurrent_workers)
|
||||
)
|
||||
|
||||
async def start(self) -> MontyFailure | None:
|
||||
if self._pool is not None:
|
||||
|
|
@ -200,23 +203,20 @@ class MontyRuntime:
|
|||
return MontyFailure(MontyErrorKind.QUEUE_FULL, "Monty execution queue is full")
|
||||
self._admitted += 1
|
||||
try:
|
||||
async with asyncio.timeout_at(asyncio.get_running_loop().time() + deadline_monotonic - time.monotonic()):
|
||||
async with self._slots:
|
||||
worker: Final = next(worker for worker in self._validators if not worker.busy)
|
||||
worker.busy = True
|
||||
try:
|
||||
validated: Final = await self._validate(
|
||||
input, compiled.input_schema, worker, deadline_monotonic
|
||||
)
|
||||
if isinstance(validated, MontyFailure):
|
||||
return validated
|
||||
return await self._execute(compiled, validated, deadline_monotonic, started, worker)
|
||||
except asyncio.CancelledError:
|
||||
await worker.close()
|
||||
raise
|
||||
finally:
|
||||
worker.busy = False
|
||||
except TimeoutError:
|
||||
async with timeout(max(0, deadline_monotonic - time.monotonic())), self._slots:
|
||||
worker: Final = next(worker for worker in self._validators if not worker.busy)
|
||||
worker.busy = True
|
||||
try:
|
||||
validated: Final = await self._validate(input, compiled.input_schema, worker, deadline_monotonic)
|
||||
if isinstance(validated, MontyFailure):
|
||||
return validated
|
||||
return await self._execute(compiled, validated, deadline_monotonic, started, worker)
|
||||
except asyncio.CancelledError:
|
||||
await worker.close()
|
||||
raise
|
||||
finally:
|
||||
worker.busy = False
|
||||
except asyncio.TimeoutError:
|
||||
return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired")
|
||||
finally:
|
||||
self._admitted -= 1
|
||||
|
|
|
|||
|
|
@ -9,9 +9,10 @@ from typing import Final
|
|||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, StrictBool, StrictStr
|
||||
|
||||
from ._compat import timeout
|
||||
from ._validator_worker import MAX_MESSAGE_BYTES
|
||||
from .compiler import escape_diagnostics
|
||||
from .models import JsonSchema, MontyErrorKind, MontyFailure
|
||||
from .models import JsonSchema, MontyErrorKind, MontyFailure, MontyLimits
|
||||
|
||||
|
||||
class _Response(BaseModel):
|
||||
|
|
@ -21,7 +22,8 @@ class _Response(BaseModel):
|
|||
|
||||
|
||||
class ValidatorProcess:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, address_space_bytes: int = MontyLimits().worker_address_space_bytes) -> None:
|
||||
self._address_space_bytes = address_space_bytes
|
||||
self.busy = False
|
||||
self.process: asyncio.subprocess.Process | None = None
|
||||
|
||||
|
|
@ -29,7 +31,7 @@ class ValidatorProcess:
|
|||
if self.process is not None and self.process.returncode is None:
|
||||
return None
|
||||
try:
|
||||
async with asyncio.timeout(5):
|
||||
async with timeout(5):
|
||||
self.process = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
"-I",
|
||||
|
|
@ -43,10 +45,16 @@ class ValidatorProcess:
|
|||
if self.process.stdout is None or await self.process.stdout.readline() != b"ready\n":
|
||||
await self.close()
|
||||
return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, "validator startup failed")
|
||||
if sys.platform == "linux":
|
||||
import resource
|
||||
|
||||
resource.prlimit(
|
||||
self.process.pid, resource.RLIMIT_AS, (self._address_space_bytes, self._address_space_bytes)
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await self.close()
|
||||
raise
|
||||
except (OSError, ValueError, TimeoutError) as exc:
|
||||
except (OSError, ValueError, asyncio.TimeoutError) as exc:
|
||||
await self.close()
|
||||
return MontyFailure(MontyErrorKind.RUNTIME_UNAVAILABLE, escape_diagnostics(str(exc)))
|
||||
return None
|
||||
|
|
@ -66,7 +74,7 @@ class ValidatorProcess:
|
|||
self, value: JsonValue, schema: JsonSchema, *, deadline: float, output: bool = False
|
||||
) -> MontyFailure | None:
|
||||
try:
|
||||
async with asyncio.timeout(max(0, deadline - time.monotonic())):
|
||||
async with timeout(max(0, deadline - time.monotonic())):
|
||||
if (failure := await self.start()) is not None:
|
||||
return failure
|
||||
process: Final = self.process
|
||||
|
|
@ -89,7 +97,7 @@ class ValidatorProcess:
|
|||
except asyncio.CancelledError:
|
||||
await self.close()
|
||||
raise
|
||||
except TimeoutError:
|
||||
except asyncio.TimeoutError:
|
||||
await self.close()
|
||||
return MontyFailure(MontyErrorKind.TIMEOUT, "node deadline expired during schema validation")
|
||||
except (OSError, ValueError, RecursionError) as exc:
|
||||
|
|
|
|||
345
litellm/proxy/workflows/runner.py
Normal file
345
litellm/proxy/workflows/runner.py
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import Request
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.models.workflow_definition import (
|
||||
Binding,
|
||||
ChatCompletionNode,
|
||||
ValueBinding,
|
||||
WorkflowBundle,
|
||||
WorkflowDefinitionStore,
|
||||
WorkflowNode,
|
||||
)
|
||||
from litellm.models.workflow_execution import (
|
||||
ChildCallResult,
|
||||
EngineRun,
|
||||
ExecutionError,
|
||||
RunInput,
|
||||
RunState,
|
||||
StepTrace,
|
||||
WorkflowExecutionStore,
|
||||
workflow_call_id,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.workflows.monty.compiler import compile_program
|
||||
from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits
|
||||
from litellm.proxy.workflows.monty.runtime import MontyRuntime
|
||||
from litellm.proxy.workflows.monty.validation import ValidatorProcess
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
class ChatExecutor(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
request: Request,
|
||||
run: EngineRun,
|
||||
node: ChatCompletionNode,
|
||||
step: StepTrace,
|
||||
) -> ChildCallResult | ExecutionError: ...
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def resolve_binding(binding: Binding, run: EngineRun, steps: tuple[StepTrace, ...]) -> JsonValue | ExecutionError:
|
||||
if isinstance(binding, ValueBinding):
|
||||
return binding.value
|
||||
if binding.source == "$input":
|
||||
return [dict(message) for message in run.input.messages]
|
||||
source: Final = next((step for step in steps if step.node_id == binding.source), None)
|
||||
return source.output if source is not None else ExecutionError("missing_input", "Node output is unavailable")
|
||||
|
||||
|
||||
def _inputs(node: WorkflowNode, run: EngineRun, steps: tuple[StepTrace, ...]) -> JsonValue | ExecutionError:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
return resolve_binding(node.messages, run, steps)
|
||||
resolved: Final = {name: resolve_binding(binding, run, steps) for name, binding in node.inputs.items()}
|
||||
failure: Final = next((value for value in resolved.values() if isinstance(value, ExecutionError)), None)
|
||||
return failure if failure is not None else {k: v for k, v in resolved.items() if not isinstance(v, ExecutionError)}
|
||||
|
||||
|
||||
def _settled(task: asyncio.Task[None]) -> None:
|
||||
if not task.cancelled() and task.exception() is not None:
|
||||
verbose_proxy_logger.error("Workflow execution could not persist its final status")
|
||||
|
||||
|
||||
class WorkflowRunner:
|
||||
def __init__(
|
||||
self,
|
||||
runs: WorkflowExecutionStore,
|
||||
definitions: WorkflowDefinitionStore,
|
||||
execute_chat: ChatExecutor,
|
||||
*,
|
||||
clock: Callable[[], datetime] = utc_now,
|
||||
) -> None:
|
||||
self.runs: Final = runs
|
||||
self.definitions: Final = definitions
|
||||
self._execute_chat: Final = execute_chat
|
||||
self._clock: Final = clock
|
||||
self._slots: Final = asyncio.Semaphore(4)
|
||||
self._monty: Final = MontyRuntime(MontyLimits())
|
||||
self._monty_start: Final = asyncio.Lock()
|
||||
self._active: tuple[asyncio.Task[None], ...] = ()
|
||||
|
||||
async def submit(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> ExecutionError | None:
|
||||
self._active = tuple(task for task in self._active if not task.done())
|
||||
if len(self._active) >= 64:
|
||||
await self.runs.finish(run.run_id, RunState.failed, None, "Workflow capacity reached")
|
||||
return ExecutionError("capacity_exceeded", "Workflow capacity reached")
|
||||
task: Final = asyncio.create_task(self._run(run, request, caller))
|
||||
task.add_done_callback(_settled)
|
||||
self._active = (*self._active, task)
|
||||
return None
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._active:
|
||||
_, pending = await asyncio.wait(self._active, timeout=30)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*self._active, return_exceptions=True)
|
||||
await self._monty.close()
|
||||
|
||||
async def _watch(self, run_id: str) -> ExecutionError:
|
||||
while True:
|
||||
await asyncio.sleep(0.5)
|
||||
current: EngineRun | None = await self.runs.get(run_id)
|
||||
if current is None or current.status in {RunState.failed, RunState.cancelled}:
|
||||
return ExecutionError("interrupted", "Execution stopped or its deadline expired")
|
||||
if current.cancel_requested:
|
||||
return ExecutionError("cancelled", "Cancellation requested")
|
||||
|
||||
async def _run(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> None:
|
||||
execution: Final = asyncio.create_task(self._execute(run, request, caller))
|
||||
watch: Final = asyncio.create_task(self._watch(run.run_id))
|
||||
try:
|
||||
done, _ = await asyncio.wait((execution, watch), return_when=asyncio.FIRST_COMPLETED)
|
||||
result: Final = execution.result() if execution in done else watch.result()
|
||||
state: Final = (
|
||||
RunState.cancelled
|
||||
if isinstance(result, ExecutionError) and result.code == "cancelled"
|
||||
else RunState.failed
|
||||
if isinstance(result, ExecutionError)
|
||||
else RunState.completed
|
||||
)
|
||||
await self.runs.finish(
|
||||
run.run_id,
|
||||
state,
|
||||
None if isinstance(result, ExecutionError) else result,
|
||||
result.message if isinstance(result, ExecutionError) else None,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await self.runs.finish(run.run_id, RunState.failed, None, "Execution worker stopped")
|
||||
raise
|
||||
except Exception: # noqa: BLE001 # settle background failures without exposing inputs or provider credentials
|
||||
await self.runs.finish(run.run_id, RunState.failed, None, "Execution could not complete")
|
||||
verbose_proxy_logger.error("Workflow execution failed")
|
||||
finally:
|
||||
execution.cancel()
|
||||
watch.cancel()
|
||||
await asyncio.gather(execution, watch, return_exceptions=True)
|
||||
|
||||
async def _execute(self, run: EngineRun, request: Request, caller: UserAPIKeyAuth) -> JsonValue | ExecutionError:
|
||||
from litellm.proxy.common_request_processing import process_workflow_request
|
||||
|
||||
async with self._slots:
|
||||
snapshot: Final = await self.definitions.get_snapshot(run.snapshot_id, run.team_id)
|
||||
if snapshot is None or snapshot.workflow_id != run.workflow_id:
|
||||
return ExecutionError("snapshot_unavailable", "Pinned definition unavailable")
|
||||
|
||||
async def execute(messages: JsonValue) -> ModelResponse | ExecutionError:
|
||||
guarded: Final = run.model_copy(
|
||||
update={
|
||||
"input": RunInput.model_validate({**run.input.model_dump(), "messages": messages}),
|
||||
}
|
||||
)
|
||||
validator: Final = ValidatorProcess()
|
||||
try:
|
||||
invalid: Final = await validator.validate(
|
||||
messages,
|
||||
dict(snapshot.bundle.input_schema),
|
||||
deadline=time.monotonic()
|
||||
+ min(5, max(0, (run.input.deadline - self._clock()).total_seconds())),
|
||||
)
|
||||
finally:
|
||||
await validator.close()
|
||||
if invalid is not None:
|
||||
return ExecutionError("invalid_input", "Messages did not satisfy the workflow input schema")
|
||||
output: Final = await self._sequence(guarded, request, snapshot.bundle, ())
|
||||
if isinstance(output, ExecutionError):
|
||||
return output
|
||||
steps: Final = await self.runs.steps(run.run_id)
|
||||
return ModelResponse(
|
||||
id=f"chatcmpl-workflow-{run.run_id}",
|
||||
model=f"workflow/{snapshot.name}",
|
||||
created=int(run.created_at.timestamp()),
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": output
|
||||
if isinstance(output, str)
|
||||
else json.dumps(output, ensure_ascii=False, allow_nan=False),
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
usage={
|
||||
"prompt_tokens": sum(step.prompt_tokens for step in steps),
|
||||
"completion_tokens": sum(step.completion_tokens for step in steps),
|
||||
"total_tokens": sum(step.prompt_tokens + step.completion_tokens for step in steps),
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = await process_workflow_request(
|
||||
request=request,
|
||||
caller=caller,
|
||||
workflow_model=f"workflow/{snapshot.name}",
|
||||
messages=[dict(message) for message in run.input.messages],
|
||||
execute=execute,
|
||||
)
|
||||
return (
|
||||
response
|
||||
if isinstance(response, ExecutionError)
|
||||
else TypeAdapter[JsonValue](JsonValue).validate_json(response.model_dump_json())
|
||||
)
|
||||
|
||||
async def _sequence(
|
||||
self,
|
||||
run: EngineRun,
|
||||
request: Request,
|
||||
bundle: WorkflowBundle,
|
||||
steps: tuple[StepTrace, ...],
|
||||
) -> JsonValue | ExecutionError:
|
||||
if len(steps) == len(bundle.nodes):
|
||||
return resolve_binding(bundle.output, run, steps)
|
||||
node: Final = bundle.nodes[len(steps)]
|
||||
value: Final = _inputs(node, run, steps)
|
||||
if isinstance(value, ExecutionError):
|
||||
return value
|
||||
started: Final = self._clock()
|
||||
step: Final = StepTrace(
|
||||
step_id=str(uuid4()),
|
||||
run_id=run.run_id,
|
||||
node_id=node.id,
|
||||
state=RunState.running,
|
||||
input=value,
|
||||
attempts=1,
|
||||
started_at=started,
|
||||
deadline=min(run.input.deadline, started + timedelta(seconds=node.timeout_seconds)),
|
||||
)
|
||||
result: Final = await self._attempt(run, request, node, step)
|
||||
return (
|
||||
result
|
||||
if isinstance(result, ExecutionError)
|
||||
else await self._sequence(run, request, bundle, (*steps, result))
|
||||
)
|
||||
|
||||
async def _attempt(
|
||||
self,
|
||||
run: EngineRun,
|
||||
request: Request,
|
||||
node: WorkflowNode,
|
||||
step: StepTrace,
|
||||
) -> StepTrace | ExecutionError:
|
||||
attempt: Final = step.model_copy(
|
||||
update={
|
||||
"call_id": workflow_call_id(run.run_id, node.id, step.attempts)
|
||||
if isinstance(node, ChatCompletionNode)
|
||||
else None,
|
||||
}
|
||||
)
|
||||
if not await self.runs.checkpoint(run.run_id, attempt):
|
||||
return ExecutionError("cancelled", "Execution stopped before dispatch")
|
||||
try:
|
||||
result: Final = await self._safe_node(run, request, node, attempt)
|
||||
except asyncio.CancelledError:
|
||||
await self.runs.checkpoint(
|
||||
run.run_id,
|
||||
attempt.model_copy(
|
||||
update={
|
||||
"state": RunState.cancelled,
|
||||
"finished_at": self._clock(),
|
||||
"error": "Node execution interrupted",
|
||||
}
|
||||
),
|
||||
)
|
||||
raise
|
||||
receipt: Final = result.receipt if isinstance(result, ExecutionError) else result
|
||||
finished: Final = attempt.model_copy(
|
||||
update={
|
||||
"state": RunState.failed if isinstance(result, ExecutionError) else RunState.completed,
|
||||
"finished_at": self._clock(),
|
||||
"error": result.message if isinstance(result, ExecutionError) else None,
|
||||
"output": None if isinstance(result, ExecutionError) else result.output,
|
||||
"prompt_tokens": receipt.prompt_tokens if receipt else 0,
|
||||
"completion_tokens": receipt.completion_tokens if receipt else 0,
|
||||
"response_cost": receipt.response_cost if receipt else None,
|
||||
"request_id": receipt.request_id if receipt else None,
|
||||
}
|
||||
)
|
||||
await self.runs.checkpoint(run.run_id, finished)
|
||||
if not isinstance(result, ExecutionError):
|
||||
return finished
|
||||
retryable: Final = (
|
||||
result.code == "admission_rate_limit"
|
||||
if isinstance(node, ChatCompletionNode)
|
||||
else result.code in {"runtime_error", "worker_crashed"}
|
||||
)
|
||||
if (
|
||||
not retryable
|
||||
or result.receipt is not None
|
||||
or step.attempts >= node.retry.max_attempts
|
||||
or (step.deadline - self._clock()).total_seconds() <= node.retry.backoff_seconds
|
||||
):
|
||||
return result
|
||||
await asyncio.sleep(node.retry.backoff_seconds)
|
||||
return await self._attempt(run, request, node, step.model_copy(update={"attempts": step.attempts + 1}))
|
||||
|
||||
async def _safe_node(
|
||||
self,
|
||||
run: EngineRun,
|
||||
request: Request,
|
||||
node: WorkflowNode,
|
||||
step: StepTrace,
|
||||
) -> ChildCallResult | ExecutionError:
|
||||
try:
|
||||
return await self._node(run, request, node, step)
|
||||
except Exception: # noqa: BLE001 # persist a sanitized terminal step for unexpected runtime failures
|
||||
return ExecutionError("execution_failed", "The node could not complete")
|
||||
|
||||
async def _node(
|
||||
self,
|
||||
run: EngineRun,
|
||||
request: Request,
|
||||
node: WorkflowNode,
|
||||
step: StepTrace,
|
||||
) -> ChildCallResult | ExecutionError:
|
||||
if isinstance(node, ChatCompletionNode):
|
||||
return await self._execute_chat(request, run, node, step)
|
||||
compiled: Final = compile_program(node.program, self._monty.limits)
|
||||
if isinstance(compiled, MontyFailure):
|
||||
return ExecutionError("invalid_program", compiled.message)
|
||||
async with self._monty_start:
|
||||
availability: Final = await self._monty.start()
|
||||
if availability is not None:
|
||||
return ExecutionError("runtime_unavailable", "The Monty runtime is unavailable")
|
||||
result: Final = await self._monty.run_program(
|
||||
compiled,
|
||||
step.input,
|
||||
deadline_monotonic=time.monotonic() + max(0, (step.deadline - self._clock()).total_seconds()),
|
||||
)
|
||||
return (
|
||||
ExecutionError(result.kind.value.lower(), result.message)
|
||||
if isinstance(result, MontyFailure)
|
||||
else ChildCallResult(result.output, 0, 0, None, None)
|
||||
)
|
||||
213
litellm/repositories/workflow_definition_repository.py
Normal file
213
litellm/repositories/workflow_definition_repository.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
from collections.abc import Mapping
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from typing import Final, Protocol, TypeVar
|
||||
from uuid import uuid4
|
||||
|
||||
from liteflow import bundle_digest # pyright: ignore[reportMissingTypeStubs] # pinned SDK omits its py.typed marker
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.models.workflow_definition import (
|
||||
DefinitionError,
|
||||
DefinitionPage,
|
||||
WorkflowBundle,
|
||||
WorkflowDefinition,
|
||||
WorkflowDraft,
|
||||
WorkflowSnapshot,
|
||||
)
|
||||
from litellm.repositories.base_repository import DbRecord, record_to_dict
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
||||
_Model: Final = TypeVar("_Model", bound=BaseModel)
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_STRING: Final = TypeAdapter(str)
|
||||
|
||||
|
||||
class DefinitionTransaction(Protocol):
|
||||
@property
|
||||
def litellm_workflowdefinition(self) -> TableActions[DbRecord]: ...
|
||||
|
||||
@property
|
||||
def litellm_workflowsnapshot(self) -> TableActions[DbRecord]: ...
|
||||
|
||||
@property
|
||||
def litellm_teamtable(self) -> TableActions[DbRecord]: ...
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
|
||||
class DefinitionDatabase(DefinitionTransaction, Protocol):
|
||||
def tx(self) -> AbstractAsyncContextManager[DefinitionTransaction]: ...
|
||||
|
||||
|
||||
class _DefinitionRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
draft: WorkflowDraft
|
||||
contributors: tuple[str, ...]
|
||||
published_etag: str | None
|
||||
|
||||
|
||||
def _decoded(row: DbRecord, model: type[_Model]) -> _Model:
|
||||
values: Final = record_to_dict(row)
|
||||
bundle: Final = _JSON.validate_json(_STRING.validate_python(values["bundle_json"]))
|
||||
return model.model_validate({field: bundle if field == "bundle" else values[field] for field in model.model_fields})
|
||||
|
||||
|
||||
def _definition(row: DbRecord) -> _DefinitionRow:
|
||||
return _DefinitionRow.model_validate({**record_to_dict(row), "draft": _decoded(row, WorkflowDraft)})
|
||||
|
||||
|
||||
def _content(bundle: WorkflowBundle) -> Mapping[str, object]:
|
||||
serialized: Final = bundle.model_dump_json(by_alias=True, exclude_unset=True)
|
||||
return {
|
||||
"etag": f'"{uuid4()}"',
|
||||
"digest": bundle_digest(_JSON.validate_json(serialized)),
|
||||
"bundle_json": serialized,
|
||||
}
|
||||
|
||||
|
||||
async def _locked_definition(
|
||||
transaction: DefinitionTransaction, workflow_id: str, team_id: str, etag: str
|
||||
) -> _DefinitionRow | DefinitionError:
|
||||
await transaction.query_raw(
|
||||
'SELECT workflow_id FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND team_id = $2 FOR UPDATE',
|
||||
workflow_id,
|
||||
team_id,
|
||||
)
|
||||
stored: Final = await transaction.litellm_workflowdefinition.find_first(
|
||||
where={"workflow_id": workflow_id, "team_id": team_id}
|
||||
)
|
||||
if stored is None:
|
||||
return DefinitionError("not_found", "Workflow not found")
|
||||
row: Final = _definition(stored)
|
||||
if row.draft.etag != etag:
|
||||
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
|
||||
return row
|
||||
|
||||
|
||||
async def _snapshot(
|
||||
transaction: DefinitionTransaction, row: _DefinitionRow, published_by: str | None
|
||||
) -> WorkflowSnapshot:
|
||||
stored: Final = await transaction.litellm_workflowsnapshot.create(
|
||||
data={
|
||||
"snapshot_id": str(uuid4()),
|
||||
"workflow_id": row.draft.workflow_id,
|
||||
"team_id": row.draft.team_id,
|
||||
"name": row.draft.name,
|
||||
"digest": row.draft.digest,
|
||||
"bundle_json": _JSON.dump_json(dict(row.draft.bundle)).decode(),
|
||||
"contributors": list(row.contributors),
|
||||
"published_by": published_by,
|
||||
}
|
||||
)
|
||||
return _decoded(stored, WorkflowSnapshot)
|
||||
|
||||
|
||||
class WorkflowDefinitionRepository:
|
||||
def __init__(self, database: DefinitionDatabase) -> None:
|
||||
self._database: Final = database
|
||||
|
||||
async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None:
|
||||
row: Final = await self._database.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
return None if row is None else LiteLLM_TeamTable.model_validate(record_to_dict(row))
|
||||
|
||||
async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None:
|
||||
stored: Final = await self._database.litellm_workflowdefinition.find_first(
|
||||
where={"workflow_id": workflow_id, **({"team_id": team_id} if team_id is not None else {})}
|
||||
)
|
||||
if stored is None:
|
||||
return None
|
||||
row: Final = _definition(stored)
|
||||
return WorkflowDefinition(draft=row.draft, contributors=row.contributors)
|
||||
|
||||
async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage:
|
||||
rows: Final = await self._database.litellm_workflowdefinition.find_many(
|
||||
where={"team_id": team_id, **({"workflow_id": {"gt": cursor}} if cursor is not None else {})},
|
||||
order={"workflow_id": "asc"},
|
||||
take=limit + 1,
|
||||
)
|
||||
drafts: Final = tuple(_decoded(row, WorkflowDraft) for row in rows[:limit])
|
||||
return DefinitionPage(workflows=drafts, next_cursor=drafts[-1].workflow_id if len(rows) > limit else None)
|
||||
|
||||
async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None:
|
||||
row: Final = await self._database.litellm_workflowsnapshot.find_first(
|
||||
where={"snapshot_id": snapshot_id, **({"team_id": team_id} if team_id is not None else {})}
|
||||
)
|
||||
return None if row is None else _decoded(row, WorkflowSnapshot)
|
||||
|
||||
async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None:
|
||||
row: Final = await self._database.litellm_workflowdefinition.find_first(
|
||||
where={"name": name, "team_id": team_id}
|
||||
)
|
||||
if row is None or (snapshot_id := _decoded(row, WorkflowDraft).published_snapshot_id) is None:
|
||||
return None
|
||||
return await self.get_snapshot(snapshot_id, team_id)
|
||||
|
||||
async def create(self, bundle: WorkflowBundle, team_id: str, contributor: str) -> WorkflowDraft | DefinitionError:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
try:
|
||||
stored: Final = await self._database.litellm_workflowdefinition.create(
|
||||
data={
|
||||
"workflow_id": str(uuid4()),
|
||||
"name": bundle.name,
|
||||
"team_id": team_id,
|
||||
"contributors": [contributor],
|
||||
**_content(bundle),
|
||||
}
|
||||
)
|
||||
except UniqueViolationError:
|
||||
return DefinitionError("conflict", "A workflow with this name already exists in the team", "bundle.name")
|
||||
return _decoded(stored, WorkflowDraft)
|
||||
|
||||
async def replace(
|
||||
self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str
|
||||
) -> WorkflowDraft | DefinitionError:
|
||||
async with self._database.tx() as transaction:
|
||||
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
|
||||
if isinstance(row, DefinitionError):
|
||||
return row
|
||||
if row.draft.name != bundle.name:
|
||||
return DefinitionError("conflict", "A workflow name cannot change", "bundle.name")
|
||||
contributors: Final = (
|
||||
(contributor,)
|
||||
if etag == row.published_etag
|
||||
else tuple(sorted(frozenset(row.contributors) | {contributor}))
|
||||
)
|
||||
stored: Final = await transaction.litellm_workflowdefinition.update(
|
||||
where={"workflow_id": workflow_id}, data={**_content(bundle), "contributors": list(contributors)}
|
||||
)
|
||||
assert stored is not None
|
||||
return _decoded(stored, WorkflowDraft)
|
||||
|
||||
async def publish(
|
||||
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
|
||||
) -> WorkflowSnapshot | DefinitionError:
|
||||
async with self._database.tx() as transaction:
|
||||
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
|
||||
if isinstance(row, DefinitionError):
|
||||
return row
|
||||
if row.draft.digest != digest:
|
||||
return DefinitionError("stale_draft", "The reviewed digest does not match the draft", "digest")
|
||||
if publisher in row.contributors:
|
||||
return DefinitionError("forbidden", "Draft contributors cannot publish their own changes")
|
||||
if row.published_etag == etag and row.draft.published_snapshot_id is not None:
|
||||
existing: Final = await transaction.litellm_workflowsnapshot.find_unique(
|
||||
where={"snapshot_id": row.draft.published_snapshot_id}
|
||||
)
|
||||
assert existing is not None
|
||||
return _decoded(existing, WorkflowSnapshot)
|
||||
snapshot: Final = await _snapshot(transaction, row, publisher)
|
||||
await transaction.litellm_workflowdefinition.update(
|
||||
where={"workflow_id": workflow_id},
|
||||
data={"published_snapshot_id": snapshot.snapshot_id, "published_etag": etag},
|
||||
)
|
||||
return snapshot
|
||||
|
||||
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError:
|
||||
async with self._database.tx() as transaction:
|
||||
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
|
||||
if isinstance(row, DefinitionError):
|
||||
return row
|
||||
return await _snapshot(transaction, row, None)
|
||||
297
litellm/repositories/workflow_execution_repository.py
Normal file
297
litellm/repositories/workflow_execution_repository.py
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
import base64
|
||||
import json
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import Final, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.models.workflow_execution import (
|
||||
EngineRun,
|
||||
ExecutionError,
|
||||
RunAdmission,
|
||||
RunInput,
|
||||
RunPage,
|
||||
RunState,
|
||||
RunView,
|
||||
StepTrace,
|
||||
)
|
||||
|
||||
|
||||
class ExecutionQueries(Protocol):
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
|
||||
class ExecutionDatabase(ExecutionQueries, Protocol):
|
||||
def tx(self) -> AbstractAsyncContextManager[ExecutionQueries]: ...
|
||||
|
||||
|
||||
class _Cursor(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
created_at: datetime
|
||||
run_id: str
|
||||
|
||||
|
||||
class _TraceRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
trace: StepTrace
|
||||
|
||||
|
||||
_RUNS: Final = TypeAdapter(tuple[EngineRun, ...])
|
||||
_VIEWS: Final = TypeAdapter(tuple[RunView, ...])
|
||||
_TRACES: Final = TypeAdapter(tuple[_TraceRow, ...])
|
||||
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
|
||||
|
||||
|
||||
def _cursor(value: str | None) -> _Cursor | ExecutionError | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return _Cursor.model_validate_json(base64.urlsafe_b64decode(value.encode("ascii")))
|
||||
except (ValueError, UnicodeError, ValidationError):
|
||||
return ExecutionError("invalid_cursor", "Invalid run cursor")
|
||||
|
||||
|
||||
def _encoded_cursor(run: RunView) -> str:
|
||||
return base64.urlsafe_b64encode(
|
||||
_Cursor(created_at=run.created_at, run_id=run.run_id).model_dump_json().encode()
|
||||
).decode("ascii")
|
||||
|
||||
|
||||
async def _first_run(database: ExecutionQueries, query: str, *args: object) -> EngineRun | None:
|
||||
rows: Final = _RUNS.validate_python(await database.query_raw(query, *args))
|
||||
return rows[0] if rows else None
|
||||
|
||||
|
||||
def _aggregate_attempts(attempts: tuple[StepTrace, ...]) -> StepTrace:
|
||||
latest: Final = max(attempts, key=lambda step: step.attempts)
|
||||
costs: Final = tuple(step.response_cost for step in attempts if step.response_cost is not None)
|
||||
ids: Final = chain.from_iterable(
|
||||
step.request_ids or ((step.request_id,) if step.request_id is not None else ()) for step in attempts
|
||||
)
|
||||
return latest.model_copy(
|
||||
update={
|
||||
"prompt_tokens": sum(step.prompt_tokens for step in attempts),
|
||||
"completion_tokens": sum(step.completion_tokens for step in attempts),
|
||||
"response_cost": sum(costs) if costs else None,
|
||||
"request_ids": tuple(dict.fromkeys(ids)),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class WorkflowExecutionRepository:
|
||||
def __init__(self, database: ExecutionDatabase) -> None:
|
||||
self._database: Final = database
|
||||
|
||||
async def find_idempotent(
|
||||
self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str
|
||||
) -> EngineRun | None:
|
||||
return await _first_run(
|
||||
self._database,
|
||||
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 '
|
||||
"AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text",
|
||||
key_hash,
|
||||
team_id,
|
||||
workflow_id,
|
||||
idempotency_key,
|
||||
)
|
||||
|
||||
async def admit(
|
||||
self,
|
||||
*,
|
||||
workflow_id: str,
|
||||
snapshot_id: str,
|
||||
team_id: str,
|
||||
key_hash: str,
|
||||
input: RunInput,
|
||||
digest: str,
|
||||
idempotency_key: str | None,
|
||||
) -> RunAdmission | ExecutionError:
|
||||
await self._expire_team(team_id)
|
||||
async with self._database.tx() as transaction:
|
||||
await transaction.query_raw(
|
||||
"SELECT true AS acquired FROM pg_advisory_xact_lock(hashtextextended($1, 0))", team_id
|
||||
)
|
||||
existing: Final = await _first_run(
|
||||
transaction,
|
||||
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 '
|
||||
"AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text",
|
||||
key_hash,
|
||||
team_id,
|
||||
workflow_id,
|
||||
idempotency_key,
|
||||
)
|
||||
if existing is not None:
|
||||
return (
|
||||
RunAdmission(existing, False)
|
||||
if existing.request_digest == digest
|
||||
else ExecutionError("idempotency_conflict", "Idempotency key was used with different input")
|
||||
)
|
||||
capacity: Final = _ROWS.validate_python(
|
||||
await transaction.query_raw(
|
||||
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND team_id = $1 '
|
||||
"AND status IN ('pending', 'running') LIMIT 64",
|
||||
team_id,
|
||||
)
|
||||
)
|
||||
if len(capacity) >= 64:
|
||||
return ExecutionError("capacity_exceeded", "The team's workflow execution queue is full")
|
||||
created: Final = await _first_run(
|
||||
transaction,
|
||||
'INSERT INTO "LiteLLM_WorkflowRun" '
|
||||
"(run_id, session_id, workflow_type, status, created_by, created_at, updated_at, "
|
||||
"workflow_id, snapshot_id, team_id, engine_owned, input, request_digest, idempotency_key) "
|
||||
"VALUES ($1, $2, 'liteflow', 'pending', $3, timezone('UTC', NOW()), timezone('UTC', NOW()), $4, $5, $6, true, $7::jsonb, $8, $9) "
|
||||
"RETURNING *",
|
||||
str(uuid4()),
|
||||
str(uuid4()),
|
||||
key_hash,
|
||||
workflow_id,
|
||||
snapshot_id,
|
||||
team_id,
|
||||
input.model_dump_json(),
|
||||
digest,
|
||||
idempotency_key,
|
||||
)
|
||||
if created is None:
|
||||
return ExecutionError("storage_error", "Workflow admission did not create a run")
|
||||
return RunAdmission(created, True)
|
||||
|
||||
async def _expire_team(self, team_id: str) -> None:
|
||||
await self._database.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowRun\" SET status = CASE WHEN cancel_requested THEN 'cancelled' ELSE 'failed' END, "
|
||||
"metadata = jsonb_build_object('error', 'Execution deadline expired'), updated_at = timezone('UTC', NOW()) "
|
||||
"WHERE team_id = $1 AND engine_owned = true AND status IN ('pending', 'running') "
|
||||
"AND (input->>'deadline')::timestamptz <= NOW() RETURNING run_id",
|
||||
team_id,
|
||||
)
|
||||
|
||||
async def get(self, run_id: str) -> EngineRun | None:
|
||||
query: Final = 'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true'
|
||||
run: Final = await _first_run(self._database, query, run_id)
|
||||
if (
|
||||
run is not None
|
||||
and run.status in {RunState.pending, RunState.running}
|
||||
and run.input.deadline <= datetime.now(timezone.utc)
|
||||
):
|
||||
await self._expire_team(run.team_id)
|
||||
return await _first_run(self._database, query, run_id)
|
||||
return run
|
||||
|
||||
async def checkpoint(self, run_id: str, step: StepTrace) -> bool:
|
||||
async with self._database.tx() as transaction:
|
||||
owned: Final = _ROWS.validate_python(
|
||||
await transaction.query_raw(
|
||||
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true '
|
||||
"AND ($2::boolean = false OR (status IN ('pending', 'running') AND NOT cancel_requested "
|
||||
"AND (input->>'deadline')::timestamptz > NOW())) FOR UPDATE",
|
||||
run_id,
|
||||
step.state is RunState.running,
|
||||
)
|
||||
)
|
||||
if not owned:
|
||||
return False
|
||||
await transaction.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowRun\" SET status = 'running', updated_at = timezone('UTC', NOW()) "
|
||||
"WHERE run_id = $1 AND status = 'pending' RETURNING run_id",
|
||||
run_id,
|
||||
)
|
||||
await transaction.query_raw(
|
||||
'INSERT INTO "LiteLLM_WorkflowEvent" '
|
||||
"(event_id, run_id, event_type, step_name, sequence_number, data, created_at) "
|
||||
"SELECT $1, $2, 'liteflow.step', $3, COALESCE(MAX(sequence_number) + 1, 0), $4::jsonb, timezone('UTC', NOW()) "
|
||||
'FROM "LiteLLM_WorkflowEvent" WHERE run_id = $2 RETURNING event_id',
|
||||
str(uuid4()),
|
||||
run_id,
|
||||
step.node_id,
|
||||
json.dumps({"trace": step.model_dump(mode="json")}, allow_nan=False),
|
||||
)
|
||||
return True
|
||||
|
||||
async def steps(self, run_id: str) -> tuple[StepTrace, ...]:
|
||||
rows: Final = _TRACES.validate_python(
|
||||
await self._database.query_raw(
|
||||
"SELECT trace FROM ("
|
||||
"SELECT CASE WHEN receipts.ids IS NULL THEN latest.trace ELSE latest.trace || "
|
||||
"jsonb_build_object('request_id', CASE WHEN cardinality(receipts.ids) = 1 THEN receipts.ids[1] END, "
|
||||
"'request_ids', to_jsonb(receipts.ids), 'response_cost', receipts.cost, "
|
||||
"'prompt_tokens', receipts.prompt_tokens, 'completion_tokens', receipts.completion_tokens) END AS trace, "
|
||||
"latest.step_name FROM ("
|
||||
"SELECT DISTINCT ON (step_name, data->'trace'->>'attempts') data->'trace' AS trace, step_name "
|
||||
"FROM \"LiteLLM_WorkflowEvent\" WHERE run_id = $1 AND event_type = 'liteflow.step' "
|
||||
"ORDER BY step_name, data->'trace'->>'attempts', sequence_number DESC"
|
||||
') latest JOIN "LiteLLM_WorkflowRun" run ON run.run_id = $1 AND run.engine_owned = true '
|
||||
"LEFT JOIN LATERAL ("
|
||||
"SELECT array_agg(log.request_id ORDER BY log.request_id) AS ids, SUM(log.spend) AS cost, "
|
||||
"SUM(log.prompt_tokens) AS prompt_tokens, SUM(log.completion_tokens) AS completion_tokens "
|
||||
"FROM \"LiteLLM_SpendLogs\" log WHERE log.litellm_call_id = latest.trace->>'call_id' "
|
||||
"AND log.api_key = run.created_by AND log.team_id = run.team_id"
|
||||
") receipts ON true"
|
||||
") latest ORDER BY (trace->>'started_at')::timestamptz, step_name",
|
||||
run_id,
|
||||
)
|
||||
)
|
||||
node_ids: Final = dict.fromkeys(row.trace.node_id for row in rows)
|
||||
return tuple(
|
||||
_aggregate_attempts(tuple(row.trace for row in rows if row.trace.node_id == node_id))
|
||||
for node_id in node_ids
|
||||
)
|
||||
|
||||
async def finish(self, run_id: str, state: RunState, output: JsonValue, error: str | None) -> None:
|
||||
await self._database.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowRun\" SET status = CASE WHEN cancel_requested THEN 'cancelled' ELSE $2 END, "
|
||||
"output = CASE WHEN cancel_requested THEN NULL ELSE $3::jsonb END, "
|
||||
"metadata = jsonb_build_object('error', $4::text), updated_at = timezone('UTC', NOW()) "
|
||||
"WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id",
|
||||
run_id,
|
||||
state.value,
|
||||
json.dumps(output, allow_nan=False),
|
||||
error,
|
||||
)
|
||||
|
||||
async def cancel(self, run_id: str) -> None:
|
||||
await self._database.query_raw(
|
||||
"UPDATE \"LiteLLM_WorkflowRun\" SET cancel_requested = true, updated_at = timezone('UTC', NOW()) "
|
||||
"WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id",
|
||||
run_id,
|
||||
)
|
||||
|
||||
async def list_runs(
|
||||
self,
|
||||
workflow_id: str,
|
||||
team_id: str,
|
||||
key_hash: str | None,
|
||||
state: RunState | None,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> RunPage | ExecutionError:
|
||||
await self._expire_team(team_id)
|
||||
position: Final = _cursor(cursor)
|
||||
if isinstance(position, ExecutionError):
|
||||
return position
|
||||
rows: Final = _VIEWS.validate_python(
|
||||
await self._database.query_raw(
|
||||
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true '
|
||||
"AND workflow_id = $1 AND team_id = $2 "
|
||||
"AND ($3::text IS NULL OR created_by = $3::text) AND ($4::text IS NULL OR status = $4::text) "
|
||||
"AND ($5::text IS NULL OR (created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)) "
|
||||
"ORDER BY created_at DESC, run_id DESC LIMIT $7",
|
||||
workflow_id,
|
||||
team_id,
|
||||
key_hash,
|
||||
state.value if state is not None else None,
|
||||
position.created_at.isoformat() if position is not None else None,
|
||||
position.run_id if position is not None else None,
|
||||
limit + 1,
|
||||
)
|
||||
)
|
||||
page: Final = rows[:limit]
|
||||
return RunPage(
|
||||
runs=page,
|
||||
count=len(page),
|
||||
next_cursor=_encoded_cursor(page[-1]) if len(rows) > limit else None,
|
||||
)
|
||||
|
|
@ -55,6 +55,7 @@ Documentation = "https://docs.litellm.ai"
|
|||
# supply-chain-pinned to a vetted, aged release.
|
||||
[project.optional-dependencies]
|
||||
proxy = [
|
||||
"liteflow @ https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz#sha256=c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072",
|
||||
"pydantic-monty==1.4.0",
|
||||
"gunicorn>=23.0.0,<24.0",
|
||||
"uvicorn>=0.33.0,<1.0",
|
||||
|
|
@ -258,6 +259,7 @@ admin-mcp = [
|
|||
"litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/31fb096904ae05684c0779110a93b135b5601182.tar.gz#sha256=5803a20a00a18177f4591706956f31be311879543049b8b38143ad90bb1cb0a6 ; python_version >= '3.12'",
|
||||
]
|
||||
proxy-dev = [
|
||||
"liteflow @ https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz#sha256=c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072",
|
||||
"pydantic-monty==1.4.0",
|
||||
{ include-group = "admin-mcp" },
|
||||
"prisma==0.11.0",
|
||||
|
|
|
|||
|
|
@ -1975,10 +1975,51 @@ model LiteLLM_ShadowEvalFunnel {
|
|||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowDefinition {
|
||||
workflow_id String @id @default(uuid())
|
||||
team_id String
|
||||
name String
|
||||
etag String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_snapshot_id String?
|
||||
published_etag String?
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
snapshots LiteLLM_WorkflowSnapshot[]
|
||||
|
||||
@@unique([team_id, name])
|
||||
@@index([team_id, workflow_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowSnapshot {
|
||||
snapshot_id String @id @default(uuid())
|
||||
workflow_id String
|
||||
team_id String
|
||||
name String
|
||||
digest String
|
||||
bundle_json String
|
||||
contributors String[]
|
||||
published_by String?
|
||||
created_at DateTime @default(now())
|
||||
definition LiteLLM_WorkflowDefinition @relation(fields: [workflow_id], references: [workflow_id], onDelete: Restrict)
|
||||
|
||||
@@index([workflow_id, created_at, snapshot_id])
|
||||
@@index([team_id, snapshot_id])
|
||||
}
|
||||
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
workflow_id String?
|
||||
snapshot_id String?
|
||||
team_id String?
|
||||
engine_owned Boolean @default(false)
|
||||
request_digest String?
|
||||
idempotency_key String?
|
||||
cancel_requested Boolean @default(false)
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
|
|
@ -1994,6 +2035,9 @@ model LiteLLM_WorkflowRun {
|
|||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
@@unique([created_by, team_id, workflow_id, idempotency_key], map: "LiteLLM_WorkflowRun_principal_idempotency_key")
|
||||
@@index([workflow_id, team_id, created_at, run_id])
|
||||
@@index([team_id, engine_owned, status])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
|
|
|
|||
|
|
@ -7,18 +7,22 @@ count them per rule the same delta-vs-base way as the other lint gates. The rule
|
|||
OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI
|
||||
what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix)
|
||||
OA002 no description (the handler has no docstring and the route no ``description=``)
|
||||
OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``)
|
||||
OA003 no 2xx success response has a schema (use ``response_model`` or a schema in ``responses``)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Final, NamedTuple
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
GATED_METHODS: Final = frozenset({"post"})
|
||||
SCHEMA_MEDIA_TYPES: Final = ("application/json", "application/sdp", "multipart/form-data")
|
||||
SUCCESS_STATUS_CODES: Final = frozenset(str(code) for code in range(200, 300))
|
||||
_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class Violation(NamedTuple):
|
||||
|
|
@ -28,7 +32,7 @@ class Violation(NamedTuple):
|
|||
|
||||
|
||||
def _as_mapping(node: object) -> Mapping[str, object]:
|
||||
return node if isinstance(node, Mapping) else {}
|
||||
return _OBJECT.validate_python(node) if isinstance(node, Mapping) else {}
|
||||
|
||||
|
||||
def _content_schema(container: object) -> object:
|
||||
|
|
@ -45,7 +49,11 @@ def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]:
|
|||
description: Final = operation.get("description")
|
||||
if not isinstance(description, str) or not description.strip():
|
||||
yield "OA002"
|
||||
if not _content_schema(_as_mapping(operation.get("responses")).get("200")):
|
||||
if not any(
|
||||
_content_schema(response)
|
||||
for status, response in _as_mapping(operation.get("responses")).items()
|
||||
if status in SUCCESS_STATUS_CODES
|
||||
):
|
||||
yield "OA003"
|
||||
|
||||
|
||||
|
|
@ -66,10 +74,10 @@ def proxy_spec() -> Mapping[str, object]:
|
|||
return _as_mapping(app.openapi())
|
||||
|
||||
|
||||
def main(argv: list[str] = sys.argv[1:]) -> int:
|
||||
def main(argv: Sequence[str] = sys.argv[1:]) -> int:
|
||||
spec: Final = json.loads(open(argv[0]).read()) if argv else proxy_spec()
|
||||
for violation in violations(spec):
|
||||
print(f"{violation.method} {violation.path} {violation.code}")
|
||||
sys.stdout.write(f"{violation.method} {violation.path} {violation.code}\n")
|
||||
return 0
|
||||
|
||||
|
||||
|
|
|
|||
427
tests/integration/database/test_workflow_execution.py
Normal file
427
tests/integration/database/test_workflow_execution.py
Normal file
|
|
@ -0,0 +1,427 @@
|
|||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Final, cast
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from fastapi import HTTPException, Request
|
||||
from openai.types.chat import ChatCompletion
|
||||
from prisma import Json, Prisma
|
||||
from prisma.types import DatasourceOverride
|
||||
from psycopg import sql
|
||||
|
||||
from litellm.models.workflow_definition import ChatCompletionNode, WorkflowDraft, WorkflowSnapshot
|
||||
from litellm.models.workflow_execution import (
|
||||
ChildCallResult,
|
||||
EngineRun,
|
||||
ExecutionError,
|
||||
RunAdmission,
|
||||
RunInput,
|
||||
RunPage,
|
||||
RunState,
|
||||
StepTrace,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
|
||||
TestRunRequest as PaidRunRequest,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
|
||||
cancel_workflow_run,
|
||||
get_run_definition,
|
||||
get_workflow_steps,
|
||||
list_definition_runs,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
|
||||
test_workflow as run_draft_test,
|
||||
)
|
||||
from litellm.proxy.workflows.runner import WorkflowRunner
|
||||
from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository
|
||||
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
|
||||
from tests.integration._support.database import scratch_database
|
||||
from tests.unit.proxy.workflows.test_runner import two_nodes
|
||||
|
||||
pytestmark: Final = pytest.mark.asyncio(loop_scope="function")
|
||||
_ROOT: Final = Path(__file__).resolve().parents[3]
|
||||
_MESSAGES: Final = ({"role": "user", "content": "plan"},)
|
||||
_CALLER: Final = UserAPIKeyAuth(token="owner", team_id="team").model_copy(update={"via_virtual_key": True})
|
||||
_REQUEST: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Storage:
|
||||
db: Prisma
|
||||
runs: WorkflowExecutionRepository
|
||||
definitions: WorkflowDefinitionRepository
|
||||
|
||||
async def admit(self, *, key: str | None = None, digest: str = "input", owner: str = "owner",
|
||||
team: str = "team", snapshot: WorkflowSnapshot | None = None,
|
||||
deadline: datetime | None = None) -> RunAdmission | ExecutionError:
|
||||
return await self.runs.admit(
|
||||
workflow_id=snapshot.workflow_id if snapshot else "workflow",
|
||||
snapshot_id=snapshot.snapshot_id if snapshot else "snapshot",
|
||||
team_id=team, key_hash=owner, input=RunInput(messages=_MESSAGES, draft=False,
|
||||
deadline=deadline or datetime.now(timezone.utc) + timedelta(minutes=5)), digest=digest, idempotency_key=key,
|
||||
)
|
||||
|
||||
async def snapshot(self, team: str = "team") -> WorkflowSnapshot:
|
||||
draft: Final = await self.definitions.create(two_nodes(), team, "author")
|
||||
assert isinstance(draft, WorkflowDraft)
|
||||
snapshot: Final = await self.definitions.publish(draft.workflow_id, team, draft.etag, draft.digest, "admin")
|
||||
assert isinstance(snapshot, WorkflowSnapshot)
|
||||
return snapshot
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def schema_sql() -> bytes:
|
||||
return subprocess.run(
|
||||
[sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel",
|
||||
str(_ROOT / "schema.prisma"), "--script"], check=True, capture_output=True, cwd=_ROOT,
|
||||
).stdout
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def store(schema_sql: bytes) -> AsyncIterator[Storage]:
|
||||
with scratch_database("workflow_execution") as url, psycopg.connect(url, autocommit=True) as connection:
|
||||
connection.execute(schema_sql)
|
||||
connection.execute(sql.SQL("ALTER DATABASE {} SET timezone TO 'Asia/Singapore'").format(
|
||||
sql.Identifier(connection.info.dbname)
|
||||
))
|
||||
async with Prisma(datasource=DatasourceOverride(url=url)) as client:
|
||||
assert await client.query_raw("SHOW TIME ZONE") == [{"TimeZone": "Asia/Singapore"}]
|
||||
for team in ("team", "other"):
|
||||
await client.litellm_teamtable.create(data={"team_id": team, "members_with_roles": Json([
|
||||
{"user_id": "admin", "role": "admin"}, {"user_id": "member", "role": "user"},
|
||||
])})
|
||||
yield Storage(client, WorkflowExecutionRepository(cast(ExecutionDatabase, client)),
|
||||
WorkflowDefinitionRepository(cast(DefinitionDatabase, client)))
|
||||
|
||||
|
||||
class ChatReplies:
|
||||
def __init__(self, store: Storage, replies: tuple[ChildCallResult | ExecutionError | Exception, ...],
|
||||
cancel_after: int | None = None) -> None:
|
||||
self.store, self.replies, self.cancel_after = store, replies, cancel_after
|
||||
self.calls: tuple[StepTrace, ...] = ()
|
||||
self.models: tuple[str, ...] = ()
|
||||
self.persisted: tuple[StepTrace, ...] = ()
|
||||
|
||||
async def __call__(self, request: Request, run: EngineRun, node: ChatCompletionNode,
|
||||
step: StepTrace) -> ChildCallResult | ExecutionError:
|
||||
self.calls = (*self.calls, step)
|
||||
self.models = (*self.models, node.model)
|
||||
self.persisted = (*self.persisted, (await self.store.runs.steps(run.run_id))[-1])
|
||||
if len(self.calls) == self.cancel_after:
|
||||
await self.store.runs.cancel(run.run_id)
|
||||
reply: Final = self.replies[len(self.calls) - 1]
|
||||
if isinstance(reply, Exception):
|
||||
raise reply
|
||||
return reply
|
||||
|
||||
|
||||
async def test_concurrent_idempotency_pins_first_snapshot_and_conflicts_on_input(store: Storage) -> None:
|
||||
original: Final = await store.snapshot()
|
||||
different: Final = original.model_copy(update={"snapshot_id": "later-snapshot"})
|
||||
results: Final = await asyncio.gather(*(
|
||||
store.admit(key="same", snapshot=snapshot) for snapshot in (original, different) * 4
|
||||
))
|
||||
assert all(isinstance(result, RunAdmission) for result in results)
|
||||
admitted: Final = tuple(result for result in results if isinstance(result, RunAdmission))
|
||||
assert sum(result.created for result in admitted) == 1
|
||||
assert len({result.run.run_id for result in admitted}) == 1
|
||||
assert len({result.run.snapshot_id for result in admitted}) == 1
|
||||
assert all(result.run == admitted[0].run for result in admitted)
|
||||
assert await store.admit(key="same", snapshot=different, digest="changed") == ExecutionError(
|
||||
"idempotency_conflict", "Idempotency key was used with different input"
|
||||
)
|
||||
assert isinstance(await store.admit(key="same", snapshot=original, owner="another"), RunAdmission)
|
||||
assert isinstance(await store.admit(key="same", snapshot=original, team="other"), RunAdmission)
|
||||
assert await store.db.litellm_workflowrun.count() == 3
|
||||
|
||||
|
||||
async def test_draft_test_retry_after_edit_recovers_only_the_same_scoped_request(store: Storage) -> None:
|
||||
draft: Final = await store.definitions.create(two_nodes(), "team", "member")
|
||||
assert isinstance(draft, WorkflowDraft)
|
||||
caller: Final = _CALLER.model_copy(update={"user_id": "member", "user_role": LitellmUserRoles.INTERNAL_USER})
|
||||
data: Final = PaidRunRequest(messages=_MESSAGES, authorize_billing=True)
|
||||
chat: Final = ChatReplies(store, (
|
||||
ChildCallResult(list(_MESSAGES), 1, 1, "first", 0.125), ChildCallResult("done", 1, 1, "last", 0.125),
|
||||
))
|
||||
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
|
||||
request: Final = Request({**_REQUEST.scope, "state": {"workflow_runner": runner}})
|
||||
try:
|
||||
accepted: Final = await run_draft_test(
|
||||
request, draft.workflow_id, data, caller, store.definitions, store.runs,
|
||||
if_match=draft.etag, idempotency_key="retry",
|
||||
)
|
||||
finally:
|
||||
await runner.close()
|
||||
changed: Final = await store.definitions.replace(
|
||||
draft.workflow_id, "team", draft.etag, two_nodes(), "member",
|
||||
)
|
||||
assert isinstance(changed, WorkflowDraft) and changed.etag != draft.etag
|
||||
replay: Final = await run_draft_test(
|
||||
_REQUEST, draft.workflow_id, data, caller, store.definitions, store.runs,
|
||||
if_match=draft.etag, idempotency_key="retry",
|
||||
)
|
||||
assert replay.status_code == accepted.status_code == 202
|
||||
assert replay.body == accepted.body and replay.headers == accepted.headers
|
||||
for body, identity, etag, key, status in (
|
||||
(data.model_copy(update={"messages": ({"role": "user", "content": "changed"},)}), caller, draft.etag, "retry", 409),
|
||||
(data, caller, changed.etag, "retry", 409),
|
||||
(data, caller, None, "retry", 412),
|
||||
(data, caller, draft.etag, "new-key", 412),
|
||||
(data, caller.model_copy(update={"token": "peer"}), draft.etag, "retry", 412),
|
||||
(data, caller.model_copy(update={"team_id": "other"}), draft.etag, "retry", 404),
|
||||
(data, caller.model_copy(update={"user_id": "removed-member"}), draft.etag, "retry", 404),
|
||||
(data, caller.model_copy(update={"via_virtual_key": False}), draft.etag, "retry", 403),
|
||||
(data.model_copy(update={"authorize_billing": False}), caller, draft.etag, "retry", 400),
|
||||
):
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
await run_draft_test(
|
||||
_REQUEST, draft.workflow_id, body, identity, store.definitions, store.runs,
|
||||
if_match=etag, idempotency_key=key,
|
||||
)
|
||||
assert rejected.value.status_code == status
|
||||
assert await store.db.litellm_workflowrun.count() == await store.db.litellm_workflowsnapshot.count() == 1
|
||||
assert len(chat.calls) == 2
|
||||
|
||||
|
||||
async def test_nullable_history_filters_cursor_ties_and_utc_in_non_utc_database(store: Storage) -> None:
|
||||
before: Final = datetime.now(timezone.utc).replace(microsecond=0)
|
||||
admissions: Final = tuple([await store.admit(owner=owner) for owner in ("owner", "foreign", "owner")])
|
||||
assert all(isinstance(item, RunAdmission) for item in admissions)
|
||||
runs: Final = tuple(item.run for item in admissions if isinstance(item, RunAdmission))
|
||||
await store.db.litellm_workflowrun.update_many(where={}, data={"created_at": runs[0].created_at})
|
||||
ordered: Final = tuple(sorted((run.run_id for run in runs), reverse=True))
|
||||
for owner, state in ((None, None), ("owner", None), (None, RunState.pending), ("owner", RunState.pending)):
|
||||
page: Final = await store.runs.list_runs("workflow", "team", owner, state, 1, None)
|
||||
assert isinstance(page, RunPage) and page.next_cursor is not None
|
||||
tail: Final = await store.runs.list_runs("workflow", "team", owner, state, 10, page.next_cursor)
|
||||
assert isinstance(tail, RunPage) and tail.next_cursor is None
|
||||
expected: Final = tuple(run_id for run_id in ordered if owner is None or run_id != runs[1].run_id)
|
||||
assert tuple(run.run_id for run in (*page.runs, *tail.runs)) == expected
|
||||
await store.runs.finish(runs[0].run_id, RunState.failed, None, "failed")
|
||||
finished: Final = await store.runs.get(runs[0].run_id)
|
||||
assert finished is not None and finished.output is None and finished.metadata == {"error": "failed"}
|
||||
filtered: Final = await store.runs.list_runs("workflow", "team", "owner", RunState.failed, 10, None)
|
||||
assert isinstance(filtered, RunPage) and tuple(run.run_id for run in filtered.runs) == (finished.run_id,)
|
||||
for timestamp in (runs[0].created_at, finished.updated_at):
|
||||
assert before <= timestamp <= datetime.now(timezone.utc)
|
||||
assert isinstance(await store.runs.list_runs("workflow", "team", None, None, 10, "bad-cursor"), ExecutionError)
|
||||
|
||||
|
||||
async def test_attempt_receipts_deduplicate_checkpoints_and_exclude_foreign_spend(store: Storage) -> None:
|
||||
admitted: Final = await store.admit()
|
||||
assert isinstance(admitted, RunAdmission)
|
||||
run: Final = admitted.run
|
||||
first: Final = StepTrace(step_id="step", run_id=run.run_id, node_id="node", state=RunState.failed, input=None,
|
||||
attempts=1, started_at=run.created_at, deadline=run.input.deadline, call_id="first",
|
||||
prompt_tokens=4, completion_tokens=2, response_cost=0.125, request_id="failed")
|
||||
second: Final = first.model_copy(update={
|
||||
"attempts": 2, "state": RunState.completed, "call_id": "second", "output": False,
|
||||
})
|
||||
for step in (first, first, second, second):
|
||||
assert await store.runs.checkpoint(run.run_id, step)
|
||||
for request_id, key, team, call_id, session in (
|
||||
("a", "owner", "team", "second", run.session_id),
|
||||
("b", "owner", "team", "second", "parent-session"),
|
||||
("wrong-key", "foreign", "team", "second", run.session_id),
|
||||
("wrong-team", "owner", "other", "second", run.session_id),
|
||||
("wrong-call", "owner", "team", "unrelated", run.session_id),
|
||||
):
|
||||
await store.db.litellm_spendlogs.create(data={
|
||||
"request_id": request_id, "call_type": "acompletion", "api_key": key, "team_id": team,
|
||||
"session_id": session, "litellm_call_id": call_id, "startTime": run.created_at, "endTime": run.created_at,
|
||||
"prompt_tokens": 3, "completion_tokens": 1, "spend": 0.25,
|
||||
})
|
||||
assert await store.runs.steps(run.run_id) == (second.model_copy(update={
|
||||
"prompt_tokens": 10, "completion_tokens": 4, "response_cost": 0.625,
|
||||
"request_id": None, "request_ids": ("failed", "a", "b"),
|
||||
}),)
|
||||
events: Final = await store.db.litellm_workflowevent.find_many(where={"run_id": run.run_id})
|
||||
assert len(events) == 4
|
||||
assert all(run.created_at <= event.created_at <= datetime.now(timezone.utc) for event in events)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("phase", ["before", "retry", "downstream"])
|
||||
async def test_cancellation_fences_every_dispatch_but_keeps_completed_receipts(store: Storage, phase: str) -> None:
|
||||
snapshot: Final = await store.snapshot()
|
||||
admitted: Final = await store.admit(snapshot=snapshot)
|
||||
assert isinstance(admitted, RunAdmission)
|
||||
run: Final = admitted.run
|
||||
receipt: Final = ChildCallResult(list(_MESSAGES), 4, 2, "receipt", 0.125)
|
||||
reply: Final = ExecutionError("admission_rate_limit", "retry") if phase == "retry" else receipt
|
||||
chat: Final = ChatReplies(store, (reply,), cancel_after=1)
|
||||
if phase == "before":
|
||||
await store.runs.cancel(run.run_id)
|
||||
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
|
||||
try:
|
||||
result: Final = await runner._sequence(run, _REQUEST, snapshot.bundle, ())
|
||||
assert isinstance(result, ExecutionError) and result.code == "cancelled"
|
||||
assert len(chat.calls) == (0 if phase == "before" else 1)
|
||||
traces: Final = await store.runs.steps(run.run_id)
|
||||
expected_ids: Final = () if phase == "before" else ((),) if phase == "retry" else (("receipt",),)
|
||||
assert tuple(trace.request_ids for trace in traces) == expected_ids
|
||||
await store.runs.finish(run.run_id, RunState.completed, {"must": "not win cancellation"}, None)
|
||||
settled: Final = await store.runs.get(run.run_id)
|
||||
assert settled is not None and settled.status is RunState.cancelled and settled.output is None
|
||||
finally:
|
||||
await runner.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error", [
|
||||
ExecutionError("invalid_input", "rejected"), ExecutionError("provider_error", "unavailable"),
|
||||
ExecutionError("rate_limit", "upstream 429"), ExecutionError("runtime_error", "runtime failed"),
|
||||
ExecutionError("admission_rate_limit", "retry budget exhausted"),
|
||||
ExecutionError("provider_error", "billed failure", ChildCallResult(None, 4, 2, "billed", 0.125)),
|
||||
ExecutionError("admission_rate_limit", "known receipt", ChildCallResult(None, 4, 2, "billed", 0.125)),
|
||||
RuntimeError("private details"),
|
||||
])
|
||||
async def test_failure_stops_downstream_and_persists_sanitized_terminal_trace(
|
||||
store: Storage, error: ExecutionError | Exception,
|
||||
) -> None:
|
||||
snapshot: Final = await store.snapshot()
|
||||
admitted: Final = await store.admit(snapshot=snapshot)
|
||||
assert isinstance(admitted, RunAdmission)
|
||||
chat: Final = ChatReplies(store, (error, error))
|
||||
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
|
||||
try:
|
||||
result: Final = await runner._sequence(admitted.run, _REQUEST, snapshot.bundle, ())
|
||||
expected: Final = error if isinstance(error, ExecutionError) else ExecutionError(
|
||||
"execution_failed", "The node could not complete"
|
||||
)
|
||||
assert result == expected
|
||||
attempts: Final = 2 if expected.code == "admission_rate_limit" and expected.receipt is None else 1
|
||||
traces: Final = await store.runs.steps(admitted.run.run_id)
|
||||
assert tuple((step.node_id, step.state, step.attempts, step.error) for step in traces) == (
|
||||
("node_1", RunState.failed, attempts, expected.message),
|
||||
)
|
||||
assert len(chat.calls) == attempts
|
||||
assert (traces[0].prompt_tokens, traces[0].completion_tokens, traces[0].response_cost, traces[0].request_ids) == (
|
||||
(4, 2, 0.125, ("billed",)) if expected.receipt is not None else (0, 0, None, ())
|
||||
)
|
||||
finally:
|
||||
await runner.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("trigger", ["get", "admit", "list"])
|
||||
async def test_deadline_expiry_is_terminal_and_fences_dispatch(store: Storage, trigger: str) -> None:
|
||||
admitted: Final = await store.admit(deadline=datetime.now(timezone.utc) - timedelta(seconds=1))
|
||||
assert isinstance(admitted, RunAdmission)
|
||||
if trigger == "admit":
|
||||
await store.admit()
|
||||
elif trigger == "list":
|
||||
await store.runs.list_runs("workflow", "team", None, None, 10, None)
|
||||
else:
|
||||
await store.runs.get(admitted.run.run_id)
|
||||
stored: Final = await store.db.litellm_workflowrun.find_unique(where={"run_id": admitted.run.run_id})
|
||||
assert stored is not None and stored.status == "failed"
|
||||
step: Final = StepTrace(
|
||||
step_id="step", run_id=stored.run_id, node_id="node", state=RunState.running,
|
||||
input=None, attempts=1, started_at=admitted.run.created_at, deadline=admitted.run.input.deadline,
|
||||
)
|
||||
assert not await store.runs.checkpoint(stored.run_id, step)
|
||||
await store.runs.finish(stored.run_id, RunState.completed, "late output", None)
|
||||
expired: Final = await store.runs.get(stored.run_id)
|
||||
assert expired is not None and expired.status is RunState.failed and expired.output is None
|
||||
assert await store.runs.steps(stored.run_id) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("role,user", [
|
||||
(LitellmUserRoles.PROXY_ADMIN, "admin"), (LitellmUserRoles.INTERNAL_USER, "admin"),
|
||||
(LitellmUserRoles.INTERNAL_USER, "member"), (None, None),
|
||||
])
|
||||
async def test_key_scope_intersects_user_scope_and_own_application_key_can_cancel(
|
||||
store: Storage, role: LitellmUserRoles | None, user: str | None,
|
||||
) -> None:
|
||||
local, foreign = await store.snapshot(), await store.snapshot("other")
|
||||
own: Final = await store.admit(snapshot=local)
|
||||
peer: Final = await store.admit(snapshot=local, owner="peer")
|
||||
other: Final = await store.admit(snapshot=foreign, team="other")
|
||||
assert isinstance(own, RunAdmission) and isinstance(peer, RunAdmission) and isinstance(other, RunAdmission)
|
||||
caller: Final = _CALLER.model_copy(update={"user_role": role, "user_id": user})
|
||||
for endpoint in (get_workflow_steps, get_run_definition, cancel_workflow_run):
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await endpoint(other.run.run_id, caller, store.definitions, store.runs)
|
||||
assert denied.value.status_code == 404
|
||||
with pytest.raises(HTTPException) as denied_history:
|
||||
await list_definition_runs(
|
||||
foreign.workflow_id, caller, store.definitions, store.runs, team_id="other", limit=10, after=None,
|
||||
)
|
||||
assert denied_history.value.status_code == 404
|
||||
history: Final = await list_definition_runs(
|
||||
local.workflow_id, caller, store.definitions, store.runs, limit=10, after=None,
|
||||
)
|
||||
expected: Final = {own.run.run_id, peer.run.run_id} if user == "admin" else {own.run.run_id}
|
||||
assert {run.run_id for run in history.runs} == expected
|
||||
for endpoint in (get_workflow_steps, get_run_definition, cancel_workflow_run):
|
||||
if user == "admin":
|
||||
assert await endpoint(peer.run.run_id, caller, store.definitions, store.runs) is not None
|
||||
else:
|
||||
with pytest.raises(HTTPException) as denied_peer:
|
||||
await endpoint(peer.run.run_id, caller, store.definitions, store.runs)
|
||||
assert denied_peer.value.status_code == 404
|
||||
pinned: Final = await get_run_definition(own.run.run_id, caller, store.definitions, store.runs)
|
||||
assert pinned.snapshot_id == local.snapshot_id
|
||||
assert (await cancel_workflow_run(own.run.run_id, caller, store.definitions, store.runs)).status_code == 202
|
||||
assert (await get_workflow_steps(own.run.run_id, caller, store.definitions, store.runs)).cancel_requested
|
||||
foreign_trace: Final = await get_workflow_steps(
|
||||
other.run.run_id, _CALLER.model_copy(update={"team_id": "other"}), store.definitions, store.runs,
|
||||
)
|
||||
assert not foreign_trace.cancel_requested
|
||||
|
||||
|
||||
async def test_submit_drains_pinned_snapshot_serializes_chat_and_never_cold_replays(store: Storage) -> None:
|
||||
snapshot: Final = await store.snapshot()
|
||||
admitted: Final = await store.admit(snapshot=snapshot, key="original")
|
||||
orphan: Final = await store.admit(snapshot=snapshot)
|
||||
assert isinstance(admitted, RunAdmission) and isinstance(orphan, RunAdmission)
|
||||
draft: Final = await store.definitions.get_draft(snapshot.workflow_id, "team")
|
||||
assert draft is not None
|
||||
changed: Final = snapshot.bundle.model_copy(update={
|
||||
"nodes": tuple(node.model_copy(update={"model": "edited"}) for node in snapshot.bundle.nodes),
|
||||
})
|
||||
replaced: Final = await store.definitions.replace(snapshot.workflow_id, "team", draft.draft.etag, changed, "author")
|
||||
assert isinstance(replaced, WorkflowDraft)
|
||||
republished: Final = await store.definitions.publish(
|
||||
snapshot.workflow_id, "team", replaced.etag, replaced.digest, "admin",
|
||||
)
|
||||
assert isinstance(republished, WorkflowSnapshot)
|
||||
chat: Final = ChatReplies(store, (
|
||||
ExecutionError("admission_rate_limit", "retry"),
|
||||
ChildCallResult(list(_MESSAGES), 4, 2, "first", 0.125),
|
||||
ChildCallResult({"answer": False}, 6, 3, "last", 0.25),
|
||||
))
|
||||
runner: Final = WorkflowRunner(store.runs, store.definitions, chat)
|
||||
try:
|
||||
assert await runner.submit(admitted.run, _REQUEST, _CALLER) is None
|
||||
finally:
|
||||
await runner.close()
|
||||
completed: Final = await store.runs.get(admitted.run.run_id)
|
||||
assert completed is not None and completed.status is RunState.completed, completed
|
||||
completion: Final = ChatCompletion.model_validate(completed.output)
|
||||
assert json.loads(completion.choices[0].message.content or "null") == {"answer": False}
|
||||
assert completion.usage is not None
|
||||
usage: Final = completion.usage
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert chat.models == ("gateway-model",) * 3
|
||||
assert tuple((step.node_id, step.attempts) for step in chat.calls) == (("node_1", 1), ("node_1", 2), ("node_2", 1))
|
||||
assert chat.calls[0].step_id == chat.calls[1].step_id and chat.calls[0].deadline == chat.calls[1].deadline
|
||||
assert len({step.call_id for step in chat.calls}) == 3
|
||||
assert tuple((step.state, step.call_id) for step in chat.persisted) == tuple(
|
||||
(RunState.running, step.call_id) for step in chat.calls
|
||||
)
|
||||
assert chat.calls[1].input == chat.calls[0].input == chat.calls[2].input == list(_MESSAGES)
|
||||
pinned: Final = await get_run_definition(completed.run_id, _CALLER, store.definitions, store.runs)
|
||||
assert pinned.bundle == snapshot.bundle
|
||||
restarted: Final = WorkflowRunner(store.runs, store.definitions, chat)
|
||||
await restarted.close()
|
||||
pending: Final = await store.runs.get(orphan.run.run_id)
|
||||
assert pending is not None and pending.status is RunState.pending and len(chat.calls) == 3
|
||||
assert await store.runs.steps(orphan.run.run_id) == ()
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import replace
|
||||
|
|
@ -13,6 +14,7 @@ import pytest_asyncio
|
|||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.proxy.workflows.monty._compat import timeout
|
||||
from litellm.proxy.workflows.monty.compiler import compile_program
|
||||
from litellm.proxy.workflows.monty.models import (
|
||||
CompiledProgram,
|
||||
|
|
@ -24,6 +26,7 @@ from litellm.proxy.workflows.monty.models import (
|
|||
PythonProgram,
|
||||
)
|
||||
from litellm.proxy.workflows.monty.runtime import BinaryArtifact, MontyRuntime, installed_binary
|
||||
from litellm.proxy.workflows.monty.validation import ValidatorProcess
|
||||
from tests.integration._support.process import stop_processes
|
||||
|
||||
pytest.importorskip("pydantic_monty")
|
||||
|
|
@ -426,6 +429,45 @@ def _slow_schema(depth: int) -> JsonSchema:
|
|||
return {**definitions, f"N{depth}": {"type": "string"}}
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "linux", reason="address-space limits require Linux prlimit")
|
||||
@pytest.mark.asyncio
|
||||
async def test_standalone_validator_bounds_child_memory_without_limiting_gateway() -> None:
|
||||
import resource
|
||||
|
||||
parent_limit: Final = resource.getrlimit(resource.RLIMIT_AS)
|
||||
validator: Final = ValidatorProcess()
|
||||
try:
|
||||
assert await validator.start() is None
|
||||
assert validator.process is not None
|
||||
soft, hard = resource.prlimit(validator.process.pid, resource.RLIMIT_AS)
|
||||
assert 0 < soft <= hard <= MontyLimits().worker_address_space_bytes
|
||||
assert resource.getrlimit(resource.RLIMIT_AS) == parent_limit
|
||||
assert await validator.validate("valid", {"type": "string"}, deadline=time.monotonic() + 5) is None
|
||||
finally:
|
||||
await validator.close()
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "linux", reason="address-space limits require Linux prlimit")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("output", [False, True])
|
||||
async def test_validator_memory_exhaustion_fails_closed_and_recovers(output: bool) -> None:
|
||||
instance: Final = MontyRuntime(MontyLimits(max_concurrent_workers=1, worker_address_space_bytes=64 * 1024 * 1024))
|
||||
assert await instance.start() is None
|
||||
ref: Final[JsonSchema] = {"$ref": "#/$defs/N0"}
|
||||
compiled: Final = _compile(
|
||||
"def transform(value): return value", ref if output else {}, {} if output else ref, _slow_schema(22)
|
||||
)
|
||||
try:
|
||||
result: Final = await instance.run_program(compiled, {"value": None}, deadline_monotonic=time.monotonic() + 15)
|
||||
assert isinstance(result, MontyFailure), result
|
||||
assert result.kind == MontyErrorKind.WORKER_CRASHED, result
|
||||
recovered: Final = await _run(instance, _compile("def transform(value): return value"), "recovered")
|
||||
assert isinstance(recovered, MontySuccess), recovered
|
||||
assert recovered.output == "recovered"
|
||||
finally:
|
||||
await instance.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("output", [False, True])
|
||||
async def test_expensive_allowed_validation_times_out_without_blocking_and_recovers(
|
||||
|
|
@ -546,7 +588,7 @@ def _worker_pids() -> frozenset[int]:
|
|||
|
||||
|
||||
async def _active_worker(previous: frozenset[int]) -> int:
|
||||
async with asyncio.timeout(5):
|
||||
async with timeout(5):
|
||||
while True:
|
||||
current: Final = _worker_pids() - previous
|
||||
if current:
|
||||
|
|
@ -599,7 +641,7 @@ async def test_fresh_worker_has_no_parent_canary_and_has_address_space_limit(
|
|||
task: Final = asyncio.create_task(_run(runtime, _compile("def transform(value):\n while True: pass"), None))
|
||||
pid: Final = await _active_worker(previous)
|
||||
assert b"LITEFLOW_ENV_CANARY" not in Path(f"/proc/{pid}/environ").read_bytes()
|
||||
async with asyncio.timeout(5):
|
||||
async with timeout(5):
|
||||
while str(runtime.limits.worker_address_space_bytes) not in Path(f"/proc/{pid}/limits").read_text():
|
||||
await asyncio.sleep(0.01)
|
||||
assert str(runtime.limits.worker_address_space_bytes) in Path(f"/proc/{pid}/limits").read_text()
|
||||
|
|
|
|||
51
tests/unit/models/test_workflow_definition.py
Normal file
51
tests/unit/models/test_workflow_definition.py
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.models.workflow_definition import WorkflowBundle
|
||||
from tests.unit.proxy.workflow_definition_rows import bundle
|
||||
|
||||
|
||||
@pytest.mark.parametrize("node_type", ["parallel", "if", "shell", "http"])
|
||||
def test_v1_contract_rejects_nodes_outside_sequential_chat_and_python(node_type: str) -> None:
|
||||
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
|
||||
unsupported: Final = {**original, "nodes": [{**original["nodes"][0], "type": node_type}]}
|
||||
with pytest.raises(ValidationError) as rejected:
|
||||
WorkflowBundle.model_validate(unsupported)
|
||||
assert rejected.value.errors(include_input=False)[0]["type"] == "union_tag_invalid"
|
||||
|
||||
|
||||
def test_binding_cannot_supply_both_literal_and_reference() -> None:
|
||||
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
|
||||
ambiguous: Final = {
|
||||
**original,
|
||||
"nodes": [
|
||||
{
|
||||
**original["nodes"][0],
|
||||
"messages": {"from": "$input", "output": "messages", "value": [{"role": "user", "content": "hello"}]},
|
||||
}
|
||||
],
|
||||
}
|
||||
with pytest.raises(ValidationError) as rejected:
|
||||
WorkflowBundle.model_validate(ambiguous)
|
||||
assert {error["loc"][-1] for error in rejected.value.errors(include_input=False) if "messages" in error["loc"]} == {
|
||||
"from",
|
||||
"output",
|
||||
"value",
|
||||
}
|
||||
|
||||
|
||||
def test_nonfinite_values_are_rejected_before_digest_serialization() -> None:
|
||||
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
|
||||
with pytest.raises(ValidationError) as rejected:
|
||||
WorkflowBundle.model_validate({**original, "input_schema": {"minimum": float("nan")}})
|
||||
assert rejected.value.errors(include_input=False)[0]["type"] == "finite_number"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["api_key", "api_base", "base_url", "headers", "callbacks", "user_config"])
|
||||
def test_definition_cannot_supply_provider_credentials_or_runtime_overrides(field: str) -> None:
|
||||
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
|
||||
with pytest.raises(ValidationError) as rejected:
|
||||
WorkflowBundle.model_validate({**original, "nodes": [{**original["nodes"][0], field: "untrusted"}]})
|
||||
assert rejected.value.errors(include_input=False)[0]["type"] == "extra_forbidden"
|
||||
|
|
@ -5130,3 +5130,69 @@ def test_member_can_reach_skill_history_visibility_checks(route: str) -> None:
|
|||
)
|
||||
except HTTPException as exc:
|
||||
pytest.fail(f"History visibility could not run: {exc.detail}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [
|
||||
"/v1/workflows",
|
||||
"/v1/workflows/definition-id",
|
||||
"/v1/workflows/definition-id/validate",
|
||||
"/v1/workflows/definition-id/publish",
|
||||
])
|
||||
def test_definition_management_routes_reach_team_authorization(route: str) -> None:
|
||||
token: Final = UserAPIKeyAuth(
|
||||
user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
allowed_routes=["management_routes"],
|
||||
)
|
||||
request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""})
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={}
|
||||
)
|
||||
assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [
|
||||
"/v1/workflows/runs",
|
||||
"/v1/workflows/runs/validate",
|
||||
"/v1/workflows/runs/publish",
|
||||
"/v1/workflows/definition-id/delete",
|
||||
"/v1/workflows/definition-id/test",
|
||||
])
|
||||
def test_definition_route_grants_do_not_grant_run_or_future_actions(route: str) -> None:
|
||||
token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_definition_routes"])
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
RouteChecks.is_virtual_key_allowed_to_call_route(route, token)
|
||||
assert rejected.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [
|
||||
"/v1/workflows/definition-id/test",
|
||||
"/v1/workflows/definition-id/runs",
|
||||
"/v1/workflows/runs/run-id/steps",
|
||||
"/v1/workflows/runs/run-id/definition",
|
||||
"/v1/workflows/runs/run-id/cancel",
|
||||
])
|
||||
def test_execution_routes_reach_their_scoped_authorization(route: str) -> None:
|
||||
token: Final = UserAPIKeyAuth(
|
||||
user_id="author", team_id="team", user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
allowed_routes=["management_routes"],
|
||||
)
|
||||
request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""})
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=None, _user_role=token.user_role, route=route, request=request, valid_token=token, request_data={}
|
||||
)
|
||||
assert RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [
|
||||
"/v1/workflows/runs",
|
||||
"/v1/workflows/runs/run-id",
|
||||
"/v1/workflows/runs/run-id/events",
|
||||
"/v1/workflows/runs/run-id/messages",
|
||||
])
|
||||
def test_execution_route_grants_preserve_legacy_tracking_admission(route: str) -> None:
|
||||
token: Final = UserAPIKeyAuth(user_id="author", allowed_routes=["workflow_execution_routes"])
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
RouteChecks.is_virtual_key_allowed_to_call_route(route, token)
|
||||
assert rejected.value.status_code == 403
|
||||
assert RouteChecks.check_route_access(route, LiteLLMRoutes.self_managed_routes.value) is False
|
||||
assert RouteChecks.check_route_access(route, LiteLLMRoutes.management_routes.value) is False
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import sys
|
|||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import datetime, timedelta, timezone, tzinfo
|
||||
from functools import partial
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
|
@ -69,7 +69,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model, safe_set_request_parsed_body
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
_check_team_member_model_budget_with_fallback,
|
||||
|
|
@ -81,6 +81,7 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
_reserve_budget_after_common_checks,
|
||||
_route_requires_auth_despite_public,
|
||||
_routing_selector_matches_claim,
|
||||
authorize_internal_virtual_key,
|
||||
run_centralized_common_checks,
|
||||
_run_post_custom_auth_checks,
|
||||
user_api_key_auth_builder,
|
||||
|
|
@ -101,6 +102,118 @@ class _RoutingRequest:
|
|||
self.state = SimpleNamespace()
|
||||
|
||||
|
||||
async def _internal_key_cache(
|
||||
monkeypatch: pytest.MonkeyPatch, *, models: tuple[str, ...] = (), max_budget: float | None = None,
|
||||
key_hash: str = "b" * 64,
|
||||
) -> UserApiKeyCache:
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache_key_object(
|
||||
hashed_token=key_hash,
|
||||
user_api_key_obj=UserAPIKeyAuth(
|
||||
token=key_hash, team_id="original-team", models=list(models), spend=2.0, max_budget=max_budget,
|
||||
),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
key="team_id:original-team", value=LiteLLM_TeamTableCachedObj(team_id="original-team")
|
||||
)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", object())
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {})
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
|
||||
counters: Final = DualCache()
|
||||
counters.in_memory_cache.set_cache(key=f"spend:key:{key_hash}", value=2.0)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "spend_counter_cache", counters)
|
||||
monkeypatch.setattr(litellm, "max_budget", 0.0)
|
||||
return cache
|
||||
|
||||
|
||||
def _internal_management_request() -> Request:
|
||||
return Request({
|
||||
"type": "http", "path": "/workflow/runs", "method": "POST", "headers": (), "query_string": b"",
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_key_cannot_execute_for_a_different_team(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
await _internal_key_cache(monkeypatch)
|
||||
request: Final = _internal_management_request()
|
||||
with pytest.raises(HTTPException, match="no longer belongs") as error:
|
||||
await authorize_internal_virtual_key(
|
||||
"b" * 64, request, {"model": "workflow-model"},
|
||||
route="/v1/chat/completions", expected_team_id="another-team",
|
||||
)
|
||||
assert error.value.status_code == 403
|
||||
assert getattr(request.state, "budget_reservation", None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("restriction", ["model", "budget"])
|
||||
async def test_internal_execution_route_enforces_inference_model_and_key_budget(
|
||||
monkeypatch: pytest.MonkeyPatch, restriction: str,
|
||||
) -> None:
|
||||
await _internal_key_cache(
|
||||
monkeypatch,
|
||||
models=("allowed-model",) if restriction == "model" else (),
|
||||
max_budget=1.0 if restriction == "budget" else None,
|
||||
)
|
||||
with pytest.raises(litellm.BudgetExceededError if restriction == "budget" else ProxyException) as error:
|
||||
await authorize_internal_virtual_key(
|
||||
"b" * 64, _internal_management_request(), {"model": "workflow-model"},
|
||||
route="/v1/chat/completions", expected_team_id="original-team",
|
||||
)
|
||||
if restriction == "budget":
|
||||
assert isinstance(error.value, litellm.BudgetExceededError)
|
||||
assert error.value.current_cost == 2.0 and error.value.max_budget == 1.0
|
||||
else:
|
||||
assert isinstance(error.value, ProxyException)
|
||||
assert error.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_execution_route_keeps_proxy_network_policy(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
await _internal_key_cache(monkeypatch)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"allowed_ips": ["192.0.2.1"]})
|
||||
with pytest.raises(HTTPException, match="not allowed") as error:
|
||||
await authorize_internal_virtual_key(
|
||||
"b" * 64, _internal_management_request(), {"model": "workflow-model"},
|
||||
route="/v1/chat/completions", expected_team_id="original-team",
|
||||
)
|
||||
assert error.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_wrapper_with_model_grant_passes_strict_pricing_but_unpriced_child_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
raw_key: Final = "sk-workflow-pricing-fixture"
|
||||
key_hash: Final = proxy_utils.hash_token(raw_key)
|
||||
await _internal_key_cache(monkeypatch, models=("workflow/example", "unpriced-child"), key_hash=key_hash)
|
||||
router: Final = litellm.Router(
|
||||
model_list=[{
|
||||
"model_name": "unpriced-child",
|
||||
"litellm_params": {"model": "openai/workflow-unpriced-fixture", "api_key": "test-only"},
|
||||
}],
|
||||
)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
||||
monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True)
|
||||
request: Final = Request({
|
||||
"type": "http", "path": "/v1/chat/completions", "method": "POST",
|
||||
"headers": ((b"authorization", f"Bearer {raw_key}".encode()),), "query_string": b"",
|
||||
})
|
||||
safe_set_request_parsed_body(request, {"model": "workflow/example", "messages": [{"role": "user", "content": "hello"}]})
|
||||
outer: Final = await user_api_key_auth(request=request, api_key=f"Bearer {raw_key}")
|
||||
assert outer.team_id == "original-team"
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await authorize_internal_virtual_key(
|
||||
key_hash, _internal_management_request(), {"model": "unpriced-child"},
|
||||
route="/v1/chat/completions", expected_team_id="original-team",
|
||||
)
|
||||
assert error.value.type == ProxyErrorTypes.model_cost_map_missing
|
||||
|
||||
|
||||
async def _real_model_budget_limiter(
|
||||
*,
|
||||
member_spend: Mapping[str, float] | None = None,
|
||||
|
|
@ -466,10 +579,17 @@ async def test_custom_auth_does_not_enforce_key_model_access_by_default():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_custom_auth_expired_key_returns_unauthorized():
|
||||
expired_token = UserAPIKeyAuth(
|
||||
async def test_post_custom_auth_expired_key_returns_unauthorized(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
class Clock(datetime):
|
||||
@classmethod
|
||||
def now(cls, tz: tzinfo | None = None) -> datetime:
|
||||
current: Final = datetime(2026, 1, 1, 12, tzinfo=timezone.utc)
|
||||
return current.astimezone(tz) if tz is not None else current.replace(tzinfo=None)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.datetime", Clock)
|
||||
expired_token: Final = UserAPIKeyAuth(
|
||||
token="test_token",
|
||||
expires=datetime.now() - timedelta(minutes=1),
|
||||
expires=Clock(2026, 1, 1, 11, 59, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,157 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from liteflow import ChatCompletion, Draft, GatewayClient, Input, Workflow, bundle_digest, compile_workflow
|
||||
|
||||
from litellm.models.workflow_definition import DefinitionPage
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.deps.repositories import get_workflow_definition_repository
|
||||
from litellm.proxy.management_endpoints.workflow_definition_endpoints import router
|
||||
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
|
||||
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
|
||||
|
||||
|
||||
def _client(database: DefinitionRows | None, identity: UserAPIKeyAuth) -> TestClient:
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: identity
|
||||
app.dependency_overrides[get_workflow_definition_repository] = lambda: (
|
||||
None if database is None else WorkflowDefinitionRepository(database.database)
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _sdk(client: TestClient) -> GatewayClient:
|
||||
def send(request: httpx.Request) -> httpx.Response:
|
||||
response: Final = client.request(
|
||||
request.method, str(request.url), content=request.content, headers=dict(request.headers)
|
||||
)
|
||||
return httpx.Response(response.status_code, headers=dict(response.headers), content=response.content)
|
||||
|
||||
return GatewayClient(
|
||||
base_url="http://testserver",
|
||||
api_key="synthetic",
|
||||
team_id="team",
|
||||
transport=httpx.MockTransport(send),
|
||||
)
|
||||
|
||||
|
||||
def _body(name: str = "example", model: str = "gateway-model") -> Mapping[str, object]:
|
||||
return {"bundle": bundle(name, model).model_dump(mode="json", by_alias=True, exclude_unset=True)}
|
||||
|
||||
|
||||
def _register(client: TestClient, name: str = "example") -> Draft:
|
||||
response: Final = client.post("/v1/workflows", json=_body(name))
|
||||
assert response.status_code == 201, response.text
|
||||
return Draft.model_validate(response.json())
|
||||
|
||||
|
||||
def test_real_sdk_registers_validates_replaces_and_publishes_with_an_independent_admin() -> None:
|
||||
database: Final = DefinitionRows()
|
||||
author_http: Final = _client(database, caller("reviewer"))
|
||||
publisher_http: Final = _client(
|
||||
database, UserAPIKeyAuth(user_id="publisher", user_role=LitellmUserRoles.PROXY_ADMIN, team_id="team")
|
||||
)
|
||||
workflow: Final = Workflow(name="example", nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)])
|
||||
with _sdk(author_http) as author, _sdk(publisher_http) as publisher:
|
||||
draft: Final = author.register(workflow)
|
||||
assert draft.bundle == compile_workflow(workflow) and draft.digest == bundle_digest(draft.bundle)
|
||||
assert author.get(draft.workflow_id) == draft
|
||||
assert author_http.get(f"/v1/workflows/{draft.workflow_id}").headers["etag"] == draft.etag
|
||||
assert author.validate(draft) == {
|
||||
"valid": True,
|
||||
"workflow_id": draft.workflow_id,
|
||||
"etag": draft.etag,
|
||||
"digest": draft.digest,
|
||||
}
|
||||
with pytest.raises(httpx.HTTPStatusError) as denied:
|
||||
author.publish(draft)
|
||||
assert denied.value.response.status_code == 403
|
||||
replacement: Final = author.register(workflow, draft=draft)
|
||||
assert replacement.etag != draft.etag
|
||||
assert publisher.publish(replacement) == {
|
||||
"workflow_id": draft.workflow_id,
|
||||
"snapshot_id": author.get(draft.workflow_id).published_snapshot_id,
|
||||
"model": "workflow/example",
|
||||
"digest": replacement.digest,
|
||||
}
|
||||
|
||||
|
||||
def test_stale_or_missing_revision_and_wrong_digest_cannot_publish_or_replace() -> None:
|
||||
database: Final = DefinitionRows()
|
||||
client: Final = _client(database, caller("reviewer"))
|
||||
draft: Final = _register(client)
|
||||
path: Final = f"/v1/workflows/{draft.workflow_id}"
|
||||
changed: Final = client.put(path, json=_body(model="other-model"), headers={"If-Match": draft.etag})
|
||||
assert changed.status_code == 200, changed.text
|
||||
assert client.put(path, json=_body()).status_code == 412
|
||||
assert client.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 412
|
||||
assert client.post(f"{path}/validate").status_code == 412
|
||||
assert (
|
||||
client.post(f"{path}/publish", json={"digest": draft.digest}, headers={"If-Match": draft.etag}).status_code
|
||||
== 412
|
||||
)
|
||||
assert (
|
||||
client.post(
|
||||
f"{path}/publish", json={"digest": "0" * 64}, headers={"If-Match": changed.json()["etag"]}
|
||||
).status_code
|
||||
== 412
|
||||
)
|
||||
assert client.get(path).json() == changed.json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body, status", [(b"", 200), (b"{}", 200), (b"null", 200), (b'{"execute":true}', 422)])
|
||||
def test_validate_accepts_only_an_optional_empty_action(body: bytes, status: int) -> None:
|
||||
client: Final = _client(DefinitionRows(), caller())
|
||||
draft: Final = _register(client)
|
||||
response: Final = client.post(
|
||||
f"/v1/workflows/{draft.workflow_id}/validate",
|
||||
content=body,
|
||||
headers={"If-Match": draft.etag, "Content-Type": "application/json"},
|
||||
)
|
||||
assert response.status_code == status, response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("user, team", [("outsider", "other-team"), ("not-a-team-member", "team")])
|
||||
def test_inaccessible_draft_cannot_be_read_edited_or_published(user: str, team: str) -> None:
|
||||
database: Final = DefinitionRows()
|
||||
draft: Final = _register(_client(database, caller()))
|
||||
client: Final = _client(database, caller(user, team))
|
||||
path: Final = f"/v1/workflows/{draft.workflow_id}"
|
||||
assert client.get(path).status_code == 404
|
||||
assert client.get(path, params={"team_id": "team"}).status_code == 404
|
||||
assert client.put(path, json=_body(), headers={"If-Match": draft.etag}).status_code == 404
|
||||
assert (
|
||||
client.post(f"{path}/publish", json={"digest": draft.digest}, headers={"If-Match": draft.etag}).status_code
|
||||
== 404
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"identity, expected",
|
||||
[
|
||||
(caller(None), 403),
|
||||
(UserAPIKeyAuth(user_id="author", user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY), 403),
|
||||
(caller(), 500),
|
||||
],
|
||||
)
|
||||
def test_no_database_preserves_authorization_precedence(identity: UserAPIKeyAuth, expected: int) -> None:
|
||||
response: Final = _client(None, identity).post("/v1/workflows", json=_body())
|
||||
assert response.status_code == expected, response.text
|
||||
|
||||
|
||||
def test_duplicate_registration_conflicts_and_cursor_pagination_returns_each_draft_once() -> None:
|
||||
client: Final = _client(DefinitionRows(), caller())
|
||||
first: Final = _register(client, "first")
|
||||
second: Final = _register(client, "second")
|
||||
assert client.post("/v1/workflows", json=_body("first")).status_code == 409
|
||||
assert client.get("/v1/workflows", params={"cursor": "malformed"}).status_code == 422
|
||||
page: Final = DefinitionPage.model_validate(client.get("/v1/workflows", params={"limit": 1}).json())
|
||||
tail: Final = DefinitionPage.model_validate(client.get("/v1/workflows", params={"cursor": page.next_cursor}).json())
|
||||
assert {page.workflows[0].workflow_id, tail.workflows[0].workflow_id} == {first.workflow_id, second.workflow_id}
|
||||
assert page.workflows[0].workflow_id != tail.workflows[0].workflow_id and tail.next_cursor is None
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.workflow_execution_endpoints import (
|
||||
TestRunRequest as PaidRunRequest, WorkflowInvocation, _execution_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("overrides", [
|
||||
{"is_session_token": True}, {"via_virtual_key": False}, {"team_id": "other"}, {"token": None},
|
||||
])
|
||||
def test_execution_requires_a_team_bound_application_key(overrides: dict[str, object]) -> None:
|
||||
owner: Final = UserAPIKeyAuth(token="owner", team_id="team").model_copy(update={"via_virtual_key": True})
|
||||
assert _execution_key(owner, "team") == "owner"
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
_execution_key(owner.model_copy(update=overrides), "team")
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extra", [{"stream": True}, {"tools": []}, {"temperature": 0}])
|
||||
def test_workflow_invocation_rejects_unsupported_chat_parameters(extra: dict[str, object]) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
WorkflowInvocation.model_validate({
|
||||
"model": "workflow/example", "messages": [{"role": "user", "content": "go"}], **extra,
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("consent", ["true", 1, None])
|
||||
def test_paid_test_requires_explicit_boolean_billing_consent(consent: object) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
PaidRunRequest.model_validate({"messages": [{"role": "user", "content": "go"}], "authorize_billing": consent})
|
||||
|
|
@ -4,8 +4,9 @@ Uses FastAPI TestClient with a mocked prisma_client.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -113,8 +114,15 @@ def _make_prisma_client() -> MagicMock:
|
|||
|
||||
|
||||
def _make_app() -> FastAPI:
|
||||
app = FastAPI()
|
||||
from litellm.proxy.deps.repositories import (
|
||||
get_workflow_definition_repository,
|
||||
get_workflow_execution_repository,
|
||||
)
|
||||
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_workflow_definition_repository] = lambda: None
|
||||
app.dependency_overrides[get_workflow_execution_repository] = lambda: None
|
||||
return app
|
||||
|
||||
|
||||
|
|
@ -223,6 +231,7 @@ class TestListWorkflowRuns:
|
|||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 1
|
||||
assert self._prisma.db.litellm_workflowrun.find_many.call_args.kwargs["where"]["engine_owned"] is False
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_filters_by_status(self, mock_pc):
|
||||
|
|
@ -734,3 +743,54 @@ class TestAdminViewerReadParity:
|
|||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(_require_run(prisma, "run-1", _override_auth_admin_viewer()))
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "suffix", "body"),
|
||||
[
|
||||
("PATCH", "", {"status": "completed", "output": {"forged": True}}),
|
||||
("POST", "/events", {"event_type": "step.failed", "step_name": "node-1"}),
|
||||
("POST", "/messages", {"role": "assistant", "content": "forged result"}),
|
||||
],
|
||||
)
|
||||
def test_engine_owned_runs_reject_tracking_writes_even_for_proxy_admin(
|
||||
method: str, suffix: str, body: Mapping[str, object],
|
||||
) -> None:
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
database: Final = _make_prisma_client()
|
||||
run: Final = _make_run()
|
||||
run.engine_owned = True
|
||||
database.db.litellm_workflowrun.find_unique = AsyncMock(return_value=run)
|
||||
database.db.litellm_workflowrun.update = AsyncMock(return_value=run)
|
||||
database.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
database.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
|
||||
database.db.litellm_workflowmessage.create = AsyncMock(return_value=_make_message())
|
||||
app: Final = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth_admin
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", database), TestClient(app) as client:
|
||||
response: Final = client.request(method, f"/v1/workflows/runs/run-1{suffix}", json=body)
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json() == {"detail": "Gateway executions cannot be changed through tracking endpoints"}
|
||||
database.db.litellm_workflowrun.update.assert_not_awaited()
|
||||
database.db.litellm_workflowmessage.create.assert_not_awaited()
|
||||
database.db.tx.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("suffix", ["", "/events", "/messages"])
|
||||
def test_engine_run_content_is_only_exposed_by_scoped_execution_routes(suffix: str) -> None:
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
database: Final = _make_prisma_client()
|
||||
run: Final = _make_run()
|
||||
run.engine_owned = True
|
||||
database.db.litellm_workflowrun.find_unique = AsyncMock(return_value=run)
|
||||
database.db.litellm_workflowevent.find_many = AsyncMock()
|
||||
database.db.litellm_workflowmessage.find_many = AsyncMock()
|
||||
app: Final = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth_admin
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", database), TestClient(app) as client:
|
||||
response: Final = client.get(f"/v1/workflows/runs/run-1{suffix}")
|
||||
assert response.status_code == 404
|
||||
database.db.litellm_workflowevent.find_many.assert_not_awaited()
|
||||
database.db.litellm_workflowmessage.find_many.assert_not_awaited()
|
||||
|
|
|
|||
122
tests/unit/proxy/workflow_definition_rows.py
Normal file
122
tests/unit/proxy/workflow_definition_rows.py
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncGenerator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, cast # noqa: TID251 # the test database implements the delegates exercised by these routes
|
||||
|
||||
from liteflow import ChatCompletion, Input, Workflow, compile_workflow
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.models.workflow_definition import WorkflowBundle
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.repositories.workflow_definition_repository import DefinitionDatabase
|
||||
|
||||
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
_COMPARISON: Final = TypeAdapter(Mapping[str, str])
|
||||
|
||||
|
||||
def bundle(name: str = "example", model: str = "gateway-model") -> WorkflowBundle:
|
||||
return WorkflowBundle.model_validate(
|
||||
compile_workflow(Workflow(name=name, nodes=[ChatCompletion(model=model, messages=Input.messages)]))
|
||||
)
|
||||
|
||||
|
||||
def caller(user_id: str | None = "author", team_id: str = "team") -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(user_id=user_id, team_id=team_id, user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
|
||||
|
||||
def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool:
|
||||
return all(
|
||||
str(row.get(key)) > _COMPARISON.validate_python(value)["gt"]
|
||||
if isinstance(value, Mapping) and "gt" in value
|
||||
else row.get(key) == value
|
||||
for key, value in where.items()
|
||||
)
|
||||
|
||||
|
||||
class Rows:
|
||||
def __init__(self, rows: tuple[Mapping[str, object], ...] = ()) -> None:
|
||||
self.rows = rows
|
||||
|
||||
async def find_first(
|
||||
self,
|
||||
skip: int | None = None,
|
||||
where: Mapping[str, object] | None = None,
|
||||
cursor: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
|
||||
distinct: Sequence[str] | None = None,
|
||||
) -> object | None:
|
||||
return next((row for row in self.rows if _matches(row, where or {})), None)
|
||||
|
||||
async def find_unique(
|
||||
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
|
||||
) -> object | None:
|
||||
return await self.find_first(where=where)
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
take: int | None = None,
|
||||
skip: int | None = None,
|
||||
where: Mapping[str, object] | None = None,
|
||||
cursor: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
|
||||
distinct: Sequence[str] | None = None,
|
||||
) -> Sequence[object]:
|
||||
matches: Final = tuple(row for row in self.rows if _matches(row, where or {}))
|
||||
return sorted(matches, key=lambda row: str(row.get("workflow_id")))[:take]
|
||||
|
||||
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> object:
|
||||
if "etag" in data and any(
|
||||
row.get("name") == data["name"] and row.get("team_id") == data["team_id"] for row in self.rows
|
||||
):
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
raise UniqueViolationError({"is_panic": False, "message": "duplicate", "error_code": "P2002"})
|
||||
row: Final = {
|
||||
"created_at": NOW,
|
||||
"published_etag": None,
|
||||
"published_snapshot_id": None,
|
||||
**data,
|
||||
}
|
||||
self.rows = (*self.rows, row)
|
||||
return row
|
||||
|
||||
async def update(
|
||||
self, data: Mapping[str, object], where: Mapping[str, object], include: Mapping[str, object] | None = None
|
||||
) -> object | None:
|
||||
self.rows = tuple({**row, **data} if _matches(row, where) else row for row in self.rows)
|
||||
return await self.find_first(where=where)
|
||||
|
||||
|
||||
class DefinitionRows:
|
||||
def __init__(self) -> None:
|
||||
self.litellm_workflowdefinition = Rows()
|
||||
self.litellm_workflowsnapshot = Rows()
|
||||
self.litellm_teamtable = Rows(
|
||||
(
|
||||
{
|
||||
"team_id": "team",
|
||||
"members_with_roles": [
|
||||
{"user_id": "author", "role": "user"},
|
||||
{"user_id": "reviewer", "role": "admin"},
|
||||
{"user_id": "other", "role": "user"},
|
||||
],
|
||||
},
|
||||
{"team_id": "other-team", "members_with_roles": [{"user_id": "outsider", "role": "admin"}]},
|
||||
)
|
||||
)
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object:
|
||||
return ()
|
||||
|
||||
@property
|
||||
def database(self) -> DefinitionDatabase:
|
||||
return cast(DefinitionDatabase, self) # cast-ok: isolated Prisma boundary implements the used table operations
|
||||
|
||||
@asynccontextmanager
|
||||
async def tx(self) -> AsyncGenerator["DefinitionRows"]:
|
||||
async with self.lock:
|
||||
yield self
|
||||
118
tests/unit/proxy/workflows/test_definitions.py
Normal file
118
tests/unit/proxy/workflows/test_definitions.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
import hashlib
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from liteflow import ChatCompletion, Input, Python, Workflow, compile_workflow
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.models.workflow_definition import (
|
||||
DefinitionError,
|
||||
PythonNode,
|
||||
WorkflowBundle,
|
||||
WorkflowDraft,
|
||||
WorkflowSnapshot,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.workflows.definitions import authorize_team, freeze_draft, published_workflow, validate_bundle
|
||||
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
|
||||
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
|
||||
|
||||
|
||||
def transform(text: str) -> str:
|
||||
return str(1 // 0)
|
||||
|
||||
|
||||
def _python_bundle(source: str) -> WorkflowBundle:
|
||||
workflow: Final = Workflow(
|
||||
name="example",
|
||||
nodes=[
|
||||
ChatCompletion(model="gateway-model", messages=Input.messages),
|
||||
Python(fn=transform),
|
||||
],
|
||||
)
|
||||
original: Final = WorkflowBundle.model_validate(compile_workflow(workflow))
|
||||
node: Final = original.nodes[-1]
|
||||
assert isinstance(node, PythonNode)
|
||||
program: Final = node.program.model_copy(
|
||||
update={"source": source, "source_sha256": hashlib.sha256(source.encode()).hexdigest()}
|
||||
)
|
||||
return original.model_copy(update={"nodes": (*original.nodes[:-1], node.model_copy(update={"program": program}))})
|
||||
|
||||
|
||||
def test_validation_compiles_python_without_executing_its_body() -> None:
|
||||
assert validate_bundle(_python_bundle("def transform(text: str) -> str:\n return str(1 // 0)\n")) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"source",
|
||||
[
|
||||
"import os\ndef transform(text: str) -> str:\n return text\n",
|
||||
"def transform(text: str) -> str:\n return open(text).read()\n",
|
||||
],
|
||||
)
|
||||
def test_validation_rejects_imports_and_host_capabilities(source: str) -> None:
|
||||
result: Final = validate_bundle(_python_bundle(source))
|
||||
assert isinstance(result, DefinitionError) and result.code == "invalid_bundle"
|
||||
assert result.param == "nodes.node_2.program"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema",
|
||||
[
|
||||
{"type": "object", "properties": {"$ref": {"type": "string"}, "$id": {"type": "number"}}},
|
||||
{"const": {"$ref": "literal"}},
|
||||
{"enum": [{"$dynamicRef": "literal"}]},
|
||||
],
|
||||
)
|
||||
def test_schema_keywords_inside_property_names_and_literals_are_data(schema: Mapping[str, JsonValue]) -> None:
|
||||
assert validate_bundle(bundle().model_copy(update={"input_schema": schema})) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema",
|
||||
[
|
||||
{"properties": {"safe": {"$ref": "https://example.com/schema"}}},
|
||||
{"patternProperties": {"^a": {"$id": "https://example.com/schema"}}},
|
||||
{"$dynamicRef": "https://example.com/schema"},
|
||||
{"$recursiveRef": "https://example.com/schema"},
|
||||
],
|
||||
)
|
||||
def test_reference_scope_is_checked_only_at_schema_positions(schema: Mapping[str, JsonValue]) -> None:
|
||||
assert validate_bundle(bundle().model_copy(update={"input_schema": schema})) == DefinitionError(
|
||||
"invalid_bundle", "Schemas must use local references and bounded nesting", "bundle.input_schema"
|
||||
)
|
||||
|
||||
|
||||
def test_forward_reference_cannot_reach_a_provider() -> None:
|
||||
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
|
||||
forward: Final = WorkflowBundle.model_validate(
|
||||
{**original, "nodes": [{**original["nodes"][0], "messages": {"from": "later", "output": "result"}}]}
|
||||
)
|
||||
assert validate_bundle(forward) == DefinitionError(
|
||||
"invalid_bundle", "Bindings must reference input messages or an earlier node result", "nodes.node_1.messages"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
|
||||
async def test_key_team_restriction_precedes_global_user_role(role: LitellmUserRoles) -> None:
|
||||
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
|
||||
identity: Final = caller("reviewer").model_copy(update={"user_role": role, "via_virtual_key": True})
|
||||
assert await authorize_team(store, identity, "other-team", "read") == DefinitionError(
|
||||
"not_found", "Workflow not found"
|
||||
)
|
||||
session: Final = identity.model_copy(update={"is_session_token": True})
|
||||
assert not isinstance(await authorize_team(store, session, "other-team", "read"), DefinitionError)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_application_identity_can_invoke_but_cannot_freeze_a_draft() -> None:
|
||||
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
|
||||
draft: Final = await store.create(bundle(), "team", "author")
|
||||
assert isinstance(draft, WorkflowDraft)
|
||||
snapshot: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
|
||||
assert isinstance(snapshot, WorkflowSnapshot)
|
||||
assert await published_workflow(store, caller(None), "example") == snapshot
|
||||
denied: Final = await freeze_draft(store, caller(None), draft.workflow_id, draft.etag)
|
||||
assert isinstance(denied, DefinitionError) and denied.code == "not_found"
|
||||
443
tests/unit/proxy/workflows/test_gateway.py
Normal file
443
tests/unit/proxy/workflows/test_gateway.py
Normal file
|
|
@ -0,0 +1,443 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from datetime import datetime, timedelta, timezone, tzinfo
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException, Request
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.models.workflow_definition import ChatCompletionNode
|
||||
from litellm.models.workflow_execution import (
|
||||
ChildCallResult,
|
||||
EngineRun,
|
||||
ExecutionError,
|
||||
RunInput,
|
||||
RunState,
|
||||
StepTrace,
|
||||
workflow_call_id,
|
||||
)
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import cache_key_object
|
||||
from litellm.proxy.auth.user_api_key_auth import authorize_internal_virtual_key
|
||||
from litellm.proxy.common_request_processing import process_workflow_request
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.deps.database import get_database
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3, get_request_stash
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy.workflows import gateway
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_NOW: Final = datetime(2026, 11, 6, tzinfo=timezone.utc)
|
||||
_KEY: Final = "a" * 64
|
||||
_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_MESSAGES: Final = [{"role": "user", "content": "private input"}]
|
||||
|
||||
|
||||
class _Clock(datetime):
|
||||
@classmethod
|
||||
def now(cls, tz: tzinfo | None = None) -> datetime:
|
||||
return _NOW
|
||||
|
||||
|
||||
def _response(content: str) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
model="workflow-model",
|
||||
choices=[{"message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
|
||||
usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
|
||||
)
|
||||
|
||||
|
||||
def _request() -> Request:
|
||||
return Request(
|
||||
{
|
||||
"type": "http",
|
||||
"app": proxy_server.app,
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "https",
|
||||
"path": "/v1/chat/completions",
|
||||
"query_string": b"",
|
||||
"server": ("gateway.example", 443),
|
||||
"client": ("192.0.2.10", 1234),
|
||||
"headers": [
|
||||
(b"x-litellm-trace-id", b"trace-root"),
|
||||
(b"x-litellm-session-id", b"caller-session"),
|
||||
(b"x-forwarded-for", b"198.51.100.8"),
|
||||
(b"content-type", b"application/json"),
|
||||
(b"authorization", b"Bearer caller-secret"),
|
||||
],
|
||||
"state": {"parent": "untouched", "database": None},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _run() -> EngineRun:
|
||||
return EngineRun(
|
||||
run_id="run",
|
||||
session_id="trace-root",
|
||||
workflow_id="review",
|
||||
snapshot_id="snapshot",
|
||||
team_id="team",
|
||||
created_by=_KEY,
|
||||
status=RunState.running,
|
||||
created_at=_NOW,
|
||||
updated_at=_NOW,
|
||||
request_digest="digest",
|
||||
input=RunInput(messages=tuple(_MESSAGES), deadline=_NOW + timedelta(minutes=2), draft=False),
|
||||
)
|
||||
|
||||
|
||||
def _node() -> ChatCompletionNode:
|
||||
return ChatCompletionNode.model_validate(
|
||||
{
|
||||
"id": "answer",
|
||||
"type": "chat_completion",
|
||||
"model": "workflow-model",
|
||||
"messages": {"from": "$input", "output": "messages"},
|
||||
"timeout_seconds": 60,
|
||||
"max_tokens": 3,
|
||||
"retry": {"max_attempts": 1, "backoff_seconds": 0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _step(attempt: int = 1) -> StepTrace:
|
||||
return StepTrace(
|
||||
step_id=f"step-{attempt}",
|
||||
run_id="run",
|
||||
node_id="answer",
|
||||
state=RunState.running,
|
||||
input=_MESSAGES,
|
||||
attempts=attempt,
|
||||
deadline=_NOW + timedelta(minutes=1),
|
||||
started_at=_NOW,
|
||||
call_id=workflow_call_id("run", "answer", attempt),
|
||||
)
|
||||
|
||||
|
||||
class _Observer(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.outcome = "success"
|
||||
self.failure_status = 500
|
||||
self.entered: Final = asyncio.Event()
|
||||
self.failure_entered: Final = asyncio.Event()
|
||||
self.completed: Final[asyncio.Queue[str]] = asyncio.Queue()
|
||||
self.admissions: Final[asyncio.Queue[UserAPIKeyAuth]] = asyncio.Queue()
|
||||
self.rows: Final[asyncio.Queue[Mapping[str, object]]] = asyncio.Queue()
|
||||
self.requests: Final[asyncio.Queue[Mapping[str, object]]] = asyncio.Queue()
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: Mapping[str, object], call_type: str
|
||||
) -> Mapping[str, object] | None:
|
||||
self.admissions.put_nowait(user_api_key_dict)
|
||||
self.requests.put_nowait(_FIELDS.validate_python(data["proxy_server_request"]))
|
||||
if self.outcome == "pre":
|
||||
raise HTTPException(self.failure_status, "private processor failure")
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, response: object
|
||||
) -> object:
|
||||
self.completed.put_nowait(str(data["model"]))
|
||||
if data["model"] == "workflow-model":
|
||||
self.entered.set()
|
||||
if self.outcome == "post":
|
||||
raise HTTPException(self.failure_status, "private post-call failure")
|
||||
if self.outcome in {"cancel", "timeout"}:
|
||||
await asyncio.Event().wait()
|
||||
if self.outcome in {"reject", "failure_timeout"}:
|
||||
raise HTTPException(422, "private guardrail detail")
|
||||
return response
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: Mapping[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
self.failure_entered.set()
|
||||
if self.outcome == "failure_timeout":
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.rows.put_nowait(
|
||||
get_logging_payload(
|
||||
kwargs=dict(kwargs),
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
await self.async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
|
||||
class _WorkflowPolicy(CustomGuardrail):
|
||||
def __init__(self, reject: bool = False) -> None:
|
||||
super().__init__(
|
||||
guardrail_name="workflow-policy",
|
||||
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
|
||||
)
|
||||
self.reject: Final = reject
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: Mapping[str, object], call_type: str
|
||||
) -> Mapping[str, object] | None:
|
||||
return {**data, "messages": [{"role": "user", "content": "guarded input"}]}
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, response: object
|
||||
) -> object:
|
||||
if self.reject:
|
||||
raise HTTPException(422, "private Python output")
|
||||
assert isinstance(response, ModelResponse)
|
||||
return response.model_copy(update={"choices": _response("guarded Python output").choices})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def proxy(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[tuple[_Observer, UserAPIKeyAuth]]:
|
||||
cache: Final = UserApiKeyCache()
|
||||
key: Final = UserAPIKeyAuth(
|
||||
token=_KEY,
|
||||
team_id="team",
|
||||
max_budget=100.0,
|
||||
max_parallel_requests=1,
|
||||
models=["workflow-model", "workflow/review"],
|
||||
)
|
||||
await cache_key_object(hashed_token=_KEY, user_api_key_obj=key, user_api_key_cache=cache, proxy_logging_obj=None)
|
||||
await cache.async_set_cache(
|
||||
key="team_id:team",
|
||||
value=LiteLLM_TeamTableCachedObj(
|
||||
team_id="team",
|
||||
metadata={"require_trace_id": True},
|
||||
),
|
||||
)
|
||||
observer: Final = _Observer()
|
||||
logging: Final = ProxyLogging(user_api_key_cache=cache)
|
||||
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(logging.internal_usage_cache)
|
||||
counters: Final = DualCache()
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "workflow-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/workflow-fixture",
|
||||
"api_key": "unit-only",
|
||||
"mock_response": _response("child answer"),
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
monkeypatch.setattr(gateway, "datetime", _Clock)
|
||||
monkeypatch.setattr(litellm, "callbacks", [observer, limiter])
|
||||
monkeypatch.setattr(litellm, "max_budget", 0.0)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_cache", counters)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace())
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"allowed_ips": ["192.0.2.10"], "cancel_on_disconnect": True})
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig())
|
||||
yield observer, key
|
||||
get_policy_registry().clear()
|
||||
get_attachment_registry().clear()
|
||||
router.reset()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("forwarded", [False, True])
|
||||
async def test_child_preserves_trusted_context_and_rechecks_model_access(
|
||||
proxy: tuple[_Observer, UserAPIKeyAuth], forwarded: bool
|
||||
) -> None:
|
||||
observer, key = proxy
|
||||
proxy_server.general_settings.update(
|
||||
{
|
||||
"use_x_forwarded_for": forwarded,
|
||||
"allowed_ips": ["198.51.100.8" if forwarded else "192.0.2.10"],
|
||||
}
|
||||
)
|
||||
request: Final = _request()
|
||||
child: Final = gateway._child_request(request, {}, workflow_call_id("run", "answer", 1), "trace-root")
|
||||
assert await get_database(child) is None
|
||||
assert "authorization" not in child.headers
|
||||
result: Final = await gateway.execute_chat_node(request, _run(), _node(), _step())
|
||||
assert isinstance(result, ChildCallResult), result
|
||||
row: Final = await asyncio.wait_for(observer.rows.get(), 5)
|
||||
seen: Final = observer.requests.get_nowait()
|
||||
assert _FIELDS.validate_python(seen["headers"])["x-litellm-trace-id"] == "trace-root"
|
||||
assert (row["session_id"], row["litellm_call_id"], row["api_key"]) == ("trace-root", _step().call_id, _KEY)
|
||||
assert result.request_id is None
|
||||
assert result.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
|
||||
assert request.scope["state"] == {"parent": "untouched", "database": None}
|
||||
assert "caller-secret" not in str(seen)
|
||||
await cache_key_object(
|
||||
hashed_token=_KEY,
|
||||
user_api_key_obj=key.model_copy(update={"models": ["another-model"]}),
|
||||
user_api_key_cache=proxy_server.user_api_key_cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
denied: Final = await gateway.execute_chat_node(request, _run(), _node(), _step(2))
|
||||
assert isinstance(denied, ExecutionError), denied
|
||||
assert denied.receipt is None and observer.requests.empty()
|
||||
assert "private" not in denied.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("phase", "status", "code"),
|
||||
[
|
||||
("auth", 429, "admission_rate_limit"),
|
||||
("auth", 500, "provider_error"),
|
||||
("pre", 429, "rate_limit"),
|
||||
("pre", 500, "provider_error"),
|
||||
("post", 429, "rate_limit"),
|
||||
("post", 500, "provider_error"),
|
||||
],
|
||||
)
|
||||
async def test_only_auth_rate_limits_are_safe_to_retry(
|
||||
proxy: tuple[_Observer, UserAPIKeyAuth], phase: str, status: int, code: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
observer, _ = proxy
|
||||
observer.outcome = phase
|
||||
observer.failure_status = status
|
||||
|
||||
async def reject_auth(
|
||||
key_hash: str,
|
||||
request: Request,
|
||||
request_data: Mapping[str, object],
|
||||
*,
|
||||
route: str | None = None,
|
||||
expected_team_id: str | None = None,
|
||||
) -> UserAPIKeyAuth:
|
||||
raise HTTPException(status, "private authentication failure")
|
||||
|
||||
if phase == "auth":
|
||||
monkeypatch.setattr(gateway, "authorize_internal_virtual_key", reject_auth)
|
||||
result: Final = await gateway.execute_chat_node(_request(), _run(), _node(), _step())
|
||||
assert isinstance(result, ExecutionError), result
|
||||
assert result.code == code and "private" not in result.message
|
||||
assert observer.admissions.qsize() == (0 if phase == "auth" else 1)
|
||||
if phase == "post":
|
||||
assert result.receipt is not None
|
||||
assert result.receipt.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
|
||||
await asyncio.wait_for(observer.rows.get(), 5)
|
||||
else:
|
||||
assert result.receipt is None
|
||||
if phase == "auth":
|
||||
assert observer.requests.empty() and observer.rows.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("outcome", ["reject", "cancel", "timeout", "failure_timeout"])
|
||||
async def test_billed_failure_and_cancellation_release_child_slot(
|
||||
proxy: tuple[_Observer, UserAPIKeyAuth], outcome: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
observer, _ = proxy
|
||||
observer.outcome = outcome
|
||||
task: Final = asyncio.create_task(gateway.execute_chat_node(_request(), _run(), _node(), _step()))
|
||||
await asyncio.wait_for((observer.failure_entered if outcome == "failure_timeout" else observer.entered).wait(), 5)
|
||||
if outcome == "cancel":
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
if outcome in {"timeout", "failure_timeout"}:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
original_time: Final = loop.time
|
||||
monkeypatch.setattr(loop, "time", lambda: original_time() + 120)
|
||||
failure: Final = await task
|
||||
assert isinstance(failure, ExecutionError) and failure.receipt is not None, failure
|
||||
assert failure.receipt.response_cost == pytest.approx(7 * 0.001 + 3 * 0.002)
|
||||
assert failure.message == (
|
||||
"The node deadline expired"
|
||||
if outcome in {"timeout", "failure_timeout"}
|
||||
else "The gateway rejected this node"
|
||||
)
|
||||
await asyncio.wait_for(observer.rows.get(), 5)
|
||||
admitted: Final = observer.admissions.get_nowait()
|
||||
if outcome in {"cancel", "timeout", "failure_timeout"}:
|
||||
assert admitted.budget_reservation is not None and admitted.budget_reservation["finalized"] is True
|
||||
observer.outcome = "success"
|
||||
retried: Final = await gateway.execute_chat_node(_request(), _run(), _node(), _step(2))
|
||||
assert isinstance(retried, ChildCallResult), retried
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("detached", [False, True])
|
||||
@pytest.mark.parametrize("reject", [False, True])
|
||||
async def test_workflow_policy_guards_modified_input_and_final_python_output_without_wrapper_charge(
|
||||
proxy: tuple[_Observer, UserAPIKeyAuth],
|
||||
detached: bool,
|
||||
reject: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
observer, key = proxy
|
||||
policy: Final = _WorkflowPolicy(reject)
|
||||
monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, policy])
|
||||
get_policy_registry().add_policy("workflow-dlp", Policy(guardrails=PolicyGuardrails(add=["workflow-policy"])))
|
||||
get_attachment_registry().add_attachment(PolicyAttachment(policy="workflow-dlp", models=["workflow/review"]))
|
||||
parent_request: Final = _request()
|
||||
parent_auth: Final = await authorize_internal_virtual_key(
|
||||
key_hash=_KEY,
|
||||
request=parent_request,
|
||||
request_data={"model": "workflow-model", "messages": _MESSAGES, "max_tokens": 3},
|
||||
expected_team_id="team",
|
||||
)
|
||||
assert parent_auth.budget_reservation is not None
|
||||
inputs: Final[asyncio.Queue[JsonValue]] = asyncio.Queue()
|
||||
|
||||
async def execute(messages: JsonValue) -> ModelResponse | ExecutionError:
|
||||
inputs.put_nowait(messages)
|
||||
assert get_request_stash() is None
|
||||
assert parent_auth.budget_reservation["finalized"] is True
|
||||
assert proxy_server.spend_counter_cache.in_memory_cache.get_cache(key=f"spend:key:{_KEY}") == 0
|
||||
child: Final = await gateway.execute_chat_node(
|
||||
_request(),
|
||||
_run(),
|
||||
_node(),
|
||||
_step().model_copy(update={"input": messages}),
|
||||
)
|
||||
assert isinstance(child, ChildCallResult), child
|
||||
return _response("private Python output")
|
||||
|
||||
invocation: Final = process_workflow_request(
|
||||
request=parent_request,
|
||||
workflow_model="workflow/review",
|
||||
caller=parent_auth,
|
||||
messages=_MESSAGES,
|
||||
execute=execute,
|
||||
)
|
||||
result: Final = await (asyncio.create_task(invocation) if detached else invocation)
|
||||
assert inputs.get_nowait() == [{"role": "user", "content": "guarded input"}]
|
||||
if reject:
|
||||
assert isinstance(result, ExecutionError), result
|
||||
assert "private" not in result.message
|
||||
else:
|
||||
assert isinstance(result, ModelResponse), result
|
||||
assert result.choices[0].message.content == "guarded Python output"
|
||||
row: Final = await asyncio.wait_for(observer.rows.get(), 5)
|
||||
assert row["litellm_call_id"] == _step().call_id
|
||||
assert observer.rows.empty()
|
||||
assert observer.admissions.qsize() == 1
|
||||
assert observer.completed.get_nowait() == "workflow-model" and observer.completed.empty()
|
||||
33
tests/unit/proxy/workflows/test_runner.py
Normal file
33
tests/unit/proxy/workflows/test_runner.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.models.workflow_definition import ReferenceBinding, ValueBinding, WorkflowBundle, WorkflowRetry
|
||||
from litellm.models.workflow_execution import ExecutionError, RunState, StepTrace
|
||||
from litellm.proxy.workflows.runner import resolve_binding
|
||||
from tests.unit.proxy.workflow_definition_rows import bundle
|
||||
from tests.unit.repositories.test_workflow_execution_repository import NOW, pending_run
|
||||
|
||||
|
||||
def two_nodes() -> WorkflowBundle:
|
||||
first: Final = bundle().nodes[0].model_copy(update={"retry": WorkflowRetry(max_attempts=2)})
|
||||
second: Final = first.model_copy(update={
|
||||
"id": "node_2", "messages": ReferenceBinding(source=first.id, output="result"),
|
||||
})
|
||||
return bundle().model_copy(update={
|
||||
"nodes": (first, second), "output": ReferenceBinding(source=second.id, output="result"),
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [None, False, 0, [], {}, [{"role": "user", "content": "next"}]])
|
||||
def test_bindings_preserve_falsy_and_structured_values(value: JsonValue) -> None:
|
||||
run: Final = pending_run()
|
||||
step: Final = StepTrace(step_id="step", run_id=run.run_id, node_id="node", state=RunState.completed,
|
||||
input=None, output=value, attempts=1, started_at=NOW, deadline=NOW)
|
||||
assert resolve_binding(ValueBinding(value=value), run, ()) == value
|
||||
assert resolve_binding(ReferenceBinding(source="node", output="result"), run, (step,)) == value
|
||||
assert resolve_binding(ReferenceBinding(source="$input", output="messages"), run, ()) == list(run.input.messages)
|
||||
assert resolve_binding(ReferenceBinding(source="missing", output="result"), run, (step,)) == ExecutionError(
|
||||
"missing_input", "Node output is unavailable"
|
||||
)
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from liteflow import Draft, bundle_digest
|
||||
|
||||
from litellm.models.workflow_definition import DefinitionError, WorkflowSnapshot, parsed_bundle
|
||||
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
|
||||
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publication_and_draft_tests_keep_immutable_snapshots_across_edits() -> None:
|
||||
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
|
||||
draft: Final = await store.create(bundle(), "team", "author")
|
||||
assert isinstance(draft, Draft) and draft.digest == bundle_digest(draft.bundle)
|
||||
published: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
|
||||
frozen: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag)
|
||||
assert isinstance(published, WorkflowSnapshot) and isinstance(frozen, WorkflowSnapshot)
|
||||
assert published.bundle == frozen.bundle == parsed_bundle(draft.bundle)
|
||||
assert frozen.published_by is None
|
||||
assert await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer") == published
|
||||
updated: Final = await store.replace(draft.workflow_id, "team", draft.etag, bundle(model="second"), "other")
|
||||
assert isinstance(updated, Draft) and updated.published_snapshot_id == published.snapshot_id
|
||||
assert await store.get_published("example", "team") == published
|
||||
stale: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag)
|
||||
assert isinstance(stale, DefinitionError) and stale.code == "stale_draft"
|
||||
second: Final = await store.publish(updated.workflow_id, "team", updated.etag, updated.digest, "reviewer")
|
||||
assert isinstance(second, WorkflowSnapshot) and second.contributors == ("other",)
|
||||
assert second.snapshot_id != published.snapshot_id and second.digest != published.digest
|
||||
assert await store.get_snapshot(published.snapshot_id) == published
|
||||
assert await store.get_snapshot(frozen.snapshot_id) == frozen
|
||||
assert await store.get_snapshot(second.snapshot_id, "other-team") is None
|
||||
assert await store.get_published("example", "team") == second
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.models.workflow_execution import EngineRun, RunInput, RunState, StepTrace
|
||||
from litellm.repositories.workflow_execution_repository import _aggregate_attempts
|
||||
|
||||
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def pending_run() -> EngineRun:
|
||||
return EngineRun(
|
||||
run_id="run", session_id="session", workflow_id="workflow", snapshot_id="snapshot", team_id="team",
|
||||
created_by="owner", status=RunState.pending, created_at=NOW, updated_at=NOW, request_digest="digest",
|
||||
input=RunInput(
|
||||
messages=({"role": "user", "content": "plan"},), deadline=NOW + timedelta(minutes=5), draft=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cost", [None, 0.0, 0.125])
|
||||
def test_attempt_aggregation_keeps_unknown_cost_and_unique_receipts(cost: float | None) -> None:
|
||||
first: Final = StepTrace(
|
||||
step_id="step", run_id="run", node_id="node", state=RunState.failed, input=None, attempts=1,
|
||||
deadline=NOW, started_at=NOW, prompt_tokens=4, completion_tokens=2, response_cost=cost,
|
||||
request_ids=("fallback-a", "fallback-b"),
|
||||
)
|
||||
latest: Final = first.model_copy(update={"attempts": 2, "state": RunState.completed, "output": False,
|
||||
"request_ids": ("fallback-b", "retry"), "response_cost": None})
|
||||
assert _aggregate_attempts((latest, first)) == latest.model_copy(update={
|
||||
"prompt_tokens": 8, "completion_tokens": 4, "response_cost": cost,
|
||||
"request_ids": ("fallback-b", "retry", "fallback-a"),
|
||||
})
|
||||
|
|
@ -1,11 +1,11 @@
|
|||
"""Tests for scripts/check_openapi_docs.py, the scan behind scripts/openapi_docs_gate.py."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import check_openapi_docs as checker
|
||||
import pytest
|
||||
|
||||
_BODY: Final = {"content": {"application/json": {"schema": {"type": "object"}}}}
|
||||
_RESPONSES: Final = {"200": {"content": {"application/json": {"schema": {"$ref": "#/components/schemas/Out"}}}}}
|
||||
|
|
@ -20,15 +20,47 @@ def _post_with_media_types(request_media_type: str, response_media_type: str) ->
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_type, response_type", [
|
||||
("application/json", "application/json"), ("application/sdp", "application/sdp"),
|
||||
("multipart/form-data", "application/json"),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"request_type, response_type",
|
||||
[
|
||||
("application/json", "application/json"),
|
||||
("application/sdp", "application/sdp"),
|
||||
("multipart/form-data", "application/json"),
|
||||
],
|
||||
)
|
||||
def test_a_fully_documented_post_has_no_findings(request_type: str, response_type: str) -> None:
|
||||
operation: Final = _post_with_media_types(request_type, response_type)
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["200", "201", "202", "206", "299"])
|
||||
def test_any_documented_success_response_satisfies_the_response_rule(status: str) -> None:
|
||||
operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}}
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["default", "100", "300", "400", "500", "2000"])
|
||||
def test_only_success_statuses_can_supply_the_response_schema(status: str) -> None:
|
||||
operation: Final = {**_DOCUMENTED, "responses": {status: _RESPONSES["200"]}}
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == (
|
||||
checker.Violation("POST", "/v1/thing", "OA003"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["201", "202"])
|
||||
def test_success_without_a_schema_is_not_documented_by_an_error_response(status: str) -> None:
|
||||
operation: Final[Mapping[str, object]] = {
|
||||
**_DOCUMENTED,
|
||||
"responses": {
|
||||
status: {"content": {"application/json": {"schema": {}}}},
|
||||
"422": _RESPONSES["200"],
|
||||
},
|
||||
}
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == (
|
||||
checker.Violation("POST", "/v1/thing", "OA003"),
|
||||
)
|
||||
|
||||
|
||||
def test_a_schema_under_another_media_type_still_counts_as_undocumented() -> None:
|
||||
operation: Final = _post_with_media_types("application/octet-stream", "audio/wav")
|
||||
assert {v.code for v in checker.violations({"paths": {"/v1/thing": {"post": operation}}})} == {"OA001", "OA003"}
|
||||
|
|
@ -44,8 +76,11 @@ def test_each_missing_piece_is_its_own_rule() -> None:
|
|||
def test_only_the_missing_piece_is_reported(body: Mapping[str, object]) -> None:
|
||||
no_body: Final = {**_DOCUMENTED, "requestBody": body}
|
||||
blank_description: Final = {**_DOCUMENTED, "description": " "}
|
||||
untyped_200: Final = {**_DOCUMENTED, "responses": {"200": {"content": {"application/json": {"schema": {}}}}}}
|
||||
spec: Final = {
|
||||
untyped_200: Final[Mapping[str, object]] = {
|
||||
**_DOCUMENTED,
|
||||
"responses": {"200": {"content": {"application/json": {"schema": {}}}}},
|
||||
}
|
||||
spec: Final[Mapping[str, object]] = {
|
||||
"paths": {
|
||||
"/a": {"post": no_body},
|
||||
"/b": {"post": blank_description},
|
||||
|
|
@ -65,11 +100,13 @@ def test_non_post_operations_are_not_gated() -> None:
|
|||
|
||||
|
||||
def test_findings_are_sorted_by_path_so_counts_are_stable() -> None:
|
||||
spec: Final = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}}
|
||||
spec: Final[Mapping[str, object]] = {"paths": {"/z": {"post": {}}, "/a": {"post": {}}}}
|
||||
assert [v.path for v in checker.violations(spec)] == ["/a", "/a", "/a", "/z", "/z", "/z"]
|
||||
|
||||
|
||||
def test_main_reads_a_spec_file_and_prints_one_line_per_finding(tmp_path, capsys) -> None:
|
||||
def test_main_reads_a_spec_file_and_prints_one_line_per_finding(
|
||||
tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
spec_file: Final = tmp_path / "openapi.json"
|
||||
spec_file.write_text('{"paths": {"/v1/thing": {"post": {"description": "x", "requestBody": {}}}}}')
|
||||
assert checker.main([str(spec_file)]) == 0
|
||||
|
|
|
|||
927
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
927
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
File diff suppressed because it is too large
Load diff
26
uv.lock
generated
26
uv.lock
generated
|
|
@ -4996,6 +4996,28 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/38/a6/800800bfed7b1fb10fc3f3d557785c3854e80d3f7a9800d784b176a1fc2d/librt-0.15.0-cp315-cp315t-win_arm64.whl", hash = "sha256:84d244b00604d17df3fc7736c327892d6bba66181254aa4087be807b6c342bdc", size = 110700, upload-time = "2026-08-07T10:49:15.499Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "liteflow"
|
||||
version = "0.1.0a1"
|
||||
source = { url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
{ name = "openai" },
|
||||
{ name = "pydantic", version = "2.13.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" },
|
||||
{ name = "pydantic", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" },
|
||||
]
|
||||
sdist = { hash = "sha256:c54fc9065d3a8462b6d6a4c3487d7ae3dd8a087a0aa8a51a46c5e62fc3e73072" }
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "httpx", specifier = ">=0.27,<1" },
|
||||
{ name = "openai", specifier = ">=1.99,<3" },
|
||||
{ name = "pydantic", specifier = ">=2.8,<3" },
|
||||
{ name = "pytest", marker = "extra == 'test'", specifier = ">=8,<9" },
|
||||
{ name = "rich", marker = "extra == 'cli'", specifier = ">=13,<15" },
|
||||
]
|
||||
provides-extras = ["test", "cli"]
|
||||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.109.0"
|
||||
|
|
@ -5083,6 +5105,7 @@ proxy = [
|
|||
{ name = "hiredis", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" },
|
||||
{ name = "httpx2" },
|
||||
{ name = "inquirerpy" },
|
||||
{ name = "liteflow" },
|
||||
{ name = "litellm-enterprise" },
|
||||
{ name = "litellm-proxy-extras" },
|
||||
{ name = "mcp" },
|
||||
|
|
@ -5254,6 +5277,7 @@ proxy-dev = [
|
|||
{ name = "backports-zstd", marker = "python_full_version < '3.14'" },
|
||||
{ name = "httptools" },
|
||||
{ name = "hypercorn" },
|
||||
{ name = "liteflow" },
|
||||
{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "opentelemetry-api" },
|
||||
{ name = "opentelemetry-exporter-otlp" },
|
||||
|
|
@ -5323,6 +5347,7 @@ requires-dist = [
|
|||
{ name = "jsonschema", specifier = ">=4.0.0,<5.0" },
|
||||
{ name = "keyring", marker = "extra == 'cli'", specifier = ">=25.6.0,<26.0" },
|
||||
{ name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=4.7,<5.0" },
|
||||
{ name = "liteflow", marker = "extra == 'proxy'", url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" },
|
||||
{ name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" },
|
||||
{ name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" },
|
||||
{ name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" },
|
||||
|
|
@ -5490,6 +5515,7 @@ proxy-dev = [
|
|||
{ name = "backports-zstd", marker = "python_full_version < '3.14'", specifier = "==1.8.0" },
|
||||
{ name = "httptools", specifier = ">=0.7.1,<0.8" },
|
||||
{ name = "hypercorn", specifier = "==0.17.3" },
|
||||
{ name = "liteflow", url = "https://github.com/BerriAI/liteflow/archive/27bcbce4e84fb73b1d4bdf2fa108df53e9b50c2d.tar.gz" },
|
||||
{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'", url = "https://github.com/BerriAI/liteadmin-mcp/archive/31fb096904ae05684c0779110a93b135b5601182.tar.gz" },
|
||||
{ name = "opentelemetry-api", specifier = "==1.45.1" },
|
||||
{ name = "opentelemetry-exporter-otlp", specifier = "==1.45.1" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue