feat: add scrapegraph-sdk integration (camel-ai#2206)

VinciGit00 · Wendong-Fan · web-flow · commit cb44039d329d · 2025-04-23T23:44:01.000+08:00
Co-authored-by: Wendong-Fan &lt;133094783+Wendong-Fan@users.noreply.github.com&gt;
diff --git a/camel/loaders/scrapegraph_reader.py b/camel/loaders/scrapegraph_reader.py
@@ -0,0 +1,96 @@
+# ========= Copyright 2023-2025@ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2025 @ CAMEL-AI.org. All Rights Reserved. =========
+
+import os
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel
+
+
+class ScrapeGraphAI:
+    r"""ScrapeGraphAI allows you to perform AI-powered web scraping and searching.
+
+    Args:
+        api_key (Optional[str]): API key for authenticating with the ScrapeGraphAI
+            API.
+
+    References:
+        https://scrapegraph.ai/
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+    ) -> None:
+        from scrapegraph_py import Client
+        from scrapegraph_py.logger import sgai_logger
+
+        self._api_key = api_key or os.environ.get("SCRAPEGRAPH_API_KEY")
+        sgai_logger.set_logging(level="INFO")
+        self.client = Client(api_key=self._api_key)
+
+    def search(
+        self,
+        user_prompt: str,
+    ) -> Dict[str, Any]:
+        r"""Perform an AI-powered web search using ScrapeGraphAI.
+
+        Args:
+            user_prompt (str): The search query or instructions.
+
+        Returns:
+            Dict[str, Any]: The search results including answer and reference URLs.
+
+        Raises:
+            RuntimeError: If the search process fails.
+        """
+        try:
+            response = self.client.searchscraper(user_prompt=user_prompt)
+            return response
+        except Exception as e:
+            raise RuntimeError(f"Failed to perform search: {e}")
+
+    def scrape(
+        self,
+        website_url: str,
+        user_prompt: str,
+        website_html: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        r"""Perform AI-powered web scraping using ScrapeGraphAI.
+
+        Args:
+            website_url (str): The URL to scrape.
+            user_prompt (str): Instructions for what data to extract.
+            website_html (Optional[str]): Optional HTML content to use instead of
+                fetching from the URL.
+
+        Returns:
+            Dict[str, Any]: The scraped data including request ID and result.
+
+        Raises:
+            RuntimeError: If the scrape process fails.
+        """
+        try:
+            response = self.client.smartscraper(
+                website_url=website_url,
+                user_prompt=user_prompt,
+                website_html=website_html,
+            )
+            return response
+        except Exception as e:
+            raise RuntimeError(f"Failed to perform scrape: {e}")
+
+    def close(self) -> None:
+        r"""Close the ScrapeGraphAI client connection."""
+        self.client.close()
diff --git a/examples/loaders/scrapegraph_example.py b/examples/loaders/scrapegraph_example.py
@@ -0,0 +1,100 @@
+# ========= Copyright 2023-2025@ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2025 @ CAMEL-AI.org. All Rights Reserved. =========
+
+"""
+Example demonstrating how to use the ScrapeGraphAI reader for web scraping and searching.
+
+This example shows:
+1. How to initialize the ScrapeGraphAI reader
+2. How to perform AI-powered web searches
+3. How to scrape websites with specific instructions
+4. How to handle errors and close the connection
+"""
+
+import os
+from typing import Dict, Any
+
+from camel.loaders.scrapegraph_reader import ScrapeGraphAI
+
+
+def search_example(api_key: str) -> Dict[str, Any]:
+    """Example of performing an AI-powered web search."""
+    # Initialize the ScrapeGraphAI reader
+    scraper = ScrapeGraphAI(api_key=api_key)
+    
+    try:
+        # Perform a search
+        search_query = "What are the latest developments in AI?"
+        result = scraper.search(user_prompt=search_query)
+        
+        print("\nSearch Results:")
+        print(f"Answer: {result.get('answer', 'No answer found')}")
+        print("References:")
+        for url in result.get('references', []):
+            print(f"- {url}")
+            
+        return result
+    finally:
+        # Always close the connection
+        scraper.close()
+
+
+def scrape_example(api_key: str) -> Dict[str, Any]:
+    """Example of scraping a website with specific instructions."""
+    # Initialize the ScrapeGraphAI reader
+    scraper = ScrapeGraphAI(api_key=api_key)
+    
+    try:
+        # Scrape a website with specific instructions
+        website_url = "https://example.com"
+        instructions = """
+        Extract the following information:
+        1. Main title of the page
+        2. All paragraph texts
+        3. Any links to other pages
+        """
+        
+        result = scraper.scrape(
+            website_url=website_url,
+            user_prompt=instructions
+        )
+        
+        print("\nScraping Results:")
+        print(f"Request ID: {result.get('request_id', 'No ID')}")
+        print("Extracted Data:")
+        print(result.get('result', {}))
+        
+        return result
+    finally:
+        # Always close the connection
+        scraper.close()
+
+
+def main():
+    # Get API key from environment variable or use a placeholder
+    api_key = os.environ.get("SCRAPEGRAPH_API_KEY", "your_api_key_here")
+    
+    if api_key == "your_api_key_here":
+        print("Please set your SCRAPEGRAPH_API_KEY environment variable")
+        return
+    
+    print("Running search example...")
+    search_example(api_key)
+    
+    print("\nRunning scrape example...")
+    scrape_example(api_key)
+
+
+if __name__ == "__main__":
+    main() 
diff --git a/pyproject.toml b/pyproject.toml
@@ -100,7 +100,8 @@ web_tools = [
     "playwright>=1.50.0",
     "html2text>=2024.2.26",
     "beautifulsoup4>=4,<5",
-    "exa-py>=1.10.0,<2"
+    "exa-py>=1.10.0,<2",
+    "scrapegraph-py>=1.12.0,<2",
 ]
 document_tools = [
     "numpy~=1.26",
diff --git a/test/loaders/test_scrapegraph_reader.py b/test/loaders/test_scrapegraph_reader.py
@@ -0,0 +1,105 @@
+# ========= Copyright 2023-2025@ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2025 @ CAMEL-AI.org. All Rights Reserved. =========
+
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from camel.loaders.scrapegraph_reader import ScrapeGraphAI
+
+
+@pytest.fixture
+def scrapegraph_ai():
+    with patch("camel.loaders.scrapegraph_reader.Client") as mock_client:
+        mock_client_instance = MagicMock()
+        mock_client.return_value = mock_client_instance
+        yield ScrapeGraphAI(api_key="test_api_key")
+
+
+def test_init_with_api_key():
+    with patch("camel.loaders.scrapegraph_reader.Client") as mock_client:
+        ScrapeGraphAI(api_key="test_api_key")
+        mock_client.assert_called_once_with(api_key="test_api_key")
+
+
+def test_init_with_env_var():
+    with patch("camel.loaders.scrapegraph_reader.Client") as mock_client, \
+         patch.dict(os.environ, {"SCRAPEGRAPH_API_KEY": "env_api_key"}):
+        ScrapeGraphAI()
+        mock_client.assert_called_once_with(api_key="env_api_key")
+
+
+def test_search_success(scrapegraph_ai):
+    mock_response = {"answer": "test answer", "references": ["url1", "url2"]}
+    scrapegraph_ai.client.searchscraper.return_value = mock_response
+
+    result = scrapegraph_ai.search("test query")
+    assert result == mock_response
+    scrapegraph_ai.client.searchscraper.assert_called_once_with(user_prompt="test query")
+
+
+def test_search_failure(scrapegraph_ai):
+    scrapegraph_ai.client.searchscraper.side_effect = Exception("Search failed")
+
+    with pytest.raises(RuntimeError, match="Failed to perform search: Search failed"):
+        scrapegraph_ai.search("test query")
+
+
+def test_scrape_success(scrapegraph_ai):
+    mock_response = {"request_id": "123", "result": {"data": "test data"}}
+    scrapegraph_ai.client.smartscraper.return_value = mock_response
+
+    result = scrapegraph_ai.scrape(
+        website_url="https://example.com",
+        user_prompt="Extract title and description"
+    )
+    assert result == mock_response
+    scrapegraph_ai.client.smartscraper.assert_called_once_with(
+        website_url="https://example.com",
+        user_prompt="Extract title and description",
+        website_html=None
+    )
+
+
+def test_scrape_with_html(scrapegraph_ai):
+    mock_response = {"request_id": "123", "result": {"data": "test data"}}
+    scrapegraph_ai.client.smartscraper.return_value = mock_response
+
+    result = scrapegraph_ai.scrape(
+        website_url="https://example.com",
+        user_prompt="Extract title and description",
+        website_html="<html>test</html>"
+    )
+    assert result == mock_response
+    scrapegraph_ai.client.smartscraper.assert_called_once_with(
+        website_url="https://example.com",
+        user_prompt="Extract title and description",
+        website_html="<html>test</html>"
+    )
+
+
+def test_scrape_failure(scrapegraph_ai):
+    scrapegraph_ai.client.smartscraper.side_effect = Exception("Scrape failed")
+
+    with pytest.raises(RuntimeError, match="Failed to perform scrape: Scrape failed"):
+        scrapegraph_ai.scrape(
+            website_url="https://example.com",
+            user_prompt="Extract title and description"
+        )
+
+
+def test_close(scrapegraph_ai):
+    scrapegraph_ai.close()
+    scrapegraph_ai.client.close.assert_called_once() 
