fix: fix generate answer node

ScrapeGraphAI · Nov 9, 2024 · d332e21 · d332e21
1 parent 6f1a88f
commit d332e21
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py
@@ -8,7 +8,6 @@
 from langchain_core.runnables import RunnableParallel
 from langchain_openai import ChatOpenAI, AzureChatOpenAI
 from langchain_aws import ChatBedrock
-from langchain_mistralai import ChatMistralAI
 from langchain_community.chat_models import ChatOllama
 from tqdm import tqdm
 from .base_node import BaseNode
@@ -77,7 +76,7 @@ def execute(self, state: dict) -> dict:
         doc = input_data[1]
 
         if self.node_config.get("schema", None) is not None:
-            if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)):
+            if isinstance(self.llm_model, ChatOpenAI):
                 self.llm_model = self.llm_model.with_structured_output(
                     schema=self.node_config["schema"]
                 )