Skip to content

feat: output parser and pydantic update #679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
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
2 changes: 1 addition & 1 deletion examples/anthropic/search_graph_schema_haiku.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/anthropic/smart_scraper_schema_haiku.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/azure/search_graph_schema_azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/azure/smart_scraper_schema_azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/bedrock/search_graph_schema_bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/bedrock/smart_scraper_schema_bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/deepseek/search_graph_schema_deepseek.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/deepseek/smart_scraper_schema_deepseek.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/extras/serch_graph_scehma.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion examples/fireworks/pdf_scraper_multi_fireworks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/fireworks/script_generator_schema_fireworks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/fireworks/search_graph_schema_fireworks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/fireworks/smart_scraper_schema_fireworks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/google_genai/search_graph_schema_gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/google_genai/smart_scraper_schema_gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/google_vertexai/search_graph_schema_gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/google_vertexai/smart_scraper_schema_gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/groq/search_graph_schema_groq.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/groq/smart_scraper_schema_groq.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/integrations/indexify_node_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/local_models/search_graph_schema_ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/local_models/smart_scraper_schema_ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/mistral/pdf_scraper_multi_mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/mistral/script_generator_schema_mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/mistral/search_graph_schema_mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/mistral/smart_scraper_schema_mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/nemotron/script_generator_schema_nemotron.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/nemotron/search_graph_schema_nemotron.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/nemotron/smart_scraper_schema_nemotron.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/oneapi/search_graph_schema_oneapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ************************************************
Expand Down
2 changes: 1 addition & 1 deletion examples/oneapi/smart_scraper_schema_oneapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/openai/pdf_scraper_multi_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/openai/script_generator_schema_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion examples/openai/search_graph_schema_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion examples/openai/smart_scraper_schema_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
13 changes: 3 additions & 10 deletions scrapegraphai/nodes/generate_answer_csv_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand Down
14 changes: 4 additions & 10 deletions scrapegraphai/nodes/generate_answer_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
Loading