Skip to content

Commit cc9fd0f

Browse files
fix formatting
1 parent 3392d57 commit cc9fd0f

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

stagehand/handlers/extract_handler.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,11 @@
77
from stagehand.a11y.utils import get_accessibility_tree
88
from stagehand.llm.inference import extract as extract_inference
99
from stagehand.metrics import StagehandFunctionName # Changed import location
10-
from stagehand.schemas import DEFAULT_EXTRACT_SCHEMA as DefaultExtractSchema, ExtractOptions, ExtractResult
10+
from stagehand.schemas import (
11+
DEFAULT_EXTRACT_SCHEMA,
12+
ExtractOptions,
13+
ExtractResult,
14+
)
1115
from stagehand.utils import inject_urls, transform_url_strings_to_ids
1216

1317
T = TypeVar("T", bound=BaseModel)
@@ -93,7 +97,7 @@ async def extract(
9397
# TODO: Remove this once we have a better way to handle URLs
9498
transformed_schema, url_paths = transform_url_strings_to_ids(schema)
9599
else:
96-
transformed_schema = DefaultExtractSchema
100+
transformed_schema = DEFAULT_EXTRACT_SCHEMA
97101

98102
# Use inference to call the LLM
99103
extraction_response = extract_inference(
@@ -149,15 +153,15 @@ async def extract(
149153
validated_model_instance = schema.model_validate(raw_data_dict)
150154
processed_data_payload = validated_model_instance # Payload is now the Pydantic model instance
151155
except Exception as e:
152-
schema_name = getattr(schema, '__name__', str(schema))
156+
schema_name = getattr(schema, "__name__", str(schema))
153157
self.logger.error(
154158
f"Failed to validate extracted data against schema {schema_name}: {e}. Keeping raw data dict in .data field."
155159
)
156160

157161
# Create ExtractResult object with extracted data as fields
158162
if isinstance(processed_data_payload, dict):
159163
result = ExtractResult(**processed_data_payload)
160-
elif hasattr(processed_data_payload, 'model_dump'):
164+
elif hasattr(processed_data_payload, "model_dump"):
161165
# For Pydantic models, convert to dict and spread as fields
162166
result = ExtractResult(**processed_data_payload.model_dump())
163167
else:

0 commit comments

Comments
 (0)