wombatu-kun commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3567769673
########## hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py: ########## @@ -0,0 +1,203 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""POST /v1/chat: the agent-loop inference endpoint (JSON or SSE).""" + +from __future__ import annotations + +import logging +import uuid +from collections.abc import AsyncIterator +from typing import Any + +from fastapi import APIRouter, HTTPException, Request +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage +from langgraph.errors import GraphRecursionError +from sse_starlette.sse import EventSourceResponse, ServerSentEvent + +from hudi_agent_gateway.api.models import ( + ChatRequest, + ChatResponse, + DoneEvent, + ErrorEvent, + TokenEvent, + ToolCallEvent, + ToolResultEvent, + ToolTraceEntry, +) +from hudi_agent_gateway.log import log_event, request_id_var, session_id_var + +logger = logging.getLogger("hudi_agent_gateway.chat") + + +def _message_text(msg: Any) -> str: + text = getattr(msg, "text", None) + if callable(text): # langchain-core < 1.x compatibility + return text() + if isinstance(text, str): + return text + return str(getattr(msg, "content", "")) + + +router = APIRouter() + +_PREVIEW_CHARS = 500 + + +def _agent_config(request: Request, session_id: str) -> dict[str, Any]: + settings = request.app.state.settings + return { + "configurable": {"thread_id": session_id}, + "recursion_limit": settings.agent_max_iterations, + } + + +def _tool_trace_from_messages(messages: list[Any]) -> list[ToolTraceEntry]: + """Reconstruct the tool trace from the messages appended this turn.""" + calls: dict[str, ToolTraceEntry] = {} + trace: list[ToolTraceEntry] = [] + for msg in messages: + if isinstance(msg, AIMessage): + for tc in msg.tool_calls or []: + entry = ToolTraceEntry(name=tc["name"], args=tc.get("args") or {}) + calls[tc.get("id") or ""] = entry + trace.append(entry) + elif isinstance(msg, ToolMessage): + matched = calls.get(msg.tool_call_id or "") + if matched is not None: + content = msg.content if isinstance(msg.content, str) else str(msg.content) + matched.result_preview = content[:_PREVIEW_CHARS] + matched.is_error = '"error"' in content[:200] Review Comment: Substring matching on the serialized payload gives false positives. `shape_result` puts the echoed SQL and then the column names at the front of the payload, so any successful result with a column named `error` (or a `describe_table` on such a table) contains `"error"` within the first 200 chars and is reported as a failed tool call, both in `tool_trace` here and in the SSE `tool_result` event in `_updates_to_events`. `_with_invocation_logging` in `registry.py` already derives the same fact correctly by `json.loads`-ing the result and checking for a top-level `error` key. Reuse that instead of the substring test. ########## hudi-agent-gateway/README.md: ########## @@ -0,0 +1,129 @@ +<!-- +Licensed to the Apache Software Foundation (ASF) under one +or more contributor license agreements. See the NOTICE file +distributed with this work for additional information +regarding copyright ownership. The ASF licenses this file +to you under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance +with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +--> +# hudi-agent-gateway + +**One deployable service that serves your Hudi lakehouse to AI.** A single +process hosts three surfaces over the same set of guarded lakehouse tools: + +| Surface | Where | What | +|---|---|---| +| Agent chat API | `POST /v1/chat` | prompt in → LangGraph agent loop (model ↔ tools) → grounded answer out; multi-turn sessions; optional SSE streaming | +| MCP server | `/mcp` (streamable HTTP) | external agents (Claude, anything MCP) call the lakehouse tools directly | +| Chat UI | `/ui/` | first-party ChatGPT-style web UI (zero third-party code) | + +The v1 tools query the lakehouse through Trino: `query_lakehouse` (guarded, +read-only SQL), `list_tables`, `describe_table`. Every model-written query +passes AST-level guardrails (single statement, SELECT-only, row cap injected +as a real `LIMIT`) and every invocation is logged as structured JSON — the +seed of the gateway's trace collection. + +## Quickstart (local process) + +```bash +cd hudi-agent-gateway +python3.12 -m venv .venv && .venv/bin/pip install -e ".[dev]" + +# point at a Trino with the Hudi connector (e.g. the local-dev stack, port-forwarded): +GATEWAY_TRINO_HOST=localhost GATEWAY_TRINO_PORT=18080 \ +GATEWAY_LLM_PROVIDER=anthropic GATEWAY_LLM_MODEL=claude-haiku-4-5-20251001 \ + .venv/bin/hudi-agent-gateway serve + +curl -X POST localhost:8000/v1/chat -H 'Content-Type: application/json' \ + -d '{"message": "How many trips per city?", "session_id": "s1"}' +open http://localhost:8000/ui/ +``` + +`GET /v1/models` lists the models the configured provider offers (live: +Anthropic and OpenAI model APIs, Ollama's local tags, vLLM's served models), +and `POST /v1/chat` accepts an optional `"model"` to pick one per request — +the chat UI exposes this as a model picker. Sessions survive model switches. + +For a fully local model, install [Ollama](https://ollama.com), pull a +tool-capable model, and use the default provider: + +```bash +ollama pull qwen3:8b +GATEWAY_LLM_PROVIDER=ollama GATEWAY_LLM_MODEL=qwen3:8b hudi-agent-gateway serve +``` + +Connect an MCP client: + +```bash +claude mcp add --transport http hudi-lakehouse http://localhost:8000/mcp/ +``` + +## Deploying on Kubernetes + +See `hudi-lakehouse/charts/hudi-agent-gateway` — the product Helm chart — +and `hudi-lakehouse/local-dev/` for a complete laptop environment +(MinIO + Hive Metastore + Trino + this gateway) where the gateway is +installed alongside Trino by default. + +## Configuration + +Environment variables (prefix `GATEWAY_` except the standard key names): + +| Variable | Default | Purpose | +|---|---|---| +| `GATEWAY_LLM_PROVIDER` | `ollama` | `anthropic` \| `openai` \| `ollama` \| `openai-compatible` | +| `GATEWAY_LLM_MODEL` | `qwen3:8b` | model name for the provider | +| `ANTHROPIC_API_KEY` / `OPENAI_API_KEY` | — | required by the matching provider | +| `GATEWAY_OLLAMA_BASE_URL` | `http://localhost:11434` | Ollama endpoint | +| `GATEWAY_OPENAI_BASE_URL` | — | endpoint for `openai-compatible` (vLLM, Together, …) | +| `GATEWAY_LLM_TIMEOUT_SECONDS` | `120` | per-model-call timeout | +| `GATEWAY_TRINO_HOST` / `_PORT` | `hudi-trino.hudi-lakehouse.svc` / `8080` | Trino coordinator | +| `GATEWAY_TRINO_CATALOG` / `_SCHEMA` / `_USER` | `hudi` / `default` / `hudi-agent-gateway` | query defaults | +| `GATEWAY_SQL_ROW_CAP` | `200` | LIMIT enforced on every query | +| `GATEWAY_SQL_TIMEOUT_SECONDS` | `120` | per-query timeout | +| `GATEWAY_TOOL_RESULT_MAX_BYTES` | `50000` | tool results truncated beyond this (with notice) | +| `GATEWAY_AGENT_MAX_ITERATIONS` | `25` | agent loop recursion limit | +| `GATEWAY_SESSION_TTL_SECONDS` / `GATEWAY_MAX_SESSIONS` | `3600` / `1000` | session store bounds | +| `GATEWAY_MAX_MESSAGES_PER_SESSION` | `40` | context window per session (trimmed pre-model) | +| `GATEWAY_SYSTEM_PROMPT_EXTRA` | — | appended to the built-in system prompt | +| `GATEWAY_MCP_ENABLED` | `true` | serve /mcp | +| `GATEWAY_HOST` / `GATEWAY_PORT` / `GATEWAY_LOG_LEVEL` | `0.0.0.0` / `8000` / `INFO` | server basics | + +Startup never depends on the LLM or Trino being reachable: `/health` is +liveness, `/ready` reports per-dependency status (and gates the Kubernetes +readiness probe). + +## Development + +```bash +.venv/bin/pytest # offline suite (fake Trino + scripted model) +.venv/bin/ruff check src tests +.venv/bin/mypy src + +# live integration (against a port-forwarded local-dev stack): +GATEWAY_IT_TRINO_HOST=localhost GATEWAY_IT_TRINO_PORT=18080 .venv/bin/pytest tests/integration +``` + +Adding a tool: write a module under `src/hudi_agent_gateway/tools/` with a +`register(registry, ...)` function and call it from +`tools/__init__.py:build_registry`. One registration exposes it to the agent +loop, the MCP server, `GET /v1/tools`, and the invocation log. + +## Design notes & limits (v1) Review Comment: The list here covers single-replica and read-only, but not that the gateway has no authentication at all. `/v1/chat` spends LLM tokens, `/mcp` executes the lakehouse tools for any caller, and `/v1/models` exercises the provider API key, and none of them is authenticated. `host` defaults to `0.0.0.0` in `GatewaySettings`, and `service.type` in the chart is one `--set` away from `LoadBalancer`. Worth stating here as an explicit v1 limit (deploy behind an authenticating proxy or on a trusted network), so an operator does not have to infer it from the absence of code. ########## hudi-agent-gateway/src/hudi_agent_gateway/api/meta.py: ########## @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Operational endpoints: /health, /ready, /v1/tools, /v1/info.""" + +from __future__ import annotations + +import time + +from fastapi import APIRouter, Request, Response + +from hudi_agent_gateway import __version__ +from hudi_agent_gateway.api.models import ( + DependencyStatus, + InfoResponse, + ModelsResponse, + ReadyResponse, +) +from hudi_agent_gateway.llm import list_models + +router = APIRouter() + + [email protected]("/health") +async def health() -> dict[str, str]: + """Liveness: the process is up and serving.""" + return {"status": "ok", "version": __version__} + + [email protected]("/ready", response_model=ReadyResponse) +async def ready(request: Request, response: Response) -> ReadyResponse: + """Readiness: configuration is valid AND dependencies are reachable.""" + state = request.app.state + checks: dict[str, DependencyStatus] = {} + + trino_ok = await state.trino_client.ping() Review Comment: `/ready` runs a live, uncached `SELECT 1` against Trino on every probe, and `TrinoClient._connect` opens a fresh connection (and a fresh worker thread) each time, while the LLM check on the next line is cached for 10s in `LLMReadiness`. With `periodSeconds: 10` on the chart's readinessProbe that is a Trino query and a new connection every ten seconds, forever, all of it landing in the coordinator's query history. The probe also sets no `timeoutSeconds`, so Kubernetes applies its 1s default to a path whose own budget is 5s (`ping`'s default). Any Trino latency above one second fails the probe, and with `failureThreshold: 3` a healthy gateway is pulled out of the Service endpoints. Cache the Trino ping the way the LLM check is cached, and set an explicit `timeoutSeconds` on the probe. ########## scripts/release/validate_source_copyright.sh: ########## @@ -50,7 +50,7 @@ echo "Performing custom Licensing Check " # Exclude the 'hudi-trino-plugin' directory. Its license checks are handled by airlift: # https://github.com/airlift/airbase/blob/823101482dbc60600d7862f0f5c93aded6190996/airbase/pom.xml#L1239 # --- -numfilesWithNoLicense=$(find . -path './hudi-trino-plugin' -prune -o -type f -iname '*' | grep -v './hudi-trino-plugin' | grep -v NOTICE | grep -v LICENSE | grep -v '.jpg' | grep -v '.json' | grep -v '.zip' | grep -v '.hfile' | grep -v '.data' | grep -v '.commit' | grep -v emptyFile | grep -v DISCLAIMER | grep -v '.sqltemplate' | grep -v KEYS | grep -v '.mailmap' | grep -v 'banner.txt' | grep -v '.txt' | grep -v "fixtures" | xargs grep -L "Licensed to the Apache Software Foundation (ASF)") +numfilesWithNoLicense=$(find . -path './hudi-trino-plugin' -prune -o -type f -iname '*' | grep -v './hudi-trino-plugin' | grep -v NOTICE | grep -v LICENSE | grep -v '.jpg' | grep -v '.png' | grep -v '.json' | grep -v '.zip' | grep -v '.hfile' | grep -v '.data' | grep -v '.commit' | grep -v emptyFile | grep -v DISCLAIMER | grep -v '.sqltemplate' | grep -v KEYS | grep -v '.mailmap' | grep -v 'banner.txt' | grep -v '.txt' | grep -v "fixtures" | xargs grep -L "Licensed to the Apache Software Foundation (ASF)") Review Comment: The PR description says the only edit to an existing file is `**/py.typed` in the root `pom.xml` rat excludes. That is not what the diff does: `pom.xml` is unchanged (commit afa0dc0d reverted it) and this line is the PR's single change to a pre-existing file. Worth correcting in the body, since a reviewer scoping "does this touch release infra?" currently gets the wrong answer. On the change itself: this permanently exempts every `*.png` in the repository, present and future, from the source-release ASF-header check, to accommodate one 9.7 KB logo (the pattern is unanchored too, so `.` matches any character). The project's established response to "my images trip the license check" has been to keep the images out of the source tarball instead: #5273 relocated a stray png, and #14023 (hudi-notebooks, the closest analog - a new top-level non-Java module) and #18939 both added an rsync exclude to `create_source_directory.sh` in the same commit. That remedy does not transfer here, because unlike rfc/ and the notebook images this logo is a runtime asset the UI serves, so excluding it would ship a source release with a broken image. The clean way out is to drop the binary: inline the logo as SVG, or as a data-URI in `style.css`. It is text, it carries an ASF header like every other UI file, it passes the existing check unmodified, and it needs no change to any release script. ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py: ########## @@ -0,0 +1,96 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Async wrapper over the (synchronous) Trino DB-API client.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + +import trino + +from hudi_agent_gateway.config import GatewaySettings + + +class TrinoQueryError(Exception): + """A query failed on the Trino side; message is safe to surface to the model.""" + + +class TrinoTimeoutError(TrinoQueryError): + pass + + +@dataclass +class QueryResult: + columns: list[str] + rows: list[list[Any]] + row_count: int = field(init=False) + + def __post_init__(self) -> None: + self.row_count = len(self.rows) + + +class TrinoClient: + """Thin async facade: each query runs the sync client in a worker thread + under a hard timeout, and fetches at most ``max_rows`` from the cursor + regardless of what the SQL says (defense in depth behind the guardrails). + """ + + def __init__(self, settings: GatewaySettings) -> None: + self._settings = settings + + def _connect(self) -> trino.dbapi.Connection: + s = self._settings + return trino.dbapi.connect( + host=s.trino_host, + port=s.trino_port, + user=s.trino_user, + catalog=s.trino_catalog, + schema=s.trino_schema, + http_scheme="http", + ) + + def _execute_sync(self, sql: str, max_rows: int) -> QueryResult: + with self._connect() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql) + rows = cursor.fetchmany(max_rows) + columns = [d[0] for d in cursor.description or []] + return QueryResult(columns=columns, rows=[list(r) for r in rows]) + finally: + cursor.close() + + async def execute(self, sql: str, *, timeout: float, max_rows: int) -> QueryResult: + try: + return await asyncio.wait_for( Review Comment: `asyncio.wait_for` cancels the awaitable, not the thread. On timeout `_execute_sync` keeps blocking inside `cursor.execute()` until Trino finishes on its own, and `cursor.close()` sits in the `finally`, so it only runs after `execute()` returns: the query is never cancelled server-side either. `asyncio.to_thread` uses the loop's default executor (`min(32, cpu + 4)` workers), and the system prompt tells the agent to fix the SQL and retry, so every timed-out query (120s by default) leaves a thread and a Trino connection behind while a fresh one starts. Once that executor is saturated, every `to_thread` queues behind the stuck threads, including `ping()` from `/ready`, so the pod stops reporting ready while Trino is still burning the abandoned work. Keep a handle on the cursor and call `Cursor.cancel()` in the timeout path, and run these on a bounded, dedicated executor rather than the shared default one. ########## hudi-lakehouse/local-dev/example/spark-app.yaml: ########## @@ -0,0 +1,61 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SparkApplication for the Apache Spark Kubernetes Operator +# (https://github.com/apache/spark-kubernetes-operator), submitted by +# local-dev/scripts/run-example.sh after uploading hudi_table_writer.py to +# s3a://warehouse/jobs/. +# +# To run your own job: copy this file + your script, adjust pyFiles (or +# mainClass/jars for JVM jobs), the env block, and any spark conf. +apiVersion: spark.apache.org/v1alpha1 +kind: SparkApplication +metadata: + name: hudi-table-writer + namespace: hudi-lakehouse +spec: + pyFiles: "s3a://warehouse/jobs/hudi_table_writer.py" + runtimeVersions: + sparkVersion: "3.5.7" + scalaVersion: "2.12" + sparkConf: + # Service account created by the spark-kubernetes-operator helm chart + # (the namespace default SA cannot manage executor pods). + spark.kubernetes.authenticate.driver.serviceAccountName: "spark" + spark.kubernetes.container.image: "hudi-lakehouse-spark:3.5" Review Comment: `quickstart.sh` accepts `--spark-version 4.1`, builds the Scala 2.13 bundle and has `build-images.sh` tag the image `hudi-lakehouse-spark:4.1`, but this manifest hardcodes the 3.5 image alongside `sparkVersion: "3.5.7"` and `scalaVersion: "2.12"` above, and neither `run-example.sh` nor `quickstart.sh` parameterizes it. So `quickstart.sh --spark-version 4.1 --run-example` submits a SparkApplication pointing at an image that was never built, and `run-example.sh` then spins its full poll timeout on `ErrImagePull`. Template the version into the manifest the way `run-example.sh` already templates the ConfigMap, or have `quickstart.sh` reject `--spark-version 4.1` together with `--run-example`. ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py: ########## @@ -0,0 +1,213 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Lakehouse tools backed by Trino. + +Handlers return JSON strings. Expected failures (guardrail rejections, query +errors, timeouts) are returned as ``{"error": ..., "hint": ...}`` payloads +rather than raised, so the agent can read the error and self-correct, and MCP +clients get a useful result instead of a protocol error. +""" + +from __future__ import annotations + +import json +import re +from typing import Annotated, Any + +from pydantic import Field + +from hudi_agent_gateway.config import GatewaySettings +from hudi_agent_gateway.tools.guardrails import enforce_guardrails +from hudi_agent_gateway.tools.registry import ToolInputError, ToolRegistry +from hudi_agent_gateway.tools.trino_client import ( + QueryResult, + TrinoClient, + TrinoQueryError, + TrinoTimeoutError, +) + +_QUERY_DESC = ( + "Run a single read-only SELECT statement (Trino SQL) against the lakehouse " + "and return the result as JSON. A server-side row cap is enforced; results " + "may be truncated (indicated by `truncated: true`)." +) +_LIST_TABLES_DESC = ( + "List tables in the lakehouse. Optionally filter by catalog and schema; " + "defaults to the gateway's configured catalog and schema." +) +_DESCRIBE_DESC = ( + "Describe a table's columns and types. Accepts `table`, `schema.table`, or " + "`catalog.schema.table`." +) + + +def shape_result(result: QueryResult, *, max_bytes: int, sql: str) -> str: + payload: dict[str, Any] = { + "sql": sql, + "columns": result.columns, + "rows": result.rows, + "row_count": result.row_count, + "truncated": False, Review Comment: `truncated` is only ever flipped by the byte-size loop below, so the truncation that actually fires in practice reports itself as a complete result. `enforce_guardrails` rewrites the model's SQL to `LIMIT {row_cap}`, and `_execute_sync` additionally caps the cursor with `fetchmany(max_rows)`, so `SELECT * FROM big_table` comes back as `row_count: 200, truncated: false`. The system prompt tells the model to warn the user only when a result has `truncated: true`, so the agent presents a silently capped result as the whole answer. That is the exact grounding failure the guardrails exist to prevent, and it fires on the most common query shape. Pass the row cap into `shape_result` and set `truncated` (or a separate `row_capped`) when `row_count == row_cap`. ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py: ########## @@ -0,0 +1,92 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SQL guardrails for model-written queries. + +AST-level (sqlglot, Trino dialect) rather than regex: read-only enforcement +survives comments, CTEs, and string literals; the row cap is injected as a +real ``LIMIT``. Fail-closed: anything that does not parse is rejected. +""" + +from __future__ import annotations + +import sqlglot +from sqlglot import exp + +from hudi_agent_gateway.tools.registry import ToolInputError + +_FORBIDDEN_NODES: tuple[type[exp.Expression], ...] = ( + exp.Insert, + exp.Update, + exp.Delete, + exp.Merge, + exp.Create, + exp.Drop, + exp.Alter, + exp.TruncateTable, + exp.Grant, + exp.Set, + exp.Command, # also covers CALL and EXPLAIN, which parse as generic commands + exp.Use, +) + +_HINT = "Provide exactly one read-only SELECT statement (Trino SQL)." + + +def enforce_guardrails(sql: str, row_cap: int) -> str: + """Validate ``sql`` as a single read-only SELECT and cap its LIMIT. + + Returns the (possibly rewritten) SQL to execute. Raises + :class:`ToolInputError` on any violation. + """ + try: + statements = sqlglot.parse(sql, read="trino") + except sqlglot.errors.SqlglotError as e: + raise ToolInputError(f"could not parse SQL: {e}", hint=_HINT) from e + + statements = [s for s in statements if s is not None] + if len(statements) != 1: + raise ToolInputError( + f"exactly one statement is allowed, got {len(statements)}", hint=_HINT + ) + stmt = statements[0] + + if not isinstance(stmt, (exp.Select, exp.SetOperation)): + raise ToolInputError( + f"read-only: only SELECT is allowed, got {type(stmt).__name__}", hint=_HINT + ) + + # Defense in depth: no write/DDL/command node anywhere in the tree + # (catches e.g. a CTE wrapping an INSERT). + for node in stmt.walk(): Review Comment: No test reaches this walk. Every statement in the rejection parametrize list in `test_guardrails.py` is a top-level non-`Select` node, so all of them are rejected by the `isinstance` gate above and never get here, and the two SELECT-shaped tests (the CTE one and the UNION one) contain no forbidden node. So the defense-in-depth layer the comment advertises, including the CTE-wrapping-a-write case it names, is entirely unproven. That matters more than a normal coverage gap here: the guardrail's safety rests wholly on sqlglot's node taxonomy, and `sqlglot>=25.0` is unpinned in `pyproject.toml`. Add a case that actually exercises the walk (a SELECT-rooted statement carrying a forbidden node), so this layer does not go silently dead if a sqlglot release reclassifies something. ########## hudi-agent-gateway/src/hudi_agent_gateway/llm.py: ########## @@ -0,0 +1,171 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLM provider factory and provider-aware model discovery. + +Construction never touches the network, so the gateway always starts; +connectivity is checked lazily and reported through ``/ready``. Model +discovery (``list_models``) asks the configured provider what it offers, so +the UI's model picker always reflects reality instead of a hardcoded list. +""" + +from __future__ import annotations + +import logging +import time + +import httpx +from langchain_core.language_models.chat_models import BaseChatModel +from pydantic import SecretStr + +from hudi_agent_gateway.config import GatewaySettings + +logger = logging.getLogger("hudi_agent_gateway.llm") + + +def build_chat_model(settings: GatewaySettings, model: str | None = None) -> BaseChatModel: + provider = settings.llm_provider + model_name = model or settings.llm_model + if provider == "anthropic": + from langchain_anthropic import ChatAnthropic + + return ChatAnthropic( # type: ignore[call-arg] # `model` is an init alias + model=model_name, + api_key=SecretStr(settings.anthropic_api_key), + timeout=settings.llm_timeout_seconds, + ) + if provider == "openai": + from langchain_openai import ChatOpenAI + + return ChatOpenAI( + model=model_name, + api_key=SecretStr(settings.openai_api_key), + timeout=settings.llm_timeout_seconds, + ) + if provider == "ollama": + from langchain_ollama import ChatOllama + + return ChatOllama(model=model_name, base_url=settings.ollama_base_url) Review Comment: This is the only provider branch that does not receive `timeout=settings.llm_timeout_seconds`, and `ollama` is the default provider in `GatewaySettings`. `GATEWAY_LLM_TIMEOUT_SECONDS` is documented in the README and wired through the chart's deployment, so on the default deployment the knob silently does nothing. Nothing else bounds the model call either: `agent.ainvoke` and `agent.astream` in the chat endpoint take no timeout, so a stalled Ollama (model load, OOM, GPU contention) hangs `/v1/chat` indefinitely with no error event on the stream. Pass `timeout=settings.llm_timeout_seconds` here like the other three branches. ########## hudi-lakehouse/scripts/build-jars.sh: ########## @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Builds the two Maven artifacts the hudi-lakehouse images need: +# 1. the Hudi Spark bundle (packaging/hudi-spark-bundle) -- JDK 11/17 +# 2. the Hudi Trino plugin (hudi-trino-plugin) -- JDK 23+ +# +# Usage: build-jars.sh [--spark-version 3.5|4.1] [--skip-bundle] [--skip-plugin] + +set -euo pipefail +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +SPARK_VERSION="3.5" +BUILD_BUNDLE=1 +BUILD_PLUGIN=1 + +while [[ $# -gt 0 ]]; do + case "$1" in + --spark-version) SPARK_VERSION="$2"; shift 2 ;; + --skip-bundle) BUILD_BUNDLE=0; shift ;; + --skip-plugin) BUILD_PLUGIN=0; shift ;; + *) echo "unknown arg: $1" >&2; exit 1 ;; + esac +done + +MVN_ARGS=(-DskipTests -Dcheckstyle.skip=true -Dscalastyle.skip=true -Drat.skip=true -Dmaven.javadoc.skip=true -Dgpg.skip=true) + +if [[ "$BUILD_BUNDLE" == 1 ]]; then + # Profiles must be activated explicitly: activating any -D profile property + # disables Maven's activeByDefault profiles, so relying on defaults is fragile. + echo ">>> Building Hudi Spark bundle (spark${SPARK_VERSION})" + (cd "$REPO_ROOT" && mvn clean package -pl packaging/hudi-spark-bundle -am \ + "-Dspark${SPARK_VERSION}" -Dflink2.1 "${MVN_ARGS[@]}") +fi + +if [[ "$BUILD_PLUGIN" == 1 ]]; then + JAVA_MAJOR=$(java -version 2>&1 | sed -n 's/.*version "\([0-9]*\).*/\1/p' | head -1) Review Comment: This probes `java` from `PATH`, not the JDK that will actually be used. `quickstart.sh` resolves (or downloads) a JDK 23 via `find_jdk23` and invokes this script as `JAVA_HOME="$JDK23" build-jars.sh --skip-bundle`. Maven honours `JAVA_HOME`, but this guard does not: it reads the developer's default JDK. So on any machine whose default JDK is 11/17/21, which is exactly the machine `find_jdk23` exists to serve, the check sees a major version below 23 and hard-exits, and the advertised one-command quickstart always dies at the Trino-plugin step. Probe the JDK that Maven will use: `"${JAVA_HOME:+$JAVA_HOME/bin/}java" -version`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
