Skip to content

Commit 5d20186

Browse files
committed
feat: add json as output
1 parent cab5f68 commit 5d20186

File tree

3 files changed

+9
-4
lines changed

3 files changed

+9
-4
lines changed

scrapegraphai/nodes/generate_answer_node.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from tqdm import tqdm
1313

1414
from ..utils.logging import get_logger
15-
from ..models import Ollama, Groq, OpenAI
15+
from ..models import Ollama
1616
# Imports from the library
1717
from .base_node import BaseNode
1818
from ..helpers import template_chunks, template_no_chunks, template_merge
@@ -44,7 +44,7 @@ def __init__(
4444
node_name: str = "GenerateAnswer",
4545
):
4646
super().__init__(node_name, "node", input, output, 2, node_config)
47-
47+
4848
self.llm_model = node_config["llm_model"]
4949

5050
if isinstance(node_config["llm_model"], Ollama):

scrapegraphai/nodes/generate_answer_omni_node.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
1111
from langchain_core.runnables import RunnableParallel
1212
from tqdm import tqdm
13-
13+
from ..models import Ollama
1414
# Imports from the library
1515
from .base_node import BaseNode
1616
from ..helpers.generate_answer_node_omni_prompts import template_no_chunk_omni, template_chunks_omni, template_merge_omni
@@ -44,6 +44,9 @@ def __init__(
4444
super().__init__(node_name, "node", input, output, 3, node_config)
4545

4646
self.llm_model = node_config["llm_model"]
47+
if isinstance(node_config["llm_model"], Ollama):
48+
self.llm_model.format="json"
49+
4750
self.verbose = (
4851
False if node_config is None else node_config.get("verbose", False)
4952
)

scrapegraphai/nodes/generate_answer_pdf_node.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
1111
from langchain_core.runnables import RunnableParallel
1212
from tqdm import tqdm
13-
13+
from ..models import Ollama
1414
from ..utils.logging import get_logger
1515

1616
# Imports from the library
@@ -59,6 +59,8 @@ def __init__(
5959
super().__init__(node_name, "node", input, output, 2, node_config)
6060

6161
self.llm_model = node_config["llm_model"]
62+
if isinstance(node_config["llm_model"], Ollama):
63+
self.llm_model.format="json"
6264
self.verbose = (
6365
False if node_config is None else node_config.get("verbose", False)
6466
)

0 commit comments

Comments
 (0)