busya
7f7a85b2c5
- Extract _optimize_images() helper to eliminate duplicate
optimization code (was duplicated between initial pass and
quality-reduction fallback)
- Move quality-reduction estimate to only apply when NOT chunking
(each chunk fits the image limit by definition)
- Fix _merge_chunk_results to return 'chunk_count' instead of 'chunks'
for consistency with plugin.py
- Simplify plugin.py: analysis.get('chunk_count', 1) instead of
fallback chain
- Document that Kilo API gateway is incompatible with AsyncOpenAI
image format (probe returns 0)
1671 lines
79 KiB
Python
1671 lines
79 KiB
Python
# #region LLMAnalysisService [C:5] [TYPE Module] [SEMANTICS llm, screenshot, playwright, openai, tenacity]
|
||
# @BRIEF Services for LLM interaction and dashboard screenshots.
|
||
# @LAYER Plugin
|
||
# @RELATION DEPENDS_ON -> [EXT:Library:tenacity]
|
||
# @INVARIANT Screenshots must be 1920px width and capture full page height.
|
||
# @DATA_CONTRACT DashboardSpec -> Screenshot + Analysis
|
||
# @RATIONALE Extracted all hardcoded timeouts into named module-level constants (PLAYWRIGHT_NAVIGATION_TIMEOUT_MS, PLAYWRIGHT_WAIT_TIMEOUT_MS, PLAYWRIGHT_SHORT_TIMEOUT_MS, HTTP_REQUEST_TIMEOUT_MS, SCREENSHOT_SERVICE_TIMEOUT_MS, LLM_HTTP_TIMEOUT_S) and DEFAULT_USER_AGENT. Zero remaining numeric timeout literals.
|
||
|
||
import asyncio
|
||
import base64
|
||
import io
|
||
import json
|
||
import os
|
||
import re
|
||
import ssl
|
||
from typing import Any
|
||
from urllib.parse import urlsplit
|
||
|
||
import httpx
|
||
from openai import AsyncOpenAI, AuthenticationError as OpenAIAuthenticationError, RateLimitError
|
||
from PIL import Image
|
||
from playwright.async_api import async_playwright
|
||
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
|
||
|
||
from ...core.config_models import Environment
|
||
from ...core.logger import belief_scope, logger
|
||
from ...services.llm_prompt_templates import DEFAULT_LLM_PROMPTS, render_prompt
|
||
from .models import LLMProviderType
|
||
|
||
# Timeout constants (milliseconds unless noted)
|
||
PLAYWRIGHT_NAVIGATION_TIMEOUT_MS = 30000
|
||
PLAYWRIGHT_WAIT_TIMEOUT_MS = 10000
|
||
PLAYWRIGHT_SHORT_TIMEOUT_MS = 5000
|
||
HTTP_REQUEST_TIMEOUT_MS = 60000
|
||
SCREENSHOT_SERVICE_TIMEOUT_MS = 120000
|
||
LLM_HTTP_TIMEOUT_S = 120 # seconds (httpx client timeout)
|
||
DEFAULT_USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
|
||
|
||
# #region ScreenshotService [TYPE Class]
|
||
# @BRIEF Handles capturing screenshots of Superset dashboards.
|
||
class ScreenshotService:
|
||
# region ScreenshotService.__init__ [TYPE Function]
|
||
# @PURPOSE Initializes the ScreenshotService with environment configuration.
|
||
# @PRE env is a valid Environment object.
|
||
def __init__(self, env: Environment):
|
||
self.env = env
|
||
# endregion ScreenshotService.__init__
|
||
|
||
# region ScreenshotService._find_first_visible_locator [TYPE Function]
|
||
# @PURPOSE Resolve the first visible locator from multiple Playwright locator strategies.
|
||
# @PRE candidates is a non-empty list of locator-like objects.
|
||
# @POST Returns a locator ready for interaction or None when nothing matches.
|
||
async def _find_first_visible_locator(self, candidates) -> Any:
|
||
for locator in candidates:
|
||
try:
|
||
match_count = await locator.count()
|
||
for index in range(match_count):
|
||
candidate = locator.nth(index)
|
||
if await candidate.is_visible():
|
||
return candidate
|
||
except Exception:
|
||
continue
|
||
return None
|
||
# endregion ScreenshotService._find_first_visible_locator
|
||
|
||
# region ScreenshotService._iter_login_roots [TYPE Function]
|
||
# @PURPOSE Enumerate page and child frames where login controls may be rendered.
|
||
# @PRE page is a Playwright page-like object.
|
||
# @POST Returns ordered roots starting with main page followed by frames.
|
||
def _iter_login_roots(self, page) -> list[Any]:
|
||
roots = [page]
|
||
page_frames = getattr(page, "frames", [])
|
||
try:
|
||
for frame in page_frames:
|
||
if frame not in roots:
|
||
roots.append(frame)
|
||
except Exception:
|
||
pass
|
||
return roots
|
||
# endregion ScreenshotService._iter_login_roots
|
||
|
||
# region ScreenshotService._extract_hidden_login_fields [TYPE Function]
|
||
# @PURPOSE Collect hidden form fields required for direct login POST fallback.
|
||
# @PRE Login page is loaded.
|
||
# @POST Returns hidden input name/value mapping aggregated from page and child frames.
|
||
async def _extract_hidden_login_fields(self, page) -> dict[str, str]:
|
||
hidden_fields: dict[str, str] = {}
|
||
for root in self._iter_login_roots(page):
|
||
try:
|
||
locator = root.locator("input[type='hidden'][name]")
|
||
count = await locator.count()
|
||
for index in range(count):
|
||
candidate = locator.nth(index)
|
||
field_name = str(await candidate.get_attribute("name") or "").strip()
|
||
if not field_name or field_name in hidden_fields:
|
||
continue
|
||
hidden_fields[field_name] = str(await candidate.input_value()).strip()
|
||
except Exception:
|
||
continue
|
||
return hidden_fields
|
||
# endregion ScreenshotService._extract_hidden_login_fields
|
||
|
||
# region ScreenshotService._extract_csrf_token [TYPE Function]
|
||
# @PURPOSE Resolve CSRF token value from main page or embedded login frame.
|
||
# @PRE Login page is loaded.
|
||
# @POST Returns first non-empty csrf token or empty string.
|
||
async def _extract_csrf_token(self, page) -> str:
|
||
hidden_fields = await self._extract_hidden_login_fields(page)
|
||
return str(hidden_fields.get("csrf_token") or "").strip()
|
||
# endregion ScreenshotService._extract_csrf_token
|
||
|
||
# region ScreenshotService._response_looks_like_login_page [TYPE Function]
|
||
# @PURPOSE Detect when fallback login POST returned the login form again instead of an authenticated page.
|
||
# @PRE response_text is normalized HTML or text from login POST response.
|
||
# @POST Returns True when login-page markers dominate the response body.
|
||
def _response_looks_like_login_page(self, response_text: str) -> bool:
|
||
normalized = str(response_text or "").strip().lower()
|
||
if not normalized:
|
||
return False
|
||
|
||
markers = [
|
||
"enter your login and password below",
|
||
"username:",
|
||
"password:",
|
||
"sign in",
|
||
'name="csrf_token"',
|
||
]
|
||
return sum(marker in normalized for marker in markers) >= 3
|
||
# endregion ScreenshotService._response_looks_like_login_page
|
||
|
||
# region ScreenshotService._redirect_looks_authenticated [TYPE Function]
|
||
# @PURPOSE Treat non-login redirects after form POST as successful authentication without waiting for redirect target.
|
||
# @PRE redirect_location may be empty or relative.
|
||
# @POST Returns True when redirect target does not point back to login flow.
|
||
def _redirect_looks_authenticated(self, redirect_location: str) -> bool:
|
||
normalized = str(redirect_location or "").strip().lower()
|
||
if not normalized:
|
||
return True
|
||
return "/login" not in normalized
|
||
# endregion ScreenshotService._redirect_looks_authenticated
|
||
|
||
# region ScreenshotService._submit_login_via_form_post [TYPE Function]
|
||
# @PURPOSE Fallback login path that submits credentials directly with csrf token.
|
||
# @PRE login_url is same-origin and csrf token can be read from DOM.
|
||
# @POST Browser context receives authenticated cookies when login succeeds.
|
||
async def _submit_login_via_form_post(self, page, login_url: str) -> bool:
|
||
hidden_fields = await self._extract_hidden_login_fields(page)
|
||
csrf_token = str(hidden_fields.get("csrf_token") or "").strip()
|
||
if not csrf_token:
|
||
logger.warning("[DEBUG] Direct form login fallback skipped: csrf_token not found")
|
||
return False
|
||
|
||
try:
|
||
request_context = page.context.request
|
||
except Exception as context_error:
|
||
logger.warning(f"[DEBUG] Direct form login fallback skipped: request context unavailable: {context_error}")
|
||
return False
|
||
|
||
parsed_url = urlsplit(login_url)
|
||
origin = f"{parsed_url.scheme}://{parsed_url.netloc}" if parsed_url.scheme and parsed_url.netloc else login_url
|
||
payload = dict(hidden_fields)
|
||
payload["username"] = self.env.username
|
||
payload["password"] = self.env.password
|
||
logger.info(
|
||
f"[DEBUG] Attempting direct form login fallback via browser context request with hidden fields: {sorted(hidden_fields.keys())}"
|
||
)
|
||
|
||
response = await request_context.post(
|
||
login_url,
|
||
form=payload,
|
||
headers={
|
||
"Origin": origin,
|
||
"Referer": login_url,
|
||
},
|
||
timeout=PLAYWRIGHT_WAIT_TIMEOUT_MS,
|
||
fail_on_status_code=False,
|
||
max_redirects=0,
|
||
)
|
||
response_url = str(getattr(response, "url", "") or "")
|
||
response_status = int(getattr(response, "status", 0) or 0)
|
||
response_headers = dict(getattr(response, "headers", {}) or {})
|
||
redirect_location = str(
|
||
response_headers.get("location")
|
||
or response_headers.get("Location")
|
||
or ""
|
||
).strip()
|
||
redirect_statuses = {301, 302, 303, 307, 308}
|
||
if response_status in redirect_statuses:
|
||
redirect_authenticated = self._redirect_looks_authenticated(redirect_location)
|
||
logger.info(
|
||
f"[DEBUG] Direct form login fallback redirect response: status={response_status} url={response_url} location={redirect_location!r} authenticated={redirect_authenticated}"
|
||
)
|
||
return redirect_authenticated
|
||
|
||
response_text = await response.text()
|
||
text_snippet = " ".join(response_text.split())[:200]
|
||
looks_like_login_page = self._response_looks_like_login_page(response_text)
|
||
logger.info(
|
||
f"[DEBUG] Direct form login fallback response: status={response_status} url={response_url} login_markup={looks_like_login_page} snippet={text_snippet!r}"
|
||
)
|
||
return not looks_like_login_page
|
||
# endregion ScreenshotService._submit_login_via_form_post
|
||
|
||
# region ScreenshotService._find_login_field_locator [TYPE Function]
|
||
# @PURPOSE Resolve login form input using semantic label text plus generic visible-input fallbacks.
|
||
# @PRE field_name is `username` or `password`.
|
||
# @POST Returns a locator for the corresponding input or None.
|
||
async def _find_login_field_locator(self, page, field_name: str) -> Any:
|
||
normalized = str(field_name or "").strip().lower()
|
||
for root in self._iter_login_roots(page):
|
||
if normalized == "username":
|
||
input_candidates = [
|
||
root.get_by_label("Username", exact=False),
|
||
root.get_by_label("Login", exact=False),
|
||
root.locator("label:text-matches('Username|Login', 'i')").locator("xpath=following::input[1]"),
|
||
root.locator("text=/Username|Login/i").locator("xpath=following::input[1]"),
|
||
root.locator("input[name='username']"),
|
||
root.locator("input#username"),
|
||
root.locator("input[placeholder*='Username']"),
|
||
root.locator("input[type='text']"),
|
||
root.locator("input:not([type='password'])"),
|
||
]
|
||
locator = await self._find_first_visible_locator(input_candidates)
|
||
if locator:
|
||
return locator
|
||
|
||
if normalized == "password":
|
||
input_candidates = [
|
||
root.get_by_label("Password", exact=False),
|
||
root.locator("label:text-matches('Password', 'i')").locator("xpath=following::input[1]"),
|
||
root.locator("text=/Password/i").locator("xpath=following::input[1]"),
|
||
root.locator("input[name='password']"),
|
||
root.locator("input#password"),
|
||
root.locator("input[placeholder*='Password']"),
|
||
root.locator("input[type='password']"),
|
||
]
|
||
locator = await self._find_first_visible_locator(input_candidates)
|
||
if locator:
|
||
return locator
|
||
|
||
return None
|
||
# endregion ScreenshotService._find_login_field_locator
|
||
|
||
# region ScreenshotService._find_submit_locator [TYPE Function]
|
||
# @PURPOSE Resolve login submit button from main page or embedded auth frame.
|
||
# @PRE page is ready for login interaction.
|
||
# @POST Returns visible submit locator or None.
|
||
async def _find_submit_locator(self, page) -> Any:
|
||
selectors = [
|
||
lambda root: root.get_by_role("button", name="Sign in", exact=False),
|
||
lambda root: root.get_by_role("button", name="Login", exact=False),
|
||
lambda root: root.locator("button[type='submit']"),
|
||
lambda root: root.locator("button#submit"),
|
||
lambda root: root.locator(".btn-primary"),
|
||
lambda root: root.locator("input[type='submit']"),
|
||
]
|
||
for root in self._iter_login_roots(page):
|
||
locator = await self._find_first_visible_locator([factory(root) for factory in selectors])
|
||
if locator:
|
||
return locator
|
||
return None
|
||
# endregion ScreenshotService._find_submit_locator
|
||
|
||
# region ScreenshotService._goto_resilient [TYPE Function]
|
||
# @PURPOSE Navigate without relying on networkidle for pages with long-polling or persistent requests.
|
||
# @PRE page is a valid Playwright page and url is non-empty.
|
||
# @POST Returns last navigation response or raises when both primary and fallback waits fail.
|
||
async def _goto_resilient(
|
||
self,
|
||
page,
|
||
url: str,
|
||
primary_wait_until: str = "domcontentloaded",
|
||
fallback_wait_until: str = "load",
|
||
timeout: int = HTTP_REQUEST_TIMEOUT_MS,
|
||
):
|
||
try:
|
||
return await page.goto(url, wait_until=primary_wait_until, timeout=timeout)
|
||
except Exception as primary_error:
|
||
logger.warning(
|
||
f"[ScreenshotService._goto_resilient] Primary navigation wait '{primary_wait_until}' failed for {url}: {primary_error}"
|
||
)
|
||
return await page.goto(url, wait_until=fallback_wait_until, timeout=timeout)
|
||
# endregion ScreenshotService._goto_resilient
|
||
|
||
# region ScreenshotService._wait_for_charts_stabilized [TYPE Function] [C:2]
|
||
# @BRIEF Wait until chart elements have non-zero dimensions, with polling.
|
||
# @PRE page is a valid Playwright page.
|
||
# @POST Waits for chart stabilization or raises on timeout (handled internally).
|
||
# @RATIONALE Polls for actual chart rendering dimensions rather than using a fixed delay — charts may load at different speeds depending on dashboard complexity and network conditions.
|
||
# @REJECTED Fixed sleep-based wait rejected — would either waste time (too long) or produce blank screenshots (too short); polling for actual canvas/svg dimensions is more reliable.
|
||
async def _wait_for_charts_stabilized(self, page, timeout_ms: int = 15000):
|
||
"""Wait until chart elements have non-zero dimensions, with polling."""
|
||
# Short initial delay for rendering pipeline to start
|
||
await asyncio.sleep(0.5)
|
||
try:
|
||
await page.wait_for_function("""() => {
|
||
const charts = document.querySelectorAll('.chart-container canvas, .slice_container svg, .grid-content canvas');
|
||
if (charts.length === 0) return true;
|
||
return Array.from(charts).some(c => {
|
||
if (c.tagName === 'CANVAS') return c.width > 10 && c.height > 10;
|
||
if (c.tagName === 'svg') {
|
||
const bbox = c.getBoundingClientRect();
|
||
return bbox.width > 10 && bbox.height > 10;
|
||
}
|
||
return false;
|
||
});
|
||
}""", timeout=timeout_ms)
|
||
except Exception:
|
||
logger.warning("[ScreenshotService] Chart stabilization wait timed out, proceeding anyway")
|
||
# endregion ScreenshotService._wait_for_charts_stabilized
|
||
|
||
# region ScreenshotService._wait_for_resize_rendered [TYPE Function] [C:2]
|
||
# @BRIEF Wait for charts to re-render after viewport resize.
|
||
# @PRE page is a valid Playwright page; chart_count_before contains pre-resize element counts.
|
||
# @POST Waits for chart content to return or timeout.
|
||
# @RATIONALE After viewport resize, Superset triggers lazy chart re-rendering — this function polls for chart elements to reappear before taking the screenshot.
|
||
# @REJECTED Single fixed wait after resize rejected — some dashboards re-render instantly while others take seconds; fixed wait is brittle across dashboard types.
|
||
async def _wait_for_resize_rendered(self, page, chart_count_before: dict, timeout_ms: int = 10000):
|
||
"""Wait for charts to re-render after viewport resize, with polling."""
|
||
try:
|
||
await page.wait_for_function("""(preCounts) => {
|
||
const currentCharts = document.querySelectorAll('.chart-container, .slice_container').length;
|
||
const currentCanvases = document.querySelectorAll('canvas').length;
|
||
const currentSvgs = document.querySelectorAll('.chart-container svg, .slice_container svg').length;
|
||
// At least one chart element must be present
|
||
return currentCharts > 0 && (currentCanvases > 0 || currentSvgs > 0);
|
||
}""", arg=chart_count_before, timeout=timeout_ms)
|
||
except Exception:
|
||
logger.warning("[ScreenshotService] Re-render wait timed out after viewport resize, proceeding anyway")
|
||
# endregion ScreenshotService._wait_for_resize_rendered
|
||
|
||
# region ScreenshotService._save_debug_screenshot [TYPE Function] [C:1]
|
||
# @BRIEF Save a debug screenshot to a temp directory for diagnostic purposes.
|
||
# @PRE debug_dir exists and is writable.
|
||
# @POST Returns the debug path or None on failure.
|
||
# @RATIONALE Uses tempfile.mkdtemp() to avoid accumulating .png files in production storage. Temp dir is cleaned up on success (rmtree) or preserved on failure for debugging.
|
||
# @REJECTED Old approach saved debug .png next to output path — accumulated _debug_failed_login.png, _preresize.png permanently in screenshots dir (F1). In-memory-only debug logging rejected — screenshot state is visual and cannot be captured in logs.
|
||
async def _save_debug_screenshot(self, page, debug_dir: str, suffix: str) -> str | None:
|
||
debug_path = os.path.join(debug_dir, suffix)
|
||
try:
|
||
await page.screenshot(path=debug_path)
|
||
return debug_path
|
||
except Exception:
|
||
return None
|
||
# endregion ScreenshotService._save_debug_screenshot
|
||
|
||
# region ScreenshotService._launch_and_login [TYPE Function] [C:4]
|
||
# @PURPOSE Launch browser, log in to Superset, navigate to dashboard URL.
|
||
# @PRE dashboard_id is valid, playwright instance is provided.
|
||
# @POST Returns (browser, context, page) tuple with active session.
|
||
# @SIDE_EFFECT Launches headless Chromium, performs UI login, navigates to dashboard.
|
||
# @RATIONALE Extracted from capture_dashboard to share login/navigation logic
|
||
# between capture_dashboard (backward compat) and capture_dashboard_chunks (multi-tab).
|
||
# @REJECTED Duplicating login logic rejected — any change to auth flow would require
|
||
# updating two code paths, leading to accidental drift.
|
||
async def _launch_and_login(
|
||
self,
|
||
playwright,
|
||
dashboard_id: str,
|
||
parsed_context: dict | None = None,
|
||
) -> tuple[Any, Any, Any]:
|
||
"""Launch browser, log in to Superset, and navigate to the dashboard.
|
||
|
||
Returns (browser, context, page).
|
||
"""
|
||
user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
|
||
base_ui_url = self.env.url.rstrip("/")
|
||
if base_ui_url.endswith("/api/v1"):
|
||
base_ui_url = base_ui_url[: -len("/api/v1")]
|
||
|
||
browser = await playwright.chromium.launch(
|
||
headless=True,
|
||
args=[
|
||
"--disable-blink-features=AutomationControlled",
|
||
"--disable-infobars",
|
||
"--no-sandbox",
|
||
],
|
||
)
|
||
context = await browser.new_context(
|
||
viewport={"width": 1280, "height": 720},
|
||
user_agent=user_agent,
|
||
extra_http_headers={
|
||
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
|
||
"Accept-Language": "ru-RU,ru;q=0.9,en-US;q=0.8,en;q=0.7",
|
||
"Upgrade-Insecure-Requests": "1",
|
||
"Sec-Fetch-Dest": "document",
|
||
"Sec-Fetch-Mode": "navigate",
|
||
"Sec-Fetch-Site": "none",
|
||
"Sec-Fetch-User": "?1",
|
||
},
|
||
)
|
||
page = await context.new_page()
|
||
await page.add_init_script("delete Object.getPrototypeOf(navigator).webdriver")
|
||
|
||
# 1. Navigate to login page and authenticate
|
||
login_url = f"{base_ui_url.rstrip('/')}/login/"
|
||
await self._goto_resilient(
|
||
page,
|
||
login_url,
|
||
primary_wait_until="domcontentloaded",
|
||
fallback_wait_until="load",
|
||
timeout=HTTP_REQUEST_TIMEOUT_MS,
|
||
)
|
||
await page.wait_for_load_state("domcontentloaded")
|
||
|
||
try:
|
||
used_direct_form_login = False
|
||
username_locator = await self._find_login_field_locator(page, "username")
|
||
|
||
if not username_locator:
|
||
used_direct_form_login = await self._submit_login_via_form_post(page, login_url)
|
||
if not used_direct_form_login:
|
||
raise RuntimeError("Could not find username input field on login page")
|
||
|
||
if username_locator is not None:
|
||
await username_locator.fill(self.env.username)
|
||
|
||
password_locator = (
|
||
await self._find_login_field_locator(page, "password")
|
||
if username_locator is not None
|
||
else None
|
||
)
|
||
if username_locator is not None and not password_locator:
|
||
raise RuntimeError("Could not find password input field on login page")
|
||
if password_locator is not None:
|
||
await password_locator.fill(self.env.password)
|
||
|
||
submit_locator = (
|
||
await self._find_submit_locator(page)
|
||
if username_locator is not None
|
||
else None
|
||
)
|
||
if username_locator is not None and not submit_locator:
|
||
raise RuntimeError("Could not find submit button on login page")
|
||
if submit_locator is not None:
|
||
await submit_locator.click()
|
||
|
||
if not used_direct_form_login:
|
||
try:
|
||
await page.wait_for_load_state("load", timeout=PLAYWRIGHT_NAVIGATION_TIMEOUT_MS)
|
||
except Exception:
|
||
pass
|
||
|
||
if not used_direct_form_login and "/login" in page.url:
|
||
error_msg = (
|
||
await page.locator(".alert-danger, .error-message").text_content()
|
||
if await page.locator(".alert-danger, .error-message").count() > 0
|
||
else "Unknown error"
|
||
)
|
||
raise RuntimeError(f"Login failed: {error_msg}")
|
||
except Exception as e:
|
||
raise RuntimeError(f"Login failed: {e!s}")
|
||
|
||
# 2. Navigate to dashboard
|
||
dashboard_url = f"{base_ui_url.rstrip('/')}/superset/dashboard/{dashboard_id}/?standalone=true"
|
||
if base_ui_url.startswith("https://") and dashboard_url.startswith("http://"):
|
||
dashboard_url = dashboard_url.replace("http://", "https://")
|
||
|
||
if parsed_context:
|
||
native_filters = parsed_context.get("native_filters")
|
||
active_tabs = parsed_context.get("activeTabs")
|
||
if native_filters:
|
||
dashboard_url += f"&native_filters={native_filters}"
|
||
if active_tabs:
|
||
dashboard_url += f"&activeTabs={active_tabs}"
|
||
|
||
await self._goto_resilient(
|
||
page,
|
||
dashboard_url,
|
||
primary_wait_until="domcontentloaded",
|
||
fallback_wait_until="load",
|
||
timeout=HTTP_REQUEST_TIMEOUT_MS,
|
||
)
|
||
|
||
if "/login" in page.url:
|
||
raise RuntimeError("Dashboard navigation redirected to login page after authentication")
|
||
|
||
# 3. Wait for dashboard content
|
||
try:
|
||
await page.wait_for_selector(
|
||
'.dashboard-component, .dashboard-header, [data-test="dashboard-grid"]',
|
||
timeout=PLAYWRIGHT_NAVIGATION_TIMEOUT_MS,
|
||
)
|
||
try:
|
||
await page.wait_for_selector(
|
||
".loading, .ant-skeleton, .spinner",
|
||
state="hidden",
|
||
timeout=HTTP_REQUEST_TIMEOUT_MS,
|
||
)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await page.wait_for_selector(
|
||
".chart-container canvas, .slice_container svg, .superset-chart-canvas, .grid-content .chart-container",
|
||
timeout=HTTP_REQUEST_TIMEOUT_MS,
|
||
)
|
||
except Exception:
|
||
pass
|
||
await page.wait_for_function(
|
||
"""() => {
|
||
const charts = document.querySelectorAll('.chart-container, .slice_container');
|
||
if (charts.length === 0) return true;
|
||
return Array.from(charts).every(chart => {
|
||
const hasCanvas = chart.querySelector('canvas') !== null;
|
||
const hasSvg = chart.querySelector('svg') !== null;
|
||
const hasContent = chart.innerText.trim().length > 0 || chart.children.length > 0;
|
||
return hasCanvas || hasSvg || hasContent;
|
||
});
|
||
}""",
|
||
timeout=HTTP_REQUEST_TIMEOUT_MS,
|
||
)
|
||
await page.evaluate("""async () => {
|
||
const delay = ms => new Promise(resolve => setTimeout(resolve, ms));
|
||
for (let i = 0; i < document.body.scrollHeight; i += 500) {
|
||
window.scrollTo(0, i);
|
||
await delay(100);
|
||
}
|
||
window.scrollTo(0, 0);
|
||
await delay(500);
|
||
}""")
|
||
except Exception:
|
||
pass
|
||
|
||
await self._wait_for_charts_stabilized(page)
|
||
logger.info(f"[_launch_and_login] Login + navigation successful for dashboard {dashboard_id}")
|
||
return browser, context, page
|
||
# endregion ScreenshotService._launch_and_login
|
||
|
||
# region ScreenshotService.capture_dashboard_chunks [TYPE Function] [C:4]
|
||
# @PURPOSE Capture per-tab screenshots: login → navigate → switch tabs → per-tab CDP screenshots.
|
||
# @PRE dashboard_id is valid, browser available.
|
||
# @POST Returns list of {tab_name, path} dicts — one per tab.
|
||
# @SIDE_EFFECT Launches browser, logs in, switches tabs, captures screenshots.
|
||
# @RATIONALE Multi-chunk: one screenshot per tab instead of one full-page.
|
||
# All screenshots written to output_dir; CDP fallback to Playwright full_page.
|
||
# @REJECTED Single full-page screenshot rejected for v2 — per-tab captures give LLM
|
||
# better visibility into individual tab content, especially for dashboards with
|
||
# many tabs where the full-page capture may miss lazy-loaded tab content.
|
||
async def capture_dashboard_chunks(
|
||
self,
|
||
dashboard_id: str,
|
||
output_dir: str,
|
||
parsed_context: dict | None = None,
|
||
) -> list[dict]:
|
||
"""Capture per-tab screenshots instead of one full-page.
|
||
|
||
Args:
|
||
dashboard_id: Superset dashboard ID
|
||
output_dir: Directory to save screenshots
|
||
parsed_context: Optional parsed context with activeTabs, native_filters from URL parse
|
||
|
||
Returns:
|
||
list of {tab_name, path} — one per tab
|
||
"""
|
||
import time as _time
|
||
|
||
timestamp = int(_time.time())
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
async with async_playwright() as p:
|
||
browser, context, page = await self._launch_and_login(p, dashboard_id, parsed_context)
|
||
try:
|
||
results: list[dict] = []
|
||
processed_tabs: set[str] = set()
|
||
|
||
async def _capture_tabs(depth: int = 0) -> None:
|
||
if depth > 3:
|
||
return
|
||
|
||
tab_selectors = [
|
||
".ant-tabs-nav-list .ant-tabs-tab",
|
||
".dashboard-component-tabs .ant-tabs-tab",
|
||
'[data-test="dashboard-component-tabs"] .ant-tabs-tab',
|
||
]
|
||
|
||
found_tabs = []
|
||
for selector in tab_selectors:
|
||
found_tabs = await page.locator(selector).all()
|
||
if found_tabs:
|
||
break
|
||
|
||
if not found_tabs:
|
||
return
|
||
|
||
logger.info(f"[capture_dashboard_chunks] Found {len(found_tabs)} tabs at depth {depth}")
|
||
for i, tab in enumerate(found_tabs):
|
||
try:
|
||
tab_text = (await tab.inner_text()).strip()
|
||
tab_id = f"{depth}_{i}_{tab_text}"
|
||
|
||
if tab_id in processed_tabs:
|
||
continue
|
||
|
||
if not await tab.is_visible():
|
||
continue
|
||
|
||
processed_tabs.add(tab_id)
|
||
logger.info(f"[capture_dashboard_chunks] Switching to tab: {tab_text}")
|
||
|
||
is_active = "ant-tabs-tab-active" in (await tab.get_attribute("class") or "")
|
||
if not is_active:
|
||
await tab.click()
|
||
try:
|
||
await page.wait_for_function(
|
||
"""() => {
|
||
const activeTab = document.querySelector('.ant-tabs-tab-active');
|
||
if (!activeTab) return true;
|
||
const tabPane = activeTab.closest('.ant-tabs')?.querySelector('.ant-tabs-content-holder');
|
||
if (!tabPane) return true;
|
||
const charts = tabPane.querySelectorAll('canvas, svg');
|
||
return charts.length > 0;
|
||
}""",
|
||
timeout=PLAYWRIGHT_WAIT_TIMEOUT_MS,
|
||
)
|
||
except Exception:
|
||
logger.warning(
|
||
f"[capture_dashboard_chunks] Content verification timed out for tab: {tab_text}"
|
||
)
|
||
|
||
# Wait for charts to stabilize
|
||
await self._wait_for_charts_stabilized(page)
|
||
|
||
# Resize viewport to 1920x1200 for consistent screenshots
|
||
await page.set_viewport_size({"width": 1920, "height": 1200})
|
||
await self._wait_for_resize_rendered(page, {})
|
||
|
||
# CDP screenshot with fallback
|
||
safe_tab = re.sub(r"[^\w\-_ ]", "", tab_text).strip().replace(" ", "_")[:40]
|
||
if not safe_tab:
|
||
safe_tab = f"tab_{depth}_{i}"
|
||
|
||
tab_filename = f"{dashboard_id}_{safe_tab}_{timestamp}.png"
|
||
tab_path = os.path.join(output_dir, tab_filename)
|
||
|
||
try:
|
||
cdp = await page.context.new_cdp_session(page)
|
||
screenshot_data = await cdp.send(
|
||
"Page.captureScreenshot",
|
||
{
|
||
"format": "png",
|
||
"fromSurface": True,
|
||
"captureBeyondViewport": True,
|
||
},
|
||
)
|
||
image_data = base64.b64decode(screenshot_data["data"])
|
||
with open(tab_path, "wb") as f:
|
||
f.write(image_data)
|
||
except Exception as cdp_err:
|
||
logger.warning(
|
||
f"[capture_dashboard_chunks] CDP failed for tab {tab_text}: {cdp_err}. "
|
||
"Falling back to Playwright full_page."
|
||
)
|
||
await page.screenshot(path=tab_path, full_page=True, timeout=PLAYWRIGHT_WAIT_TIMEOUT_MS)
|
||
|
||
logger.info(f"[capture_dashboard_chunks] Saved screenshot: {tab_path}")
|
||
results.append({"tab_name": tab_text, "path": tab_path})
|
||
|
||
# Recurse into nested tabs
|
||
await _capture_tabs(depth + 1)
|
||
|
||
except Exception as tab_e:
|
||
logger.warning(
|
||
f"[capture_dashboard_chunks] Failed to process tab {i} at depth {depth}: {tab_e}"
|
||
)
|
||
|
||
# Return to first tab
|
||
try:
|
||
first_tab = found_tabs[0]
|
||
if "ant-tabs-tab-active" not in (await first_tab.get_attribute("class") or ""):
|
||
await first_tab.click()
|
||
except Exception:
|
||
pass
|
||
|
||
await _capture_tabs()
|
||
|
||
# If no tabs found, capture the whole page as a single chunk
|
||
if not results:
|
||
logger.info("[capture_dashboard_chunks] No tabs found, capturing full-page as single chunk")
|
||
await self._wait_for_charts_stabilized(page)
|
||
await page.set_viewport_size({"width": 1920, "height": 1200})
|
||
|
||
tab_path = os.path.join(output_dir, f"{dashboard_id}_full_{timestamp}.png")
|
||
try:
|
||
cdp = await page.context.new_cdp_session(page)
|
||
screenshot_data = await cdp.send(
|
||
"Page.captureScreenshot",
|
||
{
|
||
"format": "png",
|
||
"fromSurface": True,
|
||
"captureBeyondViewport": True,
|
||
},
|
||
)
|
||
image_data = base64.b64decode(screenshot_data["data"])
|
||
with open(tab_path, "wb") as f:
|
||
f.write(image_data)
|
||
except Exception as cdp_err:
|
||
logger.warning(f"[capture_dashboard_chunks] CDP full fallback failed: {cdp_err}")
|
||
await page.screenshot(path=tab_path, full_page=True, timeout=PLAYWRIGHT_WAIT_TIMEOUT_MS)
|
||
|
||
results.append({"tab_name": "full", "path": tab_path})
|
||
|
||
return results
|
||
finally:
|
||
await browser.close()
|
||
# endregion ScreenshotService.capture_dashboard_chunks
|
||
|
||
# region ScreenshotService.capture_dashboard [TYPE Function] [C:4]
|
||
# @PURPOSE Captures multi-chunk screenshots, converts for LLM, archives to WebP.
|
||
# @PRE dashboard_id is a valid string, output_path is a writable path.
|
||
# @POST Returns list of {original, webp_path} dicts from WebP archive.
|
||
# Empty list on complete failure.
|
||
# @SIDE_EFFECT Launches browser, performs UI login, writes PNG/JPEG/WebP files;
|
||
# deletes intermediate PNG and JPEG files after conversion.
|
||
# @RATIONALE Refactored v2: delegates to capture_dashboard_chunks for per-tab PNGs,
|
||
# then converts for LLM (JPEG) and archive (WebP). Temp intermediates are cleaned up.
|
||
# @REJECTED Returning bool rejected for v2 — callers need access to archived WebP paths
|
||
# for persistence and LLM pipeline.
|
||
async def capture_dashboard(self, dashboard_id: str, output_path: str) -> tuple[list[str], list[dict]]:
|
||
"""Capture dashboard screenshots (multi-chunk), convert for LLM, archive to WebP.
|
||
|
||
Returns (jpeg_paths, archive_results) tuple.
|
||
jpeg_paths — list of JPEG paths ready for LLM analysis (caller must clean up).
|
||
archive_results — list of {original, webp_path} dicts from WebP archive.
|
||
"""
|
||
output_dir = os.path.dirname(output_path) or "."
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
with belief_scope("capture_dashboard", f"dashboard_id={dashboard_id}"):
|
||
logger.info(f"Capturing screenshots for dashboard {dashboard_id}")
|
||
|
||
# 1. Capture per-tab screenshots
|
||
chunks = await self.capture_dashboard_chunks(dashboard_id, output_dir)
|
||
png_paths = [c["path"] for c in chunks if c.get("path")]
|
||
|
||
if not png_paths:
|
||
logger.warning(f"[capture_dashboard] No screenshots captured for dashboard {dashboard_id}")
|
||
return [], []
|
||
|
||
# 2. Convert PNGs to JPEGs for LLM
|
||
jpeg_paths = self._convert_screenshots_for_llm(png_paths, output_dir)
|
||
logger.info(
|
||
f"[capture_dashboard] Converted {len(jpeg_paths)}/{len(png_paths)} PNGs to JPEGs for LLM"
|
||
)
|
||
|
||
# 3. Archive to WebP (deletes PNGs on success)
|
||
archive_results = self._archive_screenshots_as_webp(png_paths, output_dir)
|
||
logger.info(
|
||
f"[capture_dashboard] Archived {sum(1 for r in archive_results if r.get('webp_path'))} "
|
||
f"of {len(archive_results)} screenshots to WebP"
|
||
)
|
||
|
||
# 4. Return JPEGs for LLM — caller cleans up after analysis
|
||
return jpeg_paths, archive_results
|
||
# endregion ScreenshotService.capture_dashboard
|
||
|
||
# region ScreenshotService._convert_screenshots_for_llm [TYPE Function] [C:2]
|
||
# @BRIEF Convert PNG screenshots to JPEG for LLM transmission.
|
||
# @PRE png_paths is a list of existing PNG file paths.
|
||
# @POST Returns list of JPEG paths. JPEG files should be deleted after LLM call.
|
||
@staticmethod
|
||
def _convert_screenshots_for_llm(png_paths: list[str], output_dir: str) -> list[str]:
|
||
"""Convert PNG screenshots to JPEG for LLM transmission.
|
||
|
||
PNG → Pillow → JPEG quality=60, max 1024px width.
|
||
Returns list of JPEG paths.
|
||
"""
|
||
jpeg_paths: list[str] = []
|
||
for png_path in png_paths:
|
||
try:
|
||
img = Image.open(png_path)
|
||
if img.mode in ("RGBA", "P"):
|
||
img = img.convert("RGB")
|
||
|
||
max_width = 1024
|
||
if img.width > max_width:
|
||
scale = max_width / img.width
|
||
new_w = int(img.width * scale)
|
||
new_h = int(img.height * scale)
|
||
img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||
|
||
base = os.path.splitext(os.path.basename(png_path))[0]
|
||
jpeg_path = os.path.join(output_dir, f"{base}_llm.jpg")
|
||
img.save(jpeg_path, format="JPEG", quality=60, optimize=True)
|
||
jpeg_paths.append(jpeg_path)
|
||
except Exception as e:
|
||
logger.warning(f"[convert_for_llm] Failed to convert {png_path}: {e}")
|
||
|
||
return jpeg_paths
|
||
# endregion ScreenshotService._convert_screenshots_for_llm
|
||
|
||
# region ScreenshotService._archive_screenshots_as_webp [TYPE Function] [C:2]
|
||
# @BRIEF Convert PNG screenshots to WebP for archive.
|
||
# @PRE png_paths is a list of existing PNG file paths.
|
||
# @POST Returns list of {original, webp_path} dicts. PNG deleted after successful WebP save.
|
||
@staticmethod
|
||
def _archive_screenshots_as_webp(png_paths: list[str], archive_dir: str) -> list[dict]:
|
||
"""Convert PNG screenshots to WebP for archive.
|
||
|
||
PNG → Pillow → WebP lossy quality=80.
|
||
Deletes PNG after WebP saved.
|
||
Returns list of {original, webp_path} dicts.
|
||
"""
|
||
results: list[dict] = []
|
||
for png_path in png_paths:
|
||
try:
|
||
img = Image.open(png_path)
|
||
if img.mode in ("RGBA", "P"):
|
||
img = img.convert("RGB")
|
||
|
||
base = os.path.splitext(os.path.basename(png_path))[0]
|
||
webp_path = os.path.join(archive_dir, f"{base}.webp")
|
||
img.save(webp_path, format="WEBP", quality=80, lossless=False)
|
||
|
||
# Delete PNG after successful WebP save
|
||
os.remove(png_path)
|
||
|
||
results.append({"original": png_path, "webp_path": webp_path})
|
||
except Exception as e:
|
||
logger.warning(f"[archive_webp] Failed to convert {png_path}: {e}. Keeping PNG.")
|
||
results.append({"original": png_path, "webp_path": None})
|
||
|
||
return results
|
||
# endregion ScreenshotService._archive_screenshots_as_webp
|
||
|
||
# region ScreenshotService._cleanup_temp_files [TYPE Function] [C:1]
|
||
# @BRIEF Delete temporary files (PNG, JPEG intermediates).
|
||
@staticmethod
|
||
def _cleanup_temp_files(paths: list[str]) -> None:
|
||
"""Delete temporary files (PNG, JPEG intermediates)."""
|
||
for path in paths:
|
||
try:
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
except Exception as e:
|
||
logger.warning(f"[cleanup] Failed to delete {path}: {e}")
|
||
# endregion ScreenshotService._cleanup_temp_files
|
||
# #endregion ScreenshotService
|
||
|
||
# #region LLMClient [TYPE Class]
|
||
# @BRIEF Wrapper for LLM provider APIs.
|
||
class LLMClient:
|
||
# region LLMClient.__init__ [TYPE Function]
|
||
# @PURPOSE Initializes the LLMClient with provider settings.
|
||
# @PRE api_key, base_url, and default_model are non-empty strings.
|
||
def __init__(self, provider_type: LLMProviderType, api_key: str, base_url: str, default_model: str):
|
||
self.provider_type = provider_type
|
||
normalized_key = (api_key or "").strip()
|
||
if normalized_key.lower().startswith("bearer "):
|
||
normalized_key = normalized_key[7:].strip()
|
||
self.api_key = normalized_key
|
||
self.base_url = base_url
|
||
self.default_model = default_model
|
||
|
||
# DEBUG: Log initialization parameters (without exposing full API key)
|
||
logger.info("[LLMClient.__init__] Initializing LLM client:")
|
||
logger.info(f"[LLMClient.__init__] Provider Type: {provider_type}")
|
||
logger.info(f"[LLMClient.__init__] Base URL: {base_url}")
|
||
logger.info(f"[LLMClient.__init__] Default Model: {default_model}")
|
||
logger.info(f"[LLMClient.__init__] API Key (first 8 chars): {self.api_key[:8] if self.api_key and len(self.api_key) > 8 else 'EMPTY_OR_NONE'}...")
|
||
logger.info(f"[LLMClient.__init__] API Key Length: {len(self.api_key) if self.api_key else 0}")
|
||
|
||
# Some OpenAI-compatible gateways are strict about auth header naming.
|
||
default_headers = {"Authorization": f"Bearer {self.api_key}"}
|
||
if self.provider_type == LLMProviderType.OPENROUTER:
|
||
default_headers["HTTP-Referer"] = (
|
||
os.getenv("OPENROUTER_SITE_URL", "").strip()
|
||
or os.getenv("APP_BASE_URL", "").strip()
|
||
or "http://localhost:8000"
|
||
)
|
||
default_headers["X-Title"] = os.getenv("OPENROUTER_APP_NAME", "").strip() or "ss-tools"
|
||
if self.provider_type == LLMProviderType.KILO:
|
||
default_headers["Authentication"] = f"Bearer {self.api_key}"
|
||
default_headers["X-API-Key"] = self.api_key
|
||
# LiteLLM proxy uses standard OpenAI-compatible Bearer auth — no special headers needed.
|
||
# It routes to upstream providers transparently, and the default Authorization header
|
||
# is sufficient. No additional headers like HTTP-Referer or X-API-Key are required.
|
||
|
||
ssl_verify = self._get_ssl_verify()
|
||
logger.info(f"[LLMClient.__init__] SSL Verify: {ssl_verify} (set LLM_SSL_VERIFY=false to disable)")
|
||
|
||
http_client = httpx.AsyncClient(
|
||
headers=default_headers,
|
||
timeout=LLM_HTTP_TIMEOUT_S,
|
||
verify=ssl_verify,
|
||
)
|
||
self.client = AsyncOpenAI(
|
||
api_key=self.api_key,
|
||
base_url=base_url,
|
||
default_headers=default_headers,
|
||
http_client=http_client,
|
||
)
|
||
# endregion LLMClient.__init__
|
||
|
||
# region LLMClient._get_ssl_verify [TYPE Function]
|
||
# @PURPOSE Resolve SSL verification flag from environment.
|
||
# @POST Returns SSLContext with system CA dir when enabled,
|
||
# False when LLM_SSL_VERIFY env var is "false"/"0"/"no"/"off".
|
||
# @RATIONALE Используем capath=/etc/ssl/certs/ вместо cafile, потому что
|
||
# OpenSSL 3.x не использует intermediate CA сертификаты из cafile для
|
||
# построения цепочки (verify code 20). capath с хеш-симлинками работает
|
||
# корректно (verify code 0). Оба пути — cafile и capath — указывают на
|
||
# один и тот же набор сертификатов, но capath правильно обрабатывает
|
||
# цепочку Root → Policy → Issuing.
|
||
# @REJECTED verify=<string> отвергнут — httpx 0.28.x депрекейтит строковый
|
||
# путь в verify=, требует SSLContext.
|
||
# @REJECTED cafile отвергнут — OpenSSL 3.x не использует intermediate CA
|
||
# из единого bundle-файла. Только capath с хеш-симлинками даёт code 0.
|
||
@staticmethod
|
||
def _get_ssl_verify() -> ssl.SSLContext | bool:
|
||
raw = os.getenv("LLM_SSL_VERIFY", "true").strip().lower()
|
||
if raw in ("false", "0", "no", "off"):
|
||
return False
|
||
ca_dir = "/etc/ssl/certs"
|
||
if os.path.isdir(ca_dir):
|
||
return ssl.create_default_context(capath=ca_dir)
|
||
# fallback: если директории нет (редко), используем дефолтный
|
||
return ssl.create_default_context()
|
||
# endregion LLMClient._get_ssl_verify
|
||
|
||
# region LLMClient._format_connection_error [TYPE Function]
|
||
# @PURPOSE Format exception chain for diagnostics, extracting httpx cause details.
|
||
# @POST Returns a human-readable string with the full error chain.
|
||
@staticmethod
|
||
def _format_connection_error(exc: Exception) -> str:
|
||
parts = [f"{type(exc).__name__}: {exc!s}"]
|
||
cause = exc.__cause__ or exc.__context__
|
||
while cause:
|
||
parts.append(f" └─ {type(cause).__name__}: {cause!s}")
|
||
cause = cause.__cause__ or cause.__context__
|
||
return "\n".join(parts)
|
||
# endregion LLMClient._format_connection_error
|
||
|
||
# region LLMClient._supports_json_response_format [TYPE Function]
|
||
# @PURPOSE Detect whether provider/model is likely compatible with response_format=json_object.
|
||
# @PRE Client initialized with base_url and default_model.
|
||
# @POST Returns False for known-incompatible combinations to avoid avoidable 400 errors.
|
||
def _supports_json_response_format(self) -> bool:
|
||
base = (self.base_url or "").lower()
|
||
model = (self.default_model or "").lower()
|
||
|
||
# OpenRouter routes to many upstream providers; some models reject json_object mode.
|
||
if "openrouter.ai" in base:
|
||
incompatible_tokens = (
|
||
"stepfun/",
|
||
"step-",
|
||
":free",
|
||
)
|
||
if any(token in model for token in incompatible_tokens):
|
||
return False
|
||
return True
|
||
# endregion LLMClient._supports_json_response_format
|
||
|
||
# region LLMClient.get_json_completion [TYPE Function]
|
||
# @PURPOSE Helper to handle LLM calls with JSON mode and fallback parsing.
|
||
# @PRE messages is a list of valid message dictionaries.
|
||
# @POST Returns a parsed JSON dictionary.
|
||
# @SIDE_EFFECT Calls external LLM API.
|
||
def _should_retry(exception: Exception) -> bool:
|
||
"""Custom retry predicate that excludes authentication errors."""
|
||
# Don't retry on authentication errors
|
||
if isinstance(exception, OpenAIAuthenticationError):
|
||
return False
|
||
# Retry on rate limit errors and other exceptions
|
||
return isinstance(exception, (RateLimitError, Exception))
|
||
|
||
@retry(
|
||
stop=stop_after_attempt(5),
|
||
wait=wait_exponential(multiplier=2, min=5, max=60),
|
||
retry=retry_if_exception(_should_retry),
|
||
reraise=True
|
||
)
|
||
async def get_json_completion(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
|
||
with belief_scope("get_json_completion"):
|
||
response = None
|
||
try:
|
||
use_json_mode = self._supports_json_response_format()
|
||
try:
|
||
logger.info(
|
||
f"[get_json_completion] Attempting LLM call for model: {self.default_model} "
|
||
f"(json_mode={'on' if use_json_mode else 'off'})"
|
||
)
|
||
logger.info(f"[get_json_completion] Base URL being used: {self.base_url}")
|
||
logger.info(f"[get_json_completion] Number of messages: {len(messages)}")
|
||
logger.info(f"[get_json_completion] API Key present: {bool(self.api_key and len(self.api_key) > 0)}")
|
||
|
||
if use_json_mode:
|
||
response = await self.client.chat.completions.create(
|
||
model=self.default_model,
|
||
messages=messages,
|
||
response_format={"type": "json_object"}
|
||
)
|
||
else:
|
||
response = await self.client.chat.completions.create(
|
||
model=self.default_model,
|
||
messages=messages
|
||
)
|
||
except Exception as e:
|
||
if use_json_mode and (
|
||
"JSON mode is not enabled" in str(e)
|
||
or "json_object is not supported" in str(e).lower()
|
||
or "response_format" in str(e).lower()
|
||
or "400" in str(e)
|
||
):
|
||
logger.warning(f"[get_json_completion] JSON mode failed or not supported: {e!s}. Falling back to plain text response.")
|
||
response = await self.client.chat.completions.create(
|
||
model=self.default_model,
|
||
messages=messages
|
||
)
|
||
else:
|
||
raise e
|
||
|
||
logger.debug(f"[get_json_completion] LLM Response: {response}")
|
||
except OpenAIAuthenticationError as e:
|
||
logger.error(f"[get_json_completion] Authentication error: {e!s}")
|
||
# Do not retry on auth errors - re-raise to stop retry
|
||
raise
|
||
except RateLimitError as e:
|
||
logger.warning(f"[get_json_completion] Rate limit hit: {e!s}")
|
||
|
||
# Extract retry_delay from error metadata if available
|
||
retry_delay = 5.0 # Default fallback
|
||
try:
|
||
# Based on logs, the raw response is in e.body or e.response.json()
|
||
# The logs show 'metadata': {'raw': '...'} which suggests a proxy or specific client wrapper
|
||
# Let's try to find the 'retryDelay' in the error message or response
|
||
import re
|
||
|
||
# Try to find "retryDelay": "XXs" in the string representation of the error
|
||
error_str = str(e)
|
||
match = re.search(r'"retryDelay":\s*"(\d+)s"', error_str)
|
||
if match:
|
||
retry_delay = float(match.group(1))
|
||
else:
|
||
# Try to parse from response if it's a standard OpenAI-like error with body
|
||
if hasattr(e, 'body') and isinstance(e.body, dict):
|
||
# Some providers put it in details
|
||
details = e.body.get('error', {}).get('details', [])
|
||
for detail in details:
|
||
if detail.get('@type') == 'type.googleapis.com/google.rpc.RetryInfo':
|
||
delay_str = detail.get('retryDelay', '5s')
|
||
retry_delay = float(delay_str.rstrip('s'))
|
||
break
|
||
except Exception as parse_e:
|
||
logger.debug(f"[get_json_completion] Failed to parse retry delay: {parse_e}")
|
||
|
||
# Add a small safety margin (0.5s) as requested
|
||
wait_time = retry_delay + 0.5
|
||
logger.info(f"[get_json_completion] Waiting for {wait_time}s before retry...")
|
||
await asyncio.sleep(wait_time)
|
||
raise
|
||
except Exception as e:
|
||
logger.error(
|
||
f"[get_json_completion] LLM call failed.\n"
|
||
f"{self._format_connection_error(e)}"
|
||
)
|
||
raise
|
||
|
||
if not response or not hasattr(response, 'choices') or not response.choices:
|
||
raise RuntimeError(f"Invalid LLM response: {response}")
|
||
|
||
content = response.choices[0].message.content
|
||
logger.debug(f"[get_json_completion] Raw content to parse: {content}")
|
||
|
||
try:
|
||
return json.loads(content)
|
||
except json.JSONDecodeError:
|
||
logger.warning("[get_json_completion] Failed to parse JSON directly, attempting to extract from code blocks")
|
||
if "```json" in content:
|
||
json_str = content.split("```json")[1].split("```")[0].strip()
|
||
return json.loads(json_str)
|
||
elif "```" in content:
|
||
json_str = content.split("```")[1].split("```")[0].strip()
|
||
return json.loads(json_str)
|
||
else:
|
||
raise
|
||
# endregion LLMClient.get_json_completion
|
||
|
||
# region LLMClient.test_runtime_connection [TYPE Function]
|
||
# @PURPOSE Validate provider credentials using the same chat completions transport as runtime analysis.
|
||
# @PRE Client is initialized with provider credentials and default_model.
|
||
# @POST Returns lightweight JSON payload when runtime auth/model path is valid.
|
||
# @SIDE_EFFECT Calls external LLM API.
|
||
async def test_runtime_connection(self) -> dict[str, Any]:
|
||
with belief_scope("test_runtime_connection"):
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": 'Return exactly this JSON object and nothing else: {"ok": true}',
|
||
}
|
||
]
|
||
return await self.get_json_completion(messages)
|
||
# endregion LLMClient.test_runtime_connection
|
||
|
||
# region LLMClient.fetch_models [TYPE Function]
|
||
# @PURPOSE Fetch available models from the provider's API.
|
||
# @PRE Client is initialized with provider credentials.
|
||
# @POST Returns a list of model ID strings.
|
||
# @SIDE_EFFECT Calls external LLM API /v1/models endpoint.
|
||
async def fetch_models(self) -> list[str]:
|
||
with belief_scope("LLMClient.fetch_models"):
|
||
try:
|
||
response = await self.client.models.list()
|
||
model_ids = [m.id for m in response.data]
|
||
model_ids.sort()
|
||
logger.reason(
|
||
f"[LLMClient.fetch_models] Fetched {len(model_ids)} models from {self.base_url}",
|
||
extra={"src": "LLMClient.fetch_models"},
|
||
)
|
||
return model_ids
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"[LLMClient.fetch_models] Failed to fetch models.\n"
|
||
f"{self._format_connection_error(e)}",
|
||
)
|
||
raise
|
||
# endregion LLMClient.fetch_models
|
||
|
||
# region LLMClient.analyze_dashboard [TYPE Function]
|
||
# @PURPOSE Sends dashboard data (screenshot + logs) to LLM for health analysis.
|
||
# @PRE screenshot_path exists, logs is a list of strings.
|
||
# @POST Returns a structured analysis dictionary (status, summary, issues).
|
||
# @SIDE_EFFECT Reads screenshot file and calls external LLM API.
|
||
# @RATIONALE Delegates to analyze_dashboard_multimodal for single-screenshot
|
||
# backward compatibility. Keeps the same contract for v1 consumers.
|
||
async def analyze_dashboard(
|
||
self,
|
||
screenshot_path: str,
|
||
logs: list[str],
|
||
prompt_template: str = DEFAULT_LLM_PROMPTS["dashboard_validation_prompt"],
|
||
) -> dict[str, Any]:
|
||
# Delegate to multimodal variant for backward compat with v1 consumers.
|
||
return await self.analyze_dashboard_multimodal(
|
||
screenshot_paths=[screenshot_path],
|
||
logs=logs,
|
||
prompt_template=prompt_template,
|
||
)
|
||
# endregion LLMClient.analyze_dashboard
|
||
|
||
# region LLMClient._reduce_image_quality [TYPE Function] [C:2]
|
||
# @PURPOSE Open, resize, and compress a screenshot image for LLM consumption.
|
||
# @PRE path points to an existing image file.
|
||
# @POST Returns (base64_str, byte_size) tuple.
|
||
@staticmethod
|
||
def _reduce_image_quality(
|
||
path: str,
|
||
max_width: int = 1024,
|
||
image_quality: int = 60,
|
||
) -> tuple[str, int]:
|
||
"""
|
||
Open, resize, compress, and base64-encode an image.
|
||
|
||
Returns (base64_str, byte_size).
|
||
"""
|
||
with Image.open(path) as img:
|
||
if img.mode in ("RGBA", "P"):
|
||
img = img.convert("RGB")
|
||
if img.width > max_width or img.height > 2048:
|
||
scale = min(max_width / img.width, 2048 / img.height)
|
||
if scale < 1.0:
|
||
new_width = int(img.width * scale)
|
||
new_height = int(img.height * scale)
|
||
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||
buffer = io.BytesIO()
|
||
img.save(buffer, format="JPEG", quality=image_quality, optimize=True)
|
||
raw = buffer.getvalue()
|
||
return base64.b64encode(raw).decode("utf-8"), len(raw)
|
||
# endregion LLMClient._reduce_image_quality
|
||
|
||
# region LLMClient._estimate_payload_size [TYPE Function] [C:2]
|
||
# @PURPOSE Estimate LLM payload size in tokens before sending.
|
||
# @POST Returns {estimated_tokens, exceeds_limit, pct_of_limit} dict.
|
||
# @RATIONALE FR-056: if >80% of model context window, trigger quality reduction.
|
||
@staticmethod
|
||
def _estimate_payload_size(
|
||
image_paths: list[str],
|
||
text_length: int,
|
||
model_context: int = 128000,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
Estimate token usage for multimodal payload.
|
||
|
||
Rough heuristic: 1 image token ~ 258 tokens (GPT-4o), text ~4 chars/token.
|
||
Returns {estimated_tokens, exceeds_limit, pct_of_limit}
|
||
"""
|
||
image_tokens = len(image_paths) * 258 * 5 # rough upper bound for compressed images
|
||
text_tokens = text_length // 4
|
||
total_tokens = image_tokens + text_tokens
|
||
exceeds_limit = total_tokens > (model_context * 0.8)
|
||
return {
|
||
"estimated_tokens": total_tokens,
|
||
"exceeds_limit": exceeds_limit,
|
||
"pct_of_limit": round(total_tokens / model_context * 100, 1),
|
||
}
|
||
# endregion LLMClient._estimate_payload_size
|
||
|
||
# region LLMClient._deduplicate_issues [TYPE Function] [C:2]
|
||
# @PURPOSE Deduplicate issues by (severity, message, location) while preserving order.
|
||
def _deduplicate_issues(self, issues: list[dict]) -> list[dict]:
|
||
seen: set[tuple[str, str, str]] = set()
|
||
result: list[dict] = []
|
||
for issue in issues:
|
||
key = (issue.get("severity", ""), issue.get("message", ""), issue.get("location", "") or "")
|
||
if key not in seen:
|
||
seen.add(key)
|
||
result.append(issue)
|
||
return result
|
||
|
||
# endregion LLMClient._deduplicate_issues
|
||
|
||
# region LLMClient._optimize_images [TYPE Function] [C:2]
|
||
# @PURPOSE Convert screenshot paths to base64 at given quality, with fallback to raw read.
|
||
def _optimize_images(self, paths: list[str], max_width: int, quality: int) -> list[str]:
|
||
encoded: list[str] = []
|
||
for path in paths:
|
||
try:
|
||
b64, _ = self._reduce_image_quality(path, max_width, quality)
|
||
encoded.append(b64)
|
||
except Exception as e:
|
||
logger.warning(f"[_optimize_images] Optimization failed for {path}: {e}")
|
||
with open(path, "rb") as f:
|
||
raw = f.read()
|
||
b64 = base64.b64encode(raw).decode("utf-8")
|
||
encoded.append(b64)
|
||
return encoded
|
||
|
||
# endregion LLMClient._optimize_images
|
||
|
||
# region LLMClient._merge_chunk_results [TYPE Function] [C:2]
|
||
# @PURPOSE Merge multiple chunk analyses into one. Takes the worst status,
|
||
# concatenates summaries, and deduplicates issues.
|
||
# @PRE chunks is a non-empty list of {status, summary, issues} dicts.
|
||
# @POST Returns a single merged dict with chunk_count.
|
||
def _merge_chunk_results(self, chunks: list[dict[str, Any]]) -> dict[str, Any]:
|
||
STATUS_ORDER = {"FAIL": 0, "WARN": 1, "PASS": 2, "UNKNOWN": 3}
|
||
worst_status = "UNKNOWN"
|
||
all_summaries: list[str] = []
|
||
all_issues: list[dict] = []
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
s = chunk.get("status", "UNKNOWN")
|
||
if STATUS_ORDER.get(s, 3) < STATUS_ORDER.get(worst_status, 3):
|
||
worst_status = s
|
||
all_summaries.append(f"[Chunk {i + 1}/{len(chunks)}] {chunk.get('summary', 'No summary')}")
|
||
all_issues.extend(chunk.get("issues", []))
|
||
|
||
merged: dict[str, Any] = {
|
||
"status": worst_status,
|
||
"summary": " | ".join(all_summaries),
|
||
"issues": self._deduplicate_issues(all_issues),
|
||
"chunk_count": len(chunks),
|
||
}
|
||
return merged
|
||
|
||
# endregion LLMClient._merge_chunk_results
|
||
|
||
# region LLMClient._call_llm_for_images [TYPE Function] [C:2]
|
||
# @PURPOSE Send a single chunk of images to the LLM and return parsed result.
|
||
async def _call_llm_for_images(
|
||
self, encoded_images: list[str], prompt: str
|
||
) -> dict[str, Any]:
|
||
content: list[dict] = [{"type": "text", "text": prompt}]
|
||
for b64_img in encoded_images:
|
||
content.append({
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/jpeg;base64,{b64_img}"},
|
||
})
|
||
messages = [{"role": "user", "content": content}]
|
||
return await self.get_json_completion(messages)
|
||
|
||
# endregion LLMClient._call_llm_for_images
|
||
|
||
# region LLMClient.analyze_dashboard_multimodal [TYPE Function] [C:3]
|
||
# @PURPOSE Path A: send screenshots + logs to multimodal LLM, with chunking support.
|
||
# @PRE screenshot_paths is a non-empty list of paths.
|
||
# @POST Returns dict {status, summary, issues} with optional chunk_count.
|
||
# @SIDE_EFFECT Compresses images, calls external LLM API (possibly multiple times for chunks).
|
||
# @RATIONALE Screenshots are split into chunks of max_images to respect provider image limits.
|
||
# Quality reduction is skipped when chunking — each chunk fits the limit by definition.
|
||
# Results are merged via _merge_chunk_results.
|
||
async def analyze_dashboard_multimodal(
|
||
self,
|
||
screenshot_paths: list[str],
|
||
logs: list[str],
|
||
prompt_template: str = DEFAULT_LLM_PROMPTS["dashboard_validation_prompt"],
|
||
max_width: int = 1024,
|
||
image_quality: int = 60,
|
||
max_images: int | None = None,
|
||
) -> dict[str, Any]:
|
||
with belief_scope("analyze_dashboard_multimodal"):
|
||
if not screenshot_paths:
|
||
raise ValueError("screenshot_paths must be a non-empty list")
|
||
|
||
# 1. Optimize all images at requested quality
|
||
encoded_images = self._optimize_images(screenshot_paths, max_width, image_quality)
|
||
|
||
log_text = "\n".join(logs)
|
||
prompt = render_prompt(prompt_template, {"logs": log_text})
|
||
|
||
# 2. Determine chunking
|
||
n_total = len(encoded_images)
|
||
chunk_size = max_images if (max_images and max_images > 0 and max_images < n_total) else n_total
|
||
is_chunking = chunk_size < n_total
|
||
|
||
if is_chunking:
|
||
logger.reason(
|
||
f"[analyze_dashboard_multimodal] Chunking {n_total} images into "
|
||
f"{(n_total + chunk_size - 1) // chunk_size} chunks of {chunk_size}",
|
||
extra={"src": "analyze_dashboard_multimodal", "total": n_total, "chunk_size": chunk_size},
|
||
)
|
||
# Skip quality reduction: each chunk has ≤ max_images images,
|
||
# well within the context window at normal quality.
|
||
else:
|
||
# Single batch: estimate payload and reduce quality if needed
|
||
estimate = self._estimate_payload_size(
|
||
screenshot_paths, len(prompt) + len(log_text)
|
||
)
|
||
if estimate["exceeds_limit"] and image_quality > 30:
|
||
logger.info(
|
||
f"[analyze_dashboard_multimodal] Payload estimated at {estimate['pct_of_limit']}% "
|
||
f"of context window. Reducing image quality to 30."
|
||
)
|
||
encoded_images = self._optimize_images(screenshot_paths, max_width, image_quality=30)
|
||
|
||
# 3. Split into chunks
|
||
chunks: list[list[str]] = [
|
||
encoded_images[i:i + chunk_size]
|
||
for i in range(0, n_total, chunk_size)
|
||
]
|
||
|
||
# 4. Call LLM — parallel for multiple chunks, single for one
|
||
try:
|
||
if len(chunks) == 1:
|
||
result = await self._call_llm_for_images(chunks[0], prompt)
|
||
else:
|
||
tasks = [self._call_llm_for_images(chunk, prompt) for chunk in chunks]
|
||
chunk_results = await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
valid_results: list[dict] = []
|
||
for i, cr in enumerate(chunk_results):
|
||
if isinstance(cr, Exception):
|
||
logger.error(f"[analyze_dashboard_multimodal] Chunk {i + 1}/{len(chunks)} failed: {cr!s}")
|
||
valid_results.append({
|
||
"status": "UNKNOWN",
|
||
"summary": f"Chunk {i + 1} failed: {cr!s}",
|
||
"issues": [],
|
||
})
|
||
else:
|
||
valid_results.append(cr)
|
||
|
||
result = self._merge_chunk_results(valid_results)
|
||
except Exception as e:
|
||
logger.error(f"[analyze_dashboard_multimodal] Failed to get analysis: {e!s}")
|
||
return {
|
||
"status": "UNKNOWN",
|
||
"summary": f"Failed to get response from LLM: {e!s}",
|
||
"issues": [{"severity": "UNKNOWN", "message": "LLM provider returned empty or invalid response"}],
|
||
}
|
||
|
||
return result
|
||
# endregion LLMClient.analyze_dashboard_multimodal
|
||
|
||
# region LLMClient.analyze_dashboard_text_batch [TYPE Function] [C:3]
|
||
# @PURPOSE Path B batch: multiple dashboards in a single text-only LLM call.
|
||
# @PRE payloads is a non-empty list of {dashboard_id, topology, dataset_health, log_text} dicts.
|
||
# @POST Returns dict {dashboards: [{dashboard_id, status, summary, issues}]}.
|
||
# Missing/parse-error dashboard_id -> marked UNKNOWN individually.
|
||
# @RATIONALE Text-only batch avoids image token costs. Uses per-dashboard sections
|
||
# with explicit JSON response contract. Fallback ensures partial results survive
|
||
# single-dashboard parse failures.
|
||
@retry(
|
||
stop=stop_after_attempt(5),
|
||
wait=wait_exponential(multiplier=2, min=5, max=60),
|
||
retry=retry_if_exception(_should_retry),
|
||
reraise=True,
|
||
)
|
||
async def analyze_dashboard_text_batch(
|
||
self,
|
||
payloads: list[dict],
|
||
prompt_template: str,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
Batch analyze multiple dashboards in one LLM call.
|
||
|
||
payloads: list of dicts with keys:
|
||
- dashboard_id (str)
|
||
- topology (str) — dashboard structure
|
||
- dataset_health (str) — health results
|
||
- log_text (str) — execution logs
|
||
|
||
Returns dict like {dashboards: [{dashboard_id, status, summary, issues}]}
|
||
"""
|
||
if not payloads:
|
||
return {"dashboards": []}
|
||
|
||
# 1. Build per-dashboard sections
|
||
sections = []
|
||
for i, p in enumerate(payloads):
|
||
did = p.get("dashboard_id", "UNKNOWN")
|
||
top = p.get("topology", "")
|
||
first_line = top.split("\n")[0] if top else "(no topology)"
|
||
section = (
|
||
f'─── Dashboard {i + 1}: "{first_line}" (id: {did}) ───\n'
|
||
f"{top}\n\n"
|
||
f"Dataset health:\n{p.get('dataset_health', '')}\n\n"
|
||
f"Logs:\n{p.get('log_text', '')}"
|
||
)
|
||
sections.append(section)
|
||
|
||
full_prompt = prompt_template.replace("{total_dashboards}", str(len(payloads)))
|
||
full_prompt += (
|
||
"\n\nRespond with a JSON object containing EACH dashboard's results:\n"
|
||
'{"dashboards": [{"dashboard_id": "...", "status": "...", "summary": "...", "issues": [...]}]}\n\n'
|
||
)
|
||
full_prompt += "\n---\n".join(sections)
|
||
|
||
messages = [{"role": "user", "content": full_prompt}]
|
||
return await self.get_json_completion(messages)
|
||
# endregion LLMClient.analyze_dashboard_text_batch
|
||
|
||
# #endregion LLMClient
|
||
|
||
# #region DatasetHealthChecker [C:3] [TYPE Class]
|
||
# @BRIEF Checks dataset accessibility and KXD connectivity via Superset API.
|
||
# @LAYER Service
|
||
# @RELATION CALLS -> [SupersetClient]
|
||
# @INVARIANT Every unique dataset referenced by dashboard charts is checked.
|
||
# @RATIONALE Without dataset health checking, silent KXD errors (connection refused, timeout)
|
||
# are invisible to the LLM validation. Screenshot captures visual state but doesn't
|
||
# verify that data actually arrived (vs. cache).
|
||
class DatasetHealthChecker:
|
||
# region DatasetHealthChecker.__init__ [TYPE Function]
|
||
# @PURPOSE Initialize with a SupersetClient-compatible instance.
|
||
# @PRE client is a SupersetClient (sync, wrapped via asyncio.to_thread) or AsyncSupersetClient.
|
||
# @POST self.client is ready for health checks.
|
||
def __init__(self, client: Any):
|
||
self.client = client
|
||
# endregion DatasetHealthChecker.__init__
|
||
|
||
# region DatasetHealthChecker._call_sync [TYPE Function]
|
||
# @PURPOSE Wrap a sync client method call in asyncio.to_thread for async compat.
|
||
# @PRE method is a callable on self.client.
|
||
# @POST Returns the result of method(*args, **kwargs) executed in a thread.
|
||
@staticmethod
|
||
async def _call_sync(method, *args: Any, **kwargs: Any) -> Any:
|
||
"""Call a potentially sync method in a thread, or await if already async."""
|
||
if asyncio.iscoroutinefunction(method):
|
||
return await method(*args, **kwargs)
|
||
return await asyncio.to_thread(method, *args, **kwargs)
|
||
# endregion DatasetHealthChecker._call_sync
|
||
|
||
# region DatasetHealthChecker.check_dataset_health [TYPE Function]
|
||
# @PURPOSE Fetch dataset metadata and verify level 1-2 accessibility.
|
||
# @PRE dataset_id is a valid Superset dataset ID.
|
||
# @POST Returns dict with level 1-2 health fields.
|
||
# @SIDE_EFFECT Calls GET /api/v1/dataset/{id} via client.get_dataset
|
||
async def check_dataset_health(self, dataset_id: int) -> dict:
|
||
"""
|
||
Check a single dataset's accessibility (levels 1-2 per FR-044).
|
||
|
||
Level 1: metadata_accessible — HTTP 200 from GET /api/v1/dataset/{id}
|
||
Level 2: datasource_resolvable — database info available
|
||
|
||
Returns dict with:
|
||
dataset_id, dataset_name, database_name, backend, kind,
|
||
metadata_accessible (bool), error (str|None)
|
||
"""
|
||
try:
|
||
dataset = await self._call_sync(self.client.get_dataset, dataset_id)
|
||
# The response from Superset may have a 'result' wrapper or be flat
|
||
result_data = dataset.get("result", dataset) if isinstance(dataset, dict) else {}
|
||
# Extract database info
|
||
database = result_data.get("database", {}) or {}
|
||
result = {
|
||
"dataset_id": dataset_id,
|
||
"dataset_name": result_data.get("table_name", f"dataset_{dataset_id}"),
|
||
"database_name": database.get("database_name", "unknown"),
|
||
"backend": database.get("backend", "unknown"),
|
||
"kind": result_data.get("kind", "physical"),
|
||
"metadata_accessible": True,
|
||
"error": None,
|
||
}
|
||
return result
|
||
except Exception as e:
|
||
return {
|
||
"dataset_id": dataset_id,
|
||
"dataset_name": f"dataset_{dataset_id}",
|
||
"database_name": "unknown",
|
||
"backend": "unknown",
|
||
"kind": "unknown",
|
||
"metadata_accessible": False,
|
||
"error": str(e),
|
||
}
|
||
# endregion DatasetHealthChecker.check_dataset_health
|
||
|
||
# region DatasetHealthChecker.check_chart_data [TYPE Function]
|
||
# @PURPOSE Execute chart data query (level 3-4 per FR-044).
|
||
# @PRE chart_id is valid, form_data is constructed from chart params.
|
||
# @POST Returns dict with execution result.
|
||
# @SIDE_EFFECT Calls POST /api/v1/chart/data via client.network.request
|
||
async def check_chart_data(self, chart_id: int, form_data: dict) -> dict:
|
||
"""
|
||
Execute a chart query to verify data returns.
|
||
|
||
Level 3: query_executable — POST /api/v1/chart/data succeeds
|
||
Level 4: data_returned — row_count > 0 or no error
|
||
|
||
Returns dict with:
|
||
chart_id, executed (bool), duration_ms (int|None),
|
||
row_count (int|None), error (str|None)
|
||
"""
|
||
import time
|
||
start = time.time()
|
||
try:
|
||
# Use the client's network layer for the chart data POST.
|
||
# For sync SupersetClient: network.request(...) is synchronous.
|
||
# We wrap it via asyncio.to_thread if it's a sync method.
|
||
payload = json.dumps(form_data)
|
||
headers = {"Content-Type": "application/json"}
|
||
network_request = self.client.network.request
|
||
result = await self._call_sync(
|
||
network_request,
|
||
"POST", "/chart/data",
|
||
data=payload,
|
||
headers=headers,
|
||
)
|
||
duration_ms = int((time.time() - start) * 1000)
|
||
|
||
# Normalize response — may have 'result' wrapper
|
||
rows = []
|
||
if isinstance(result, dict):
|
||
rows = result.get("result", []) or []
|
||
elif isinstance(result, list):
|
||
rows = result
|
||
|
||
return {
|
||
"chart_id": chart_id,
|
||
"executed": True,
|
||
"duration_ms": duration_ms,
|
||
"row_count": len(rows),
|
||
"error": None,
|
||
}
|
||
except Exception as e:
|
||
duration_ms = int((time.time() - start) * 1000)
|
||
return {
|
||
"chart_id": chart_id,
|
||
"executed": False,
|
||
"duration_ms": duration_ms,
|
||
"row_count": None,
|
||
"error": str(e),
|
||
}
|
||
# endregion DatasetHealthChecker.check_chart_data
|
||
|
||
# region DatasetHealthChecker.check_dashboard_datasets [TYPE Function]
|
||
# @PURPOSE For every unique dataset in dashboard charts, check health.
|
||
# @PRE chart_list has chart dicts with slice_id and datasource_id.
|
||
# @POST Returns dict with datasets and optional chart_data lists.
|
||
async def check_dashboard_datasets(
|
||
self,
|
||
chart_list: list[dict],
|
||
execute_chart_data: bool = False,
|
||
) -> dict:
|
||
"""
|
||
Check all unique datasets referenced by dashboard charts.
|
||
|
||
Args:
|
||
chart_list: list of chart dicts with at least
|
||
{'slice_id', 'datasource_id', 'viz_type', 'params'}
|
||
execute_chart_data: if True, also execute chart queries (level 3-4)
|
||
|
||
Returns:
|
||
{datasets: [...], chart_data: [...]}
|
||
"""
|
||
# Collect unique datasource_ids
|
||
unique_ds_ids: set[int] = set()
|
||
for chart in chart_list:
|
||
ds_id = chart.get("datasource_id")
|
||
if ds_id is not None:
|
||
unique_ds_ids.add(int(ds_id))
|
||
|
||
# Check each dataset
|
||
dataset_results: list[dict] = []
|
||
for ds_id in sorted(unique_ds_ids):
|
||
result = await self.check_dataset_health(ds_id)
|
||
# Map affected charts
|
||
affected_charts = [
|
||
{"chart_id": c.get("slice_id"), "chart_name": c.get("slice_name", f"chart_{c.get('slice_id')}")}
|
||
for c in chart_list
|
||
if c.get("datasource_id") == ds_id
|
||
]
|
||
result["affected_charts"] = affected_charts
|
||
dataset_results.append(result)
|
||
|
||
# Optionally execute chart data
|
||
chart_data_results: list[dict] = []
|
||
if execute_chart_data:
|
||
for chart in chart_list:
|
||
chart_id = chart.get("slice_id")
|
||
params = chart.get("params", {})
|
||
if isinstance(params, str):
|
||
params = json.loads(params)
|
||
form_data = {
|
||
"slice_id": chart_id,
|
||
"viz_type": chart.get("viz_type", "table"),
|
||
"datasource_id": chart.get("datasource_id"),
|
||
"datasource_type": chart.get("datasource_type", "table"),
|
||
"granularity_sqla": params.get("granularity_sqla"),
|
||
"time_range": params.get("time_range", "Last 30 days"),
|
||
"metrics": params.get("metrics", []),
|
||
"groupby": params.get("groupby", []),
|
||
"adhoc_filters": params.get("adhoc_filters", []),
|
||
}
|
||
result = await self.check_chart_data(chart_id, form_data)
|
||
chart_data_results.append(result)
|
||
|
||
return {
|
||
"datasets": dataset_results,
|
||
"chart_data": chart_data_results,
|
||
}
|
||
# endregion DatasetHealthChecker.check_dashboard_datasets
|
||
# #endregion DatasetHealthChecker
|
||
|
||
# #region RedactionService [C:2] [TYPE Module]
|
||
# @BRIEF Redacts PII, credentials, and sensitive data from logs and LLM responses.
|
||
# @LAYER Service
|
||
# @RATIONALE FR-029: sensitive data must be filtered BEFORE external LLM send and BEFORE persistence.
|
||
class RedactionService:
|
||
"""Redacts PII, credentials, and sensitive data."""
|
||
|
||
# Common patterns to redact
|
||
PATTERNS = [
|
||
(r"password[=:]\s*\S+", "password=***"),
|
||
(r"secret[=:]\s*\S+", "secret=***"),
|
||
(r"token[=:]\s*\S+", "token=***"),
|
||
(r"api_key[=:]\s*\S+", "api_key=***"),
|
||
(r"apikey[=:]\s*\S+", "apikey=***"),
|
||
(r"Authorization:\s*\S+", "Authorization: ***"),
|
||
(r"Bearer\s+\S+\.\S+\.\S+", "Bearer ***"),
|
||
(r"[A-Za-z0-9+/=]{40,}", "***"), # base64 or long tokens
|
||
(r"[\w.+-]+@[\w-]+\.[\w.-]+", "***@***"), # emails
|
||
]
|
||
|
||
# region RedactionService.redact_logs [TYPE Function] [C:2]
|
||
# @BRIEF Redact PII/credentials from log lines.
|
||
# @PRE logs is a list of strings.
|
||
# @POST Returns redacted list preserving structure.
|
||
@staticmethod
|
||
def redact_logs(logs: list[str]) -> list[str]:
|
||
"""Redact PII/credentials from log lines."""
|
||
redacted = []
|
||
for line in logs:
|
||
for pattern, replacement in RedactionService.PATTERNS:
|
||
line = re.sub(pattern, replacement, line, flags=re.IGNORECASE)
|
||
redacted.append(line)
|
||
return redacted
|
||
# endregion RedactionService.redact_logs
|
||
|
||
# region RedactionService.redact_raw_response [TYPE Function] [C:2]
|
||
# @BRIEF Redact sensitive data from LLM raw response.
|
||
# @PRE raw is a string.
|
||
# @POST Returns redacted string.
|
||
@staticmethod
|
||
def redact_raw_response(raw: str) -> str:
|
||
"""Redact sensitive data from LLM raw response."""
|
||
for pattern, replacement in RedactionService.PATTERNS:
|
||
raw = re.sub(pattern, replacement, raw, flags=re.IGNORECASE)
|
||
return raw
|
||
# endregion RedactionService.redact_raw_response
|
||
# #endregion RedactionService
|
||
|
||
# #endregion LLMAnalysisService
|