import logging import os from pathlib import Path from typing import Any, Dict, List, Optional import pandas as pd import httpx from agno.tools import Toolkit from agno.utils.log import logger as agno_logger try: from backend.ml_module.services.storage_service import MLStorageService from backend.ml_module.core.constants import StoragePaths except ImportError: from ml_module.services.storage_service import MLStorageService from ml_module.core.constants import StoragePaths from agno.run import RunContext logger = logging.getLogger(__name__) DATA_SOURCES_API_BASE_URL = os.environ.get("DATA_SOURCES_API_BASE_URL", "http://127.0.0.1:8000") DEFAULT_REQUEST_TIMEOUT = float(os.environ.get("TENANT_FILES_TOOL_TIMEOUT", "120")) class TenantFileToolkit(Toolkit): """ Toolkit for listing, reading, and writing files stored in MinIO, scoped securely to a tenant. This toolkit relies on `session_state` to inject the `tenant_id` to ensure that an agent cannot access files belonging to another tenant. """ def __init__(self, storage_service: Optional[MLStorageService] = None): super().__init__(name="tenant_file_toolkit") self.storage = storage_service or MLStorageService() self.api_base_url = DATA_SOURCES_API_BASE_URL self.register(self.list_tenant_assets_structured) self.register(self.list_tenant_assets) self.register(self.load_tenant_file_to_dataframe) self.register(self.stage_tenant_asset_for_ml) def _get_tenant_id(self, run_context: Optional["RunContext"] = None) -> str: """Extracts the tenant ID securely from the agent's run session.""" if run_context and run_context.session_state and "tenant_id" in run_context.session_state: return run_context.session_state["tenant_id"] logger.warning("No tenant_id found in Agent run_context or session_state. Falling back to default 'unknown_tenant'.") return "unknown_tenant" def _get_session_state(self, run_context: Optional["RunContext"] = None) -> Dict[str, str]: if run_context and run_context.session_state: return run_context.session_state return {} def _fetch_assets_from_api( self, tenant_id: str, jwt_token: str, page_size: int = 200, max_pages: int = 5, ) -> Dict[str, List[Dict[str, str]]]: headers = {"Authorization": f"Bearer {jwt_token}"} datasets: List[Dict[str, str]] = [] models: List[Dict[str, str]] = [] reports: List[Dict[str, str]] = [] other: List[Dict[str, str]] = [] timeout = httpx.Timeout(DEFAULT_REQUEST_TIMEOUT) try: with httpx.Client(base_url=self.api_base_url, timeout=timeout, follow_redirects=True) as client: page = 1 total_pages = 1 while page <= total_pages and page <= max_pages: try: resp = client.get( f"/api/v1/tenant-files/assets?page={page}&page_size={page_size}", headers=headers, ) except httpx.TimeoutException as exc: return { "error": ( f"Tenant files API timed out for tenant={tenant_id} " f"base_url={self.api_base_url} timeout={DEFAULT_REQUEST_TIMEOUT}s error={exc}" ) } except httpx.HTTPError as exc: return { "error": ( f"Tenant files API HTTP error for tenant={tenant_id} " f"base_url={self.api_base_url} error={exc}" ) } if resp.status_code != 200: return { "error": f"Tenant files API request failed ({resp.status_code}): {resp.text[:300]}" } payload = resp.json() items = payload.get("items", []) total_pages = int(payload.get("total_pages", 1) or 1) for item in items: filename = item.get("filename", "unknown") file_type = (item.get("file_type") or "").lower() created_at = item.get("created_at", "") size_mb = round(float(item.get("size_bytes", 0)) / (1024 * 1024), 2) asset_id = item.get("asset_id", "") record = { "path": f"{tenant_id}/files/{filename}", "asset_id": asset_id, "filename": filename, "file_type": file_type, "size_mb": size_mb, "last_modified": created_at, } if file_type in {"csv", "xlsx", "xls", "parquet"}: datasets.append(record) elif file_type in {"joblib", "pkl", "onnx"}: models.append(record) elif file_type in {"json", "md", "txt", "html"}: reports.append(record) else: other.append(record) page += 1 except Exception as exc: return { "error": ( f"Unexpected tenant files API error for tenant={tenant_id} " f"base_url={self.api_base_url} error={exc}" ) } return { "search_prefix": f"{tenant_id}/", "datasets": datasets, "models": models, "reports": reports, "other": other, } def _fetch_asset_preview_from_api( self, asset_id: str, jwt_token: str, page_size: int = 200, sheet_name: Optional[str] = None, ) -> Dict[str, Any]: headers = {"Authorization": f"Bearer {jwt_token}"} timeout = httpx.Timeout(DEFAULT_REQUEST_TIMEOUT) params: Dict[str, Any] = {"page": 1, "page_size": max(1, min(page_size, 500))} if sheet_name: params["sheet_name"] = sheet_name try: with httpx.Client(base_url=self.api_base_url, timeout=timeout, follow_redirects=True) as client: resp = client.get( f"/api/v1/tenant-files/assets/{asset_id}/preview", headers=headers, params=params, ) if resp.status_code != 200: return { "error": f"Tenant files preview failed ({resp.status_code}): {resp.text[:300]}" } return resp.json() except httpx.TimeoutException as exc: return { "error": ( f"Tenant files preview timed out asset_id={asset_id} " f"base_url={self.api_base_url} timeout={DEFAULT_REQUEST_TIMEOUT}s error={exc}" ) } except httpx.HTTPError as exc: return { "error": ( f"Tenant files preview HTTP error asset_id={asset_id} " f"base_url={self.api_base_url} error={exc}" ) } def _resolve_asset_from_path( self, file_path: str, tenant_id: str, jwt_token: str, run_context: Optional["RunContext"] = None, ) -> Optional[Dict[str, Any]]: catalog = self.list_tenant_assets_structured(prefix="", run_context=run_context) if not isinstance(catalog, dict): return None candidates: List[Dict[str, Any]] = [] for group in ("datasets", "models", "reports", "other"): candidates.extend(catalog.get(group, []) or []) normalized = file_path.strip().lower() basename = file_path.split("/")[-1].strip().lower() for item in candidates: item_path = str(item.get("path", "")).strip().lower() item_name = str(item.get("filename", "")).strip().lower() if normalized and (normalized == item_path or normalized == item_name): return item if basename and (basename == item_name or basename == item_path.split("/")[-1]): return item # Final safety: if caller passed tenant-prefixed path without filename metadata, # try exact filename match from tail segment. if basename: for item in candidates: if str(item.get("filename", "")).strip().lower() == basename: return item return None def _sanitize_local_filename(self, filename: str) -> str: candidate = Path(filename).name.strip() if not candidate or candidate in {".", ".."}: raise ValueError("Invalid filename") return candidate def _resolve_workspace_dir(self, run_context: Optional["RunContext"] = None) -> Path: session_state = self._get_session_state(run_context) workspace_value = str(session_state.get("workspace") or "").strip() if not workspace_value: raise ValueError("Missing workspace in run context") workspace = Path(workspace_value).resolve() workspace.mkdir(parents=True, exist_ok=True) return workspace def _download_all_rows_from_api( self, asset_id: str, jwt_token: str, ) -> pd.DataFrame: """Download all rows for a tenant asset by paginating the preview API. User-uploaded files live in tenant-files (not ml-projects), so direct MinIO reads fail with NoSuchKey. The preview endpoint is the correct access path. """ PAGE_SIZE = 500 headers = {"Authorization": f"Bearer {jwt_token}"} timeout = httpx.Timeout(DEFAULT_REQUEST_TIMEOUT) all_rows: list = [] columns: list = [] page = 1 with httpx.Client(base_url=self.api_base_url, timeout=timeout, follow_redirects=True) as client: while True: params: Dict[str, Any] = {"page": page, "page_size": PAGE_SIZE} resp = client.get( f"/api/v1/tenant-files/assets/{asset_id}/preview", headers=headers, params=params, ) if resp.status_code != 200: raise RuntimeError( f"Preview API failed ({resp.status_code}): {resp.text[:300]}" ) payload = resp.json() rows = payload.get("rows") or [] if not columns: columns = payload.get("columns") or [] all_rows.extend(rows) if len(rows) < PAGE_SIZE: break page += 1 return pd.DataFrame(all_rows) if all_rows else pd.DataFrame(columns=columns) def list_tenant_assets( self, prefix: str = "", run_context: Optional["RunContext"] = None ) -> str: """ Lists available CSV files, datasets, and reports for the tenant in the file storage cluster (MinIO). This tool should be used first to explore what data files are available before querying them. Args: prefix (str, optional): A specific folder prefix to list. If empty, it lists the root of the tenant's workspace. run_context: Agno RunContext (auto-injected). Returns: str: A formatted markdown representation of the available files and their sizes. """ structured = self.list_tenant_assets_structured(prefix=prefix, run_context=run_context) if isinstance(structured, str): return structured if structured.get("error"): return ( "Unable to reliably list tenant assets right now. " f"Tenant-files API error: {structured['error']}" ) search_prefix = structured.get("search_prefix", "") output = [f"## Assets for Tenant Workspace (`{search_prefix}`)"] def _format_lines(items): return [ f"- `{item['path']}` ({item['size_mb']} MB) [Last Modified: {item['last_modified']}]" for item in items ] if structured.get("datasets"): output.append("### Datasets (CSV/Excel/Parquet)") output.extend(_format_lines(structured["datasets"])) if structured.get("models"): output.append("\n### Models (.joblib/.pkl/.onnx)") output.extend(_format_lines(structured["models"])) if structured.get("reports"): output.append("\n### Reports (.json/.md/.txt/.html)") output.extend(_format_lines(structured["reports"])) if structured.get("other"): output.append("\n### Other Artifacts") output.extend(_format_lines(structured["other"])) if len(output) == 1: return f"No assets found for prefix: `{search_prefix}`." return "\n".join(output) def list_tenant_assets_structured( self, prefix: str = "", run_context: Optional["RunContext"] = None, ) -> Dict[str, Any]: """Return machine-friendly grouped asset catalog for a tenant.""" session_state = self._get_session_state(run_context) tenant_id = self._get_tenant_id(run_context) jwt_token = (session_state.get("supabase_jwt") or "").strip() api_error: Optional[str] = None # Primary path: use tenant-files API so agent sees the same assets as UI. if jwt_token: api_result = self._fetch_assets_from_api(tenant_id=tenant_id, jwt_token=jwt_token) if "error" not in api_result: return api_result api_error = str(api_result["error"]) logger.warning(api_error) search_prefix = f"{tenant_id}/" if prefix: search_prefix = f"{tenant_id}/{prefix.lstrip('/')}" if not self.storage.client: if api_error: return { "search_prefix": search_prefix, "datasets": [], "models": [], "reports": [], "other": [], "error": api_error, } return "Storage client is unavailable." try: objects = self.storage.client.list_objects( self.storage.bucket_name, prefix=search_prefix, recursive=True, ) datasets: List[Dict[str, str]] = [] models: List[Dict[str, str]] = [] reports: List[Dict[str, str]] = [] other: List[Dict[str, str]] = [] for obj in objects: path = obj.object_name file_info = { "path": path, "size_mb": round(obj.size / (1024 * 1024), 2), "last_modified": obj.last_modified.strftime('%Y-%m-%d %H:%M:%S'), } lower_path = path.lower() if lower_path.endswith((".csv", ".xlsx", ".xls", ".parquet")): datasets.append(file_info) elif lower_path.endswith((".joblib", ".pkl", ".onnx")): models.append(file_info) elif lower_path.endswith((".json", ".md", ".txt", ".html")): reports.append(file_info) else: other.append(file_info) result = { "search_prefix": search_prefix, "datasets": datasets, "models": models, "reports": reports, "other": other, } if api_error and not any([datasets, models, reports, other]): result["error"] = api_error return result except Exception as e: error_msg = f"Failed to list tenant assets: {str(e)}" logger.error(error_msg) return error_msg def load_tenant_file_to_dataframe( self, file_path: str, chunksize: Optional[int] = 10000, run_context: Optional["RunContext"] = None ) -> str: """ Reads a tenant dataset (CSV/XLSX/XLS/Parquet) into memory safely. Prefers tenant-files API preview so agent sees exactly what UI uploaded assets expose. Args: file_path (str): The full path to the file in MinIO (e.g., 'tenant_123/files/my_data.csv'). chunksize (int, optional): The number of rows to load at a time to prevent memory overflow. Defaults to 10000. run_context: Agno RunContext (auto-injected). Returns: str: A summary of the loaded DataFrame (columns, memory usage, head of the data), or an error message. """ tenant_id = self._get_tenant_id(run_context) session_state = self._get_session_state(run_context) jwt_token = (session_state.get("supabase_jwt") or "").strip() logger.info(f"Loading tenant file as DataFrame: {file_path}") # Primary path: resolve asset from tenant-files API and preview it. # This keeps behavior aligned with UI uploads and supports XLSX correctly. if jwt_token: resolved_asset = self._resolve_asset_from_path( file_path=file_path, tenant_id=tenant_id, jwt_token=jwt_token, run_context=run_context, ) if resolved_asset: asset_id = resolved_asset.get("asset_id") preview_payload = self._fetch_asset_preview_from_api( asset_id=asset_id, jwt_token=jwt_token, page_size=(chunksize or 100), ) if "error" in preview_payload: logger.error(preview_payload["error"]) return f"Failed to preview file `{file_path}` via tenant-files API: {preview_payload['error']}" rows = preview_payload.get("rows", []) or [] columns = preview_payload.get("columns", []) or [] if rows: df = pd.DataFrame(rows) else: df = pd.DataFrame(columns=columns) df_info = self._get_dataframe_summary(df, is_chunk=True) return ( f"Successfully previewed tenant asset `{resolved_asset.get('filename', file_path)}` " f"(asset_id: `{asset_id}`):\n{df_info}" ) # Fallback path: direct object read (legacy behavior) # Keep strict tenant-prefix check here for safety. if not file_path.startswith(f"{tenant_id}/"): return ( f"Access Denied: file path `{file_path}` is not in tenant scope `{tenant_id}/...`. " f"Try passing the exact filename from list_tenant_assets output." ) try: df_or_iterator = self.storage.load_dataframe(file_path, chunksize=chunksize) if chunksize: first_chunk = next(df_or_iterator) df_info = self._get_dataframe_summary(first_chunk, is_chunk=True) return f"Successfully read FIRST CHUNK of file `{file_path}`:\n{df_info}" df_info = self._get_dataframe_summary(df_or_iterator) return f"Successfully loaded file `{file_path}`:\n{df_info}" except Exception as e: error_msg = f"Failed to load file `{file_path}`: {str(e)}" logger.error(error_msg) return error_msg def stage_tenant_asset_for_ml( self, file_path: str, asset_id: Optional[str] = None, version: int = 1, run_context: Optional["RunContext"] = None, ) -> str: """ Resolves a tenant asset and materializes it into the ML workspace so downstream ML tools can load it without relying on arbitrary filesystem discovery. Args: file_path: Filename or tenant-scoped path from list_tenant_assets. asset_id: Optional explicit asset identifier. version: Asset version to stage. run_context: Agno RunContext (auto-injected). Returns: str: Success message with the staged local path. """ tenant_id = self._get_tenant_id(run_context) session_state = self._get_session_state(run_context) jwt_token = (session_state.get("supabase_jwt") or "").strip() if not jwt_token: return "Failed to stage tenant asset: missing authenticated tenant session." resolved_asset: Optional[Dict[str, Any]] = None if asset_id: catalog = self.list_tenant_assets_structured(prefix="", run_context=run_context) if isinstance(catalog, dict): for group in ("datasets", "models", "reports", "other"): for item in catalog.get(group, []) or []: if str(item.get("asset_id", "")).strip() == asset_id: resolved_asset = item break if resolved_asset: break if resolved_asset is None: resolved_asset = self._resolve_asset_from_path( file_path=file_path, tenant_id=tenant_id, jwt_token=jwt_token, run_context=run_context, ) if resolved_asset is None: return ( f"Failed to stage tenant asset `{file_path}`. " "Call list_tenant_assets first and pass the exact filename or asset_id." ) resolved_asset_id = str(resolved_asset.get("asset_id") or asset_id or "").strip() filename = self._sanitize_local_filename(str(resolved_asset.get("filename") or file_path)) if not resolved_asset_id: return f"Failed to stage tenant asset `{filename}`: missing asset_id metadata." try: # Download via the tenant-files preview API (paginated) — user files # live in the tenant-files service, NOT in the ml-projects MinIO bucket. df = self._download_all_rows_from_api( asset_id=resolved_asset_id, jwt_token=jwt_token, ) workspace = self._resolve_workspace_dir(run_context) processed_dir = workspace / "processed" processed_dir.mkdir(parents=True, exist_ok=True) local_path = processed_dir / filename suffix = local_path.suffix.lower() if suffix == ".csv": df.to_csv(local_path, index=False) elif suffix in {".xlsx", ".xls"}: df.to_excel(local_path, index=False) elif suffix == ".parquet": df.to_parquet(local_path, index=False) else: local_path = local_path.with_suffix(".csv") df.to_csv(local_path, index=False) return ( f"Successfully staged tenant asset `{filename}` to `{local_path.as_posix()}` " f"for ML use. Rows: {len(df)}, Columns: {len(df.columns)}, asset_id: `{resolved_asset_id}`." ) except Exception as exc: logger.error("Failed to stage tenant asset %s for tenant %s: %s", filename, tenant_id, exc) return f"Failed to stage tenant asset `{filename}`: {exc}" def _get_dataframe_summary(self, df: pd.DataFrame, is_chunk: bool = False) -> str: """Generates a markdown summary of a pandas DataFrame.""" import io buffer = io.StringIO() df.info(buf=buffer) info_str = buffer.getvalue() try: head_md = df.head(5).to_markdown() except ImportError: head_md = f"```text\n{df.head(5).to_string(index=False)}\n```" summary = [ f"**Rows:** {len(df)}{' (in chunk)' if is_chunk else ''}", f"**Columns:** {len(df.columns)}", f"\n**Data Info:**", f"```text\n{info_str}\n```", f"\n**Sample Data (Head):**", head_md ] return "\n".join(summary)