mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix: use Prisma workflow persistence and bound draft retries
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
ca2486e4f2
commit
17ec43986f
4 changed files with 412 additions and 318 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue