Skip to content

Commit e6e2ce6

Browse files
authored
Merge pull request #679 from LorenzoPaleari/output-parser-update
feat: output parser and pydantic update
2 parents 4f8b55d + eb89549 commit e6e2ce6

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+137
-138
lines changed

examples/anthropic/search_graph_schema_haiku.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SearchGraph
1010

1111
load_dotenv()

examples/anthropic/smart_scraper_schema_haiku.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import os
66
from typing import List
7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from dotenv import load_dotenv
99
from scrapegraphai.graphs import SmartScraperGraph
1010
from scrapegraphai.utils import prettify_exec_info

examples/azure/search_graph_schema_azure.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/azure/smart_scraper_schema_azure.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
import json
77
from typing import List
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from dotenv import load_dotenv
1010
from scrapegraphai.graphs import SmartScraperGraph
1111

examples/bedrock/search_graph_schema_bedrock.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from scrapegraphai.graphs import SearchGraph
55
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
66

7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from typing import List
99

1010
# ************************************************

examples/bedrock/smart_scraper_schema_bedrock.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
Basic example of scraping pipeline using SmartScraper
33
"""
44
from typing import List
5-
from langchain_core.pydantic_v1 import BaseModel, Field
5+
from pydantic import BaseModel, Field
66
from scrapegraphai.graphs import SmartScraperGraph
77
from scrapegraphai.utils import prettify_exec_info
88

examples/deepseek/search_graph_schema_deepseek.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/deepseek/smart_scraper_schema_deepseek.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import os
66
from typing import List
7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from dotenv import load_dotenv
99
from scrapegraphai.graphs import SmartScraperGraph
1010
from scrapegraphai.utils import prettify_exec_info

examples/extras/serch_graph_scehma.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
from dotenv import load_dotenv
77
from scrapegraphai.graphs import SearchGraph
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from typing import List
1010
load_dotenv()
1111

examples/fireworks/pdf_scraper_multi_fireworks.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import PdfScraperMultiGraph
1010

1111
load_dotenv()

examples/fireworks/script_generator_schema_fireworks.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import ScriptCreatorGraph
1010
from scrapegraphai.utils import prettify_exec_info
1111

examples/fireworks/search_graph_schema_fireworks.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/fireworks/smart_scraper_schema_fireworks.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os, json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SmartScraperGraph
1010

1111
load_dotenv()

examples/google_genai/search_graph_schema_gemini.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/google_genai/smart_scraper_schema_gemini.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import os
66
from typing import List
7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from dotenv import load_dotenv
99
from scrapegraphai.utils import prettify_exec_info
1010
from scrapegraphai.graphs import SmartScraperGraph

examples/google_vertexai/search_graph_schema_gemini.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/google_vertexai/smart_scraper_schema_gemini.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import os
66
from typing import List
7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from dotenv import load_dotenv
99
from scrapegraphai.utils import prettify_exec_info
1010
from scrapegraphai.graphs import SmartScraperGraph

examples/groq/search_graph_schema_groq.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/groq/smart_scraper_schema_groq.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import os, json
66
from typing import List
7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from dotenv import load_dotenv
99
from scrapegraphai.graphs import SmartScraperGraph
1010
from scrapegraphai.utils import prettify_exec_info

examples/integrations/indexify_node_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from dotenv import load_dotenv
99
load_dotenv()
1010

11-
from langchain_core.pydantic_v1 import BaseModel, Field
11+
from pydantic import BaseModel, Field
1212
from scrapegraphai.graphs import SmartScraperGraph
1313
from scrapegraphai.integrations import IndexifyNode
1414

examples/local_models/search_graph_schema_ollama.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from scrapegraphai.graphs import SearchGraph
55
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
66

7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from typing import List
99

1010
# ************************************************

examples/local_models/smart_scraper_schema_ollama.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
"""
44
import json
55
from typing import List
6-
from langchain_core.pydantic_v1 import BaseModel, Field
6+
from pydantic import BaseModel, Field
77
from scrapegraphai.graphs import SmartScraperGraph
88
from scrapegraphai.utils import prettify_exec_info
99

examples/mistral/pdf_scraper_multi_mistral.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import PdfScraperMultiGraph
1010

1111
load_dotenv()

examples/mistral/script_generator_schema_mistral.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from dotenv import load_dotenv
77
from scrapegraphai.graphs import ScriptCreatorGraph
88
from scrapegraphai.utils import prettify_exec_info
9-
from langchain_core.pydantic_v1 import BaseModel, Field
9+
from pydantic import BaseModel, Field
1010
from typing import List
1111

1212
load_dotenv()

examples/mistral/search_graph_schema_mistral.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

examples/mistral/smart_scraper_schema_mistral.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os, json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SmartScraperGraph
1010

1111
load_dotenv()

examples/nemotron/script_generator_schema_nemotron.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from scrapegraphai.graphs import ScriptCreatorGraph
88
from scrapegraphai.utils import prettify_exec_info
99

10-
from langchain_core.pydantic_v1 import BaseModel, Field
10+
from pydantic import BaseModel, Field
1111
from typing import List
1212

1313
load_dotenv()

examples/nemotron/search_graph_schema_nemotron.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

12-
from langchain_core.pydantic_v1 import BaseModel, Field
12+
from pydantic import BaseModel, Field
1313
from typing import List
1414

1515
# ************************************************

examples/nemotron/smart_scraper_schema_nemotron.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os, json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SmartScraperGraph
1010

