Relocated

Scrapybara · Dec 21, 2024 · dbb7afa · dbb7afa
1 parent 5d584ff
commit dbb7afa
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 125 deletions.
diff --git 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
@@ -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):