Skip to content

Commit 7ddb90e

Browse files
committed
feat: reserved invocation_settings envelope field on generated /invoke
Every wrap_in_fastapi-generated /invoke now accepts an optional invocation_settings object as a reserved envelope field. It is kept out of the function's /schema (no namespace collision with data inputs) and delivered either as a kwarg when the wrapped function declares an invocation_settings parameter, or via the get_invocation_settings() ContextVar accessor otherwise (context propagated across the executor thread). The ContextVar is set even on the kwarg path so helper code deep in a plugin's call stack can always use the accessor. GET /capabilities advertises support so callers (e.g. the ETL controller) can gate injection without schema sniffing. The signature check for the kwarg path is hoisted to wrap time; requests pay a boolean, not reflection. This lets one standing plugin process serve work items with differing node settings (warm pools, worker fleets, per-invoke reconfiguration) instead of freezing settings from a boot-time file. 75 tests passed (5 new in test/test_invocation_settings.py).
1 parent 4aea423 commit 7ddb90e

3 files changed

Lines changed: 120 additions & 3 deletions

File tree

test/test_invocation_settings.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
from typing import Any, Optional
2+
3+
from fastapi.testclient import TestClient
4+
5+
from unstructured_platform_plugins.etl_uvicorn.api_generator import wrap_in_fastapi
6+
from unstructured_platform_plugins.etl_uvicorn.invocation_context import (
7+
get_invocation_settings,
8+
)
9+
from pydantic import BaseModel
10+
11+
12+
class Out(BaseModel):
13+
content: str
14+
settings_seen: Optional[dict[str, Any]] = None
15+
16+
17+
def fn_declares(content: str, invocation_settings: Optional[dict[str, Any]] = None) -> Out:
18+
return Out(content=content, settings_seen=invocation_settings)
19+
20+
21+
def fn_uses_context(content: str) -> Out:
22+
return Out(content=content, settings_seen=get_invocation_settings())
23+
24+
25+
def test_declared_param_receives_settings():
26+
client = TestClient(wrap_in_fastapi(func=fn_declares, plugin_id="t1"))
27+
resp = client.post(
28+
"/invoke",
29+
json={"content": "x", "invocation_settings": {"strategy": "fast"}},
30+
)
31+
assert resp.status_code == 200
32+
assert resp.json()["output"]["settings_seen"] == {"strategy": "fast"}
33+
34+
35+
def test_context_accessor_receives_settings():
36+
client = TestClient(wrap_in_fastapi(func=fn_uses_context, plugin_id="t2"))
37+
resp = client.post(
38+
"/invoke",
39+
json={"content": "x", "invocation_settings": {"k": 1}},
40+
)
41+
assert resp.status_code == 200
42+
assert resp.json()["output"]["settings_seen"] == {"k": 1}
43+
44+
45+
def test_omitted_settings_is_none_and_context_resets():
46+
client = TestClient(wrap_in_fastapi(func=fn_uses_context, plugin_id="t3"))
47+
with_settings = client.post("/invoke", json={"content": "x", "invocation_settings": {"k": 1}})
48+
without = client.post("/invoke", json={"content": "x"})
49+
assert with_settings.json()["output"]["settings_seen"] == {"k": 1}
50+
assert without.json()["output"]["settings_seen"] is None
51+
52+
53+
def test_capabilities_endpoint():
54+
client = TestClient(wrap_in_fastapi(func=fn_uses_context, plugin_id="t4"))
55+
resp = client.get("/capabilities")
56+
assert resp.status_code == 200
57+
assert resp.json()["invocation_settings"] is True
58+
59+
60+
def test_schema_unpolluted_by_envelope_field():
61+
client = TestClient(wrap_in_fastapi(func=fn_uses_context, plugin_id="t5"))
62+
schema = client.get("/schema").json()
63+
assert "invocation_settings" not in schema["inputs"].get("properties", {})

unstructured_platform_plugins/etl_uvicorn/api_generator.py

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import contextvars
23
import hashlib
34
import inspect
45
import json
@@ -18,6 +19,9 @@
1819
from uvicorn.importer import import_from_string
1920

