ss-tools/backend/src/agent/tools.py

# backend/src/agent/tools.py
# #region AgentChat.Tools [C:4] [TYPE Module] [SEMANTICS agent-chat,tools,langchain]
# @DEFGROUP AgentChat Native LangChain @tool functions.
# @REJECTED StructuredTool wrapping — native @tool is the single source of truth.
# @REJECTED StructuredTool wrapping — native @tool is the single source of truth.
# @RATIONALE LangChain @tool decorator chosen over direct FastAPI calls for LangGraph compatibility — tools are auto-registered in the agent's tool-calling loop.

import json
import os
from typing import Any

import httpx
from langchain_core.tools import tool
from pydantic import BaseModel, Field

from src.agent.context import get_service_jwt, get_user_jwt

FASTAPI_URL = os.getenv("FASTAPI_URL", "http://localhost:8000")
TOOL_RESPONSE_LIMIT = 4000


def _dual_auth_headers() -> dict[str, str]:
    """Build dual-identity auth headers for tool→FastAPI calls per FR-007/FR-019.

    Authorization: Bearer {service_jwt} — authenticates the agent container.
    X-User-JWT: {user_jwt} — authorizes the operation as the end user.
    Falls back to user JWT in Authorization if service JWT unavailable.
    """
    user_jwt = get_user_jwt() or ""
    svc_jwt = get_service_jwt() or os.getenv("SERVICE_JWT", "")
    headers = {}
    if svc_jwt:
        headers["Authorization"] = f"Bearer {svc_jwt}"
        if user_jwt:
            headers["X-User-JWT"] = user_jwt
    elif user_jwt:
        headers["Authorization"] = f"Bearer {user_jwt}"
    return headers


def _trim_response(text: str, limit: int = TOOL_RESPONSE_LIMIT) -> str:
    if len(text) <= limit:
        return text
    return f"{text[:limit]}\n... response truncated ..."


async def _get(path: str, params: dict[str, Any] | None = None) -> httpx.Response:
    async with httpx.AsyncClient() as client:
        return await client.get(
            f"{FASTAPI_URL}{path}",
            params=params,
            headers=_dual_auth_headers(),
        )


async def _post(
    path: str,
    payload: dict[str, Any] | None = None,
    params: dict[str, Any] | None = None,
) -> httpx.Response:
    async with httpx.AsyncClient() as client:
        return await client.post(
            f"{FASTAPI_URL}{path}",
            json=payload or {},
            params=params,
            headers=_dual_auth_headers(),
        )


def _api_result(resp: httpx.Response, ok_statuses: set[int] | None = None) -> str:
    ok_statuses = ok_statuses or {200, 201, 202}
    if resp.status_code not in ok_statuses:
        return f"Error {resp.status_code}: {resp.text}"
    return _trim_response(resp.text)


async def _resolve_active_provider_id(provider_id: str | None = None) -> str | None:
    if provider_id:
        return provider_id
    resp = await _get("/api/llm/providers")
    if resp.status_code != 200:
        return None
    try:
        providers = resp.json()
    except json.JSONDecodeError:
        return None
    active = next((p for p in providers if p.get("is_active")), None)
    if active:
        return str(active.get("id"))
    if providers:
        return str(providers[0].get("id"))
    return None


def _dashboard_ids(dashboard_id: int | str | None, dashboard_ids: list[int] | list[str] | None) -> list[int]:
    values: list[int] = []
    for raw in dashboard_ids or []:
        values.append(int(raw))
    if dashboard_id is not None and str(dashboard_id).strip():
        values.append(int(dashboard_id))
    return list(dict.fromkeys(values))


# ── Tool: search_dashboards ──
class SearchDashboardsInput(BaseModel):
    query: str = Field(description="Search query for dashboard name")
    env_id: str | None = Field(default=None, description="Environment ID (e.g. 'prod', 'ss-dev')")


# @ingroup AgentChat
# @PRE User authenticated via dual-identity JWT
# @POST Returns JSON string result from FastAPI
# @SIDE_EFFECT HTTP call to FastAPI backend

