blaxel 0.2.30__py3-none-any.whl → 0.2.31rc120__py3-none-any.whl

blaxel/core/sandbox/client/models/fuzzy_search_response.py

@@ -0,0 +1,88 @@
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.fuzzy_search_match import FuzzySearchMatch
+
+
+T = TypeVar("T", bound="FuzzySearchResponse")
+
+
+@_attrs_define
+class FuzzySearchResponse:
+    """
+    Attributes:
+        matches (Union[Unset, list['FuzzySearchMatch']]):
+        total (Union[Unset, int]):  Example: 5.
+    """
+
+    matches: Union[Unset, list["FuzzySearchMatch"]] = UNSET
+    total: Union[Unset, int] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        matches: Union[Unset, list[dict[str, Any]]] = UNSET
+        if not isinstance(self.matches, Unset):
+            matches = []
+            for matches_item_data in self.matches:
+                if type(matches_item_data) is dict:
+                    matches_item = matches_item_data
+                else:
+                    matches_item = matches_item_data.to_dict()
+                matches.append(matches_item)
+
+        total = self.total
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if matches is not UNSET:
+            field_dict["matches"] = matches
+        if total is not UNSET:
+            field_dict["total"] = total
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        from ..models.fuzzy_search_match import FuzzySearchMatch
+
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        matches = []
+        _matches = d.pop("matches", UNSET)
+        for matches_item_data in _matches or []:
+            matches_item = FuzzySearchMatch.from_dict(matches_item_data)
+
+            matches.append(matches_item)
+
+        total = d.pop("total", UNSET)
+
+        fuzzy_search_response = cls(
+            matches=matches,
+            total=total,
+        )
+
+        fuzzy_search_response.additional_properties = d
+        return fuzzy_search_response
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/core/sandbox/client/models/tree_request.py

@@ -0,0 +1,76 @@
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.tree_request_files import TreeRequestFiles
+
+
+T = TypeVar("T", bound="TreeRequest")
+
+
+@_attrs_define
+class TreeRequest:
+    """
+    Attributes:
+        files (Union[Unset, TreeRequestFiles]):  Example: {'"dir/file2.txt"': '"content2"}', '{"file1.txt"':
+            '"content1"'}.
+    """
+
+    files: Union[Unset, "TreeRequestFiles"] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        files: Union[Unset, dict[str, Any]] = UNSET
+        if self.files and not isinstance(self.files, Unset) and not isinstance(self.files, dict):
+            files = self.files.to_dict()
+        elif self.files and isinstance(self.files, dict):
+            files = self.files
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if files is not UNSET:
+            field_dict["files"] = files
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        from ..models.tree_request_files import TreeRequestFiles
+
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        _files = d.pop("files", UNSET)
+        files: Union[Unset, TreeRequestFiles]
+        if isinstance(_files, Unset):
+            files = UNSET
+        else:
+            files = TreeRequestFiles.from_dict(_files)
+
+        tree_request = cls(
+            files=files,
+        )
+
+        tree_request.additional_properties = d
+        return tree_request
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/core/sandbox/client/models/tree_request_files.py

@@ -0,0 +1,49 @@
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="TreeRequestFiles")
+
+
+@_attrs_define
+class TreeRequestFiles:
+    """
+    Example:
+        {'"dir/file2.txt"': '"content2"}', '{"file1.txt"': '"content1"'}
+
+    """
+
+    additional_properties: dict[str, str] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        tree_request_files = cls()
+
+        tree_request_files.additional_properties = d
+        return tree_request_files
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> str:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: str) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/core/sandbox/default/action.py

@@ -1,5 +1,6 @@
 
 import httpx
+from contextlib import asynccontextmanager
 
 from ...common.internal import get_forced_url, get_global_unique_hash
 from ...common.settings import settings
@@ -9,6 +10,7 @@ from ..types import ResponseError, SandboxConfiguration
 class SandboxAction:
     def __init__(self, sandbox_config: SandboxConfiguration):
         self.sandbox_config = sandbox_config
+        self._client: httpx.AsyncClient | None = None
 
     @property
     def name(self) -> str:
@@ -54,15 +56,17 @@ class SandboxAction:
         return None
 
     def get_client(self) -> httpx.AsyncClient:
