add examples and tests for sdk

Author: VinciGit00

examples/.env.examples

@@ -0,0 +1 @@
+SCRAPEGRAPH_API_KEY="YOUR_SCRAPEGRAPH_API_KEY_HERE"

examples/scrape_example.py

@@ -0,0 +1,39 @@
+import os
+from dotenv import load_dotenv
+from scrapegraphaiapisdk.scrape import scrape
+from pydantic import BaseModel
+from typing import List
+
+# Load environment variables from .env file
+load_dotenv()
+
+class Product(BaseModel):
+    name: str
+    price: float
+    description: str
+
+class ProductList(BaseModel):
+    products: List[Product]
+
+def main():
+    # Get API key from environment variables
+    api_key = os.getenv("SCRAPEGRAPH_API_KEY")
+    
+    # URL to scrape
+    url = "https://example.com/products"
+    
+    # Natural language prompt
+    prompt = "Extract all products from this page including their names, prices, and descriptions"
+    
+    # Create schema
+    schema = ProductList
+    
+    # Make the request
+    try:
+        result = scrape(api_key, url, prompt, schema)
+        print(f"Scraped data: {result}")
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+if __name__ == "__main__":
+    main() 

examples/status_example.py

@@ -0,0 +1,20 @@
+import os
+from dotenv import load_dotenv
+from scrapegraphaiapisdk.status import status
+
+# Load environment variables from .env file
+load_dotenv()
+
+def main():
+    # Get API key from environment variables
+    api_key = os.getenv("SCRAPEGRAPH_API_KEY")
+    
+    # Check API status
+    try:
+        result = status(api_key)
+        print(f"API Status: {result}")
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+if __name__ == "__main__":
+    main() 

scrapegraphaiapisdk/scrape.py

@@ -1,25 +1,24 @@
 """
 This module provides a function to scrape and extract structured data from a webpage
 using the ScrapeGraph AI API. It allows specifying a schema for the output structure
-using either a dictionary or a Pydantic model.
+using a Pydantic model.
 """
 
-from typing import Union
 from pydantic import BaseModel
 import requests
 
-def scrape(api_key: str, url: str, prompt: str, schema: Union[dict, BaseModel, None] = None) -> str:
+def scrape(api_key: str, url: str, prompt: str, schema: BaseModel) -> str:
     """Scrape and extract structured data from a webpage using ScrapeGraph AI.
 
     Args:
         api_key (str): Your ScrapeGraph AI API key
         url (str): The URL of the webpage to scrape
         prompt (str): Natural language prompt describing what data to extract
-        schema (Union[dict, BaseModel, None], optional): Schema definition for the output structure.
-            Can be either a dictionary or a Pydantic model. If None, returns raw JSON.
+        schema (BaseModel): Pydantic model defining the output structure.
+            The model will be converted to JSON schema before making the request.
 
     Returns:
-        str: Extracted data in JSON format matching the provided schema (if specified)
+        str: Extracted data in JSON format matching the provided schema
     """
     endpoint = "https://api.scrapegraph.ai/v1/scrape"
     headers = {
@@ -30,10 +29,10 @@ def scrape(api_key: str, url: str, prompt: str, schema: Union[dict, BaseModel, N
     payload = {
         "url": url,
         "prompt": prompt,
-        "schema": schema.dict() if isinstance(schema, BaseModel) else schema if schema is not None else None
+        "schema": schema.model_json_schema()
     }
 
     response = requests.post(endpoint, headers=headers, json=payload)
-    response.raise_for_status()  # Raise an exception for bad status codes
+    response.raise_for_status()
 
     return response.text

tests/.env.examples

@@ -0,0 +1 @@
+SCRAPEGRAPH_API_KEY="YOUR_SCRAPEGRAPH_API_KEY_HERE"

tests/test_scrape.py

@@ -0,0 +1,44 @@
+import pytest
+import os
+from dotenv import load_dotenv
+from pydantic import BaseModel
+from typing import List
+from scrapegraphaiapisdk.scrape import scrape
+
+# Load environment variables from .env file
+load_dotenv()
+
+class Product(BaseModel):
+    name: str
+    price: float
+    description: str
+
+class ProductList(BaseModel):
+    products: List[Product]
+
+def test_scrape_successful():
+    api_key = os.getenv("SCRAPEGRAPH_API_KEY")
+    url = "https://example.com/products"
+    prompt = "Extract all products"
+    schema = ProductList
+    
+    # Mock the response
+    with pytest.mock.patch('requests.post') as mock_post:
+        mock_post.return_value.text = '{"products": [{"name": "Test Product", "price": 99.99, "description": "Test Description"}]}'
+        mock_post.return_value.raise_for_status.return_value = None
+        
+        result = scrape(api_key, url, prompt, schema)
+        assert isinstance(result, str)
+        assert "Test Product" in result
+
+def test_scrape_invalid_api_key():
+    api_key = "invalid_key"
+    url = "https://example.com/products"
+    prompt = "Extract all products"
+    schema = ProductList
+    
+    with pytest.mock.patch('requests.post') as mock_post:
+        mock_post.side_effect = Exception("Invalid API key")
+        
+        with pytest.raises(Exception):
+            scrape(api_key, url, prompt, schema) 

tests/test_status.py

@@ -0,0 +1,28 @@
+import pytest
+import os
+from dotenv import load_dotenv
+from scrapegraphaiapisdk.status import status
+
+# Load environment variables from .env file
+load_dotenv()
+
+def test_status_successful():
+    api_key = os.getenv("SCRAPEGRAPH_API_KEY")
+    
+    # Mock the response
+    with pytest.mock.patch('requests.get') as mock_get:
+        mock_get.return_value.text = '{"status": "ok"}'
+        mock_get.return_value.raise_for_status.return_value = None
+        
+        result = status(api_key)
+        assert isinstance(result, str)
+        assert "ok" in result
+
+def test_status_invalid_api_key():
+    api_key = "invalid_key"
+    
+    with pytest.mock.patch('requests.get') as mock_get:
+        mock_get.side_effect = Exception("Invalid API key")
+        
+        with pytest.raises(Exception):
+            status(api_key)