@tool(args_schema=SearchDashboardsInput)
async def search_dashboards(query: str, env_id: str | None = None) -> str:
    """Search and list dashboards by name, with optional environment filter.
    Pass env_id like 'prod', 'ss-dev', or 'ss-preprod' to filter by environment.
    Returns a formatted list of dashboards with their details.
    """
    params = {"q": query, "env_id": env_id or ""}
    resp = await _get("/api/dashboards", params=params)
    if resp.status_code != 200:
        return f"Error {resp.status_code}: Unable to fetch dashboards. {resp.text}"

    try:
        data = resp.json()
        dashboards = data.get("dashboards", [])
        total = data.get("total", 0)

        if total == 0:
            return f"No dashboards found matching '{query}' in environment '{env_id or 'default'}'."

        lines = [f"Found {total} dashboard(s) in environment '{env_id or 'default'}':"]
        for db in dashboards:
            title = db.get("title", "Untitled")
            owners = ", ".join(db.get("owners", [])) or "N/A"
            modified = db.get("last_modified", "")[:10] if db.get("last_modified") else "unknown"
            lines.append(f"  - {title}")
            lines.append(f"    Owners: {owners}")
            lines.append(f"    Last modified: {modified}")
        return "\n".join(lines)
    except (json.JSONDecodeError, KeyError) as e:
        return f"Could not parse dashboards response: {e}"


# ── Tool: get_health_summary ──
class HealthSummaryInput(BaseModel):
    env_id: str | None = Field(default=None, description="Environment ID (e.g. 'prod', 'ss-dev')")


# @ingroup AgentChat
# @PRE User authenticated via dual-identity JWT
# @POST Returns JSON string result from FastAPI
# @SIDE_EFFECT HTTP call to FastAPI backend
@tool(args_schema=HealthSummaryInput)
async def get_health_summary(env_id: str | None = None) -> str:
    """Get system health summary — dashboard validation status, recent failures.
    Pass env_id like 'prod', 'ss-dev', or 'ss-preprod' to filter by environment.
    """
    resp = await _get("/api/health/summary", params={"environment_id": env_id or ""})
    if resp.status_code != 200:
        return f"Health check failed: HTTP {resp.status_code}"
    return _trim_response(resp.text, 2000)


# ── Tool: list_environments ──
# @ingroup AgentChat
# @PRE User authenticated via dual-identity JWT
# @POST Returns JSON string result from FastAPI
# @SIDE_EFFECT HTTP call to FastAPI backend
@tool
async def list_environments() -> str:
    """List configured deployment environments."""
    resp = await _get("/api/settings/environments")
    return _api_result(resp)


# ── Tool: get_task_status ──
# @ingroup AgentChat
# @PRE User authenticated via dual-identity JWT
# @POST Returns JSON string result from FastAPI
# @SIDE_EFFECT HTTP call to FastAPI backend
@tool
async def get_task_status(task_id: str) -> str:
    """Check the status of a background task by its task_id."""
    resp = await _get(f"/api/tasks/{task_id}")
    return _api_result(resp)


@tool
async def show_capabilities() -> str:
    """Show all tools available to the agent chat."""
    lines = ["Available LangChain tools:"]
    for tool_obj in get_all_tools():
        lines.append(f"- {tool_obj.name}: {tool_obj.description or ''}")
    return "\n".join(lines)


@tool
async def list_llm_providers() -> str:
    """List configured LLM providers."""
    resp = await _get("/api/llm/providers")
    return _api_result(resp)


@tool
async def get_llm_status() -> str:
    """Check whether the LLM runtime is configured and usable."""
    resp = await _get("/api/llm/status")
    return _api_result(resp)


class RunBackupInput(BaseModel):
    environment_id: str = Field(description="Environment ID for the backup")
    dashboard_id: int | None = Field(default=None, description="Optional single dashboard ID")
    dashboard_ids: list[int] | None = Field(default=None, description="Optional list of dashboard IDs")


