362 lines
11 KiB
Python
362 lines
11 KiB
Python
"""Power BI Query Tool for LangGraph.
|
|
|
|
This tool provides DAX query capabilities for Power BI datasets
|
|
via the Power BI REST API. Only read-only queries are allowed.
|
|
"""
|
|
|
|
import logging
|
|
import asyncio
|
|
import os
|
|
import re
|
|
import functools
|
|
from typing import Annotated
|
|
|
|
import anyio
|
|
import httpx
|
|
from langchain_core.tools import tool
|
|
from msal import ConfidentialClientApplication, SerializableTokenCache
|
|
|
|
from modules.shared.configuration import APP_CONFIG
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Configuration constants - encapsulated in this file
|
|
POWERBI_DATASET_ID = APP_CONFIG.get("VALUEON_POWERBI_DATASET_ID", "")
|
|
POWERBI_CLIENT_ID = APP_CONFIG.get("VALUEON_POWERBI_CLIENT_ID", "")
|
|
POWERBI_CLIENT_SECRET = APP_CONFIG.get("VALUEON_POWERBI_CLIENT_SECRET", "")
|
|
POWERBI_TENANT_ID = APP_CONFIG.get("VALUEON_POWERBI_TENANT_ID", "")
|
|
POWERBI_BASE_URL = "https://api.powerbi.com/v1.0/myorg"
|
|
POWERBI_AUTHORITY_BASE = "https://login.microsoftonline.com"
|
|
POWERBI_SCOPE = ["https://analysis.windows.net/powerbi/api/.default"]
|
|
|
|
# Limit results to prevent excessive context usage
|
|
MAX_ROWS_LIMIT = 100
|
|
|
|
|
|
def _validate_environment() -> tuple[bool, str]:
|
|
"""Validate that all required environment variables are set.
|
|
|
|
Returns:
|
|
A tuple of (is_valid, error_message)
|
|
"""
|
|
missing = []
|
|
if not POWERBI_DATASET_ID:
|
|
missing.append("POWERBI_DATASET_ID")
|
|
if not POWERBI_CLIENT_ID:
|
|
missing.append("POWERBI_CLIENT_ID")
|
|
if not POWERBI_CLIENT_SECRET:
|
|
missing.append("POWERBI_CLIENT_SECRET")
|
|
if not POWERBI_TENANT_ID:
|
|
missing.append("POWERBI_TENANT_ID")
|
|
|
|
if missing:
|
|
return False, f"Missing required environment variables: {', '.join(missing)}"
|
|
|
|
return True, ""
|
|
|
|
|
|
def _validate_dax_query(*, dax_query: str) -> tuple[bool, str]:
|
|
"""Validate that the query is a valid DAX query.
|
|
|
|
Args:
|
|
dax_query: The DAX query to validate
|
|
|
|
Returns:
|
|
A tuple of (is_valid, error_message)
|
|
"""
|
|
# Remove leading/trailing whitespace
|
|
normalized_query = dax_query.strip()
|
|
|
|
if not normalized_query:
|
|
return False, "Query cannot be empty"
|
|
|
|
# DAX queries typically start with EVALUATE, DEFINE, or are table expressions
|
|
# We'll be lenient and just check it's not trying to do something dangerous
|
|
# DAX is read-only by nature, but we validate structure
|
|
|
|
# Check for minimum length
|
|
if len(normalized_query) < 5:
|
|
return False, "Query is too short to be valid"
|
|
|
|
return True, ""
|
|
|
|
|
|
def _get_access_token_sync(
|
|
*,
|
|
tenant_id: str,
|
|
client_id: str,
|
|
client_secret: str,
|
|
authority_base: str = POWERBI_AUTHORITY_BASE,
|
|
cache: SerializableTokenCache | None = None,
|
|
) -> str:
|
|
"""Get Power BI access token using MSAL (synchronous).
|
|
|
|
Args:
|
|
tenant_id: Azure AD tenant ID
|
|
client_id: Application client ID
|
|
client_secret: Application client secret
|
|
authority_base: Azure AD authority base URL
|
|
cache: Optional token cache for reuse
|
|
|
|
Returns:
|
|
Access token string
|
|
|
|
Raises:
|
|
RuntimeError: If token acquisition fails
|
|
"""
|
|
authority = f"{authority_base}/{tenant_id}"
|
|
|
|
app = ConfidentialClientApplication(
|
|
client_id=client_id,
|
|
authority=authority,
|
|
client_credential=client_secret,
|
|
token_cache=cache,
|
|
)
|
|
|
|
# Try cache first; fall back to client credentials
|
|
result = app.acquire_token_silent(
|
|
POWERBI_SCOPE, account=None
|
|
) or app.acquire_token_for_client(scopes=POWERBI_SCOPE)
|
|
|
|
if "access_token" not in result:
|
|
raise RuntimeError(
|
|
f"MSAL token error: {result.get('error')} - {result.get('error_description')}"
|
|
)
|
|
|
|
return result["access_token"]
|
|
|
|
|
|
async def _get_access_token_async(
|
|
*,
|
|
tenant_id: str,
|
|
client_id: str,
|
|
client_secret: str,
|
|
**kwargs,
|
|
) -> str:
|
|
"""Get Power BI access token using MSAL (asynchronous).
|
|
|
|
Args:
|
|
tenant_id: Azure AD tenant ID
|
|
client_id: Application client ID
|
|
client_secret: Application client secret
|
|
**kwargs: Additional arguments for _get_access_token_sync
|
|
|
|
Returns:
|
|
Access token string
|
|
"""
|
|
# Create a partial function with arguments pre-filled
|
|
func = functools.partial(
|
|
_get_access_token_sync,
|
|
tenant_id=tenant_id,
|
|
client_id=client_id,
|
|
client_secret=client_secret,
|
|
**kwargs,
|
|
)
|
|
# Offload the blocking MSAL HTTP call to a worker thread
|
|
return await anyio.to_thread.run_sync(func)
|
|
|
|
|
|
async def _execute_dax_query(
|
|
*, dax_query: str, dataset_id: str, access_token: str
|
|
) -> dict:
|
|
"""Execute a DAX query against Power BI dataset.
|
|
|
|
Args:
|
|
dax_query: The DAX query to execute
|
|
dataset_id: Power BI dataset ID
|
|
access_token: Access token for authentication
|
|
|
|
Returns:
|
|
Dictionary containing query results
|
|
|
|
Raises:
|
|
RuntimeError: If query execution fails
|
|
"""
|
|
url = f"{POWERBI_BASE_URL}/datasets/{dataset_id}/executeQueries"
|
|
|
|
body = {
|
|
"queries": [{"query": dax_query}],
|
|
"serializerSettings": {"includeNulls": True},
|
|
}
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
resp = await client.post(url, headers=headers, json=body)
|
|
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(
|
|
f"Power BI executeQueries failed: {resp.status_code} - {resp.text}"
|
|
)
|
|
|
|
payload = resp.json()
|
|
|
|
try:
|
|
rows = payload["results"][0]["tables"][0]["rows"]
|
|
except (KeyError, IndexError) as e:
|
|
raise RuntimeError("Unexpected executeQueries response structure") from e
|
|
|
|
# Extract column names from the first row if available
|
|
if rows:
|
|
columns = list(rows[0].keys())
|
|
else:
|
|
columns = []
|
|
|
|
return {"columns": columns, "rows": rows}
|
|
|
|
|
|
def _strip_table_qualifier(*, column_name: str) -> str:
|
|
"""Strip table qualifier from column name.
|
|
|
|
Power BI often returns columns as 'Table[Column]'. This strips to 'Column'.
|
|
|
|
Args:
|
|
column_name: The column name to process
|
|
|
|
Returns:
|
|
Processed column name
|
|
"""
|
|
if "[" in column_name and column_name.endswith("]"):
|
|
return column_name.split("[", 1)[1][:-1]
|
|
return column_name
|
|
|
|
|
|
def _format_results(*, columns: list[str], rows: list[dict], max_rows: int) -> str:
|
|
"""Format query results into a readable string.
|
|
|
|
Args:
|
|
columns: List of column names
|
|
rows: List of row data (as dictionaries)
|
|
max_rows: Maximum number of rows to display
|
|
|
|
Returns:
|
|
Formatted string representation of the results
|
|
"""
|
|
total_rows = len(rows)
|
|
|
|
if total_rows == 0:
|
|
return "Query executed successfully but returned no results."
|
|
|
|
# Strip table qualifiers from column names
|
|
clean_columns = [_strip_table_qualifier(column_name=col) for col in columns]
|
|
|
|
# Limit rows to max_rows
|
|
display_rows = rows[:max_rows]
|
|
truncated = total_rows > max_rows
|
|
|
|
# Calculate column widths
|
|
col_widths = [len(str(col)) for col in clean_columns]
|
|
for row in display_rows:
|
|
for i, col in enumerate(columns):
|
|
value = row.get(col, "")
|
|
col_widths[i] = max(col_widths[i], len(str(value)))
|
|
|
|
# Build header
|
|
header_parts = []
|
|
for col, width in zip(clean_columns, col_widths):
|
|
header_parts.append(str(col).ljust(width))
|
|
header = " | ".join(header_parts)
|
|
separator = "-" * len(header)
|
|
|
|
# Build rows
|
|
row_lines = []
|
|
for row in display_rows:
|
|
row_parts = []
|
|
for col, width in zip(columns, col_widths):
|
|
value = row.get(col, "")
|
|
row_parts.append(str(value).ljust(width))
|
|
row_lines.append(" | ".join(row_parts))
|
|
|
|
# Combine all parts
|
|
result_parts = [
|
|
f"Query returned {total_rows} row(s):",
|
|
]
|
|
|
|
if truncated:
|
|
result_parts.append(
|
|
f"(Results limited to {max_rows} rows for context efficiency)\n"
|
|
)
|
|
else:
|
|
result_parts.append("")
|
|
|
|
result_parts.extend([header, separator, "\n".join(row_lines)])
|
|
|
|
return "\n".join(result_parts)
|
|
|
|
|
|
@tool
|
|
async def query_powerbi_data(
|
|
dax_query: Annotated[str, "The DAX query to execute against the Power BI dataset"],
|
|
) -> str:
|
|
"""Execute a DAX query against the Power BI dataset to access warehouse inventory data.
|
|
|
|
This tool provides access to a Power BI table called 'data_full' which contains
|
|
articles available in the warehouse of the user. Use DAX (Data Analysis Expressions)
|
|
queries to retrieve and analyze this inventory data.
|
|
|
|
Available table:
|
|
- 'data_full': Contains warehouse inventory articles and their details
|
|
|
|
Common query patterns:
|
|
- View all data: EVALUATE 'data_full'
|
|
- With filter: EVALUATE FILTER('data_full', [Column] = "Value")
|
|
- Top N rows: EVALUATE TOPN(10, 'data_full', [Column], DESC)
|
|
- Calculated: EVALUATE SUMMARIZE('data_full', [Column1], "Total", SUM([Column2]))
|
|
|
|
Results are limited to 100 rows maximum for efficiency.
|
|
|
|
Args:
|
|
dax_query: The DAX query to execute (e.g., "EVALUATE 'data_full'")
|
|
|
|
Returns:
|
|
A formatted string containing the query results with columns and rows
|
|
"""
|
|
try:
|
|
# Validate environment configuration
|
|
is_valid_env, error_msg = _validate_environment()
|
|
if not is_valid_env:
|
|
logger.error(f"Environment validation failed: {error_msg}")
|
|
return f"Configuration Error: {error_msg}"
|
|
|
|
# Validate the query
|
|
is_valid_query, error_msg = _validate_dax_query(dax_query=dax_query)
|
|
if not is_valid_query:
|
|
logger.warning(f"Invalid query attempt: {dax_query[:100]}...")
|
|
return f"Query Validation Error: {error_msg}"
|
|
|
|
logger.info(f"Executing Power BI query: {dax_query[:100]}...")
|
|
|
|
# Get access token
|
|
access_token = await _get_access_token_async(
|
|
tenant_id=POWERBI_TENANT_ID,
|
|
client_id=POWERBI_CLIENT_ID,
|
|
client_secret=POWERBI_CLIENT_SECRET,
|
|
)
|
|
|
|
# Execute the query
|
|
result = await _execute_dax_query(
|
|
dax_query=dax_query,
|
|
dataset_id=POWERBI_DATASET_ID,
|
|
access_token=access_token,
|
|
)
|
|
|
|
# Format and return results
|
|
formatted_output = _format_results(
|
|
columns=result["columns"], rows=result["rows"], max_rows=MAX_ROWS_LIMIT
|
|
)
|
|
|
|
logger.info(
|
|
f"Query completed successfully, returned {len(result['rows'])} row(s)"
|
|
)
|
|
return formatted_output
|
|
|
|
except RuntimeError as e:
|
|
logger.error(f"Runtime error in query_powerbi_data tool: {str(e)}")
|
|
return f"Error executing query: {str(e)}"
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error in query_powerbi_data tool: {str(e)}")
|
|
return f"Unexpected error: {str(e)}"
|