Skip to content

Commit 49897c4

Browse files
committed
fix: generate answer node
1 parent 0255007 commit 49897c4

File tree

1 file changed

+6
-9
lines changed

1 file changed

+6
-9
lines changed

scrapegraphai/nodes/generate_answer_node.py

+6-9
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44
from typing import List, Optional
55
from json.decoder import JSONDecodeError
6+
import time
67
from langchain.prompts import PromptTemplate
78
from langchain_core.output_parsers import JsonOutputParser
89
from langchain_core.runnables import RunnableParallel
@@ -12,14 +13,13 @@
1213
from tqdm import tqdm
1314
from .base_node import BaseNode
1415
from ..utils.output_parser import get_structured_output_parser, get_pydantic_output_parser
16+
from requests.exceptions import Timeout
17+
from langchain.callbacks.manager import CallbackManager
18+
from langchain.callbacks import get_openai_callback
1519
from ..prompts import (
1620
TEMPLATE_CHUNKS, TEMPLATE_NO_CHUNKS, TEMPLATE_MERGE,
1721
TEMPLATE_CHUNKS_MD, TEMPLATE_NO_CHUNKS_MD, TEMPLATE_MERGE_MD
1822
)
19-
from langchain.callbacks.manager import CallbackManager
20-
from langchain.callbacks import get_openai_callback
21-
from requests.exceptions import Timeout
22-
import time
2323

2424
class GenerateAnswerNode(BaseNode):
2525
"""
@@ -82,11 +82,8 @@ def execute(self, state: dict) -> dict:
8282

8383
if self.node_config.get("schema", None) is not None:
8484
if isinstance(self.llm_model, ChatOpenAI):
85-
self.llm_model = self.llm_model.with_structured_output(
86-
schema=self.node_config["schema"]
87-
)
88-
output_parser = get_structured_output_parser(self.node_config["schema"])
89-
format_instructions = "NA"
85+
output_parser = get_pydantic_output_parser(self.node_config["schema"])
86+
format_instructions = output_parser.get_format_instructions()
9087
else:
9188
if not isinstance(self.llm_model, ChatBedrock):
9289
output_parser = get_pydantic_output_parser(self.node_config["schema"])

0 commit comments

Comments
 (0)