@tool(args_schema=RunBackupInput)
async def run_backup(
    environment_id: str,
    dashboard_id: int | None = None,
    dashboard_ids: list[int] | None = None,
) -> str:
    """Run a Superset backup for an environment, optionally scoped to dashboard IDs."""
    params: dict[str, Any] = {"environment_id": environment_id}
    ids = _dashboard_ids(dashboard_id, dashboard_ids)
    if ids:
        params["dashboard_ids"] = ids
    resp = await _post("/api/tasks", {"plugin_id": "superset-backup", "params": params})
    return _api_result(resp, {201})


class ExecuteMigrationInput(BaseModel):
    source_env_id: str = Field(description="Source environment ID")
    target_env_id: str = Field(description="Target environment ID")
    selected_ids: list[int] = Field(description="Dashboard IDs to migrate")
    replace_db_config: bool = Field(default=False, description="Replace DB configuration during migration")
    fix_cross_filters: bool = Field(default=True, description="Fix cross filters during migration")


@tool(args_schema=ExecuteMigrationInput)
async def execute_migration(
    source_env_id: str,
    target_env_id: str,
    selected_ids: list[int],
    replace_db_config: bool = False,
    fix_cross_filters: bool = True,
) -> str:
    """Execute dashboard migration between two environments."""
    payload = {
        "source_env_id": source_env_id,
        "target_env_id": target_env_id,
        "selected_ids": selected_ids,
        "replace_db_config": replace_db_config,
        "fix_cross_filters": fix_cross_filters,
    }
    resp = await _post("/api/migration/execute", payload)
    return _api_result(resp)


class GitDashboardInput(BaseModel):
    dashboard_ref: str = Field(description="Dashboard ID, slug, or title resolvable by the Git API")
    env_id: str | None = Field(default=None, description="Optional environment ID used to resolve dashboard_ref")


class CreateBranchInput(GitDashboardInput):
    branch_name: str = Field(description="New branch name")
    from_branch: str = Field(default="dev", description="Source branch")


@tool(args_schema=CreateBranchInput)
async def create_branch(
    dashboard_ref: str,
    branch_name: str,
    from_branch: str = "dev",
    env_id: str | None = None,
) -> str:
    """Create a branch in a dashboard Git repository."""
    resp = await _post(
        f"/api/git/repositories/{dashboard_ref}/branches",
        {"name": branch_name, "from_branch": from_branch},
        params={"env_id": env_id} if env_id else None,
    )
    return _api_result(resp)


class CommitChangesInput(GitDashboardInput):
    message: str = Field(description="Commit message")
    files: list[str] = Field(default_factory=list, description="Files to stage and commit; empty means backend default")


@tool(args_schema=CommitChangesInput)
async def commit_changes(
    dashboard_ref: str,
    message: str,
    files: list[str] | None = None,
    env_id: str | None = None,
) -> str:
    """Stage and commit changes in a dashboard Git repository."""
    resp = await _post(
        f"/api/git/repositories/{dashboard_ref}/commit",
        {"message": message, "files": files or []},
        params={"env_id": env_id} if env_id else None,
    )
    return _api_result(resp)


class DeployDashboardInput(GitDashboardInput):
    environment_id: str = Field(description="Target deployment environment ID")


@tool(args_schema=DeployDashboardInput)
async def deploy_dashboard(
    dashboard_ref: str,
    environment_id: str,
    env_id: str | None = None,
) -> str:
    """Deploy a dashboard from Git to a target environment."""
    resp = await _post(
        f"/api/git/repositories/{dashboard_ref}/deploy",
        {"environment_id": environment_id},
        params={"env_id": env_id} if env_id else None,
    )
    return _api_result(resp)


class RunLlmDocumentationInput(BaseModel):
    dataset_id: str = Field(description="Dataset ID to document")
    environment_id: str = Field(description="Environment ID")
    provider_id: str | None = Field(default=None, description="Optional LLM provider ID; active provider is used if omitted")


