77from stagehand .a11y .utils import get_accessibility_tree
88from stagehand .llm .inference import extract as extract_inference
99from 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
1511from stagehand .utils import inject_urls , transform_url_strings_to_ids
1612
1713T = TypeVar ("T" , bound = BaseModel )
@@ -153,15 +149,14 @@ async def extract(
153149 validated_model_instance = schema .model_validate (raw_data_dict )
154150 processed_data_payload = validated_model_instance # Payload is now the Pydantic model instance
155151 except Exception as e :
156- schema_name = getattr (schema , "__name__" , str (schema ))
157152 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."
159154 )
160155
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+ )
165160
166161 return result
167162
@@ -171,4 +166,4 @@ async def _extract_page_text(self) -> ExtractResult:
171166
172167 tree = await get_accessibility_tree (self .stagehand_page , self .logger )
173168 output_string = tree ["simplified" ]
174- return ExtractResult (data = { "extraction" : output_string } )
169+ return ExtractResult (data = output_string )
0 commit comments