This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e46d19386190b201bdfa166907226406956fb3cc Author: Kaxil Naik <[email protected]> AuthorDate: Wed Sep 30 07:04:24 2026 +0100 Add `ObjectStorageToolset` for reading files on object storage (#73899) * Add ObjectStorageToolset for reading files on object storage The toolset gives an agent read-only access to the files under one ObjectStoragePath root: list a directory, get a file's size and modification time, and read a text file a window of lines at a time or a Parquet or Avro file as its schema and first rows. Paths are resolved under the root, and absolute paths, schemes, '..' and symlinks that leave a local root are refused. Like the other toolsets it works in AgentOperator, a Pydantic AI agent, Strands and Google ADK. * Test corrupt columnar files and paths on an object-store root A corrupt Parquet or Avro file is refused rather than failing the run: pyarrow's ArrowInvalid is a ValueError, which the refusal already catches. On an object store only the toolset's own path check applies, so encoded dots and backslashes are tested there to show they stay part of a key under the root. --- providers/common/ai/docs/agent_security.rst | 12 +- providers/common/ai/docs/stability.rst | 3 + providers/common/ai/docs/toolsets/index.rst | 13 +- .../common/ai/docs/toolsets/object_storage.rst | 123 +++++++ providers/common/ai/provider.yaml | 1 + .../example_dags/example_object_storage_toolset.py | 48 +++ .../airflow/providers/common/ai/sandbox/output.py | 6 +- .../providers/common/ai/toolsets/__init__.py | 2 + .../providers/common/ai/toolsets/object_storage.py | 367 +++++++++++++++++++++ .../providers/common/ai/utils/file_analysis.py | 53 ++- .../unit/common/ai/toolsets/test_object_storage.py | 359 ++++++++++++++++++++ .../unit/common/ai/utils/test_file_analysis.py | 6 +- 12 files changed, 965 insertions(+), 28 deletions(-) diff --git a/providers/common/ai/docs/agent_security.rst b/providers/common/ai/docs/agent_security.rst index ed1bd46b22a..5a169a869e5 100644 --- a/providers/common/ai/docs/agent_security.rst +++ b/providers/common/ai/docs/agent_security.rst @@ -82,8 +82,8 @@ No single layer is sufficient on its own. They work together. up in a database error or a hook's return value reaches the model, the model provider and any trace as ``***``. ``AgentOperator`` applies this to the toolsets you pass in ``toolsets``, in ``agent_params["toolsets"]`` and in a ``Toolset`` - capability, including your own. The SQL, hook, DataFusion, MCP, sandbox and - managed-agent toolsets apply it wherever they run, including in a Pydantic AI agent + capability, including your own. The SQL, hook, object storage, DataFusion, MCP, + sandbox and managed-agent toolsets apply it wherever they run, including in a Pydantic AI agent you build yourself. - Masks only secrets Airflow has registered, such as connection passwords and sensitive connection extras. A credential that exists only in the data itself is @@ -92,6 +92,14 @@ No single layer is sufficient on its own. They work together. framework's own tools, and MCP servers a framework connects to itself. In a sandbox, the model writes the commands, so it can print a secret in a form the masker does not recognize; masking there guards against accidents only. + * - **ObjectStorageToolset: one read-only root** + - The agent can list and read files under ``path`` and nothing else: it cannot write, + and a path that is absolute, carries a scheme or climbs out with ``..`` is refused. + On a local root, so is a symbolic link that leads out of it. + - The check is on the path the model supplies. What the connection may read is the + real limit, so scope its role or key to the prefix you pass as ``path``. On a + filesystem that follows symlinks on the server side, such as SFTP, the path check + cannot see where a link leads. * - **HookToolset: explicit allow-list** - Only methods listed in ``allowed_methods`` are exposed as tools. Auto-discovery is not supported. Methods are validated at Dag parse diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index c7e1f254bae..ab798c360d7 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -172,3 +172,6 @@ Everything this provider ships that is not in the table above is experimental. - Written against Strands 1.56 and ADK 2.9.1. CI does not run the tests of the two adapters, because both frameworks exclude dependency versions that Airflow's development environment uses. + * - :class:`~airflow.providers.common.ai.toolsets.object_storage.ObjectStorageToolset` + (:doc:`toolsets/object_storage`) + - New; its tools and read limits may change after first use. diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst index ed6f44a7fb4..c2e91ec63fd 100644 --- a/providers/common/ai/docs/toolsets/index.rst +++ b/providers/common/ai/docs/toolsets/index.rst @@ -64,8 +64,11 @@ Start with what you have - ``HookToolset`` * - A question that is a query, against a DBAPI database - ``SQLToolset`` + * - Files on an object store that the agent should read the way a person would: + browse a directory, open a report, look at the first rows of a Parquet file + - ``ObjectStorageToolset`` * - Files on an object store (Parquet, CSV, Avro) or a catalog-managed - table format such as Iceberg, rather than rows in a database + table format such as Iceberg, to query with SQL rather than read - ``DataFusionToolset`` * - A vendor that already ships a server built for agents, whose tools you would otherwise re-wrap by hand @@ -78,7 +81,7 @@ Start with what you have * - Reasoning that should happen on the vendor's own infrastructure - A subclass of ``BaseManagedAgentToolset`` that you write -The hook, SQL, DataFusion, MCP, Agent Skills and managed-agent guides each have a +The hook, SQL, object storage, DataFusion, MCP, Agent Skills and managed-agent guides each have a *When to choose it* section giving the case for choosing it, what it cannot do, an example that exists in this repository, and where its credentials and its work come from. :doc:`../sandbox/index` carries the same section for ``SandboxToolset``. @@ -91,6 +94,7 @@ Toolset guides Airflow hooks as tools <hook> SQL databases <sql> + Files on object storage <object_storage> Files with DataFusion <datafusion> MCP servers <mcp> Agent Skills <skills> @@ -106,13 +110,16 @@ Airflow's 350+ provider hooks already have typed methods, rich docstrings, and managed credentials. Toolsets expose them as pydantic-ai tools so that LLM agents can call them during multi-turn reasoning. -Six toolsets are exported directly from the ``airflow.providers.common.ai.toolsets`` +Seven toolsets are exported directly from the ``airflow.providers.common.ai.toolsets`` package root: - :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset`: generic adapter for any Airflow Hook. Guide: :doc:`hook`. - :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`: curated 4-tool database toolset. Guide: :doc:`sql`. +- :class:`~airflow.providers.common.ai.toolsets.object_storage.ObjectStorageToolset`: + read-only access to the files under one object-storage path. Guide: + :doc:`object_storage`. - :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset`: connect to `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow connections. Guide: :doc:`mcp`. diff --git a/providers/common/ai/docs/toolsets/object_storage.rst b/providers/common/ai/docs/toolsets/object_storage.rst new file mode 100644 index 00000000000..7dae611ecc2 --- /dev/null +++ b/providers/common/ai/docs/toolsets/object_storage.rst @@ -0,0 +1,123 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/toolset:object_storage: + +Files on object storage: ``ObjectStorageToolset`` +================================================= + +.. note:: + + Experimental: this can change or be removed in a minor release of this provider. + See :ref:`howto/stability`. + +Give an agent the files under one location in S3, GCS, Azure Blob Storage or any other +store Airflow's :class:`~airflow.sdk.ObjectStoragePath` can open, and let it find and +read what it needs: a month's reports, the config files of a failing job, the first +rows of a Parquet extract. The agent can only read, and only under the path you give +it. + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py + :language: python + :start-after: [START howto_toolset_object_storage] + :end-before: [END howto_toolset_object_storage] + +The credentials come from ``conn_id``, the same connection your other tasks use for +that store. With ``conn_id=None`` the store's default credentials apply, such as the +worker's AWS role. + +The three tools +--------------- + +``list_files`` + Lists one directory, sorted by name: files with their size in bytes, and + subdirectories with a trailing slash. The model goes deeper by listing a + subdirectory. A directory with more than ``max_files`` entries is listed a page at a + time, with a note giving the ``offset`` to continue from. Every call lists the whole + directory from storage, so a very large one is slow to page through. + +``get_file_info`` + The size and last-modified time of one file, so the model can decide whether it is + worth reading. + +``read_file`` + A text file a window of lines at a time. The model passes ``offset`` and ``limit`` + as line numbers, and a result that stops early says which ``offset`` to continue + from, the same shape as the sandbox's ``read_file``. A Parquet or Avro file comes + back as its schema and its first 20 rows. Text compressed as ``.gz``, ``.bz2`` or + ``.xz`` is decompressed first. + +Every path the model supplies is relative to the root. An absolute path, a path with a +scheme such as ``s3://``, and a path that climbs out with ``..`` are refused, and on a +local root a symbolic link that points outside the root is refused too, and left out of +listings. + +What the model cannot read is refused with a message it can act on, rather than failing +the task: an image or PDF, a binary file, a file larger than ``max_read_bytes`` (10 MiB +by default), a corrupt file, a path that does not exist, or one the connection may not +read. + +Parameters +---------- + +``path`` + The root, for example ``"s3://acme-reports/finance/"``. Templated when the toolset + is passed to ``AgentOperator`` or ``@task.agent``. +``conn_id`` + The connection for the store. Templated like ``path``. +``max_files`` + The most entries one ``list_files`` result holds. Default ``200``. +``max_read_bytes`` + The largest file ``read_file`` opens, after decompression. Default 10 MiB. +``max_output_bytes`` + The most bytes one ``read_file`` result holds. Default 50 KiB. +``tool_prefix`` + A prefix for the tool names, needed when the agent has another toolset with the same + tool names: a second ``ObjectStorageToolset``, or a ``SandboxToolset``, which has a + ``read_file`` of its own. + +When to choose it +----------------- + +**Choose it when** the agent should read files the way a person would: browse a +directory, open the report that matters, look at a file's first lines or rows. Text +files need nothing beyond the filesystem package for your store, which Airflow's object +storage already uses; Parquet and Avro need this provider's ``parquet`` or ``avro`` +extra. To hand the model a known file rather than let it find one, ``@task.llm_file_analysis`` +reads the file for it. For read-only access through a hook's own methods, +``HookToolset(S3Hook(), allowed_methods=["list_keys", "read_key"])`` works too, but the +model then chooses the bucket and key, where this toolset holds it under one root. + +**What it cannot do** + +- It cannot write, delete, move or copy anything. +- It does not search inside files. The model finds a file by listing directories, so a + store with thousands of files per directory is slow to explore. Give it a narrower root. +- It reads one file at a time. To ask a question across many files, such as a sum over + a month of Parquet files, use :doc:`datafusion`, which runs SQL over them. +- The root is the boundary only as far as the path check goes. The connection's own + permissions are the real limit on what can be read, so scope its role or key to the + prefix you pass as ``path``. + +Using it with other agent frameworks +------------------------------------ + +``ObjectStorageToolset`` implements +:class:`~airflow.providers.common.ai.tools.ToolProvider`, so the same three tools work in +a Strands or Google ADK agent through ``AirflowTools``; see :doc:`../frameworks/index`. +Outside ``AgentOperator``, ``path`` and ``conn_id`` are used as given: they are not +rendered as templates. diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 77872c5e8b1..023d1bb1e4c 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -551,6 +551,7 @@ toolsets: - airflow.providers.common.ai.toolsets.datafusion - airflow.providers.common.ai.toolsets.logging - airflow.providers.common.ai.toolsets.mcp + - airflow.providers.common.ai.toolsets.object_storage - airflow.providers.common.ai.toolsets.sandbox - airflow.providers.common.ai.toolsets.skills - airflow.providers.common.ai.toolsets.langchain_bridge diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py new file mode 100644 index 00000000000..3cd86eb9fb5 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Example Dag: an agent that reads the files under one object-storage path.""" + +from __future__ import annotations + +from airflow.providers.common.ai.toolsets import ObjectStorageToolset +from airflow.providers.common.compat.sdk import dag, task + + +# [START howto_toolset_object_storage] +@dag(tags=["example"]) +def example_object_storage_toolset(): + @task.agent( + llm_conn_id="pydanticai_default", + system_prompt=( + "You answer questions about the files in a reports bucket. List directories " + "to find what exists, then read the files you need. Cite the files you used." + ), + toolsets=[ + # Read-only, and every path the model names is resolved under this root. + ObjectStorageToolset("s3://acme-reports/finance/", conn_id="aws_default"), + ], + ) + def summarize_month(month: str) -> str: + return f"Summarize revenue and the biggest changes for {month}." + + summarize_month("2026-09") + + +# [END howto_toolset_object_storage] + + +example_object_storage_toolset() diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/output.py b/providers/common/ai/src/airflow/providers/common/ai/sandbox/output.py index 867aaa91679..678c9b6f589 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/output.py +++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/output.py @@ -119,6 +119,7 @@ def render_file_window( limit: int | None, max_lines: int, max_bytes: int, + long_line_hint: str = "Read part of it with a shell command instead (e.g. cut, head -c, or sed).", ) -> str: """ Render a window of a text file, head-first, with a continuation offset. @@ -140,10 +141,7 @@ def render_file_window( kept, truncated = _keep_lines(window, max_lines=max_lines, max_bytes=max_bytes, from_tail=False) if not kept and truncated: - return ( - f"Line {start} is longer than the {format_size(max_bytes)} output limit. " - "Read part of it with a shell command instead (e.g. cut, head -c, or sed)." - ) + return f"Line {start} is longer than the {format_size(max_bytes)} output limit. {long_line_hint}" body = "\n".join(kept) next_offset = start + len(kept) diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py index a3adfd7f1ce..df797bfe623 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py @@ -23,6 +23,7 @@ from airflow.providers.common.ai.toolsets.managed_agent import ( BaseManagedAgentToolset, FailoverManagedAgentToolset, ) +from airflow.providers.common.ai.toolsets.object_storage import ObjectStorageToolset from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset __all__ = [ @@ -30,6 +31,7 @@ __all__ = [ "FailoverManagedAgentToolset", "HookToolset", "MCPToolset", + "ObjectStorageToolset", "SQLToolset", "SandboxToolset", "airflow_toolset_to_langchain_tools", diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py new file mode 100644 index 00000000000..1eeee652a28 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py @@ -0,0 +1,367 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Read-only toolset giving an agent the files under one object-storage path.""" + +from __future__ import annotations + +import os +from datetime import datetime, timezone +from pathlib import PurePosixPath +from typing import TYPE_CHECKING, Any, Literal + +from fsspec.implementations.local import LocalFileSystem +from pydantic_ai.exceptions import ToolFailed +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets.abstract import ToolsetTool + +from airflow.providers.common.ai.exceptions import LLMFileAnalysisError, LLMFileAnalysisLimitExceededError +from airflow.providers.common.ai.sandbox.output import format_size, render_file_window +from airflow.providers.common.ai.utils.file_analysis import ( + detect_compression, + detect_file_format, + read_bytes, + sample_columnar_file, +) +from airflow.providers.common.ai.utils.masking import mask_secrets +from airflow.providers.common.ai.utils.tool_definition import ( + build_args_validator, + return_schema_kwargs, + serialize_for_llm, +) +from airflow.providers.common.ai.utils.toolset_base import AirflowToolset +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, ObjectStoragePath + +if TYPE_CHECKING: + from collections.abc import Sequence + + from pydantic_ai._run_context import RunContext + +LIST_FILES = "list_files" +GET_FILE_INFO = "get_file_info" +READ_FILE = "read_file" + +_PATH_DESCRIPTION = "Path relative to the storage root, using / between parts. Omit for the root itself." + +_SCHEMAS: dict[str, dict[str, Any]] = { + LIST_FILES: { + "type": "object", + "properties": { + "path": {"type": "string", "description": _PATH_DESCRIPTION}, + "offset": { + "type": ["integer", "null"], + "description": "Entry to start listing from (0-indexed), to page through a large directory.", + }, + }, + "required": [], + }, + GET_FILE_INFO: { + "type": "object", + "properties": {"path": {"type": "string", "description": _PATH_DESCRIPTION}}, + "required": ["path"], + }, + READ_FILE: { + "type": "object", + "properties": { + "path": {"type": "string", "description": _PATH_DESCRIPTION}, + "offset": { + "type": ["integer", "null"], + "description": "Line number to start reading from (1-indexed).", + }, + "limit": {"type": ["integer", "null"], "description": "Maximum number of lines to read."}, + }, + "required": ["path"], + }, +} + +_DESCRIPTIONS = { + LIST_FILES: ( + "List the files and directories in one directory of the storage, sorted by name. " + "Directories end with a slash; list one of them to go deeper. A large directory is listed " + "a page at a time, and the result tells you the offset to continue from." + ), + GET_FILE_INFO: "Get the size and last-modified time of one file or directory.", + READ_FILE: ( + "Read a text file. Long files are returned a window at a time and the result tells you the " + "offset to continue from. For a Parquet or Avro file, returns its schema and first rows." + ), +} + +_MAX_OUTPUT_LINES = 2000 +_SAMPLE_ROWS = 20 +_COLUMNAR_FORMATS: tuple[Literal["parquet", "avro"], ...] = ("parquet", "avro") +_MEDIA_FORMATS = frozenset({"jpeg", "jpg", "pdf", "png"}) + + +class ObjectStorageToolset(AirflowToolset): + """ + Give an agent read-only access to the files under one object-storage path. + + .. note:: + + Experimental: this can change or be removed in a minor release of this provider. + See :ref:`howto/stability`. + + Exposes three tools, ``list_files``, ``get_file_info`` and ``read_file``, rooted at + ``path``, which is any location Airflow's + :class:`~airflow.sdk.ObjectStoragePath` can open: ``s3://``, ``gs://``, ``abfs://``, + ``file://`` and the rest, with credentials from ``conn_id``. The model names files by + paths relative to that root. It cannot write, delete or move anything, and a path that + is absolute, carries a scheme, or climbs out of the root with ``..`` is refused. + + ``read_file`` returns a text file a window of lines at a time, like the sandbox's own + ``read_file``, and a Parquet or Avro file as its schema and first rows. Compressed text + (``.gz``, ``.bz2``, ``.xz``) is decompressed. Images, PDFs and other binary files are + refused, as is any file larger than ``max_read_bytes``. So is a file that cannot be read, + such as a corrupt one, or one the connection may not open: the model is told why, and the + run goes on. On a local root, a symlink that leads out of the root is refused too. + + :param path: Root the agent may read under. Templated when the toolset is passed to + ``AgentOperator`` / ``@task.agent``. + :param conn_id: Airflow connection for the storage, or ``None`` for the default + credentials of its protocol. Templated like ``path``. + :param max_files: Most entries one ``list_files`` result holds; the model pages through + a larger directory. The whole directory is still listed from storage on each call. + Default ``200``. + :param max_read_bytes: Largest file ``read_file`` will open, after decompression. + Default 10 MiB. + :param max_output_bytes: Most bytes one ``read_file`` result holds; the model reads on + from the offset it is given. Default 50 KiB. + :param tool_prefix: Prefix for the three tool names, e.g. ``"reports"`` gives + ``reports_read_file``. Set this when one agent has another toolset with the same tool + names, such as a second ``ObjectStorageToolset`` or a ``SandboxToolset``, whose + ``read_file`` would collide, since duplicate tool names are rejected. + """ + + # Rendered, on a copy, by AgentOperator. Deliberately not ``template_fields``, which + # Airflow's templater would render in place wherever the toolset is nested. + agent_template_fields: Sequence[str] = ("_path", "_conn_id") + + def __init__( + self, + path: str, + *, + conn_id: str | None = None, + max_files: int = 200, + max_read_bytes: int = 10 * 1024 * 1024, + max_output_bytes: int = 50 * 1024, + tool_prefix: str = "", + ) -> None: + for name, value in ( + ("max_files", max_files), + ("max_read_bytes", max_read_bytes), + ("max_output_bytes", max_output_bytes), + ): + if value < 1: + raise ValueError(f"{name} must be at least 1, got {value}.") + if tool_prefix and not tool_prefix.isidentifier(): + raise ValueError(f"tool_prefix must be a valid Python identifier, got {tool_prefix!r}.") + self._path = path + self._conn_id = conn_id + self._max_files = max_files + self._max_read_bytes = max_read_bytes + self._max_output_bytes = max_output_bytes + self._tool_prefix = tool_prefix + + @property + def id(self) -> str: + suffix = f"-{self._tool_prefix}" if self._tool_prefix else "" + return f"object-storage-{self._conn_id or 'default'}{suffix}" + + def _tool_name(self, base: str) -> str: + return f"{self._tool_prefix}_{base}" if self._tool_prefix else base + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + tools: dict[str, ToolsetTool[Any]] = {} + for base, schema in _SCHEMAS.items(): + name = self._tool_name(base) + tools[name] = ToolsetTool( + toolset=self, + tool_def=ToolDefinition( + name=name, + description=_DESCRIPTIONS[base], + parameters_json_schema=schema, + **return_schema_kwargs({"type": "string"}), + ), + max_retries=1, + args_validator=build_args_validator(schema), + ) + return tools + + async def _execute_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[Any], + tool: ToolsetTool[Any], + ) -> str: + base = name.removeprefix(f"{self._tool_prefix}_") if self._tool_prefix else name + relative = tool_args.get("path") or "" + try: + if base == LIST_FILES: + return await self.run_blocking( + self._list_files, relative, offset=tool_args.get("offset") or 0 + ) + if base == GET_FILE_INFO: + return await self.run_blocking(self._get_file_info, relative) + if base == READ_FILE: + return await self.run_blocking( + self._read_file, relative, offset=tool_args.get("offset"), limit=tool_args.get("limit") + ) + except ToolFailed: + raise + except (OSError, EOFError, ValueError, AirflowOptionalProviderFeatureException) as e: + # Storage the connection may not read, a corrupt or mislabelled file, a codec this + # Python build lacks: final for this path, so the model is told, not the task failed. + raise ToolFailed(f"{relative or '/'!r} cannot be read: {type(e).__name__}: {e}") from None + raise ValueError(f"Unknown tool: {name!r}") + + # ------------------------------------------------------------------ + # Tool implementations. Each runs in a worker thread. + # ------------------------------------------------------------------ + + def _root(self) -> ObjectStoragePath: + # Built on use, not in __init__, so AgentOperator's rendering of _path and _conn_id applies. + return ObjectStoragePath(self._path, conn_id=self._conn_id) + + def _resolve(self, relative: str) -> ObjectStoragePath: + """Turn the model's relative path into a path under the root, or refuse it.""" + if relative.startswith("/") or "://" in relative: + raise ToolFailed(f"{relative!r} is not a relative path. Name files relative to the storage root.") + parts = [part for part in PurePosixPath(relative).parts if part != "."] + if ".." in parts: + raise ToolFailed(f"{relative!r} leaves the storage root; '..' is not allowed.") + root = self._root() + target = root.joinpath(*parts) if parts else root + if not _inside_local_root(root, target.path): + raise ToolFailed(f"{relative!r} resolves outside the storage root.") + return target + + def _list_files(self, relative: str, *, offset: int) -> str: + directory = self._resolve(relative) + if not directory.is_dir(): + raise ToolFailed(f"{relative or '/'!r} is not a directory.") + root = self._root() + # One listing call with details, rather than a stat per entry, which on S3 or GCS is a + # request each. Some stores list a directory's own placeholder object; skip it. + own_path = directory.path.rstrip("/") + entries = sorted( + ( + info + for info in directory.fs.ls(directory.path, detail=True) + if info["name"].rstrip("/") != own_path and _inside_local_root(root, info["name"]) + ), + key=_entry_name, + ) + page = entries[offset : offset + self._max_files] + listed = [ + {"name": f"{_entry_name(info)}/"} + if info["type"] == "directory" + else {"name": _entry_name(info), "size_bytes": info.get("size")} + for info in page + ] + result: dict[str, Any] = {"path": relative or "/", "entries": listed} + if offset + len(page) < len(entries): + result["note"] = ( + f"Showing entries {offset + 1} to {offset + len(page)} of {len(entries)}; list again " + f"with offset={offset + len(page)} for more." + ) + return serialize_for_llm(result) + + def _get_file_info(self, relative: str) -> str: + target = self._resolve(relative) + # One request: on S3 or GCS, exists(), is_dir() and stat() would each be a round trip. + try: + stat = target.stat() + except FileNotFoundError: + raise ToolFailed(f"{relative!r} does not exist.") from None + if stat.get("type") == "directory": + return serialize_for_llm({"path": relative, "type": "directory"}) + info: dict[str, Any] = {"path": relative, "type": "file", "size_bytes": stat.st_size} + if modified := _as_iso(stat.st_mtime): + info["modified"] = modified + return serialize_for_llm(info) + + def _read_file(self, relative: str, *, offset: int | None, limit: int | None) -> str: + target = self._resolve(relative) + if not target.is_file(): + raise ToolFailed(f"{relative!r} is not a file.") + try: + file_format, compression = detect_file_format(target) + except LLMFileAnalysisError: + # An extension file analysis does not know, such as .sql or .yaml.gz: read it as text. + file_format, compression = "txt", detect_compression(target) + if file_format in _MEDIA_FORMATS: + raise ToolFailed(f"{relative!r} is a {file_format} file, which this tool cannot read as text.") + try: + for columnar in _COLUMNAR_FORMATS: + if file_format == columnar: + sample = sample_columnar_file( + target, file_format=columnar, sample_rows=_SAMPLE_ROWS, max_bytes=self._max_read_bytes + ) + return _cut(sample, self._max_output_bytes) + data = read_bytes(target, compression=compression, max_bytes=self._max_read_bytes) + except LLMFileAnalysisLimitExceededError: + raise ToolFailed( + f"{relative!r} is larger than the {format_size(self._max_read_bytes)} this tool reads." + ) from None + if b"\x00" in data: + raise ToolFailed(f"{relative!r} is a binary file, which this tool cannot read as text.") + # Masked whole, before a window is cut: a secret spanning lines would otherwise leak a + # line at a time. + return render_file_window( + mask_secrets(data), + offset=offset, + limit=limit, + max_lines=_MAX_OUTPUT_LINES, + max_bytes=self._max_output_bytes, + long_line_hint="It cannot be read a line at a time.", + ) + + +def _entry_name(info: dict[str, Any]) -> str: + return PurePosixPath(info["name"].rstrip("/")).name + + +def _inside_local_root(root: ObjectStoragePath, path: str) -> bool: + """ + Whether ``path`` stays inside ``root`` once symlinks are resolved, for a root on local disk. + + Checked by filesystem, not scheme: a root written as a plain path has no scheme but is + still on the worker's disk, where a symlink under it could point anywhere. Object stores + have no symlinks, so any other root is not checked here. + """ + if not isinstance(root.fs, LocalFileSystem): + return True + real_root = os.path.realpath(root.path) + return os.path.commonpath([real_root, os.path.realpath(path)]) == real_root + + +def _as_iso(modified: Any) -> str | None: + """Return a modification time as ISO 8601; stores report a timestamp, a datetime or nothing.""" + if isinstance(modified, datetime): + return modified.isoformat() + if isinstance(modified, (int, float)) and modified > 0: + return datetime.fromtimestamp(modified, tz=timezone.utc).isoformat() + return None + + +def _cut(text: str, max_bytes: int) -> str: + encoded = text.encode("utf-8") + if len(encoded) <= max_bytes: + return text + kept = encoded[:max_bytes].decode("utf-8", "ignore") + return f"{kept}\n[... cut at {format_size(max_bytes)}]" diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py index 9f0fbc4e241..930a3fae0f6 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py @@ -26,7 +26,7 @@ import logging from bisect import insort from dataclasses import dataclass from pathlib import PurePosixPath -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal # bz2/lzma are optional CPython extensions and may be missing from some interpreter builds try: @@ -347,7 +347,7 @@ def _prepare_file( f"File {path} has format {file_format!r}; set multi_modal=True to analyze images or PDFs." ) prepared.attachment = BinaryContent( - data=_read_raw_bytes(path, compression=compression, max_bytes=max_content_bytes), + data=read_bytes(path, compression=compression, max_bytes=max_content_bytes), media_type=_MEDIA_TYPES[file_format], identifier=str(path), ) @@ -378,6 +378,13 @@ def _prepare_file( return prepared +def detect_compression(path: ObjectStoragePath) -> str | None: + """Return the codec a path's last suffix names, if this Python build can decompress it.""" + suffixes = path.suffixes + codec = _COMPRESSION_SUFFIXES.get(suffixes[-1].removeprefix(".").lower()) if suffixes else None + return codec if codec in _DECOMPRESSORS else None + + def detect_file_format(path: ObjectStoragePath) -> tuple[str, str | None]: """Detect the logical file format and compression codec from a path suffix.""" suffixes = [suffix.removeprefix(".").lower() for suffix in path.suffixes] @@ -431,7 +438,7 @@ def _render_text_content( def _render_text_like( path: ObjectStoragePath, *, compression: str | None, max_content_bytes: int ) -> _RenderResult: - raw_bytes = _read_raw_bytes(path, compression=compression, max_bytes=max_content_bytes) + raw_bytes = read_bytes(path, compression=compression, max_bytes=max_content_bytes) text = _decode_text(raw_bytes) return _RenderResult(text=_truncate_text(text), estimated_rows=None, content_size_bytes=len(raw_bytes)) @@ -439,7 +446,7 @@ def _render_text_like( def _render_json( path: ObjectStoragePath, *, compression: str | None, max_content_bytes: int ) -> _RenderResult: - raw_bytes = _read_raw_bytes(path, compression=compression, max_bytes=max_content_bytes) + raw_bytes = read_bytes(path, compression=compression, max_bytes=max_content_bytes) decoded = _decode_text(raw_bytes) document = json.loads(decoded) if isinstance(document, list): @@ -457,7 +464,7 @@ def _render_json( def _render_csv( path: ObjectStoragePath, *, compression: str | None, sample_rows: int, max_content_bytes: int ) -> _RenderResult: - raw_bytes = _read_raw_bytes(path, compression=compression, max_bytes=max_content_bytes) + raw_bytes = read_bytes(path, compression=compression, max_bytes=max_content_bytes) decoded = _decode_text(raw_bytes) reader = list(csv.reader(io.StringIO(decoded))) if not reader: @@ -475,6 +482,18 @@ def _render_csv( ) +def sample_columnar_file( + path: ObjectStoragePath, *, file_format: Literal["parquet", "avro"], sample_rows: int, max_bytes: int +) -> str: + """ + Describe a Parquet or Avro file for a model: its schema and its first ``sample_rows`` rows. + + :raises LLMFileAnalysisLimitExceededError: if the file is larger than ``max_bytes``. + """ + render = _render_parquet if file_format == "parquet" else _render_avro + return render(path, sample_rows=sample_rows, max_content_bytes=max_bytes).text + + def _render_parquet(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes: int) -> _RenderResult: try: import pyarrow.parquet as pq @@ -498,17 +517,14 @@ def _render_parquet(path: ObjectStoragePath, *, sample_rows: int, max_content_by schema = ", ".join(f"{field.name}: {field.type}" for field in parquet_file.schema_arrow) sampled_rows: list[dict[str, Any]] = [] - if sample_rows > 0 and num_rows > 0 and parquet_file.num_row_groups > 0: - remaining_rows = sample_rows - for row_group_index in range(parquet_file.num_row_groups): - if remaining_rows <= 0: + if sample_rows > 0 and num_rows > 0: + # Decode only the first rows: a whole row group can decompress to many times the + # file's size, which the size limit above does not bound. + for batch in parquet_file.iter_batches(batch_size=sample_rows): + sampled_rows.extend(batch.to_pylist()) + if len(sampled_rows) >= sample_rows: break - row_group = parquet_file.read_row_group(row_group_index) - if row_group.num_rows == 0: - continue - group_rows = row_group.slice(0, remaining_rows).to_pylist() - sampled_rows.extend(group_rows) - remaining_rows -= len(group_rows) + sampled_rows = sampled_rows[:sample_rows] payload = [f"Schema: {schema}", "Sample rows:", dumps_masked(sampled_rows, indent=2)] return _RenderResult( text=_truncate_text("\n".join(payload)), @@ -559,7 +575,12 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes ) -def _read_raw_bytes(path: ObjectStoragePath, *, compression: str | None, max_bytes: int) -> bytes: +def read_bytes(path: ObjectStoragePath, *, compression: str | None, max_bytes: int) -> bytes: + """ + Read ``path``, decompressing it with ``compression``, and refuse more than ``max_bytes``. + + :raises LLMFileAnalysisLimitExceededError: if the content is larger than ``max_bytes``. + """ with path.open("rb") as handle: if compression is None: return _read_limited_bytes(handle, path=path, max_bytes=max_bytes) diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py new file mode 100644 index 00000000000..a02bb133b7a --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py @@ -0,0 +1,359 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import gzip +import json +import uuid +from typing import Any +from unittest.mock import patch + +import pytest +from fsspec.implementations.memory import MemoryFileSystem +from pydantic_ai import RunContext +from pydantic_ai.exceptions import ToolFailed +from pydantic_ai.models.test import TestModel +from pydantic_ai.usage import RunUsage + +from airflow.providers.common.ai.toolsets.object_storage import ObjectStorageToolset +from airflow.sdk.io.store import _STORE_CACHE, ObjectStore + + [email protected] +def storage(tmp_path): + """A small tree of files to read, with a sibling directory the agent must not reach.""" + root = tmp_path / "reports" + (root / "2026" / "09").mkdir(parents=True) + (root / "2026" / "09" / "summary.md").write_text("# September\n\nRevenue up 4%.\n") + (root / "orders.csv").write_text("id,total\n1,10\n2,20\n") + (root / "config.yaml").write_text("retries: 3\n") + (root / "notes.txt.gz").write_bytes(gzip.compress(b"compressed notes\n")) + (root / "logo.png").write_bytes(b"\x89PNG\r\n\x1a\n") + (root / "blob.bin").write_bytes(b"\x00\x01\x02") + secret = tmp_path / "secrets" + secret.mkdir() + (secret / "key.txt").write_text("do not read\n") + (root / "link").symlink_to(secret / "key.txt") + return root + + [email protected] +def object_store(monkeypatch): + """ + A root on an in-memory object store, with a file above it the agent must not reach. + + No symlinks and no local disk, so only the toolset's own path check keeps a path inside. + """ + conn_id = f"memory-{uuid.uuid4().hex}" + fs = MemoryFileSystem() + monkeypatch.setitem(_STORE_CACHE, f"memory-{conn_id}", ObjectStore("memory", conn_id, fs=fs)) + bucket = f"/{conn_id}" + fs.pipe(f"{bucket}/reports/a/ok.txt", b"inside\n") + fs.pipe(f"{bucket}/secret.txt", b"do not read\n") + fs.pipe(f"{bucket}/reports/..%2Fsecret.txt", b"a key inside the root\n") + yield ObjectStorageToolset(f"memory:/{bucket}/reports", conn_id=conn_id) + fs.rm(bucket, recursive=True) + + +def _call(toolset: ObjectStorageToolset, name: str, arguments: dict[str, Any]) -> Any: + async def call() -> Any: + ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage()) + tools = await toolset.get_tools(ctx) + validated = tools[name].args_validator.validate_python(arguments) + return await toolset.call_tool(name, validated, ctx, tools[name]) + + return asyncio.run(call()) + + +class TestTools: + def test_exposes_three_read_only_tools(self, storage): + ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage()) + + tools = asyncio.run(ObjectStorageToolset(f"file://{storage}").get_tools(ctx)) + + assert list(tools) == ["list_files", "get_file_info", "read_file"] + + def test_the_id_names_the_connection(self, storage): + assert ( + ObjectStorageToolset(f"file://{storage}", conn_id="reports_s3").id == "object-storage-reports_s3" + ) + + def test_a_prefix_renames_the_tools(self, storage): + ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage()) + + tools = asyncio.run(ObjectStorageToolset(f"file://{storage}", tool_prefix="reports").get_tools(ctx)) + + assert list(tools) == ["reports_list_files", "reports_get_file_info", "reports_read_file"] + listing = json.loads( + _call(ObjectStorageToolset(f"file://{storage}", tool_prefix="reports"), "reports_list_files", {}) + ) + assert listing["path"] == "/" + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"max_files": 0}, "max_files"), + ({"max_read_bytes": 0}, "max_read_bytes"), + ({"tool_prefix": "my-files"}, "tool_prefix"), + ], + ) + def test_rejects_an_invalid_setting(self, kwargs, match): + with pytest.raises(ValueError, match=match): + ObjectStorageToolset("file:///tmp", **kwargs) + + +class TestListFiles: + def test_lists_one_directory_with_sizes(self, storage): + listing = json.loads(_call(ObjectStorageToolset(f"file://{storage}"), "list_files", {})) + + names = [entry["name"] for entry in listing["entries"]] + assert names == sorted(names) + assert "2026/" in names + assert {"name": "orders.csv", "size_bytes": 19} in listing["entries"] + + def test_lists_a_subdirectory(self, storage): + listing = json.loads( + _call(ObjectStorageToolset(f"file://{storage}"), "list_files", {"path": "2026/09"}) + ) + + assert listing["entries"] == [{"name": "summary.md", "size_bytes": 28}] + + def test_pages_through_a_directory_larger_than_max_files(self, storage): + toolset = ObjectStorageToolset(f"file://{storage}", max_files=2) + + first = json.loads(_call(toolset, "list_files", {})) + second = json.loads(_call(toolset, "list_files", {"offset": 2})) + + assert len(first["entries"]) == 2 + assert first["note"].endswith("list again with offset=2 for more.") + assert {e["name"] for e in first["entries"]}.isdisjoint(e["name"] for e in second["entries"]) + + def test_leaves_out_a_symlink_that_leads_out_of_the_root(self, storage): + listing = json.loads(_call(ObjectStorageToolset(f"file://{storage}"), "list_files", {})) + + assert "link" not in {entry["name"] for entry in listing["entries"]} + + def test_lists_a_child_with_the_same_name_as_its_directory(self, storage): + (storage / "data").mkdir() + (storage / "data" / "data").write_text("x") + + listing = json.loads(_call(ObjectStorageToolset(f"file://{storage}"), "list_files", {"path": "data"})) + + assert listing["entries"] == [{"name": "data", "size_bytes": 1}] + + def test_refuses_a_file_as_a_directory(self, storage): + with pytest.raises(ToolFailed, match="is not a directory"): + _call(ObjectStorageToolset(f"file://{storage}"), "list_files", {"path": "orders.csv"}) + + +class TestPathsStayUnderTheRoot: + @pytest.mark.parametrize( + ("path", "match"), + [ + ("../secrets/key.txt", "leaves the storage root"), + ("2026/../../secrets/key.txt", "leaves the storage root"), + ("/etc/passwd", "is not a relative path"), + ("file:///etc/passwd", "is not a relative path"), + ("link", "resolves outside the storage root"), + ], + ) + def test_refuses_a_path_outside_the_root(self, storage, path, match): + with pytest.raises(ToolFailed, match=match): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": path}) + + @pytest.mark.parametrize("tool", ["read_file", "get_file_info"]) + def test_a_root_written_without_a_scheme_is_checked_for_symlinks_too(self, storage, tool): + """A plain path is still on the worker's disk, where the symlink leads anywhere.""" + with pytest.raises(ToolFailed, match="resolves outside the storage root"): + _call(ObjectStorageToolset(str(storage)), tool, {"path": "link"}) + + def test_a_sibling_whose_name_starts_with_the_root_is_outside_it(self, storage): + (storage.parent / "reports-private").mkdir() + (storage.parent / "reports-private" / "k.txt").write_text("x") + (storage / "near").symlink_to(storage.parent / "reports-private" / "k.txt") + + with pytest.raises(ToolFailed, match="resolves outside the storage root"): + _call(ObjectStorageToolset(str(storage)), "read_file", {"path": "near"}) + + +class TestPathsStayUnderAnObjectStoreRoot: + """An object store has no directories to climb out of: a key is only ever a string.""" + + def test_dot_dot_is_refused(self, object_store): + with pytest.raises(ToolFailed, match="leaves the storage root"): + _call(object_store, "read_file", {"path": "../secret.txt"}) + + @pytest.mark.parametrize( + "path", + ["%2e%2e/secret.txt", "a/%2e%2e/%2e%2e/secret.txt", "..\\secret.txt", "a\\..\\..\\secret.txt"], + ) + def test_encoded_dots_and_backslashes_do_not_climb_out(self, object_store, path): + """Neither is decoded or treated as a separator, so the key sits under the root and is missing.""" + with pytest.raises(ToolFailed, match="is not a file"): + _call(object_store, "read_file", {"path": path}) + + def test_a_key_that_looks_encoded_is_read_under_the_root(self, object_store): + assert _call(object_store, "read_file", {"path": "..%2Fsecret.txt"}) == "a key inside the root" + + +class TestReadFile: + def test_reads_a_text_file(self, storage): + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "2026/09/summary.md"}) + + assert result == "# September\n\nRevenue up 4%." + + def test_reads_a_file_with_an_extension_file_analysis_does_not_know(self, storage): + assert ( + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "config.yaml"}) + == "retries: 3" + ) + + def test_decompresses_a_compressed_text_file(self, storage): + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "notes.txt.gz"}) + + assert result == "compressed notes" + + def test_decompresses_text_with_an_extension_file_analysis_does_not_know(self, storage): + (storage / "schema.sql.gz").write_bytes(gzip.compress(b"CREATE TABLE t (id INT);\n")) + + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "schema.sql.gz"}) + + assert result == "CREATE TABLE t (id INT);" + + def test_a_corrupt_file_is_refused_instead_of_failing_the_run(self, storage): + (storage / "broken.txt.gz").write_bytes(b"not gzip at all") + + with pytest.raises(ToolFailed, match="cannot be read"): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "broken.txt.gz"}) + + @pytest.mark.parametrize( + ("path", "module", "error"), + [ + pytest.param("broken.parquet", "pyarrow.parquet", "ArrowInvalid", id="parquet"), + pytest.param("broken.avro", "fastavro", "ValueError", id="avro"), + ], + ) + def test_a_corrupt_columnar_file_is_refused_instead_of_failing_the_run( + self, storage, path, module, error + ): + """pyarrow's ArrowInvalid is a ValueError, so the refusal covers Parquet as well as gzip.""" + pytest.importorskip(module) + (storage / path).write_bytes(b"neither parquet nor avro") + + with pytest.raises(ToolFailed, match=f"cannot be read: {error}"): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": path}) + + def test_reads_the_schema_and_first_rows_of_a_parquet_file(self, storage): + pq = pytest.importorskip("pyarrow.parquet") + pa = pytest.importorskip("pyarrow") + pq.write_table( + pa.table({"id": list(range(100)), "total": [i * 10 for i in range(100)]}), + storage / "orders.parquet", + ) + + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "orders.parquet"}) + + assert result.startswith("Schema: id: int64, total: int64") + assert '"id": 19' in result + assert '"id": 20' not in result + + def test_reads_a_window_and_says_where_to_continue(self, storage): + result = _call( + ObjectStorageToolset(f"file://{storage}"), + "read_file", + {"path": "orders.csv", "offset": 2, "limit": 1}, + ) + + assert result == "1,10\n[... 1 more line; read on with offset=3]" + + @pytest.mark.parametrize( + ("path", "match"), + [ + ("logo.png", "png file"), + ("blob.bin", "binary file"), + ("missing.txt", "not a file"), + ("2026", "not a file"), + ], + ) + def test_refuses_what_it_cannot_read_as_text(self, storage, path, match): + with pytest.raises(ToolFailed, match=match): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": path}) + + def test_refuses_a_file_over_the_read_limit(self, storage): + with pytest.raises(ToolFailed, match="larger than the 8B"): + _call( + ObjectStorageToolset(f"file://{storage}", max_read_bytes=8), + "read_file", + {"path": "orders.csv"}, + ) + + @pytest.mark.enable_redact + def test_a_secret_spanning_lines_is_masked_before_a_window_is_cut(self, storage, register_secret): + secret = register_secret("pem-part-one-91c3\npem-part-two-91c3") + (storage / "key.pem").write_text(f"header\n{secret}\nfooter\n") + + first_line_of_secret = _call( + ObjectStorageToolset(f"file://{storage}"), + "read_file", + {"path": "key.pem", "offset": 2, "limit": 1}, + ) + + assert "pem-part-one-91c3" not in first_line_of_secret + + +class TestGetFileInfo: + def test_describes_a_file(self, storage): + info = json.loads( + _call(ObjectStorageToolset(f"file://{storage}"), "get_file_info", {"path": "orders.csv"}) + ) + + assert info["type"] == "file" + assert info["size_bytes"] == 19 + + def test_a_missing_path_is_refused(self, storage): + with pytest.raises(ToolFailed, match="does not exist"): + _call(ObjectStorageToolset(f"file://{storage}"), "get_file_info", {"path": "missing.txt"}) + + def test_the_modification_time_is_iso_8601(self, storage): + info = json.loads( + _call(ObjectStorageToolset(f"file://{storage}"), "get_file_info", {"path": "orders.csv"}) + ) + + assert info["modified"].endswith("+00:00") + + def test_describes_a_directory(self, storage): + info = json.loads(_call(ObjectStorageToolset(f"file://{storage}"), "get_file_info", {"path": "2026"})) + + assert info == {"path": "2026", "type": "directory"} + + +class TestOtherFrameworks: + def test_a_refusal_reaches_a_native_agent_as_an_error_result(self, storage): + tools = {tool.name: tool for tool in ObjectStorageToolset(f"file://{storage}").airflow_tools()} + + result = asyncio.run(tools["read_file"].call({"path": "../secrets/key.txt"})) + + assert result.is_error + assert "leaves the storage root" in result.content + + +def test_a_file_the_connection_may_not_read_is_refused_not_fatal(storage): + with patch("airflow.providers.common.ai.toolsets.object_storage.read_bytes", autospec=True) as read: + read.side_effect = PermissionError("AccessDenied") + with pytest.raises(ToolFailed, match="cannot be read: PermissionError"): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "orders.csv"}) diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py index f2fbbaf0162..c7bfb7d9389 100644 --- a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py @@ -34,13 +34,13 @@ from airflow.providers.common.ai.utils.file_analysis import ( _DECOMPRESSORS, FileAnalysisRequest, _infer_partitions, - _read_raw_bytes, _render_avro, _render_parquet, _resolve_paths, _truncate_text, build_file_analysis_request, detect_file_format, + read_bytes, ) from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, ObjectStoragePath @@ -477,7 +477,7 @@ class TestFileAnalysisHelpers: path = tmp_path / f"events.log.{suffix}" path.write_bytes(codec.compress(b"line one\nline two\n")) - content = _read_raw_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) + content = read_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) assert content == b"line one\nline two\n" @@ -498,7 +498,7 @@ class TestFileAnalysisHelpers: path = tmp_path / f"events.log.{suffix}" path.write_bytes(codec.compress(b"first\n") + separator + codec.compress(b"second\n")) - content = _read_raw_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) + content = read_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) assert content == expected