@tool(args_schema=RunLlmDocumentationInput)
async def run_llm_documentation(
    dataset_id: str,
    environment_id: str,
    provider_id: str | None = None,
) -> str:
    """Generate dataset documentation via the LLM documentation task."""
    resolved_provider = await _resolve_active_provider_id(provider_id)
    params = {
        "dataset_id": str(dataset_id),
        "environment_id": environment_id,
    }
    if resolved_provider:
        params["provider_id"] = resolved_provider
    resp = await _post("/api/tasks", {"plugin_id": "llm_documentation", "params": params})
    return _api_result(resp, {201})


class RunLlmValidationInput(BaseModel):
    environment_id: str = Field(description="Environment ID")
    dashboard_id: int | None = Field(default=None, description="Single dashboard ID")
    dashboard_ids: list[int] | None = Field(default=None, description="Dashboard IDs to validate")
    dashboard_url: str | None = Field(default=None, description="Optional full Superset dashboard URL")
    provider_id: str | None = Field(default=None, description="Optional LLM provider ID; active provider is used if omitted")
    name: str | None = Field(default=None, description="Optional validation policy name")
    policy_id: str | None = Field(default=None, description="Existing validation policy ID to run immediately")


@tool(args_schema=RunLlmValidationInput)
async def run_llm_validation(
    environment_id: str,
    dashboard_id: int | None = None,
    dashboard_ids: list[int] | None = None,
    dashboard_url: str | None = None,
    provider_id: str | None = None,
    name: str | None = None,
    policy_id: str | None = None,
) -> str:
    """Create and immediately run an LLM dashboard validation policy, or run an existing policy."""
    if policy_id:
        resp = await _post(f"/api/validation-tasks/{policy_id}/run")
        return _api_result(resp)

    ids = _dashboard_ids(dashboard_id, dashboard_ids)
    sources = []
    if dashboard_url:
        sources.append({"type": "dashboard_url", "value": dashboard_url})
    if not ids and not sources:
        return "Error: provide dashboard_id, dashboard_ids, dashboard_url, or policy_id."

    resolved_provider = await _resolve_active_provider_id(provider_id)
    if not resolved_provider:
        return "Error: no LLM provider configured. Configure or pass provider_id."

    payload = {
        "name": name or f"Agent LLM validation - {environment_id}",
        "environment_id": environment_id,
        "dashboard_ids": [str(i) for i in ids],
        "provider_id": resolved_provider,
        "is_active": True,
        "screenshot_enabled": True,
        "logs_enabled": True,
        "execute_chart_data": False,
    }
    if sources:
        payload["sources"] = sources

    create_resp = await _post("/api/validation-tasks", payload)
    if create_resp.status_code != 201:
        return _api_result(create_resp, {201})
    try:
        created = create_resp.json()
        created_id = created["id"]
    except (json.JSONDecodeError, KeyError) as exc:
        return f"Validation policy created but response could not be parsed: {exc}. Raw: {create_resp.text}"
    run_resp = await _post(f"/api/validation-tasks/{created_id}/run")
    return _api_result(run_resp)


@tool
async def list_maintenance_events() -> str:
    """List active and completed maintenance events."""
    resp = await _get("/api/maintenance/events")
    return _api_result(resp)


class StartMaintenanceInput(BaseModel):
    tables: list[str] = Field(description="Affected table names")
    environment_id: str = Field(description="Target environment ID")
    start_time: str = Field(description="ISO datetime for maintenance start")
    end_time: str | None = Field(default=None, description="Optional ISO datetime for maintenance end")
    message: str | None = Field(default=None, description="Optional banner message")


@tool(args_schema=StartMaintenanceInput)
async def start_maintenance(
    tables: list[str],
    environment_id: str,
    start_time: str,
    end_time: str | None = None,
    message: str | None = None,
) -> str:
    """Start a maintenance event and apply banners to affected dashboards."""
    payload = {
        "tables": tables,
        "environment_id": environment_id,
        "start_time": start_time,
        "end_time": end_time,
        "message": message,
    }
    resp = await _post("/api/maintenance/start", payload)
    return _api_result(resp, {200, 202})


