From 5ad563b4bc6c03c8495bf2bd361df90829bebf81 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Tue, 4 Aug 2026 01:12:21 -0700 Subject: [PATCH 01/12] Add GSF structured query tools Signed-off-by: Soumili Nandi --- mcp/uv.lock | 1 + pyproject.toml | 2 + scripts/setup.sh | 1 + sources/gsf/README.md | 46 +++++ sources/gsf/pyproject.toml | 25 +++ sources/gsf/src/gsf/__init__.py | 10 ++ sources/gsf/src/gsf/client.py | 279 +++++++++++++++++++++++++++++ sources/gsf/src/gsf/errors.py | 63 +++++++ sources/gsf/src/gsf/models.py | 102 +++++++++++ sources/gsf/src/gsf/provenance.py | 12 ++ sources/gsf/src/gsf/register.py | 160 +++++++++++++++++ sources/gsf/tests/__init__.py | 2 + sources/gsf/tests/conftest.py | 47 +++++ sources/gsf/tests/test_client.py | 138 ++++++++++++++ sources/gsf/tests/test_models.py | 44 +++++ sources/gsf/tests/test_register.py | 102 +++++++++++ uv.lock | 20 +++ 17 files changed, 1054 insertions(+) create mode 100644 sources/gsf/README.md create mode 100644 sources/gsf/pyproject.toml create mode 100644 sources/gsf/src/gsf/__init__.py create mode 100644 sources/gsf/src/gsf/client.py create mode 100644 sources/gsf/src/gsf/errors.py create mode 100644 sources/gsf/src/gsf/models.py create mode 100644 sources/gsf/src/gsf/provenance.py create mode 100644 sources/gsf/src/gsf/register.py create mode 100644 sources/gsf/tests/__init__.py create mode 100644 sources/gsf/tests/conftest.py create mode 100644 sources/gsf/tests/test_client.py create mode 100644 sources/gsf/tests/test_models.py create mode 100644 sources/gsf/tests/test_register.py diff --git a/mcp/uv.lock b/mcp/uv.lock index 89971737b..ffa003de6 100644 --- a/mcp/uv.lock +++ b/mcp/uv.lock @@ -245,6 +245,7 @@ provides-extras = ["pii", "s3", "dev", "docs", "viz"] dev = [ { name = "aiq-api", editable = "../frontends/aiq_api" }, { name = "aiq-debug", editable = "../frontends/debug" }, + { name = "aiq-gsf", editable = "../sources/gsf" }, { name = "aiq-research-cli", editable = "../frontends/cli" }, { name = "dask", extras = ["distributed"], specifier = ">=2024.1.0" }, { name = "duckduckgo-news-search", editable = "../sources/duckduckgo_news_search" }, diff --git a/pyproject.toml b/pyproject.toml index 64ebd52f5..9a00a87fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -210,6 +210,7 @@ dev = [ "mypy>=1.5.0", "dask[distributed]>=2024.1.0", # Workspace packages for testing + "aiq-gsf", "google-scholar-paper-search", "tavily-web-search", "exa-web-search", @@ -259,6 +260,7 @@ exclude = ["mcp"] [tool.uv.sources] aiq-agent = { workspace = true } +aiq-gsf = { workspace = true } google-scholar-paper-search = { workspace = true } tavily-web-search = { workspace = true } exa-web-search = { workspace = true } diff --git a/scripts/setup.sh b/scripts/setup.sh index 118a3d9af..4737d7894 100755 --- a/scripts/setup.sh +++ b/scripts/setup.sh @@ -114,6 +114,7 @@ echo "Benchmarks installed" # Install data sources echo "" echo "Installing data sources..." +"${UV_BIN}" pip install -e ./sources/gsf "${UV_BIN}" pip install -e ./sources/tavily_web_search "${UV_BIN}" pip install -e ./sources/exa_web_search "${UV_BIN}" pip install -e ./sources/nimble_web_search diff --git a/sources/gsf/README.md b/sources/gsf/README.md new file mode 100644 index 000000000..54ce94b4b --- /dev/null +++ b/sources/gsf/README.md @@ -0,0 +1,46 @@ + + +# AI-Q GSF source + +This package exposes NVIDIA Generative Semantic Fabric (GSF) capabilities as a +NeMo Agent Toolkit function group. The current implementation provides: + +- `gsf__text_to_sql` +- `gsf__query_context` + +`gsf__catalog_search` is registered as an explicit +`capability_unavailable` placeholder until its GSF API contract is ready. + +The function group owns one shared HTTP connection pool. Authentication remains +request-scoped: each tool invocation obtains the current AI-Q user token and +passes it to GSF without storing it on the client. + +```yaml +function_groups: + gsf: + _type: gsf + base_url: ${GSF_BASE_URL:-http://gsf:3001} + include: + - query_context + - text_to_sql + +functions: + data_sources: + _type: data_source_registry + sources: + - id: gsf + name: "Enterprise Structured Data" + description: >- + Build authorized semantic context and execute bounded structured-data + queries through GSF. + default_enabled: true + requires_auth: true + tools: + - gsf +``` + +The `/api/v1/query-context` contract is provisional while GSF enriches the +existing `/api/text-to-data` capability. diff --git a/sources/gsf/pyproject.toml b/sources/gsf/pyproject.toml new file mode 100644 index 000000000..ca2298bd3 --- /dev/null +++ b/sources/gsf/pyproject.toml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools >= 83", "setuptools-scm>=8"] + +[tool.setuptools.packages.find] +where = ["src"] + +[project] +name = "aiq-gsf" +version = "0.1.0" +description = "NAT function group for NVIDIA GSF structured-data capabilities" +readme = "README.md" +requires-python = ">=3.11,<3.14" +license = {text = "Apache-2.0"} +dependencies = [ + "httpx>=0.27.0,<1", + "nvidia-nat-core==1.8.0", + "pydantic>=2.0.0", +] + +[project.entry-points."nat.plugins"] +gsf = "gsf.register" diff --git a/sources/gsf/src/gsf/__init__.py b/sources/gsf/src/gsf/__init__.py new file mode 100644 index 000000000..14d7fd043 --- /dev/null +++ b/sources/gsf/src/gsf/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AI-Q integration for NVIDIA Generative Semantic Fabric.""" + +from .client import GSFClient +from .errors import GSFError +from .errors import GSFErrorCode + +__all__ = ["GSFClient", "GSFError", "GSFErrorCode"] diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py new file mode 100644 index 000000000..89675275e --- /dev/null +++ b/sources/gsf/src/gsf/client.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed asynchronous client for GSF HTTP capabilities.""" + +import asyncio +import json +from collections.abc import Mapping +from typing import Any +from typing import TypeVar + +import httpx +from pydantic import BaseModel +from pydantic import ValidationError + +from .errors import GSFError +from .errors import GSFErrorCode +from .models import QueryContextRequest +from .models import QueryContextResponse +from .models import TextToSQLRequest +from .models import TextToSQLResponse + +ResponseT = TypeVar("ResponseT", bound=BaseModel) + +_FORWARDED_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) +_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) + + +class GSFClient: + """NAT-independent client sharing one bounded HTTP connection pool.""" + + def __init__( + self, + *, + base_url: str, + api_version: str = "v1", + connect_timeout_seconds: float = 5.0, + read_timeout_seconds: float = 60.0, + max_retries: int = 2, + max_response_bytes: int = 5_000_000, + default_max_rows: int = 1_000, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + self._api_base_url = f"{base_url.rstrip('/')}/api/{api_version.strip('/')}" + self._max_retries = max_retries + self._max_response_bytes = max_response_bytes + self._default_max_rows = default_max_rows + self._transport = transport + self._timeout = httpx.Timeout( + connect=connect_timeout_seconds, + read=read_timeout_seconds, + write=read_timeout_seconds, + pool=connect_timeout_seconds, + ) + self._client: httpx.AsyncClient | None = None + + @classmethod + def from_config(cls, config: Any) -> "GSFClient": + """Construct a client from the GSF function-group config without importing NAT.""" + + return cls( + base_url=str(config.base_url), + api_version=config.api_version, + connect_timeout_seconds=config.connect_timeout_seconds, + read_timeout_seconds=config.read_timeout_seconds, + max_retries=config.max_retries, + max_response_bytes=config.max_response_bytes, + default_max_rows=config.default_max_rows, + ) + + async def __aenter__(self) -> "GSFClient": + self._client = httpx.AsyncClient(timeout=self._timeout, transport=self._transport) + return self + + async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: + if self._client is not None: + await self._client.aclose() + self._client = None + + async def text_to_sql( + self, + request: TextToSQLRequest, + *, + token: str, + trace_headers: Mapping[str, str] | None = None, + ) -> TextToSQLResponse: + """Call GSF text-to-SQL and enforce AI-Q's configured row ceiling.""" + + max_rows = min(request.max_rows, self._default_max_rows) + payload = request.model_dump(exclude_none=True) + payload["max_rows"] = max_rows + result = await self._post( + "text-to-sql", + payload, + response_model=TextToSQLResponse, + token=token, + trace_headers=trace_headers, + capability="GSF text-to-SQL", + ) + if len(result.rows) > max_rows: + result.rows = result.rows[:max_rows] + result.truncated = True + return result + + async def query_context( + self, + request: QueryContextRequest, + *, + token: str, + trace_headers: Mapping[str, str] | None = None, + ) -> QueryContextResponse: + """Call GSF query-context and validate its token-budgeted metadata.""" + + return await self._post( + "query-context", + request.model_dump(exclude_none=True), + response_model=QueryContextResponse, + token=token, + trace_headers=trace_headers, + capability="GSF query context", + ) + + async def _post( + self, + endpoint: str, + payload: dict[str, Any], + *, + response_model: type[ResponseT], + token: str, + trace_headers: Mapping[str, str] | None, + capability: str, + ) -> ResponseT: + client = self._require_client() + headers = self._build_headers(token, trace_headers) + attempts = self._max_retries + 1 + + for attempt in range(attempts): + response: httpx.Response | None = None + try: + request = client.build_request( + "POST", + f"{self._api_base_url}/{endpoint}", + json=payload, + headers=headers, + ) + response = await client.send(request, stream=True) + request_id = response.headers.get("x-request-id") + if response.status_code >= 400: + error = self._http_error(response.status_code, capability, request_id) + if error.retryable and attempt + 1 < attempts: + await response.aclose() + response = None + await asyncio.sleep(2**attempt) + continue + raise error + + body = await self._read_bounded(response, request_id=request_id) + return self._validate_response(body, response_model, request_id=request_id) + except GSFError: + raise + except httpx.TimeoutException as exc: + if attempt + 1 < attempts: + await asyncio.sleep(2**attempt) + continue + raise GSFError( + GSFErrorCode.TIMEOUT, + f"{capability} timed out.", + retryable=True, + ) from exc + except httpx.TransportError as exc: + if attempt + 1 < attempts: + await asyncio.sleep(2**attempt) + continue + raise GSFError( + GSFErrorCode.UPSTREAM_ERROR, + f"{capability} could not reach GSF.", + retryable=True, + ) from exc + finally: + if response is not None: + await response.aclose() + + raise AssertionError("unreachable") + + def _require_client(self) -> httpx.AsyncClient: + if self._client is None: + raise RuntimeError("GSFClient must be used as an async context manager") + return self._client + + @staticmethod + def _build_headers(token: str, trace_headers: Mapping[str, str] | None) -> dict[str, str]: + headers = { + "Accept": "application/json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + for name, value in (trace_headers or {}).items(): + if name.lower() in _FORWARDED_HEADER_NAMES and value: + headers[name] = value + return headers + + async def _read_bounded(self, response: httpx.Response, *, request_id: str | None) -> bytes: + content_length = response.headers.get("content-length") + if content_length is not None: + try: + if int(content_length) > self._max_response_bytes: + raise self._response_too_large(request_id) + except ValueError: + pass + + body = bytearray() + async for chunk in response.aiter_bytes(): + body.extend(chunk) + if len(body) > self._max_response_bytes: + raise self._response_too_large(request_id) + return bytes(body) + + @staticmethod + def _validate_response(body: bytes, response_model: type[ResponseT], *, request_id: str | None) -> ResponseT: + try: + payload = json.loads(body) + if isinstance(payload, dict) and set(payload) == {"data"}: + payload = payload["data"] + return response_model.model_validate(payload) + except (json.JSONDecodeError, UnicodeDecodeError, ValidationError, TypeError) as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned an invalid response.", + request_id=request_id, + ) from exc + + def _response_too_large(self, request_id: str | None) -> GSFError: + return GSFError( + GSFErrorCode.RESPONSE_TOO_LARGE, + "GSF response exceeded the configured size limit.", + request_id=request_id, + ) + + @staticmethod + def _http_error(status_code: int, capability: str, request_id: str | None) -> GSFError: + if status_code == 401: + return GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF authentication is required.", + request_id=request_id, + ) + if status_code == 403: + return GSFError(GSFErrorCode.FORBIDDEN, "GSF access is forbidden.", request_id=request_id) + if status_code == 404: + return GSFError( + GSFErrorCode.CAPABILITY_UNAVAILABLE, + f"{capability} is unavailable.", + request_id=request_id, + ) + if status_code in {400, 422}: + return GSFError( + GSFErrorCode.INVALID_REQUEST, + "GSF rejected the request.", + request_id=request_id, + ) + if status_code == 429: + return GSFError( + GSFErrorCode.RATE_LIMITED, + "GSF rate limit was reached.", + retryable=True, + request_id=request_id, + ) + if status_code in _RETRYABLE_STATUS_CODES: + return GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF is temporarily unavailable.", + retryable=True, + request_id=request_id, + ) + return GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF request failed.", + request_id=request_id, + ) diff --git a/sources/gsf/src/gsf/errors.py b/sources/gsf/src/gsf/errors.py new file mode 100644 index 000000000..bf05d5c0c --- /dev/null +++ b/sources/gsf/src/gsf/errors.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stable, non-secret errors returned by the GSF integration.""" + +from enum import StrEnum +from typing import Literal + +from pydantic import BaseModel +from pydantic import ConfigDict + + +class GSFErrorCode(StrEnum): + """Error codes exposed by GSF-backed AI-Q tools.""" + + AUTHENTICATION_REQUIRED = "authentication_required" + FORBIDDEN = "forbidden" + INVALID_REQUEST = "invalid_request" + CAPABILITY_UNAVAILABLE = "capability_unavailable" + TIMEOUT = "timeout" + RATE_LIMITED = "rate_limited" + INVALID_RESPONSE = "invalid_response" + RESPONSE_TOO_LARGE = "response_too_large" + UPSTREAM_ERROR = "upstream_error" + + +class GSFError(Exception): + """Internal exception containing only caller-safe GSF failure details.""" + + def __init__( + self, + code: GSFErrorCode, + message: str, + *, + retryable: bool = False, + request_id: str | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.message = message + self.retryable = retryable + self.request_id = request_id + + +class GSFToolError(BaseModel): + """Serialized error envelope returned to an AI-Q agent.""" + + model_config = ConfigDict(extra="forbid") + + status: Literal["error"] = "error" + code: GSFErrorCode + retryable: bool + request_id: str | None = None + message: str + + @classmethod + def from_exception(cls, error: GSFError) -> "GSFToolError": + return cls( + code=error.code, + retryable=error.retryable, + request_id=error.request_id, + message=error.message, + ) diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py new file mode 100644 index 000000000..b9ae355e0 --- /dev/null +++ b/sources/gsf/src/gsf/models.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed, NAT-independent contracts for GSF capabilities.""" + +from typing import Any + +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field + + +class GSFRequest(BaseModel): + """Base model for data sent from AI-Q to GSF.""" + + model_config = ConfigDict(extra="forbid") + + +class GSFResponse(BaseModel): + """Base model for the validated subset of data returned by GSF.""" + + model_config = ConfigDict(extra="ignore") + + +class CatalogSearchRequest(GSFRequest): + """Provisional catalog-search input retained for the unavailable placeholder.""" + + question: str = Field(min_length=1, max_length=4_096) + database_name: str | None = None + max_results: int = Field(default=10, ge=1, le=100) + token_budget: int | None = Field(default=None, ge=1) + + +class ResultColumn(GSFResponse): + """A column in a bounded SQL result.""" + + name: str + data_type: str + + +class SemanticContext(GSFResponse): + """Semantic provenance used to produce a SQL query.""" + + metrics: list[dict[str, Any]] = Field(default_factory=list) + grain: str | None = None + units: list[str] = Field(default_factory=list) + filters: list[str] = Field(default_factory=list) + rules: list[str] = Field(default_factory=list) + omissions: list[str] = Field(default_factory=list) + + +class TextToSQLRequest(GSFRequest): + """Generate validated SQL and optionally execute it with bounded results.""" + + question: str = Field(min_length=1, max_length=4_096) + database_name: str | None = None + execute: bool = True + object_ids: list[str] = Field(default_factory=list) + max_rows: int = Field(default=1_000, ge=1) + + +class TextToSQLResponse(GSFResponse): + """Validated SQL, bounded rows, and semantic provenance returned by GSF.""" + + request_id: str + sql: str + columns: list[ResultColumn] + rows: list[dict[str, Any]] + truncated: bool + objects_used: list[str] = Field(default_factory=list) + joins_used: list[dict[str, Any]] = Field(default_factory=list) + semantic_context: SemanticContext + validation_attempts: list[dict[str, Any]] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + timings: dict[str, int] = Field(default_factory=dict) + + +class QueryContextRequest(GSFRequest): + """Build compact, authorized context for a later SQL-generation step.""" + + question: str = Field(min_length=1, max_length=4_096) + database_name: str | None = None + object_ids: list[str] = Field(default_factory=list) + token_budget: int | None = Field(default=None, ge=1) + + +class QueryContextResponse(GSFResponse): + """Token-budgeted semantic and physical metadata relevant to a question.""" + + request_id: str + tables: list[dict[str, Any]] = Field(default_factory=list) + columns: list[dict[str, Any]] = Field(default_factory=list) + keys: list[dict[str, Any]] = Field(default_factory=list) + join_paths: list[dict[str, Any]] = Field(default_factory=list) + values: list[dict[str, Any]] = Field(default_factory=list) + metrics: list[dict[str, Any]] = Field(default_factory=list) + grain: str | None = None + units: list[str] = Field(default_factory=list) + rules: list[str] = Field(default_factory=list) + omissions: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + truncated: bool = False diff --git a/sources/gsf/src/gsf/provenance.py b/sources/gsf/src/gsf/provenance.py new file mode 100644 index 000000000..297b25b7c --- /dev/null +++ b/sources/gsf/src/gsf/provenance.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Non-sensitive helpers for GSF result provenance.""" + +import hashlib + + +def sql_sha256(sql: str) -> str: + """Return a stable SQL digest without logging or persisting the SQL text.""" + + return hashlib.sha256(sql.encode("utf-8")).hexdigest() diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py new file mode 100644 index 000000000..31433ae13 --- /dev/null +++ b/sources/gsf/src/gsf/register.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Register GSF capabilities as one NAT function group.""" + +import logging +from collections.abc import Mapping + +from pydantic import Field +from pydantic import HttpUrl + +from aiq_agent.auth.utils import get_auth_token +from nat.builder.builder import Builder +from nat.builder.context import Context +from nat.builder.function import FunctionGroup +from nat.cli.register_workflow import register_function_group +from nat.data_models.function import FunctionGroupBaseConfig + +from .client import GSFClient +from .errors import GSFError +from .errors import GSFErrorCode +from .errors import GSFToolError +from .models import CatalogSearchRequest +from .models import QueryContextRequest +from .models import TextToSQLRequest + +logger = logging.getLogger(__name__) + +_TRACE_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) + + +class GSFFunctionGroupConfig(FunctionGroupBaseConfig, name="gsf"): + """Shared configuration for AI-Q's GSF tools.""" + + base_url: HttpUrl + api_version: str = "v1" + connect_timeout_seconds: float = Field(default=5.0, gt=0) + read_timeout_seconds: float = Field(default=60.0, gt=0) + max_retries: int = Field(default=2, ge=0, le=5) + max_response_bytes: int = Field(default=5_000_000, ge=1) + default_max_rows: int = Field(default=1_000, ge=1) + + +def _tool_error(error: GSFError) -> str: + return GSFToolError.from_exception(error).model_dump_json(exclude_none=True) + + +def _authentication_error() -> str: + return _tool_error( + GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF authentication is required.", + ) + ) + + +def _request_trace_headers() -> Mapping[str, str]: + try: + metadata = Context.get().metadata + incoming = metadata.headers if metadata else None + except Exception: + return {} + if not incoming: + return {} + return {name: value for name, value in incoming.items() if name.lower() in _TRACE_HEADER_NAMES and value} + + +@register_function_group(config_type=GSFFunctionGroupConfig) +async def gsf_function_group(config: GSFFunctionGroupConfig, _builder: Builder): + """Build namespaced GSF tools around one shared HTTP client.""" + + async with GSFClient.from_config(config) as client: + + async def catalog_search(request: CatalogSearchRequest) -> str: + """Search the GSF enterprise catalog (not available in this integration yet).""" + + del request + return _tool_error( + GSFError( + GSFErrorCode.CAPABILITY_UNAVAILABLE, + "GSF catalog search is unavailable.", + ) + ) + + async def text_to_sql(request: TextToSQLRequest) -> str: + """Generate validated SQL and optionally return bounded rows from authorized enterprise data. + + Use for an analytical question after the relevant structured-data scope is known. The result contains SQL, + rows, semantic context, warnings, and provenance; AI-Q remains responsible for analysis and synthesis. + """ + + token = get_auth_token() + if not token: + return _authentication_error() + try: + result = await client.text_to_sql( + request, + token=token, + trace_headers=_request_trace_headers(), + ) + return result.model_dump_json(exclude_none=True) + except GSFError as error: + return _tool_error(error) + except Exception: + logger.error("Unexpected GSF text-to-SQL failure") + return _tool_error( + GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF text-to-SQL failed.", + ) + ) + + async def query_context(request: QueryContextRequest) -> str: + """Build compact, authorized, token-budgeted semantic context for SQL generation. + + Use when a downstream SQL-generation step needs relevant tables, columns, keys, joins, metrics, grain, + rules, omissions, and warnings without executing a query. + """ + + token = get_auth_token() + if not token: + return _authentication_error() + try: + result = await client.query_context( + request, + token=token, + trace_headers=_request_trace_headers(), + ) + return result.model_dump_json(exclude_none=True) + except GSFError as error: + return _tool_error(error) + except Exception: + logger.error("Unexpected GSF query-context failure") + return _tool_error( + GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF query context failed.", + ) + ) + + group = FunctionGroup(config=config) + group.add_function( + "catalog_search", + catalog_search, + input_schema=CatalogSearchRequest, + description=catalog_search.__doc__, + ) + group.add_function( + "text_to_sql", + text_to_sql, + input_schema=TextToSQLRequest, + description=text_to_sql.__doc__, + ) + group.add_function( + "query_context", + query_context, + input_schema=QueryContextRequest, + description=query_context.__doc__, + ) + yield group diff --git a/sources/gsf/tests/__init__.py b/sources/gsf/tests/__init__.py new file mode 100644 index 000000000..d51c4fe1e --- /dev/null +++ b/sources/gsf/tests/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py new file mode 100644 index 000000000..eb510256b --- /dev/null +++ b/sources/gsf/tests/conftest.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + + +@pytest.fixture +def text_to_sql_response() -> dict: + return { + "request_id": "gsf-request-1", + "sql": "SELECT revenue FROM quarterly_results", + "columns": [{"name": "revenue", "data_type": "numeric"}], + "rows": [{"revenue": 100}, {"revenue": 200}], + "truncated": False, + "objects_used": ["metric:revenue"], + "joins_used": [], + "semantic_context": { + "metrics": [{"id": "metric:revenue"}], + "grain": "quarter", + "units": ["USD"], + "filters": [], + "rules": [], + "omissions": [], + }, + "validation_attempts": [], + "warnings": [], + "timings": {"total_ms": 25}, + } + + +@pytest.fixture +def query_context_response() -> dict: + return { + "request_id": "gsf-request-2", + "tables": [{"id": "table:quarterly_results", "grain": "quarter"}], + "columns": [{"table_id": "table:quarterly_results", "name": "revenue", "data_type": "numeric"}], + "keys": [{"table_id": "table:quarterly_results", "columns": ["quarter"]}], + "join_paths": [], + "values": [], + "metrics": [{"id": "metric:revenue", "unit": "USD"}], + "grain": "quarter", + "units": ["USD"], + "rules": [], + "omissions": [], + "warnings": [], + "truncated": False, + } diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py new file mode 100644 index 000000000..0e751f970 --- /dev/null +++ b/sources/gsf/tests/test_client.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from unittest.mock import AsyncMock +from unittest.mock import patch + +import httpx +import pytest +from gsf.client import GSFClient +from gsf.errors import GSFError +from gsf.errors import GSFErrorCode +from gsf.models import QueryContextRequest +from gsf.models import TextToSQLRequest + + +@pytest.mark.asyncio +async def test_text_to_sql_sends_scoped_request_and_bounds_rows(text_to_sql_response: dict) -> None: + seen_request: httpx.Request | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal seen_request + seen_request = request + return httpx.Response(200, json=text_to_sql_response, headers={"x-request-id": "header-request"}) + + client = GSFClient( + base_url="https://gsf.example", + default_max_rows=1, + transport=httpx.MockTransport(handler), + ) + async with client: + result = await client.text_to_sql( + TextToSQLRequest(question="Show revenue", database_name="benchmark_db", max_rows=20), + token="user-token", + trace_headers={"traceparent": "00-trace", "authorization": "do-not-forward"}, + ) + + assert seen_request is not None + assert seen_request.url == "https://gsf.example/api/v1/text-to-sql" + assert seen_request.headers["authorization"] == "Bearer user-token" + assert seen_request.headers["traceparent"] == "00-trace" + assert json.loads(seen_request.content) == { + "question": "Show revenue", + "database_name": "benchmark_db", + "execute": True, + "object_ids": [], + "max_rows": 1, + } + assert result.rows == [{"revenue": 100}] + assert result.truncated is True + + +@pytest.mark.asyncio +async def test_query_context_omits_database_and_unwraps_data(query_context_response: dict) -> None: + seen_payload: dict | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal seen_payload + seen_payload = json.loads(request.content) + return httpx.Response(200, json={"data": query_context_response}) + + client = GSFClient(base_url="https://gsf.example/", transport=httpx.MockTransport(handler)) + async with client: + result = await client.query_context( + QueryContextRequest(question="What revenue data is available?", token_budget=2_000), + token="user-token", + ) + + assert seen_payload == { + "question": "What revenue data is available?", + "object_ids": [], + "token_budget": 2_000, + } + assert result.request_id == "gsf-request-2" + + +@pytest.mark.asyncio +async def test_client_normalizes_forbidden_without_leaking_body() -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(403, text="secret database details") + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + with pytest.raises(GSFError) as raised: + await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + + assert raised.value.code is GSFErrorCode.FORBIDDEN + assert "secret" not in raised.value.message + + +@pytest.mark.asyncio +async def test_client_retries_rate_limit_then_succeeds(query_context_response: dict) -> None: + attempts = 0 + + async def handler(_request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(429) + return httpx.Response(200, json=query_context_response) + + client = GSFClient(base_url="https://gsf.example", max_retries=1, transport=httpx.MockTransport(handler)) + with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: + async with client: + await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + + assert attempts == 2 + sleep.assert_awaited_once_with(1) + + +@pytest.mark.asyncio +async def test_client_rejects_oversized_response(query_context_response: dict) -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=query_context_response) + + client = GSFClient( + base_url="https://gsf.example", + max_response_bytes=10, + transport=httpx.MockTransport(handler), + ) + async with client: + with pytest.raises(GSFError) as raised: + await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + + assert raised.value.code is GSFErrorCode.RESPONSE_TOO_LARGE + + +@pytest.mark.asyncio +async def test_client_rejects_malformed_response() -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=b"not-json") + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + with pytest.raises(GSFError) as raised: + await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + + assert raised.value.code is GSFErrorCode.INVALID_RESPONSE diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py new file mode 100644 index 000000000..ca935171b --- /dev/null +++ b/sources/gsf/tests/test_models.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from gsf.models import QueryContextRequest +from gsf.models import QueryContextResponse +from gsf.models import TextToSQLRequest +from gsf.models import TextToSQLResponse +from pydantic import ValidationError + + +def test_text_to_sql_request_supports_optional_database_name() -> None: + request = TextToSQLRequest(question="Show quarterly revenue", database_name="benchmark_db") + + assert request.database_name == "benchmark_db" + assert request.execute is True + assert request.max_rows == 1_000 + + +def test_query_context_request_omits_optional_database_name() -> None: + payload = QueryContextRequest(question="What revenue data is available?").model_dump(exclude_none=True) + + assert "database_name" not in payload + + +def test_requests_reject_unknown_fields() -> None: + with pytest.raises(ValidationError): + TextToSQLRequest.model_validate({"question": "Show revenue", "unknown": True}) + + +def test_text_to_sql_response_requires_provenance(text_to_sql_response: dict) -> None: + text_to_sql_response.pop("request_id") + + with pytest.raises(ValidationError): + TextToSQLResponse.model_validate(text_to_sql_response) + + +def test_query_context_response_accepts_future_fields(query_context_response: dict) -> None: + query_context_response["future_gsf_metadata"] = {"enabled": True} + + result = QueryContextResponse.model_validate(query_context_response) + + assert result.request_id == "gsf-request-2" + assert not hasattr(result, "future_gsf_metadata") diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py new file mode 100644 index 000000000..6710090d1 --- /dev/null +++ b/sources/gsf/tests/test_register.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest +from gsf.models import QueryContextResponse +from gsf.models import TextToSQLResponse +from gsf.register import GSFFunctionGroupConfig +from gsf.register import gsf_function_group + + +class FakeClientContext: + def __init__(self, client: MagicMock) -> None: + self.client = client + + async def __aenter__(self) -> MagicMock: + return self.client + + async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: + return None + + +@pytest.mark.asyncio +async def test_group_exposes_only_requested_tools(text_to_sql_response: dict, query_context_response: dict) -> None: + client = MagicMock() + client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) + client.query_context = AsyncMock(return_value=QueryContextResponse.model_validate(query_context_response)) + config = GSFFunctionGroupConfig( + base_url="https://gsf.example", + include=["text_to_sql", "query_context"], + ) + + with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): + async with gsf_function_group(config, MagicMock()) as group: + tools = await group.get_accessible_functions() + + assert set(tools) == {"gsf__query_context", "gsf__text_to_sql"} + + +@pytest.mark.asyncio +async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: dict) -> None: + client = MagicMock() + client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) + client.query_context = AsyncMock() + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) + + with ( + patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), + patch("gsf.register.get_auth_token", side_effect=["token-one", "token-two"]), + patch("gsf.register._request_trace_headers", return_value={}), + ): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__text_to_sql"] + first = json.loads(await tool.ainvoke({"question": "First question"})) + second = json.loads(await tool.ainvoke({"question": "Second question"})) + + assert first["request_id"] == "gsf-request-1" + assert second["request_id"] == "gsf-request-1" + assert client.text_to_sql.await_args_list[0].kwargs["token"] == "token-one" + assert client.text_to_sql.await_args_list[1].kwargs["token"] == "token-two" + assert "token" not in client.__dict__ + + +@pytest.mark.asyncio +async def test_missing_authentication_fails_closed() -> None: + client = MagicMock() + client.text_to_sql = AsyncMock() + client.query_context = AsyncMock() + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["query_context"]) + + with ( + patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), + patch("gsf.register.get_auth_token", return_value=None), + ): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__query_context"] + result = json.loads(await tool.ainvoke({"question": "Show data"})) + + assert result["code"] == "authentication_required" + client.query_context.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_catalog_search_is_explicitly_unavailable() -> None: + client = MagicMock() + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["catalog_search"]) + + with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__catalog_search"] + result = json.loads(await tool.ainvoke({"question": "Find revenue metrics"})) + + assert result == { + "status": "error", + "code": "capability_unavailable", + "retryable": False, + "message": "GSF catalog search is unavailable.", + } diff --git a/uv.lock b/uv.lock index 86c3c5e7c..51eb8db87 100644 --- a/uv.lock +++ b/uv.lock @@ -15,6 +15,7 @@ members = [ "aiq-agent", "aiq-api", "aiq-debug", + "aiq-gsf", "aiq-research-cli", "deepsearch-qa-evaluator", "duckduckgo-news-search", @@ -261,6 +262,7 @@ viz = [ dev = [ { name = "aiq-api" }, { name = "aiq-debug" }, + { name = "aiq-gsf" }, { name = "aiq-research-cli" }, { name = "dask", extra = ["distributed"] }, { name = "duckduckgo-news-search" }, @@ -331,6 +333,7 @@ provides-extras = ["pii", "s3", "dev", "docs", "viz"] dev = [ { name = "aiq-api", editable = "frontends/aiq_api" }, { name = "aiq-debug", editable = "frontends/debug" }, + { name = "aiq-gsf", editable = "sources/gsf" }, { name = "aiq-research-cli", editable = "frontends/cli" }, { name = "dask", extras = ["distributed"], specifier = ">=2024.1.0" }, { name = "duckduckgo-news-search", editable = "sources/duckduckgo_news_search" }, @@ -417,6 +420,23 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "aiq-gsf" +version = "0.1.0" +source = { editable = "sources/gsf" } +dependencies = [ + { name = "httpx" }, + { name = "nvidia-nat-core" }, + { name = "pydantic" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.27.0,<1" }, + { name = "nvidia-nat-core", specifier = "==1.8.0" }, + { name = "pydantic", specifier = ">=2.0.0" }, +] + [[package]] name = "aiq-research-cli" version = "2.0.1" From 2d13f7b5da0c3a80468e2888921ad0a69e510079 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Tue, 4 Aug 2026 14:00:52 -0700 Subject: [PATCH 02/12] Align GSF adapter with chat completions Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 17 +- sources/gsf/src/gsf/client.py | 243 ++++++++++++++++++++++++----- sources/gsf/src/gsf/models.py | 63 ++++---- sources/gsf/src/gsf/register.py | 39 ++--- sources/gsf/tests/conftest.py | 35 ++--- sources/gsf/tests/test_client.py | 87 +++++++---- sources/gsf/tests/test_models.py | 39 +++-- sources/gsf/tests/test_register.py | 36 ++--- 8 files changed, 375 insertions(+), 184 deletions(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index 54ce94b4b..f1de4cb90 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -9,22 +9,21 @@ This package exposes NVIDIA Generative Semantic Fabric (GSF) capabilities as a NeMo Agent Toolkit function group. The current implementation provides: - `gsf__text_to_sql` -- `gsf__query_context` -`gsf__catalog_search` is registered as an explicit -`capability_unavailable` placeholder until its GSF API contract is ready. +`gsf__catalog_search` and `gsf__query_context` are registered as explicit +`capability_unavailable` placeholders until their GSF API contracts are ready. The function group owns one shared HTTP connection pool. Authentication remains request-scoped: each tool invocation obtains the current AI-Q user token and passes it to GSF without storing it on the client. +`GSF_BASE_URL` must point to GSF's auth-aware API origin. ```yaml function_groups: gsf: _type: gsf - base_url: ${GSF_BASE_URL:-http://gsf:3001} + base_url: ${GSF_BASE_URL:-http://gsf:3000} include: - - query_context - text_to_sql functions: @@ -42,5 +41,9 @@ functions: - gsf ``` -The `/api/v1/query-context` contract is provisional while GSF enriches the -existing `/api/text-to-data` capability. +Text-to-SQL uses GSF's `/api/chat/completions` SSE endpoint with +`prediction: false`. Its optional AI-Q `database_name` input is sent to GSF as +`target_db`, selecting an existing GSF connection rather than creating one. +Future prediction tools can reuse the client transport with `prediction: true`. +The adapter normalizes GSF's current response fields while preserving optional +semantic and benchmarking fields as they become available. diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index 89675275e..d825ea640 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -7,21 +7,18 @@ import json from collections.abc import Mapping from typing import Any -from typing import TypeVar import httpx -from pydantic import BaseModel from pydantic import ValidationError from .errors import GSFError from .errors import GSFErrorCode -from .models import QueryContextRequest -from .models import QueryContextResponse +from .models import ChatCompletionResult +from .models import ChatCompletionsRequest +from .models import ResultColumn from .models import TextToSQLRequest from .models import TextToSQLResponse -ResponseT = TypeVar("ResponseT", bound=BaseModel) - _FORWARDED_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) _RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) @@ -33,7 +30,6 @@ def __init__( self, *, base_url: str, - api_version: str = "v1", connect_timeout_seconds: float = 5.0, read_timeout_seconds: float = 60.0, max_retries: int = 2, @@ -41,7 +37,7 @@ def __init__( default_max_rows: int = 1_000, transport: httpx.AsyncBaseTransport | None = None, ) -> None: - self._api_base_url = f"{base_url.rstrip('/')}/api/{api_version.strip('/')}" + self._api_base_url = f"{base_url.rstrip('/')}/api" self._max_retries = max_retries self._max_response_bytes = max_response_bytes self._default_max_rows = default_max_rows @@ -60,7 +56,6 @@ def from_config(cls, config: Any) -> "GSFClient": return cls( base_url=str(config.base_url), - api_version=config.api_version, connect_timeout_seconds=config.connect_timeout_seconds, read_timeout_seconds=config.read_timeout_seconds, max_retries=config.max_retries, @@ -84,54 +79,52 @@ async def text_to_sql( token: str, trace_headers: Mapping[str, str] | None = None, ) -> TextToSQLResponse: - """Call GSF text-to-SQL and enforce AI-Q's configured row ceiling.""" + """Run the SQL branch of GSF chat completions and normalize its answer.""" max_rows = min(request.max_rows, self._default_max_rows) - payload = request.model_dump(exclude_none=True) - payload["max_rows"] = max_rows - result = await self._post( - "text-to-sql", - payload, - response_model=TextToSQLResponse, + result = await self.chat_completions( + ChatCompletionsRequest( + question=request.question, + prediction=False, + target_db=request.database_name, + ), token=token, trace_headers=trace_headers, - capability="GSF text-to-SQL", ) - if len(result.rows) > max_rows: - result.rows = result.rows[:max_rows] - result.truncated = True - return result + return self._normalize_text_to_sql(result, max_rows=max_rows) - async def query_context( + async def chat_completions( self, - request: QueryContextRequest, + request: ChatCompletionsRequest, *, token: str, trace_headers: Mapping[str, str] | None = None, - ) -> QueryContextResponse: - """Call GSF query-context and validate its token-budgeted metadata.""" + ) -> ChatCompletionResult: + """Call the shared GSF chat endpoint and extract its final SSE result.""" - return await self._post( - "query-context", + body, request_id, content_type = await self._post( + "chat/completions", request.model_dump(exclude_none=True), - response_model=QueryContextResponse, token=token, trace_headers=trace_headers, - capability="GSF query context", + capability="GSF chat completions", + accept="text/event-stream", ) + answer = self._parse_chat_answer(body, content_type=content_type, request_id=request_id) + return ChatCompletionResult(answer=answer, request_id=request_id) async def _post( self, endpoint: str, payload: dict[str, Any], *, - response_model: type[ResponseT], token: str, trace_headers: Mapping[str, str] | None, capability: str, - ) -> ResponseT: + accept: str = "application/json", + ) -> tuple[bytes, str | None, str]: client = self._require_client() - headers = self._build_headers(token, trace_headers) + headers = self._build_headers(token, trace_headers, accept=accept) attempts = self._max_retries + 1 for attempt in range(attempts): @@ -155,7 +148,7 @@ async def _post( raise error body = await self._read_bounded(response, request_id=request_id) - return self._validate_response(body, response_model, request_id=request_id) + return body, request_id, response.headers.get("content-type", "") except GSFError: raise except httpx.TimeoutException as exc: @@ -188,9 +181,14 @@ def _require_client(self) -> httpx.AsyncClient: return self._client @staticmethod - def _build_headers(token: str, trace_headers: Mapping[str, str] | None) -> dict[str, str]: + def _build_headers( + token: str, + trace_headers: Mapping[str, str] | None, + *, + accept: str, + ) -> dict[str, str]: headers = { - "Accept": "application/json", + "Accept": accept, "Authorization": f"Bearer {token}", "Content-Type": "application/json", } @@ -215,20 +213,181 @@ async def _read_bounded(self, response: httpx.Response, *, request_id: str | Non raise self._response_too_large(request_id) return bytes(body) - @staticmethod - def _validate_response(body: bytes, response_model: type[ResponseT], *, request_id: str | None) -> ResponseT: + @classmethod + def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | None) -> dict[str, Any]: try: - payload = json.loads(body) - if isinstance(payload, dict) and set(payload) == {"data"}: - payload = payload["data"] - return response_model.model_validate(payload) - except (json.JSONDecodeError, UnicodeDecodeError, ValidationError, TypeError) as exc: + text = body.decode("utf-8") + if "text/event-stream" not in content_type and not text.lstrip().startswith(("data:", ":")): + payload = json.loads(text) + return cls._answer_from_event(payload, request_id=request_id) + + data_lines: list[str] = [] + for line in [*text.splitlines(), ""]: + if line.startswith(":"): + continue + if line.startswith("data:"): + data_lines.append(line[5:].lstrip()) + continue + if line or not data_lines: + continue + event_data = "\n".join(data_lines) + data_lines.clear() + if event_data == "[DONE]": + continue + event = json.loads(event_data) + if not isinstance(event, dict): + continue + if event.get("type") == "error": + raise GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF chat completions failed.", + request_id=request_id, + ) + if event.get("type") == "result": + return cls._answer_from_event(event, request_id=request_id) + except GSFError: + raise + except (json.JSONDecodeError, UnicodeDecodeError, TypeError) as exc: raise GSFError( GSFErrorCode.INVALID_RESPONSE, "GSF returned an invalid response.", request_id=request_id, ) from exc + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF response did not contain a final result.", + request_id=request_id, + ) + + @staticmethod + def _answer_from_event(payload: Any, *, request_id: str | None) -> dict[str, Any]: + if isinstance(payload, dict) and isinstance(payload.get("answer"), dict): + return payload["answer"] + if isinstance(payload, dict) and payload.get("type") is None: + return payload + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned an invalid final answer.", + request_id=request_id, + ) + + @classmethod + def _normalize_text_to_sql(cls, result: ChatCompletionResult, *, max_rows: int) -> TextToSQLResponse: + answer = result.answer + sql = answer.get("sql") or answer.get("sql_code") + if not isinstance(sql, str) or not sql.strip(): + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF did not return validated SQL.", + request_id=result.request_id, + ) + + rows = cls._normalize_rows(answer.get("rows", answer.get("sql_response_from_db")), result.request_id) + upstream_truncated = bool(answer.get("truncated", False)) + truncated = upstream_truncated or len(rows) > max_rows + rows = rows[:max_rows] + columns = cls._normalize_columns(answer.get("columns", answer.get("sql_columns"))) + if not columns and rows: + columns = [ResultColumn(name=str(name)) for name in rows[0]] + + try: + return TextToSQLResponse( + request_id=answer.get("request_id") or result.request_id, + response=answer.get("response"), + sql=sql, + columns=columns, + rows=rows, + truncated=truncated, + custom_analyses_used=answer.get("custom_analyses_used"), + objects_used=answer.get("objects_used"), + joins_used=answer.get("joins_used"), + semantic_context=answer.get("semantic_context"), + validation_attempts=answer.get("validation_attempts"), + assumptions=answer.get("assumptions"), + warnings=answer.get("warnings"), + timings=answer.get("timings"), + ) + except ValidationError as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid text-to-SQL data.", + request_id=result.request_id, + ) from exc + + @staticmethod + def _normalize_columns(value: Any) -> list[ResultColumn]: + if not isinstance(value, list): + return [] + columns: list[ResultColumn] = [] + for column in value: + if isinstance(column, dict): + name = column.get("name") or column.get("column_name") or column.get("id") + if name is not None: + columns.append( + ResultColumn( + name=str(name), + data_type=column.get("data_type") or column.get("type"), + ) + ) + elif column is not None: + columns.append(ResultColumn(name=str(column))) + return columns + + @staticmethod + def _normalize_rows(value: Any, request_id: str | None) -> list[dict[str, Any]]: + if value is None: + return [] + if isinstance(value, str): + try: + value = json.loads(value) + except json.JSONDecodeError as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid query rows.", + request_id=request_id, + ) from exc + if isinstance(value, dict): + return [value] + if not isinstance(value, list): + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid query rows.", + request_id=request_id, + ) + + rows: list[dict[str, Any]] = [] + for item in value: + if isinstance(item, dict): + rows.append(item) + continue + if isinstance(item, str): + try: + parsed = json.loads(item) + except json.JSONDecodeError as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid query rows.", + request_id=request_id, + ) from exc + if isinstance(parsed, dict): + rows.append(parsed) + elif isinstance(parsed, list) and all(isinstance(row, dict) for row in parsed): + rows.extend(parsed) + else: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid query rows.", + request_id=request_id, + ) + else: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid query rows.", + request_id=request_id, + ) + return rows + def _response_too_large(self, request_id: str | None) -> GSFError: return GSFError( GSFErrorCode.RESPONSE_TOO_LARGE, diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index b9ae355e0..956787449 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -31,11 +31,27 @@ class CatalogSearchRequest(GSFRequest): token_budget: int | None = Field(default=None, ge=1) +class ChatCompletionsRequest(GSFRequest): + """Current GSF chat-completions request shared by SQL and prediction flows.""" + + question: str = Field(min_length=1, max_length=4_096) + conversation_id: str | None = None + prediction: bool | None = None + target_db: str | None = None + + +class ChatCompletionResult(GSFResponse): + """Final structured answer extracted from the GSF SSE event stream.""" + + answer: dict[str, Any] + request_id: str | None = None + + class ResultColumn(GSFResponse): """A column in a bounded SQL result.""" name: str - data_type: str + data_type: str | None = None class SemanticContext(GSFResponse): @@ -50,29 +66,30 @@ class SemanticContext(GSFResponse): class TextToSQLRequest(GSFRequest): - """Generate validated SQL and optionally execute it with bounded results.""" + """Generate and execute validated SQL with bounded results.""" question: str = Field(min_length=1, max_length=4_096) database_name: str | None = None - execute: bool = True - object_ids: list[str] = Field(default_factory=list) max_rows: int = Field(default=1_000, ge=1) class TextToSQLResponse(GSFResponse): """Validated SQL, bounded rows, and semantic provenance returned by GSF.""" - request_id: str + request_id: str | None = None + response: str | None = None sql: str - columns: list[ResultColumn] - rows: list[dict[str, Any]] - truncated: bool - objects_used: list[str] = Field(default_factory=list) - joins_used: list[dict[str, Any]] = Field(default_factory=list) - semantic_context: SemanticContext - validation_attempts: list[dict[str, Any]] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - timings: dict[str, int] = Field(default_factory=dict) + columns: list[ResultColumn] = Field(default_factory=list) + rows: list[dict[str, Any]] = Field(default_factory=list) + truncated: bool = False + custom_analyses_used: list[Any] | None = None + objects_used: list[str] | None = None + joins_used: list[dict[str, Any]] | None = None + semantic_context: SemanticContext | None = None + validation_attempts: list[dict[str, Any]] | None = None + assumptions: list[str] | None = None + warnings: list[str] | None = None + timings: dict[str, int | float] | None = None class QueryContextRequest(GSFRequest): @@ -82,21 +99,3 @@ class QueryContextRequest(GSFRequest): database_name: str | None = None object_ids: list[str] = Field(default_factory=list) token_budget: int | None = Field(default=None, ge=1) - - -class QueryContextResponse(GSFResponse): - """Token-budgeted semantic and physical metadata relevant to a question.""" - - request_id: str - tables: list[dict[str, Any]] = Field(default_factory=list) - columns: list[dict[str, Any]] = Field(default_factory=list) - keys: list[dict[str, Any]] = Field(default_factory=list) - join_paths: list[dict[str, Any]] = Field(default_factory=list) - values: list[dict[str, Any]] = Field(default_factory=list) - metrics: list[dict[str, Any]] = Field(default_factory=list) - grain: str | None = None - units: list[str] = Field(default_factory=list) - rules: list[str] = Field(default_factory=list) - omissions: list[str] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - truncated: bool = False diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py index 31433ae13..5eb9e904f 100644 --- a/sources/gsf/src/gsf/register.py +++ b/sources/gsf/src/gsf/register.py @@ -33,7 +33,6 @@ class GSFFunctionGroupConfig(FunctionGroupBaseConfig, name="gsf"): """Shared configuration for AI-Q's GSF tools.""" base_url: HttpUrl - api_version: str = "v1" connect_timeout_seconds: float = Field(default=5.0, gt=0) read_timeout_seconds: float = Field(default=60.0, gt=0) max_retries: int = Field(default=2, ge=0, le=5) @@ -83,10 +82,11 @@ async def catalog_search(request: CatalogSearchRequest) -> str: ) async def text_to_sql(request: TextToSQLRequest) -> str: - """Generate validated SQL and optionally return bounded rows from authorized enterprise data. + """Generate validated SQL and return bounded rows from authorized enterprise data. - Use for an analytical question after the relevant structured-data scope is known. The result contains SQL, - rows, semantic context, warnings, and provenance; AI-Q remains responsible for analysis and synthesis. + Use for an analytical question after the relevant structured-data scope is known. The result contains SQL + and rows, plus semantic context, warnings, and provenance when GSF provides them. AI-Q remains responsible + for analysis and synthesis. """ token = get_auth_token() @@ -111,32 +111,15 @@ async def text_to_sql(request: TextToSQLRequest) -> str: ) async def query_context(request: QueryContextRequest) -> str: - """Build compact, authorized, token-budgeted semantic context for SQL generation. + """Build GSF query context (not available in this integration yet).""" - Use when a downstream SQL-generation step needs relevant tables, columns, keys, joins, metrics, grain, - rules, omissions, and warnings without executing a query. - """ - - token = get_auth_token() - if not token: - return _authentication_error() - try: - result = await client.query_context( - request, - token=token, - trace_headers=_request_trace_headers(), - ) - return result.model_dump_json(exclude_none=True) - except GSFError as error: - return _tool_error(error) - except Exception: - logger.error("Unexpected GSF query-context failure") - return _tool_error( - GSFError( - GSFErrorCode.UPSTREAM_ERROR, - "GSF query context failed.", - ) + del request + return _tool_error( + GSFError( + GSFErrorCode.CAPABILITY_UNAVAILABLE, + "GSF query context is unavailable.", ) + ) group = FunctionGroup(config=config) group.add_function( diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index eb510256b..471cd4d68 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -4,10 +4,26 @@ import pytest +@pytest.fixture +def chat_sql_answer() -> dict: + """Current GSF chat-completions SQL answer envelope.""" + + return { + "response": "Revenue was returned for two quarters.", + "sql_code": "SELECT revenue FROM quarterly_results", + "sql_columns": [], + "custom_analyses_used": [], + "sql_response_from_db": ['[{"revenue":100},{"revenue":200}]'], + } + + @pytest.fixture def text_to_sql_response() -> dict: + """Normalized response used by the NAT registration tests.""" + return { "request_id": "gsf-request-1", + "response": "Revenue was returned for two quarters.", "sql": "SELECT revenue FROM quarterly_results", "columns": [{"name": "revenue", "data_type": "numeric"}], "rows": [{"revenue": 100}, {"revenue": 200}], @@ -26,22 +42,3 @@ def text_to_sql_response() -> dict: "warnings": [], "timings": {"total_ms": 25}, } - - -@pytest.fixture -def query_context_response() -> dict: - return { - "request_id": "gsf-request-2", - "tables": [{"id": "table:quarterly_results", "grain": "quarter"}], - "columns": [{"table_id": "table:quarterly_results", "name": "revenue", "data_type": "numeric"}], - "keys": [{"table_id": "table:quarterly_results", "columns": ["quarter"]}], - "join_paths": [], - "values": [], - "metrics": [{"id": "metric:revenue", "unit": "USD"}], - "grain": "quarter", - "units": ["USD"], - "rules": [], - "omissions": [], - "warnings": [], - "truncated": False, - } diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index 0e751f970..ef5bb8164 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -10,21 +10,37 @@ from gsf.client import GSFClient from gsf.errors import GSFError from gsf.errors import GSFErrorCode -from gsf.models import QueryContextRequest +from gsf.models import ChatCompletionsRequest from gsf.models import TextToSQLRequest +def _sse_response(answer: dict) -> httpx.Response: + events = [ + 'data: {"type":"step","node":"construct_sql_from_candidates"}', + "", + f"data: {json.dumps({'type': 'result', 'answer': answer})}", + "", + "data: [DONE]", + "", + ] + return httpx.Response( + 200, + content="\n".join(events).encode(), + headers={"content-type": "text/event-stream", "x-request-id": "header-request"}, + ) + + @pytest.mark.asyncio -async def test_text_to_sql_sends_scoped_request_and_bounds_rows(text_to_sql_response: dict) -> None: +async def test_text_to_sql_maps_database_to_target_db_and_bounds_rows(chat_sql_answer: dict) -> None: seen_request: httpx.Request | None = None async def handler(request: httpx.Request) -> httpx.Response: nonlocal seen_request seen_request = request - return httpx.Response(200, json=text_to_sql_response, headers={"x-request-id": "header-request"}) + return _sse_response(chat_sql_answer) client = GSFClient( - base_url="https://gsf.example", + base_url="https://gsf.example/", default_max_rows=1, transport=httpx.MockTransport(handler), ) @@ -36,42 +52,39 @@ async def handler(request: httpx.Request) -> httpx.Response: ) assert seen_request is not None - assert seen_request.url == "https://gsf.example/api/v1/text-to-sql" + assert seen_request.url == "https://gsf.example/api/chat/completions" assert seen_request.headers["authorization"] == "Bearer user-token" + assert seen_request.headers["accept"] == "text/event-stream" assert seen_request.headers["traceparent"] == "00-trace" assert json.loads(seen_request.content) == { "question": "Show revenue", - "database_name": "benchmark_db", - "execute": True, - "object_ids": [], - "max_rows": 1, + "prediction": False, + "target_db": "benchmark_db", } + assert [column.name for column in result.columns] == ["revenue"] assert result.rows == [{"revenue": 100}] assert result.truncated is True + assert result.response == "Revenue was returned for two quarters." @pytest.mark.asyncio -async def test_query_context_omits_database_and_unwraps_data(query_context_response: dict) -> None: +async def test_chat_completions_omits_target_db_and_supports_prediction(chat_sql_answer: dict) -> None: seen_payload: dict | None = None async def handler(request: httpx.Request) -> httpx.Response: nonlocal seen_payload seen_payload = json.loads(request.content) - return httpx.Response(200, json={"data": query_context_response}) + return _sse_response(chat_sql_answer) - client = GSFClient(base_url="https://gsf.example/", transport=httpx.MockTransport(handler)) + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: - result = await client.query_context( - QueryContextRequest(question="What revenue data is available?", token_budget=2_000), + result = await client.chat_completions( + ChatCompletionsRequest(question="Who will purchase next?", prediction=True), token="user-token", ) - assert seen_payload == { - "question": "What revenue data is available?", - "object_ids": [], - "token_budget": 2_000, - } - assert result.request_id == "gsf-request-2" + assert seen_payload == {"question": "Who will purchase next?", "prediction": True} + assert result.answer == chat_sql_answer @pytest.mark.asyncio @@ -82,14 +95,14 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: with pytest.raises(GSFError) as raised: - await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.FORBIDDEN assert "secret" not in raised.value.message @pytest.mark.asyncio -async def test_client_retries_rate_limit_then_succeeds(query_context_response: dict) -> None: +async def test_client_retries_rate_limit_then_succeeds(chat_sql_answer: dict) -> None: attempts = 0 async def handler(_request: httpx.Request) -> httpx.Response: @@ -97,21 +110,21 @@ async def handler(_request: httpx.Request) -> httpx.Response: attempts += 1 if attempts == 1: return httpx.Response(429) - return httpx.Response(200, json=query_context_response) + return _sse_response(chat_sql_answer) client = GSFClient(base_url="https://gsf.example", max_retries=1, transport=httpx.MockTransport(handler)) with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: async with client: - await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") assert attempts == 2 sleep.assert_awaited_once_with(1) @pytest.mark.asyncio -async def test_client_rejects_oversized_response(query_context_response: dict) -> None: +async def test_client_rejects_oversized_response(chat_sql_answer: dict) -> None: async def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json=query_context_response) + return _sse_response(chat_sql_answer) client = GSFClient( base_url="https://gsf.example", @@ -120,7 +133,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: ) async with client: with pytest.raises(GSFError) as raised: - await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.RESPONSE_TOO_LARGE @@ -133,6 +146,24 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: with pytest.raises(GSFError) as raised: - await client.query_context(QueryContextRequest(question="Show data"), token="user-token") + await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.INVALID_RESPONSE + + +@pytest.mark.asyncio +async def test_client_maps_sse_error_without_leaking_message() -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content=b'data: {"type":"error","message":"secret database details"}\n\n', + headers={"content-type": "text/event-stream"}, + ) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + with pytest.raises(GSFError) as raised: + await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + + assert raised.value.code is GSFErrorCode.UPSTREAM_ERROR + assert "secret" not in raised.value.message diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py index ca935171b..52a725352 100644 --- a/sources/gsf/tests/test_models.py +++ b/sources/gsf/tests/test_models.py @@ -2,8 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from gsf.models import ChatCompletionsRequest from gsf.models import QueryContextRequest -from gsf.models import QueryContextResponse from gsf.models import TextToSQLRequest from gsf.models import TextToSQLResponse from pydantic import ValidationError @@ -13,10 +13,23 @@ def test_text_to_sql_request_supports_optional_database_name() -> None: request = TextToSQLRequest(question="Show quarterly revenue", database_name="benchmark_db") assert request.database_name == "benchmark_db" - assert request.execute is True assert request.max_rows == 1_000 +def test_chat_request_uses_gsf_target_db_contract() -> None: + payload = ChatCompletionsRequest( + question="Show quarterly revenue", + prediction=False, + target_db="benchmark_db", + ).model_dump(exclude_none=True) + + assert payload == { + "question": "Show quarterly revenue", + "prediction": False, + "target_db": "benchmark_db", + } + + def test_query_context_request_omits_optional_database_name() -> None: payload = QueryContextRequest(question="What revenue data is available?").model_dump(exclude_none=True) @@ -28,17 +41,23 @@ def test_requests_reject_unknown_fields() -> None: TextToSQLRequest.model_validate({"question": "Show revenue", "unknown": True}) -def test_text_to_sql_response_requires_provenance(text_to_sql_response: dict) -> None: - text_to_sql_response.pop("request_id") +def test_text_to_sql_response_accepts_missing_future_enrichments() -> None: + result = TextToSQLResponse.model_validate( + { + "sql": "SELECT revenue FROM quarterly_results", + "rows": [{"revenue": 100}], + } + ) - with pytest.raises(ValidationError): - TextToSQLResponse.model_validate(text_to_sql_response) + assert result.request_id is None + assert result.semantic_context is None + assert result.warnings is None -def test_query_context_response_accepts_future_fields(query_context_response: dict) -> None: - query_context_response["future_gsf_metadata"] = {"enabled": True} +def test_text_to_sql_response_ignores_future_fields(text_to_sql_response: dict) -> None: + text_to_sql_response["future_gsf_metadata"] = {"enabled": True} - result = QueryContextResponse.model_validate(query_context_response) + result = TextToSQLResponse.model_validate(text_to_sql_response) - assert result.request_id == "gsf-request-2" + assert result.request_id == "gsf-request-1" assert not hasattr(result, "future_gsf_metadata") diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index 6710090d1..f8b992aa3 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -7,7 +7,6 @@ from unittest.mock import patch import pytest -from gsf.models import QueryContextResponse from gsf.models import TextToSQLResponse from gsf.register import GSFFunctionGroupConfig from gsf.register import gsf_function_group @@ -25,27 +24,22 @@ async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) - @pytest.mark.asyncio -async def test_group_exposes_only_requested_tools(text_to_sql_response: dict, query_context_response: dict) -> None: +async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> None: client = MagicMock() client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) - client.query_context = AsyncMock(return_value=QueryContextResponse.model_validate(query_context_response)) - config = GSFFunctionGroupConfig( - base_url="https://gsf.example", - include=["text_to_sql", "query_context"], - ) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): async with gsf_function_group(config, MagicMock()) as group: tools = await group.get_accessible_functions() - assert set(tools) == {"gsf__query_context", "gsf__text_to_sql"} + assert set(tools) == {"gsf__text_to_sql"} @pytest.mark.asyncio async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: dict) -> None: client = MagicMock() client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) - client.query_context = AsyncMock() config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) with ( @@ -69,34 +63,40 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d async def test_missing_authentication_fails_closed() -> None: client = MagicMock() client.text_to_sql = AsyncMock() - client.query_context = AsyncMock() - config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["query_context"]) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) with ( patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), patch("gsf.register.get_auth_token", return_value=None), ): async with gsf_function_group(config, MagicMock()) as group: - tool = (await group.get_accessible_functions())["gsf__query_context"] + tool = (await group.get_accessible_functions())["gsf__text_to_sql"] result = json.loads(await tool.ainvoke({"question": "Show data"})) assert result["code"] == "authentication_required" - client.query_context.assert_not_awaited() + client.text_to_sql.assert_not_awaited() @pytest.mark.asyncio -async def test_catalog_search_is_explicitly_unavailable() -> None: +@pytest.mark.parametrize( + ("tool_name", "tool_input", "message"), + [ + ("catalog_search", {"question": "Find revenue metrics"}, "GSF catalog search is unavailable."), + ("query_context", {"question": "Show data"}, "GSF query context is unavailable."), + ], +) +async def test_placeholder_tools_are_explicitly_unavailable(tool_name: str, tool_input: dict, message: str) -> None: client = MagicMock() - config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["catalog_search"]) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=[tool_name]) with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): async with gsf_function_group(config, MagicMock()) as group: - tool = (await group.get_accessible_functions())["gsf__catalog_search"] - result = json.loads(await tool.ainvoke({"question": "Find revenue metrics"})) + tool = (await group.get_accessible_functions())[f"gsf__{tool_name}"] + result = json.loads(await tool.ainvoke(tool_input)) assert result == { "status": "error", "code": "capability_unavailable", "retryable": False, - "message": "GSF catalog search is unavailable.", + "message": message, } From 1fcb85d7a8df967db1d6f81f1283291b09d2d16f Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Wed, 5 Aug 2026 16:19:26 -0700 Subject: [PATCH 03/12] Extend GSF query tools and local authentication Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 39 ++++-- sources/gsf/src/gsf/client.py | 190 ++++++++++++++++++++++++----- sources/gsf/src/gsf/models.py | 36 +++--- sources/gsf/src/gsf/register.py | 66 ++++++++-- sources/gsf/tests/conftest.py | 43 +++++++ sources/gsf/tests/test_client.py | 148 ++++++++++++++++++++-- sources/gsf/tests/test_models.py | 25 ++-- sources/gsf/tests/test_register.py | 81 ++++++++++++ 8 files changed, 544 insertions(+), 84 deletions(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index f1de4cb90..748da115a 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -9,13 +9,14 @@ This package exposes NVIDIA Generative Semantic Fabric (GSF) capabilities as a NeMo Agent Toolkit function group. The current implementation provides: - `gsf__text_to_sql` +- `gsf__text_to_pql` `gsf__catalog_search` and `gsf__query_context` are registered as explicit `capability_unavailable` placeholders until their GSF API contracts are ready. -The function group owns one shared HTTP connection pool. Authentication remains -request-scoped: each tool invocation obtains the current AI-Q user token and -passes it to GSF without storing it on the client. +By default, the function group owns one shared HTTP connection pool and keeps +authentication request-scoped: each tool invocation obtains the current AI-Q +user token and passes it to GSF without storing it on the client. `GSF_BASE_URL` must point to GSF's auth-aware API origin. ```yaml @@ -25,6 +26,7 @@ function_groups: base_url: ${GSF_BASE_URL:-http://gsf:3000} include: - text_to_sql + - text_to_pql functions: data_sources: @@ -41,9 +43,32 @@ functions: - gsf ``` -Text-to-SQL uses GSF's `/api/chat/completions` SSE endpoint with -`prediction: false`. Its optional AI-Q `database_name` input is sent to GSF as -`target_db`, selecting an existing GSF connection rather than creating one. -Future prediction tools can reuse the client transport with `prediction: true`. +For local development and automated evaluation without an incoming AI-Q user +token, explicitly configure a GSF password session. The credentials must come +from environment variables and are used only when a tool is invoked: + +```yaml +function_groups: + gsf: + _type: gsf + base_url: ${GSF_BASE_URL:-http://gsf:3000} + auth: + mode: password + email: ${GSF_EMAIL} + password: ${GSF_PASSWORD} + include: + - text_to_sql + - text_to_pql +``` + +When `auth` is omitted, the existing request-scoped AI-Q user-token flow is +used. Password mode creates one GSF session when the function group starts, +reuses its cookie for local development or evaluation calls, and signs out when +the group closes. The client does not fall back between authentication methods. + +Text-to-SQL and text-to-PQL use GSF's `/api/chat/completions` SSE endpoint with +`prediction: false` for SQL and `prediction: true` for PQL. Their optional AI-Q +`database_name` input is sent to GSF as `target_db`, selecting an existing GSF +connection rather than creating one. The adapter normalizes GSF's current response fields while preserving optional semantic and benchmarking fields as they become available. diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index d825ea640..3f24e307f 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -5,22 +5,28 @@ import asyncio import json +import logging from collections.abc import Mapping from typing import Any import httpx +from pydantic import SecretStr from pydantic import ValidationError from .errors import GSFError from .errors import GSFErrorCode -from .models import ChatCompletionResult -from .models import ChatCompletionsRequest from .models import ResultColumn +from .models import TextToPQLRequest +from .models import TextToPQLResponse from .models import TextToSQLRequest from .models import TextToSQLResponse _FORWARDED_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) _RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) +_PASSWORD_SIGN_IN_PATH = "api/auth/sign-in/email" # pragma: allowlist secret +_PASSWORD_SIGN_OUT_PATH = "api/auth/sign-out" # pragma: allowlist secret + +logger = logging.getLogger(__name__) class GSFClient: @@ -35,12 +41,19 @@ def __init__( max_retries: int = 2, max_response_bytes: int = 5_000_000, default_max_rows: int = 1_000, + password_auth_email: str | None = None, + password_auth_password: SecretStr | None = None, transport: httpx.AsyncBaseTransport | None = None, ) -> None: - self._api_base_url = f"{base_url.rstrip('/')}/api" + if (password_auth_email is None) != (password_auth_password is None): + raise ValueError("GSF password authentication requires both email and password") + self._base_url = base_url.rstrip("/") + self._api_base_url = f"{self._base_url}/api" self._max_retries = max_retries self._max_response_bytes = max_response_bytes self._default_max_rows = default_max_rows + self._password_auth_email = password_auth_email + self._password_auth_password = password_auth_password self._transport = transport self._timeout = httpx.Timeout( connect=connect_timeout_seconds, @@ -54,6 +67,7 @@ def __init__( def from_config(cls, config: Any) -> "GSFClient": """Construct a client from the GSF function-group config without importing NAT.""" + password_auth = getattr(config, "auth", None) return cls( base_url=str(config.base_url), connect_timeout_seconds=config.connect_timeout_seconds, @@ -61,69 +75,111 @@ def from_config(cls, config: Any) -> "GSFClient": max_retries=config.max_retries, max_response_bytes=config.max_response_bytes, default_max_rows=config.default_max_rows, + password_auth_email=password_auth.email if password_auth is not None else None, + password_auth_password=password_auth.password if password_auth is not None else None, ) async def __aenter__(self) -> "GSFClient": self._client = httpx.AsyncClient(timeout=self._timeout, transport=self._transport) + if self._password_auth_email is not None: + try: + await self._sign_in_with_password(self._client) + except BaseException: + await self._client.aclose() + self._client = None + raise return self async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: - if self._client is not None: - await self._client.aclose() + client = self._client + if client is not None: + try: + if self._password_auth_email is not None: + await self._sign_out_password_session(client) + finally: + await client.aclose() self._client = None async def text_to_sql( self, request: TextToSQLRequest, *, - token: str, + token: str | None, trace_headers: Mapping[str, str] | None = None, ) -> TextToSQLResponse: """Run the SQL branch of GSF chat completions and normalize its answer.""" max_rows = min(request.max_rows, self._default_max_rows) - result = await self.chat_completions( - ChatCompletionsRequest( - question=request.question, - prediction=False, - target_db=request.database_name, - ), + payload: dict[str, Any] = { + "question": request.question, + "prediction": False, + } + if request.database_name is not None: + payload["target_db"] = request.database_name + + answer, request_id = await self._chat_completions( + payload, token=token, trace_headers=trace_headers, ) - return self._normalize_text_to_sql(result, max_rows=max_rows) + return self._normalize_text_to_sql(answer, request_id=request_id, max_rows=max_rows) - async def chat_completions( + async def text_to_pql( self, - request: ChatCompletionsRequest, + request: TextToPQLRequest, *, - token: str, + token: str | None, trace_headers: Mapping[str, str] | None = None, - ) -> ChatCompletionResult: - """Call the shared GSF chat endpoint and extract its final SSE result.""" + ) -> TextToPQLResponse: + """Run the prediction branch of GSF chat completions and normalize its answer.""" + + payload: dict[str, Any] = { + "question": request.question, + "prediction": True, + } + if request.database_name is not None: + payload["target_db"] = request.database_name + + answer, request_id = await self._chat_completions( + payload, + token=token, + trace_headers=trace_headers, + ) + return self._normalize_text_to_pql(answer, request_id=request_id) + async def _chat_completions( + self, + payload: dict[str, Any], + *, + token: str | None, + trace_headers: Mapping[str, str] | None = None, + ) -> tuple[dict[str, Any], str | None]: body, request_id, content_type = await self._post( "chat/completions", - request.model_dump(exclude_none=True), + payload, token=token, trace_headers=trace_headers, capability="GSF chat completions", accept="text/event-stream", ) - answer = self._parse_chat_answer(body, content_type=content_type, request_id=request_id) - return ChatCompletionResult(answer=answer, request_id=request_id) + return self._parse_chat_answer(body, content_type=content_type, request_id=request_id), request_id async def _post( self, endpoint: str, payload: dict[str, Any], *, - token: str, + token: str | None, trace_headers: Mapping[str, str] | None, capability: str, accept: str = "application/json", ) -> tuple[bytes, str | None, str]: client = self._require_client() + if self._password_auth_email is None and not token: + raise GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF authentication is required.", + ) headers = self._build_headers(token, trace_headers, accept=accept) attempts = self._max_retries + 1 @@ -175,6 +231,46 @@ async def _post( raise AssertionError("unreachable") + async def _sign_in_with_password(self, client: httpx.AsyncClient) -> None: + try: + response = await client.post( + f"{self._base_url}/{_PASSWORD_SIGN_IN_PATH}", + json={ + "email": self._password_auth_email, + "password": self._password_auth_password.get_secret_value(), + }, + ) + except httpx.TimeoutException as exc: + raise GSFError( + GSFErrorCode.TIMEOUT, + "GSF password sign-in timed out.", + retryable=True, + ) from exc + except httpx.TransportError as exc: + raise GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF password sign-in could not reach GSF.", + retryable=True, + ) from exc + + if response.status_code >= 400: + raise GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF password sign-in was rejected.", + ) + + async def _sign_out_password_session(self, client: httpx.AsyncClient) -> None: + try: + response = await client.post( + f"{self._base_url}/{_PASSWORD_SIGN_OUT_PATH}", + json={}, + headers={"Origin": self._base_url, "Referer": f"{self._base_url}/"}, + ) + if response.status_code >= 400: + logger.warning("GSF password session cleanup returned HTTP %s", response.status_code) + except httpx.HTTPError: + logger.warning("GSF password session cleanup did not complete") + def _require_client(self) -> httpx.AsyncClient: if self._client is None: raise RuntimeError("GSFClient must be used as an async context manager") @@ -182,16 +278,17 @@ def _require_client(self) -> httpx.AsyncClient: @staticmethod def _build_headers( - token: str, + token: str | None, trace_headers: Mapping[str, str] | None, *, accept: str, ) -> dict[str, str]: headers = { "Accept": accept, - "Authorization": f"Bearer {token}", "Content-Type": "application/json", } + if token: + headers["Authorization"] = f"Bearer {token}" for name, value in (trace_headers or {}).items(): if name.lower() in _FORWARDED_HEADER_NAMES and value: headers[name] = value @@ -273,17 +370,22 @@ def _answer_from_event(payload: Any, *, request_id: str | None) -> dict[str, Any ) @classmethod - def _normalize_text_to_sql(cls, result: ChatCompletionResult, *, max_rows: int) -> TextToSQLResponse: - answer = result.answer + def _normalize_text_to_sql( + cls, + answer: Mapping[str, Any], + *, + request_id: str | None, + max_rows: int, + ) -> TextToSQLResponse: sql = answer.get("sql") or answer.get("sql_code") if not isinstance(sql, str) or not sql.strip(): raise GSFError( GSFErrorCode.INVALID_RESPONSE, "GSF did not return validated SQL.", - request_id=result.request_id, + request_id=request_id, ) - rows = cls._normalize_rows(answer.get("rows", answer.get("sql_response_from_db")), result.request_id) + rows = cls._normalize_rows(answer.get("rows", answer.get("sql_response_from_db")), request_id) upstream_truncated = bool(answer.get("truncated", False)) truncated = upstream_truncated or len(rows) > max_rows rows = rows[:max_rows] @@ -293,7 +395,7 @@ def _normalize_text_to_sql(cls, result: ChatCompletionResult, *, max_rows: int) try: return TextToSQLResponse( - request_id=answer.get("request_id") or result.request_id, + request_id=answer.get("request_id") or request_id, response=answer.get("response"), sql=sql, columns=columns, @@ -312,7 +414,35 @@ def _normalize_text_to_sql(cls, result: ChatCompletionResult, *, max_rows: int) raise GSFError( GSFErrorCode.INVALID_RESPONSE, "GSF returned invalid text-to-SQL data.", - request_id=result.request_id, + request_id=request_id, + ) from exc + + @classmethod + def _normalize_text_to_pql(cls, answer: Mapping[str, Any], *, request_id: str | None) -> TextToPQLResponse: + pql = answer.get("pql") or answer.get("pql_code") + if not isinstance(pql, str) or not pql.strip(): + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF did not return validated PQL.", + request_id=request_id, + ) + + try: + return TextToPQLResponse( + request_id=answer.get("request_id") or request_id, + response=answer.get("response"), + pql=pql, + objects_used=answer.get("objects_used"), + semantic_context=answer.get("semantic_context"), + assumptions=answer.get("assumptions"), + warnings=answer.get("warnings"), + timings=answer.get("timings"), + ) + except ValidationError as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid text-to-PQL data.", + request_id=request_id, ) from exc @staticmethod diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index 956787449..10e16c397 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -31,22 +31,6 @@ class CatalogSearchRequest(GSFRequest): token_budget: int | None = Field(default=None, ge=1) -class ChatCompletionsRequest(GSFRequest): - """Current GSF chat-completions request shared by SQL and prediction flows.""" - - question: str = Field(min_length=1, max_length=4_096) - conversation_id: str | None = None - prediction: bool | None = None - target_db: str | None = None - - -class ChatCompletionResult(GSFResponse): - """Final structured answer extracted from the GSF SSE event stream.""" - - answer: dict[str, Any] - request_id: str | None = None - - class ResultColumn(GSFResponse): """A column in a bounded SQL result.""" @@ -73,6 +57,13 @@ class TextToSQLRequest(GSFRequest): max_rows: int = Field(default=1_000, ge=1) +class TextToPQLRequest(GSFRequest): + """Generate validated PQL for prediction workflows.""" + + question: str = Field(min_length=1, max_length=4_096) + database_name: str | None = None + + class TextToSQLResponse(GSFResponse): """Validated SQL, bounded rows, and semantic provenance returned by GSF.""" @@ -92,6 +83,19 @@ class TextToSQLResponse(GSFResponse): timings: dict[str, int | float] | None = None +class TextToPQLResponse(GSFResponse): + """Validated PQL and semantic provenance returned by GSF.""" + + request_id: str | None = None + response: str | None = None + pql: str + objects_used: list[str] | None = None + semantic_context: SemanticContext | None = None + assumptions: list[str] | None = None + warnings: list[str] | None = None + timings: dict[str, int | float] | None = None + + class QueryContextRequest(GSFRequest): """Build compact, authorized context for a later SQL-generation step.""" diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py index 5eb9e904f..81862f6b3 100644 --- a/sources/gsf/src/gsf/register.py +++ b/sources/gsf/src/gsf/register.py @@ -5,9 +5,13 @@ import logging from collections.abc import Mapping +from typing import Literal +from pydantic import BaseModel +from pydantic import ConfigDict from pydantic import Field from pydantic import HttpUrl +from pydantic import SecretStr from aiq_agent.auth.utils import get_auth_token from nat.builder.builder import Builder @@ -22,6 +26,7 @@ from .errors import GSFToolError from .models import CatalogSearchRequest from .models import QueryContextRequest +from .models import TextToPQLRequest from .models import TextToSQLRequest logger = logging.getLogger(__name__) @@ -29,10 +34,21 @@ _TRACE_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) +class GSFPasswordAuthConfig(BaseModel): + """Explicit GSF password-session configuration for development and evaluation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + mode: Literal["password"] + email: str = Field(min_length=1) + password: SecretStr + + class GSFFunctionGroupConfig(FunctionGroupBaseConfig, name="gsf"): """Shared configuration for AI-Q's GSF tools.""" base_url: HttpUrl + auth: GSFPasswordAuthConfig | None = None connect_timeout_seconds: float = Field(default=5.0, gt=0) read_timeout_seconds: float = Field(default=60.0, gt=0) max_retries: int = Field(default=2, ge=0, le=5) @@ -44,13 +60,18 @@ def _tool_error(error: GSFError) -> str: return GSFToolError.from_exception(error).model_dump_json(exclude_none=True) -def _authentication_error() -> str: - return _tool_error( - GSFError( +def _resolve_request_token(config: GSFFunctionGroupConfig) -> str | None: + """Resolve a bearer token unless an explicit password session is configured.""" + + if config.auth is not None: + return None + token = get_auth_token() + if not token: + raise GSFError( GSFErrorCode.AUTHENTICATION_REQUIRED, "GSF authentication is required.", ) - ) + return token def _request_trace_headers() -> Mapping[str, str]: @@ -89,13 +110,10 @@ async def text_to_sql(request: TextToSQLRequest) -> str: for analysis and synthesis. """ - token = get_auth_token() - if not token: - return _authentication_error() try: result = await client.text_to_sql( request, - token=token, + token=_resolve_request_token(config), trace_headers=_request_trace_headers(), ) return result.model_dump_json(exclude_none=True) @@ -110,6 +128,32 @@ async def text_to_sql(request: TextToSQLRequest) -> str: ) ) + async def text_to_pql(request: TextToPQLRequest) -> str: + """Generate validated PQL from an authorized enterprise-data prediction question. + + Use for prediction-style analytical questions after the relevant structured-data scope is known. The + result contains PQL plus semantic context, warnings, and provenance when GSF provides them. AI-Q remains + responsible for analysis and synthesis. + """ + + try: + result = await client.text_to_pql( + request, + token=_resolve_request_token(config), + trace_headers=_request_trace_headers(), + ) + return result.model_dump_json(exclude_none=True) + except GSFError as error: + return _tool_error(error) + except Exception: + logger.error("Unexpected GSF text-to-PQL failure") + return _tool_error( + GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF text-to-PQL failed.", + ) + ) + async def query_context(request: QueryContextRequest) -> str: """Build GSF query context (not available in this integration yet).""" @@ -134,6 +178,12 @@ async def query_context(request: QueryContextRequest) -> str: input_schema=TextToSQLRequest, description=text_to_sql.__doc__, ) + group.add_function( + "text_to_pql", + text_to_pql, + input_schema=TextToPQLRequest, + description=text_to_pql.__doc__, + ) group.add_function( "query_context", query_context, diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index 471cd4d68..275797a6a 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -17,6 +17,27 @@ def chat_sql_answer() -> dict: } +@pytest.fixture +def chat_pql_answer() -> dict: + """Current GSF chat-completions PQL answer envelope.""" + + return { + "response": "A churn prediction query was generated.", + "pql_code": "PREDICT churn FOR customers NEXT 30 DAYS", + "objects_used": ["prediction:churn"], + "semantic_context": { + "metrics": [{"id": "prediction:churn"}], + "grain": "customer", + "units": [], + "filters": [], + "rules": [], + "omissions": [], + }, + "warnings": [], + "timings": {"total_ms": 20}, + } + + @pytest.fixture def text_to_sql_response() -> dict: """Normalized response used by the NAT registration tests.""" @@ -42,3 +63,25 @@ def text_to_sql_response() -> dict: "warnings": [], "timings": {"total_ms": 25}, } + + +@pytest.fixture +def text_to_pql_response() -> dict: + """Normalized PQL response used by the NAT registration tests.""" + + return { + "request_id": "gsf-request-2", + "response": "A churn prediction query was generated.", + "pql": "PREDICT churn FOR customers NEXT 30 DAYS", + "objects_used": ["prediction:churn"], + "semantic_context": { + "metrics": [{"id": "prediction:churn"}], + "grain": "customer", + "units": [], + "filters": [], + "rules": [], + "omissions": [], + }, + "warnings": [], + "timings": {"total_ms": 20}, + } diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index ef5bb8164..e69d290a2 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -10,8 +10,11 @@ from gsf.client import GSFClient from gsf.errors import GSFError from gsf.errors import GSFErrorCode -from gsf.models import ChatCompletionsRequest +from gsf.models import TextToPQLRequest from gsf.models import TextToSQLRequest +from pydantic import SecretStr + +_TEST_PASSWORD = "${TEST_GSF_PASSWORD}" def _sse_response(answer: dict) -> httpx.Response: @@ -68,7 +71,7 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio -async def test_chat_completions_omits_target_db_and_supports_prediction(chat_sql_answer: dict) -> None: +async def test_text_to_sql_omits_optional_target_db(chat_sql_answer: dict) -> None: seen_payload: dict | None = None async def handler(request: httpx.Request) -> httpx.Response: @@ -78,13 +81,136 @@ async def handler(request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: - result = await client.chat_completions( - ChatCompletionsRequest(question="Who will purchase next?", prediction=True), + result = await client.text_to_sql( + TextToSQLRequest(question="Show revenue"), + token="user-token", + ) + + assert seen_payload == {"question": "Show revenue", "prediction": False} + assert result.sql == "SELECT revenue FROM quarterly_results" + + +@pytest.mark.asyncio +async def test_text_to_pql_maps_database_to_target_db(chat_pql_answer: dict) -> None: + seen_payload: dict | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal seen_payload + seen_payload = json.loads(request.content) + return _sse_response(chat_pql_answer) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + result = await client.text_to_pql( + TextToPQLRequest(question="Predict churn risk", database_name="benchmark_db"), token="user-token", ) - assert seen_payload == {"question": "Who will purchase next?", "prediction": True} - assert result.answer == chat_sql_answer + assert seen_payload == { + "question": "Predict churn risk", + "prediction": True, + "target_db": "benchmark_db", + } + assert result.pql == "PREDICT churn FOR customers NEXT 30 DAYS" + assert result.response == "A churn prediction query was generated." + + +@pytest.mark.asyncio +async def test_password_auth_logs_in_uses_cookie_and_signs_out(chat_sql_answer: dict) -> None: + seen_paths: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + seen_paths.append(request.url.path) + if request.url.path == "/api/auth/sign-in/email": + assert json.loads(request.content) == { + "email": "developer@example.com", + "password": _TEST_PASSWORD, + } + return httpx.Response( + 200, + json={"user": {"email": "developer@example.com"}}, + headers={"set-cookie": "better-auth.session_token=session-value; Path=/; HttpOnly; SameSite=Lax"}, + ) + if request.url.path == "/api/auth/sign-out": + assert "better-auth.session_token=session-value" in request.headers["cookie"] + assert request.headers["origin"] == "https://gsf.example" + assert request.headers["referer"] == "https://gsf.example/" + return httpx.Response(200, json={"success": True}) + + assert request.url.path == "/api/chat/completions" + assert "better-auth.session_token=session-value" in request.headers["cookie"] + assert "authorization" not in request.headers + return _sse_response(chat_sql_answer) + + client = GSFClient( + base_url="https://gsf.example", + password_auth_email="developer@example.com", # pragma: allowlist secret + password_auth_password=SecretStr(_TEST_PASSWORD), + transport=httpx.MockTransport(handler), + ) + async with client: + result = await client.text_to_sql(TextToSQLRequest(question="Show data"), token=None) + + assert seen_paths == [ + "/api/auth/sign-in/email", + "/api/chat/completions", + "/api/auth/sign-out", + ] + assert result.sql == "SELECT revenue FROM quarterly_results" + + +@pytest.mark.asyncio +async def test_password_session_is_reused_across_tool_calls(chat_sql_answer: dict) -> None: + seen_paths: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + seen_paths.append(request.url.path) + if request.url.path == "/api/auth/sign-in/email": + return httpx.Response( + 200, + json={"user": {"email": "developer@example.com"}}, + headers={"set-cookie": "better-auth.session_token=session-value; Path=/; HttpOnly; SameSite=Lax"}, + ) + if request.url.path == "/api/auth/sign-out": + return httpx.Response(200, json={"success": True}) + + assert "better-auth.session_token=session-value" in request.headers["cookie"] + return _sse_response(chat_sql_answer) + + client = GSFClient( + base_url="https://gsf.example", + password_auth_email="developer@example.com", # pragma: allowlist secret + password_auth_password=SecretStr(_TEST_PASSWORD), + transport=httpx.MockTransport(handler), + ) + async with client: + await client.text_to_sql(TextToSQLRequest(question="First question"), token=None) + await client.text_to_sql(TextToSQLRequest(question="Second question"), token=None) + + assert seen_paths == [ + "/api/auth/sign-in/email", + "/api/chat/completions", + "/api/chat/completions", + "/api/auth/sign-out", + ] + + +@pytest.mark.asyncio +async def test_client_without_password_auth_requires_bearer_before_http() -> None: + calls = 0 + + async def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(500) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + with pytest.raises(GSFError) as raised: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token=None) + + assert raised.value.code is GSFErrorCode.AUTHENTICATION_REQUIRED + assert calls == 0 @pytest.mark.asyncio @@ -95,7 +221,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: with pytest.raises(GSFError) as raised: - await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.FORBIDDEN assert "secret" not in raised.value.message @@ -115,7 +241,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", max_retries=1, transport=httpx.MockTransport(handler)) with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: async with client: - await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert attempts == 2 sleep.assert_awaited_once_with(1) @@ -133,7 +259,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: ) async with client: with pytest.raises(GSFError) as raised: - await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.RESPONSE_TOO_LARGE @@ -146,7 +272,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: with pytest.raises(GSFError) as raised: - await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.INVALID_RESPONSE @@ -163,7 +289,7 @@ async def handler(_request: httpx.Request) -> httpx.Response: client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) async with client: with pytest.raises(GSFError) as raised: - await client.chat_completions(ChatCompletionsRequest(question="Show data"), token="user-token") + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert raised.value.code is GSFErrorCode.UPSTREAM_ERROR assert "secret" not in raised.value.message diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py index 52a725352..3d44fce37 100644 --- a/sources/gsf/tests/test_models.py +++ b/sources/gsf/tests/test_models.py @@ -2,8 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import pytest -from gsf.models import ChatCompletionsRequest from gsf.models import QueryContextRequest +from gsf.models import TextToPQLRequest +from gsf.models import TextToPQLResponse from gsf.models import TextToSQLRequest from gsf.models import TextToSQLResponse from pydantic import ValidationError @@ -16,18 +17,10 @@ def test_text_to_sql_request_supports_optional_database_name() -> None: assert request.max_rows == 1_000 -def test_chat_request_uses_gsf_target_db_contract() -> None: - payload = ChatCompletionsRequest( - question="Show quarterly revenue", - prediction=False, - target_db="benchmark_db", - ).model_dump(exclude_none=True) +def test_text_to_pql_request_supports_optional_database_name() -> None: + request = TextToPQLRequest(question="Predict churn", database_name="benchmark_db") - assert payload == { - "question": "Show quarterly revenue", - "prediction": False, - "target_db": "benchmark_db", - } + assert request.database_name == "benchmark_db" def test_query_context_request_omits_optional_database_name() -> None: @@ -54,6 +47,14 @@ def test_text_to_sql_response_accepts_missing_future_enrichments() -> None: assert result.warnings is None +def test_text_to_pql_response_accepts_missing_future_enrichments() -> None: + result = TextToPQLResponse.model_validate({"pql": "PREDICT churn FOR customers NEXT 30 DAYS"}) + + assert result.request_id is None + assert result.semantic_context is None + assert result.warnings is None + + def test_text_to_sql_response_ignores_future_fields(text_to_sql_response: dict) -> None: text_to_sql_response["future_gsf_metadata"] = {"enabled": True} diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index f8b992aa3..7f2efec5d 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -7,10 +7,14 @@ from unittest.mock import patch import pytest +from gsf.models import TextToPQLResponse from gsf.models import TextToSQLResponse from gsf.register import GSFFunctionGroupConfig +from gsf.register import GSFPasswordAuthConfig from gsf.register import gsf_function_group +_TEST_PASSWORD = "${TEST_GSF_PASSWORD}" + class FakeClientContext: def __init__(self, client: MagicMock) -> None: @@ -23,6 +27,23 @@ async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) - return None +def test_password_auth_is_optional_and_keeps_secret_wrapped() -> None: + default_config = GSFFunctionGroupConfig(base_url="https://gsf.example") + password_config = GSFFunctionGroupConfig( + base_url="https://gsf.example", + auth={ + "mode": "password", + "email": "developer@example.com", + "password": _TEST_PASSWORD, + }, + ) + + assert default_config.auth is None + assert isinstance(password_config.auth, GSFPasswordAuthConfig) + assert password_config.auth.password.get_secret_value() == _TEST_PASSWORD + assert _TEST_PASSWORD not in repr(password_config) + + @pytest.mark.asyncio async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> None: client = MagicMock() @@ -36,6 +57,19 @@ async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> assert set(tools) == {"gsf__text_to_sql"} +@pytest.mark.asyncio +async def test_group_exposes_text_to_pql_when_requested(text_to_pql_response: dict) -> None: + client = MagicMock() + client.text_to_pql = AsyncMock(return_value=TextToPQLResponse.model_validate(text_to_pql_response)) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_pql"]) + + with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): + async with gsf_function_group(config, MagicMock()) as group: + tools = await group.get_accessible_functions() + + assert set(tools) == {"gsf__text_to_pql"} + + @pytest.mark.asyncio async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: dict) -> None: client = MagicMock() @@ -59,6 +93,53 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d assert "token" not in client.__dict__ +@pytest.mark.asyncio +async def test_text_to_pql_resolves_token_per_invocation(text_to_pql_response: dict) -> None: + client = MagicMock() + client.text_to_pql = AsyncMock(return_value=TextToPQLResponse.model_validate(text_to_pql_response)) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_pql"]) + + with ( + patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), + patch("gsf.register.get_auth_token", return_value="token-one"), + patch("gsf.register._request_trace_headers", return_value={}), + ): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__text_to_pql"] + result = json.loads(await tool.ainvoke({"question": "Predict churn"})) + + assert result["request_id"] == "gsf-request-2" + assert client.text_to_pql.await_args.kwargs["token"] == "token-one" + + +@pytest.mark.asyncio +async def test_explicit_password_auth_does_not_resolve_user_token(text_to_sql_response: dict) -> None: + client = MagicMock() + client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) + config = GSFFunctionGroupConfig( + base_url="https://gsf.example", + include=["text_to_sql"], + auth={ + "mode": "password", + "email": "developer@example.com", + "password": _TEST_PASSWORD, + }, + ) + + with ( + patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), + patch("gsf.register.get_auth_token") as get_auth_token, + patch("gsf.register._request_trace_headers", return_value={}), + ): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__text_to_sql"] + result = json.loads(await tool.ainvoke({"question": "Show data"})) + + assert result["request_id"] == "gsf-request-1" + get_auth_token.assert_not_called() + assert client.text_to_sql.await_args.kwargs["token"] is None + + @pytest.mark.asyncio async def test_missing_authentication_fails_closed() -> None: client = MagicMock() From d4c15eeb6373c646b1c9310641cdf86585888713 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Wed, 5 Aug 2026 16:19:42 -0700 Subject: [PATCH 04/12] Validate function group registry references Signed-off-by: Soumili Nandi --- .../aiq-configure-workflow/scripts/validate_config.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.agents/skills/aiq-configure-workflow/scripts/validate_config.py b/.agents/skills/aiq-configure-workflow/scripts/validate_config.py index f26013931..bb7651916 100644 --- a/.agents/skills/aiq-configure-workflow/scripts/validate_config.py +++ b/.agents/skills/aiq-configure-workflow/scripts/validate_config.py @@ -119,7 +119,8 @@ def _validate_registry(registry: dict, declared_functions: set[str], errors: lis for tool in tools: if tool not in declared_functions: errors.append( - f"source '{label}' lists tool '{tool}' in its tools:, but '{tool}' is not declared under functions:" + f"source '{label}' lists tool '{tool}' in its tools:, but '{tool}' is not declared under " + "functions: or function_groups:" ) for bool_field in ("default_enabled", "requires_auth"): if bool_field in source and not isinstance(source[bool_field], bool): @@ -228,6 +229,10 @@ def validate(path: str) -> int: functions = {} declared_functions = set(functions.keys()) + function_groups = data.get("function_groups") or {} + if not isinstance(function_groups, dict): + function_groups = {} + for field, alias in _iter_refs(functions): if alias not in defined_aliases: defined = ", ".join(sorted(defined_aliases)) or "none" @@ -247,7 +252,7 @@ def validate(path: str) -> int: if registry is None: warnings.append("no data_source_registry function found (fine for minimal configs).") else: - _validate_registry(registry, declared_functions, errors, warnings) + _validate_registry(registry, declared_functions | set(function_groups), errors, warnings) workflow = data.get("workflow") for agent_name in REQUIRED_WORKFLOW_AGENTS: From d528a731cf99445657076a8b8b7042700df1f1c4 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Wed, 5 Aug 2026 17:15:01 -0700 Subject: [PATCH 05/12] Handle GSF prediction PQL response Signed-off-by: Soumili Nandi --- sources/gsf/src/gsf/client.py | 4 +++- sources/gsf/tests/conftest.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index 3f24e307f..ac95f2141 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -419,7 +419,9 @@ def _normalize_text_to_sql( @classmethod def _normalize_text_to_pql(cls, answer: Mapping[str, Any], *, request_id: str | None) -> TextToPQLResponse: - pql = answer.get("pql") or answer.get("pql_code") + # GSF's prediction branch currently returns the PQL in ``sql_code`` so + # its frontend can render SQL and prediction results with one shape. + pql = answer.get("pql") or answer.get("pql_code") or answer.get("sql_code") if not isinstance(pql, str) or not pql.strip(): raise GSFError( GSFErrorCode.INVALID_RESPONSE, diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index 275797a6a..c1b084ced 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -23,7 +23,7 @@ def chat_pql_answer() -> dict: return { "response": "A churn prediction query was generated.", - "pql_code": "PREDICT churn FOR customers NEXT 30 DAYS", + "sql_code": "PREDICT churn FOR customers NEXT 30 DAYS", "objects_used": ["prediction:churn"], "semantic_context": { "metrics": [{"id": "prediction:churn"}], From ef0b8d409981a95547710529a95587a3cb3d6a32 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 12:58:23 -0700 Subject: [PATCH 06/12] Add GSF catalog search capability Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 13 ++++- sources/gsf/src/gsf/client.py | 80 ++++++++++++++++++++++++++++++ sources/gsf/src/gsf/models.py | 23 ++++++++- sources/gsf/src/gsf/register.py | 28 ++++++++--- sources/gsf/tests/conftest.py | 44 ++++++++++++++++ sources/gsf/tests/test_client.py | 59 ++++++++++++++++++++++ sources/gsf/tests/test_models.py | 40 +++++++++++++++ sources/gsf/tests/test_register.py | 23 ++++++++- 8 files changed, 297 insertions(+), 13 deletions(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index 748da115a..94a60d91c 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -10,9 +10,10 @@ NeMo Agent Toolkit function group. The current implementation provides: - `gsf__text_to_sql` - `gsf__text_to_pql` +- `gsf__catalog_search` -`gsf__catalog_search` and `gsf__query_context` are registered as explicit -`capability_unavailable` placeholders until their GSF API contracts are ready. +`gsf__query_context` is registered as an explicit `capability_unavailable` +placeholder until its GSF API contract is ready. By default, the function group owns one shared HTTP connection pool and keeps authentication request-scoped: each tool invocation obtains the current AI-Q @@ -25,6 +26,7 @@ function_groups: _type: gsf base_url: ${GSF_BASE_URL:-http://gsf:3000} include: + - catalog_search - text_to_sql - text_to_pql @@ -57,6 +59,7 @@ function_groups: email: ${GSF_EMAIL} password: ${GSF_PASSWORD} include: + - catalog_search - text_to_sql - text_to_pql ``` @@ -72,3 +75,9 @@ Text-to-SQL and text-to-PQL use GSF's `/api/chat/completions` SSE endpoint with connection rather than creating one. The adapter normalizes GSF's current response fields while preserving optional semantic and benchmarking fields as they become available. + +Catalog search uses `POST /api/question-entity-coverage` and returns entity +coverage plus ranked semantic candidates for DS-agent grounding and routing. +Its optional `database_name` is sent as `target_db`. The current GSF endpoint +does not enforce that selector yet, so database-scoped search depends on a GSF +contract update. diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index ac95f2141..f9c83a85d 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -15,6 +15,8 @@ from .errors import GSFError from .errors import GSFErrorCode +from .models import CatalogSearchRequest +from .models import CatalogSearchResponse from .models import ResultColumn from .models import TextToPQLRequest from .models import TextToPQLResponse @@ -124,6 +126,32 @@ async def text_to_sql( ) return self._normalize_text_to_sql(answer, request_id=request_id, max_rows=max_rows) + async def catalog_search( + self, + request: CatalogSearchRequest, + *, + token: str | None, + trace_headers: Mapping[str, str] | None = None, + ) -> CatalogSearchResponse: + """Find GSF semantic candidates and measure entity coverage.""" + + payload: dict[str, Any] = { + "question": request.question, + "max_distance": request.max_distance, + } + if request.database_name is not None: + payload["target_db"] = request.database_name + + body, request_id, _content_type = await self._post( + "question-entity-coverage", + payload, + token=token, + trace_headers=trace_headers, + capability="GSF entity coverage", + ) + data = self._parse_json_data(body, request_id=request_id) + return self._normalize_catalog_search(data, request_id=request_id, max_results=request.max_results) + async def text_to_pql( self, request: TextToPQLRequest, @@ -357,6 +385,27 @@ def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | request_id=request_id, ) + @staticmethod + def _parse_json_data(body: bytes, *, request_id: str | None) -> dict[str, Any]: + try: + payload = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned an invalid response.", + request_id=request_id, + ) from exc + + if isinstance(payload, dict) and isinstance(payload.get("data"), dict): + payload = payload["data"] + if not isinstance(payload, dict): + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned an invalid response.", + request_id=request_id, + ) + return payload + @staticmethod def _answer_from_event(payload: Any, *, request_id: str | None) -> dict[str, Any]: if isinstance(payload, dict) and isinstance(payload.get("answer"), dict): @@ -417,6 +466,37 @@ def _normalize_text_to_sql( request_id=request_id, ) from exc + @staticmethod + def _normalize_catalog_search( + data: Mapping[str, Any], + *, + request_id: str | None, + max_results: int, + ) -> CatalogSearchResponse: + candidates = data.get("candidates") + if not isinstance(candidates, list): + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid catalog-search data.", + request_id=request_id, + ) + + truncated = len(candidates) > max_results + try: + return CatalogSearchResponse( + request_id=data.get("request_id") or request_id, + coverage=data.get("coverage"), + candidates=candidates[:max_results], + uncovered_entities=data.get("uncovered_entities"), + truncated=truncated, + ) + except ValidationError as exc: + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned invalid catalog-search data.", + request_id=request_id, + ) from exc + @classmethod def _normalize_text_to_pql(cls, answer: Mapping[str, Any], *, request_id: str | None) -> TextToPQLResponse: # GSF's prediction branch currently returns the PQL in ``sql_code`` so diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index 10e16c397..870d18623 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -23,12 +23,31 @@ class GSFResponse(BaseModel): class CatalogSearchRequest(GSFRequest): - """Provisional catalog-search input retained for the unavailable placeholder.""" + """Find semantic candidates relevant to an enterprise-data question.""" question: str = Field(min_length=1, max_length=4_096) database_name: str | None = None max_results: int = Field(default=10, ge=1, le=100) - token_budget: int | None = Field(default=None, ge=1) + max_distance: float = Field(default=0.75, gt=0) + + +class CatalogCandidate(GSFResponse): + """A semantic candidate returned by GSF entity-coverage search.""" + + label: str + attribute: str + term: str + id: str + + +class CatalogSearchResponse(GSFResponse): + """Coverage and ranked semantic candidates returned by GSF.""" + + request_id: str | None = None + coverage: float = Field(ge=0, le=1) + candidates: list[CatalogCandidate] + uncovered_entities: list[str] | None = None + truncated: bool = False class ResultColumn(GSFResponse): diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py index 81862f6b3..e87233b43 100644 --- a/sources/gsf/src/gsf/register.py +++ b/sources/gsf/src/gsf/register.py @@ -92,15 +92,29 @@ async def gsf_function_group(config: GSFFunctionGroupConfig, _builder: Builder): async with GSFClient.from_config(config) as client: async def catalog_search(request: CatalogSearchRequest) -> str: - """Search the GSF enterprise catalog (not available in this integration yet).""" + """Find GSF semantic candidates relevant to an enterprise-data question. - del request - return _tool_error( - GSFError( - GSFErrorCode.CAPABILITY_UNAVAILABLE, - "GSF catalog search is unavailable.", + Use the returned entity coverage and ranked candidates to ground routing and later analytical calls. This + returns catalog context, not a final answer to the user's question. + """ + + try: + result = await client.catalog_search( + request, + token=_resolve_request_token(config), + trace_headers=_request_trace_headers(), + ) + return result.model_dump_json(exclude_none=True) + except GSFError as error: + return _tool_error(error) + except Exception: + logger.error("Unexpected GSF catalog-search failure") + return _tool_error( + GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF catalog search failed.", + ) ) - ) async def text_to_sql(request: TextToSQLRequest) -> str: """Generate validated SQL and return bounded rows from authorized enterprise data. diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index c1b084ced..978d51354 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -4,6 +4,50 @@ import pytest +@pytest.fixture +def catalog_search_api_response() -> dict: + """Current GSF question-entity-coverage response envelope.""" + + return { + "data": { + "coverage": 0.5, + "candidates": [ + { + "label": "ColumnAttribute", + "attribute": "recognized_revenue", + "term": "Revenue", + "id": "attr:revenue", + }, + { + "label": "SqlAttribute", + "attribute": "net_revenue_sql", + "term": "Net Revenue", + "id": "sql-attr:net-revenue", + }, + ], + } + } + + +@pytest.fixture +def catalog_search_response() -> dict: + """Normalized catalog response used by NAT registration tests.""" + + return { + "request_id": "gsf-catalog-request-1", + "coverage": 0.5, + "candidates": [ + { + "label": "ColumnAttribute", + "attribute": "recognized_revenue", + "term": "Revenue", + "id": "attr:revenue", + } + ], + "truncated": False, + } + + @pytest.fixture def chat_sql_answer() -> dict: """Current GSF chat-completions SQL answer envelope.""" diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index e69d290a2..d407f0ada 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -10,6 +10,7 @@ from gsf.client import GSFClient from gsf.errors import GSFError from gsf.errors import GSFErrorCode +from gsf.models import CatalogSearchRequest from gsf.models import TextToPQLRequest from gsf.models import TextToSQLRequest from pydantic import SecretStr @@ -33,6 +34,64 @@ def _sse_response(answer: dict) -> httpx.Response: ) +@pytest.mark.asyncio +async def test_catalog_search_uses_entity_coverage_path_maps_scope_and_bounds_candidates( + catalog_search_api_response: dict, +) -> None: + seen_request: httpx.Request | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal seen_request + seen_request = request + return httpx.Response( + 200, + json=catalog_search_api_response, + headers={"x-request-id": "header-request"}, + ) + + client = GSFClient(base_url="https://gsf.example/", transport=httpx.MockTransport(handler)) + async with client: + result = await client.catalog_search( + CatalogSearchRequest( + question="Find revenue metrics", + database_name="benchmark_db", + max_results=1, + max_distance=0.5, + ), + token="user-token", + trace_headers={"traceparent": "00-trace", "authorization": "do-not-forward"}, + ) + + assert seen_request is not None + assert seen_request.url == "https://gsf.example/api/question-entity-coverage" + assert seen_request.headers["authorization"] == "Bearer user-token" + assert seen_request.headers["accept"] == "application/json" + assert seen_request.headers["traceparent"] == "00-trace" + assert json.loads(seen_request.content) == { + "question": "Find revenue metrics", + "max_distance": 0.5, + "target_db": "benchmark_db", + } + assert result.request_id == "header-request" + assert result.coverage == 0.5 + assert [candidate.id for candidate in result.candidates] == ["attr:revenue"] + assert result.uncovered_entities is None + assert result.truncated is True + + +@pytest.mark.asyncio +async def test_catalog_search_maps_missing_endpoint_to_capability_unavailable() -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(404) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + with pytest.raises(GSFError) as raised: + await client.catalog_search(CatalogSearchRequest(question="Find revenue"), token="user-token") + + assert raised.value.code is GSFErrorCode.CAPABILITY_UNAVAILABLE + + @pytest.mark.asyncio async def test_text_to_sql_maps_database_to_target_db_and_bounds_rows(chat_sql_answer: dict) -> None: seen_request: httpx.Request | None = None diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py index 3d44fce37..2358608f5 100644 --- a/sources/gsf/tests/test_models.py +++ b/sources/gsf/tests/test_models.py @@ -2,6 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from gsf.models import CatalogCandidate +from gsf.models import CatalogSearchRequest +from gsf.models import CatalogSearchResponse from gsf.models import QueryContextRequest from gsf.models import TextToPQLRequest from gsf.models import TextToPQLResponse @@ -10,6 +13,43 @@ from pydantic import ValidationError +def test_catalog_search_request_supports_optional_scope_and_search_controls() -> None: + request = CatalogSearchRequest( + question="Find revenue metrics", + database_name="benchmark_db", + max_results=20, + max_distance=0.5, + ) + + assert request.database_name == "benchmark_db" + assert request.max_results == 20 + assert request.max_distance == 0.5 + + +def test_catalog_search_response_validates_coverage() -> None: + with pytest.raises(ValidationError): + CatalogSearchResponse( + coverage=1.5, + candidates=[ + CatalogCandidate( + label="ColumnAttribute", + attribute="revenue", + term="Revenue", + id="attr:revenue", + ) + ], + ) + + +def test_catalog_search_response_ignores_future_fields(catalog_search_response: dict) -> None: + catalog_search_response["future_gsf_metadata"] = {"enabled": True} + + result = CatalogSearchResponse.model_validate(catalog_search_response) + + assert result.request_id == "gsf-catalog-request-1" + assert not hasattr(result, "future_gsf_metadata") + + def test_text_to_sql_request_supports_optional_database_name() -> None: request = TextToSQLRequest(question="Show quarterly revenue", database_name="benchmark_db") diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index 7f2efec5d..a361f30f0 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -7,6 +7,7 @@ from unittest.mock import patch import pytest +from gsf.models import CatalogSearchResponse from gsf.models import TextToPQLResponse from gsf.models import TextToSQLResponse from gsf.register import GSFFunctionGroupConfig @@ -24,7 +25,7 @@ async def __aenter__(self) -> MagicMock: return self.client async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: - return None + pass def test_password_auth_is_optional_and_keeps_secret_wrapped() -> None: @@ -70,6 +71,25 @@ async def test_group_exposes_text_to_pql_when_requested(text_to_pql_response: di assert set(tools) == {"gsf__text_to_pql"} +@pytest.mark.asyncio +async def test_catalog_search_resolves_token_per_invocation(catalog_search_response: dict) -> None: + client = MagicMock() + client.catalog_search = AsyncMock(return_value=CatalogSearchResponse.model_validate(catalog_search_response)) + config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["catalog_search"]) + + with ( + patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), + patch("gsf.register.get_auth_token", return_value="token-one"), + patch("gsf.register._request_trace_headers", return_value={}), + ): + async with gsf_function_group(config, MagicMock()) as group: + tool = (await group.get_accessible_functions())["gsf__catalog_search"] + result = json.loads(await tool.ainvoke({"question": "Find revenue metrics"})) + + assert result["request_id"] == "gsf-catalog-request-1" + assert client.catalog_search.await_args.kwargs["token"] == "token-one" + + @pytest.mark.asyncio async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: dict) -> None: client = MagicMock() @@ -162,7 +182,6 @@ async def test_missing_authentication_fails_closed() -> None: @pytest.mark.parametrize( ("tool_name", "tool_input", "message"), [ - ("catalog_search", {"question": "Find revenue metrics"}, "GSF catalog search is unavailable."), ("query_context", {"question": "Show data"}, "GSF query context is unavailable."), ], ) From be9a8346e49fb863055085b4216b556e308bcb39 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 13:44:30 -0700 Subject: [PATCH 07/12] Exclude GSF prose from SQL tool output Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 2 ++ sources/gsf/src/gsf/client.py | 1 - sources/gsf/src/gsf/models.py | 1 - sources/gsf/tests/test_client.py | 2 +- sources/gsf/tests/test_register.py | 2 ++ 5 files changed, 5 insertions(+), 3 deletions(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index 94a60d91c..3e9d9c9e9 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -75,6 +75,8 @@ Text-to-SQL and text-to-PQL use GSF's `/api/chat/completions` SSE endpoint with connection rather than creating one. The adapter normalizes GSF's current response fields while preserving optional semantic and benchmarking fields as they become available. +For text-to-SQL, GSF's compatibility prose is discarded; AI-Q consumes the +generated SQL, bounded rows, and any structured semantic provenance instead. Catalog search uses `POST /api/question-entity-coverage` and returns entity coverage plus ranked semantic candidates for DS-agent grounding and routing. diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index f9c83a85d..8cbdd9232 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -445,7 +445,6 @@ def _normalize_text_to_sql( try: return TextToSQLResponse( request_id=answer.get("request_id") or request_id, - response=answer.get("response"), sql=sql, columns=columns, rows=rows, diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index 870d18623..e80bfd607 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -87,7 +87,6 @@ class TextToSQLResponse(GSFResponse): """Validated SQL, bounded rows, and semantic provenance returned by GSF.""" request_id: str | None = None - response: str | None = None sql: str columns: list[ResultColumn] = Field(default_factory=list) rows: list[dict[str, Any]] = Field(default_factory=list) diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index d407f0ada..023c708a1 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -126,7 +126,7 @@ async def handler(request: httpx.Request) -> httpx.Response: assert [column.name for column in result.columns] == ["revenue"] assert result.rows == [{"revenue": 100}] assert result.truncated is True - assert result.response == "Revenue was returned for two quarters." + assert "response" not in result.model_dump() @pytest.mark.asyncio diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index a361f30f0..63a751d62 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -108,6 +108,8 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d assert first["request_id"] == "gsf-request-1" assert second["request_id"] == "gsf-request-1" + assert "response" not in first + assert "response" not in second assert client.text_to_sql.await_args_list[0].kwargs["token"] == "token-one" assert client.text_to_sql.await_args_list[1].kwargs["token"] == "token-two" assert "token" not in client.__dict__ From ed4c1da5175f3b310fc111322d1eacda11083196 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 15:52:29 -0700 Subject: [PATCH 08/12] Preserve GSF SQL thought summaries Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 2 ++ sources/gsf/src/gsf/client.py | 1 + sources/gsf/src/gsf/models.py | 1 + sources/gsf/tests/conftest.py | 3 ++- sources/gsf/tests/test_client.py | 1 + sources/gsf/tests/test_models.py | 1 + sources/gsf/tests/test_register.py | 2 ++ 7 files changed, 10 insertions(+), 1 deletion(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index 3e9d9c9e9..e469c4209 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -77,6 +77,8 @@ The adapter normalizes GSF's current response fields while preserving optional semantic and benchmarking fields as they become available. For text-to-SQL, GSF's compatibility prose is discarded; AI-Q consumes the generated SQL, bounded rows, and any structured semantic provenance instead. +GSF's optional `thoughts` summary is retained as diagnostic context, not as +authoritative evidence. Catalog search uses `POST /api/question-entity-coverage` and returns entity coverage plus ranked semantic candidates for DS-agent grounding and routing. diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index 8cbdd9232..e4492ca10 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -445,6 +445,7 @@ def _normalize_text_to_sql( try: return TextToSQLResponse( request_id=answer.get("request_id") or request_id, + thoughts=answer.get("thoughts"), sql=sql, columns=columns, rows=rows, diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index e80bfd607..58a8ab2bd 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -87,6 +87,7 @@ class TextToSQLResponse(GSFResponse): """Validated SQL, bounded rows, and semantic provenance returned by GSF.""" request_id: str | None = None + thoughts: str | None = None sql: str columns: list[ResultColumn] = Field(default_factory=list) rows: list[dict[str, Any]] = Field(default_factory=list) diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index 978d51354..15e0743cd 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -54,6 +54,7 @@ def chat_sql_answer() -> dict: return { "response": "Revenue was returned for two quarters.", + "thoughts": "- Constructing SQL: Used quarterly_results.", "sql_code": "SELECT revenue FROM quarterly_results", "sql_columns": [], "custom_analyses_used": [], @@ -88,7 +89,7 @@ def text_to_sql_response() -> dict: return { "request_id": "gsf-request-1", - "response": "Revenue was returned for two quarters.", + "thoughts": "- Constructing SQL: Used quarterly_results.", "sql": "SELECT revenue FROM quarterly_results", "columns": [{"name": "revenue", "data_type": "numeric"}], "rows": [{"revenue": 100}, {"revenue": 200}], diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index 023c708a1..87056146c 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -126,6 +126,7 @@ async def handler(request: httpx.Request) -> httpx.Response: assert [column.name for column in result.columns] == ["revenue"] assert result.rows == [{"revenue": 100}] assert result.truncated is True + assert result.thoughts == "- Constructing SQL: Used quarterly_results." assert "response" not in result.model_dump() diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py index 2358608f5..d3cc84152 100644 --- a/sources/gsf/tests/test_models.py +++ b/sources/gsf/tests/test_models.py @@ -83,6 +83,7 @@ def test_text_to_sql_response_accepts_missing_future_enrichments() -> None: ) assert result.request_id is None + assert result.thoughts is None assert result.semantic_context is None assert result.warnings is None diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index 63a751d62..bc00ce37a 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -108,6 +108,8 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d assert first["request_id"] == "gsf-request-1" assert second["request_id"] == "gsf-request-1" + assert first["thoughts"] == "- Constructing SQL: Used quarterly_results." + assert second["thoughts"] == "- Constructing SQL: Used quarterly_results." assert "response" not in first assert "response" not in second assert client.text_to_sql.await_args_list[0].kwargs["token"] == "token-one" From 0d42fd659ba3ed637f05657d6feaa4123caf66f8 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 16:05:28 -0700 Subject: [PATCH 09/12] Remove unused GSF provenance helper Signed-off-by: Soumili Nandi --- sources/gsf/src/gsf/provenance.py | 12 ------------ 1 file changed, 12 deletions(-) delete mode 100644 sources/gsf/src/gsf/provenance.py diff --git a/sources/gsf/src/gsf/provenance.py b/sources/gsf/src/gsf/provenance.py deleted file mode 100644 index 297b25b7c..000000000 --- a/sources/gsf/src/gsf/provenance.py +++ /dev/null @@ -1,12 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Non-sensitive helpers for GSF result provenance.""" - -import hashlib - - -def sql_sha256(sql: str) -> str: - """Return a stable SQL digest without logging or persisting the SQL text.""" - - return hashlib.sha256(sql.encode("utf-8")).hexdigest() From 185af46dcd9b7795c64e40e4514a24064b267d41 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 17:04:26 -0700 Subject: [PATCH 10/12] Disable unvalidated GSF PQL tool Signed-off-by: Soumili Nandi --- sources/gsf/README.md | 12 +++++------ sources/gsf/src/gsf/register.py | 33 ------------------------------ sources/gsf/tests/test_register.py | 33 ------------------------------ 3 files changed, 5 insertions(+), 73 deletions(-) diff --git a/sources/gsf/README.md b/sources/gsf/README.md index e469c4209..272bf0cb1 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -9,11 +9,12 @@ This package exposes NVIDIA Generative Semantic Fabric (GSF) capabilities as a NeMo Agent Toolkit function group. The current implementation provides: - `gsf__text_to_sql` -- `gsf__text_to_pql` - `gsf__catalog_search` `gsf__query_context` is registered as an explicit `capability_unavailable` placeholder until its GSF API contract is ready. +PQL client and model groundwork remains internal, but no PQL tool is registered +until its GSF contract and integration behavior are validated. By default, the function group owns one shared HTTP connection pool and keeps authentication request-scoped: each tool invocation obtains the current AI-Q @@ -28,7 +29,6 @@ function_groups: include: - catalog_search - text_to_sql - - text_to_pql functions: data_sources: @@ -61,7 +61,6 @@ function_groups: include: - catalog_search - text_to_sql - - text_to_pql ``` When `auth` is omitted, the existing request-scoped AI-Q user-token flow is @@ -69,10 +68,9 @@ used. Password mode creates one GSF session when the function group starts, reuses its cookie for local development or evaluation calls, and signs out when the group closes. The client does not fall back between authentication methods. -Text-to-SQL and text-to-PQL use GSF's `/api/chat/completions` SSE endpoint with -`prediction: false` for SQL and `prediction: true` for PQL. Their optional AI-Q -`database_name` input is sent to GSF as `target_db`, selecting an existing GSF -connection rather than creating one. +Text-to-SQL uses GSF's `/api/chat/completions` SSE endpoint with +`prediction: false`. Its optional AI-Q `database_name` input is sent to GSF as +`target_db`, selecting an existing GSF connection rather than creating one. The adapter normalizes GSF's current response fields while preserving optional semantic and benchmarking fields as they become available. For text-to-SQL, GSF's compatibility prose is discarded; AI-Q consumes the diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py index e87233b43..3144baa11 100644 --- a/sources/gsf/src/gsf/register.py +++ b/sources/gsf/src/gsf/register.py @@ -26,7 +26,6 @@ from .errors import GSFToolError from .models import CatalogSearchRequest from .models import QueryContextRequest -from .models import TextToPQLRequest from .models import TextToSQLRequest logger = logging.getLogger(__name__) @@ -142,32 +141,6 @@ async def text_to_sql(request: TextToSQLRequest) -> str: ) ) - async def text_to_pql(request: TextToPQLRequest) -> str: - """Generate validated PQL from an authorized enterprise-data prediction question. - - Use for prediction-style analytical questions after the relevant structured-data scope is known. The - result contains PQL plus semantic context, warnings, and provenance when GSF provides them. AI-Q remains - responsible for analysis and synthesis. - """ - - try: - result = await client.text_to_pql( - request, - token=_resolve_request_token(config), - trace_headers=_request_trace_headers(), - ) - return result.model_dump_json(exclude_none=True) - except GSFError as error: - return _tool_error(error) - except Exception: - logger.error("Unexpected GSF text-to-PQL failure") - return _tool_error( - GSFError( - GSFErrorCode.UPSTREAM_ERROR, - "GSF text-to-PQL failed.", - ) - ) - async def query_context(request: QueryContextRequest) -> str: """Build GSF query context (not available in this integration yet).""" @@ -192,12 +165,6 @@ async def query_context(request: QueryContextRequest) -> str: input_schema=TextToSQLRequest, description=text_to_sql.__doc__, ) - group.add_function( - "text_to_pql", - text_to_pql, - input_schema=TextToPQLRequest, - description=text_to_pql.__doc__, - ) group.add_function( "query_context", query_context, diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index bc00ce37a..d68961e0f 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -8,7 +8,6 @@ import pytest from gsf.models import CatalogSearchResponse -from gsf.models import TextToPQLResponse from gsf.models import TextToSQLResponse from gsf.register import GSFFunctionGroupConfig from gsf.register import GSFPasswordAuthConfig @@ -58,19 +57,6 @@ async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> assert set(tools) == {"gsf__text_to_sql"} -@pytest.mark.asyncio -async def test_group_exposes_text_to_pql_when_requested(text_to_pql_response: dict) -> None: - client = MagicMock() - client.text_to_pql = AsyncMock(return_value=TextToPQLResponse.model_validate(text_to_pql_response)) - config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_pql"]) - - with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): - async with gsf_function_group(config, MagicMock()) as group: - tools = await group.get_accessible_functions() - - assert set(tools) == {"gsf__text_to_pql"} - - @pytest.mark.asyncio async def test_catalog_search_resolves_token_per_invocation(catalog_search_response: dict) -> None: client = MagicMock() @@ -117,25 +103,6 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d assert "token" not in client.__dict__ -@pytest.mark.asyncio -async def test_text_to_pql_resolves_token_per_invocation(text_to_pql_response: dict) -> None: - client = MagicMock() - client.text_to_pql = AsyncMock(return_value=TextToPQLResponse.model_validate(text_to_pql_response)) - config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_pql"]) - - with ( - patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), - patch("gsf.register.get_auth_token", return_value="token-one"), - patch("gsf.register._request_trace_headers", return_value={}), - ): - async with gsf_function_group(config, MagicMock()) as group: - tool = (await group.get_accessible_functions())["gsf__text_to_pql"] - result = json.loads(await tool.ainvoke({"question": "Predict churn"})) - - assert result["request_id"] == "gsf-request-2" - assert client.text_to_pql.await_args.kwargs["token"] == "token-one" - - @pytest.mark.asyncio async def test_explicit_password_auth_does_not_resolve_user_token(text_to_sql_response: dict) -> None: client = MagicMock() From f0a1a6b5f341f14d90d7f361260a89c39d0527ac Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 20:04:46 -0700 Subject: [PATCH 11/12] Harden GSF integration behavior Signed-off-by: Soumili Nandi --- .../scripts/validate_config.py | 3 +- sources/gsf/README.md | 12 +- sources/gsf/pyproject.toml | 2 +- sources/gsf/src/gsf/client.py | 97 ++++++++-- sources/gsf/src/gsf/errors.py | 4 + sources/gsf/src/gsf/models.py | 2 +- sources/gsf/src/gsf/register.py | 48 +++-- sources/gsf/tests/__init__.py | 2 - sources/gsf/tests/conftest.py | 2 + sources/gsf/tests/test_client.py | 178 +++++++++++++++++- sources/gsf/tests/test_models.py | 33 ++++ sources/gsf/tests/test_register.py | 76 +++++--- 12 files changed, 378 insertions(+), 81 deletions(-) delete mode 100644 sources/gsf/tests/__init__.py diff --git a/.agents/skills/aiq-configure-workflow/scripts/validate_config.py b/.agents/skills/aiq-configure-workflow/scripts/validate_config.py index bb7651916..cf6c66692 100644 --- a/.agents/skills/aiq-configure-workflow/scripts/validate_config.py +++ b/.agents/skills/aiq-configure-workflow/scripts/validate_config.py @@ -229,8 +229,9 @@ def validate(path: str) -> int: functions = {} declared_functions = set(functions.keys()) - function_groups = data.get("function_groups") or {} + function_groups = data.get("function_groups", {}) if not isinstance(function_groups, dict): + errors.append("`function_groups:` must be a mapping.") function_groups = {} for field, alias in _iter_refs(functions): diff --git a/sources/gsf/README.md b/sources/gsf/README.md index 272bf0cb1..0fc906c3c 100644 --- a/sources/gsf/README.md +++ b/sources/gsf/README.md @@ -11,8 +11,6 @@ NeMo Agent Toolkit function group. The current implementation provides: - `gsf__text_to_sql` - `gsf__catalog_search` -`gsf__query_context` is registered as an explicit `capability_unavailable` -placeholder until its GSF API contract is ready. PQL client and model groundwork remains internal, but no PQL tool is registered until its GSF contract and integration behavior are validated. @@ -25,7 +23,7 @@ user token and passes it to GSF without storing it on the client. function_groups: gsf: _type: gsf - base_url: ${GSF_BASE_URL:-http://gsf:3000} + base_url: ${GSF_BASE_URL} include: - catalog_search - text_to_sql @@ -47,13 +45,13 @@ functions: For local development and automated evaluation without an incoming AI-Q user token, explicitly configure a GSF password session. The credentials must come -from environment variables and are used only when a tool is invoked: +from environment variables: ```yaml function_groups: gsf: _type: gsf - base_url: ${GSF_BASE_URL:-http://gsf:3000} + base_url: ${GSF_BASE_URL} auth: mode: password email: ${GSF_EMAIL} @@ -80,6 +78,4 @@ authoritative evidence. Catalog search uses `POST /api/question-entity-coverage` and returns entity coverage plus ranked semantic candidates for DS-agent grounding and routing. -Its optional `database_name` is sent as `target_db`. The current GSF endpoint -does not enforce that selector yet, so database-scoped search depends on a GSF -contract update. +Its optional `database_name` is sent as `target_db`. diff --git a/sources/gsf/pyproject.toml b/sources/gsf/pyproject.toml index ca2298bd3..459109fab 100644 --- a/sources/gsf/pyproject.toml +++ b/sources/gsf/pyproject.toml @@ -3,7 +3,7 @@ [build-system] build-backend = "setuptools.build_meta" -requires = ["setuptools >= 83", "setuptools-scm>=8"] +requires = ["setuptools >= 83"] [tool.setuptools.packages.find] where = ["src"] diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index e4492ca10..269205c40 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -6,6 +6,8 @@ import asyncio import json import logging +import random +import re from collections.abc import Mapping from typing import Any @@ -23,10 +25,14 @@ from .models import TextToSQLRequest from .models import TextToSQLResponse -_FORWARDED_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) +FORWARDED_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) _RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) _PASSWORD_SIGN_IN_PATH = "api/auth/sign-in/email" # pragma: allowlist secret _PASSWORD_SIGN_OUT_PATH = "api/auth/sign-out" # pragma: allowlist secret +_HTTP_MAX_CONNECTIONS = 100 +_HTTP_MAX_KEEPALIVE_CONNECTIONS = 20 +_MAX_RETRY_DELAY_SECONDS = 30.0 +_SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n") logger = logging.getLogger(__name__) @@ -47,6 +53,8 @@ def __init__( password_auth_password: SecretStr | None = None, transport: httpx.AsyncBaseTransport | None = None, ) -> None: + """Initialize bounded HTTP behavior and optional password authentication.""" + if (password_auth_email is None) != (password_auth_password is None): raise ValueError("GSF password authentication requires both email and password") self._base_url = base_url.rstrip("/") @@ -82,7 +90,13 @@ def from_config(cls, config: Any) -> "GSFClient": ) async def __aenter__(self) -> "GSFClient": - self._client = httpx.AsyncClient(timeout=self._timeout, transport=self._transport) + """Open the HTTP client and establish a configured password session.""" + + limits = httpx.Limits( + max_connections=_HTTP_MAX_CONNECTIONS, + max_keepalive_connections=_HTTP_MAX_KEEPALIVE_CONNECTIONS, + ) + self._client = httpx.AsyncClient(timeout=self._timeout, limits=limits, transport=self._transport) if self._password_auth_email is not None: try: await self._sign_in_with_password(self._client) @@ -93,6 +107,8 @@ async def __aenter__(self) -> "GSFClient": return self async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: + """Close the password session and shared HTTP client.""" + client = self._client if client is not None: try: @@ -182,6 +198,8 @@ async def _chat_completions( token: str | None, trace_headers: Mapping[str, str] | None = None, ) -> tuple[dict[str, Any], str | None]: + """Call GSF chat completions and return its normalized final event.""" + body, request_id, content_type = await self._post( "chat/completions", payload, @@ -202,6 +220,8 @@ async def _post( capability: str, accept: str = "application/json", ) -> tuple[bytes, str | None, str]: + """Send one authenticated bounded POST with safe retry behavior.""" + client = self._require_client() if self._password_auth_email is None and not token: raise GSFError( @@ -225,9 +245,10 @@ async def _post( if response.status_code >= 400: error = self._http_error(response.status_code, capability, request_id) if error.retryable and attempt + 1 < attempts: + delay = self._retry_delay(attempt, response.headers.get("retry-after")) await response.aclose() response = None - await asyncio.sleep(2**attempt) + await asyncio.sleep(delay) continue raise error @@ -235,19 +256,22 @@ async def _post( return body, request_id, response.headers.get("content-type", "") except GSFError: raise - except httpx.TimeoutException as exc: + except httpx.ConnectTimeout as exc: if attempt + 1 < attempts: - await asyncio.sleep(2**attempt) + await asyncio.sleep(self._retry_delay(attempt, None)) continue raise GSFError( GSFErrorCode.TIMEOUT, f"{capability} timed out.", retryable=True, ) from exc + except httpx.TimeoutException as exc: + raise GSFError( + GSFErrorCode.TIMEOUT, + f"{capability} timed out.", + retryable=True, + ) from exc except httpx.TransportError as exc: - if attempt + 1 < attempts: - await asyncio.sleep(2**attempt) - continue raise GSFError( GSFErrorCode.UPSTREAM_ERROR, f"{capability} could not reach GSF.", @@ -260,6 +284,8 @@ async def _post( raise AssertionError("unreachable") async def _sign_in_with_password(self, client: httpx.AsyncClient) -> None: + """Establish a Better Auth password session on the provided client.""" + try: response = await client.post( f"{self._base_url}/{_PASSWORD_SIGN_IN_PATH}", @@ -267,6 +293,7 @@ async def _sign_in_with_password(self, client: httpx.AsyncClient) -> None: "email": self._password_auth_email, "password": self._password_auth_password.get_secret_value(), }, + headers=self._auth_origin_headers(), ) except httpx.TimeoutException as exc: raise GSFError( @@ -288,11 +315,13 @@ async def _sign_in_with_password(self, client: httpx.AsyncClient) -> None: ) async def _sign_out_password_session(self, client: httpx.AsyncClient) -> None: + """Best-effort sign out of the active Better Auth password session.""" + try: response = await client.post( f"{self._base_url}/{_PASSWORD_SIGN_OUT_PATH}", json={}, - headers={"Origin": self._base_url, "Referer": f"{self._base_url}/"}, + headers=self._auth_origin_headers(), ) if response.status_code >= 400: logger.warning("GSF password session cleanup returned HTTP %s", response.status_code) @@ -300,10 +329,17 @@ async def _sign_out_password_session(self, client: httpx.AsyncClient) -> None: logger.warning("GSF password session cleanup did not complete") def _require_client(self) -> httpx.AsyncClient: + """Return the active client or reject use outside its context.""" + if self._client is None: raise RuntimeError("GSFClient must be used as an async context manager") return self._client + def _auth_origin_headers(self) -> dict[str, str]: + """Build consistent origin headers for Better Auth mutations.""" + + return {"Origin": self._base_url, "Referer": f"{self._base_url}/"} + @staticmethod def _build_headers( token: str | None, @@ -311,6 +347,8 @@ def _build_headers( *, accept: str, ) -> dict[str, str]: + """Build headers while forwarding only approved tracing metadata.""" + headers = { "Accept": accept, "Content-Type": "application/json", @@ -318,11 +356,24 @@ def _build_headers( if token: headers["Authorization"] = f"Bearer {token}" for name, value in (trace_headers or {}).items(): - if name.lower() in _FORWARDED_HEADER_NAMES and value: + if name.lower() in FORWARDED_HEADER_NAMES and value: headers[name] = value return headers + @staticmethod + def _retry_delay(attempt: int, retry_after: str | None) -> float: + """Return a capped server-directed or jittered retry delay.""" + + if retry_after: + try: + return max(0.0, min(float(retry_after), _MAX_RETRY_DELAY_SECONDS)) + except ValueError: + pass + return min(2**attempt, _MAX_RETRY_DELAY_SECONDS) * (0.5 + random.random() / 2) + async def _read_bounded(self, response: httpx.Response, *, request_id: str | None) -> bytes: + """Read a response without exceeding the configured byte ceiling.""" + content_length = response.headers.get("content-length") if content_length is not None: try: @@ -340,6 +391,8 @@ async def _read_bounded(self, response: httpx.Response, *, request_id: str | Non @classmethod def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | None) -> dict[str, Any]: + """Extract the final answer object from JSON or an SSE response.""" + try: text = body.decode("utf-8") if "text/event-stream" not in content_type and not text.lstrip().startswith(("data:", ":")): @@ -347,7 +400,7 @@ def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | return cls._answer_from_event(payload, request_id=request_id) data_lines: list[str] = [] - for line in [*text.splitlines(), ""]: + for line in [*_SSE_LINE_SPLIT.split(text), ""]: if line.startswith(":"): continue if line.startswith("data:"): @@ -387,6 +440,8 @@ def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | @staticmethod def _parse_json_data(body: bytes, *, request_id: str | None) -> dict[str, Any]: + """Parse a JSON response and unwrap its optional data envelope.""" + try: payload = json.loads(body) except (json.JSONDecodeError, UnicodeDecodeError) as exc: @@ -408,6 +463,8 @@ def _parse_json_data(body: bytes, *, request_id: str | None) -> dict[str, Any]: @staticmethod def _answer_from_event(payload: Any, *, request_id: str | None) -> dict[str, Any]: + """Validate and extract a final chat answer from an event payload.""" + if isinstance(payload, dict) and isinstance(payload.get("answer"), dict): return payload["answer"] if isinstance(payload, dict) and payload.get("type") is None: @@ -426,6 +483,8 @@ def _normalize_text_to_sql( request_id: str | None, max_rows: int, ) -> TextToSQLResponse: + """Normalize current and legacy GSF SQL fields into AI-Q output.""" + sql = answer.get("sql") or answer.get("sql_code") if not isinstance(sql, str) or not sql.strip(): raise GSFError( @@ -434,11 +493,11 @@ def _normalize_text_to_sql( request_id=request_id, ) - rows = cls._normalize_rows(answer.get("rows", answer.get("sql_response_from_db")), request_id) + rows = cls._normalize_rows(answer.get("rows") or answer.get("sql_response_from_db"), request_id) upstream_truncated = bool(answer.get("truncated", False)) truncated = upstream_truncated or len(rows) > max_rows rows = rows[:max_rows] - columns = cls._normalize_columns(answer.get("columns", answer.get("sql_columns"))) + columns = cls._normalize_columns(answer.get("columns") or answer.get("sql_columns")) if not columns and rows: columns = [ResultColumn(name=str(name)) for name in rows[0]] @@ -473,6 +532,8 @@ def _normalize_catalog_search( request_id: str | None, max_results: int, ) -> CatalogSearchResponse: + """Validate catalog candidates and enforce the result ceiling.""" + candidates = data.get("candidates") if not isinstance(candidates, list): raise GSFError( @@ -499,6 +560,8 @@ def _normalize_catalog_search( @classmethod def _normalize_text_to_pql(cls, answer: Mapping[str, Any], *, request_id: str | None) -> TextToPQLResponse: + """Normalize current and legacy GSF prediction fields into PQL output.""" + # GSF's prediction branch currently returns the PQL in ``sql_code`` so # its frontend can render SQL and prediction results with one shape. pql = answer.get("pql") or answer.get("pql_code") or answer.get("sql_code") @@ -529,6 +592,8 @@ def _normalize_text_to_pql(cls, answer: Mapping[str, Any], *, request_id: str | @staticmethod def _normalize_columns(value: Any) -> list[ResultColumn]: + """Normalize structured or name-only column metadata.""" + if not isinstance(value, list): return [] columns: list[ResultColumn] = [] @@ -548,6 +613,8 @@ def _normalize_columns(value: Any) -> list[ResultColumn]: @staticmethod def _normalize_rows(value: Any, request_id: str | None) -> list[dict[str, Any]]: + """Normalize GSF row payload variants into dictionaries.""" + if value is None: return [] if isinstance(value, str): @@ -601,6 +668,8 @@ def _normalize_rows(value: Any, request_id: str | None) -> list[dict[str, Any]]: return rows def _response_too_large(self, request_id: str | None) -> GSFError: + """Build an error for an oversized GSF response.""" + return GSFError( GSFErrorCode.RESPONSE_TOO_LARGE, "GSF response exceeded the configured size limit.", @@ -609,6 +678,8 @@ def _response_too_large(self, request_id: str | None) -> GSFError: @staticmethod def _http_error(status_code: int, capability: str, request_id: str | None) -> GSFError: + """Map an HTTP failure to the stable GSF error contract.""" + if status_code == 401: return GSFError( GSFErrorCode.AUTHENTICATION_REQUIRED, diff --git a/sources/gsf/src/gsf/errors.py b/sources/gsf/src/gsf/errors.py index bf05d5c0c..3ed254b17 100644 --- a/sources/gsf/src/gsf/errors.py +++ b/sources/gsf/src/gsf/errors.py @@ -35,6 +35,8 @@ def __init__( retryable: bool = False, request_id: str | None = None, ) -> None: + """Initialize a normalized failure without retaining response data.""" + super().__init__(message) self.code = code self.message = message @@ -55,6 +57,8 @@ class GSFToolError(BaseModel): @classmethod def from_exception(cls, error: GSFError) -> "GSFToolError": + """Convert an internal exception to the serialized tool error shape.""" + return cls( code=error.code, retryable=error.retryable, diff --git a/sources/gsf/src/gsf/models.py b/sources/gsf/src/gsf/models.py index 58a8ab2bd..e67b85744 100644 --- a/sources/gsf/src/gsf/models.py +++ b/sources/gsf/src/gsf/models.py @@ -44,7 +44,7 @@ class CatalogSearchResponse(GSFResponse): """Coverage and ranked semantic candidates returned by GSF.""" request_id: str | None = None - coverage: float = Field(ge=0, le=1) + coverage: float | None = Field(default=None, ge=0, le=1) candidates: list[CatalogCandidate] uncovered_entities: list[str] | None = None truncated: bool = False diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py index 3144baa11..da933817e 100644 --- a/sources/gsf/src/gsf/register.py +++ b/sources/gsf/src/gsf/register.py @@ -13,25 +13,22 @@ from pydantic import HttpUrl from pydantic import SecretStr -from aiq_agent.auth.utils import get_auth_token from nat.builder.builder import Builder from nat.builder.context import Context from nat.builder.function import FunctionGroup from nat.cli.register_workflow import register_function_group from nat.data_models.function import FunctionGroupBaseConfig +from .client import FORWARDED_HEADER_NAMES from .client import GSFClient from .errors import GSFError from .errors import GSFErrorCode from .errors import GSFToolError from .models import CatalogSearchRequest -from .models import QueryContextRequest from .models import TextToSQLRequest logger = logging.getLogger(__name__) -_TRACE_HEADER_NAMES = frozenset({"baggage", "traceparent", "tracestate", "x-correlation-id", "x-request-id"}) - class GSFPasswordAuthConfig(BaseModel): """Explicit GSF password-session configuration for development and evaluation.""" @@ -56,15 +53,30 @@ class GSFFunctionGroupConfig(FunctionGroupBaseConfig, name="gsf"): def _tool_error(error: GSFError) -> str: + """Serialize an internal GSF exception for safe tool output.""" + return GSFToolError.from_exception(error).model_dump_json(exclude_none=True) +def _get_auth_token() -> str | None: + """Lazily resolve the current AI-Q token for bearer authentication.""" + + try: + from aiq_agent.auth.utils import get_auth_token + except ImportError as exc: + raise GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "AI-Q authentication support is unavailable.", + ) from exc + return get_auth_token() + + def _resolve_request_token(config: GSFFunctionGroupConfig) -> str | None: """Resolve a bearer token unless an explicit password session is configured.""" if config.auth is not None: return None - token = get_auth_token() + token = _get_auth_token() if not token: raise GSFError( GSFErrorCode.AUTHENTICATION_REQUIRED, @@ -74,14 +86,17 @@ def _resolve_request_token(config: GSFFunctionGroupConfig) -> str | None: def _request_trace_headers() -> Mapping[str, str]: + """Read and filter tracing headers from the current NAT context.""" + try: metadata = Context.get().metadata incoming = metadata.headers if metadata else None except Exception: + logger.debug("Unable to read request trace headers from NAT context", exc_info=True) return {} if not incoming: return {} - return {name: value for name, value in incoming.items() if name.lower() in _TRACE_HEADER_NAMES and value} + return {name: value for name, value in incoming.items() if name.lower() in FORWARDED_HEADER_NAMES and value} @register_function_group(config_type=GSFFunctionGroupConfig) @@ -107,7 +122,7 @@ async def catalog_search(request: CatalogSearchRequest) -> str: except GSFError as error: return _tool_error(error) except Exception: - logger.error("Unexpected GSF catalog-search failure") + logger.exception("Unexpected GSF catalog-search failure") return _tool_error( GSFError( GSFErrorCode.UPSTREAM_ERROR, @@ -133,7 +148,7 @@ async def text_to_sql(request: TextToSQLRequest) -> str: except GSFError as error: return _tool_error(error) except Exception: - logger.error("Unexpected GSF text-to-SQL failure") + logger.exception("Unexpected GSF text-to-SQL failure") return _tool_error( GSFError( GSFErrorCode.UPSTREAM_ERROR, @@ -141,17 +156,6 @@ async def text_to_sql(request: TextToSQLRequest) -> str: ) ) - async def query_context(request: QueryContextRequest) -> str: - """Build GSF query context (not available in this integration yet).""" - - del request - return _tool_error( - GSFError( - GSFErrorCode.CAPABILITY_UNAVAILABLE, - "GSF query context is unavailable.", - ) - ) - group = FunctionGroup(config=config) group.add_function( "catalog_search", @@ -165,10 +169,4 @@ async def query_context(request: QueryContextRequest) -> str: input_schema=TextToSQLRequest, description=text_to_sql.__doc__, ) - group.add_function( - "query_context", - query_context, - input_schema=QueryContextRequest, - description=query_context.__doc__, - ) yield group diff --git a/sources/gsf/tests/__init__.py b/sources/gsf/tests/__init__.py deleted file mode 100644 index d51c4fe1e..000000000 --- a/sources/gsf/tests/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py index 15e0743cd..a78f0554a 100644 --- a/sources/gsf/tests/conftest.py +++ b/sources/gsf/tests/conftest.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Shared GSF response fixtures.""" + import pytest diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index 87056146c..8212581fe 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Tests for the typed GSF HTTP client.""" + import json from unittest.mock import AsyncMock from unittest.mock import patch @@ -19,6 +21,8 @@ def _sse_response(answer: dict) -> httpx.Response: + """Build a representative GSF SSE result response.""" + events = [ 'data: {"type":"step","node":"construct_sql_from_candidates"}', "", @@ -34,10 +38,30 @@ def _sse_response(answer: dict) -> httpx.Response: ) +@pytest.mark.parametrize( + ("email", "password"), + [ + ("developer@example.com", None), + (None, SecretStr(_TEST_PASSWORD)), + ], +) +def test_password_auth_requires_both_email_and_password(email: str | None, password: SecretStr | None) -> None: + """Reject incomplete password-session credentials.""" + + with pytest.raises(ValueError): + GSFClient( + base_url="https://gsf.example", + password_auth_email=email, + password_auth_password=password, + ) + + @pytest.mark.asyncio async def test_catalog_search_uses_entity_coverage_path_maps_scope_and_bounds_candidates( catalog_search_api_response: dict, ) -> None: + """Map catalog scope and enforce the candidate limit.""" + seen_request: httpx.Request | None = None async def handler(request: httpx.Request) -> httpx.Response: @@ -81,6 +105,8 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_catalog_search_maps_missing_endpoint_to_capability_unavailable() -> None: + """Map a missing catalog endpoint to an unavailable capability.""" + async def handler(_request: httpx.Request) -> httpx.Response: return httpx.Response(404) @@ -94,7 +120,11 @@ async def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_text_to_sql_maps_database_to_target_db_and_bounds_rows(chat_sql_answer: dict) -> None: + """Map SQL scope and enforce the configured row limit.""" + seen_request: httpx.Request | None = None + chat_sql_answer["rows"] = None + chat_sql_answer["columns"] = None async def handler(request: httpx.Request) -> httpx.Response: nonlocal seen_request @@ -132,6 +162,8 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_text_to_sql_omits_optional_target_db(chat_sql_answer: dict) -> None: + """Omit target_db when no database scope is requested.""" + seen_payload: dict | None = None async def handler(request: httpx.Request) -> httpx.Response: @@ -150,8 +182,32 @@ async def handler(request: httpx.Request) -> httpx.Response: assert result.sql == "SELECT revenue FROM quarterly_results" +@pytest.mark.asyncio +async def test_text_to_sql_preserves_unicode_line_separator_in_sse(chat_sql_answer: dict) -> None: + """Preserve Unicode separators inside SSE JSON string values.""" + + thoughts = "First decision.\u2028Second decision." + answer = {**chat_sql_answer, "thoughts": thoughts} + + async def handler(_request: httpx.Request) -> httpx.Response: + result_event = json.dumps({"type": "result", "answer": answer}, ensure_ascii=False) + return httpx.Response( + 200, + content=f"data: {result_event}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + ) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + result = await client.text_to_sql(TextToSQLRequest(question="Show revenue"), token="user-token") + + assert result.thoughts == thoughts + + @pytest.mark.asyncio async def test_text_to_pql_maps_database_to_target_db(chat_pql_answer: dict) -> None: + """Map prediction database scope to the GSF target_db field.""" + seen_payload: dict | None = None async def handler(request: httpx.Request) -> httpx.Response: @@ -177,6 +233,8 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_password_auth_logs_in_uses_cookie_and_signs_out(chat_sql_answer: dict) -> None: + """Use one Better Auth cookie from sign-in through sign-out.""" + seen_paths: list[str] = [] async def handler(request: httpx.Request) -> httpx.Response: @@ -186,6 +244,8 @@ async def handler(request: httpx.Request) -> httpx.Response: "email": "developer@example.com", "password": _TEST_PASSWORD, } + assert request.headers["origin"] == "https://gsf.example" + assert request.headers["referer"] == "https://gsf.example/" return httpx.Response( 200, json={"user": {"email": "developer@example.com"}}, @@ -221,6 +281,8 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_password_session_is_reused_across_tool_calls(chat_sql_answer: dict) -> None: + """Reuse one password session across multiple GSF calls.""" + seen_paths: list[str] = [] async def handler(request: httpx.Request) -> httpx.Response: @@ -257,6 +319,8 @@ async def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_client_without_password_auth_requires_bearer_before_http() -> None: + """Fail before HTTP when neither password nor bearer auth exists.""" + calls = 0 async def handler(_request: httpx.Request) -> httpx.Response: @@ -275,6 +339,8 @@ async def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_client_normalizes_forbidden_without_leaking_body() -> None: + """Map forbidden responses without exposing upstream response text.""" + async def handler(_request: httpx.Request) -> httpx.Response: return httpx.Response(403, text="secret database details") @@ -289,6 +355,8 @@ async def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_client_retries_rate_limit_then_succeeds(chat_sql_answer: dict) -> None: + """Retry a rate-limited request with deterministic jitter.""" + attempts = 0 async def handler(_request: httpx.Request) -> httpx.Response: @@ -299,16 +367,118 @@ async def handler(_request: httpx.Request) -> httpx.Response: return _sse_response(chat_sql_answer) client = GSFClient(base_url="https://gsf.example", max_retries=1, transport=httpx.MockTransport(handler)) - with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: + with ( + patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep, + patch("gsf.client.random.random", return_value=0.5), + ): async with client: await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") assert attempts == 2 - sleep.assert_awaited_once_with(1) + sleep.assert_awaited_once_with(0.75) + + +@pytest.mark.asyncio +async def test_client_raises_rate_limited_after_retries_are_exhausted() -> None: + """Raise RATE_LIMITED after exhausting all configured attempts.""" + + attempts = 0 + + async def handler(_request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + return httpx.Response(429) + + client = GSFClient(base_url="https://gsf.example", max_retries=2, transport=httpx.MockTransport(handler)) + with ( + patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep, + patch("gsf.client.random.random", return_value=0.5), + ): + async with client: + with pytest.raises(GSFError) as raised: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") + + assert attempts == 3 + assert sleep.await_count == 2 + assert raised.value.code is GSFErrorCode.RATE_LIMITED + assert raised.value.retryable is True + + +@pytest.mark.asyncio +async def test_client_retries_connect_timeout_then_raises_timeout() -> None: + """Retry connect timeouts before returning the normalized timeout.""" + + attempts = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + raise httpx.ConnectTimeout("GSF connection timed out", request=request) + + client = GSFClient(base_url="https://gsf.example", max_retries=2, transport=httpx.MockTransport(handler)) + with ( + patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep, + patch("gsf.client.random.random", return_value=0.5), + ): + async with client: + with pytest.raises(GSFError) as raised: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") + + assert attempts == 3 + assert sleep.await_count == 2 + assert raised.value.code is GSFErrorCode.TIMEOUT + assert raised.value.retryable is True + + +@pytest.mark.asyncio +async def test_client_does_not_retry_read_timeout() -> None: + """Avoid replaying a request after a read timeout.""" + + attempts = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + raise httpx.ReadTimeout("GSF response timed out", request=request) + + client = GSFClient(base_url="https://gsf.example", max_retries=2, transport=httpx.MockTransport(handler)) + with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: + async with client: + with pytest.raises(GSFError) as raised: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") + + assert raised.value.code is GSFErrorCode.TIMEOUT + assert attempts == 1 + sleep.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_client_does_not_retry_transport_error() -> None: + """Avoid replaying a request after a generic transport failure.""" + + attempts = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + raise httpx.ConnectError("GSF connection failed", request=request) + + client = GSFClient(base_url="https://gsf.example", max_retries=2, transport=httpx.MockTransport(handler)) + with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: + async with client: + with pytest.raises(GSFError) as raised: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") + + assert attempts == 1 + assert raised.value.code is GSFErrorCode.UPSTREAM_ERROR + assert raised.value.retryable is True + sleep.assert_not_awaited() @pytest.mark.asyncio async def test_client_rejects_oversized_response(chat_sql_answer: dict) -> None: + """Reject a response that exceeds the configured byte ceiling.""" + async def handler(_request: httpx.Request) -> httpx.Response: return _sse_response(chat_sql_answer) @@ -326,6 +496,8 @@ async def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_client_rejects_malformed_response() -> None: + """Reject a malformed non-SSE response.""" + async def handler(_request: httpx.Request) -> httpx.Response: return httpx.Response(200, content=b"not-json") @@ -339,6 +511,8 @@ async def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_client_maps_sse_error_without_leaking_message() -> None: + """Map SSE errors without exposing upstream error text.""" + async def handler(_request: httpx.Request) -> httpx.Response: return httpx.Response( 200, diff --git a/sources/gsf/tests/test_models.py b/sources/gsf/tests/test_models.py index d3cc84152..417bef732 100644 --- a/sources/gsf/tests/test_models.py +++ b/sources/gsf/tests/test_models.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Tests for GSF request and response models.""" + import pytest from gsf.models import CatalogCandidate from gsf.models import CatalogSearchRequest @@ -14,6 +16,8 @@ def test_catalog_search_request_supports_optional_scope_and_search_controls() -> None: + """Accept optional catalog scope and bounded search controls.""" + request = CatalogSearchRequest( question="Find revenue metrics", database_name="benchmark_db", @@ -27,6 +31,8 @@ def test_catalog_search_request_supports_optional_scope_and_search_controls() -> def test_catalog_search_response_validates_coverage() -> None: + """Reject catalog coverage outside the normalized range.""" + with pytest.raises(ValidationError): CatalogSearchResponse( coverage=1.5, @@ -41,7 +47,20 @@ def test_catalog_search_response_validates_coverage() -> None: ) +def test_catalog_search_response_accepts_missing_coverage(catalog_search_response: dict) -> None: + """Preserve candidates when catalog coverage is absent.""" + + catalog_search_response.pop("coverage") + + result = CatalogSearchResponse.model_validate(catalog_search_response) + + assert result.coverage is None + assert result.candidates + + def test_catalog_search_response_ignores_future_fields(catalog_search_response: dict) -> None: + """Ignore unknown catalog enrichments from newer GSF versions.""" + catalog_search_response["future_gsf_metadata"] = {"enabled": True} result = CatalogSearchResponse.model_validate(catalog_search_response) @@ -51,6 +70,8 @@ def test_catalog_search_response_ignores_future_fields(catalog_search_response: def test_text_to_sql_request_supports_optional_database_name() -> None: + """Accept an optional database scope for SQL requests.""" + request = TextToSQLRequest(question="Show quarterly revenue", database_name="benchmark_db") assert request.database_name == "benchmark_db" @@ -58,23 +79,31 @@ def test_text_to_sql_request_supports_optional_database_name() -> None: def test_text_to_pql_request_supports_optional_database_name() -> None: + """Accept an optional database scope for prediction requests.""" + request = TextToPQLRequest(question="Predict churn", database_name="benchmark_db") assert request.database_name == "benchmark_db" def test_query_context_request_omits_optional_database_name() -> None: + """Omit unset query-context database scope during serialization.""" + payload = QueryContextRequest(question="What revenue data is available?").model_dump(exclude_none=True) assert "database_name" not in payload def test_requests_reject_unknown_fields() -> None: + """Reject unknown fields in outbound GSF request models.""" + with pytest.raises(ValidationError): TextToSQLRequest.model_validate({"question": "Show revenue", "unknown": True}) def test_text_to_sql_response_accepts_missing_future_enrichments() -> None: + """Accept SQL responses that omit optional enrichment fields.""" + result = TextToSQLResponse.model_validate( { "sql": "SELECT revenue FROM quarterly_results", @@ -89,6 +118,8 @@ def test_text_to_sql_response_accepts_missing_future_enrichments() -> None: def test_text_to_pql_response_accepts_missing_future_enrichments() -> None: + """Accept PQL responses that omit optional enrichment fields.""" + result = TextToPQLResponse.model_validate({"pql": "PREDICT churn FOR customers NEXT 30 DAYS"}) assert result.request_id is None @@ -97,6 +128,8 @@ def test_text_to_pql_response_accepts_missing_future_enrichments() -> None: def test_text_to_sql_response_ignores_future_fields(text_to_sql_response: dict) -> None: + """Ignore unknown SQL enrichments from newer GSF versions.""" + text_to_sql_response["future_gsf_metadata"] = {"enabled": True} result = TextToSQLResponse.model_validate(text_to_sql_response) diff --git a/sources/gsf/tests/test_register.py b/sources/gsf/tests/test_register.py index d68961e0f..6c08370b2 100644 --- a/sources/gsf/tests/test_register.py +++ b/sources/gsf/tests/test_register.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Tests for GSF NAT function-group registration.""" + import json from unittest.mock import AsyncMock from unittest.mock import MagicMock @@ -11,23 +13,34 @@ from gsf.models import TextToSQLResponse from gsf.register import GSFFunctionGroupConfig from gsf.register import GSFPasswordAuthConfig +from gsf.register import _request_trace_headers from gsf.register import gsf_function_group _TEST_PASSWORD = "${TEST_GSF_PASSWORD}" class FakeClientContext: + """Expose a mocked client through an asynchronous context manager.""" + def __init__(self, client: MagicMock) -> None: + """Store the mocked GSF client.""" + self.client = client async def __aenter__(self) -> MagicMock: + """Return the mocked GSF client.""" + return self.client async def __aexit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: + """Complete context cleanup without suppressing failures.""" + pass def test_password_auth_is_optional_and_keeps_secret_wrapped() -> None: + """Keep optional password configuration secret-wrapped.""" + default_config = GSFFunctionGroupConfig(base_url="https://gsf.example") password_config = GSFFunctionGroupConfig( base_url="https://gsf.example", @@ -44,8 +57,31 @@ def test_password_auth_is_optional_and_keeps_secret_wrapped() -> None: assert _TEST_PASSWORD not in repr(password_config) +def test_request_trace_headers_forwards_only_allowlisted_nonempty_values() -> None: + """Forward only nonempty request headers on the tracing allowlist.""" + + context = MagicMock() + context.metadata.headers = { + "Traceparent": "00-trace", + "x-request-id": "request-1", + "baggage": "", + "Authorization": "Bearer secret", + "Cookie": "session=secret", + } + + with patch("gsf.register.Context.get", return_value=context): + headers = _request_trace_headers() + + assert headers == { + "Traceparent": "00-trace", + "x-request-id": "request-1", + } + + @pytest.mark.asyncio async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> None: + """Expose only tools selected by the function-group include list.""" + client = MagicMock() client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) @@ -59,13 +95,15 @@ async def test_group_exposes_only_requested_tools(text_to_sql_response: dict) -> @pytest.mark.asyncio async def test_catalog_search_resolves_token_per_invocation(catalog_search_response: dict) -> None: + """Resolve a fresh bearer token for a catalog invocation.""" + client = MagicMock() client.catalog_search = AsyncMock(return_value=CatalogSearchResponse.model_validate(catalog_search_response)) config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["catalog_search"]) with ( patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), - patch("gsf.register.get_auth_token", return_value="token-one"), + patch("gsf.register._get_auth_token", return_value="token-one"), patch("gsf.register._request_trace_headers", return_value={}), ): async with gsf_function_group(config, MagicMock()) as group: @@ -78,13 +116,15 @@ async def test_catalog_search_resolves_token_per_invocation(catalog_search_respo @pytest.mark.asyncio async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: dict) -> None: + """Resolve bearer tokens independently for consecutive SQL calls.""" + client = MagicMock() client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) with ( patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), - patch("gsf.register.get_auth_token", side_effect=["token-one", "token-two"]), + patch("gsf.register._get_auth_token", side_effect=["token-one", "token-two"]), patch("gsf.register._request_trace_headers", return_value={}), ): async with gsf_function_group(config, MagicMock()) as group: @@ -105,6 +145,8 @@ async def test_text_to_sql_resolves_token_per_invocation(text_to_sql_response: d @pytest.mark.asyncio async def test_explicit_password_auth_does_not_resolve_user_token(text_to_sql_response: dict) -> None: + """Skip user-token resolution when password mode is explicit.""" + client = MagicMock() client.text_to_sql = AsyncMock(return_value=TextToSQLResponse.model_validate(text_to_sql_response)) config = GSFFunctionGroupConfig( @@ -119,7 +161,7 @@ async def test_explicit_password_auth_does_not_resolve_user_token(text_to_sql_re with ( patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), - patch("gsf.register.get_auth_token") as get_auth_token, + patch("gsf.register._get_auth_token") as get_auth_token, patch("gsf.register._request_trace_headers", return_value={}), ): async with gsf_function_group(config, MagicMock()) as group: @@ -133,13 +175,15 @@ async def test_explicit_password_auth_does_not_resolve_user_token(text_to_sql_re @pytest.mark.asyncio async def test_missing_authentication_fails_closed() -> None: + """Return an authentication error without invoking the GSF client.""" + client = MagicMock() client.text_to_sql = AsyncMock() config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=["text_to_sql"]) with ( patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)), - patch("gsf.register.get_auth_token", return_value=None), + patch("gsf.register._get_auth_token", return_value=None), ): async with gsf_function_group(config, MagicMock()) as group: tool = (await group.get_accessible_functions())["gsf__text_to_sql"] @@ -147,27 +191,3 @@ async def test_missing_authentication_fails_closed() -> None: assert result["code"] == "authentication_required" client.text_to_sql.assert_not_awaited() - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("tool_name", "tool_input", "message"), - [ - ("query_context", {"question": "Show data"}, "GSF query context is unavailable."), - ], -) -async def test_placeholder_tools_are_explicitly_unavailable(tool_name: str, tool_input: dict, message: str) -> None: - client = MagicMock() - config = GSFFunctionGroupConfig(base_url="https://gsf.example", include=[tool_name]) - - with patch("gsf.register.GSFClient.from_config", return_value=FakeClientContext(client)): - async with gsf_function_group(config, MagicMock()) as group: - tool = (await group.get_accessible_functions())[f"gsf__{tool_name}"] - result = json.loads(await tool.ainvoke(tool_input)) - - assert result == { - "status": "error", - "code": "capability_unavailable", - "retryable": False, - "message": message, - } From 9872f90a0359e131c874fb479db82f49d4b0eaf5 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Thu, 6 Aug 2026 20:16:51 -0700 Subject: [PATCH 12/12] Harden GSF SSE response handling Signed-off-by: Soumili Nandi --- sources/gsf/src/gsf/client.py | 7 +++++- sources/gsf/tests/test_client.py | 41 ++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/sources/gsf/src/gsf/client.py b/sources/gsf/src/gsf/client.py index 269205c40..22ac713dc 100644 --- a/sources/gsf/src/gsf/client.py +++ b/sources/gsf/src/gsf/client.py @@ -412,7 +412,12 @@ def _parse_chat_answer(cls, body: bytes, *, content_type: str, request_id: str | data_lines.clear() if event_data == "[DONE]": continue - event = json.loads(event_data) + try: + event = json.loads(event_data) + except json.JSONDecodeError: + # Step/progress events are observational. A malformed one + # must not hide a later, valid terminal result event. + continue if not isinstance(event, dict): continue if event.get("type") == "error": diff --git a/sources/gsf/tests/test_client.py b/sources/gsf/tests/test_client.py index 8212581fe..bb9e3d7e0 100644 --- a/sources/gsf/tests/test_client.py +++ b/sources/gsf/tests/test_client.py @@ -204,6 +204,25 @@ async def handler(_request: httpx.Request) -> httpx.Response: assert result.thoughts == thoughts +@pytest.mark.asyncio +async def test_text_to_sql_skips_malformed_intermediate_sse_event(chat_sql_answer: dict) -> None: + """Ignore a malformed progress event before a valid terminal result.""" + + async def handler(_request: httpx.Request) -> httpx.Response: + result_event = json.dumps({"type": "result", "answer": chat_sql_answer}) + return httpx.Response( + 200, + content=f"data: {{malformed-step\n\ndata: {result_event}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + ) + + client = GSFClient(base_url="https://gsf.example", transport=httpx.MockTransport(handler)) + async with client: + result = await client.text_to_sql(TextToSQLRequest(question="Show revenue"), token="user-token") + + assert result.sql == "SELECT revenue FROM quarterly_results" + + @pytest.mark.asyncio async def test_text_to_pql_maps_database_to_target_db(chat_pql_answer: dict) -> None: """Map prediction database scope to the GSF target_db field.""" @@ -378,6 +397,28 @@ async def handler(_request: httpx.Request) -> httpx.Response: sleep.assert_awaited_once_with(0.75) +@pytest.mark.asyncio +async def test_client_honors_retry_after_for_rate_limit(chat_sql_answer: dict) -> None: + """Prefer GSF's bounded Retry-After delay over local jitter.""" + + attempts = 0 + + async def handler(_request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(429, headers={"Retry-After": "12"}) + return _sse_response(chat_sql_answer) + + client = GSFClient(base_url="https://gsf.example", max_retries=1, transport=httpx.MockTransport(handler)) + with patch("gsf.client.asyncio.sleep", new_callable=AsyncMock) as sleep: + async with client: + await client.text_to_sql(TextToSQLRequest(question="Show data"), token="user-token") + + assert attempts == 2 + sleep.assert_awaited_once_with(12.0) + + @pytest.mark.asyncio async def test_client_raises_rate_limited_after_retries_are_exhausted() -> None: """Raise RATE_LIMITED after exhausting all configured attempts."""