-        if self.sandbox_config.force_url:
-            return httpx.AsyncClient(
-                base_url=self.sandbox_config.force_url, headers=self.sandbox_config.headers
+        """Get persistent HTTP client for this sandbox instance."""
+        if self._client is None:
+            base_url = self.sandbox_config.force_url or self.url
+            self._client = httpx.AsyncClient(
+                base_url=base_url,
+                headers=self.sandbox_config.headers if self.sandbox_config.force_url else {**settings.headers, **self.sandbox_config.headers},
+                http2=False,
+                limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
+                timeout=httpx.Timeout(300.0, connect=10.0),
             )
-        # Create a new client instance each time to avoid "Cannot open a client instance more than once" error
-        return httpx.AsyncClient(
-            base_url=self.url,
-            headers={**settings.headers, **self.sandbox_config.headers},
-        )
+        return self._client
 
     def handle_response_error(self, response: httpx.Response):
         if not response.is_success:

blaxel/core/sandbox/default/filesystem.py

@@ -34,10 +34,14 @@ class SandboxFileSystem(SandboxAction):
         path = self.format_path(path)
         body = FileRequest(is_directory=True, permissions=permissions)
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.put(f"/filesystem/{path}", json=body.to_dict())
+        client = self.get_client()
+        response = await client.put(f"/filesystem/{path}", json=body.to_dict())
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(response.json())
+            return SuccessResponse.from_dict(data)
+        finally:
+            await response.aclose()
 
     async def write(self, path: str, content: str) -> SuccessResponse:
         path = self.format_path(path)
@@ -53,10 +57,14 @@ class SandboxFileSystem(SandboxAction):
         # Use regular upload for small files
         body = FileRequest(content=content)
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.put(f"/filesystem/{path}", json=body.to_dict())
+        client = self.get_client()
+        response = await client.put(f"/filesystem/{path}", json=body.to_dict())
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(response.json())
+            return SuccessResponse.from_dict(data)
+        finally:
+            await response.aclose()
 
     async def write_binary(self, path: str, content: Union[bytes, bytearray, str]) -> SuccessResponse:
         """Write binary content to a file.
@@ -96,13 +104,16 @@ class SandboxFileSystem(SandboxAction):
         url = f"{self.url}/filesystem/{path}"
         headers = {**settings.headers, **self.sandbox_config.headers}
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.put(url, files=files, data=data, headers=headers)
-
+        client = self.get_client()
+        response = await client.put(url, files=files, data=data, headers=headers)
+        try:
+            content_bytes = await response.aread()
             if not response.is_success:
-                raise Exception(f"Failed to write binary: {response.status_code} {response.text}")
-
-            return SuccessResponse.from_dict(response.json())
+                error_text = content_bytes.decode('utf-8', errors='ignore')
+                raise Exception(f"Failed to write binary: {response.status_code} {error_text}")
+            return SuccessResponse.from_dict(json.loads(content_bytes))
+        finally:
+            await response.aclose()
 
     async def write_tree(
         self,
@@ -118,26 +129,32 @@ class SandboxFileSystem(SandboxAction):
 
         path = destination_path or ""
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.put(
-                f"/filesystem/tree/{path}",
-                json={"files": files_dict},
-                headers={"Content-Type": "application/json"},
-            )
+        client = self.get_client()
+        response = await client.put(
+            f"/filesystem/tree/{path}",
+            json={"files": files_dict},
+            headers={"Content-Type": "application/json"},
+        )
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return Directory.from_dict(response.json())
+            return Directory.from_dict(data)
+        finally:
+            await response.aclose()
 
     async def read(self, path: str) -> str:
         path = self.format_path(path)
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.get(f"/filesystem/{path}")
+        client = self.get_client()
+        response = await client.get(f"/filesystem/{path}")
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-
-            data = response.json()
             if "content" in data:
                 return data["content"]
             raise Exception("Unsupported file type")
+        finally:
+            await response.aclose()
 
     async def read_binary(self, path: str) -> bytes:
         """Read binary content from a file.
@@ -157,10 +174,14 @@ class SandboxFileSystem(SandboxAction):
             "Accept": "application/octet-stream",
         }
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.get(url, headers=headers)
+        client = self.get_client()
+        response = await client.get(url, headers=headers)
+        try:
+            content = await response.aread()
             self.handle_response_error(response)
