diff --git a/.agents/skills/aiq-configure-workflow/scripts/validate_config.py b/.agents/skills/aiq-configure-workflow/scripts/validate_config.py index f26013931..cf6c66692 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,11 @@ def validate(path: str) -> int: functions = {} declared_functions = set(functions.keys()) + 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): if alias not in defined_aliases: defined = ", ".join(sorted(defined_aliases)) or "none" @@ -247,7 +253,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: 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..0fc906c3c --- /dev/null +++ b/sources/gsf/README.md @@ -0,0 +1,81 @@ + + +# 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__catalog_search` + +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 +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} + include: + - catalog_search + - 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 +``` + +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: + +```yaml +function_groups: + gsf: + _type: gsf + base_url: ${GSF_BASE_URL} + auth: + mode: password + email: ${GSF_EMAIL} + password: ${GSF_PASSWORD} + include: + - catalog_search + - text_to_sql +``` + +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 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 +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. +Its optional `database_name` is sent as `target_db`. diff --git a/sources/gsf/pyproject.toml b/sources/gsf/pyproject.toml new file mode 100644 index 000000000..459109fab --- /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"] + +[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..22ac713dc --- /dev/null +++ b/sources/gsf/src/gsf/client.py @@ -0,0 +1,726 @@ +# 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 +import logging +import random +import re +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 CatalogSearchRequest +from .models import CatalogSearchResponse +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 +_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__) + + +class GSFClient: + """NAT-independent client sharing one bounded HTTP connection pool.""" + + def __init__( + self, + *, + base_url: str, + 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, + password_auth_email: str | None = None, + 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("/") + 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, + 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.""" + + password_auth = getattr(config, "auth", None) + return cls( + base_url=str(config.base_url), + 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, + 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": + """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) + except BaseException: + await self._client.aclose() + self._client = None + raise + 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: + 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 | 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) + 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(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, + *, + token: str | None, + trace_headers: Mapping[str, str] | None = None, + ) -> 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]: + """Call GSF chat completions and return its normalized final event.""" + + body, request_id, content_type = await self._post( + "chat/completions", + payload, + token=token, + trace_headers=trace_headers, + capability="GSF chat completions", + accept="text/event-stream", + ) + 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 | None, + trace_headers: Mapping[str, str] | None, + 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( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF authentication is required.", + ) + headers = self._build_headers(token, trace_headers, accept=accept) + 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: + delay = self._retry_delay(attempt, response.headers.get("retry-after")) + await response.aclose() + response = None + await asyncio.sleep(delay) + continue + raise error + + body = await self._read_bounded(response, request_id=request_id) + return body, request_id, response.headers.get("content-type", "") + except GSFError: + raise + except httpx.ConnectTimeout as exc: + if attempt + 1 < attempts: + 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: + 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") + + 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}", + json={ + "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( + 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: + """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=self._auth_origin_headers(), + ) + 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: + """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, + trace_headers: Mapping[str, str] | None, + *, + accept: str, + ) -> dict[str, str]: + """Build headers while forwarding only approved tracing metadata.""" + + headers = { + "Accept": accept, + "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 + 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: + 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) + + @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:", ":")): + payload = json.loads(text) + return cls._answer_from_event(payload, request_id=request_id) + + data_lines: list[str] = [] + for line in [*_SSE_LINE_SPLIT.split(text), ""]: + 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 + 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": + 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 _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: + 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]: + """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: + return payload + raise GSFError( + GSFErrorCode.INVALID_RESPONSE, + "GSF returned an invalid final answer.", + request_id=request_id, + ) + + @classmethod + def _normalize_text_to_sql( + cls, + answer: Mapping[str, Any], + *, + 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( + GSFErrorCode.INVALID_RESPONSE, + "GSF did not return validated SQL.", + request_id=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") or 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 request_id, + thoughts=answer.get("thoughts"), + 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=request_id, + ) from exc + + @staticmethod + def _normalize_catalog_search( + data: Mapping[str, Any], + *, + 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( + 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: + """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") + 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 + def _normalize_columns(value: Any) -> list[ResultColumn]: + """Normalize structured or name-only column metadata.""" + + 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]]: + """Normalize GSF row payload variants into dictionaries.""" + + 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: + """Build an error for an oversized GSF response.""" + + 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: + """Map an HTTP failure to the stable GSF error contract.""" + + 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..3ed254b17 --- /dev/null +++ b/sources/gsf/src/gsf/errors.py @@ -0,0 +1,67 @@ +# 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: + """Initialize a normalized failure without retaining response data.""" + + 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": + """Convert an internal exception to the serialized tool error shape.""" + + 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..e67b85744 --- /dev/null +++ b/sources/gsf/src/gsf/models.py @@ -0,0 +1,124 @@ +# 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): + """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) + 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 | None = Field(default=None, ge=0, le=1) + candidates: list[CatalogCandidate] + uncovered_entities: list[str] | None = None + truncated: bool = False + + +class ResultColumn(GSFResponse): + """A column in a bounded SQL result.""" + + name: str + data_type: str | None = None + + +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 and execute validated SQL with bounded results.""" + + question: str = Field(min_length=1, max_length=4_096) + database_name: str | None = None + 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.""" + + 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) + 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 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.""" + + 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) diff --git a/sources/gsf/src/gsf/register.py b/sources/gsf/src/gsf/register.py new file mode 100644 index 000000000..da933817e --- /dev/null +++ b/sources/gsf/src/gsf/register.py @@ -0,0 +1,172 @@ +# 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 typing import Literal + +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field +from pydantic import HttpUrl +from pydantic import SecretStr + +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 TextToSQLRequest + +logger = logging.getLogger(__name__) + + +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) + 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: + """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() + if not token: + raise GSFError( + GSFErrorCode.AUTHENTICATION_REQUIRED, + "GSF authentication is required.", + ) + return token + + +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 FORWARDED_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: + """Find GSF semantic candidates relevant to an enterprise-data question. + + 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.exception("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. + + 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. + """ + + try: + result = await client.text_to_sql( + 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.exception("Unexpected GSF text-to-SQL failure") + return _tool_error( + GSFError( + GSFErrorCode.UPSTREAM_ERROR, + "GSF text-to-SQL 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__, + ) + yield group diff --git a/sources/gsf/tests/conftest.py b/sources/gsf/tests/conftest.py new file mode 100644 index 000000000..a78f0554a --- /dev/null +++ b/sources/gsf/tests/conftest.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared GSF response fixtures.""" + +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.""" + + 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": [], + "sql_response_from_db": ['[{"revenue":100},{"revenue":200}]'], + } + + +@pytest.fixture +def chat_pql_answer() -> dict: + """Current GSF chat-completions PQL answer envelope.""" + + return { + "response": "A churn prediction query was generated.", + "sql_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.""" + + return { + "request_id": "gsf-request-1", + "thoughts": "- Constructing SQL: Used quarterly_results.", + "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 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 new file mode 100644 index 000000000..bb9e3d7e0 --- /dev/null +++ b/sources/gsf/tests/test_client.py @@ -0,0 +1,570 @@ +# 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 + +import httpx +import pytest +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 + +_TEST_PASSWORD = "${TEST_GSF_PASSWORD}" + + +def _sse_response(answer: dict) -> httpx.Response: + """Build a representative GSF SSE result 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.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: + 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: + """Map a missing catalog endpoint to an unavailable capability.""" + + 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: + """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 + seen_request = request + return _sse_response(chat_sql_answer) + + 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/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", + "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.thoughts == "- Constructing SQL: Used quarterly_results." + assert "response" not in result.model_dump() + + +@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: + nonlocal seen_payload + seen_payload = json.loads(request.content) + return _sse_response(chat_sql_answer) + + 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 seen_payload == {"question": "Show revenue", "prediction": False} + 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_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.""" + + 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": "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: + """Use one Better Auth cookie from sign-in through sign-out.""" + + 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, + } + assert request.headers["origin"] == "https://gsf.example" + assert request.headers["referer"] == "https://gsf.example/" + 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: + """Reuse one password session across multiple GSF calls.""" + + 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: + """Fail before HTTP when neither password nor bearer auth exists.""" + + 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 +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") + + 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="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(chat_sql_answer: dict) -> None: + """Retry a rate-limited request with deterministic jitter.""" + + attempts = 0 + + async def handler(_request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(429) + 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, + 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(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.""" + + 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) + + 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.text_to_sql(TextToSQLRequest(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: + """Reject a malformed non-SSE response.""" + + 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.text_to_sql(TextToSQLRequest(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: + """Map SSE errors without exposing upstream error text.""" + + 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.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 new file mode 100644 index 000000000..417bef732 --- /dev/null +++ b/sources/gsf/tests/test_models.py @@ -0,0 +1,138 @@ +# 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 +from gsf.models import CatalogSearchResponse +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 + + +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", + 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: + """Reject catalog coverage outside the normalized range.""" + + with pytest.raises(ValidationError): + CatalogSearchResponse( + coverage=1.5, + candidates=[ + CatalogCandidate( + label="ColumnAttribute", + attribute="revenue", + term="Revenue", + id="attr:revenue", + ) + ], + ) + + +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) + + 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: + """Accept an optional database scope for SQL requests.""" + + request = TextToSQLRequest(question="Show quarterly revenue", database_name="benchmark_db") + + assert request.database_name == "benchmark_db" + assert request.max_rows == 1_000 + + +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", + "rows": [{"revenue": 100}], + } + ) + + assert result.request_id is None + assert result.thoughts is None + assert result.semantic_context is None + assert result.warnings is 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 + 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: + """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) + + 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 new file mode 100644 index 000000000..6c08370b2 --- /dev/null +++ b/sources/gsf/tests/test_register.py @@ -0,0 +1,193 @@ +# 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 +from unittest.mock import patch + +import pytest +from gsf.models import CatalogSearchResponse +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", + 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) + + +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"]) + + 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_sql"} + + +@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._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: + """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._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 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" + 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_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( + 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: + """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), + ): + 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["code"] == "authentication_required" + client.text_to_sql.assert_not_awaited() 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"