1111
load_dotenv()

examples/oneapi/search_graph_schema_oneapi.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from scrapegraphai.graphs import SearchGraph
55
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
66

7-
from langchain_core.pydantic_v1 import BaseModel, Field
7+
from pydantic import BaseModel, Field
88
from typing import List
99

1010
# ************************************************

examples/oneapi/smart_scraper_schema_oneapi.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
Basic example of scraping pipeline using SmartScraper and OneAPI
33
"""
44
from typing import List
5-
from langchain_core.pydantic_v1 import BaseModel, Field
5+
from pydantic import BaseModel, Field
66
from scrapegraphai.graphs import SmartScraperGraph
77
from scrapegraphai.utils import prettify_exec_info
88

examples/openai/pdf_scraper_multi_openai.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import PdfScraperMultiGraph
1010

1111
load_dotenv()

examples/openai/script_generator_schema_openai.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from scrapegraphai.graphs import ScriptCreatorGraph
88
from scrapegraphai.utils import prettify_exec_info
99

10-
from langchain_core.pydantic_v1 import BaseModel, Field
10+
from pydantic import BaseModel, Field
1111
from typing import List
1212

1313
load_dotenv()

examples/openai/search_graph_schema_openai.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SearchGraph
1010
from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
1111

examples/openai/smart_scraper_schema_openai.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os, json
66
from typing import List
77
from dotenv import load_dotenv
8-
from langchain_core.pydantic_v1 import BaseModel, Field
8+
from pydantic import BaseModel, Field
99
from scrapegraphai.graphs import SmartScraperGraph
1010

1111
load_dotenv()

scrapegraphai/nodes/generate_answer_csv_node.py

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,14 @@
33
"""
44

55
from typing import List, Optional
6-
from pydantic.v1 import BaseModel as BaseModelV1
76
from langchain.prompts import PromptTemplate
87
from langchain_core.output_parsers import JsonOutputParser
98
from langchain_core.runnables import RunnableParallel
10-
from langchain_core.utils.pydantic import is_basemodel_subclass
119
from langchain_openai import ChatOpenAI
1210
from langchain_mistralai import ChatMistralAI
1311
from tqdm import tqdm
14-
from ..utils.logging import get_logger
1512
from .base_node import BaseNode
16-
from ..utils.llm_output_parser import typed_dict_output_parser, base_model_v2_output_parser, base_model_v1_output_parser
13+
from ..utils.output_parser import get_structured_output_parser, get_pydantic_output_parser
1714
from ..prompts import TEMPLATE_CHUKS_CSV, TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV
1815

1916
class GenerateAnswerCSVNode(BaseNode):
@@ -101,14 +98,10 @@ def execute(self, state):
10198
self.llm_model = self.llm_model.with_structured_output(
10299
schema = self.node_config["schema"]) # json schema works only on specific models
103100

104-
output_parser = typed_dict_output_parser
105-
if is_basemodel_subclass(self.node_config["schema"]):
106-
output_parser = base_model_v2_output_parser
107-
if issubclass(self.node_config["schema"], BaseModelV1):
108-
output_parser = base_model_v1_output_parser
101+
output_parser = get_structured_output_parser(self.node_config["schema"])
109102
format_instructions = "NA"
110103
else:
111-
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
104+
output_parser = get_pydantic_output_parser(self.node_config["schema"])
112105
format_instructions = output_parser.get_format_instructions()
113106

114107
else:

scrapegraphai/nodes/generate_answer_node.py

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,15 @@
22
GenerateAnswerNode Module
33
"""
44
from typing import List, Optional
5-
from pydantic.v1 import BaseModel as BaseModelV1
65
from langchain.prompts import PromptTemplate
76
from langchain_core.output_parsers import JsonOutputParser
87
from langchain_core.runnables import RunnableParallel
9-
from langchain_core.utils.pydantic import is_basemodel_subclass
108
from langchain_openai import ChatOpenAI, AzureChatOpenAI
119
from langchain_mistralai import ChatMistralAI
1210
from langchain_community.chat_models import ChatOllama
1311
from tqdm import tqdm
1412
from .base_node import BaseNode
15-
from ..utils.llm_output_parser import base_model_v1_output_parser, base_model_v2_output_parser, typed_dict_output_parser
13+
from ..utils.output_parser import get_structured_output_parser, get_pydantic_output_parser
1614
from ..prompts import (TEMPLATE_CHUNKS,
1715
TEMPLATE_NO_CHUNKS, TEMPLATE_MERGE,
1816
TEMPLATE_CHUNKS_MD, TEMPLATE_NO_CHUNKS_MD,
@@ -95,15 +93,11 @@ def execute(self, state: dict) -> dict:
9593
if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
9694
self.llm_model = self.llm_model.with_structured_output(
9795
schema = self.node_config["schema"]) # json schema works only on specific models
98-
99-
output_parser = typed_dict_output_parser
100-
if is_basemodel_subclass(self.node_config["schema"]):
101-
output_parser = base_model_v2_output_parser
102-
if issubclass(self.node_config["schema"], BaseModelV1):
103-
output_parser = base_model_v1_output_parser
96+
97+
output_parser = get_structured_output_parser(self.node_config["schema"])
10498
format_instructions = "NA"
10599
else:
106-
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
100+
output_parser = get_pydantic_output_parser(self.node_config["schema"])
107101
format_instructions = output_parser.get_format_instructions()
108102

109103
else:

0 commit comments

Comments
 (0)