Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion scrapegraphai/utils/output_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
from typing import Any, Callable, Dict, List, Type, Union

from langchain_core.exceptions import OutputParserException
from langchain_core.outputs import Generation
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.outputs import Generation
from pydantic import BaseModel as BaseModelV2
from pydantic.v1 import BaseModel as BaseModelV1

Expand Down Expand Up @@ -58,6 +58,9 @@ def get_structured_output_parser(
Returns:
Callable: The output parser function.
"""
if isinstance(schema, dict):
return _dict_output_parser

if issubclass(schema, BaseModelV1):
return _base_model_v1_output_parser

Expand Down
14 changes: 14 additions & 0 deletions tests/utils/output_parser_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from scrapegraphai.utils.output_parser import (
TolerantJsonOutputParser,
_strip_doubled_braces,
get_structured_output_parser,
)


Expand All @@ -22,6 +23,19 @@ def test_strip_doubled_braces_ignores_unbalanced():
assert _strip_doubled_braces(text) == text


def test_structured_output_parser_accepts_json_schema():
schema = {
"title": "Person",
"type": "object",
"properties": {"name": {"type": "string"}},
}
output = {"name": "Ada"}

parser = get_structured_output_parser(schema)

assert parser(output) == output


def test_tolerant_parser_parses_clean_json_unchanged():
parser = TolerantJsonOutputParser()
assert parser.parse('{"content": "hi"}') == {"content": "hi"}
Expand Down