ktx/python/ktx-daemon/src/ktx_daemon/__main__.py

"""Command entry point for one-shot KTX daemon compute operations."""

from __future__ import annotations

import argparse
import json
import sys
import time
from collections.abc import Callable
from types import TracebackType
from typing import Any

from pydantic import ValidationError

from ktx_daemon.code_execution import ExecuteCodeRequest, execute_code_response
from ktx_daemon.database_introspection import (
    DatabaseIntrospectionRequest,
    introspect_database_response,
)
from ktx_daemon.embeddings import (
    ComputeEmbeddingBulkRequest,
    ComputeEmbeddingRequest,
    compute_embedding_bulk_response,
    compute_embedding_response,
)
from ktx_daemon.lookml import ParseLookMLRequest, parse_lookml_project
from ktx_daemon.semantic_layer import (
    SemanticLayerQueryRequest,
    ValidateSourcesRequest,
    query_semantic_layer,
    validate_semantic_layer,
)
from ktx_daemon.source_generation import (
    GenerateSourcesRequest,
    generate_sources_response,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="ktx-daemon")
    subcommands = parser.add_subparsers(dest="command", required=True)
    subcommands.add_parser("semantic-query", help="Compile a semantic-layer query")
    subcommands.add_parser("semantic-validate", help="Validate semantic-layer sources")
    subcommands.add_parser(
        "semantic-generate-sources",
        help="Generate semantic-layer sources from schema scan data",
    )
    subcommands.add_parser(
        "database-introspect",
        help="Introspect a Postgres database schema",
    )
    subcommands.add_parser(
        "lookml-parse",
        help="Parse LookML files into KSL-ready structures",
    )
    subcommands.add_parser(
        "embedding-compute",
        help="Compute one local text embedding",
    )
    subcommands.add_parser(
        "embedding-compute-bulk",
        help="Compute local text embeddings in bulk",
    )
    subcommands.add_parser(
        "code-execute",
        help="Execute Python code with the current in-process boundary",
    )
    serve_http = subcommands.add_parser(
        "serve-http",
        help="Run the KTX daemon portable compute HTTP server",
    )
    serve_http.add_argument("--host", default="127.0.0.1")
    serve_http.add_argument("--port", type=int, default=8765)
    serve_http.add_argument(
        "--log-level",
        default="info",
        choices=["critical", "error", "warning", "info", "debug", "trace"],
    )
    serve_http.add_argument(
        "--enable-code-execution",
        action="store_true",
        help="Expose POST /code/execute on the HTTP server",
    )
    return parser


def _read_stdin_json() -> dict[str, Any]:
    raw = sys.stdin.read()
    parsed = json.loads(raw)
    if not isinstance(parsed, dict):
        raise ValueError("stdin JSON must be an object")
    return parsed


def install_serve_http_exception_hooks(started_at: float) -> Callable[[], None]:
    original_hook = sys.excepthook

    def hook(
        exc_type: type[BaseException],
        exc: BaseException,
        tb: TracebackType | None,
    ) -> None:
        report_serve_http_crash(exc, started_at=started_at)
        original_hook(exc_type, exc, tb)

    sys.excepthook = hook

    def dispose() -> None:
        sys.excepthook = original_hook

    return dispose


def report_serve_http_crash(error: BaseException, *, started_at: float) -> None:
    from ktx_daemon.telemetry import report_exception
    from ktx_daemon.telemetry.daemon_lifecycle import emit_daemon_stopped_once

    report_exception(
        error,
        source="serve-http",
        handled=False,
        fatal=True,
    )
    emit_daemon_stopped_once(
        reason="crash",
        uptime_ms=max(0, (time.perf_counter() - started_at) * 1000),
    )


def run_http_server(
    *,
    host: str,
    port: int,
    log_level: str,
    enable_code_execution: bool,
) -> None:
    import uvicorn

    from ktx_daemon.app import create_app

    started_at = time.perf_counter()
    dispose_hooks = install_serve_http_exception_hooks(started_at)
    try:
        try:
            uvicorn.run(
                create_app(
                    enable_code_execution=enable_code_execution,
                    telemetry_started_at=started_at,
                ),
                host=host,
                port=port,
                log_level=log_level,
            )
        except Exception as error:
            report_serve_http_crash(error, started_at=started_at)
            raise
    finally:
        dispose_hooks()


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "serve-http":
        run_http_server(
            host=args.host,
            port=args.port,
            log_level=args.log_level,
            enable_code_execution=args.enable_code_execution,
        )
        return 0

    try:
        payload = _read_stdin_json()
        if args.command == "semantic-query":
            response = query_semantic_layer(
                SemanticLayerQueryRequest.model_validate(payload)
            )
        elif args.command == "semantic-validate":
            response = validate_semantic_layer(
                ValidateSourcesRequest.model_validate(payload)
            )
        elif args.command == "semantic-generate-sources":
            response = generate_sources_response(
                GenerateSourcesRequest.model_validate(payload)
            )
        elif args.command == "database-introspect":
            response = introspect_database_response(
                DatabaseIntrospectionRequest.model_validate(payload)
            )
        elif args.command == "lookml-parse":
            response = parse_lookml_project(ParseLookMLRequest.model_validate(payload))
        elif args.command == "embedding-compute":
            response = compute_embedding_response(
                ComputeEmbeddingRequest.model_validate(payload)
            )
        elif args.command == "embedding-compute-bulk":
            response = compute_embedding_bulk_response(
                ComputeEmbeddingBulkRequest.model_validate(payload)
            )
        elif args.command == "code-execute":
            response = execute_code_response(
                ExecuteCodeRequest.model_validate(payload),
                nest_api_url=None,
                auth_header=None,
            )
        else:
            parser.error(f"Unknown command: {args.command}")
            return 2
        sys.stdout.write(response.model_dump_json() + "\n")
        return 0
    except (json.JSONDecodeError, ValidationError, ValueError) as error:
        sys.stderr.write(f"{error}\n")
        return 1
    except Exception as error:
        from ktx_daemon.telemetry import report_exception

        report_exception(
            error,
            source=str(args.command),
            handled=True,
            fatal=False,
        )
        sys.stderr.write(f"{type(error).__name__}: {error}\n")
        return 1


if __name__ == "__main__":
    raise SystemExit(main())