Reuse shared schema injection helper across payload builders

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/managers/browser_use_payload_utils.py

@@ -2,7 +2,7 @@
 
 from hyperbrowser.models.agents.browser_use import StartBrowserUseTaskParams
 
-from ..schema_utils import resolve_schema_input
+from ..schema_utils import inject_resolved_schema
 from .serialization_utils import serialize_model_dump_to_dict
 
 
@@ -11,6 +11,9 @@ def build_browser_use_start_payload(params: StartBrowserUseTaskParams) -> Dict[s
         params,
         error_message="Failed to serialize browser-use start params",
     )
-    if params.output_model_schema:
-        payload["outputModelSchema"] = resolve_schema_input(params.output_model_schema)
+    inject_resolved_schema(
+        payload,
+        key="outputModelSchema",
+        schema_input=params.output_model_schema,
+    )
     return payload

hyperbrowser/client/managers/extract_payload_utils.py

@@ -3,7 +3,7 @@
 from hyperbrowser.exceptions import HyperbrowserError
 from hyperbrowser.models.extract import StartExtractJobParams
 
-from ..schema_utils import resolve_schema_input
+from ..schema_utils import inject_resolved_schema
 from .serialization_utils import serialize_model_dump_to_dict
 
 
@@ -15,6 +15,5 @@ def build_extract_start_payload(params: StartExtractJobParams) -> Dict[str, Any]
         params,
         error_message="Failed to serialize extract start params",
     )
-    if params.schema_:
-        payload["schema"] = resolve_schema_input(params.schema_)
+    inject_resolved_schema(payload, key="schema", schema_input=params.schema_)
     return payload

hyperbrowser/client/schema_utils.py

@@ -13,6 +13,12 @@ def resolve_schema_input(schema_input: Any) -> Any:
     return schema_input
 
 
+def inject_resolved_schema(payload: dict, *, key: str, schema_input: Any) -> None:
+    if schema_input is None:
+        return
+    payload[key] = resolve_schema_input(schema_input)
+
+
 def inject_web_output_schemas(payload: dict, formats: Optional[List[Any]]) -> None:
     if not formats:
         return
@@ -21,6 +27,8 @@ def inject_web_output_schemas(payload: dict, formats: Optional[List[Any]]) -> No
         schema_input = getattr(output_format, "schema_", None)
         if schema_input is None:
             continue
-        payload["outputs"]["formats"][index]["schema"] = resolve_schema_input(
-            schema_input
+        inject_resolved_schema(
+            payload["outputs"]["formats"][index],
+            key="schema",
+            schema_input=schema_input,
         )

tests/test_browser_use_payload_utils.py

@@ -29,8 +29,11 @@ class _SchemaModel(BaseModel):
 
     monkeypatch.setattr(
         browser_use_payload_utils,
-        "resolve_schema_input",
-        lambda schema_input: {"resolvedSchema": schema_input.__name__},
+        "inject_resolved_schema",
+        lambda payload, *, key, schema_input: payload.__setitem__(
+            key,
+            {"resolvedSchema": schema_input.__name__},
+        ),
     )
 
     payload = browser_use_payload_utils.build_browser_use_start_payload(params)

tests/test_extract_payload_utils.py

@@ -38,8 +38,11 @@ class _SchemaModel(BaseModel):
 
     monkeypatch.setattr(
         extract_payload_utils,
-        "resolve_schema_input",
-        lambda schema_input: {"resolvedSchema": schema_input.__name__},
+        "inject_resolved_schema",
+        lambda payload, *, key, schema_input: payload.__setitem__(
+            key,
+            {"resolvedSchema": schema_input.__name__},
+        ),
     )
 
     payload = extract_payload_utils.build_extract_start_payload(params)

tests/test_schema_utils.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel
+
+from hyperbrowser.client.schema_utils import inject_resolved_schema
+
+
+class _SchemaModel(BaseModel):
+    value: str
+
+
+def test_inject_resolved_schema_sets_resolved_schema_value():
+    payload = {"a": 1}
+
+    inject_resolved_schema(payload, key="schema", schema_input=_SchemaModel)
+
+    assert payload["schema"]["type"] == "object"
+    assert "value" in payload["schema"]["properties"]
+
+
+def test_inject_resolved_schema_ignores_none_schema_inputs():
+    payload = {"a": 1}
+
+    inject_resolved_schema(payload, key="schema", schema_input=None)
+
+    assert payload == {"a": 1}