scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

add metadata Json? column to LiteLLM_ScheduledTaskTable

Browse source 
Lets agents stash arbitrary state on a task — session ids, tags,
correlation refs, anything they want surfaced back through /due.
Same shape as action_args (Json?), same encoder path, same omit-when-
None semantics on create + update.

Surfaced in:
- POST /v1/tasks                accepts optional metadata
- PATCH /v1/tasks/{task_id}     can update metadata
- GET /v1/tasks{,/{task_id}}    returns metadata
- GET /v1/tasks/due             returns metadata in claim payload

Schema applied to all three schema.prisma files plus the migration.
Tests cover metadata round-trip and omit-when-not-supplied.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia 2026-04-29 16:20:09 -07:00
parent 254352f28e
commit cda2a9cc80
8 changed files with 57 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260429000000_add_scheduled_tasks/migration.sql
View file

@ -10,6 +10,7 @@ CREATE TABLE "LiteLLM_ScheduledTaskTable" (
"action_args" JSONB,
"check_prompt" TEXT,
"format_prompt" TEXT,
"metadata" JSONB,
"schedule_kind" TEXT NOT NULL,
"schedule_spec" TEXT NOT NULL,
"schedule_tz" TEXT,

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable {
action_args Json?
check_prompt String?
format_prompt String?
metadata Json?
schedule_kind String
schedule_spec String

3
litellm/proxy/scheduled_tasks/endpoints.py
View file

@ -57,6 +57,7 @@ def _row_to_response(row) -> ScheduledTaskResponse:
action_args=row.action_args,
check_prompt=row.check_prompt,
format_prompt=row.format_prompt,
metadata=getattr(row, "metadata", None),
schedule_kind=row.schedule_kind,
schedule_spec=row.schedule_spec,
schedule_tz=row.schedule_tz,
@ -136,6 +137,7 @@ async def create_scheduled_task(
action_args=data.action_args,
check_prompt=data.check_prompt,
format_prompt=data.format_prompt,
metadata=data.metadata,
schedule_kind=data.schedule_kind,
schedule_spec=data.schedule_spec,
schedule_tz=data.schedule_tz,
@ -207,6 +209,7 @@ async def get_due_tasks(
action_args=r["action_args"],
check_prompt=r["check_prompt"],
format_prompt=r["format_prompt"],
metadata=r.get("metadata"),
scheduled_for=r["scheduled_for"],
)
for r in rows

28
litellm/proxy/scheduled_tasks/store.py
View file

@ -40,9 +40,14 @@ UPDATABLE_FIELDS = frozenset(
"action",
"action_args",
"format_prompt",
"metadata",
}
)
# Json? columns that go through Prisma — must be JSON-encoded if present,
# omitted entirely if None.
JSON_FIELDS = frozenset({"action_args", "metadata"})
MAX_ACTIVE_TASKS_PER_KEY = 10
TERMINAL_STATUSES = ("fired", "expired", "cancelled")
@ -65,6 +70,7 @@ async def create_task(
action_args: Optional[Dict[str, Any]],
check_prompt: Optional[str],
format_prompt: Optional[str],
metadata: Optional[Any],
schedule_kind: str,
schedule_spec: str,
schedule_tz: Optional[str],
@ -90,6 +96,8 @@ async def create_task(
}
if action_args is not None:
data["action_args"] = _encode_json(action_args)
if metadata is not None:
data["metadata"] = _encode_json(metadata)
return await prisma_client.db.litellm_scheduledtasktable.create(data=data)
@ -149,13 +157,17 @@ async def update_task_for_owner(
# Mirror the create path: Json columns must be json-encoded for
# prisma-client-python. Drop the key entirely when caller sent null;
# passing None to a Json? column is rejected.
if "action_args" in fields:
if fields["action_args"] is None:
fields = {k: v for k, v in fields.items() if k != "action_args"}
if not fields:
raise ValueError("no fields to update")
cleaned: Dict[str, Any] = {}
for k, v in fields.items():
if k in JSON_FIELDS:
if v is None:
continue # skip — Prisma rejects None on Json?
cleaned[k] = _encode_json(v)
else:
fields = {**fields, "action_args": _encode_json(fields["action_args"])}
cleaned[k] = v
if not cleaned:
raise ValueError("no fields to update")
fields = cleaned
return await prisma_client.db.litellm_scheduledtasktable.update(
where={"task_id": task_id},
@ -205,7 +217,8 @@ async def claim_due(
"""
SELECT task_id, schedule_kind, schedule_spec, schedule_tz,
fire_once, expires_at, next_run_at,
action, action_args, check_prompt, format_prompt, title
action, action_args, check_prompt, format_prompt,
metadata, title
FROM "LiteLLM_ScheduledTaskTable"
WHERE status = 'pending'
AND next_run_at <= now()
@ -264,6 +277,7 @@ async def claim_due(
"action_args": r["action_args"],
"check_prompt": r["check_prompt"],
"format_prompt": r["format_prompt"],
"metadata": r["metadata"],
"scheduled_for": r["next_run_at"],
}
for r in rows

4
litellm/proxy/scheduled_tasks/types.py
View file

@ -13,6 +13,7 @@ class CreateScheduledTaskRequest(BaseModel):
action_args: Optional[Dict[str, Any]] = None
check_prompt: Optional[str] = None
format_prompt: Optional[str] = None
metadata: Optional[Any] = None
schedule_kind: ScheduleKind
schedule_spec: str
@ -34,6 +35,7 @@ class UpdateScheduledTaskRequest(BaseModel):
action: Optional[str] = None
action_args: Optional[Dict[str, Any]] = None
format_prompt: Optional[str] = None
metadata: Optional[Any] = None
class ScheduledTaskResponse(BaseModel):
@ -48,6 +50,7 @@ class ScheduledTaskResponse(BaseModel):
action_args: Optional[Dict[str, Any]]
check_prompt: Optional[str]
format_prompt: Optional[str]
metadata: Optional[Any]
schedule_kind: str
schedule_spec: str
@ -73,6 +76,7 @@ class DueTaskResponse(BaseModel):
action_args: Optional[Dict[str, Any]]
check_prompt: Optional[str]
format_prompt: Optional[str]
metadata: Optional[Any]
scheduled_for: datetime

1
litellm/proxy/schema.prisma
View file

@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable {
action_args Json?
check_prompt String?
format_prompt String?
metadata Json?
schedule_kind String
schedule_spec String

1
schema.prisma
View file

@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable {
action_args Json?
check_prompt String?
format_prompt String?
metadata Json?
schedule_kind String
schedule_spec String

33
tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py
View file

@ -44,6 +44,7 @@ def _make_row(**kwargs) -> MagicMock:
"action_args": None,
"check_prompt": "is it done?",
"format_prompt": None,
"metadata": None,
"schedule_kind": "interval",
"schedule_spec": "5m",
"schedule_tz": None,
@ -84,18 +85,19 @@ class _FakeScheduledTaskTable:
# Mirror real Prisma write semantics:
# - relation field `owner_key: {connect: {token: X}}` becomes the
# scalar `owner_token` column on the row;
# - Json columns (`action_args`) come in as JSON strings and are
# deserialized back to Python on read.
# - Json? columns (`action_args`, `metadata`) come in as JSON
# strings and are deserialised back to Python on read.
normalized = dict(data)
if "owner_key" in normalized:
connect = normalized.pop("owner_key", {}).get("connect", {})
if "token" in connect:
normalized["owner_token"] = connect["token"]
if isinstance(normalized.get("action_args"), str):
try:
normalized["action_args"] = json.loads(normalized["action_args"])
except ValueError:
pass
for json_field in ("action_args", "metadata"):
if isinstance(normalized.get(json_field), str):
try:
normalized[json_field] = json.loads(normalized[json_field])
except ValueError:
pass
self._counter += 1
row = _make_row(task_id=f"task-{self._counter}", **normalized)
self.rows.append(row)
@ -122,7 +124,7 @@ class _FakeScheduledTaskTable:
for r in self.rows:
if r.task_id == where["task_id"]:
for k, v in data.items():
if k == "action_args" and isinstance(v, str):
if k in ("action_args", "metadata") and isinstance(v, str):
try:
v = json.loads(v)
except ValueError:
@ -193,6 +195,7 @@ class _FakeTx:
"action_args": r.action_args,
"check_prompt": r.check_prompt,
"format_prompt": r.format_prompt,
"metadata": r.metadata,
"title": r.title,
}
)
@ -340,6 +343,20 @@ class TestCreate:
r = self.client.post("/v1/tasks", json=_create_payload(title="overflow"))
assert r.status_code == 429
def test_metadata_roundtrips(self):
payload = _create_payload(metadata={"agent_session": "abc", "tags": [1, 2, 3]})
with _patch_prisma(self.prisma):
r = self.client.post("/v1/tasks", json=payload)
assert r.status_code == 200, r.text
body = r.json()
assert body["metadata"] == {"agent_session": "abc", "tags": [1, 2, 3]}
def test_metadata_omitted_when_not_supplied(self):
with _patch_prisma(self.prisma):
r = self.client.post("/v1/tasks", json=_create_payload())
assert r.status_code == 200, r.text
assert r.json()["metadata"] is None
def test_expires_in_past_rejected(self):
payload = _create_payload(
expires_at=(_now() - timedelta(seconds=1)).isoformat(),