-            return response.content
+            return content
+        finally:
+            await response.aclose()
 
     async def download(self, src: str, destination_path: str, mode: int = 0o644) -> None:
         """Download a file from the sandbox to the local filesystem.
@@ -178,23 +199,177 @@ class SandboxFileSystem(SandboxAction):
     async def rm(self, path: str, recursive: bool = False) -> SuccessResponse:
         path = self.format_path(path)
 
-        async with self.get_client() as client_instance:
-            params = {"recursive": "true"} if recursive else {}
-            response = await client_instance.delete(f"/filesystem/{path}", params=params)
+        client = self.get_client()
+        params = {"recursive": "true"} if recursive else {}
+        response = await client.delete(f"/filesystem/{path}", params=params)
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(response.json())
+            return SuccessResponse.from_dict(data)
+        finally:
+            await response.aclose()
 
     async def ls(self, path: str) -> Directory:
         path = self.format_path(path)
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.get(f"/filesystem/{path}")
+        client = self.get_client()
+        response = await client.get(f"/filesystem/{path}")
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-
-            data = response.json()
             if not ("files" in data or "subdirectories" in data):
                 raise Exception('{"error": "Directory not found"}')
             return Directory.from_dict(data)
+        finally:
+            await response.aclose()
+
+    async def find(
+        self,
+        path: str,
+        type: str | None = None,
+        patterns: List[str] | None = None,
+        max_results: int | None = None,
+        exclude_dirs: List[str] | None = None,
+        exclude_hidden: bool | None = None,
+    ):
+        """Find files and directories.
+
+        Args:
+            path: Path to search in
+            type: Type of search ('file' or 'directory')
+            patterns: File patterns to include (e.g., ['*.py', '*.json'])
+            max_results: Maximum number of results to return
+            exclude_dirs: Directory names to skip
+            exclude_hidden: Exclude hidden files and directories
+
+        Returns:
+            FindResponse with matching files/directories
+        """
+        path = self.format_path(path)
+        
+        params = {}
+        if type is not None:
+            params['type'] = type
+        if patterns is not None and len(patterns) > 0:
+            params['patterns'] = ','.join(patterns)
+        if max_results is not None:
+            params['maxResults'] = max_results
+        if exclude_dirs is not None and len(exclude_dirs) > 0:
+            params['excludeDirs'] = ','.join(exclude_dirs)
+        if exclude_hidden is not None:
+            params['excludeHidden'] = exclude_hidden
+
+        url = f"{self.url}/filesystem-find/{path}"
+        headers = {**settings.headers, **self.sandbox_config.headers}
+
+        client = self.get_client()
+        response = await client.get(url, params=params, headers=headers)
+        try:
+            data = json.loads(await response.aread())
+            self.handle_response_error(response)
+                
+            from ..client.models.find_response import FindResponse
+            return FindResponse.from_dict(data)
+        finally:
+            await response.aclose()
+
+    async def grep(
+        self,
+        query: str,
+        path: str = "/",
+        case_sensitive: bool | None = None,
+        context_lines: int | None = None,
+        max_results: int | None = None,
+        file_pattern: str | None = None,
+        exclude_dirs: List[str] | None = None,
+    ):
+        """Search for text content inside files using ripgrep.
+
+        Args:
+            query: Text to search for
+            path: Directory path to search in
+            case_sensitive: Case sensitive search (default: false)
+            context_lines: Number of context lines to include (default: 0)
+            max_results: Maximum number of results to return (default: 100)
+            file_pattern: File pattern to include (e.g., '*.py')
+            exclude_dirs: Directory names to skip
+
+        Returns:
+            ContentSearchResponse with matching lines
+        """
+        path = self.format_path(path)
+        
+        params = {'query': query}
+        if case_sensitive is not None:
+            params['caseSensitive'] = case_sensitive
+        if context_lines is not None:
+            params['contextLines'] = context_lines
+        if max_results is not None:
+            params['maxResults'] = max_results
+        if file_pattern is not None:
+            params['filePattern'] = file_pattern
+        if exclude_dirs is not None and len(exclude_dirs) > 0:
+            params['excludeDirs'] = ','.join(exclude_dirs)
+
+        url = f"{self.url}/filesystem-content-search/{path}"
+        headers = {**settings.headers, **self.sandbox_config.headers}
+
+        client = self.get_client()
+        response = await client.get(url, params=params, headers=headers)
+        try:
+            data = json.loads(await response.aread())
+            self.handle_response_error(response)
+                
+            from ..client.models.content_search_response import ContentSearchResponse
+            return ContentSearchResponse.from_dict(data)
+        finally:
+            await response.aclose()
+
+    async def find(
+        self,
+        path: str,
+        type: str | None = None,
+        patterns: List[str] | None = None,
+        max_results: int | None = None,
+        exclude_dirs: List[str] | None = None,
+        exclude_hidden: bool | None = None,
+    ):
+        """Find files and directories.
+
+        Args:
+            path: Path to search in
+            type: Type of search ('file' or 'directory')
+            patterns: File patterns to include (e.g., ['*.py', '*.json'])
+            max_results: Maximum number of results to return
+            exclude_dirs: Directory names to skip
+            exclude_hidden: Exclude hidden files and directories
+
+        Returns:
+            FindResponse with matching files/directories
+        """
+        path = self.format_path(path)
+        
+        params = {}
+        if type is not None:
+            params['type'] = type
+        if patterns is not None and len(patterns) > 0:
+            params['patterns'] = ','.join(patterns)
+        if max_results is not None:
+            params['maxResults'] = max_results
+        if exclude_dirs is not None and len(exclude_dirs) > 0:
+            params['excludeDirs'] = ','.join(exclude_dirs)
+        if exclude_hidden is not None:
+            params['excludeHidden'] = exclude_hidden
+
+        url = f"{self.url}/filesystem-find/{path}"
+        headers = {**settings.headers, **self.sandbox_config.headers}
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(url, params=params, headers=headers)
+            self.handle_response_error(response)
+            
+            from ..client.models.find_response import FindResponse
+            return FindResponse.from_dict(response.json())
 
     async def cp(self, source: str, destination: str, max_wait: int = 180000) -> CopyResponse:
         """Copy files or directories using the cp command.
