Can't pass JSON schema to the extract function

[kushalpatil07]

Before submitting an issue, please:

• Check the documentation for relevant information
• Search existing issues to avoid duplicates

Environment Information

Please provide the following information to help us reproduce and resolve your issue:

Stagehand:

• Language/SDK: Python
• Stagehand version: 0.5.1

AI Provider:
OpenAI

• Provider: OpenAI
• Model: gpt-4o

Issue Description

According to the code, I expected to pass a json Schema as a argument to extract using the ExtractOptions.schema_definition, which accepts a dict as well. 

But in ExtractHandler.extract() the json Schema is passed on as a BaseModel type, and the code stops working in the validate path.

Steps to Reproduce

Minimal Reproduction Code

This is what I am using to try it out. The schema might be wrong, I'll fix that, if I get past this issue.

### Steps to Reproduce
# Define JSON schema directly without Pydantic classes
    json_schema = {
        "type": "object",
        "properties": {
            "list_of_links": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "link": {
                            "type": "string",
                            "format": "uri"
                        }
                    },
                    "required": ["link"]
                }
            }
        },
        "required": ["list_of_links"]
    }

    extraction = await stagehand.page.extract(
        options_or_instruction=ExtractOptions(schema_definition=json_schema, instruction ="extract the links of all videos that are present in this page.")
    )

Error Messages / Log trace

Traceback (most recent call last):
  File "/Users/kushal/Desktop/muscle-mem/.venv/lib/python3.11/site-packages/stagehand/handlers/extract_handler.py", line 155, in extract
    validated_model_instance = schema.model_validate(raw_data_dict)
                               ^^^^^^
  File "/Users/kushal/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/bdb.py", line 90, in trace_dispatch
    return self.dispatch_line(frame)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kushal/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/bdb.py", line 115, in dispatch_line
    if self.quitting: raise BdbQuit
                      ^^^^^^^^^^^^^
bdb.BdbQuit

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/kushal/Desktop/muscle-mem/test.py", line 75, in <module>
    asyncio.run(run())
  File "/Users/kushal/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/Users/kushal/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kushal/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/asyncio/base_events.py", line 654, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/Users/kushal/Desktop/muscle-mem/test.py", line 61, in run
    extraction = await stagehand.page.extract(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kushal/Desktop/muscle-mem/.venv/lib/python3.11/site-packages/stagehand/main.py", line 86, in wrapped
    return await attr(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kushal/Desktop/muscle-mem/.venv/lib/python3.11/site-packages/stagehand/page.py", line 384, in extract
    result = await self._extract_handler.extract(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kushal/Desktop/muscle-mem/.venv/lib/python3.11/site-packages/stagehand/handlers/extract_handler.py", line 159, in extract
    f"Failed to validate extracted data against schema {schema.__name__}: {e}. Keeping raw data dict in .data field."
                                                        ^^^^^^^^^^^^^^^
AttributeError: 'dict' object has no attribute '__name__'

[kushalpatil07]

My use case requires to define the schema of a object at runtime and then capture it. How can I achieve this? Imagining I can create a schema in json at anytime

[miguelg719]

Hi @kushalpatil07 yes unfortunately litellm doesn't work well with json schemas so we default to Pydantic objects. We opened an issue and are waiting to hear back BerriAI/litellm#10674

Working on a workaround, but for now, it's probably best to pass a Pydantic object

[kushalpatil07]

Was able to get something working atleast for me like this.
from pydantic import BaseModel, create_model, HttpUrl
from typing import List

# Your JSON schema
json_schema = {
    "type": "object",
    "properties": {
        "list_of_links": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "link": {
                        "type": "string",
                        "format": "uri"
                    }
                },
                "required": ["link"]
            }
        }
    },
    "required": ["list_of_links"]
}

# Mapping JSON schema types to Pydantic field types
type_map = {
    ("string", "uri"): HttpUrl,
    ("string", None): str,
    ("array", None): list,
    ("object", None): dict,
    ("integer", None): int,
    ("number", None): float,
    ("boolean", None): bool,
}

def schema_to_pydantic(schema: dict, name: str = "Model") -> type[BaseModel]:
    """Convert (a subset of) JSON Schema to a Pydantic model dynamically."""

    fields = {}
    for prop, details in schema.get("properties", {}).items():
        schema_type = details.get("type")
        schema_format = details.get("format")
        py_type = type_map.get((schema_type, schema_format)) or type_map.get((schema_type, None))

        if schema_type == "array" and "items" in details:
            # recursively create model for array items
            item_schema = details["items"]
            item_model = schema_to_pydantic(item_schema, name=f"{prop.capitalize()}Item")
            py_type = List[item_model]

        elif schema_type == "object":
            # recursively create model for nested object
            py_type = schema_to_pydantic(details, name=prop.capitalize())

        fields[prop] = (py_type, ... if prop in schema.get("required", []) else None)

    return create_model(name, **fields)

[miguelg719]

@kushalpatil07 want to push a PR?

[kushalpatil07]

Yes I can. Does this look good? A little bit of guidance on where to place the files. And I can open a good PR