Initial open-source release

python/klo-daemon/src/klo_daemon/app.py

@@ -0,0 +1,228 @@
+"""FastAPI app factory for the KLO daemon semantic compute server."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import Response
+
+from klo_daemon.code_execution import (
+    ExecuteCodeRequest,
+    ExecuteCodeResponse,
+    dumps_numpy_json,
+    execute_code_response,
+)
+from klo_daemon.database_introspection import (
+    DatabaseIntrospectionRequest,
+    DatabaseIntrospectionResponse,
+    introspect_database_response,
+)
+from klo_daemon.embeddings import (
+    ComputeEmbeddingBulkRequest,
+    ComputeEmbeddingBulkResponse,
+    ComputeEmbeddingRequest,
+    ComputeEmbeddingResponse,
+    EmbeddingProvider,
+    compute_embedding_bulk_response,
+    compute_embedding_response,
+)
+from klo_daemon.lookml import (
+    ParseLookMLRequest,
+    ParseLookMLResponse,
+    parse_lookml_project,
+)
+from klo_daemon.semantic_layer import (
+    SemanticLayerQueryRequest,
+    SemanticLayerQueryResponse,
+    ValidateSourcesRequest,
+    ValidateSourcesResponse,
+    query_semantic_layer,
+    validate_semantic_layer,
+)
+from klo_daemon.source_generation import (
+    GenerateSourcesRequest,
+    GenerateSourcesResponse,
+    generate_sources_response,
+)
+from klo_daemon.table_identifier import (
+    ParseTableIdentifierBatchRequest,
+    ParseTableIdentifierBatchResponse,
+    parse_table_identifier_response,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class NumpyORJSONResponse(Response):
+    media_type = "application/json"
+
+    def render(self, content: Any) -> bytes:
+        return dumps_numpy_json(content)
+
+
+def create_app(
+    *,
+    embedding_provider: EmbeddingProvider | None = None,
+    database_introspector: Callable[
+        [DatabaseIntrospectionRequest], DatabaseIntrospectionResponse
+    ]
+    | None = None,
+    enable_code_execution: bool = False,
+) -> FastAPI:
+    app = FastAPI(
+        title="KLO Daemon",
+        description="Stateless portable compute server for KLO.",
+        version="0.1.0",
+    )
+
+    @app.get("/health")
+    async def health() -> dict[str, str]:
+        return {"status": "healthy"}
+
+    @app.post("/database/introspect", response_model=DatabaseIntrospectionResponse)
+    async def database_introspect(
+        request: DatabaseIntrospectionRequest,
+    ) -> DatabaseIntrospectionResponse:
+        try:
+            introspector = database_introspector or introspect_database_response
+            return introspector(request)
+        except ValueError as error:
+            logger.warning("Database introspection rejected: %s", error)
+            raise HTTPException(status_code=400, detail=str(error)) from error
+        except Exception as error:
+            logger.exception("Database introspection failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Database introspection failed: {error}",
+            ) from error
+
+    @app.post("/embeddings/compute", response_model=ComputeEmbeddingResponse)
+    async def embedding_compute(
+        request: ComputeEmbeddingRequest,
+    ) -> ComputeEmbeddingResponse:
+        try:
+            return compute_embedding_response(
+                request,
+                provider=embedding_provider,
+            )
+        except ValueError as error:
+            logger.warning("Embedding compute rejected: %s", error)
+            raise HTTPException(status_code=400, detail=str(error)) from error
+        except Exception as error:
+            logger.exception("Embedding compute failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Embedding compute failed: {error}",
+            ) from error
+
+    @app.post(
+        "/embeddings/compute-bulk",
+        response_model=ComputeEmbeddingBulkResponse,
+    )
+    async def embedding_compute_bulk(
+        request: ComputeEmbeddingBulkRequest,
+    ) -> ComputeEmbeddingBulkResponse:
+        try:
+            return compute_embedding_bulk_response(
+                request,
+                provider=embedding_provider,
+            )
+        except ValueError as error:
+            logger.warning("Bulk embedding compute rejected: %s", error)
+            raise HTTPException(status_code=400, detail=str(error)) from error
+        except Exception as error:
+            logger.exception("Bulk embedding compute failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Bulk embedding compute failed: {error}",
+            ) from error
+
+    if enable_code_execution:
+
+        @app.post(
+            "/code/execute",
+            response_model=ExecuteCodeResponse,
+            response_class=NumpyORJSONResponse,
+        )
+        async def code_execute(request: ExecuteCodeRequest) -> ExecuteCodeResponse:
+            try:
+                return execute_code_response(
+                    request,
+                    nest_api_url=None,
+                    auth_header=None,
+                )
+            except Exception as error:
+                logger.exception("Code execution failed: %s", error)
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Code execution failed: {error}",
+                ) from error
+
+    @app.post("/lookml/parse", response_model=ParseLookMLResponse)
+    async def lookml_parse(request: ParseLookMLRequest) -> ParseLookMLResponse:
+        try:
+            return parse_lookml_project(request)
+        except Exception as error:
+            logger.exception("LookML parsing failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"LookML parsing failed: {error}",
+            ) from error
+
+    @app.post(
+        "/sql/parse-table-identifier",
+        response_model=ParseTableIdentifierBatchResponse,
+    )
+    async def sql_parse_table_identifier(
+        request: ParseTableIdentifierBatchRequest,
+    ) -> ParseTableIdentifierBatchResponse:
+        try:
+            return parse_table_identifier_response(request)
+        except Exception as error:
+            logger.exception("Table identifier parsing failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Table identifier parsing failed: {error}",
+            ) from error
+
+    @app.post(
+        "/semantic-layer/generate-sources", response_model=GenerateSourcesResponse
+    )
+    async def semantic_generate_sources(
+        request: GenerateSourcesRequest,
+    ) -> GenerateSourcesResponse:
+        try:
+            return generate_sources_response(request)
+        except Exception as error:
+            logger.exception("Semantic source generation failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Semantic source generation failed: {error}",
+            ) from error
+
+    @app.post("/semantic-layer/query", response_model=SemanticLayerQueryResponse)
+    async def semantic_query(
+        request: SemanticLayerQueryRequest,
+    ) -> SemanticLayerQueryResponse:
+        try:
+            return query_semantic_layer(request)
+        except ValueError as error:
+            logger.warning("Semantic query rejected: %s", error)
+            raise HTTPException(status_code=400, detail=str(error)) from error
+        except Exception as error:
+            logger.exception("Semantic query failed: %s", error)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Semantic layer query failed: {error}",
+            ) from error
+
+    @app.post("/semantic-layer/validate", response_model=ValidateSourcesResponse)
+    async def semantic_validate(
+        request: ValidateSourcesRequest,
+    ) -> ValidateSourcesResponse:
+        return validate_semantic_layer(request)
+
+    return app