Data Access

Data Manager

agora_workbench.code_execution.data_access.manager

DataLakeDataManager for fetching and caching data assets.

This module provides the core manager that handles fetching assets from DataLake-cataloged sources and caching them to disk for tool access.

DataLakeDataManager(allowed_local_roots=None, extra_fetchers=None, credential=None, artifact_resolver=None)

Manages data asset retrieval and caching for code execution sessions.

Fetches data assets from DataLake-cataloged sources and caches them to disk in their original format. Tools receive Path objects and handle loading the data as needed.

Authentication uses a credential chain (create_storage_credential): the mounted az login MSAL cache for local development, falling back to managed identity in production, for downstream Azure resources (Storage, AI Search).

Blob artifact IDs are resolved through a pluggable ArtifactResolver, defaulting to Azure AI Search; see artifact_resolvers.py.

Supports: - Azure Blob Storage (abfss://, az://, https://) - Local filesystem (absolute paths, relative paths, file:// URIs)

Initialize the data manager.

Uses a credential chain (az login MSAL cache locally, managed identity in production) for Azure Blob access unless a credential is provided. Blob URL fetching is available whenever a credential can be initialized. Blob artifact ID resolution is delegated to an ArtifactResolver; the default one queries Azure AI Search and so additionally requires DATA_LAKE_SEARCH_ENDPOINT.

Parameters:

Name	Type	Description	Default
allowed_local_roots	list[str] | None	Optional list of directory paths the local file fetcher is allowed to read from. If None or empty, all paths are permitted (suitable for sandboxed containers).	None
extra_fetchers	list[AssetFetcher] | None	Optional list of additional AssetFetcher instances to register. These are checked before the built-in fetchers (local file, blob), allowing custom fetchers to override default handling for specific URL schemes or patterns.	None
credential	AsyncTokenCredential | None	Optional async token credential to use for Azure Blob Storage and Azure AI Search access. When omitted, the manager creates the same storage credential chain as before, resolving AZURE_CLIENT_ID for user-assigned managed identity binding.	None
artifact_resolver	ArtifactResolver | None	Optional resolver turning <blob>id</blob> identifiers into fetchable URLs, for deployments whose catalog is not an Azure AI Search index. When omitted, a SearchIndexArtifactResolver is built from the environment, preserving existing behavior. A supplied resolver is used as-is and is not given the manager's credential, so it must arrange its own authentication.	None

Raises:

Type	Description
TypeError	If artifact_resolver does not implement the ArtifactResolver protocol.

Source code in src/agora_workbench/code_execution/data_access/manager.py

def __init__(
    self,
    allowed_local_roots: list[str] | None = None,
    extra_fetchers: list[AssetFetcher] | None = None,
    credential: AsyncTokenCredential | None = None,
    artifact_resolver: ArtifactResolver | None = None,
):
    """
    Initialize the data manager.

    Uses a credential chain (``az login`` MSAL cache locally, managed
    identity in production) for Azure Blob access unless a credential is
    provided. Blob URL fetching is available whenever a credential can be
    initialized. Blob artifact ID resolution is delegated to an
    ``ArtifactResolver``; the default one queries Azure AI Search and so
    additionally requires ``DATA_LAKE_SEARCH_ENDPOINT``.

    Args:
        allowed_local_roots: Optional list of directory paths the local
            file fetcher is allowed to read from. If ``None`` or empty,
            all paths are permitted (suitable for sandboxed containers).
        extra_fetchers: Optional list of additional ``AssetFetcher``
            instances to register. These are checked *before* the
            built-in fetchers (local file, blob), allowing custom
            fetchers to override default handling for specific URL
            schemes or patterns.
        credential: Optional async token credential to use for Azure Blob
            Storage and Azure AI Search access. When omitted, the manager
            creates the same storage credential chain as before, resolving
            ``AZURE_CLIENT_ID`` for user-assigned managed identity binding.
        artifact_resolver: Optional resolver turning ``<blob>id</blob>``
            identifiers into fetchable URLs, for deployments whose catalog
            is not an Azure AI Search index. When omitted, a
            ``SearchIndexArtifactResolver`` is built from the environment,
            preserving existing behavior. A supplied resolver is used as-is
            and is *not* given the manager's credential, so it must arrange
            its own authentication.

    Raises:
        TypeError: If ``artifact_resolver`` does not implement the
            ``ArtifactResolver`` protocol.
    """
    self._cache_dir = Path(tempfile.mkdtemp(prefix="data_lake_cache_"))
    self._cache_index = {}  # Maps artifact_id -> cache file path

    self._credential_init_error: str | None = None
    self._credential: AsyncTokenCredential | None = None
    self._owns_credential = credential is None

    # Initialize fetchers — custom fetchers take priority over built-ins
    self._fetchers: list[AssetFetcher] = list(extra_fetchers or [])
    self._fetchers.append(LocalFileFetcher(allowed_roots=allowed_local_roots))

    try:
        if credential is not None:
            self._credential = credential
        else:
            mi_client_id = (os.getenv("AZURE_CLIENT_ID") or "").strip() or None
            # MSAL cache (mounted az login) locally, managed identity in
            # production. Pass the AZURE_CLIENT_ID-resolved id through so
            # prod keeps binding to the same user-assigned identity.
            self._credential = create_storage_credential(client_id=mi_client_id)
    except (ImportError, RuntimeError, TypeError, ValueError) as e:
        self._credential_init_error = f"{type(e).__name__}: {e}"
        LOGGER.warning(f"Failed to initialize Azure storage credential: {e}")

    if self._credential is not None:
        self._fetchers.append(BlobFetcher(credential=self._credential))

    if artifact_resolver is not None:
        _validate_artifact_resolver(artifact_resolver)
        self._artifact_resolver: ArtifactResolver = artifact_resolver
    else:
        # The deferred credential error is handed over so the resolver can
        # explain *why* it is unavailable rather than blaming configuration.
        self._artifact_resolver = SearchIndexArtifactResolver.from_env(
            credential=self._credential,
            credential_init_error=self._credential_init_error,
        )

get_cache_path(qualified_name) async

Get the filesystem path where the asset is cached.

Ensures the asset is fetched and cached to disk in its original format, then returns the path. This allows kernel subprocesses to load the asset directly.

Parameters:

Name	Type	Description	Default
qualified_name	AssetId	Type-tagged artifact format base64_id Examples: id, id	required

Returns:

Type	Description
Path	Path to the cached file

Raises:

Type	Description
ValueError	If not in tagged format, unsupported type, or artifact not found

Source code in src/agora_workbench/code_execution/data_access/manager.py

async def get_cache_path(self, qualified_name: "AssetId") -> Path:
    """
    Get the filesystem path where the asset is cached.

    Ensures the asset is fetched and cached to disk in its original format,
    then returns the path. This allows kernel subprocesses to load the
    asset directly.

    Args:
        qualified_name: Type-tagged artifact format <type>base64_id</type>
                      Examples: <blob>id</blob>, <sql>id</sql>

    Returns:
        Path to the cached file

    Raises:
        ValueError: If not in tagged format, unsupported type, or artifact not found
    """
    # Extract artifact type and ID from tagged format
    artifact_match = re.match(r"^<(\w+)>([^<>]+)</\1>$", qualified_name.strip())
    if not artifact_match:
        # Fallback: accept unclosed tags like "<blob>id" (LLM sometimes omits closing tag)
        artifact_match = re.match(r"^<(\w+)>([^<>]+)$", qualified_name.strip())
    if not artifact_match:
        raise ValueError(
            f"Invalid artifact format - expected <type>id</type>, got: {qualified_name}. "
            f"{self._asset_tag_guidance()} {agent_guidance.DISCOVER_DATA}"
        )

    artifact_type = artifact_match.group(1)
    artifact_id = artifact_match.group(2)

    LOGGER.info(f"Resolving {artifact_type} artifact: {artifact_id}")

    # Check if already cached (use artifact_id as cache key)
    if artifact_id in self._cache_index:
        cache_path = self._cache_index[artifact_id]
        if cache_path.exists():
            LOGGER.debug(f"Asset already cached: {cache_path}")
            return cache_path

    # Route to appropriate resolver based on artifact type
    if artifact_type == "blob":
        resource_url = await self._get_blob_url_from_artifact_id(artifact_id)
    elif artifact_type == "local":
        # Local artifacts: the artifact_id is the file path itself
        resource_url = artifact_id
    else:
        raise ValueError(f"Unsupported artifact type: {artifact_type}. {self._asset_tag_guidance()}")

    # Fetch and cache the asset
    LOGGER.debug(f"Fetching and caching {artifact_type} asset")
    cache_path = self._get_cache_file_path(resource_url)

    # Stream asset directly to file to avoid loading into memory
    bytes_written = await self._fetch_asset_to_file(resource_url, cache_path)

    # Update index (use artifact_id as key)
    self._cache_index[artifact_id] = cache_path

    LOGGER.debug(f"Cached asset to disk ({bytes_written} bytes)")
    return cache_path

get_asset_info(qualified_name)

Get information about a cached asset.

Parameters:

Name	Type	Description	Default
qualified_name	str	Asset identifier	required

Returns:

Type	Description
dict	Dict with asset metadata

Source code in src/agora_workbench/code_execution/data_access/manager.py

def get_asset_info(self, qualified_name: str) -> dict:
    """
    Get information about a cached asset.

    Args:
        qualified_name: Asset identifier

    Returns:
        Dict with asset metadata
    """
    info = {
        "qualified_name": qualified_name,
        "cached": qualified_name in self._cache_index,
    }

    if qualified_name in self._cache_index:
        cache_path = self._cache_index[qualified_name]
        if cache_path.exists():
            info["size_bytes"] = cache_path.stat().st_size
            info["cache_location"] = str(cache_path)

    return info

list_available()

List assets available in this session's cache.

Returns:

Type	Description
list[str]	List of qualified names currently cached

Source code in src/agora_workbench/code_execution/data_access/manager.py

def list_available(self) -> list[str]:
    """
    List assets available in this session's cache.

    Returns:
        List of qualified names currently cached
    """
    return list(self._cache_index.keys())

cleanup()

Clean up cache directory, credentials, and resources.

Removes the temporary cache directory if it was created by this manager. Call this when the session is ending to free up disk space.

Source code in src/agora_workbench/code_execution/data_access/manager.py

def cleanup(self) -> None:
    """
    Clean up cache directory, credentials, and resources.

    Removes the temporary cache directory if it was created by this manager.
    Call this when the session is ending to free up disk space.
    """
    # Clear cache index
    self._cache_index.clear()

    # Close the artifact resolver (releases any catalog client it owns)
    resolver_close = _resolver_aclose(getattr(self, "_artifact_resolver", None))
    if resolver_close is not None:
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(resolver_close())
        except RuntimeError:
            try:
                loop = asyncio.new_event_loop()
                loop.run_until_complete(resolver_close())
                loop.close()
            except Exception as e:
                LOGGER.debug(f"Error closing artifact resolver: {e}")

    # Close managed identity credential
    if hasattr(self, "_credential") and self._credential is not None and self._owns_credential:
        try:
            loop = asyncio.get_running_loop()
            # We're inside a running loop — schedule close as a task
            loop.create_task(self._credential.close())
        except RuntimeError:
            # No running loop — safe to create a temporary one
            try:
                loop = asyncio.new_event_loop()
                loop.run_until_complete(self._credential.close())
                loop.close()
            except Exception as e:
                LOGGER.debug(f"Error closing credential: {e}")

    # Remove temp directory
    if self._cache_dir and self._cache_dir.exists():
        try:
            shutil.rmtree(self._cache_dir)
            LOGGER.info(f"Cleaned up cache directory: {self._cache_dir}")
        except Exception as e:
            LOGGER.warning(f"Failed to clean up cache directory: {e}")

aclose() async

Async cleanup — preferred over sync cleanup() when inside an event loop.

Source code in src/agora_workbench/code_execution/data_access/manager.py

async def aclose(self) -> None:
    """Async cleanup — preferred over sync cleanup() when inside an event loop."""
    self._cache_index.clear()

    # Close fetchers (releases pooled connections)
    for fetcher in self._fetchers:
        if hasattr(fetcher, "close"):
            try:
                await fetcher.close()
            except Exception as e:
                LOGGER.debug(f"Error closing fetcher {fetcher.__class__.__name__}: {e}")

    resolver_close = _resolver_aclose(getattr(self, "_artifact_resolver", None))
    if resolver_close is not None:
        try:
            await resolver_close()
        except Exception as e:
            LOGGER.debug(f"Error closing artifact resolver: {e}")

    if hasattr(self, "_credential") and self._credential is not None and self._owns_credential:
        try:
            await self._credential.close()
        except Exception as e:
            LOGGER.debug(f"Error closing credential: {e}")

    if self._cache_dir and self._cache_dir.exists():
        try:
            shutil.rmtree(self._cache_dir)
            LOGGER.info(f"Cleaned up cache directory: {self._cache_dir}")
        except Exception as e:
            LOGGER.warning(f"Failed to clean up cache directory: {e}")

__del__()

Cleanup on garbage collection.

Source code in src/agora_workbench/code_execution/data_access/manager.py

def __del__(self):
    """Cleanup on garbage collection."""
    try:
        self.cleanup()
    except Exception as e:
        LOGGER.exception(f"Error during DataLakeDataManager cleanup: {e}")

Asset Resolution

agora_workbench.code_execution.data_access.resolution

Utilities for detecting and resolving DataLake-cataloged asset references.

AssetResolutionMiddleware(server)

Bases: Middleware

FastMCP middleware that resolves DataLake asset references before Pydantic validation.

When the agent sends a tool call with tagged asset references like <blob>base64_id</blob>, this middleware intercepts the arguments, resolves each reference to a local cache path via the session's data manager, and replaces the argument value with the path string. This happens before FastMCP/Pydantic coerces arguments against the function signature, so parameters with their natural types (Path, bool, int, etc.) receive properly typed values instead of being mangled by premature coercion.

Resolution metadata (qualified name, cache path, parameter name) is stored in a ContextVar so that the tool callback can inject the asset into the kernel for the generated execution code.

Source code in src/agora_workbench/code_execution/data_access/resolution.py

def __init__(self, server: "CodeExecutionServer"):
    self.server = server

looks_like_qualified_name(value)

Check if a string looks like a DataLake qualified name.

Recognizes type-tagged artifact format: id Examples: abc123, xyz789

Parameters:

Name	Type	Description	Default
value	str	String to check	required

Returns:

Type	Description
bool	True if value matches the type-tagged artifact pattern

Source code in src/agora_workbench/code_execution/data_access/resolution.py

def looks_like_qualified_name(value: str) -> bool:
    """
    Check if a string looks like a DataLake qualified name.

    Recognizes type-tagged artifact format: <type>id</type>
    Examples: <blob>abc123</blob>, <sql>xyz789</sql>

    Args:
        value: String to check

    Returns:
        True if value matches the type-tagged artifact pattern
    """
    if not isinstance(value, str):
        return False

    stripped = value.strip()
    if not stripped:
        return False

    # Accept both <type>id</type> and unclosed <type>id (LLM sometimes omits closing tag)
    return bool(re.match(r"^<(\w+)>([^<>]+)</\1>$", stripped)) or bool(re.match(r"^<(\w+)>([^<>]+)$", stripped))

should_resolve_as_asset(value)

Determine if a parameter value should be resolved as a DataLake asset.

Detection is purely value-based: any string matching the type-tagged format <type>id</type> will be resolved.

Parameters:

Name	Type	Description	Default
value	Any	Parameter value to check	required

Returns:

Type	Description
bool	True if value is a type-tagged asset reference

Source code in src/agora_workbench/code_execution/data_access/resolution.py

def should_resolve_as_asset(value: Any) -> bool:
    """
    Determine if a parameter value should be resolved as a DataLake asset.

    Detection is purely value-based: any string matching the type-tagged
    format ``<type>id</type>`` will be resolved.

    Args:
        value: Parameter value to check

    Returns:
        True if value is a type-tagged asset reference
    """
    if not isinstance(value, str):
        return False

    return looks_like_qualified_name(value)