7
7
from stagehand .a11y .utils import get_accessibility_tree
8
8
from stagehand .llm .inference import extract as extract_inference
9
9
from stagehand .metrics import StagehandFunctionName # Changed import location
10
- from stagehand .types import (
11
- DefaultExtractSchema ,
12
- ExtractOptions ,
13
- ExtractResult ,
14
- )
10
+ from stagehand .types import DefaultExtractSchema , ExtractOptions , ExtractResult
15
11
from stagehand .utils import inject_urls , transform_url_strings_to_ids
16
12
17
13
T = TypeVar ("T" , bound = BaseModel )
@@ -153,15 +149,14 @@ async def extract(
153
149
validated_model_instance = schema .model_validate (raw_data_dict )
154
150
processed_data_payload = validated_model_instance # Payload is now the Pydantic model instance
155
151
except Exception as e :
156
- schema_name = getattr (schema , "__name__" , str (schema ))
157
152
self .logger .error (
158
- f"Failed to validate extracted data against schema { schema_name } : { e } . Keeping raw data dict in .data field."
153
+ f"Failed to validate extracted data against schema { schema . __name__ } : { e } . Keeping raw data dict in .data field."
159
154
)
160
155
161
- # Create ExtractResult object with extracted data as fields
162
- # Instead of trying to spread dict fields, always use the data field approach
163
- # This ensures result.data is properly set for the page.extract() method
164
- result = ExtractResult ( data = processed_data_payload )
156
+ # Create ExtractResult object
157
+ result = ExtractResult (
158
+ data = processed_data_payload ,
159
+ )
165
160
166
161
return result
167
162
@@ -171,4 +166,4 @@ async def _extract_page_text(self) -> ExtractResult:
171
166
172
167
tree = await get_accessibility_tree (self .stagehand_page , self .logger )
173
168
output_string = tree ["simplified" ]
174
- return ExtractResult (data = { "extraction" : output_string } )
169
+ return ExtractResult (data = output_string )
0 commit comments