Merge pull request #679 from LorenzoPaleari/output-parser-update

VinciGit00 · web-flow · commit e6e2ce6db460 · 2024-09-18T09:54:37.000+02:00
feat: output parser and pydantic update
diff --git a/examples/anthropic/search_graph_schema_haiku.py b/examples/anthropic/search_graph_schema_haiku.py
@@ -5,7 +5,7 @@
 import os
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SearchGraph
 
 load_dotenv()
diff --git a/examples/anthropic/smart_scraper_schema_haiku.py b/examples/anthropic/smart_scraper_schema_haiku.py
@@ -4,7 +4,7 @@
 
 import os
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
diff --git a/examples/azure/search_graph_schema_azure.py b/examples/azure/search_graph_schema_azure.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/azure/smart_scraper_schema_azure.py b/examples/azure/smart_scraper_schema_azure.py
@@ -5,7 +5,7 @@
 import os
 import json
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 
diff --git a/examples/bedrock/search_graph_schema_bedrock.py b/examples/bedrock/search_graph_schema_bedrock.py
@@ -4,7 +4,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/bedrock/smart_scraper_schema_bedrock.py b/examples/bedrock/smart_scraper_schema_bedrock.py
@@ -2,7 +2,7 @@
 Basic example of scraping pipeline using SmartScraper
 """
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
 
diff --git a/examples/deepseek/search_graph_schema_deepseek.py b/examples/deepseek/search_graph_schema_deepseek.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/deepseek/smart_scraper_schema_deepseek.py b/examples/deepseek/smart_scraper_schema_deepseek.py
@@ -4,7 +4,7 @@
 
 import os
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
diff --git a/examples/extras/serch_graph_scehma.py b/examples/extras/serch_graph_scehma.py
@@ -5,7 +5,7 @@
 import os
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SearchGraph
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 load_dotenv()
 
diff --git a/examples/fireworks/pdf_scraper_multi_fireworks.py b/examples/fireworks/pdf_scraper_multi_fireworks.py
@@ -5,7 +5,7 @@
 import json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import PdfScraperMultiGraph
 
 load_dotenv()
diff --git a/examples/fireworks/script_generator_schema_fireworks.py b/examples/fireworks/script_generator_schema_fireworks.py
@@ -5,7 +5,7 @@
 import os
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import ScriptCreatorGraph
 from scrapegraphai.utils import prettify_exec_info
 
diff --git a/examples/fireworks/search_graph_schema_fireworks.py b/examples/fireworks/search_graph_schema_fireworks.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/fireworks/smart_scraper_schema_fireworks.py b/examples/fireworks/smart_scraper_schema_fireworks.py
@@ -5,7 +5,7 @@
 import os, json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 
 load_dotenv()
diff --git a/examples/google_genai/search_graph_schema_gemini.py b/examples/google_genai/search_graph_schema_gemini.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/google_genai/smart_scraper_schema_gemini.py b/examples/google_genai/smart_scraper_schema_gemini.py
@@ -4,7 +4,7 @@
 
 import os
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.utils import prettify_exec_info
 from scrapegraphai.graphs import SmartScraperGraph
diff --git a/examples/google_vertexai/search_graph_schema_gemini.py b/examples/google_vertexai/search_graph_schema_gemini.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/google_vertexai/smart_scraper_schema_gemini.py b/examples/google_vertexai/smart_scraper_schema_gemini.py
@@ -4,7 +4,7 @@
 
 import os
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.utils import prettify_exec_info
 from scrapegraphai.graphs import SmartScraperGraph
diff --git a/examples/groq/search_graph_schema_groq.py b/examples/groq/search_graph_schema_groq.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/groq/smart_scraper_schema_groq.py b/examples/groq/smart_scraper_schema_groq.py
@@ -4,7 +4,7 @@
 
 import os, json
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
diff --git a/examples/integrations/indexify_node_example.py b/examples/integrations/indexify_node_example.py
@@ -8,7 +8,7 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.integrations import IndexifyNode
 
diff --git a/examples/local_models/search_graph_schema_ollama.py b/examples/local_models/search_graph_schema_ollama.py
@@ -4,7 +4,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/local_models/smart_scraper_schema_ollama.py b/examples/local_models/smart_scraper_schema_ollama.py
@@ -3,7 +3,7 @@
 """
 import json
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
 
