From 5d20186bf20fb2384f2a9e7e81c2e875ff50a4f3 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Wed, 5 Jun 2024 09:20:20 +0200 Subject: [PATCH] feat: add json as output --- scrapegraphai/nodes/generate_answer_node.py | 4 ++-- scrapegraphai/nodes/generate_answer_omni_node.py | 5 ++++- scrapegraphai/nodes/generate_answer_pdf_node.py | 4 +++- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index 0db2d9fb..3aeb9ec5 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -12,7 +12,7 @@ from langchain_core.runnables import RunnableParallel from tqdm import tqdm from ..utils.logging import get_logger -from ..models import Ollama, Groq, OpenAI +from ..models import Ollama # Imports from the library from .base_node import BaseNode from ..helpers import template_chunks, template_no_chunks, template_merge @@ -44,7 +44,7 @@ class GenerateAnswerNode(BaseNode): node_name: str = "GenerateAnswer", ): super().__init__(node_name, "node", input, output, 2, node_config) - + self.llm_model = node_config["llm_model"] if isinstance(node_config["llm_model"], Ollama): diff --git a/scrapegraphai/nodes/generate_answer_omni_node.py b/scrapegraphai/nodes/generate_answer_omni_node.py index 12b8b90b..13eed843 100644 --- a/scrapegraphai/nodes/generate_answer_omni_node.py +++ b/scrapegraphai/nodes/generate_answer_omni_node.py @@ -10,7 +10,7 @@ from langchain.prompts import PromptTemplate from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser from langchain_core.runnables import RunnableParallel from tqdm import tqdm - +from ..models import Ollama # Imports from the library from .base_node import BaseNode from ..helpers.generate_answer_node_omni_prompts import template_no_chunk_omni, template_chunks_omni, template_merge_omni @@ -44,6 +44,9 @@ class GenerateAnswerOmniNode(BaseNode): super().__init__(node_name, "node", input, output, 3, node_config) self.llm_model = node_config["llm_model"] + if isinstance(node_config["llm_model"], Ollama): + self.llm_model.format="json" + self.verbose = ( False if node_config is None else node_config.get("verbose", False) ) diff --git a/scrapegraphai/nodes/generate_answer_pdf_node.py b/scrapegraphai/nodes/generate_answer_pdf_node.py index 4f055390..4f7de770 100644 --- a/scrapegraphai/nodes/generate_answer_pdf_node.py +++ b/scrapegraphai/nodes/generate_answer_pdf_node.py @@ -10,7 +10,7 @@ from langchain.prompts import PromptTemplate from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser from langchain_core.runnables import RunnableParallel from tqdm import tqdm - +from ..models import Ollama from ..utils.logging import get_logger # Imports from the library @@ -59,6 +59,8 @@ class GenerateAnswerPDFNode(BaseNode): super().__init__(node_name, "node", input, output, 2, node_config) self.llm_model = node_config["llm_model"] + if isinstance(node_config["llm_model"], Ollama): + self.llm_model.format="json" self.verbose = ( False if node_config is None else node_config.get("verbose", False) )