Merge pull request #626 from LorenzoPaleari/598-fix-pydantic-errors

Fixed pydantic errors when using `with_strctured_output`
ScrapeGraphAI · Sep 2, 2024 · 8442700 · 8442700
2 parents 5e99071 + 8e74ac5
commit 8442700
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 18 deletions.
diff --git a/scrapegraphai/nodes/generate_answer_csv_node.py b/scrapegraphai/nodes/generate_answer_csv_node.py
@@ -6,11 +6,13 @@
 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 ..prompts.generate_answer_node_csv_prompts import (TEMPLATE_CHUKS_CSV,
-                                                        TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV)
+from ..prompts import TEMPLATE_CHUKS_CSV, TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV
 
 class GenerateAnswerCSVNode(BaseNode):
     """
@@ -92,9 +94,24 @@ def execute(self, state):
 
         # Initialize the output parser
         if self.node_config.get("schema", None) is not None:
-            output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+
+            if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
+                self.llm_model = self.llm_model.with_structured_output(
+                    schema = self.node_config["schema"],
+                    method="function_calling") # json schema works only on specific models
+
+                # default parser to empty lambda function
+                output_parser = lambda x: x
+                if is_basemodel_subclass(self.node_config["schema"]):
+                    output_parser = dict
+                format_instructions = "NA"
+            else:
+                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                format_instructions = output_parser.get_format_instructions()
+
         else:
             output_parser = JsonOutputParser()
+            format_instructions = output_parser.get_format_instructions()
 
         TEMPLATE_NO_CHUKS_CSV_PROMPT = TEMPLATE_NO_CHUKS_CSV
         TEMPLATE_CHUKS_CSV_PROMPT = TEMPLATE_CHUKS_CSV
@@ -105,8 +122,6 @@ def execute(self, state):
             TEMPLATE_CHUKS_CSV_PROMPT = self.additional_info + TEMPLATE_CHUKS_CSV
             TEMPLATE_MERGE_CSV_PROMPT = self.additional_info + TEMPLATE_MERGE_CSV
 
-        format_instructions = output_parser.get_format_instructions()
-
         chains_dict = {}
 
         if len(doc) == 1:

diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py
@@ -1,16 +1,15 @@
 """
 GenerateAnswerNode Module
 """
-from sys import modules
 from typing import List, Optional
 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 ..utils.logging import get_logger
 from .base_node import BaseNode
 from ..prompts import TEMPLATE_CHUNKS, TEMPLATE_NO_CHUNKS, TEMPLATE_MERGE, TEMPLATE_CHUNKS_MD, TEMPLATE_NO_CHUNKS_MD, TEMPLATE_MERGE_MD
 
@@ -91,14 +90,20 @@ 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"],
-                    method="json_schema")
+                    method="function_calling") # json schema works only on specific models
+
+                # default parser to empty lambda function
+                output_parser = lambda x: x
+                if is_basemodel_subclass(self.node_config["schema"]):
+                    output_parser = dict
+                format_instructions = "NA"
             else:
                 output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                format_instructions = output_parser.get_format_instructions()
 
         else:
             output_parser = JsonOutputParser()
-
-        format_instructions = output_parser.get_format_instructions()
+            format_instructions = output_parser.get_format_instructions()
 
         if isinstance(self.llm_model, (ChatOpenAI, AzureChatOpenAI)) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper:
             template_no_chunks_prompt  = TEMPLATE_NO_CHUNKS_MD

diff --git a/scrapegraphai/nodes/generate_answer_omni_node.py b/scrapegraphai/nodes/generate_answer_omni_node.py
@@ -5,6 +5,9 @@
 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 langchain_community.chat_models import ChatOllama
 from .base_node import BaseNode
@@ -78,9 +81,25 @@ def execute(self, state: dict) -> dict:
 
         # Initialize the output parser
         if self.node_config.get("schema", None) is not None:
-            output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+
+            if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
+                self.llm_model = self.llm_model.with_structured_output(
+                    schema = self.node_config["schema"],
+                    method="function_calling") # json schema works only on specific models
+
+                # default parser to empty lambda function
+                output_parser = lambda x: x
+                if is_basemodel_subclass(self.node_config["schema"]):
+                    output_parser = dict
+                format_instructions = "NA"
+            else:
+                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                format_instructions = output_parser.get_format_instructions()
+
         else:
             output_parser = JsonOutputParser()
+            format_instructions = output_parser.get_format_instructions()
+
         TEMPLATE_NO_CHUNKS_OMNI_prompt = TEMPLATE_NO_CHUNKS_OMNI
         TEMPLATE_CHUNKS_OMNI_prompt = TEMPLATE_CHUNKS_OMNI
         TEMPLATE_MERGE_OMNI_prompt= TEMPLATE_MERGE_OMNI
@@ -90,7 +109,6 @@ def execute(self, state: dict) -> dict:
             TEMPLATE_CHUNKS_OMNI_prompt = self.additional_info + TEMPLATE_CHUNKS_OMNI_prompt
             TEMPLATE_MERGE_OMNI_prompt = self.additional_info + TEMPLATE_MERGE_OMNI_prompt
 
-        format_instructions = output_parser.get_format_instructions()
 
 
         chains_dict = {}

diff --git a/scrapegraphai/nodes/generate_answer_pdf_node.py b/scrapegraphai/nodes/generate_answer_pdf_node.py
@@ -5,6 +5,9 @@
 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 langchain_community.chat_models import ChatOllama
 from ..utils.logging import get_logger
@@ -93,9 +96,25 @@ def execute(self, state):
 
         # Initialize the output parser
         if self.node_config.get("schema", None) is not None:
-            output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+
+            if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
+                self.llm_model = self.llm_model.with_structured_output(
+                    schema = self.node_config["schema"],
+                    method="function_calling") # json schema works only on specific models
+
+                # default parser to empty lambda function
+                output_parser = lambda x: x
+                if is_basemodel_subclass(self.node_config["schema"]):
+                    output_parser = dict
+                format_instructions = "NA"
+            else:
+                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                format_instructions = output_parser.get_format_instructions()
+
         else:
             output_parser = JsonOutputParser()
+            format_instructions = output_parser.get_format_instructions()
+
         TEMPLATE_NO_CHUNKS_PDF_prompt = TEMPLATE_NO_CHUNKS_PDF
         TEMPLATE_CHUNKS_PDF_prompt = TEMPLATE_CHUNKS_PDF
         TEMPLATE_MERGE_PDF_prompt = TEMPLATE_MERGE_PDF
@@ -105,8 +124,6 @@ def execute(self, state):
             TEMPLATE_CHUNKS_PDF_prompt = self.additional_info + TEMPLATE_CHUNKS_PDF_prompt
             TEMPLATE_MERGE_PDF_prompt = self.additional_info + TEMPLATE_MERGE_PDF_prompt
 
-        format_instructions = output_parser.get_format_instructions()
-
         if len(doc) == 1:
             prompt = PromptTemplate(
                 template=TEMPLATE_NO_CHUNKS_PDF_prompt,

diff --git a/scrapegraphai/nodes/merge_answers_node.py b/scrapegraphai/nodes/merge_answers_node.py
@@ -4,6 +4,9 @@
 from typing import List, Optional
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.utils.pydantic import is_basemodel_subclass
+from langchain_openai import ChatOpenAI
+from langchain_mistralai import ChatMistralAI
 from ..utils.logging import get_logger
 from .base_node import BaseNode
 from ..prompts import TEMPLATE_COMBINED
@@ -68,11 +71,24 @@ def execute(self, state: dict) -> dict:
             answers_str += f"CONTENT WEBSITE {i+1}: {answer}\n"
 
         if self.node_config.get("schema", None) is not None:
-            output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+
+            if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
+                self.llm_model = self.llm_model.with_structured_output(
+                    schema = self.node_config["schema"],
+                    method="function_calling") # json schema works only on specific models
+
+                # default parser to empty lambda function
+                output_parser = lambda x: x
+                if is_basemodel_subclass(self.node_config["schema"]):
+                    output_parser = dict
+                format_instructions = "NA"
+            else:
+                output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
+                format_instructions = output_parser.get_format_instructions()
+
         else:
             output_parser = JsonOutputParser()
-
-        format_instructions = output_parser.get_format_instructions()
+            format_instructions = output_parser.get_format_instructions()
 
         prompt_template = PromptTemplate(
             template=TEMPLATE_COMBINED,