Skip to content

Commit 1db164e

Browse files
committed
feat: fix tests
1 parent 88710f1 commit 1db164e

File tree

4 files changed

+7
-8
lines changed

4 files changed

+7
-8
lines changed

examples/single_node/robot_node.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
Example of custom graph using existing nodes
33
"""
44

5-
from scrapegraphai.models import Ollama
5+
from langchain_community.chat_models import ChatOllama
66
from scrapegraphai.nodes import RobotsNode
77

88
# ************************************************
@@ -26,7 +26,7 @@
2626
# Define the node
2727
# ************************************************
2828

29-
llm_model = Ollama(graph_config["llm"])
29+
llm_model = ChatOllama(graph_config["llm"])
3030

3131
robots_node = RobotsNode(
3232
input="url",

tests/nodes/robot_node_test.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import pytest
22
from unittest.mock import MagicMock
3-
4-
from scrapegraphai.models import Ollama
3+
from langchain_community.chat_models import ChatOllama
54
from scrapegraphai.nodes import RobotsNode
65

76
@pytest.fixture

tests/nodes/search_internet_node_test.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import unittest
2-
from scrapegraphai.models import Ollama
2+
from langchain_community.chat_models import ChatOllama
33
from scrapegraphai.nodes import SearchInternetNode
44

55
class TestSearchInternetNode(unittest.TestCase):
@@ -18,7 +18,7 @@ def setUp(self):
1818
}
1919

2020
# Define the model
21-
self.llm_model = Ollama(self.graph_config["llm"])
21+
self.llm_model = ChatOllama(self.graph_config["llm"])
2222

2323
# Initialize the SearchInternetNode
2424
self.search_node = SearchInternetNode(

tests/nodes/search_link_node_test.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import pytest
2-
from scrapegraphai.models import Ollama
2+
from langchain_community.chat_models import ChatOllama
33
from scrapegraphai.nodes import SearchLinkNode
44
from unittest.mock import patch, MagicMock
55

@@ -18,7 +18,7 @@ def setup():
1818
}
1919

2020
# Instantiate the LLM model with the configuration
21-
llm_model = Ollama(graph_config["llm"])
21+
llm_model = ChatOllama(graph_config["llm"])
2222

2323
# Define the SearchLinkNode with necessary configurations
2424
search_link_node = SearchLinkNode(

0 commit comments

Comments
 (0)