feat: fix tests

examples/single_node/robot_node.py

@@ -2,7 +2,7 @@
 Example of custom graph using existing nodes
 """
 
-from scrapegraphai.models import Ollama
+from langchain_community.chat_models import ChatOllama
 from scrapegraphai.nodes import RobotsNode
 
 # ************************************************
@@ -26,7 +26,7 @@ graph_config = {
 # Define the node
 # ************************************************
 
-llm_model = Ollama(graph_config["llm"])
+llm_model = ChatOllama(graph_config["llm"])
 
 robots_node = RobotsNode(
     input="url",

tests/nodes/robot_node_test.py

@@ -1,7 +1,6 @@
 import pytest
 from unittest.mock import MagicMock
-
-from scrapegraphai.models import Ollama
+from langchain_community.chat_models import ChatOllama
 from scrapegraphai.nodes import RobotsNode
 
 @pytest.fixture

tests/nodes/search_internet_node_test.py

@@ -1,5 +1,5 @@
 import unittest
-from scrapegraphai.models import Ollama
+from langchain_community.chat_models import ChatOllama
 from scrapegraphai.nodes import SearchInternetNode
 
 class TestSearchInternetNode(unittest.TestCase):
@@ -18,7 +18,7 @@ class TestSearchInternetNode(unittest.TestCase):
         }
 
         # Define the model
-        self.llm_model = Ollama(self.graph_config["llm"])
+        self.llm_model = ChatOllama(self.graph_config["llm"])
 
         # Initialize the SearchInternetNode
         self.search_node = SearchInternetNode(

tests/nodes/search_link_node_test.py

@@ -1,5 +1,5 @@
 import pytest
-from scrapegraphai.models import Ollama
+from langchain_community.chat_models import ChatOllama
 from scrapegraphai.nodes import SearchLinkNode
 from unittest.mock import patch, MagicMock
 
@@ -18,7 +18,7 @@ def setup():
     }
 
     # Instantiate the LLM model with the configuration
-    llm_model = Ollama(graph_config["llm"])
+    llm_model = ChatOllama(graph_config["llm"])
 
     # Define the SearchLinkNode with necessary configurations
     search_link_node = SearchLinkNode(