diff --git a/litellm/proxy/workflows/child.py b/litellm/proxy/workflows/child.py index 3b9a0f3ac90..b4af50110cb 100644 --- a/litellm/proxy/workflows/child.py +++ b/litellm/proxy/workflows/child.py @@ -69,6 +69,7 @@ class GatewayChildCaller: param=node.id, ) original_request: Final = _OBJECT.validate_json(execution.request_json) + metadata: Final = original_request.get("metadata") payload: Final = { **node.model_dump( mode="json", @@ -82,6 +83,7 @@ class GatewayChildCaller: "timeout": remaining, **({"user": original_request["user"]} if "user" in original_request else {}), "metadata": { + **({"tags": metadata["tags"]} if isinstance(metadata, dict) and "tags" in metadata else {}), "workflow_run_id": execution.run_id, "workflow_node_id": node.id, "workflow_snapshot_id": execution.snapshot_id, diff --git a/litellm/proxy/workflows/validation.py b/litellm/proxy/workflows/validation.py index 01cac0fcb43..c94740246a1 100644 --- a/litellm/proxy/workflows/validation.py +++ b/litellm/proxy/workflows/validation.py @@ -1,10 +1,11 @@ import json -from collections.abc import Mapping -from itertools import chain +from collections.abc import Callable, Mapping +from functools import cache +from itertools import chain, count from typing import Final from jsonschema import Draft202012Validator, SchemaError -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm.models.workflow_bundle import ( Binding, @@ -21,6 +22,7 @@ from litellm.proxy.workflows.monty.compiler import compile_program from litellm.proxy.workflows.monty.models import MontyFailure, MontyLimits DEFAULT_MONTY_LIMITS: Final = MontyLimits() +_JSON: Final = TypeAdapter[JsonValue](JsonValue) def schema_is_local(value: JsonValue, depth: int = 0) -> bool: @@ -82,15 +84,32 @@ def validate_bundle(bundle: WorkflowBundle) -> WorkflowFailure | None: def _types( - schema: JsonValue | Mapping[str, JsonValue], root: Mapping[str, JsonValue], seen: frozenset[str] = frozenset() -) -> frozenset[str]: + schema: JsonValue | Mapping[str, JsonValue], root: Mapping[str, JsonValue] +) -> frozenset[str] | WorkflowFailure: + work: Final = count() + + @cache + def visit(encoded: str, seen: frozenset[str]) -> frozenset[str] | WorkflowFailure: + if next(work) >= 4096 or len(seen) > 48: + return WorkflowFailure(code="schema_complexity", message="Schema compatibility traversal limit exceeded") + return _schema_types(_JSON.validate_json(encoded), root, seen, visit) + + return visit(json.dumps(schema, sort_keys=True), frozenset[str]()) + + +def _schema_types( + schema: JsonValue, + root: Mapping[str, JsonValue], + seen: frozenset[str], + visit: Callable[[str, frozenset[str]], frozenset[str] | WorkflowFailure], +) -> frozenset[str] | WorkflowFailure: if not isinstance(schema, Mapping): return frozenset() reference: Final = schema.get("$ref") if isinstance(reference, str) and reference.startswith("#/$defs/") and reference not in seen: definitions: Final = root.get("$defs") return ( - _types(definitions.get(reference[8:]), root, seen | {reference}) + visit(json.dumps(definitions.get(reference[8:]), sort_keys=True), seen | {reference}) if isinstance(definitions, dict) else frozenset() ) @@ -99,7 +118,11 @@ def _types( return frozenset({type_}) alternatives: Final = schema.get("anyOf", schema.get("oneOf")) if isinstance(alternatives, list): - return frozenset(chain.from_iterable(_types(alternative, root, seen) for alternative in alternatives)) + results: Final = tuple(visit(json.dumps(alternative, sort_keys=True), seen) for alternative in alternatives) + failures: Final = tuple(result for result in results if isinstance(result, WorkflowFailure)) + if failures: + return failures[0] + return frozenset(chain.from_iterable(result for result in results if isinstance(result, frozenset))) return frozenset() @@ -122,6 +145,10 @@ def _binding_compatibility( actual: Final = _result_schema(bundle, binding.source) actual_types: Final = _types(actual, actual) expected_types: Final = _types(expected, expected_root) + if isinstance(actual_types, WorkflowFailure): + return actual_types + if isinstance(expected_types, WorkflowFailure): + return expected_types if ( actual_types and expected_types diff --git a/litellm/repositories/workflow_repository.py b/litellm/repositories/workflow_repository.py index c6e94022987..b7db52a07d9 100644 --- a/litellm/repositories/workflow_repository.py +++ b/litellm/repositories/workflow_repository.py @@ -1,15 +1,16 @@ import json -from collections.abc import Sequence +from collections.abc import Mapping from contextlib import AbstractAsyncContextManager -from datetime import datetime +from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, Final, Protocol, - TypeVar, - cast, # noqa: TID251 # PrismaWrapper forwards transaction operations dynamically + cast, # noqa: TID251 # PrismaWrapper exposes generated delegates dynamically ) +from uuid import NAMESPACE_URL, uuid5 +from prisma import Json from pydantic import JsonValue, TypeAdapter from litellm._uuid import uuid @@ -22,376 +23,362 @@ from litellm.models.workflow import ( WorkflowSnapshot, WorkflowStep, ) +from litellm.repositories.prisma_protocols import TableActions if TYPE_CHECKING: + from prisma import models as prisma_models + from litellm.proxy.utils import PrismaClient -class WorkflowTransaction(Protocol): - async def query_raw(self, query: str, *args: object) -> object: ... +_JSON: Final = TypeAdapter[JsonValue](JsonValue) - async def execute_raw(self, query: str, *args: object) -> int: ... + +def _json(encoded: str) -> Json: + return Json(_JSON.validate_json(encoded)) + + +class WorkflowTransaction(Protocol): + @property + def litellm_workflowdefinition(self) -> TableActions["prisma_models.LiteLLM_WorkflowDefinition"]: ... + + @property + def litellm_workflowsnapshot(self) -> TableActions["prisma_models.LiteLLM_WorkflowSnapshot"]: ... + + @property + def litellm_workflowexecution(self) -> TableActions["prisma_models.LiteLLM_WorkflowExecution"]: ... + + @property + def litellm_workflowstep(self) -> TableActions["prisma_models.LiteLLM_WorkflowStep"]: ... + + @property + def litellm_workflowrun(self) -> TableActions["prisma_models.LiteLLM_WorkflowRun"]: ... + + @property + def litellm_workflowevent(self) -> TableActions["prisma_models.LiteLLM_WorkflowEvent"]: ... class WorkflowDatabase(WorkflowTransaction, Protocol): def tx(self) -> AbstractAsyncContextManager[WorkflowTransaction]: ... -_DRAFTS: Final = TypeAdapter(tuple[WorkflowDraft, ...]) -_SNAPSHOTS: Final = TypeAdapter(tuple[WorkflowSnapshot, ...]) -_EXECUTIONS: Final = TypeAdapter(tuple[WorkflowExecution, ...]) -_CHECKPOINTS: Final = TypeAdapter(tuple[WorkflowCheckpoint, ...]) -_STEPS: Final = TypeAdapter(tuple[WorkflowStep, ...]) -_DRAFT_COLUMNS: Final = "workflow_id, team_id, name, etag, bundle_json, contributors, published_snapshot_id" -_SNAPSHOT_COLUMNS: Final = "snapshot_id, workflow_id, team_id, name, bundle_json, published_by, created_at, origin" -_EXECUTION_COLUMNS: Final = ( - "run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, state, " - "deadline, lease_owner, lease_until, cancel_requested, checkpoints, response_json, error" -) -_STEP_COLUMNS: Final = ( - "step_id, run_id, node_id, state, input, output, attempts, deadline, started_at, finished_at, error, " - "prompt_tokens, completion_tokens, response_cost, request_id" -) +def _lease_where(run_id: str, owner: str, *, allow_cancel: bool = False) -> Mapping[str, object]: + return { + "run_id": run_id, + "lease_owner": owner, + "lease_until": {"gt": datetime.now(timezone.utc)}, + "state": "running", + **({} if allow_cancel else {"cancel_requested": False}), + } -_RowT: Final = TypeVar("_RowT") +async def _lock_execution( + database: WorkflowTransaction, run_id: str, owner: str, *, allow_cancel: bool = False +) -> bool: + return bool( + await database.litellm_workflowexecution.update_many( + where=_lease_where(run_id, owner, allow_cancel=allow_cancel), data={"lease_owner": owner} + ) + ) -def _first(rows: Sequence[_RowT]) -> _RowT | None: - return rows[0] if rows else None +def _step_data(step: WorkflowStep) -> Mapping[str, object]: + return { + "state": step.state.value, + "output": _json(json.dumps(step.output, allow_nan=False)), + "finished_at": step.finished_at, + "error": step.error, + "prompt_tokens": step.prompt_tokens, + "completion_tokens": step.completion_tokens, + "response_cost": step.response_cost, + "request_id": step.request_id, + } class WorkflowRepository: def __init__(self, prisma_client: "PrismaClient") -> None: - self._database: Final = cast( # cast-ok: Prisma exposes raw SQL and transaction methods dynamically + self._database: Final = cast( # cast-ok: Prisma exposes generated model delegates and transactions dynamically WorkflowDatabase, prisma_client.writer_db ) async def create(self, team_id: str, name: str, bundle_json: str, author: str) -> WorkflowDraft | WorkflowConflict: - rows: Final = _DRAFTS.validate_python( - await self._database.query_raw( - 'INSERT INTO "LiteLLM_WorkflowDefinition" ' - "(workflow_id, team_id, name, etag, bundle_json, contributors) " - "VALUES ($1, $2, $3, $4, $5, ARRAY[$6]::text[]) ON CONFLICT (team_id, name) DO NOTHING " - f"RETURNING {_DRAFT_COLUMNS}", - str(uuid.uuid4()), - team_id, - name, - str(uuid.uuid4()), - bundle_json, - author, - ) + workflow_id: Final = str(uuid.uuid4()) + inserted: Final = await self._database.litellm_workflowdefinition.create_many( + data=[ + { + "workflow_id": workflow_id, + "team_id": team_id, + "name": name, + "etag": str(uuid.uuid4()), + "bundle_json": bundle_json, + "contributors": [author], + } + ], + skip_duplicates=True, ) - return rows[0] if rows else WorkflowConflict(message="Workflow already exists") + if not inserted: + return WorkflowConflict(message="Workflow already exists") + row: Final = await self._database.litellm_workflowdefinition.find_unique(where={"workflow_id": workflow_id}) + return WorkflowDraft.model_validate(row, from_attributes=True) async def get(self, team_id: str, name: str) -> WorkflowDraft | None: - return _first( - _DRAFTS.validate_python( - await self._database.query_raw( - f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2', - team_id, - name, - ) - ) + row: Final = await self._database.litellm_workflowdefinition.find_unique( + where={"team_id_name": {"team_id": team_id, "name": name}} ) + return WorkflowDraft.model_validate(row, from_attributes=True) if row is not None else None + + async def get_id(self, workflow_id: str) -> WorkflowDraft | None: + row: Final = await self._database.litellm_workflowdefinition.find_unique(where={"workflow_id": workflow_id}) + return WorkflowDraft.model_validate(row, from_attributes=True) if row is not None else None + + async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]: + rows: Final = await self._database.litellm_workflowdefinition.find_many( + where={"team_id": team_id, "workflow_id": {"gt": after or ""}}, + order={"workflow_id": "asc"}, + take=limit, + ) + return tuple(WorkflowDraft.model_validate(row, from_attributes=True) for row in rows) async def replace( self, team_id: str, name: str, etag: str, bundle_json: str, author: str ) -> WorkflowDraft | WorkflowConflict: - rows: Final = _DRAFTS.validate_python( - await self._database.query_raw( - 'UPDATE "LiteLLM_WorkflowDefinition" SET etag = $4, bundle_json = $5, ' - "contributors = ARRAY(SELECT DISTINCT unnest(contributors || ARRAY[$6]::text[])) " - f"WHERE team_id = $1 AND name = $2 AND etag = $3 RETURNING {_DRAFT_COLUMNS}", - team_id, - name, - etag, - str(uuid.uuid4()), - bundle_json, - author, + async with self._database.tx() as tx: + previous: Final = await tx.litellm_workflowdefinition.find_first( + where={"team_id": team_id, "name": name, "etag": etag} ) - ) - return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412) - - async def get_id(self, workflow_id: str) -> WorkflowDraft | None: - return _first( - _DRAFTS.validate_python( - await self._database.query_raw( - f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1', - workflow_id, - ) + if previous is None: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + updated: Final = await tx.litellm_workflowdefinition.update_many( + where={"workflow_id": previous.workflow_id, "etag": etag}, + data={ + "etag": str(uuid.uuid4()), + "bundle_json": bundle_json, + "contributors": list(dict.fromkeys((*previous.contributors, author))), + }, ) - ) - - async def list(self, team_id: str, after: str | None, limit: int) -> tuple[WorkflowDraft, ...]: - return _DRAFTS.validate_python( - await self._database.query_raw( - f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" ' - "WHERE team_id = $1 AND workflow_id > $2 ORDER BY workflow_id LIMIT $3", - team_id, - after or "", - limit, - ) - ) + if not updated: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + row: Final = await tx.litellm_workflowdefinition.find_unique(where={"workflow_id": previous.workflow_id}) + return WorkflowDraft.model_validate(row, from_attributes=True) async def publish(self, team_id: str, name: str, etag: str, publisher: str) -> WorkflowSnapshot | WorkflowConflict: async with self._database.tx() as tx: - drafts: Final = _DRAFTS.validate_python( - await tx.query_raw( - f'SELECT {_DRAFT_COLUMNS} FROM "LiteLLM_WorkflowDefinition" ' - "WHERE team_id = $1 AND name = $2 FOR UPDATE", - team_id, - name, - ) + draft: Final = await tx.litellm_workflowdefinition.find_first( + where={"team_id": team_id, "name": name, "etag": etag} ) - if not drafts or drafts[0].etag != etag: + if draft is None: return WorkflowConflict(message="Draft changed or does not exist", status_code=412) - if publisher in drafts[0].contributors: + if publisher in draft.contributors: return WorkflowConflict(message="Draft contributors cannot publish their own draft", status_code=403) - rows: Final = _SNAPSHOTS.validate_python( - await tx.query_raw( - 'INSERT INTO "LiteLLM_WorkflowSnapshot" ' - "(snapshot_id, workflow_id, team_id, name, bundle_json, published_by) " - f"VALUES ($1, $2, $3, $4, $5, $6) RETURNING {_SNAPSHOT_COLUMNS}", - str(uuid.uuid4()), - drafts[0].workflow_id, - team_id, - name, - drafts[0].bundle_json, - publisher, - ) + snapshot_id: Final = str(uuid.uuid4()) + updated: Final = await tx.litellm_workflowdefinition.update_many( + where={"workflow_id": draft.workflow_id, "etag": etag}, + data={"etag": str(uuid.uuid4()), "published_snapshot_id": snapshot_id}, ) - await tx.execute_raw( - 'UPDATE "LiteLLM_WorkflowDefinition" SET published_snapshot_id = $3, etag = $4 ' - "WHERE team_id = $1 AND name = $2", - team_id, - name, - rows[0].snapshot_id, - str(uuid.uuid4()), + if not updated: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + row: Final = await tx.litellm_workflowsnapshot.create( + data={ + "snapshot_id": snapshot_id, + "workflow_id": draft.workflow_id, + "team_id": team_id, + "name": name, + "bundle_json": draft.bundle_json, + "published_by": publisher, + } ) - return rows[0] + return WorkflowSnapshot.model_validate(row, from_attributes=True) async def snapshot(self, snapshot_id: str) -> WorkflowSnapshot | None: - return _first( - _SNAPSHOTS.validate_python( - await self._database.query_raw( - f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = $1', - snapshot_id, - ) - ) - ) + row: Final = await self._database.litellm_workflowsnapshot.find_unique(where={"snapshot_id": snapshot_id}) + return WorkflowSnapshot.model_validate(row, from_attributes=True) if row is not None else None async def published(self, team_id: str, name: str) -> WorkflowSnapshot | None: - return _first( - _SNAPSHOTS.validate_python( - await self._database.query_raw( - f'SELECT {_SNAPSHOT_COLUMNS} FROM "LiteLLM_WorkflowSnapshot" WHERE snapshot_id = ' - '(SELECT published_snapshot_id FROM "LiteLLM_WorkflowDefinition" WHERE team_id = $1 AND name = $2)', - team_id, - name, - ) - ) - ) + draft: Final = await self.get(team_id, name) + return await self.snapshot(draft.published_snapshot_id) if draft and draft.published_snapshot_id else None async def pin_draft(self, workflow_id: str, etag: str, requested_by: str) -> WorkflowSnapshot | WorkflowConflict: - rows: Final = _SNAPSHOTS.validate_python( - await self._database.query_raw( - 'INSERT INTO "LiteLLM_WorkflowSnapshot" ' - "(snapshot_id, workflow_id, team_id, name, bundle_json, published_by, origin) " - "SELECT $3, workflow_id, team_id, name, bundle_json, $4, 'draft_test' " - 'FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND etag = $2 ' - f"RETURNING {_SNAPSHOT_COLUMNS}", - workflow_id, - etag, - str(uuid.uuid4()), - requested_by, + async with self._database.tx() as tx: + locked: Final = await tx.litellm_workflowdefinition.update_many( + where={"workflow_id": workflow_id, "etag": etag}, data={"etag": etag} ) - ) - return rows[0] if rows else WorkflowConflict(message="Draft changed or does not exist", status_code=412) + if not locked: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + draft: Final = await tx.litellm_workflowdefinition.find_unique(where={"workflow_id": workflow_id}) + if draft is None: + return WorkflowConflict(message="Draft changed or does not exist", status_code=412) + snapshot_id: Final = str(uuid5(NAMESPACE_URL, f"draft_test/{workflow_id}/{etag}")) + existing: Final = await tx.litellm_workflowsnapshot.find_unique(where={"snapshot_id": snapshot_id}) + if existing is not None: + return WorkflowSnapshot.model_validate(existing, from_attributes=True) + row: Final = await tx.litellm_workflowsnapshot.create( + data={ + "snapshot_id": snapshot_id, + "workflow_id": workflow_id, + "team_id": draft.team_id, + "name": draft.name, + "bundle_json": draft.bundle_json, + "published_by": requested_by, + "origin": "draft_test", + } + ) + return WorkflowSnapshot.model_validate(row, from_attributes=True) async def by_dispatch_key(self, dispatch_key: str) -> WorkflowExecution | None: - return _first( - _EXECUTIONS.validate_python( - await self._database.query_raw( - f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1', - dispatch_key, - ) - ) - ) + row: Final = await self._database.litellm_workflowexecution.find_unique(where={"dispatch_key": dispatch_key}) + return WorkflowExecution.model_validate(row, from_attributes=True) if row is not None else None async def execution(self, run_id: str) -> WorkflowExecution | None: - return _first( - _EXECUTIONS.validate_python( - await self._database.query_raw( - f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE run_id = $1', - run_id, - ) - ) - ) + row: Final = await self._database.litellm_workflowexecution.find_unique(where={"run_id": run_id}) + return WorkflowExecution.model_validate(row, from_attributes=True) if row is not None else None async def dispatch(self, execution: WorkflowExecution) -> WorkflowExecution: async with self._database.tx() as tx: - await tx.query_raw( - "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended($1, 0))", execution.dispatch_key - ) - existing: Final = _EXECUTIONS.validate_python( - await tx.query_raw( - f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" WHERE dispatch_key = $1', - execution.dispatch_key, - ) - ) - if existing: - return existing[0] - await tx.execute_raw( - 'INSERT INTO "LiteLLM_WorkflowRun" ' - "(run_id, session_id, workflow_type, created_by, input, metadata, engine_owned, updated_at) " - "VALUES ($1, $1, 'liteflow', $2, $3::jsonb, $4::jsonb, true, NOW())", - execution.run_id, - execution.principal, - execution.request_json, - json.dumps( + inserted: Final = await tx.litellm_workflowexecution.create_many( + data=[ { + "run_id": execution.run_id, "snapshot_id": execution.snapshot_id, "team_id": execution.team_id, - "workflow_name": execution.name, + "name": execution.name, + "principal": execution.principal, + "dispatch_key": execution.dispatch_key, + "request_digest": execution.request_digest, + "request_json": execution.request_json, + "deadline": execution.deadline, } - ), + ], + skip_duplicates=True, ) - rows: Final = _EXECUTIONS.validate_python( - await tx.query_raw( - 'INSERT INTO "LiteLLM_WorkflowExecution" ' - "(run_id, snapshot_id, team_id, name, principal, dispatch_key, request_digest, request_json, deadline) " - f"VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9::timestamptz AT TIME ZONE 'UTC') RETURNING {_EXECUTION_COLUMNS}", - execution.run_id, - execution.snapshot_id, - execution.team_id, - execution.name, - execution.principal, - execution.dispatch_key, - execution.request_digest, - execution.request_json, - execution.deadline, + if inserted: + await tx.litellm_workflowrun.create( + data={ + "run_id": execution.run_id, + "session_id": execution.run_id, + "workflow_type": "liteflow", + "created_by": execution.principal, + "input": _json(execution.request_json), + "engine_owned": True, + "metadata": _json( + json.dumps( + { + "snapshot_id": execution.snapshot_id, + "team_id": execution.team_id, + "workflow_name": execution.name, + } + ) + ), + } ) - ) - return rows[0] + row: Final = await tx.litellm_workflowexecution.find_unique(where={"dispatch_key": execution.dispatch_key}) + return WorkflowExecution.model_validate(row, from_attributes=True) async def lease_next(self, owner: str) -> WorkflowExecution | None: - return _first( - _EXECUTIONS.validate_python( - await self._database.query_raw( - "UPDATE \"LiteLLM_WorkflowExecution\" SET lease_owner = $1, lease_until = NOW() + INTERVAL '30 seconds', " - "state = 'running' WHERE run_id = (SELECT run_id FROM \"LiteLLM_WorkflowExecution\" " - "WHERE state IN ('pending', 'running') AND (lease_until IS NULL OR lease_until < NOW()) " - f"ORDER BY deadline FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING {_EXECUTION_COLUMNS}", - owner, - ) + now: Final = datetime.now(timezone.utc) + eligible: Final = { + "state": {"in": ["pending", "running"]}, + "OR": [{"lease_until": None}, {"lease_until": {"lt": now}}], + } + async with self._database.tx() as tx: + candidate: Final = await tx.litellm_workflowexecution.find_first(where=eligible, order={"deadline": "asc"}) + if candidate is None: + return None + updated: Final = await tx.litellm_workflowexecution.update_many( + where={**eligible, "run_id": candidate.run_id}, + data={"lease_owner": owner, "lease_until": now + timedelta(seconds=30), "state": "running"}, ) - ) + if not updated: + return None + row: Final = await tx.litellm_workflowexecution.find_unique(where={"run_id": candidate.run_id}) + return WorkflowExecution.model_validate(row, from_attributes=True) async def renew(self, run_id: str, owner: str) -> bool: return bool( - await self._database.execute_raw( - "UPDATE \"LiteLLM_WorkflowExecution\" SET lease_until = NOW() + INTERVAL '30 seconds' " - "WHERE run_id = $1 AND lease_owner = $2 AND lease_until > NOW() AND state = 'running'", - run_id, - owner, + await self._database.litellm_workflowexecution.update_many( + where=_lease_where(run_id, owner, allow_cancel=True), + data={"lease_until": datetime.now(timezone.utc) + timedelta(seconds=30)}, ) ) async def steps(self, run_id: str) -> tuple[WorkflowStep, ...]: - return _STEPS.validate_python( - await self._database.query_raw( - f'SELECT {_STEP_COLUMNS} FROM "LiteLLM_WorkflowStep" WHERE run_id = $1 ORDER BY started_at, step_id', - run_id, - ) + rows: Final = await self._database.litellm_workflowstep.find_many( + where={"run_id": run_id}, order=[{"started_at": "asc"}, {"step_id": "asc"}] ) + return tuple(WorkflowStep.model_validate(row, from_attributes=True) for row in rows) async def start_step( self, run_id: str, node_id: str, input: JsonValue, deadline: datetime, owner: str ) -> WorkflowStep | WorkflowConflict: - rows: Final = _STEPS.validate_python( - await self._database.query_raw( - 'INSERT INTO "LiteLLM_WorkflowStep" (step_id, run_id, node_id, input, deadline) ' - "SELECT $1,$2,$3,$4::jsonb,$5::timestamptz AT TIME ZONE 'UTC' WHERE EXISTS (SELECT 1 FROM \"LiteLLM_WorkflowExecution\" " - "WHERE run_id=$2 AND lease_owner=$6 AND lease_until>NOW() AND NOT cancel_requested) " - "ON CONFLICT (run_id, node_id) DO UPDATE SET node_id = EXCLUDED.node_id " - f"RETURNING {_STEP_COLUMNS}", - str(uuid.uuid4()), - run_id, - node_id, - json.dumps(input, allow_nan=False), - deadline, - owner, + async with self._database.tx() as tx: + if not await _lock_execution(tx, run_id, owner): + return WorkflowConflict(message="Execution lease lost") + existing: Final = await tx.litellm_workflowstep.find_unique( + where={"run_id_node_id": {"run_id": run_id, "node_id": node_id}} ) - ) - return rows[0] if rows else WorkflowConflict(message="Execution lease lost") + if existing is not None: + return WorkflowStep.model_validate(existing, from_attributes=True) + row: Final = await tx.litellm_workflowstep.create( + data={ + "step_id": str(uuid.uuid4()), + "run_id": run_id, + "node_id": node_id, + "input": _json(json.dumps(input, allow_nan=False)), + "deadline": deadline, + } + ) + return WorkflowStep.model_validate(row, from_attributes=True) async def attempt_step(self, step_id: str, owner: str) -> WorkflowStep | WorkflowConflict: - rows: Final = _STEPS.validate_python( - await self._database.query_raw( - 'UPDATE "LiteLLM_WorkflowStep" SET attempts = attempts + 1 ' - 'WHERE step_id = $1 AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e ' - 'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$2 ' - f"AND e.lease_until>NOW() AND NOT e.cancel_requested) RETURNING {_STEP_COLUMNS}", - step_id, - owner, + async with self._database.tx() as tx: + previous: Final = await tx.litellm_workflowstep.find_unique(where={"step_id": step_id}) + if previous is None or not await _lock_execution(tx, previous.run_id, owner): + return WorkflowConflict(message="Execution lease lost") + row: Final = await tx.litellm_workflowstep.update( + where={"step_id": step_id}, data={"attempts": {"increment": 1}} ) - ) - return rows[0] if rows else WorkflowConflict(message="Execution lease lost") - - async def _finish_step(self, database: WorkflowTransaction, step: WorkflowStep, owner: str) -> bool: - return bool( - await database.execute_raw( - "UPDATE \"LiteLLM_WorkflowStep\" SET state=$2, output=$3::jsonb, finished_at=$4::timestamptz AT TIME ZONE 'UTC', error=$5, " - "prompt_tokens=$6, completion_tokens=$7, response_cost=$8, request_id=$9 WHERE step_id=$1 " - 'AND EXISTS (SELECT 1 FROM "LiteLLM_WorkflowExecution" e ' - 'WHERE e.run_id="LiteLLM_WorkflowStep".run_id AND e.lease_owner=$10 ' - "AND e.lease_until>NOW() AND NOT e.cancel_requested)", - step.step_id, - step.state.value, - json.dumps(step.output, allow_nan=False), - step.finished_at, - step.error, - step.prompt_tokens, - step.completion_tokens, - step.response_cost, - step.request_id, - owner, - ) - ) + return WorkflowStep.model_validate(row, from_attributes=True) async def finish_step(self, step: WorkflowStep, owner: str) -> bool: - return await self._finish_step(self._database, step, owner) + async with self._database.tx() as tx: + if not await _lock_execution(tx, step.run_id, owner): + return False + return bool( + await tx.litellm_workflowstep.update_many( + where={"step_id": step.step_id, "run_id": step.run_id}, data=_step_data(step) + ) + ) async def checkpoint(self, run_id: str, checkpoint: WorkflowCheckpoint, step: WorkflowStep, owner: str) -> bool: async with self._database.tx() as tx: - rows: Final = _EXECUTIONS.validate_python( - await tx.query_raw( - f'SELECT {_EXECUTION_COLUMNS} FROM "LiteLLM_WorkflowExecution" ' - "WHERE run_id = $1 AND lease_owner=$2 AND lease_until>NOW() AND NOT cancel_requested FOR UPDATE", - run_id, - owner, - ) - ) - if not rows or not await self._finish_step(tx, step, owner): + if step.run_id != run_id or not await _lock_execution(tx, run_id, owner): return False - if any(item.node_path == checkpoint.node_path for item in rows[0].checkpoints): + current: Final = await tx.litellm_workflowexecution.find_unique(where={"run_id": run_id}) + execution: Final = WorkflowExecution.model_validate(current, from_attributes=True) + if any(item.node_path == checkpoint.node_path for item in execution.checkpoints): return True - checkpoints: Final = (*rows[0].checkpoints, checkpoint) - await tx.execute_raw( - 'UPDATE "LiteLLM_WorkflowExecution" SET checkpoints = $2::jsonb WHERE run_id = $1', - run_id, - _CHECKPOINTS.dump_json(checkpoints).decode(), + updated: Final = await tx.litellm_workflowstep.update_many( + where={"step_id": step.step_id, "run_id": run_id}, data=_step_data(step) ) - await tx.execute_raw( - 'INSERT INTO "LiteLLM_WorkflowEvent" ' - "(event_id, run_id, event_type, step_name, sequence_number, data) " - "VALUES ($1,$2,'node.completed',$3,$4,$5::jsonb)", - str(uuid.uuid4()), - run_id, - checkpoint.node_path, - len(checkpoints) - 1, - checkpoint.model_dump_json(), + if not updated: + return False + checkpoints: Final = (*execution.checkpoints, checkpoint) + await tx.litellm_workflowexecution.update( + where={"run_id": run_id}, + data={ + "checkpoints": _json( + json.dumps([item.model_dump(mode="json") for item in checkpoints], allow_nan=False) + ), + }, + ) + await tx.litellm_workflowevent.create( + data={ + "event_id": str(uuid.uuid4()), + "run_id": run_id, + "event_type": "node.completed", + "step_name": checkpoint.node_path, + "sequence_number": len(checkpoints) - 1, + "data": _json(checkpoint.model_dump_json()), + } ) return True @@ -405,34 +392,31 @@ class WorkflowRepository: owner: str, ) -> bool: async with self._database.tx() as tx: - updated: Final = _EXECUTIONS.validate_python( - await tx.query_raw( - "UPDATE \"LiteLLM_WorkflowExecution\" SET state = CASE WHEN cancel_requested THEN 'cancelled' ELSE $2 END, " - "response_json = CASE WHEN cancel_requested THEN NULL ELSE $3 END, error = $4 " - f"WHERE run_id = $1 AND state IN ($5,$6) AND lease_owner=$7 AND lease_until>NOW() RETURNING {_EXECUTION_COLUMNS}", - run_id, - state.value, - response_json, - error, - WorkflowRunState.RUNNING.value, - WorkflowRunState.PENDING.value, - owner, - ) - ) - if not updated: + if not await _lock_execution(tx, run_id, owner, allow_cancel=True): return False - await tx.execute_raw( - 'UPDATE "LiteLLM_WorkflowRun" SET status = $2, output = $3::jsonb, updated_at = NOW() WHERE run_id = $1', - run_id, - updated[0].state.value, - json.dumps({"response": json.loads(updated[0].response_json) if updated[0].response_json else None}), + previous: Final = await tx.litellm_workflowexecution.find_unique(where={"run_id": run_id}) + if previous is None: + return False + final_state: Final = WorkflowRunState.CANCELLED if previous.cancel_requested else state + final_response: Final = None if previous.cancel_requested else response_json + await tx.litellm_workflowexecution.update( + where={"run_id": run_id}, + data={ + "state": final_state.value, + "response_json": final_response, + "error": error, + }, + ) + await tx.litellm_workflowrun.update( + where={"run_id": run_id}, + data={ + "status": final_state.value, + "output": _json(json.dumps({"response": json.loads(final_response) if final_response else None})), + }, ) return True async def cancel(self, run_id: str) -> None: - await self._database.execute_raw( - 'UPDATE "LiteLLM_WorkflowExecution" SET cancel_requested = true WHERE run_id = $1 AND state IN ($2,$3)', - run_id, - WorkflowRunState.PENDING.value, - WorkflowRunState.RUNNING.value, + await self._database.litellm_workflowexecution.update_many( + where={"run_id": run_id, "state": {"in": ["pending", "running"]}}, data={"cancel_requested": True} ) diff --git a/tests/integration/management/test_workflow_execution.py b/tests/integration/management/test_workflow_execution.py index d68271bdcbb..27ea782b430 100644 --- a/tests/integration/management/test_workflow_execution.py +++ b/tests/integration/management/test_workflow_execution.py @@ -3,6 +3,7 @@ import json import threading import uuid from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from typing import Final @@ -220,7 +221,7 @@ def test_draft_execution_requires_its_owning_team_key_even_for_multiteam_members with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") author: Final = scenario.user(user_role="internal_user") - members: Final = [{"role": "admin", "user_id": author}] + members: Final[JsonValue] = [{"role": "admin", "user_id": author}] owning_team: Final = scenario.team(models=[model], members_with_roles=members) other_team: Final = scenario.team(models=[model], members_with_roles=members) owning_key: Final = scenario.key(team_id=owning_team, user_id=author, models=[model]) @@ -299,6 +300,86 @@ def test_etags_and_independent_publication_do_not_execute_drafts(workflow_gatewa assert len(wire.drain()) == 2 +def test_shared_schema_references_do_not_expand_repeatedly(workflow_gateway: Gateway) -> None: + with workflow_gateway.scenario() as scenario: + model: Final = scenario.model() + actors: Final = _actors(scenario, model) + definitions: Final = _OBJECT.validate_python( + { + "level_0": dict(_MESSAGES), + **{f"level_{index}": {"anyOf": [{"$ref": f"#/$defs/level_{index - 1}"}] * 2} for index in range(1, 31)}, + } + ) + response: Final = workflow_gateway.client.post( + "/v1/workflows", + json={ + "bundle": { + **_bundle(model), + "input_schema": { + "$defs": definitions, + "$ref": "#/$defs/level_30", + }, + } + }, + headers={"Authorization": f"Bearer {actors.author}"}, + timeout=10, + ) + assert response.status_code == 201, response.text + + +def test_draft_test_retries_reuse_one_snapshot(workflow_gateway: Gateway) -> None: + with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: + model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") + actors: Final = _actors(scenario, model) + draft: Final = _register(workflow_gateway, actors, model) + with ThreadPoolExecutor(max_workers=4) as pool: + responses: Final = tuple( + pool.map( + lambda _: workflow_gateway.request( + "POST", + f"/v1/workflows/{draft.workflow_id}/test", + {"messages": _request()["messages"], "authorize_billing": True}, + key=actors.author, + headers={"If-Match": draft.etag, "Idempotency-Key": "draft-retry"}, + ), + range(4), + ) + ) + assert all(response.status_code == 202 for response in responses) + assert all(response.json() == responses[0].json() for response in responses) + run_id: Final = str(responses[0].json()["run_id"]) + trace: Final = eventually( + lambda: _trace(workflow_gateway, actors, run_id), lambda value: value.state == "completed", seconds=30 + ) + snapshots: Final = workflow_gateway.rows( + 'SELECT snapshot_id FROM "LiteLLM_WorkflowSnapshot" WHERE workflow_id=%s', (draft.workflow_id,) + ) + assert snapshots == [{"snapshot_id": trace.snapshot_id}] + assert len(wire.drain()) == 2 + + +def test_workflow_child_spend_keeps_request_budget_tags(workflow_gateway: Gateway) -> None: + with workflow_gateway.scenario() as scenario, wire_server(_completion) as wire: + model: Final = scenario.model(model="openai/gpt-6-sol", api_base=wire.url + "/v1") + actors: Final = _actors(scenario, model) + draft: Final = _register(workflow_gateway, actors, model) + _publish(workflow_gateway, actors, draft) + tag: Final = f"workflow-tag-{uuid.uuid4().hex}" + response: Final = workflow_gateway.request( + "POST", "/v1/chat/completions", {**_request(), "metadata": {"tags": [tag]}}, key=actors.caller + ) + assert response.status_code == 200, response.text + run_id: Final = response.headers["x-litellm-workflow-run-id"] + charges: Final = eventually( + lambda: workflow_gateway.rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (run_id,) + ), + lambda rows: len(rows) == 2, + seconds=70, + ) + assert all(tag in TypeAdapter(tuple[str, ...]).validate_python(row["request_tags"]) for row in charges) + + @pytest.mark.parametrize("structured", [False, True]) def test_typed_checkpoints_idempotency_bounded_delegation_and_single_child_billing( workflow_gateway: Gateway,