Centralize schema preprocessing for manager payloads

cursoragent · shrisukhani · cursoragent · commit a8dcf4731601 · 2026-02-13T04:54:33.000Z
Co-authored-by: Shri Sukhani &lt;shrisukhani@users.noreply.github.com&gt;
diff --git a/hyperbrowser/client/managers/async_manager/agents/browser_use.py b/hyperbrowser/client/managers/async_manager/agents/browser_use.py
@@ -1,8 +1,8 @@
-import jsonref
 from typing import Optional
 
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import poll_until_terminal_status_async, retry_operation_async
+from ....schema_utils import resolve_schema_input
 
 from .....models import (
     POLLING_ATTEMPTS,
@@ -22,13 +22,9 @@ async def start(
         self, params: StartBrowserUseTaskParams
     ) -> StartBrowserUseTaskResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.output_model_schema and hasattr(
-            params.output_model_schema, "model_json_schema"
-        ):
-            payload["outputModelSchema"] = jsonref.replace_refs(
-                params.output_model_schema.model_json_schema(),
-                proxies=False,
-                lazy_load=False,
+        if params.output_model_schema:
+            payload["outputModelSchema"] = resolve_schema_input(
+                params.output_model_schema
             )
         response = await self._client.transport.post(
             self._client._build_url("/task/browser-use"),
diff --git a/hyperbrowser/client/managers/async_manager/extract.py b/hyperbrowser/client/managers/async_manager/extract.py
@@ -8,8 +8,8 @@
     StartExtractJobParams,
     StartExtractJobResponse,
 )
-import jsonref
 from ...polling import poll_until_terminal_status_async, retry_operation_async
+from ...schema_utils import resolve_schema_input
 
 
 class ExtractManager:
@@ -21,10 +21,8 @@ async def start(self, params: StartExtractJobParams) -> StartExtractJobResponse:
             raise HyperbrowserError("Either schema or prompt must be provided")
 
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.schema_ and hasattr(params.schema_, "model_json_schema"):
-            payload["schema"] = jsonref.replace_refs(
-                params.schema_.model_json_schema(), proxies=False, lazy_load=False
-            )
+        if params.schema_:
+            payload["schema"] = resolve_schema_input(params.schema_)
 
         response = await self._client.transport.post(
             self._client._build_url("/extract"),
diff --git a/hyperbrowser/client/managers/async_manager/web/__init__.py b/hyperbrowser/client/managers/async_manager/web/__init__.py
@@ -3,11 +3,10 @@
 from hyperbrowser.models import (
     FetchParams,
     FetchResponse,
-    FetchOutputJson,
     WebSearchParams,
     WebSearchResponse,
 )
-import jsonref
+from ....schema_utils import inject_web_output_schemas
 
 
 class WebManager:
@@ -18,15 +17,9 @@ def __init__(self, client):
 
     async def fetch(self, params: FetchParams) -> FetchResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = await self._client.transport.post(
             self._client._build_url("/web/fetch"),
diff --git a/hyperbrowser/client/managers/async_manager/web/batch_fetch.py b/hyperbrowser/client/managers/async_manager/web/batch_fetch.py
@@ -7,17 +7,16 @@
     GetBatchFetchJobParams,
     BatchFetchJobResponse,
     POLLING_ATTEMPTS,
-    FetchOutputJson,
 )
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import (
     has_exceeded_max_wait,
     poll_until_terminal_status_async,
     retry_operation_async,
 )
+from ....schema_utils import inject_web_output_schemas
 import asyncio
 import time
-import jsonref
 
 
 class BatchFetchManager:
@@ -28,15 +27,9 @@ async def start(
         self, params: StartBatchFetchJobParams
     ) -> StartBatchFetchJobResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = await self._client.transport.post(
             self._client._build_url("/web/batch-fetch"),
diff --git a/hyperbrowser/client/managers/async_manager/web/crawl.py b/hyperbrowser/client/managers/async_manager/web/crawl.py
@@ -7,17 +7,16 @@
     GetWebCrawlJobParams,
     WebCrawlJobResponse,
     POLLING_ATTEMPTS,
-    FetchOutputJson,
 )
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import (
     has_exceeded_max_wait,
     poll_until_terminal_status_async,
     retry_operation_async,
 )
+from ....schema_utils import inject_web_output_schemas
 import asyncio
 import time
-import jsonref
 
 
 class WebCrawlManager:
@@ -26,15 +25,9 @@ def __init__(self, client):
 
     async def start(self, params: StartWebCrawlJobParams) -> StartWebCrawlJobResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = await self._client.transport.post(
             self._client._build_url("/web/crawl"),
diff --git a/hyperbrowser/client/managers/sync_manager/agents/browser_use.py b/hyperbrowser/client/managers/sync_manager/agents/browser_use.py
@@ -1,8 +1,8 @@
-import jsonref
 from typing import Optional
 
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import poll_until_terminal_status, retry_operation
+from ....schema_utils import resolve_schema_input
 
 from .....models import (
     POLLING_ATTEMPTS,
@@ -20,13 +20,9 @@ def __init__(self, client):
 
     def start(self, params: StartBrowserUseTaskParams) -> StartBrowserUseTaskResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.output_model_schema and hasattr(
-            params.output_model_schema, "model_json_schema"
-        ):
-            payload["outputModelSchema"] = jsonref.replace_refs(
-                params.output_model_schema.model_json_schema(),
-                proxies=False,
-                lazy_load=False,
+        if params.output_model_schema:
+            payload["outputModelSchema"] = resolve_schema_input(
+                params.output_model_schema
             )
         response = self._client.transport.post(
             self._client._build_url("/task/browser-use"),
diff --git a/hyperbrowser/client/managers/sync_manager/extract.py b/hyperbrowser/client/managers/sync_manager/extract.py
@@ -8,8 +8,8 @@
     StartExtractJobParams,
     StartExtractJobResponse,
 )
-import jsonref
 from ...polling import poll_until_terminal_status, retry_operation
+from ...schema_utils import resolve_schema_input
 
 
 class ExtractManager:
@@ -21,10 +21,8 @@ def start(self, params: StartExtractJobParams) -> StartExtractJobResponse:
             raise HyperbrowserError("Either schema or prompt must be provided")
 
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.schema_ and hasattr(params.schema_, "model_json_schema"):
-            payload["schema"] = jsonref.replace_refs(
-                params.schema_.model_json_schema(), proxies=False, lazy_load=False
-            )
+        if params.schema_:
+            payload["schema"] = resolve_schema_input(params.schema_)
 
         response = self._client.transport.post(
             self._client._build_url("/extract"),
diff --git a/hyperbrowser/client/managers/sync_manager/web/__init__.py b/hyperbrowser/client/managers/sync_manager/web/__init__.py
@@ -3,11 +3,10 @@
 from hyperbrowser.models import (
     FetchParams,
     FetchResponse,
-    FetchOutputJson,
     WebSearchParams,
     WebSearchResponse,
 )
-import jsonref
+from ....schema_utils import inject_web_output_schemas
 
 
 class WebManager:
@@ -18,15 +17,9 @@ def __init__(self, client):
 
     def fetch(self, params: FetchParams) -> FetchResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = self._client.transport.post(
             self._client._build_url("/web/fetch"),
diff --git a/hyperbrowser/client/managers/sync_manager/web/batch_fetch.py b/hyperbrowser/client/managers/sync_manager/web/batch_fetch.py
@@ -7,16 +7,15 @@
     GetBatchFetchJobParams,
     BatchFetchJobResponse,
     POLLING_ATTEMPTS,
-    FetchOutputJson,
 )
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import (
     has_exceeded_max_wait,
     poll_until_terminal_status,
     retry_operation,
 )
+from ....schema_utils import inject_web_output_schemas
 import time
-import jsonref
 
 
 class BatchFetchManager:
@@ -25,15 +24,9 @@ def __init__(self, client):
 
     def start(self, params: StartBatchFetchJobParams) -> StartBatchFetchJobResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = self._client.transport.post(
             self._client._build_url("/web/batch-fetch"),
diff --git a/hyperbrowser/client/managers/sync_manager/web/crawl.py b/hyperbrowser/client/managers/sync_manager/web/crawl.py
@@ -7,16 +7,15 @@
     GetWebCrawlJobParams,
     WebCrawlJobResponse,
     POLLING_ATTEMPTS,
-    FetchOutputJson,
 )
 from hyperbrowser.exceptions import HyperbrowserError
 from ....polling import (
     has_exceeded_max_wait,
     poll_until_terminal_status,
     retry_operation,
 )
+from ....schema_utils import inject_web_output_schemas
 import time
-import jsonref
 
 
 class WebCrawlManager:
@@ -25,15 +24,9 @@ def __init__(self, client):
 
     def start(self, params: StartWebCrawlJobParams) -> StartWebCrawlJobResponse:
         payload = params.model_dump(exclude_none=True, by_alias=True)
-        if params.outputs and params.outputs.formats:
-            for index, output in enumerate(params.outputs.formats):
-                if isinstance(output, FetchOutputJson) and output.schema_:
-                    if hasattr(output.schema_, "model_json_schema"):
-                        payload["outputs"]["formats"][index]["schema"] = jsonref.replace_refs(
-                            output.schema_.model_json_schema(),
-                            proxies=False,
-                            lazy_load=False,
-                        )
+        inject_web_output_schemas(
+            payload, params.outputs.formats if params.outputs else None
+        )
 
         response = self._client.transport.post(
             self._client._build_url("/web/crawl"),
diff --git a/hyperbrowser/client/schema_utils.py b/hyperbrowser/client/schema_utils.py
@@ -0,0 +1,26 @@
+from typing import Any, List, Optional
+
+import jsonref
+
+
+def resolve_schema_input(schema_input: Any) -> Any:
+    if hasattr(schema_input, "model_json_schema"):
+        return jsonref.replace_refs(
+            schema_input.model_json_schema(),
+            proxies=False,
+            lazy_load=False,
+        )
+    return schema_input
+
+
+def inject_web_output_schemas(payload: dict, formats: Optional[List[Any]]) -> None:
+    if not formats:
+        return
+
+    for index, output_format in enumerate(formats):
+        schema_input = getattr(output_format, "schema_", None)
+        if schema_input is None:
+            continue
+        payload["outputs"]["formats"][index]["schema"] = resolve_schema_input(
+            schema_input
+        )
