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)

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).

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 downstream resource access when Azure services are configured. Falls back to local-only mode when DATA_LAKE_SEARCH_ENDPOINT is not set.

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

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,
):
    """
    Initialize the data manager.

    Uses a credential chain (``az login`` MSAL cache locally, managed
    identity in production) for downstream resource access when Azure
    services are configured. Falls back to local-only mode when
    ``DATA_LAKE_SEARCH_ENDPOINT`` is not set.

    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.
    """
    self._cache_dir = Path(tempfile.mkdtemp(prefix="data_lake_cache_"))
    self._cache_index = {}  # Maps artifact_id -> cache file path
    self._url_cache: dict[str, str] = {}  # Maps artifact_id -> resolved blob URL

    search_endpoint = os.getenv("DATA_LAKE_SEARCH_ENDPOINT")
    self._credential_init_error: str | None = 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))

    if search_endpoint:
        # Azure mode: add blob fetcher and search client
        self._blob_details_index = os.getenv("DATA_LAKE_BLOB_DETAILS_INDEX", "blob-details")
        try:
            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)
            self._fetchers.append(BlobFetcher(credential=self._credential))

            self._search_client = SearchClient(
                endpoint=search_endpoint,
                index_name=self._blob_details_index,
                credential=self._credential,
            )
            LOGGER.info(f"Initialized blob-details search client: {search_endpoint}/{self._blob_details_index}")
        except (ImportError, RuntimeError, TypeError, ValueError) as e:
            self._credential_init_error = f"{type(e).__name__}: {e}"
            self._credential = None
            self._search_client = None
            LOGGER.warning(f"Failed to initialize Azure data access components: {e}")
    else:
        self._credential = None
        self._search_client = None
        self._blob_details_index = None
        LOGGER.info("DataLakeDataManager running in local-only mode (no Azure Search endpoint configured)")

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"{agent_guidance.ASSET_TAG_FORMAT} {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}. {agent_guidance.ASSET_TAG_FORMAT}")

    # 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 search client
    if hasattr(self, "_search_client") and self._search_client:
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._search_client.close())
        except RuntimeError:
            try:
                loop = asyncio.new_event_loop()
                loop.run_until_complete(self._search_client.close())
                loop.close()
            except Exception as e:
                LOGGER.debug(f"Error closing search client: {e}")

    # Close managed identity credential
    if hasattr(self, "_credential") and self._credential is not None:
        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()
    self._url_cache.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}")

    if hasattr(self, "_search_client") and self._search_client:
        try:
            await self._search_client.close()
        except Exception as e:
            LOGGER.debug(f"Error closing search client: {e}")

    if hasattr(self, "_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)