wehub-resource-sync
110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
"""
|
|
Compatibility layer for scrapegraph-py SDK.
|
|
|
|
Supports both the v2 `Client` API (PR #82) and the newer `ScrapeGraphAI`
|
|
API (PR #84) which uses Pydantic request models and an ApiResult wrapper.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional, Type
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
def _detect_api() -> str:
|
|
try:
|
|
from scrapegraph_py import ScrapeGraphAI # noqa: F401
|
|
|
|
return "v3"
|
|
except ImportError:
|
|
pass
|
|
try:
|
|
from scrapegraph_py import Client # noqa: F401
|
|
|
|
return "v2"
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"scrapegraph_py is not installed. Install it with 'pip install scrapegraph-py'."
|
|
) from e
|
|
|
|
|
|
def _schema_to_dict(schema: Optional[Type[BaseModel]]) -> Optional[dict]:
|
|
if schema is None:
|
|
return None
|
|
if isinstance(schema, dict):
|
|
return schema
|
|
if isinstance(schema, type) and issubclass(schema, BaseModel):
|
|
return schema.model_json_schema()
|
|
return None
|
|
|
|
|
|
def _unwrap_result(result: Any) -> dict:
|
|
if hasattr(result, "status") and hasattr(result, "data"):
|
|
if result.status != "success":
|
|
raise RuntimeError(
|
|
getattr(result, "error", "scrapegraph-py request failed")
|
|
)
|
|
data = result.data
|
|
if hasattr(data, "model_dump"):
|
|
return data.model_dump(by_alias=True, exclude_none=True)
|
|
return data if isinstance(data, dict) else {"data": data}
|
|
return result
|
|
|
|
|
|
def extract(
|
|
api_key: Optional[str],
|
|
url: str,
|
|
prompt: str,
|
|
schema: Optional[Type[BaseModel]] = None,
|
|
) -> dict:
|
|
"""Call the scrapegraph-py extract endpoint across SDK versions."""
|
|
api = _detect_api()
|
|
|
|
if api == "v3":
|
|
from scrapegraph_py import ExtractRequest, ScrapeGraphAI
|
|
|
|
kwargs: dict[str, Any] = {"url": url, "prompt": prompt}
|
|
schema_dict = _schema_to_dict(schema)
|
|
if schema_dict is not None:
|
|
kwargs["schema_"] = schema_dict
|
|
with ScrapeGraphAI(api_key=api_key) as client:
|
|
return _unwrap_result(client.extract(ExtractRequest(**kwargs)))
|
|
|
|
from scrapegraph_py import Client
|
|
|
|
with Client(api_key=api_key) as client:
|
|
return client.extract(url=url, prompt=prompt, output_schema=schema)
|
|
|
|
|
|
def scrape(api_key: Optional[str], url: str) -> dict:
|
|
"""Call the scrapegraph-py scrape endpoint across SDK versions."""
|
|
api = _detect_api()
|
|
|
|
if api == "v3":
|
|
from scrapegraph_py import ScrapeGraphAI, ScrapeRequest
|
|
|
|
with ScrapeGraphAI(api_key=api_key) as client:
|
|
return _unwrap_result(client.scrape(ScrapeRequest(url=url)))
|
|
|
|
from scrapegraph_py import Client
|
|
|
|
with Client(api_key=api_key) as client:
|
|
return client.scrape(url=url)
|
|
|
|
|
|
def search(api_key: Optional[str], query: str) -> dict:
|
|
"""Call the scrapegraph-py search endpoint across SDK versions."""
|
|
api = _detect_api()
|
|
|
|
if api == "v3":
|
|
from scrapegraph_py import ScrapeGraphAI, SearchRequest
|
|
|
|
with ScrapeGraphAI(api_key=api_key) as client:
|
|
return _unwrap_result(client.search(SearchRequest(query=query)))
|
|
|
|
from scrapegraph_py import Client
|
|
|
|
with Client(api_key=api_key) as client:
|
|
return client.search(query=query)
|