"""Command entry point for one-shot KTX daemon compute operations.""" from __future__ import annotations import argparse import json import sys 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 run_http_server( *, host: str, port: int, log_level: str, enable_code_execution: bool, ) -> None: import uvicorn from ktx_daemon.app import create_app uvicorn.run( create_app(enable_code_execution=enable_code_execution), host=host, port=port, log_level=log_level, ) 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: sys.stderr.write(f"{type(error).__name__}: {error}\n") return 1 if __name__ == "__main__": raise SystemExit(main())