vinothchandar commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3634640870
########## 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: added to the limits list in ba979e3. ########## 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: fixed in ba979e3 — both places now parse the JSON and share the one check (`is_tool_error` in registry.py), per your suggestion. ########## 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: no need for an issue, took both in ba979e3: the trino ping is cached 10s like the llm check, and the readinessProbe sets an explicit `timeoutSeconds` now. ########## 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: fixed in ba979e3 with the `client_kwargs` route you landed on. ########## 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: `WITH t AS (INSERT INTO x VALUES (1)) SELECT * FROM t` does parse SELECT-rooted in the trino dialect, so the walk is live code. added exactly that as a rejection test in ba979e3. -- 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]
