fix: updated ollama structured output

ScrapeGraphAI · Jan 12, 2025 · 3b95911 · 3b95911
1 parent 7fee217
commit 3b95911
Show file tree

Hide file tree

Showing 7 changed files with 9 additions and 7 deletions.
diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_lite_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_lite_ollama.py
@@ -12,7 +12,6 @@
     "llm": {
         "model": "ollama/llama3.1",
         "temperature": 0,
-        "format": "json",
         "base_url": "http://localhost:11434",
     },
     "verbose": True,

diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_multi_concat_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_multi_concat_ollama.py
@@ -18,7 +18,6 @@
     "llm": {
         "model": "ollama/llama3.1",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         "base_url": "http://localhost:11434",  # set ollama URL arbitrarily
     },
     "verbose": True,

diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_multi_lite_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_multi_lite_ollama.py
@@ -15,7 +15,6 @@
     "llm": {
         "model": "ollama/llama3.1",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         "base_url": "http://localhost:11434",  # set ollama URL arbitrarily
     },
     "verbose": True,

diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_multi_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_multi_ollama.py
@@ -13,7 +13,6 @@
     "llm": {
         "model": "ollama/llama3.1",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         # "base_url": "http://localhost:11434", # set ollama URL arbitrarily
     },
     "verbose": True,

diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_ollama.py
@@ -13,9 +13,8 @@
     "llm": {
         "model": "ollama/llama3.2:3b",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         # "base_url": "http://localhost:11434", # set ollama URL arbitrarily
-        "model_tokens": 1024,
+        "model_tokens": 4096,
     },
     "verbose": True,
     "headless": False,

diff --git a/examples/smart_scraper_graph/ollama/smart_scraper_schema_ollama.py b/examples/smart_scraper_graph/ollama/smart_scraper_schema_ollama.py
@@ -27,7 +27,6 @@ class Projects(BaseModel):
     "llm": {
         "model": "ollama/llama3.1",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         # "base_url": "http://localhost:11434", # set ollama URL arbitrarily
     },
     "verbose": True,

diff --git a/scrapegraphai/nodes/merge_answers_node.py b/scrapegraphai/nodes/merge_answers_node.py
@@ -5,6 +5,7 @@
 from typing import List, Optional
 
 from langchain.prompts import PromptTemplate
+from langchain_community.chat_models import ChatOllama
 from langchain_core.output_parsers import JsonOutputParser
 from langchain_mistralai import ChatMistralAI
 from langchain_openai import ChatOpenAI
@@ -42,6 +43,13 @@ def __init__(
         super().__init__(node_name, "node", input, output, 2, node_config)
 
         self.llm_model = node_config["llm_model"]
+
+        if isinstance(self.llm_model, ChatOllama):
+            if self.node_config.get("schema", None) is None:
+                self.llm_model.format = "json"
+            else:
+                self.llm_model.format = self.node_config["schema"].model_json_schema()
+
         self.verbose = (
             False if node_config is None else node_config.get("verbose", False)
         )