stash

Author: maticzav

examples/async_create.py

@@ -12,7 +12,7 @@
 
 
 # Regular Task
-async def create_regular_task():
+async def create_regular_task() -> None:
     res = await client.tasks.create(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -23,7 +23,7 @@ async def create_regular_task():
 
 
 # Structured Output
-async def create_structured_task():
+async def create_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str
@@ -44,7 +44,7 @@ class SearchResult(BaseModel):
 # Main
 
 
-async def main():
+async def main() -> None:
     await asyncio.gather(
         #
         create_regular_task(),

examples/async_retrieve.py

@@ -12,7 +12,7 @@
 
 
 # Regular Task
-async def retrieve_regular_task():
+async def retrieve_regular_task() -> None:
     """
     Retrieves a regular task and waits for it to finish.
     """
@@ -39,7 +39,7 @@ async def retrieve_regular_task():
     print("Done")
 
 
-async def retrieve_structured_task():
+async def retrieve_structured_task() -> None:
     """
     Retrieves a structured task and waits for it to finish.
     """
@@ -84,7 +84,7 @@ class SearchResult(BaseModel):
 # Main
 
 
-async def main():
+async def main() -> None:
     await asyncio.gather(
         #
         retrieve_regular_task(),

examples/async_run.py

@@ -12,7 +12,7 @@
 
 
 # Regular Task
-async def run_regular_task():
+async def run_regular_task() -> None:
     regular_result = await client.tasks.run(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -27,7 +27,7 @@ async def run_regular_task():
 
 
 # Structured Output
-async def run_structured_task():
+async def run_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str
@@ -52,7 +52,7 @@ class SearchResult(BaseModel):
     print("Structured Task Done")
 
 
-async def main():
+async def main() -> None:
     await asyncio.gather(
         #
         run_regular_task(),

examples/async_stream.py

@@ -13,7 +13,7 @@
 
 
 # Regular Task
-async def stream_regular_task():
+async def stream_regular_task() -> None:
     regular_task = await client.tasks.create(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -36,7 +36,7 @@ async def stream_regular_task():
 
 
 # Structured Output
-async def stream_structured_task():
+async def stream_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str
@@ -70,7 +70,7 @@ class SearchResult(BaseModel):
 # Main
 
 
-async def main():
+async def main() -> None:
     await asyncio.gather(
         #
         stream_regular_task(),

examples/create.py

@@ -11,7 +11,7 @@
 
 
 # Regular Task
-def create_regular_task():
+def create_regular_task() -> None:
     res = client.tasks.create(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -25,7 +25,7 @@ def create_regular_task():
 
 
 # Structured Output
-def create_structured_task():
+def create_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str

examples/retrieve.py

@@ -12,7 +12,7 @@
 
 
 # Regular Task
-def retrieve_regular_task():
+def retrieve_regular_task() -> None:
     """
     Retrieves a regular task and waits for it to finish.
     """
@@ -42,7 +42,7 @@ def retrieve_regular_task():
 retrieve_regular_task()
 
 
-def retrieve_structured_task():
+def retrieve_structured_task() -> None:
     """
     Retrieves a structured task and waits for it to finish.
     """

examples/run.py

@@ -11,7 +11,7 @@
 
 
 # Regular Task
-def run_regular_task():
+def run_regular_task() -> None:
     regular_result = client.tasks.run(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -29,7 +29,7 @@ def run_regular_task():
 
 
 # Structured Output
-def run_structured_task():
+def run_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str

examples/stream.py

@@ -12,7 +12,7 @@
 
 
 # Regular Task
-def stream_regular_task():
+def stream_regular_task() -> None:
     regular_task = client.tasks.create(
         task="""
         Find top 10 Hacker News articles and return the title and url.
@@ -38,7 +38,7 @@ def stream_regular_task():
 
 
 # Structured Output
-def stream_structured_task():
+def stream_structured_task() -> None:
     class HackerNewsPost(BaseModel):
         title: str
         url: str

src/browser_use_sdk/resources/tasks.py

@@ -417,42 +417,42 @@ def retrieve(
     def stream(
         self,
         task_id: str,
-        structured_output_json: None | NotGiven = NOT_GIVEN,
+        structured_output_json: type[T],
         *,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-    ) -> Iterator[TaskView]: ...
+    ) -> Iterator[TaskViewWithOutput[T]]: ...
 
     @overload
     def stream(
         self,
         task_id: str,
-        structured_output_json: type[T],
+        structured_output_json: None | NotGiven = NOT_GIVEN,
         *,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-    ) -> Iterator[TaskViewWithOutput[T]]: ...
+    ) -> Iterator[TaskView]: ...
 
     def stream(
         self,
         task_id: str,
-        structured_output_json: type[BaseModel] | None | NotGiven = NOT_GIVEN,
+        structured_output_json: type[T] | None | NotGiven = NOT_GIVEN,
         *,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-    ) -> Iterator[Union[TaskView, TaskViewWithOutput[BaseModel]]]:
+    ) -> Iterator[TaskView | TaskViewWithOutput[T]]:
         """
         Stream the task view as it is updated until the task is finished.
         """
@@ -466,14 +466,15 @@ def stream(
         ):
             if structured_output_json is not None and isinstance(structured_output_json, type):
                 if res.done_output is None:
-                    yield TaskViewWithOutput[BaseModel](
+                    yield TaskViewWithOutput[T](
                         **res.model_dump(),
                         parsed_output=None,
                     )
                 else:
-                    parsed_output = structured_output_json.model_validate_json(res.done_output)
+                    schema: type[T] = structured_output_json
+                    parsed_output: T = schema.model_validate_json(res.done_output)
 
-                    yield TaskViewWithOutput[BaseModel](
+                    yield TaskViewWithOutput[T](
                         **res.model_dump(),
                         parsed_output=parsed_output,
                     )
@@ -1246,7 +1247,7 @@ def stream(
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-    ) -> AsyncIterator[TaskView] | AsyncIterator[TaskViewWithOutput[T]]:
+    ) -> AsyncIterator[TaskView | TaskViewWithOutput[T]]:
         """
         Stream the task view as it is updated until the task is finished.
         """
@@ -1259,7 +1260,6 @@ async def _gen() -> AsyncIterator[TaskView | TaskViewWithOutput[T]]:
                 extra_body=extra_body,
                 timeout=timeout,
             ):
-                # If a schema (type[T]) is passed, wrap with parsed_output[T]
                 if structured_output_json is not None and isinstance(structured_output_json, type):
                     if res.done_output is None:
                         yield TaskViewWithOutput[T](