@@ -328,10 +503,14 @@ class SandboxFileSystem(SandboxAction):
         headers = {**settings.headers, **self.sandbox_config.headers}
         body = {"permissions": permissions}
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.post(url, json=body, headers=headers)
+        client = self.get_client()
+        response = await client.post(url, json=body, headers=headers)
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return response.json()
+            return data
+        finally:
+            await response.aclose()
 
     async def _upload_part(
         self, upload_id: str, part_number: int, data: bytes
@@ -344,12 +523,16 @@ class SandboxFileSystem(SandboxAction):
         # Prepare multipart form data with the file chunk
         files = {"file": ("part", io.BytesIO(data), "application/octet-stream")}
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.put(
-                url, files=files, params=params, headers=headers
-            )
+        client = self.get_client()
+        response = await client.put(
+            url, files=files, params=params, headers=headers
+        )
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return response.json()
+            return data
+        finally:
+            await response.aclose()
 
     async def _complete_multipart_upload(
         self, upload_id: str, parts: List[Dict[str, Any]]
@@ -359,21 +542,28 @@ class SandboxFileSystem(SandboxAction):
         headers = {**settings.headers, **self.sandbox_config.headers}
         body = {"parts": parts}
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.post(url, json=body, headers=headers)
+        client = self.get_client()
+        response = await client.post(url, json=body, headers=headers)
+        try:
+            data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(response.json())
+            return SuccessResponse.from_dict(data)
+        finally:
+            await response.aclose()
 
     async def _abort_multipart_upload(self, upload_id: str) -> None:
         """Abort a multipart upload and clean up all parts."""
         url = f"{self.url}/filesystem-multipart/{upload_id}/abort"
         headers = {**settings.headers, **self.sandbox_config.headers}
 
-        async with self.get_client() as client_instance:
-            response = await client_instance.delete(url, headers=headers)
+        client = self.get_client()
+        response = await client.delete(url, headers=headers)
+        try:
             # Don't raise error if abort fails - we want to throw the original error
             if not response.is_success:
-                logger.warning(f"Warning: Failed to abort multipart upload: {response.status_code}")
+                print(f"Warning: Failed to abort multipart upload: {response.status_code}")
+        finally:
+            await response.aclose()
 
     async def _upload_with_multipart(
         self, path: str, data: bytes, permissions: str = "0644"