diff --git a/examples/mistral/pdf_scraper_multi_mistral.py b/examples/mistral/pdf_scraper_multi_mistral.py
@@ -5,7 +5,7 @@
 import json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import PdfScraperMultiGraph
 
 load_dotenv()
diff --git a/examples/mistral/script_generator_schema_mistral.py b/examples/mistral/script_generator_schema_mistral.py
@@ -6,7 +6,7 @@
 from dotenv import load_dotenv
 from scrapegraphai.graphs import ScriptCreatorGraph
 from scrapegraphai.utils import prettify_exec_info
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 load_dotenv()
diff --git a/examples/mistral/search_graph_schema_mistral.py b/examples/mistral/search_graph_schema_mistral.py
@@ -5,7 +5,7 @@
 import os
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
diff --git a/examples/mistral/smart_scraper_schema_mistral.py b/examples/mistral/smart_scraper_schema_mistral.py
@@ -5,7 +5,7 @@
 import os, json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 
 load_dotenv()
diff --git a/examples/nemotron/script_generator_schema_nemotron.py b/examples/nemotron/script_generator_schema_nemotron.py
@@ -7,7 +7,7 @@
 from scrapegraphai.graphs import ScriptCreatorGraph
 from scrapegraphai.utils import prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 load_dotenv()
diff --git a/examples/nemotron/search_graph_schema_nemotron.py b/examples/nemotron/search_graph_schema_nemotron.py
@@ -9,7 +9,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/nemotron/smart_scraper_schema_nemotron.py b/examples/nemotron/smart_scraper_schema_nemotron.py
@@ -5,7 +5,7 @@
 import os, json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 
 load_dotenv()
diff --git a/examples/oneapi/search_graph_schema_oneapi.py b/examples/oneapi/search_graph_schema_oneapi.py
@@ -4,7 +4,7 @@
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 # ************************************************
diff --git a/examples/oneapi/smart_scraper_schema_oneapi.py b/examples/oneapi/smart_scraper_schema_oneapi.py
@@ -2,7 +2,7 @@
 Basic example of scraping pipeline using SmartScraper and OneAPI
 """
 from typing import List
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
 
diff --git a/examples/openai/pdf_scraper_multi_openai.py b/examples/openai/pdf_scraper_multi_openai.py
@@ -5,7 +5,7 @@
 import json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import PdfScraperMultiGraph
 
 load_dotenv()
diff --git a/examples/openai/script_generator_schema_openai.py b/examples/openai/script_generator_schema_openai.py
@@ -7,7 +7,7 @@
 from scrapegraphai.graphs import ScriptCreatorGraph
 from scrapegraphai.utils import prettify_exec_info
 
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from typing import List
 
 load_dotenv()
diff --git a/examples/openai/search_graph_schema_openai.py b/examples/openai/search_graph_schema_openai.py
@@ -5,7 +5,7 @@
 import os
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
 
diff --git a/examples/openai/smart_scraper_schema_openai.py b/examples/openai/smart_scraper_schema_openai.py
@@ -5,7 +5,7 @@
 import os, json
 from typing import List
 from dotenv import load_dotenv
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from scrapegraphai.graphs import SmartScraperGraph
 
 load_dotenv()
diff --git a/scrapegraphai/nodes/generate_answer_csv_node.py b/scrapegraphai/nodes/generate_answer_csv_node.py
@@ -3,17 +3,14 @@
 """
 
 from typing import List, Optional
