Skip to content

Commit d29338b

Browse files
f-aguzziVinciGit00
andcommitted
feat(GenerateAnswerNode): built-in structured output through LangChain
Co-Authored-By: Marco Vinciguerra <[email protected]>
1 parent b48ee82 commit d29338b

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

scrapegraphai/nodes/generate_answer_node.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,12 @@
55
from langchain.prompts import PromptTemplate
66
from langchain_core.output_parsers import JsonOutputParser
77
from langchain_core.runnables import RunnableParallel
8-
from langchain_openai import ChatOpenAI
8+
from langchain_openai import ChatOpenAI, AzureChatOpenAI
9+
from langchain_mistralai import ChatMistralAI
10+
from langchain_anthropic import ChatAnthropic
11+
from langchain_groq import ChatGroq
12+
from langchain_fireworks import ChatFireworks
13+
from langchain_google_vertexai import ChatVertexAI
914
from langchain_community.chat_models import ChatOllama
1015
from tqdm import tqdm
1116
from ..utils.logging import get_logger
@@ -88,7 +93,9 @@ def execute(self, state: dict) -> dict:
8893
# Initialize the output parser
8994
if self.node_config.get("schema", None) is not None:
9095
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
91-
if isinstance(self.llm_model, ChatOpenAI) and (self.llm_model.model_name=="gpt-4o-mini" or self.llm_model.model_name=="gpt-4o-2024-08-06"):
96+
97+
# Use built-in structured output for providers that allow it
98+
if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI, ChatAnthropic, ChatFireworks, ChatGroq, ChatVertexAI)):
9299
self.llm_model = self.llm_model.with_structured_output(
93100
schema = self.node_config["schema"],
94101
method="json_schema")
@@ -98,7 +105,7 @@ def execute(self, state: dict) -> dict:
98105

99106
format_instructions = output_parser.get_format_instructions()
100107

101-
if isinstance(self.llm_model, ChatOpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper:
108+
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:
102109
template_no_chunks_prompt = TEMPLATE_NO_CHUNKS_MD
103110
template_chunks_prompt = TEMPLATE_CHUNKS_MD
104111
template_merge_prompt = TEMPLATE_MERGE_MD

0 commit comments

Comments
 (0)