Fix incorrect return annotation on StructuredData.get()

examples/scrape-logo/agent.py

@@ -32,10 +32,11 @@ def scrape_logo_url(url: str) -> str | None:
             # Do not raise on failure, we want to handle the case where no logo is found
             raise_on_failure=False,
         )
-        if data.success:
+        logo = data.get() if data.success else None
+        if isinstance(logo, Logo):
             # Case 1: structured output worked
-            logger.info(f"Logo found for {url}: {data.get().get_url(url)}")
-            return data.get().get_url(url)
+            logger.info(f"Logo found for {url}: {logo.get_url(url)}")
+            return logo.get_url(url)
         images = session.scrape(only_images=True)
         for image in images:
             # Case 2: there is a logo image in data.images

packages/notte-browser/src/notte_browser/session.py

@@ -1038,11 +1038,11 @@ async def ascrape(
         **params: Unpack[ScrapeMarkdownParamsDict],
     ) -> StructuredData[TBaseModel]: ...
 
-    # instructions only, raise_on_failure=True (default) -> unwrapped BaseModel
+    # instructions only, raise_on_failure=True (default) -> unwrapped BaseModel as dict
     @overload
     async def ascrape(
         self, *, instructions: str, raise_on_failure: Literal[True] = ..., **params: Unpack[ScrapeMarkdownParamsDict]
-    ) -> BaseModel: ...
+    ) -> dict[str, Any]: ...
 
     # instructions only, raise_on_failure=False -> wrapped StructuredData[BaseModel]
     @overload
@@ -1058,7 +1058,7 @@ async def ascrape(self, /, *, raise_on_failure: bool = True, **params: Unpack[Sc
     @track_usage("local.session.scrape")
     async def ascrape(
         self, *, raise_on_failure: bool = True, **params: Unpack[ScrapeParamsDict]
-    ) -> StructuredData[BaseModel] | BaseModel | str | list[ImageData]:
+    ) -> StructuredData[BaseModel] | BaseModel | dict[str, Any] | str | list[ImageData]:
         # Extract and convert response_format for the action (store as JSON schema)
         response_format = params.get("response_format")
         instructions = params.get("instructions")

packages/notte-core/src/notte_core/data/space.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Annotated, Any, Generic, Self, TypeVar
+from typing import Annotated, Any, Generic, Self, TypeVar, cast
 
 import requests
 from pydantic import BaseModel, Field, RootModel, model_serializer, model_validator
@@ -91,13 +91,22 @@ def serialize_model(self):
             result["data"] = self.data
         return result
 
-    def get(self) -> TBaseModel:
-        """Get the extracted data, raising ScrapeFailedError if extraction failed."""
+    def get(self) -> TBaseModel | dict[str, Any]:
+        """Get the extracted data, raising ScrapeFailedError if extraction failed.
+
+        Returns the validated model when the data was extracted against a user-provided
+        schema. When no schema was provided the payload is stored as a `DictBaseModel`
+        (i.e. `RootModel`) wrapper and this returns the unwrapped plain JSON object, which
+        is *not* a `BaseModel`. Callers must handle both shapes, e.g. by guarding with
+        `isinstance(data, BaseModel)` before calling `model_dump()`.
+        """
         if not self.success or self.data is None:
             raise ScrapeFailedError(self.error or "Unknown extraction error")
-        if isinstance(self.data, RootModel):
-            return self.data.root  # type: ignore[attr-defined]
-        return self.data
+        # local alias: keeps the narrowed type checkable instead of `RootModel[Unknown]`
+        data: TBaseModel | DictBaseModel = self.data
+        if isinstance(data, RootModel):
+            return cast(dict[str, Any], data.root)
+        return data
 
 
 class DataSpace(BaseModel):

packages/notte-sdk/src/notte_sdk/endpoints/page.py

@@ -208,7 +208,7 @@ def scrape(
                 # Validate against response_format if provided
                 if request.response_format is not None:
                     extracted_data_dict = (
-                        extracted_data.model_dump() if isinstance(extracted_data, BaseModel) else extracted_data  # pyright: ignore[reportUnnecessaryIsInstance]
+                        extracted_data.model_dump() if isinstance(extracted_data, BaseModel) else extracted_data
                     )
                     extracted_data = request.response_format.model_validate(extracted_data_dict)
                 return extracted_data

tests/test_structured_data.py

@@ -0,0 +1,51 @@
+import pytest
+from notte_core.data.space import DataSpace, StructuredData
+from notte_core.errors.processing import ScrapeFailedError
+from pydantic import BaseModel, RootModel
+
+
+class Verification(BaseModel):
+    status: str
+    code: str | None = None
+
+
+def test_get_returns_plain_dict_for_dict_payload():
+    """A dict payload is wrapped in a RootModel, so `get()` returns the unwrapped dict."""
+    structured = StructuredData(success=True, data={"status": "found", "code": "463092"})
+
+    assert isinstance(structured.data, RootModel)
+    data = structured.get()
+    assert isinstance(data, dict)
+    assert not isinstance(data, BaseModel)
+    assert data == {"status": "found", "code": "463092"}
+
+
+def test_get_returns_model_for_model_payload():
+    """A model payload is returned as-is, so `get()` returns a `BaseModel`."""
+    structured = StructuredData[Verification](success=True, data=Verification(status="found", code="463092"))
+
+    data = structured.get()
+    assert isinstance(data, Verification)
+    assert data.model_dump() == {"status": "found", "code": "463092"}
+
+
+def test_get_returns_plain_dict_for_deserialized_response():
+    """Structured data deserialized from a JSON response also yields a plain dict."""
+    space = DataSpace.model_validate(
+        {
+            "markdown": "# page",
+            "structured": {"success": True, "error": None, "data": {"status": "found", "code": "463092"}},
+        }
+    )
+
+    assert space.structured is not None
+    data = space.structured.get()
+    assert isinstance(data, dict)
+    assert data == {"status": "found", "code": "463092"}
+
+
+def test_get_raises_on_failed_extraction():
+    structured = StructuredData[Verification](success=False, error="boom", data=None)
+
+    with pytest.raises(ScrapeFailedError):
+        _ = structured.get()