feat: add new scrapegraph endpoint (#17709)

Author: VinciGit00

llama-index-integrations/tools/llama-index-tools-scrapegraph/examples/scrapegraph-search-llama-index.py

@@ -0,0 +1,46 @@
+"""Example of using scrapegraph search with LlamaIndex."""
+
+from llama_index.core.agent import ReActAgent
+from llama_index.tools.scrapegraph import ScrapegraphToolSpec
+from llama_index.core.llms import OpenAI
+
+# Initialize the tool spec with your API key
+SCRAPEGRAPH_API_KEY = "your-api-key-here"  # Replace with your actual API key
+tool_spec = ScrapegraphToolSpec()
+
+# Create a list of tools from the tool spec
+tools = tool_spec.to_tool_list()
+
+# Initialize the LLM
+llm = OpenAI(temperature=0)
+
+# Create an agent with the tools
+agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)
+
+
+# Example queries to demonstrate usage
+def run_search_example():
+    """Run an example search query using the scrapegraph tool."""
+    # Example 1: Basic search
+    query = "What are the latest developments in artificial intelligence?"
+    response = agent.chat(
+        f"Use the scrapegraph search tool to find information about: {query}"
+    )
+    print("\nSearch Results:")
+    print(f"Query: {query}")
+    print(f"Response: {response}")
+
+    # Example 2: More specific search
+    query = "What are the key features of Python 3.12?"
+    response = agent.chat(
+        f"Use the scrapegraph search tool to find detailed information about: {query}"
+    )
+    print("\nSearch Results:")
+    print(f"Query: {query}")
+    print(f"Response: {response}")
+
+
+if __name__ == "__main__":
+    # Make sure to set your API keys in environment variables
+    # os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
+    run_search_example()

llama-index-integrations/tools/llama-index-tools-scrapegraph/llama_index/tools/scrapegraph/base.py

@@ -14,7 +14,7 @@ class ScrapegraphToolSpec(BaseToolSpec):
     spec_functions = [
         "scrapegraph_smartscraper",
         "scrapegraph_markdownify",
-        "scrapegraph_local_scrape",
+        "scrapegraph_search",
     ]
 
     def scrapegraph_smartscraper(
@@ -56,16 +56,15 @@ def scrapegraph_markdownify(self, url: str, api_key: str) -> str:
 
         return client.markdownify(website_url=url)
 
-    def scrapegraph_local_scrape(self, text: str, api_key: str) -> str:
-        """Extract structured data from raw text using scrapegraph.
+    def scrapegraph_search(self, query: str, api_key: str) -> str:
+        """Perform a search query using scrapegraph.
 
         Args:
-            text (str): Raw text to process and extract data from
+            query (str): Search query to execute
             api_key (str): scrapegraph API key
 
         Returns:
-            str: Structured data extracted from the input text
+            str: Search results from scrapegraph
         """
         client = Client(api_key=api_key)
-
-        return client.local_scrape(text=text)
+        return client.search(query=query)

llama-index-integrations/tools/llama-index-tools-scrapegraph/pyproject.toml

@@ -29,7 +29,7 @@ license = "MIT"
 maintainers = ["Vincigit00"]
 name = "llama-index-tools-scrapegraphai"
 readme = "README.md"
-version = "0.1.0"
+version = "0.1.1"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"

llama-index-integrations/tools/llama-index-tools-scrapegraph/tests/test_tools_scrapegraph.py

@@ -52,3 +52,21 @@ def test_sync_scraping(tool_spec: ScrapegraphToolSpec, mock_sync_client: Mock):
         website_url=url, user_prompt=prompt, output_schema=schema
     )
     assert response == expected_response
+
+
+def test_search(tool_spec: ScrapegraphToolSpec, mock_sync_client: Mock):
+    """Test search functionality."""
+    # Test data
+    query = "test search query"
+    api_key = "sgai-0000-0000-0000-0000-0000-0000-0000-0000"
+    expected_response = "Search results"
+
+    # Configure mock
+    mock_sync_client.search.return_value = expected_response
+
+    # Execute test
+    response = tool_spec.scrapegraph_search(query=query, api_key=api_key)
+
+    # Verify
+    mock_sync_client.search.assert_called_once_with(query=query)
+    assert response == expected_response