rename klo to ktx

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

@@ -0,0 +1,172 @@
+"""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())