diff --git a/README.md b/README.md index 2a10a6c..cae0399 100644 --- a/README.md +++ b/README.md @@ -152,9 +152,15 @@ res = sgai.search( schema={...}, # optional time_range="past_week", # optional location_geo_code="us", # optional + allowed_types=["text/html", "application/pdf"], # optional MIME allowlist ) ``` +By default `search` accepts every supported content type, including PDFs, and processes up to 25 +pages per PDF. You do not need to send `processors` or `max_pages` for this default. Use +`allowed_types` to restrict accepted MIME types. Only configure `processors` to override the cap; +`PdfProcessor()` also defaults to 25, while `max_pages` accepts `1`–`500`, or `-1` for no page limit. + ### crawl Crawl a website and its linked pages. diff --git a/src/scrapegraph_py/__init__.py b/src/scrapegraph_py/__init__.py index 0e88ba0..063b19b 100644 --- a/src/scrapegraph_py/__init__.py +++ b/src/scrapegraph_py/__init__.py @@ -39,6 +39,7 @@ MonitorTickEntry, MonitorTickStatus, MonitorUpdateRequest, + PdfProcessor, ScrapeBrandingFormatEntry, ScrapeCaptureFormat, ScrapeContentFormat, @@ -104,6 +105,7 @@ "FetchConfig", "FetchContentType", "FetchMode", + "PdfProcessor", "FormatConfig", "ScrapeFormat", "ScrapeContentFormat", diff --git a/src/scrapegraph_py/async_client.py b/src/scrapegraph_py/async_client.py index f3afa2f..71dd637 100644 --- a/src/scrapegraph_py/async_client.py +++ b/src/scrapegraph_py/async_client.py @@ -34,6 +34,7 @@ MonitorCreateRequest, MonitorResponse, MonitorUpdateRequest, + PdfProcessor, ScrapeRequest, ScrapeResponse, SearchRequest, @@ -96,7 +97,8 @@ async def start( allow_external: bool | None = None, include_patterns: list[str] | None = None, exclude_patterns: list[str] | None = None, - content_types: list[FetchContentType] | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, fetch_config: FetchConfig | None = None, ) -> ApiResult[CrawlResponse]: req = CrawlRequest( @@ -109,7 +111,8 @@ async def start( allow_external=allow_external, include_patterns=include_patterns, exclude_patterns=exclude_patterns, - content_types=content_types, + allowed_types=allowed_types, + processors=processors, fetch_config=fetch_config, ) ) @@ -347,6 +350,8 @@ async def scrape( formats: list[FormatConfig] | None = None, fetch_config: FetchConfig | None = None, content_type: FetchContentType | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, ) -> ApiResult[ScrapeResponse]: req = ScrapeRequest( **_compact( @@ -354,6 +359,8 @@ async def scrape( formats=formats, fetch_config=fetch_config, content_type=content_type, + allowed_types=allowed_types, + processors=processors, ) ) return await self._post("/scrape", req, ScrapeResponse) @@ -369,6 +376,8 @@ async def extract( mode: HtmlMode | None = None, fetch_config: FetchConfig | None = None, content_type: FetchContentType | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, ) -> ApiResult[ExtractResponse]: req = ExtractRequest( **_compact( @@ -380,6 +389,8 @@ async def extract( mode=mode, fetch_config=fetch_config, content_type=content_type, + allowed_types=allowed_types, + processors=processors, ) ) return await self._post("/extract", req, ExtractResponse) @@ -395,6 +406,8 @@ async def search( schema: dict[str, object] | None = None, fetch_config: FetchConfig | None = None, location_geo_code: str | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, time_range: TimeRange | None = None, ) -> ApiResult[SearchResponse]: req = SearchRequest( @@ -407,6 +420,8 @@ async def search( schema=schema, fetch_config=fetch_config, location_geo_code=location_geo_code, + allowed_types=allowed_types, + processors=processors, time_range=time_range, ) ) diff --git a/src/scrapegraph_py/client.py b/src/scrapegraph_py/client.py index b6bf53b..b583e7f 100644 --- a/src/scrapegraph_py/client.py +++ b/src/scrapegraph_py/client.py @@ -34,6 +34,7 @@ MonitorCreateRequest, MonitorResponse, MonitorUpdateRequest, + PdfProcessor, ScrapeRequest, ScrapeResponse, SearchRequest, @@ -96,7 +97,8 @@ def start( allow_external: bool | None = None, include_patterns: list[str] | None = None, exclude_patterns: list[str] | None = None, - content_types: list[FetchContentType] | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, fetch_config: FetchConfig | None = None, ) -> ApiResult[CrawlResponse]: req = CrawlRequest( @@ -109,7 +111,8 @@ def start( allow_external=allow_external, include_patterns=include_patterns, exclude_patterns=exclude_patterns, - content_types=content_types, + allowed_types=allowed_types, + processors=processors, fetch_config=fetch_config, ) ) @@ -347,6 +350,8 @@ def scrape( formats: list[FormatConfig] | None = None, fetch_config: FetchConfig | None = None, content_type: FetchContentType | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, ) -> ApiResult[ScrapeResponse]: req = ScrapeRequest( **_compact( @@ -354,6 +359,8 @@ def scrape( formats=formats, fetch_config=fetch_config, content_type=content_type, + allowed_types=allowed_types, + processors=processors, ) ) return self._post("/scrape", req, ScrapeResponse) @@ -369,6 +376,8 @@ def extract( mode: HtmlMode | None = None, fetch_config: FetchConfig | None = None, content_type: FetchContentType | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, ) -> ApiResult[ExtractResponse]: req = ExtractRequest( **_compact( @@ -380,6 +389,8 @@ def extract( mode=mode, fetch_config=fetch_config, content_type=content_type, + allowed_types=allowed_types, + processors=processors, ) ) return self._post("/extract", req, ExtractResponse) @@ -395,6 +406,8 @@ def search( schema: dict[str, object] | None = None, fetch_config: FetchConfig | None = None, location_geo_code: str | None = None, + allowed_types: list[FetchContentType] | None = None, + processors: list[PdfProcessor] | None = None, time_range: TimeRange | None = None, ) -> ApiResult[SearchResponse]: req = SearchRequest( @@ -407,6 +420,8 @@ def search( schema=schema, fetch_config=fetch_config, location_geo_code=location_geo_code, + allowed_types=allowed_types, + processors=processors, time_range=time_range, ) ) diff --git a/src/scrapegraph_py/schemas.py b/src/scrapegraph_py/schemas.py index 20301e6..971dbfc 100644 --- a/src/scrapegraph_py/schemas.py +++ b/src/scrapegraph_py/schemas.py @@ -2,7 +2,7 @@ from typing import Annotated, Generic, Literal, TypeVar -from pydantic import BaseModel, ConfigDict, Field, HttpUrl, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, HttpUrl, model_validator from pydantic.alias_generators import to_camel T = TypeVar("T") @@ -27,6 +27,10 @@ class ResponseModel(BaseModel): FetchContentType = Literal[ "text/html", "application/json", + "text/markdown", + "text/plain", + "text/csv", + "application/x-latex", "application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/vnd.openxmlformats-officedocument.presentationml.presentation", @@ -42,11 +46,33 @@ class ResponseModel(BaseModel): "application/epub+zip", "application/rtf", "application/vnd.oasis.opendocument.text", - "text/csv", - "text/plain", - "application/x-latex", ] + +class PdfProcessor(CamelModel): + type: Literal["pdf"] = "pdf" + max_pages: Literal[-1] | Annotated[int, Field(ge=1, le=500)] = 25 + + +def _unique_allowed_types(values: list[FetchContentType]) -> list[FetchContentType]: + if len(values) != len(set(values)): + raise ValueError("duplicate allowed types not allowed") + return values + + +def _unique_processors(values: list[PdfProcessor]) -> list[PdfProcessor]: + types = [processor.type for processor in values] + if len(types) != len(set(types)): + raise ValueError("duplicate processor types not allowed") + return values + + +AllowedTypes = Annotated[ + list[FetchContentType], Field(min_length=1), AfterValidator(_unique_allowed_types) +] +Processors = Annotated[list[PdfProcessor], Field(min_length=1), AfterValidator(_unique_processors)] + + ScrapeContentFormat = Literal["markdown", "html", "links", "images", "summary", "json", "branding"] ScrapeCaptureFormat = Literal["screenshot"] ScrapeFormat = Literal[ @@ -147,6 +173,8 @@ class ScrapeBrandingFormatEntry(CamelModel): class ScrapeRequest(CamelModel): url: HttpUrl content_type: FetchContentType | None = None + allowed_types: AllowedTypes | None = None + processors: Processors | None = None fetch_config: FetchConfig | None = None formats: list[ScrapeFormatEntry] = Field(default_factory=lambda: [ScrapeMarkdownFormatEntry()]) @@ -166,6 +194,8 @@ class ExtractRequestBase(CamelModel): prompt: Annotated[str, Field(min_length=1, max_length=10000)] schema_: dict[str, object] | None = Field(default=None, alias="schema") content_type: FetchContentType | None = None + allowed_types: AllowedTypes | None = None + processors: Processors | None = None fetch_config: FetchConfig | None = None @model_validator(mode="after") @@ -187,6 +217,8 @@ class SearchRequest(CamelModel): prompt: Annotated[str, Field(min_length=1, max_length=10000)] | None = None schema_: dict[str, object] | None = Field(default=None, alias="schema") location_geo_code: Annotated[str, Field(max_length=10)] | None = None + allowed_types: AllowedTypes | None = None + processors: Processors | None = None time_range: TimeRange | None = None @model_validator(mode="after") @@ -237,7 +269,8 @@ class CrawlRequest(CamelModel): allow_external: bool = False include_patterns: list[str] | None = None exclude_patterns: list[str] | None = None - content_types: list[FetchContentType] | None = None + allowed_types: AllowedTypes | None = None + processors: Processors | None = None fetch_config: FetchConfig | None = None @model_validator(mode="after") diff --git a/tests/test_client.py b/tests/test_client.py index b889c13..deabdcf 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -10,6 +10,7 @@ JsonFormatConfig, LinksFormatConfig, MarkdownFormatConfig, + PdfProcessor, ScrapeGraphAI, ScreenshotFormatConfig, ) @@ -260,6 +261,21 @@ def test_with_extraction(self): _, kwargs = mock.call_args assert kwargs["json"]["prompt"] == "Summarize results" + def test_with_pdf_options(self): + body = {"results": [], "metadata": {"search": {}, "pages": {}}} + with patch.object(httpx.Client, "request", return_value=mock_response(body)) as mock: + sgai = ScrapeGraphAI(api_key=API_KEY) + res = sgai.search( + "papers", + allowed_types=["application/pdf"], + processors=[PdfProcessor(max_pages=10)], + ) + + assert res.status == "success" + _, kwargs = mock.call_args + assert kwargs["json"]["allowedTypes"] == ["application/pdf"] + assert kwargs["json"]["processors"] == [{"type": "pdf", "maxPages": 10}] + class TestCrawl: def test_start(self): diff --git a/tests/test_schemas.py b/tests/test_schemas.py new file mode 100644 index 0000000..0bfcf90 --- /dev/null +++ b/tests/test_schemas.py @@ -0,0 +1,73 @@ +import pytest +from pydantic import ValidationError + +from scrapegraph_py import ( + CrawlRequest, + ExtractRequest, + ScrapeRequest, + SearchRequest, +) + +REQUESTS = [ + (ScrapeRequest, {"url": "https://example.com"}), + (ExtractRequest, {"url": "https://example.com", "prompt": "title"}), + (SearchRequest, {"query": "example"}), + (CrawlRequest, {"url": "https://example.com"}), +] + + +@pytest.mark.parametrize(("request_model", "base"), REQUESTS) +def test_documented_pdf_configuration(request_model, base): + request = request_model( + **base, + allowed_types=["application/pdf"], + processors=[{"type": "pdf", "max_pages": 10}], + ) + + assert request.model_dump(by_alias=True, mode="json", exclude_none=True)["processors"] == [ + {"type": "pdf", "maxPages": 10} + ] + + +@pytest.mark.parametrize(("request_model", "base"), REQUESTS) +def test_default_pdf_cap_can_be_omitted(request_model, base): + request = request_model(**base) + assert "processors" not in request.model_dump(by_alias=True, mode="json", exclude_none=True) + + request = request_model(**base, processors=[{"type": "pdf"}]) + assert request.model_dump(by_alias=True, mode="json", exclude_none=True)["processors"] == [ + {"type": "pdf", "maxPages": 25} + ] + + +@pytest.mark.parametrize(("request_model", "base"), REQUESTS) +@pytest.mark.parametrize( + "invalid", + [ + {"allowed_types": []}, + {"allowed_types": ["application/pdf", "application/pdf"]}, + {"processors": []}, + { + "processors": [ + {"type": "pdf", "max_pages": 1}, + {"type": "pdf", "max_pages": 10}, + ] + }, + ], +) +def test_rejects_empty_and_duplicate_arrays(request_model, base, invalid): + with pytest.raises(ValidationError): + request_model(**base, **invalid) + + +@pytest.mark.parametrize(("request_model", "base"), REQUESTS) +@pytest.mark.parametrize("max_pages", [1, 500, -1]) +def test_accepts_documented_pdf_page_limits(request_model, base, max_pages): + request_model(**base, processors=[{"type": "pdf", "max_pages": max_pages}]) + + +@pytest.mark.parametrize(("request_model", "base"), REQUESTS) +@pytest.mark.parametrize("max_pages", [0, -2, 501, 1.5]) +def test_rejects_invalid_pdf_page_limits(request_model, base, max_pages): + with pytest.raises(ValidationError): + request_model(**base, processors=[{"type": "pdf", "max_pages": max_pages}])