2021
from unstructured_platform_plugins.etl_uvicorn.otel import get_metric_provider, get_trace_provider
22+
from unstructured_platform_plugins.etl_uvicorn.invocation_context import (
23+
_invocation_settings_var,
24+
)
2125
from unstructured_platform_plugins.etl_uvicorn.utils import (
2226
get_func,
2327
get_input_schema,
@@ -67,7 +71,12 @@ async def invoke_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -
6771
if inspect.iscoroutinefunction(func):
6872
return await func(**kwargs)
6973
else:
70-
return await asyncio.get_event_loop().run_in_executor(None, partial(func, **kwargs))
74+
# copy_context() so ContextVars (e.g. invocation_settings) reach the
75+
# worker thread; run_in_executor does not propagate them by itself.
76+
ctx = contextvars.copy_context()
77+
return await asyncio.get_event_loop().run_in_executor(
78+
None, partial(ctx.run, partial(func, **kwargs))
79+
)
7180

7281

7382
def check_precheck_func(precheck_func: Callable):
@@ -149,8 +158,22 @@ class InvokeResponse(BaseModel):
149158
output: Optional[response_type] = None
150159
message_channels: MessageChannels = Field(default_factory=MessageChannels)
151160

152-
input_schema = get_input_schema(func, omit=["usage", "filedata_meta", "message_channels"])
161+
input_schema = get_input_schema(
162+
func, omit=["usage", "filedata_meta", "message_channels", "invocation_settings"]
163+
)
153164
input_schema_model = schema_to_base_model(input_schema)
165+
# First-class per-invoke settings: every generated /invoke accepts an optional
166+
# invocation_settings object, independent of the wrapped function's signature.
167+
# Delivery: passed as a kwarg when the function declares the parameter,
168+
# otherwise exposed through get_invocation_settings() for the call's duration.
169+
input_schema_model = create_model(
170+
"InvokeEnvelope",
171+
__base__=input_schema_model,
172+
invocation_settings=(Optional[dict[str, Any]], None),
173+
)
174+
func_declares_invocation_settings = (
175+
"invocation_settings" in inspect.signature(func).parameters
176+
)
154177

155178
logging.getLogger("etl_uvicorn.fastapi")
156179

@@ -261,6 +284,10 @@ async def run_job(request: input_schema_model) -> ResponseType:
261284
log_func_and_body(func=func, body=request.json())
262285
# Create dictionary from pydantic model while preserving underlying types
263286
request_dict = {f: getattr(request, f) for f in request.model_fields}
287+
invocation_settings = request_dict.pop("invocation_settings", None)
288+
if func_declares_invocation_settings:
289+
request_dict["invocation_settings"] = invocation_settings
290+
settings_token = _invocation_settings_var.set(invocation_settings)
264291
# Make sure nested classes get instantiated correctly
265292
if "file_data" in request_dict:
266293
request_dict["file_data"] = file_data_from_dict(
@@ -269,7 +296,10 @@ async def run_job(request: input_schema_model) -> ResponseType:
269296
map_inputs(func=func, raw_inputs=request_dict)
270297
if logger.level == LOG_LEVELS.get("trace", logging.NOTSET):
271298
logger.log(level=logger.level, msg=f"passing inputs to function: {request_dict}")
272-
return await wrap_fn(func=func, kwargs=request_dict)
299+
try:
300+
return await wrap_fn(func=func, kwargs=request_dict)
301+
finally:
302+
_invocation_settings_var.reset(settings_token)
273303

274304
else:
275305

@@ -315,6 +345,10 @@ async def run_precheck() -> InvokePrecheckResponse:
315345
async def get_id() -> str:
316346
return plugin_id
317347

348+
@fastapi_app.get("/capabilities")
349+
async def get_capabilities() -> dict[str, Any]:
350+
return {"invocation_settings": True}
351+
318352
# Run initial schema validation
319353
try:
320354
asyncio.run(get_schema())
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
"""Per-invocation context for first-class /invoke envelope fields.
2+
3+
``invocation_settings`` is a reserved, always-accepted field of every generated
4+
``/invoke`` request. Functions that declare an ``invocation_settings`` parameter
5+
receive it as a kwarg; all other code (including plugins with hand-rolled apps
6+
that adopt the same contract) can read it for the duration of the call via
7+
:func:`get_invocation_settings`.
8+
"""
9+
10+
from contextvars import ContextVar
11+
from typing import Any, Optional
12+
13+
_invocation_settings_var: ContextVar[Optional[dict[str, Any]]] = ContextVar(
14+
"invocation_settings", default=None
15+
)
16+
17+
18+
def get_invocation_settings() -> Optional[dict[str, Any]]:
19+
"""The invocation_settings of the in-flight /invoke call, if any."""
20+
return _invocation_settings_var.get()

0 commit comments

Comments
 (0)