Skip to content

Commit 073d226

Browse files
committed
feat: add new search engine avaiability and new tests
1 parent a8251bd commit 073d226

File tree

4 files changed

+113
-2
lines changed

4 files changed

+113
-2
lines changed
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
"""
2+
Example of custom graph using existing nodes
3+
"""
4+
5+
from scrapegraphai.models import Ollama
6+
from scrapegraphai.nodes import SearchInternetNode
7+
8+
# ************************************************
9+
# Define the configuration for the graph
10+
# ************************************************
11+
12+
graph_config = {
13+
"llm": {
14+
"model": "llama3",
15+
"temperature": 0,
16+
"streaming": True
17+
},
18+
"search_engine": "google",
19+
"max_results": 3,
20+
"verbose": True
21+
}
22+
23+
# ************************************************
24+
# Define the node
25+
# ************************************************
26+
27+
llm_model = Ollama(graph_config["llm"])
28+
29+
search_node = SearchInternetNode(
30+
input="user_input",
31+
output=["search_results"],
32+
node_config={
33+
"llm_model": llm_model,
34+
"search_engine": graph_config["search_engine"],
35+
"max_results": graph_config["max_results"],
36+
"verbose": graph_config["verbose"]
37+
}
38+
)
39+
40+
# ************************************************
41+
# Test the node
42+
# ************************************************
43+
44+
state = {
45+
"user_input": "What is the capital of France?"
46+
}
47+
48+
result = search_node.execute(state)
49+
50+
print(result)

scrapegraphai/nodes/search_internet_node.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def __init__(
4343
self.verbose = (
4444
False if node_config is None else node_config.get("verbose", False)
4545
)
46+
self.search_engine = node_config.get("search_engine", "google")
4647
self.max_results = node_config.get("max_results", 3)
4748

4849
def execute(self, state: dict) -> dict:
@@ -97,7 +98,8 @@ def execute(self, state: dict) -> dict:
9798

9899
self.logger.info(f"Search Query: {search_query}")
99100

100-
answer = search_on_web(query=search_query, max_results=self.max_results)
101+
answer = search_on_web(query=search_query, max_results=self.max_results,
102+
search_engine=self.search_engine)
101103

102104
if len(answer) == 0:
103105
# raise an exception if no answer is found

scrapegraphai/utils/research_web.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ def search_on_web(query: str, search_engine: str = "Google", max_results: int =
2626
>>> search_on_web("example query", search_engine="Google", max_results=5)
2727
['http://example.com', 'http://example.org', ...]
2828
29-
This function allows switching between Google and DuckDuckGo to perform internet searches, returning a list of result URLs.
29+
This function allows switching between Google and DuckDuckGo to perform
30+
internet searches, returning a list of result URLs.
3031
"""
3132

3233
if search_engine.lower() == "google":
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import unittest
2+
from scrapegraphai.models import Ollama
3+
from scrapegraphai.nodes import SearchInternetNode
4+
5+
class TestSearchInternetNode(unittest.TestCase):
6+
7+
def setUp(self):
8+
# Configuration for the graph
9+
self.graph_config = {
10+
"llm": {
11+
"model": "llama3",
12+
"temperature": 0,
13+
"streaming": True
14+
},
15+
"search_engine": "google",
16+
"max_results": 3,
17+
"verbose": True
18+
}
19+
20+
# Define the model
21+
self.llm_model = Ollama(self.graph_config["llm"])
22+
23+
# Initialize the SearchInternetNode
24+
self.search_node = SearchInternetNode(
25+
input="user_input",
26+
output=["search_results"],
27+
node_config={
28+
"llm_model": self.llm_model,
29+
"search_engine": self.graph_config["search_engine"],
30+
"max_results": self.graph_config["max_results"],
31+
"verbose": self.graph_config["verbose"]
32+
}
33+
)
34+
35+
def test_execute_search_node(self):
36+
# Initial state
37+
state = {
38+
"user_input": "What is the capital of France?"
39+
}
40+
41+
# Expected output
42+
expected_output = {
43+
"user_input": "What is the capital of France?",
44+
"search_results": [
45+
"https://en.wikipedia.org/wiki/Paris",
46+
"https://en.wikipedia.org/wiki/France",
47+
"https://en.wikipedia.org/wiki/%C3%8Ele-de-France"
48+
]
49+
}
50+
51+
# Execute the node
52+
result = self.search_node.execute(state)
53+
54+
# Assert the results
55+
self.assertEqual(result, expected_output)
56+
57+
if __name__ == "__main__":
58+
unittest.main()

0 commit comments

Comments
 (0)