feat: add new search engine avaiability and new tests

This commit is contained in:
Marco Vinciguerra 2024-06-18 14:35:13 +02:00
parent a8251bdb85
commit 073d226723
4 changed files with 113 additions and 2 deletions

View 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)

View File

@ -43,6 +43,7 @@ class SearchInternetNode(BaseNode):
self.verbose = (
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)
def execute(self, state: dict) -> dict:
@ -97,7 +98,8 @@ class SearchInternetNode(BaseNode):
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:
# raise an exception if no answer is found

View File

@ -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)
['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":

View 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()