class EndMaintenanceInput(BaseModel):
    maintenance_id: str | None = Field(default=None, description="Maintenance event ID; omit when end_all=True")
    environment_id: str | None = Field(default=None, description="Optional environment scope for end_all")
    end_all: bool = Field(default=False, description="End all active maintenance events")


@tool(args_schema=EndMaintenanceInput)
async def end_maintenance(
    maintenance_id: str | None = None,
    environment_id: str | None = None,
    end_all: bool = False,
) -> str:
    """End one maintenance event, or end all active events when end_all is true."""
    if end_all:
        payload = {"environment_id": environment_id} if environment_id else {}
        resp = await _post("/api/maintenance/end-all", payload)
        return _api_result(resp, {202})
    if not maintenance_id:
        return "Error: maintenance_id is required unless end_all=true."
    resp = await _post(f"/api/maintenance/{maintenance_id}/end")
    return _api_result(resp, {202})


# ── All available tools for the agent ──
def get_all_tools() -> list:
    return [
        show_capabilities,
        search_dashboards,
        get_health_summary,
        list_environments,
        get_task_status,
        list_llm_providers,
        get_llm_status,
        create_branch,
        commit_changes,
        deploy_dashboard,
        execute_migration,
        run_backup,
        run_llm_validation,
        run_llm_documentation,
        list_maintenance_events,
        start_maintenance,
        end_maintenance,
    ]


def get_tools_for_query(query: str, *, prefetch_available: bool = False) -> list:
    """Return a compact, intent-scoped tool set to keep small-context models usable."""
    text = (query or "").lower()
    selected = [show_capabilities]
    matched_intent = False

    if any(word in text for word in ["инструмент", "tool", "capabilit", "умеешь", "можешь"]):
        return selected

    if any(word in text for word in ["дашборд", "dashboard", "dashboards", "дашборды"]):
        matched_intent = True
        if not prefetch_available:
            selected.append(search_dashboards)
    if any(word in text for word in ["здоров", "health", "статус системы", "system status"]):
        matched_intent = True
        selected.append(get_health_summary)
    if any(word in text for word in ["окруж", "environment", "env"]):
        matched_intent = True
        selected.append(list_environments)
    if any(word in text for word in ["задач", "task", "таск"]):
        matched_intent = True
        selected.append(get_task_status)
    if any(word in text for word in ["llm", "provider", "провайдер", "модель"]):
        matched_intent = True
        selected.extend([list_llm_providers, get_llm_status])
    if any(word in text for word in ["branch", "ветк"]):
        matched_intent = True
        selected.append(create_branch)
    if any(word in text for word in ["commit", "коммит"]):
        matched_intent = True
        selected.append(commit_changes)
    if any(word in text for word in ["deploy", "депло", "разверн"]):
        matched_intent = True
        selected.append(deploy_dashboard)
    if any(word in text for word in ["миграц", "migration", "migrate"]):
        matched_intent = True
        selected.append(execute_migration)
    if any(word in text for word in ["backup", "бэкап", "резерв"]):
        matched_intent = True
        selected.append(run_backup)
    if any(word in text for word in ["валидац", "validation", "validate"]):
        matched_intent = True
        selected.append(run_llm_validation)
    if any(word in text for word in ["документ", "documentation", "docs"]):
        matched_intent = True
        selected.append(run_llm_documentation)
    if any(word in text for word in ["maintenance", "обслуж", "баннер"]):
        matched_intent = True
        selected.extend([list_maintenance_events, start_maintenance, end_maintenance])

    if len(selected) == 1 and not matched_intent:
        selected.extend([search_dashboards, get_health_summary, list_environments, get_task_status])

    unique = {}
    for tool_obj in selected:
        unique[tool_obj.name] = tool_obj
    return list(unique.values())
# #endregion AgentChat.Tools