From 0f92a7e79229e749849faebde8f5fd1483cc1593 Mon Sep 17 00:00:00 2001 From: Bohdan Agarkov Date: Sat, 21 Dec 2024 23:05:08 +0100 Subject: [PATCH 1/4] Added pydantic query wrapper code --- src/scrapybara/agent/client.py | 125 +++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/src/scrapybara/agent/client.py b/src/scrapybara/agent/client.py index 73b3e70..ac5cb01 100644 --- a/src/scrapybara/agent/client.py +++ b/src/scrapybara/agent/client.py @@ -1,6 +1,7 @@ # This file was auto-generated by Fern from our API Definition. import typing +from pydantic import BaseModel, ValidationError from ..core.client_wrapper import SyncClientWrapper from ..core.request_options import RequestOptions from ..types.act_response import ActResponse @@ -16,6 +17,8 @@ # this is used as the default value for optional parameters OMIT = typing.cast(typing.Any, ...) +PydanticModelT = typing.TypeVar("PydanticModelT", bound=BaseModel) + class AgentClient: def __init__(self, *, client_wrapper: SyncClientWrapper): @@ -181,6 +184,67 @@ def scrape( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) + def query( + self, + instance_id: str, + *, + cmd: typing.Optional[str] = OMIT, + schema: PydanticModelT, + model: typing.Optional[typing.Literal["claude"]] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> PydanticModelT: + """ + Parameters + ---------- + instance_id : str + + cmd : typing.Optional[str] + Command to execute. If not provided, will use schema class docstring. + + schema : PydanticModelT + Pydantic model class defining the expected response structure. + + model : typing.Optional[typing.Literal["claude"]] + Model to use for scraping. + + request_options : typing.Optional[RequestOptions] + Request-specific configuration. + + Returns + ------- + PydanticModelT + Response data validated against the provided schema. + + Raises + ------ + ValueError + If no command is provided at least via cmd parameter or schema docstring. + ValidationError + If response data fails validation against schema. + """ + cmd = ( + cmd + if cmd + else (str(schema.__doc__) if schema.__doc__ else "No command provided") + ) + if cmd == "No command provided": + raise ValueError( + "No command provided, please provide a 'cmd' parameter or docstring in schema class." + ) + + response = self.scrape( + 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 AsyncAgentClient: def __init__(self, *, client_wrapper: AsyncClientWrapper): @@ -361,3 +425,64 @@ async def main() -> None: except JSONDecodeError: raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) + + async def query( + self, + instance_id: str, + *, + cmd: typing.Optional[str] = OMIT, + schema: PydanticModelT, + model: typing.Optional[typing.Literal["claude"]] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> PydanticModelT: + """ + Parameters + ---------- + instance_id : str + + cmd : typing.Optional[str] + Command to execute. If not provided, will use schema class docstring. + + schema : PydanticModelT + Pydantic model class defining the expected response structure. + + model : typing.Optional[typing.Literal["claude"]] + Model to use for scraping. + + request_options : typing.Optional[RequestOptions] + Request-specific configuration. + + Returns + ------- + PydanticModelT + Response data validated against the provided schema. + + Raises + ------ + ValueError + If no command is provided at least via cmd parameter or schema docstring. + ValidationError + If response data fails validation against schema. + """ + cmd = ( + cmd + if cmd + else (str(schema.__doc__) if schema.__doc__ else "No command provided") + ) + if cmd == "No command provided": + raise ValueError( + "No command provided, please provide a 'cmd' parameter or docstring in schema class." + ) + + response = await self.scrape( + 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 From 5d584ffa6c196d604a451b6fcd704fd7da2b8b4d Mon Sep 17 00:00:00 2001 From: Bohdan Agarkov Date: Sat, 21 Dec 2024 23:11:51 +0100 Subject: [PATCH 2/4] Little code quality fix --- src/scrapybara/agent/client.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/scrapybara/agent/client.py b/src/scrapybara/agent/client.py index ac5cb01..4424091 100644 --- a/src/scrapybara/agent/client.py +++ b/src/scrapybara/agent/client.py @@ -225,9 +225,9 @@ def query( cmd = ( cmd if cmd - else (str(schema.__doc__) if schema.__doc__ else "No command provided") + else (str(schema.__doc__) if schema.__doc__ else None) ) - if cmd == "No command provided": + if cmd is None: raise ValueError( "No command provided, please provide a 'cmd' parameter or docstring in schema class." ) @@ -467,9 +467,9 @@ async def query( cmd = ( cmd if cmd - else (str(schema.__doc__) if schema.__doc__ else "No command provided") + else (str(schema.__doc__) if schema.__doc__ else None) ) - if cmd == "No command provided": + if cmd is None: raise ValueError( "No command provided, please provide a 'cmd' parameter or docstring in schema class." ) From dbb7afa65e5b2592559790162d16ace3bf447983 Mon Sep 17 00:00:00 2001 From: Bohdan Agarkov Date: Sat, 21 Dec 2024 23:40:03 +0100 Subject: [PATCH 3/4] Relocated --- src/scrapybara/agent/client.py | 125 --------------------------------- src/scrapybara/client.py | 58 +++++++++++++++ 2 files changed, 58 insertions(+), 125 deletions(-) diff --git a/src/scrapybara/agent/client.py b/src/scrapybara/agent/client.py index 4424091..73b3e70 100644 --- a/src/scrapybara/agent/client.py +++ b/src/scrapybara/agent/client.py @@ -1,7 +1,6 @@ # This file was auto-generated by Fern from our API Definition. import typing -from pydantic import BaseModel, ValidationError from ..core.client_wrapper import SyncClientWrapper from ..core.request_options import RequestOptions from ..types.act_response import ActResponse @@ -17,8 +16,6 @@ # this is used as the default value for optional parameters OMIT = typing.cast(typing.Any, ...) -PydanticModelT = typing.TypeVar("PydanticModelT", bound=BaseModel) - class AgentClient: def __init__(self, *, client_wrapper: SyncClientWrapper): @@ -184,67 +181,6 @@ def scrape( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - def query( - self, - instance_id: str, - *, - cmd: typing.Optional[str] = OMIT, - schema: PydanticModelT, - model: typing.Optional[typing.Literal["claude"]] = OMIT, - request_options: typing.Optional[RequestOptions] = None, - ) -> PydanticModelT: - """ - Parameters - ---------- - instance_id : str - - cmd : typing.Optional[str] - Command to execute. If not provided, will use schema class docstring. - - schema : PydanticModelT - Pydantic model class defining the expected response structure. - - model : typing.Optional[typing.Literal["claude"]] - Model to use for scraping. - - request_options : typing.Optional[RequestOptions] - Request-specific configuration. - - Returns - ------- - PydanticModelT - Response data validated against the provided schema. - - Raises - ------ - ValueError - If no command is provided at least via cmd parameter or schema docstring. - ValidationError - If response data fails validation against schema. - """ - 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.scrape( - 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 AsyncAgentClient: def __init__(self, *, client_wrapper: AsyncClientWrapper): @@ -425,64 +361,3 @@ async def main() -> None: except JSONDecodeError: raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - - async def query( - self, - instance_id: str, - *, - cmd: typing.Optional[str] = OMIT, - schema: PydanticModelT, - model: typing.Optional[typing.Literal["claude"]] = OMIT, - request_options: typing.Optional[RequestOptions] = None, - ) -> PydanticModelT: - """ - Parameters - ---------- - instance_id : str - - cmd : typing.Optional[str] - Command to execute. If not provided, will use schema class docstring. - - schema : PydanticModelT - Pydantic model class defining the expected response structure. - - model : typing.Optional[typing.Literal["claude"]] - Model to use for scraping. - - request_options : typing.Optional[RequestOptions] - Request-specific configuration. - - Returns - ------- - PydanticModelT - Response data validated against the provided schema. - - Raises - ------ - ValueError - If no command is provided at least via cmd parameter or schema docstring. - ValidationError - If response data fails validation against schema. - """ - 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.scrape( - 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 diff --git a/src/scrapybara/client.py b/src/scrapybara/client.py index e02a3e1..39e3027 100644 --- a/src/scrapybara/client.py +++ b/src/scrapybara/client.py @@ -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 ( @@ -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): @@ -71,6 +75,33 @@ def scrape( request_options=request_options, ) + def query( + 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): @@ -111,6 +142,33 @@ async def scrape( request_options=request_options, ) + async def query( + 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): From 0297f17ef4309911777ee5c402a460c6c67235ef Mon Sep 17 00:00:00 2001 From: Bohdan Agarkov Date: Sat, 21 Dec 2024 23:46:53 +0100 Subject: [PATCH 4/4] Changed name --- src/scrapybara/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/scrapybara/client.py b/src/scrapybara/client.py index 39e3027..334c8b6 100644 --- a/src/scrapybara/client.py +++ b/src/scrapybara/client.py @@ -75,7 +75,7 @@ def scrape( request_options=request_options, ) - def query( + def scrape_to_pydantic( self, *, cmd: typing.Optional[str] = OMIT, @@ -142,7 +142,7 @@ async def scrape( request_options=request_options, ) - async def query( + async def scrape_to_pydantic( self, *, cmd: typing.Optional[str] = OMIT,