Spaces:
Running
Running
File size: 61,693 Bytes
b8277c4 a8c9ee8 b8277c4 a8c9ee8 b8277c4 a8c9ee8 b8277c4 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 b3f6656 a8c9ee8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 | """DataSourcesSQLToolkit: Multi-tenant SQL execution toolkit for SQL Agents.
This toolkit provides agents with tools to interact with the centralized
data_sources API for fetching schemas, instructions, and executing raw SQL
queries securely across multiple tenants and data sources.
"""
import httpx
import asyncio
import json
import logging
from datetime import datetime
from typing import Dict, Any, List, Optional, Union
from agno.tools import Toolkit
from agno.run import RunContext
from agno.utils.log import logger
# Import hybrid keyword utilities with defensive fallback
try:
from hybrid_keyword_utils import extract_hybrid_keywords, create_gemini_semantic_client
_hybrid_utils_available = True
except ImportError:
try:
from backend.SQL_Agent.hybrid_keyword_utils import extract_hybrid_keywords, create_gemini_semantic_client
_hybrid_utils_available = True
except ImportError:
logger.warning(
"Could not import hybrid_keyword_utils. Semantic extraction will be unavailable. "
"This is expected if running without the hybrid_keyword_utils module."
)
extract_hybrid_keywords = None
create_gemini_semantic_client = None
_hybrid_utils_available = False
# Base URL for your data_sources API.
# Ensure this matches where your backend API is running.
DATA_SOURCES_API_BASE_URL = "http://127.0.0.1:8000"
# Default timeout for API calls (seconds). Set higher for potentially long queries.
DEFAULT_API_TIMEOUT = 120.0
# Per-call timeout budget (can be overridden per method)
DEFAULT_REQUEST_TIMEOUT = 120.0
class DataSourcesSQLToolkit(Toolkit):
"""
A multi-tenant toolkit for SQL Agents using the Data Sources API
to fetch schemas/instructions and execute RAW SQL queries synchronously.
MODIFIED: This toolkit is now tenant-safe. It reads tenant_id, source_name,
and api_key directly from the agent's session_state for every call,
ensuring proper data isolation and security.
"""
def __init__(
self,
api_base_url: str = DATA_SOURCES_API_BASE_URL,
timeout: float = DEFAULT_API_TIMEOUT,
api_key: Optional[str] = None # This is a FALLBACK key only
):
"""
Initialize the DataSourcesSQLToolkit.
Args:
api_base_url: Base URL for the data sources API
timeout: Default timeout for API client
api_key: Optional FALLBACK API key (e.g., from env vars)
"""
self.api_base_url = api_base_url
self.api_key = api_key # Store fallback
# --- CRITICAL FIX ---
# DO NOT set the API key in the default client headers.
# This client is shared across all requests and must be stateless.
self.client = httpx.Client(
base_url=self.api_base_url,
timeout=timeout,
headers={} # Headers will be provided per-request
)
# --- Async client for non-blocking calls (preferred) ---
# Keep a separate sync client for backwards compatibility.
try:
self._async_client = httpx.AsyncClient(
base_url=self.api_base_url,
timeout=timeout,
headers={}
)
except Exception:
# If AsyncClient cannot be created for some reason, fall back to sync client.
self._async_client = None
# Cache for search results per (tenant_id, keywords_tuple, source_names_tuple)
self._search_cache: Dict[tuple, Dict[str, Any]] = {}
# Initialize semantic client for hybrid keyword extraction
self._semantic_client = None
if _hybrid_utils_available and create_gemini_semantic_client:
try:
self._semantic_client = create_gemini_semantic_client()
logger.info("Gemini semantic client initialized for hybrid keyword extraction")
except Exception as e:
logger.warning(f"Could not initialize semantic client: {e}")
# Reference to the agent for session_state injection
self._agent_ref = None
logger.info(f"DataSourcesSQLToolkit initialized for API: {self.api_base_url} with timeout {timeout}s")
# Define the tools this toolkit provides to the agent
super().__init__(
name="data_sources_sql_tools",
tools=[
self.find_relevant_tables, # NEW: High-level intelligent table discovery
self.search_schema,
self.list_sources, # NEW: Explicit source listing
self.get_available_sources_and_schema,
self.get_source_instructions,
self.execute_sql_query, # The primary execution tool
self.save_query_to_tenant_csv, # NEW: Bridge tool to ML files
self.compose_dataset_workflow,
],
instructions="""CRITICAL WORKFLOW:
1. After calling find_relevant_tables, you MUST immediately call execute_sql_query.
2. NEVER stop after just searching - always execute the SQL query.
3. ALWAYS use execute_sql_query when user asks for data counts or lists.
4. NEVER just describe what you would do - actually call the tool and explain the results.
5. After getting results, provide a clear natural language answer to the user."""
)
def set_agent_ref(self, agent_ref):
"""
Set a reference to the agent for accessing session_state during tool execution.
This is crucial for multi-tenant context injection.
Args:
agent_ref: Reference to the Agno agent
"""
self._agent_ref = agent_ref
logger.info("Agent reference set in toolkit for session_state injection")
def _get_session_state_from_agent(self) -> Optional[Dict[str, Any]]:
"""
Get the current session state from the agent reference.
Returns:
session_state dict from agent if available, None otherwise
"""
if self._agent_ref and hasattr(self._agent_ref, 'session_state'):
return self._agent_ref.session_state
return None
# --- NEW HELPER: Get Per-Request Headers ---
def _get_request_headers(self, session_state: Optional[Dict[str, Any]]) -> Dict[str, str]:
"""
Safely builds headers for a single request.
It prioritizes the supabase_jwt from the current session_state.
"""
headers = {}
jwt_token = None
if session_state:
jwt_token = session_state.get("supabase_jwt")
if jwt_token:
logger.debug("Using JWT token from session_state")
if jwt_token:
headers["Authorization"] = f"Bearer {jwt_token}"
else:
logger.warning("No JWT token found in session_state. API call may fail.")
return headers
# Helper to run async coroutines from sync context when safe
def _run_coro_sync(self, coro):
"""
Run an async coroutine from sync code when no event loop is running.
If an event loop is already running, return the coroutine (caller must await).
"""
try:
loop = asyncio.get_running_loop()
except RuntimeError:
# No running loop, safe to run
return asyncio.run(coro)
else:
# An event loop is running; return the coroutine so caller can await it.
return coro
# --- Helper Methods ---
def _normalize_error(self, error_msg: str, code: str = "UNKNOWN_ERROR", hint: Optional[str] = None) -> Dict[str, Any]:
"""
Normalize error responses to a consistent format.
Args:
error_msg: The error message
code: Error code for categorization
hint: Optional hint for resolution
Returns:
Normalized error dict with error, code, and optional hint
"""
result = {
"error": error_msg,
"code": code
}
if hint:
result["hint"] = hint
return result
# --- CRITICAL FIX: Force session_state ---
def _resolve_tenant(self, tenant_id: Optional[str], session_state: Optional[Dict[str, Any]]) -> Union[str, Dict[str, Any]]:
"""
Resolve tenant_id ONLY from session state.
Ignores any tenant_id passed by the LLM to prevent confusion.
"""
if not session_state:
return self._normalize_error("Session state not found.", "SESSION_ERROR")
resolved = (session_state.get("tenant_id", "") or "").strip()
if not resolved:
return self._normalize_error(
"Tenant ID not found in session_state.",
code="MISSING_TENANT_CONTEXT",
hint="The /tenant-run endpoint must inject 'tenant_id'."
)
return resolved
# --- CRITICAL FIX: Force session_state ---
def _resolve_source(self, source_name: Optional[str], session_state: Optional[Dict[str, Any]]) -> Union[str, Dict[str, Any]]:
"""
Resolve source_name ONLY from session state.
Ignores any source_name passed by the LLM.
"""
if not session_state:
return self._normalize_error("Session state not found.", "SESSION_ERROR")
resolved = (session_state.get("source_name", "") or "").strip()
if not resolved:
return self._normalize_error(
"Source name not found in session_state.",
code="MISSING_SOURCE_CONTEXT",
hint="The /tenant-run endpoint must inject 'source_name'."
)
return resolved
# --- Helper Function ---
def _format_schema_for_llm(self, source_name: str, schema_json_string: Optional[str]) -> str:
"""Formats a single source's JSON schema string (received from API) for the LLM."""
if not schema_json_string:
return f"**Source: `{source_name}`**\n Schema: <Not Available or Empty>"
try:
# The API returns the raw schema string, which should be JSON parsable
schema_data_list = json.loads(schema_json_string)
# Expecting the format: [{'schema_name': '...', 'tables': [...]}]
if not schema_data_list or not isinstance(schema_data_list, list) or not isinstance(schema_data_list[0], dict):
logger.warning(f"Received schema for '{source_name}' is not in expected list-of-dict format.")
# Attempt to format directly if it's just a dict (less robust)
if isinstance(schema_data_list, dict):
schema_info = schema_data_list
else:
return f"**Source: `{source_name}`**\n Schema: <Invalid Format Received>"
else:
schema_info = schema_data_list[0] # Assuming one schema object per source from API
output_parts = [f"**Source: `{source_name}`** (DB/Schema Name: {schema_info.get('schema_name', 'N/A')})\n"]
tables = schema_info.get("tables", [])
if not tables:
output_parts.append(" *No tables found or schema details missing.*")
else:
for table in tables:
table_name = table.get('table_name', 'N/A')
output_parts.append(f" **Table: `{table_name}`**")
fields = table.get('fields', [])
if not fields:
output_parts.append(" *No column details available.*")
else:
output_parts.append(" Columns:")
for field in fields:
col_name = field.get('name', 'N/A')
col_type = field.get('type', 'Unknown')
col_example = field.get('example', '') # Example might be missing or empty
# Add example only if present and non-empty to avoid clutter
example_str = f", example: '{col_example}'" if col_example else ""
output_parts.append(f" - `{col_name}` (type: {col_type}{example_str})")
# Add sample rows if available
example_rows = table.get('example_rows', [])
if example_rows:
output_parts.append(" Sample Data:")
for row in example_rows[:3]: # Limit to 3 rows
row_str = ', '.join(f"{k}={v}" for k, v in list(row.items())[:5]) # Limit columns shown
output_parts.append(f" {row_str}")
output_parts.append("") # Add a blank line after each table definition for readability
return "\n".join(output_parts)
except json.JSONDecodeError:
logger.error(f"Failed to parse schema JSON for source '{source_name}'. Received: {schema_json_string[:200]}...")
return f"**Source: `{source_name}`**\n Schema: <Error Parsing JSON>"
except (IndexError, KeyError, TypeError) as e:
logger.error(f"Error processing schema structure for source '{source_name}': {e}")
return f"**Source: `{source_name}`**\n Schema: <Error Processing Schema Structure>"
except Exception as e: # Catch any other unexpected errors during formatting
logger.exception(f"Unexpected error formatting schema for '{source_name}': {e}")
return f"**Source: `{source_name}`**\n Schema: <Unexpected Formatting Error>"
# --- Agent Tools ---
def list_sources(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None
) -> Dict[str, Any]:
"""List all available data sources for a tenant."""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context
request_headers = self._get_request_headers(session_state)
logger.info(f"Tool: list_sources called for tenant '{tenant_id}'")
try:
url = f"/api/v1/data-sources/my-tenant/list"
response = self.client.get(url, timeout=DEFAULT_REQUEST_TIMEOUT, headers=request_headers)
if response.status_code == 200:
data = response.json()
result = {
"available_sources": data.get("available_sources", []),
"count": data.get("count", 0),
"tenant_id": data.get("tenant_id", tenant_id),
"assets": [],
"datasets": [],
}
# Unified catalog extension: include tenant file assets when available
try:
assets_resp = self.client.get(
"/api/v1/tenant-files/assets?page=1&page_size=200",
timeout=DEFAULT_REQUEST_TIMEOUT,
headers=request_headers,
)
if assets_resp.status_code == 200:
assets_payload = assets_resp.json()
result["assets"] = assets_payload.get("items", [])
result["datasets"] = [
item for item in result["assets"]
if item.get("file_type") in {"csv", "parquet", "xlsx", "xls"}
]
except Exception as assets_err:
logger.debug(f"Tenant files catalog fetch skipped/failed: {assets_err}")
logger.info(f"Found {result['count']} source(s) for tenant '{tenant_id}': {result['available_sources']}")
if session_state is not None:
cache = session_state.setdefault("available_sources_cache", {})
cache[tenant_id] = result["available_sources"]
return result
else:
return self._normalize_error(f"Failed to list sources: {response.text[:200]}", code="API_ERROR")
except Exception as e:
logger.exception(f"Error in list_sources: {e}")
return self._normalize_error(f"Error: {str(e)}", code="UNEXPECTED_ERROR")
def find_relevant_tables(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
question: str = "",
concepts: Optional[List[str]] = None,
source_names: Optional[List[str]] = None,
include_samples: bool = False
) -> Dict[str, Any]:
"""
Intelligently find relevant tables for a user question using hybrid keyword extraction.
This is the PRIMARY tool for schema discovery - it combines deterministic keywords,
semantic hints from LLM, and agent-provided concepts.
:param run_context: Agno RunContext containing session_state (auto-injected by framework).
:param tenant_id: The ID of the tenant. MANDATORY.
:param question: The original user question. MANDATORY.
:param concepts: Optional list of high-level concepts identified by the agent (e.g., ["revenue", "customers"])
:param source_names: Optional list to filter specific sources. None = search all.
:param include_samples: Whether to include example rows in results (default: False).
:return: Dict with 'formatted_schema_string', 'matches', 'total_matches', 'keyword_breakdown'
"""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context # type: ignore[assignment]
if not question or not question.strip():
return self._normalize_error("Question is required.", "missing_parameter", "Provide a non-empty question string.")
logger.info(f"Tool: find_relevant_tables called for tenant '{tenant_id}' with question: '{question[:100]}...'")
logger.info(f" Agent concepts: {concepts}")
# --- FIX for [None] Pydantic Error ---
if source_names and all(s is None for s in source_names):
source_names = None
if source_names is None and session_state:
preferred_source = session_state.get("source_name") or session_state.get("preferred_source")
if preferred_source:
source_names = [preferred_source]
# Check if we have cached results for this exact question in session state
if session_state:
cache_key = f"find_relevant|{tenant_id}|{question}"
if "keyword_extraction_cache" not in session_state:
session_state["keyword_extraction_cache"] = {}
if cache_key in session_state["keyword_extraction_cache"]:
logger.info(f"Using cached keyword extraction for question")
keyword_result = session_state["keyword_extraction_cache"][cache_key]
merged_keywords = keyword_result["combined"]
else:
# Perform hybrid keyword extraction
if extract_hybrid_keywords:
keyword_result = extract_hybrid_keywords(
question=question,
llm_concepts=concepts,
semantic_client=self._semantic_client
)
# Cache the extraction result
session_state["keyword_extraction_cache"][cache_key] = keyword_result
merged_keywords = keyword_result["combined"]
logger.info(f"Keyword extraction breakdown:")
logger.info(f" Base (deterministic): {keyword_result['base']}")
logger.info(f" Semantic (LLM hints): {keyword_result['semantic']}")
logger.info(f" Concepts (agent): {keyword_result['concepts']}")
logger.info(f" Combined: {merged_keywords}")
else:
# Fallback if hybrid extraction not available
logger.warning("Hybrid keyword extraction not available, using simple fallback")
# Simple fallback: split question and filter stopwords
words = question.lower().split()
stopwords = {'what', 'when', 'where', 'who', 'which', 'how', 'show', 'give', 'the', 'a', 'an', 'is', 'are'}
base_kw = [w for w in words if w not in stopwords and len(w) > 2]
merged_keywords = base_kw + (concepts or [])
keyword_result = {
'base': base_kw,
'semantic': [],
'concepts': concepts or [],
'combined': merged_keywords
}
session_state["keyword_extraction_cache"][cache_key] = keyword_result
else:
# No session state, extract without caching
if extract_hybrid_keywords:
keyword_result = extract_hybrid_keywords(
question=question,
llm_concepts=concepts,
semantic_client=self._semantic_client
)
merged_keywords = keyword_result["combined"]
else:
# Simple fallback
words = question.lower().split()
stopwords = {'what', 'when', 'where', 'who', 'which', 'how', 'show', 'give', 'the', 'a', 'an', 'is', 'are'}
base_kw = [w for w in words if w not in stopwords and len(w) > 2]
merged_keywords = base_kw + (concepts or [])
keyword_result = {
'base': base_kw,
'semantic': [],
'concepts': concepts or [],
'combined': merged_keywords
}
# Use the existing search_schema with merged keywords
search_result = self.search_schema(
run_context=run_context,
tenant_id=tenant_id,
keywords=merged_keywords,
source_names=source_names,
include_samples=include_samples,
original_question=question,
keyword_metadata=keyword_result
)
# Add keyword breakdown to result for observability
if "error" not in search_result:
search_result["keyword_breakdown"] = keyword_result
search_result["original_question"] = question
if session_state is not None:
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_question"] = question
metadata["last_keyword_breakdown"] = keyword_result
if search_result.get("matches"):
metadata["last_schema_matches"] = search_result["matches"]
logger.info(f"find_relevant_tables completed: {search_result.get('total_matches', 0)} matches found")
return search_result
def search_schema(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
keywords: Optional[List[str]] = None,
source_names: Optional[List[str]] = None,
include_samples: bool = False,
original_question: Optional[str] = None,
keyword_metadata: Optional[Dict[str, List[str]]] = None
) -> Dict[str, Any]:
"""
Search for relevant tables across tenant data sources using keywords.
Returns ranked tables based on keyword matches in table/column names and descriptions.
Results are cached in session_state for reuse in follow-up queries.
:param run_context: Agno RunContext containing session_state (auto-injected by framework).
:param tenant_id: The ID of the tenant. MANDATORY.
:param keywords: List of search keywords (e.g., ['user', 'order', 'payment']). MANDATORY.
:param source_names: Optional list to filter specific sources. None = search all.
:param include_samples: Whether to include example rows in results (default: False).
:param original_question: Optional original user question for logging/analytics.
:param keyword_metadata: Optional dict with keyword breakdown (base, semantic, concepts) for observability.
:return: Dict with 'formatted_schema_string', 'matches', 'available_sources', 'total_matches'
"""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context # type: ignore[assignment]
# CRITICAL: Get per-request headers
request_headers = self._get_request_headers(session_state)
if not keywords or not isinstance(keywords, list):
return self._normalize_error("Keywords must be a non-empty list.", "invalid_parameter", "Provide keywords as a list of strings.")
# Normalize and cache key
normalized_keywords = tuple(sorted([k.lower().strip() for k in keywords if k.strip()]))
cache_key = (tenant_id, normalized_keywords, tuple(source_names or []))
# Check session_state cache first (across agent calls in same session)
if session_state and "schema_search_cache" in session_state:
cache_key_str = f"{tenant_id}|{normalized_keywords}|{source_names or []}"
if cache_key_str in session_state["schema_search_cache"]:
logger.info(f"Returning session-cached search results for tenant '{tenant_id}', keywords: {keywords}")
cached = session_state["schema_search_cache"][cache_key_str].copy()
cached['cache_hit'] = True
cached['cache_source'] = 'session_state'
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_schema_search"] = {
"tenant_id": tenant_id,
"keywords": list(normalized_keywords),
"include_samples": include_samples,
"total_matches": cached.get("total_matches", 0),
"cache_hit": True,
"cache_source": 'session_state'
}
metadata["last_schema_available_sources"] = cached.get("available_sources", [])
return cached
# Check toolkit-level cache (within single agent execution)
if cache_key in self._search_cache:
logger.info(f"Returning toolkit-cached search results for tenant '{tenant_id}', keywords: {keywords}")
cached = self._search_cache[cache_key].copy()
cached['cache_hit'] = True
cached['cache_source'] = 'toolkit'
if session_state is not None:
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_schema_search"] = {
"tenant_id": tenant_id,
"keywords": list(normalized_keywords),
"include_samples": include_samples,
"total_matches": cached.get("total_matches", 0),
"cache_hit": True,
"cache_source": 'toolkit'
}
metadata["last_schema_available_sources"] = cached.get("available_sources", [])
return cached
logger.info(f"Tool: search_schema called for tenant '{tenant_id}' with keywords: {keywords}")
try:
url = "/api/v1/data-sources/schema/search"
payload = {
"tenant_id": tenant_id,
"keywords": list(keywords),
"include_samples": include_samples
}
if source_names:
payload["source_names"] = source_names
response = self.client.post(url, json=payload, timeout=DEFAULT_REQUEST_TIMEOUT, headers=request_headers)
if response.status_code == 200:
data = response.json()
result = {
"formatted_schema_string": data.get("formatted_schema_string", ""),
"matches": data.get("matches", []),
"available_sources": data.get("available_sources", []),
"total_matches": data.get("total_matches", 0),
"cache_hit": False
}
# Cache in toolkit-level cache
self._search_cache[cache_key] = result.copy()
# Cache in session_state if available
if session_state:
if "schema_search_cache" not in session_state:
session_state["schema_search_cache"] = {}
cache_key_str = f"{tenant_id}|{normalized_keywords}|{source_names or []}"
session_state["schema_search_cache"][cache_key_str] = result.copy()
logger.info(f"Schema search completed: {result['total_matches']} matches")
return result
else:
error_text = response.text[:200]
logger.error(f"Schema search failed: {error_text}")
return self._normalize_error(f"Schema search failed: {error_text}", "api_error")
except Exception as e:
logger.exception(f"Search schema error: {e}")
return self._normalize_error(f"Search error: {str(e)}", "UNEXPECTED_ERROR")
def _fallback_schema_search(self, run_context: Optional[RunContext], tenant_id: str, keywords: List[str], source_names: Optional[List[str]]) -> Dict[str, Any]:
"""Fallback when /schema/search endpoint unavailable - uses basic filtering."""
logger.info("Using fallback schema search with client-side filtering")
# Get all schemas
full_result = self.get_available_sources_and_schema(run_context=run_context, tenant_id=tenant_id, keywords=None)
if "error" in full_result:
return full_result
# Simple keyword matching on the formatted string
schema_str = full_result.get("formatted_schema_string", "")
available_sources = full_result.get("available_sources", [])
# Basic filtering: check if any keyword appears in the schema
matched = any(kw.lower() in schema_str.lower() for kw in keywords)
if matched:
return {
"formatted_schema_string": schema_str,
"matches": [], # No detailed matches in fallback
"available_sources": available_sources,
"total_matches": len(available_sources),
"cache_hit": False,
"fallback_mode": True
}
else:
return {
"formatted_schema_string": f"No tables found matching keywords: {', '.join(keywords)}\n\nAvailable sources: {', '.join(available_sources)}",
"matches": [],
"available_sources": available_sources,
"total_matches": 0,
"cache_hit": False,
"fallback_mode": True
}
# --- Agent Tools ---
def get_available_sources_and_schema(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
keywords: Optional[List[str]] = None
) -> Dict[str, Any]:
"""
Retrieves the list of available data source names and their schemas for the specified tenant.
Use this first to understand which data sources ('source_name') and tables are available.
:param run_context: Agno RunContext containing session_state (auto-injected by framework).
:param tenant_id: The ID of the tenant (e.g., 'acme-corp'). MANDATORY.
:param keywords: Optional list of keywords to filter relevant tables. If provided, uses search_schema.
:return: Dict containing 'formatted_schema_string' (all schemas combined) and 'available_sources' (list of names). Returns 'error' on failure.
"""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context # type: ignore[assignment]
# CRITICAL: Get per-request headers
request_headers = self._get_request_headers(session_state)
# If keywords provided, delegate to search_schema with run_context
if keywords:
logger.info(f"Tool: get_available_sources_and_schema called with keywords - delegating to search_schema")
return self.search_schema(run_context=run_context, tenant_id=tenant_id, keywords=keywords)
logger.info(f"Tool: get_available_sources_and_schema called for tenant '{tenant_id}'")
# Use the dedicated list_sources endpoint to get available sources
list_result = self.list_sources(run_context=run_context, tenant_id=tenant_id)
# Check if list_sources returned an error
if "error" in list_result:
logger.error(f"Failed to list sources for tenant '{tenant_id}': {list_result.get('error')}")
return list_result # Return the normalized error directly
# Extract source names from the response
source_names = list_result.get("available_sources", [])
if not source_names:
logger.warning(f"No source configurations found for tenant '{tenant_id}'.")
return {
"formatted_schema_string": "No data sources configured for this tenant.",
"available_sources": []
}
logger.info(f"Found {len(source_names)} sources for tenant '{tenant_id}': {source_names}")
# Fetch schema for each identified source name
all_schemas_formatted = []
fetch_errors = []
for source_name in source_names:
try:
schema_url = f"/api/v1/data-sources/schema/{source_name}"
logger.debug(f"Fetching schema from {schema_url}")
schema_response = self.client.get(
schema_url,
timeout=DEFAULT_REQUEST_TIMEOUT,
headers=request_headers # Use per-request headers
)
if schema_response.status_code == 200:
schema_data = schema_response.json()
schema_json_string = schema_data.get("schema_data")
formatted_schema = self._format_schema_for_llm(source_name, schema_json_string)
all_schemas_formatted.append(formatted_schema)
elif schema_response.status_code == 404:
logger.warning(f"Schema API returned 404 for source '{source_name}', tenant '{tenant_id}'.")
all_schemas_formatted.append(f"**Source: `{source_name}`**\n Schema: <Not Found for Tenant>")
fetch_errors.append(source_name)
else:
error_text = schema_response.text[:200] # Limit error text length
logger.warning(f"Failed to get schema for source '{source_name}' (Status {schema_response.status_code}): {error_text}")
all_schemas_formatted.append(f"**Source: `{source_name}`**\n Schema: <Error: {schema_response.status_code}>")
fetch_errors.append(source_name)
except httpx.RequestError as e:
logger.error(f"API connection error fetching schema for '{source_name}': {e}")
all_schemas_formatted.append(f"**Source: `{source_name}`**\n Schema: <API Connection Error>")
fetch_errors.append(source_name)
except Exception as e: # Catch other errors during processing for a single source
logger.exception(f"Unexpected error processing schema for '{source_name}': {e}")
all_schemas_formatted.append(f"**Source: `{source_name}`**\n Schema: <Unexpected Error Processing Schema>")
fetch_errors.append(source_name)
# Combine successfully fetched schemas into one string
# Separate each source's schema clearly
final_schema_string = "\n\n---\n\n".join(all_schemas_formatted)
# Report any errors clearly
if fetch_errors:
error_message = f"Note: Failed to retrieve schema for the following sources: {', '.join(fetch_errors)}."
final_schema_string += f"\n\n**Warning:** {error_message}"
# Return the combined string and the list of sources attempted/found
return {
"formatted_schema_string": final_schema_string,
"available_sources": source_names # Return the names found via config
}
def get_source_instructions(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
source_name: Optional[str] = None
) -> Dict[str, str]:
"""
Retrieves special instructions (like SQL dialect, syntax rules, function usage)
for a specific data source. Call this *after* choosing a source_name
from get_available_sources_and_schema and *before* writing SQL.
**IMPORTANT**: Instructions are automatically cached in session_state for the agent
to use during SQL generation. The agent can access them via:
session_state['source_instructions'][source_name]
:param run_context: Agno RunContext containing session_state (auto-injected by framework).
:param tenant_id: The ID of the tenant (e.g., 'acme-corp'). MANDATORY.
:param source_name: The name of the data source (e.g., 'production_db'). MANDATORY.
:return: Dictionary containing 'instructions' string or 'error'.
"""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context # type: ignore[assignment]
source_context = self._resolve_source(source_name, session_state)
if isinstance(source_context, dict) and source_context.get("error"):
return source_context
source_name = source_context # type: ignore[assignment]
# CRITICAL: Get per-request headers
request_headers = self._get_request_headers(session_state)
logger.info(f"Tool: get_source_instructions called for tenant '{tenant_id}', source '{source_name}'")
# Check session_state cache first
if session_state and "source_instructions" in session_state:
cached_instructions = session_state["source_instructions"].get(source_name)
if cached_instructions:
logger.info(f"Returning session-cached instructions for source '{source_name}'")
return {"instructions": cached_instructions}
try:
url = f"/api/v1/data-sources/instructions/{source_name}"
logger.debug(f"Fetching instructions from {url}")
response = self.client.get(
url,
timeout=DEFAULT_REQUEST_TIMEOUT,
headers=request_headers # Use per-request headers
)
if response.status_code == 200:
data = response.json()
instructions = data.get("instructions", "No specific instructions provided for this source.")
connector_type = data.get("connector_type", "unknown")
logger.info(f"Successfully fetched instructions for '{source_name}' ({connector_type}).")
# Store in session_state for agent reuse
if session_state is not None:
# Initialize source_instructions dict if needed
if "source_instructions" not in session_state:
session_state["source_instructions"] = {}
# Cache instructions by source name
session_state["source_instructions"][source_name] = instructions
# Also track metadata
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_source_instructions"] = {
"tenant_id": tenant_id,
"source_name": source_name,
"connector_type": connector_type,
"timestamp": datetime.now().isoformat()
}
logger.debug(f"Cached instructions for '{source_name}' in session_state")
return {"instructions": instructions}
elif response.status_code == 404:
logger.warning(f"Instructions API returned 404 for source '{source_name}', tenant '{tenant_id}'.")
return self._normalize_error(
f"Instructions not found for source '{source_name}' for this tenant.",
"not_found",
"Verify the source_name exists for this tenant via list_sources()."
)
else:
error_text = response.text[:200]
logger.error(f"Failed to get instructions for '{source_name}' (Status {response.status_code}): {error_text}")
return self._normalize_error(
f"Failed to get instructions (Status {response.status_code}): {error_text}",
"api_error",
"Check API logs and verify the /instructions endpoint is operational."
)
except httpx.RequestError as e:
logger.error(f"API connection error getting instructions for '{source_name}': {e}")
return self._normalize_error(f"API connection error: {e}", "connection_error", "Ensure the data sources API is running and accessible.")
except Exception as e:
logger.exception(f"Unexpected error in get_source_instructions for '{source_name}': {e}")
return self._normalize_error(f"An unexpected error occurred while fetching instructions: {e}", "internal_error", "Check logs for detailed stack trace.")
def execute_sql_query(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
source_name: Optional[str] = None,
sql_query: str = "",
async_mode: bool = False,
max_rows: Optional[int] = None,
timeout_seconds: Optional[int] = None
) -> Dict[str, Any]:
"""
Executes a raw SQL query against a specific data source for the tenant.
Supports synchronous execution (default) and async job submission.
:param run_context: Agno RunContext containing session_state (auto-injected by framework).
:param tenant_id: Tenant identifier. Resolved from session_state when omitted.
:param source_name: Target data source. Resolved from session_state when omitted.
:param sql_query: Raw SQL string to execute.
:param async_mode: When True, enqueue query and return job metadata.
:param max_rows: Optional server-side row limit override.
:param timeout_seconds: Optional execution timeout override.
:return: Dictionary with execution payload or normalized error.
"""
# Get session_state from run_context (Agno v2 pattern)
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
tenant_id = tenant_context # type: ignore[assignment]
source_context = self._resolve_source(source_name, session_state)
if isinstance(source_context, dict) and source_context.get("error"):
return source_context
source_name = source_context # type: ignore[assignment]
# CRITICAL: Get per-request headers
request_headers = self._get_request_headers(session_state)
if not sql_query or not sql_query.strip():
return self._normalize_error("SQL query cannot be empty.", "missing_parameter", "Provide a non-empty sql_query string.")
logger.info(f"Tool: execute_sql_query called for tenant '{tenant_id}' on source '{source_name}'. Query: {sql_query[:100]}...")
# Basic client-side check for obviously disallowed operations (optional layer)
# The primary security relies on the DB user permissions configured in the backend.
disallowed_keywords = ['drop ', 'delete ', 'truncate ', 'alter ', 'insert ', 'update ', 'grant ', 'revoke ']
normalized_query = sql_query.strip().lower()
if any(keyword in normalized_query for keyword in disallowed_keywords):
# Allow specific commands like SELECT, SHOW, DESCRIBE, EXPLAIN etc.
allowed_starts = ('select', 'with', 'show', 'describe', 'explain')
if not normalized_query.startswith(allowed_starts):
logger.warning(f"Potentially unsafe SQL query blocked client-side for tenant {tenant_id}: {sql_query[:100]}...")
return self._normalize_error(
"Query blocked: Operation potentially modifies data or structure. Only read-only queries (SELECT, SHOW, etc.) are allowed.",
"forbidden_operation",
"Rewrite query using SELECT, SHOW, DESCRIBE, or EXPLAIN commands only."
)
try:
url = "/api/v1/data-sources/execute-raw-sql"
payload = {
"tenant_id": tenant_id,
"source_name": source_name,
"sql_query": sql_query # Send the raw SQL
}
if async_mode:
payload["async_mode"] = True
if max_rows is not None:
payload["max_rows"] = max_rows
if timeout_seconds is not None:
payload["timeout_seconds"] = timeout_seconds
logger.debug(f"Posting raw SQL to {url}")
response = self.client.post(
url,
json=payload,
timeout=DEFAULT_API_TIMEOUT, # Use longer timeout for query execution
headers=request_headers # Use per-request headers
)
# --- Handle API Response ---
if response.status_code == 200:
data = response.json()
status = data.get("status", "success")
if status.lower() == "success":
results = data.get("results", [])
# --- START: Large Data Handling (Pandas Summary) ---
if isinstance(results, str) and (results.startswith("minio://") or results.startswith("s3://")):
logger.info(f"Large result detected at {results}. Generating summary for LLM...")
try:
import pandas as pd
import io
from backend.core.minio.config import get_minio_config
from minio import Minio
mconf = get_minio_config()
minio_client = Minio(
endpoint=mconf['endpoint'],
access_key=mconf['access_key'],
secret_key=mconf['secret_key'],
secure=mconf.get('secure', False)
)
bucket, obj_name = results.replace("minio://", "").split("/", 1)
response_obj = minio_client.get_object(bucket, obj_name)
file_content = response_obj.read()
df = pd.read_json(io.BytesIO(file_content))
row_count = len(df)
preview_rows = 200
description = df.describe(include='all').to_markdown()
head_data = df.head(preview_rows).to_dict(orient='records')
if session_state is not None:
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_sql_execution"] = {
"tenant_id": tenant_id,
"source_name": source_name,
"row_count": row_count,
"async_mode": async_mode
}
metadata["last_sql_query"] = sql_query.strip()
return {
"status": "success",
"results": head_data,
"summary": description,
"message": (
f"Result too large ({row_count} rows). "
f"Returned first {preview_rows} rows and statistical summary. "
f"Full data stored at {results}."
),
"row_count": row_count,
"rows_limited": True
}
except Exception as fetch_err:
logger.error(f"Failed to fetch/summarize large result: {fetch_err}")
return {
"status": "success",
"results": [],
"message": f"Result stored at {results}, but could not be downloaded for summary. ({str(fetch_err)})",
"row_count": "Unknown",
"rows_limited": True
}
# --- END: Large Data Handling ---
row_count = len(results) if isinstance(results, list) else data.get("rows_returned", 0)
logger.info(f"SQL query executed successfully for tenant '{tenant_id}', returned {row_count} rows.")
if session_state is not None:
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_sql_execution"] = {
"tenant_id": tenant_id,
"source_name": source_name,
"row_count": row_count,
"async_mode": async_mode
}
metadata["last_sql_query"] = sql_query.strip()
return {
"status": status,
"results": results,
"rows_returned": data.get("rows_returned", row_count),
"rows_limited": data.get("rows_limited", False),
"execution_time_ms": data.get("execution_time_ms")
}
logger.info(f"SQL execution acknowledged with status '{status}' for tenant '{tenant_id}'")
if session_state is not None:
metadata = session_state.setdefault("analysis_metadata", {})
metadata["last_sql_execution"] = {
"tenant_id": tenant_id,
"source_name": source_name,
"status": status,
"async_mode": data.get("async_mode", async_mode),
"job_id": data.get("job_id")
}
metadata["last_sql_query"] = sql_query.strip()
return data
else:
# Attempt to get detailed error message from API response
error_detail = f"Request failed with status {response.status_code}."
try:
error_json = response.json()
# FastAPI often puts errors in {"detail": "..."}
if "detail" in error_json:
error_detail = error_json["detail"]
else: # Otherwise, convert whole JSON to string
error_detail = json.dumps(error_json)
except json.JSONDecodeError:
# If response is not JSON, use raw text
error_detail = response.text[:500] # Limit length
logger.error(f"SQL execution API failed for tenant '{tenant_id}' (Status {response.status_code}): {error_detail}")
# Provide a structured error back to the agent
return self._normalize_error(
f"SQL Execution Failed (Status {response.status_code}): {error_detail}",
"execution_error",
"Check query syntax and ensure the source is accessible for this tenant."
)
except httpx.TimeoutException:
logger.error(f"API timeout executing SQL for tenant '{tenant_id}' on source '{source_name}'.")
return self._normalize_error(
f"API request timed out after {DEFAULT_API_TIMEOUT} seconds. The query might be too long-running for synchronous execution.",
"timeout_error",
"Consider optimizing the query or using async job submission for long-running queries."
)
except httpx.RequestError as e:
logger.error(f"API connection error executing SQL for tenant '{tenant_id}': {e}")
return self._normalize_error(f"API connection error during SQL execution: {e}", "connection_error", "Ensure the data sources API is running and accessible.")
except Exception as e:
logger.exception(f"Unexpected error in execute_sql_query for tenant '{tenant_id}': {e}")
return {"error": f"An unexpected client-side error occurred during SQL execution: {e}"}
def compose_dataset_workflow(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
name: str = "dataset",
target_format: str = "parquet",
steps: Optional[List[Dict[str, Any]]] = None,
keep_intermediates: bool = False,
) -> Dict[str, Any]:
"""Compose a dataset from SQL/file workflow steps via tenant-files API."""
session_state = run_context.session_state if run_context else None
tenant_context = self._resolve_tenant(tenant_id, session_state)
if isinstance(tenant_context, dict) and tenant_context.get("error"):
return tenant_context
if not steps:
return self._normalize_error("steps must be provided", code="missing_parameter")
request_headers = self._get_request_headers(session_state)
payload = {
"name": name,
"target_format": target_format,
"steps": steps,
"keep_intermediates": keep_intermediates,
}
try:
response = self.client.post(
"/api/v1/tenant-files/datasets/compose",
json=payload,
timeout=DEFAULT_REQUEST_TIMEOUT,
headers=request_headers,
)
if response.status_code in (200, 201):
return response.json()
return self._normalize_error(
f"Failed to compose dataset: {response.text[:200]}",
code="API_ERROR",
)
except Exception as e:
logger.exception(f"Error in compose_dataset_workflow: {e}")
return self._normalize_error(f"Error: {str(e)}", code="UNEXPECTED_ERROR")
def save_query_to_tenant_csv(
self,
run_context: Optional[RunContext] = None,
tenant_id: Optional[str] = None,
source_name: Optional[str] = None,
sql_query: str = "",
dataset_name: str = "query_export"
) -> str:
"""
Executes a raw SQL query and directly saves the full results as a CSV in the
tenant's MinIO storage. This bridges Data Sources with ML capabilities!
Use this tool when you need to run Machine Learning, Data Profiling, or deep
pandas analysis on SQL data, because this returns a file path that can be
given to ML/pandas tools.
:param run_context: Agno RunContext (auto-injected).
:param tenant_id: Tenant ID (auto-resolved from session).
:param source_name: Source name (auto-resolved from session).
:param sql_query: The SQL query to run.
:param dataset_name: A descriptive name for the resulting CSV file.
:return: A success string containing the MinIO `minio://...` path.
"""
# Execute the query
logger.info(f"save_query_to_tenant_csv called for dataset '{dataset_name}'")
result = self.execute_sql_query(
run_context=run_context,
tenant_id=tenant_id,
source_name=source_name,
sql_query=sql_query,
async_mode=False
# Let it run without strict max_rows so we get the real dataset
)
if "error" in result:
return f"Failed to execute SQL: {result['error']}"
# Extract results
raw_data = result.get("results")
# If the result is already a minio string (because execute_sql_query intercepted a giant payload)
if isinstance(raw_data, str) and (raw_data.startswith("minio://") or raw_data.startswith("s3://")):
return f"Data successfully saved to object storage. Path to use for ML/Analysis tools: `{raw_data}`"
if not raw_data or not isinstance(raw_data, list):
return "Query returned no data or an invalid format. Cannot save to CSV."
try:
import pandas as pd
import io
import json
# Resolve tenant explicitly
session_state = run_context.session_state if run_context else None
resolved_tenant = self._resolve_tenant(tenant_id, session_state)
if isinstance(resolved_tenant, dict): return "Could not resolve tenant_id for saving file."
df = pd.DataFrame(raw_data)
csv_buffer = io.BytesIO()
df.to_csv(csv_buffer, index=False)
csv_buffer.seek(0)
timestamp = int(datetime.now().timestamp())
safe_name = dataset_name.replace(" ", "_").lower()
filename = f"{safe_name}_{timestamp}.csv"
request_headers = self._get_request_headers(session_state)
# Remove Content-Type so httpx sets the multipart boundary correctly
headers = {k: v for k, v in request_headers.items() if k.lower() != "content-type"}
files = {"file": (filename, csv_buffer, "text/csv")}
metadata = json.dumps({"source": "sql_agent", "label": "query_export"})
data = {"metadata": metadata}
response = self.client.post(
"/api/v1/tenant-files/assets",
files=files,
data=data,
headers=headers,
timeout=120.0
)
if response.status_code == 201:
resp_data = response.json()
asset_id = resp_data.get("asset", {}).get("asset_id", "")
final_path = resp_data.get("asset", {}).get("path", filename)
return f"Data ({len(df)} rows, {len(df.columns)} columns) successfully saved to tenant storage. Path/Asset ID to use for ML/Analysis tools: `{final_path}` or `{asset_id}`"
else:
error_text = response.text[:200]
return f"Query succeeded, but saving to tenant-files API failed (HTTP {response.status_code}): {error_text}"
except ImportError:
return "Failed to import pandas. Cannot save CSV locally."
except Exception as e:
logger.exception(f"Error saving query to CSV: {e}")
return f"Query succeeded, but saving to CSV failed: {str(e)}" |