-from pydantic.v1 import BaseModel as BaseModelV1
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.runnables import RunnableParallel
-from langchain_core.utils.pydantic import is_basemodel_subclass
 from langchain_openai import ChatOpenAI
 from langchain_mistralai import ChatMistralAI
 from tqdm import tqdm
-from ..utils.logging import get_logger
 from .base_node import BaseNode
-from ..utils.llm_output_parser import typed_dict_output_parser, base_model_v2_output_parser, base_model_v1_output_parser
+from ..utils.output_parser import get_structured_output_parser, get_pydantic_output_parser
 from ..prompts import TEMPLATE_CHUKS_CSV, TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV
 
 class GenerateAnswerCSVNode(BaseNode):
@@ -101,14 +98,10 @@ def execute(self, state):
                 self.llm_model = self.llm_model.with_structured_output(
                     schema = self.node_config["schema"]) # json schema works only on specific models
 
-                output_parser = typed_dict_output_parser
-                if is_basemodel_subclass(self.node_config["schema"]):
-                    output_parser = base_model_v2_output_parser
-                    if issubclass(self.node_config["schema"], BaseModelV1):
-                        output_parser = base_model_v1_output_parser
+                output_parser = get_structured_output_parser(self.node_config["schema"])
                 format_instructions = "NA"
             else:
-                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                output_parser = get_pydantic_output_parser(self.node_config["schema"])
                 format_instructions = output_parser.get_format_instructions()
 
         else:
diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py
@@ -2,17 +2,15 @@
 GenerateAnswerNode Module
 """
 from typing import List, Optional
-from pydantic.v1 import BaseModel as BaseModelV1
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.runnables import RunnableParallel
-from langchain_core.utils.pydantic import is_basemodel_subclass
 from langchain_openai import ChatOpenAI, AzureChatOpenAI
 from langchain_mistralai import ChatMistralAI
 from langchain_community.chat_models import ChatOllama
 from tqdm import tqdm
 from .base_node import BaseNode
-from ..utils.llm_output_parser import base_model_v1_output_parser, base_model_v2_output_parser, typed_dict_output_parser
+from ..utils.output_parser import get_structured_output_parser, get_pydantic_output_parser
 from ..prompts import (TEMPLATE_CHUNKS,
                        TEMPLATE_NO_CHUNKS, TEMPLATE_MERGE,
                        TEMPLATE_CHUNKS_MD, TEMPLATE_NO_CHUNKS_MD,
@@ -95,15 +93,11 @@ def execute(self, state: dict) -> dict:
             if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
                 self.llm_model = self.llm_model.with_structured_output(
                     schema = self.node_config["schema"]) # json schema works only on specific models
-
-                output_parser = typed_dict_output_parser
-                if is_basemodel_subclass(self.node_config["schema"]):
-                    output_parser = base_model_v2_output_parser
-                    if issubclass(self.node_config["schema"], BaseModelV1):
-                        output_parser = base_model_v1_output_parser
+                
+                output_parser = get_structured_output_parser(self.node_config["schema"])
                 format_instructions = "NA"
             else:
-                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                output_parser = get_pydantic_output_parser(self.node_config["schema"])
                 format_instructions = output_parser.get_format_instructions()
 
         else:
diff --git a/scrapegraphai/nodes/generate_answer_omni_node.py b/scrapegraphai/nodes/generate_answer_omni_node.py
diff --git a/scrapegraphai/nodes/generate_answer_pdf_node.py b/scrapegraphai/nodes/generate_answer_pdf_node.py
diff --git a/scrapegraphai/nodes/graph_iterator_node.py b/scrapegraphai/nodes/graph_iterator_node.py
diff --git a/scrapegraphai/nodes/merge_answers_node.py b/scrapegraphai/nodes/merge_answers_node.py
diff --git a/scrapegraphai/utils/llm_output_parser.py b/scrapegraphai/utils/llm_output_parser.py
diff --git a/scrapegraphai/utils/output_parser.py b/scrapegraphai/utils/output_parser.py
