Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions src/scrapybara/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import httpx
import os

import typing
from pydantic import BaseModel, ValidationError
from scrapybara.environment import ScrapybaraEnvironment
from .core.request_options import RequestOptions
from .types import (
Expand Down Expand Up @@ -31,6 +33,8 @@

OMIT = typing.cast(typing.Any, ...)

PydanticModelT = typing.TypeVar("PydanticModelT", bound=BaseModel)


class Agent:
def __init__(self, instance_id: str, client: BaseClient):
Expand Down Expand Up @@ -71,6 +75,33 @@ def scrape(
request_options=request_options,
)

def scrape_to_pydantic(
self,
*,
cmd: typing.Optional[str] = OMIT,
schema: PydanticModelT,
model: typing.Optional[typing.Literal["claude"]] = OMIT,
request_options: typing.Optional[RequestOptions] = None,
) -> PydanticModelT:
cmd = cmd if cmd else (str(schema.__doc__) if schema.__doc__ else None)
if cmd is None:
raise ValueError(
"No command provided, please provide a 'cmd' parameter or docstring in schema class."
)

response = self._client.agent.scrape(
self.instance_id,
cmd=cmd,
schema=schema.model_json_schema(),
model=model,
request_options=request_options,
)

try:
return schema.model_validate(response.data)
except ValidationError as e:
raise ValidationError(f"Validation error at client side: {e}") from e


class AsyncAgent:
def __init__(self, instance_id: str, client: AsyncBaseClient):
Expand Down Expand Up @@ -111,6 +142,33 @@ async def scrape(
request_options=request_options,
)

async def scrape_to_pydantic(
self,
*,
cmd: typing.Optional[str] = OMIT,
schema: PydanticModelT,
model: typing.Optional[typing.Literal["claude"]] = OMIT,
request_options: typing.Optional[RequestOptions] = None,
) -> PydanticModelT:
cmd = cmd if cmd else (str(schema.__doc__) if schema.__doc__ else None)
if cmd is None:
raise ValueError(
"No command provided, please provide a 'cmd' parameter or docstring in schema class."
)

response = await self._client.agent.scrape(
self.instance_id,
cmd=cmd,
schema=schema.model_json_schema(),
model=model,
request_options=request_options,
)

try:
return schema.model_validate(response.data)
except ValidationError as e:
raise ValidationError(f"Validation error at client side: {e}") from e


class Browser:
def __init__(self, instance_id: str, client: BaseClient):
Expand Down