mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-06-23 21:00:30 +08:00
feat: add new search engine avaiability and new tests
This commit is contained in:
parent
a8251bdb85
commit
073d226723
50
examples/single_node/search_internet_node.py
Normal file
50
examples/single_node/search_internet_node.py
Normal file
@ -0,0 +1,50 @@
|
|||||||
|
"""
|
||||||
|
Example of custom graph using existing nodes
|
||||||
|
"""
|
||||||
|
|
||||||
|
from scrapegraphai.models import Ollama
|
||||||
|
from scrapegraphai.nodes import SearchInternetNode
|
||||||
|
|
||||||
|
# ************************************************
|
||||||
|
# Define the configuration for the graph
|
||||||
|
# ************************************************
|
||||||
|
|
||||||
|
graph_config = {
|
||||||
|
"llm": {
|
||||||
|
"model": "llama3",
|
||||||
|
"temperature": 0,
|
||||||
|
"streaming": True
|
||||||
|
},
|
||||||
|
"search_engine": "google",
|
||||||
|
"max_results": 3,
|
||||||
|
"verbose": True
|
||||||
|
}
|
||||||
|
|
||||||
|
# ************************************************
|
||||||
|
# Define the node
|
||||||
|
# ************************************************
|
||||||
|
|
||||||
|
llm_model = Ollama(graph_config["llm"])
|
||||||
|
|
||||||
|
search_node = SearchInternetNode(
|
||||||
|
input="user_input",
|
||||||
|
output=["search_results"],
|
||||||
|
node_config={
|
||||||
|
"llm_model": llm_model,
|
||||||
|
"search_engine": graph_config["search_engine"],
|
||||||
|
"max_results": graph_config["max_results"],
|
||||||
|
"verbose": graph_config["verbose"]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# ************************************************
|
||||||
|
# Test the node
|
||||||
|
# ************************************************
|
||||||
|
|
||||||
|
state = {
|
||||||
|
"user_input": "What is the capital of France?"
|
||||||
|
}
|
||||||
|
|
||||||
|
result = search_node.execute(state)
|
||||||
|
|
||||||
|
print(result)
|
||||||
@ -43,6 +43,7 @@ class SearchInternetNode(BaseNode):
|
|||||||
self.verbose = (
|
self.verbose = (
|
||||||
False if node_config is None else node_config.get("verbose", False)
|
False if node_config is None else node_config.get("verbose", False)
|
||||||
)
|
)
|
||||||
|
self.search_engine = node_config.get("search_engine", "google")
|
||||||
self.max_results = node_config.get("max_results", 3)
|
self.max_results = node_config.get("max_results", 3)
|
||||||
|
|
||||||
def execute(self, state: dict) -> dict:
|
def execute(self, state: dict) -> dict:
|
||||||
@ -97,7 +98,8 @@ class SearchInternetNode(BaseNode):
|
|||||||
|
|
||||||
self.logger.info(f"Search Query: {search_query}")
|
self.logger.info(f"Search Query: {search_query}")
|
||||||
|
|
||||||
answer = search_on_web(query=search_query, max_results=self.max_results)
|
answer = search_on_web(query=search_query, max_results=self.max_results,
|
||||||
|
search_engine=self.search_engine)
|
||||||
|
|
||||||
if len(answer) == 0:
|
if len(answer) == 0:
|
||||||
# raise an exception if no answer is found
|
# raise an exception if no answer is found
|
||||||
|
|||||||
@ -26,7 +26,8 @@ def search_on_web(query: str, search_engine: str = "Google", max_results: int =
|
|||||||
>>> search_on_web("example query", search_engine="Google", max_results=5)
|
>>> search_on_web("example query", search_engine="Google", max_results=5)
|
||||||
['http://example.com', 'http://example.org', ...]
|
['http://example.com', 'http://example.org', ...]
|
||||||
|
|
||||||
This function allows switching between Google and DuckDuckGo to perform internet searches, returning a list of result URLs.
|
This function allows switching between Google and DuckDuckGo to perform
|
||||||
|
internet searches, returning a list of result URLs.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
if search_engine.lower() == "google":
|
if search_engine.lower() == "google":
|
||||||
|
|||||||
58
tests/nodes/search_internet_node_test.py
Normal file
58
tests/nodes/search_internet_node_test.py
Normal file
@ -0,0 +1,58 @@
|
|||||||
|
import unittest
|
||||||
|
from scrapegraphai.models import Ollama
|
||||||
|
from scrapegraphai.nodes import SearchInternetNode
|
||||||
|
|
||||||
|
class TestSearchInternetNode(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
# Configuration for the graph
|
||||||
|
self.graph_config = {
|
||||||
|
"llm": {
|
||||||
|
"model": "llama3",
|
||||||
|
"temperature": 0,
|
||||||
|
"streaming": True
|
||||||
|
},
|
||||||
|
"search_engine": "google",
|
||||||
|
"max_results": 3,
|
||||||
|
"verbose": True
|
||||||
|
}
|
||||||
|
|
||||||
|
# Define the model
|
||||||
|
self.llm_model = Ollama(self.graph_config["llm"])
|
||||||
|
|
||||||
|
# Initialize the SearchInternetNode
|
||||||
|
self.search_node = SearchInternetNode(
|
||||||
|
input="user_input",
|
||||||
|
output=["search_results"],
|
||||||
|
node_config={
|
||||||
|
"llm_model": self.llm_model,
|
||||||
|
"search_engine": self.graph_config["search_engine"],
|
||||||
|
"max_results": self.graph_config["max_results"],
|
||||||
|
"verbose": self.graph_config["verbose"]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_execute_search_node(self):
|
||||||
|
# Initial state
|
||||||
|
state = {
|
||||||
|
"user_input": "What is the capital of France?"
|
||||||
|
}
|
||||||
|
|
||||||
|
# Expected output
|
||||||
|
expected_output = {
|
||||||
|
"user_input": "What is the capital of France?",
|
||||||
|
"search_results": [
|
||||||
|
"https://en.wikipedia.org/wiki/Paris",
|
||||||
|
"https://en.wikipedia.org/wiki/France",
|
||||||
|
"https://en.wikipedia.org/wiki/%C3%8Ele-de-France"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
# Execute the node
|
||||||
|
result = self.search_node.execute(state)
|
||||||
|
|
||||||
|
# Assert the results
|
||||||
|
self.assertEqual(result, expected_output)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Loading…
Reference in New Issue
Block a user