wombatu-kun commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3567769670
########## 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 limits list omits the biggest one: there is no authn/authz at all. `/v1/chat`, `/mcp` and `/v1/models` are open to anyone who can reach the pod, and `host` defaults to `0.0.0.0`. Fix: add a bullet, e.g. "No authentication in v1: deploy behind an authenticating proxy or on a trusted network." ########## 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: `is_error` is a substring test (`'"error"' in content[:200]`), and `shape_result` puts the column names near the front of the payload, so a *successful* result with a column named `error` is reported as a failed tool call, both in `tool_trace` here and in the SSE `tool_result` event. Fix: reuse the check `_with_invocation_logging` in registry.py already does, `json.loads(content)` plus a top-level `error` key test, and share it between the two. ########## 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: Two things: 1. `/ready` runs an uncached `SELECT 1` on a fresh Trino connection on every probe, while the LLM check next to it is cached for 10s. With `periodSeconds: 10` in the chart that is a new connection and a query in Trino's history every 10s, forever. 2. The readinessProbe sets no `timeoutSeconds`, so Kubernetes applies its 1s default to a check whose own budget is `ping`'s 5s. Any Trino latency over 1s fails the probe, and `failureThreshold: 3` then pulls a healthy gateway out of the Service. Fix: cache the Trino ping the way `LLMReadiness` caches the LLM check, and set an explicit `timeoutSeconds` on the readinessProbe in deployment.yaml. Follow-up material, not a v1 blocker; I can file it as an issue. ########## 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: `ollama` is the default provider and the only branch with no timeout, so `GATEWAY_LLM_TIMEOUT_SECONDS` (documented in the README, wired through the chart as `llmTimeoutSeconds`) does nothing on the default deployment, and a stalled Ollama hangs `/v1/chat` unboundedly (`ainvoke`/`astream` in chat.py add no timeout either). Correcting my own earlier suggestion: `ChatOllama` has no `timeout` field, unlike the other three, so it has to go to the httpx client: ```python return ChatOllama( model=model_name, base_url=settings.ollama_base_url, client_kwargs={"timeout": settings.llm_timeout_seconds}, ) ``` ########## 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: Nothing reaches this walk: all 12 rejection cases in test_guardrails.py are non-Select at the root, so the isinstance gate above catches them first, and the two SELECT-shaped tests (CTE, UNION) carry no forbidden node. Fix: add one rejection case that is SELECT-rooted but carries a forbidden node in the tree, so this layer cannot go silently dead when sqlglot (unpinned, `>=25.0`) reshuffles its node taxonomy. If no such statement can be built in the Trino dialect, then the loop is dead code and can go instead. ########## 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 await, not the thread. On timeout `_execute_sync` stays blocked in `cursor.execute()` until Trino finishes on its own, the query is never cancelled server-side, and the `finally`'s `cursor.close()` only runs after that. These run on the loop's shared default executor and the system prompt tells the agent to fix the SQL and retry, so repeated timeouts (120s each by default) pile up stuck threads and eventually starve every `to_thread` call, including `ping()` from `/ready`. Fix: keep a handle on the cursor and call `Cursor.cancel()` in the timeout path, and run queries on a dedicated bounded executor instead of the default one. Not a v1 blocker; happy to file it as an issue. ########## 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: Fair on the hardening parts. This one is a 2-line bug though, so restating it concretely: `truncated` is only ever set by the byte-size loop, so the truncation that actually fires in practice, the row cap, reports itself as a complete result. `enforce_guardrails` rewrites the SQL to `LIMIT {row_cap}` and `fetchmany(max_rows)` caps the cursor, so `SELECT * FROM big_table` comes back as `row_count: 200, truncated: false`, and the system prompt tells the model to warn the user only when `truncated: true`. Fix: pass `row_cap` into `shape_result` and set `truncated` (or a separate `row_capped`) when `row_count == row_cap`. Say the word and I file it as an issue instead. -- 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]
