From bbfb17038bb705773dee20fc11424d0d282f171f Mon Sep 17 00:00:00 2001 From: Jay Patel <21jaypatel@gmail.com> Date: Mon, 30 Mar 2026 14:17:35 -0400 Subject: [PATCH 1/3] feat(context-grounding): overhaul CLI with named options and ephemeral support Rewrites the context-grounding CLI with explicit named options, ephemeral index support, and full command coverage. Co-Authored-By: Claude Opus 4.6 (1M context) --- packages/uipath-platform/pyproject.toml | 2 +- .../_context_grounding_service.py | 196 ++++ packages/uipath-platform/uv.lock | 2 +- packages/uipath/pyproject.toml | 2 +- .../_cli/services/cli_context_grounding.py | 854 ++++++++++---- .../cli/contract/test_sdk_cli_alignment.py | 104 +- .../test_context_grounding_commands.py | 1045 ++++++++--------- packages/uipath/uv.lock | 4 +- 8 files changed, 1415 insertions(+), 794 deletions(-) diff --git a/packages/uipath-platform/pyproject.toml b/packages/uipath-platform/pyproject.toml index 5bf170dce..bbec9bc57 100644 --- a/packages/uipath-platform/pyproject.toml +++ b/packages/uipath-platform/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-platform" -version = "0.1.15" +version = "0.1.16" description = "HTTP client library for programmatic access to UiPath Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py b/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py index 601de3964..627143429 100644 --- a/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -1690,6 +1690,184 @@ async def delete_index_async( headers=spec.headers, ) + @resource_override(resource_type="index") + @traced(name="contextgrounding_list", run_type="uipath") + def list_indexes( + self, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> List[ContextGroundingIndex]: + """List all context grounding indexes in a folder. + + If no folder_key or folder_path is provided and no folder context is + configured, falls back to listing across all folders. + + Args: + folder_key (Optional[str]): The key of the folder to list indexes from. + folder_path (Optional[str]): The path of the folder to list indexes from. + + Returns: + List[ContextGroundingIndex]: A list of all indexes in the folder. + """ + resolved_folder_key = self._resolve_folder_key(folder_key, folder_path) + if not resolved_folder_key: + return self.retrieve_across_folders() + + spec = self._list_spec(folder_key=resolved_folder_key) + + response = self.request( + spec.method, + spec.endpoint, + params=spec.params, + headers=spec.headers, + ).json() + + return [ + ContextGroundingIndex.model_validate(item) for item in response["value"] + ] + + @resource_override(resource_type="index") + @traced(name="contextgrounding_list", run_type="uipath") + async def list_indexes_async( + self, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> List[ContextGroundingIndex]: + """Asynchronously list all context grounding indexes in a folder. + + If no folder_key or folder_path is provided and no folder context is + configured, falls back to listing across all folders. + + Args: + folder_key (Optional[str]): The key of the folder to list indexes from. + folder_path (Optional[str]): The path of the folder to list indexes from. + + Returns: + List[ContextGroundingIndex]: A list of all indexes in the folder. + """ + resolved_folder_key = self._resolve_folder_key(folder_key, folder_path) + if not resolved_folder_key: + return await self.retrieve_across_folders_async() + + spec = self._list_spec(folder_key=resolved_folder_key) + + response = ( + await self.request_async( + spec.method, + spec.endpoint, + params=spec.params, + headers=spec.headers, + ) + ).json() + + return [ + ContextGroundingIndex.model_validate(item) for item in response["value"] + ] + + @resource_override(resource_type="index") + @traced(name="contextgrounding_delete_by_name", run_type="uipath") + def delete_by_name( + self, + name: str, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> None: + """Delete a context grounding index by its name. + + This method retrieves the index by name and then deletes it. + + Args: + name (str): The name of the context index to delete. + folder_key (Optional[str]): The key of the folder where the index resides. + folder_path (Optional[str]): The path of the folder where the index resides. + + Raises: + Exception: If no index with the given name is found. + """ + index = self.retrieve(name, folder_key=folder_key, folder_path=folder_path) + self.delete_index(index, folder_key=folder_key, folder_path=folder_path) + + @resource_override(resource_type="index") + @traced(name="contextgrounding_delete_by_name", run_type="uipath") + async def delete_by_name_async( + self, + name: str, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> None: + """Asynchronously delete a context grounding index by its name. + + This method retrieves the index by name and then deletes it. + + Args: + name (str): The name of the context index to delete. + folder_key (Optional[str]): The key of the folder where the index resides. + folder_path (Optional[str]): The path of the folder where the index resides. + + Raises: + Exception: If no index with the given name is found. + """ + index = await self.retrieve_async( + name, folder_key=folder_key, folder_path=folder_path + ) + await self.delete_index_async( + index, folder_key=folder_key, folder_path=folder_path + ) + + @resource_override(resource_type="index") + @traced(name="contextgrounding_exists", run_type="uipath") + @resource_override(resource_type="index") + @traced(name="contextgrounding_ingest_by_name", run_type="uipath") + def ingest_by_name( + self, + name: str, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> None: + """Trigger ingestion on a context grounding index by its name. + + This method retrieves the index by name and then triggers data ingestion. + + Args: + name (str): The name of the context index to ingest. + folder_key (Optional[str]): The key of the folder where the index resides. + folder_path (Optional[str]): The path of the folder where the index resides. + + Raises: + Exception: If no index with the given name is found. + IngestionInProgressException: If ingestion is already in progress. + """ + index = self.retrieve(name, folder_key=folder_key, folder_path=folder_path) + self.ingest_data(index, folder_key=folder_key, folder_path=folder_path) + + @resource_override(resource_type="index") + @traced(name="contextgrounding_ingest_by_name", run_type="uipath") + async def ingest_by_name_async( + self, + name: str, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> None: + """Asynchronously trigger ingestion on a context grounding index by its name. + + This method retrieves the index by name and then triggers data ingestion. + + Args: + name (str): The name of the context index to ingest. + folder_key (Optional[str]): The key of the folder where the index resides. + folder_path (Optional[str]): The path of the folder where the index resides. + + Raises: + Exception: If no index with the given name is found. + IngestionInProgressException: If ingestion is already in progress. + """ + index = await self.retrieve_async( + name, folder_key=folder_key, folder_path=folder_path + ) + await self.ingest_data_async( + index, folder_key=folder_key, folder_path=folder_path + ) + def _ingest_spec( self, key: str, @@ -1722,6 +1900,24 @@ def _retrieve_across_folders_spec( params=params, ) + def _list_spec( + self, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> RequestSpec: + folder_key = self._resolve_folder_key(folder_key, folder_path) + + return RequestSpec( + method="GET", + endpoint=Endpoint("/ecs_/v2/indexes"), + params={ + "$expand": "dataSource", + }, + headers={ + **header_folder(folder_key, None), + }, + ) + def _retrieve_spec( self, name: str, diff --git a/packages/uipath-platform/uv.lock b/packages/uipath-platform/uv.lock index 1883bda06..9065b5ada 100644 --- a/packages/uipath-platform/uv.lock +++ b/packages/uipath-platform/uv.lock @@ -1088,7 +1088,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.1.15" +version = "0.1.16" source = { editable = "." } dependencies = [ { name = "httpx" }, diff --git a/packages/uipath/pyproject.toml b/packages/uipath/pyproject.toml index ae3ebae89..380ae367d 100644 --- a/packages/uipath/pyproject.toml +++ b/packages/uipath/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath" -version = "2.10.38" +version = "2.10.39" description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools." readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py b/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py index b086120fc..53e225c02 100644 --- a/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py +++ b/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py @@ -1,26 +1,14 @@ """Context Grounding service commands for UiPath CLI. -Context Grounding provides semantic search over indexed document collections, -enabling RAG (Retrieval-Augmented Generation) for automation processes. - -Commands: - list - List all indexes in a folder - retrieve - Get details of a specific index by name - search - Perform semantic search against an index - ingest - Trigger re-ingestion of an index - delete - Delete an index by name - -Note: - create_index is intentionally not exposed here — its source configuration - (bucket, Google Drive, OneDrive, Dropbox, Confluence) is too complex to - express cleanly as CLI flags. Use the Python SDK directly for index creation. +Context Grounding provides capabilities for indexing, retrieving, and searching +through contextual information used to enhance AI-enabled automation. """ -# ruff: noqa: D301 - Using regular """ strings (not r""") for Click \b formatting +# ruff: noqa: D301 - Click uses \b in docstrings for formatting -from typing import Any, NoReturn, Optional +import json as json_module +from typing import Any, Optional import click -from httpx import HTTPStatusError from .._utils._service_base import ( ServiceCommandBase, @@ -29,60 +17,27 @@ service_command, ) -# The SDK raises a bare Exception with this message when an index name doesn't exist. -_INDEX_NOT_FOUND_MSG = "ContextGroundingIndex not found" - - -def _handle_retrieve_error(index_name: str, e: Exception) -> NoReturn: - """Convert a retrieve() exception to a clean ClickException. - - Mirrors the pattern used in cli_buckets.py: - - bare Exception with the SDK's not-found message → structured not-found error - - HTTPStatusError 404 → structured not-found error - - anything else → re-raise so service_command's handler deals with it - """ - if isinstance(e, HTTPStatusError): - if e.response.status_code == 404: - handle_not_found_error("Index", index_name, e) - raise e - if _INDEX_NOT_FOUND_MSG.lower() in str(e).lower(): - handle_not_found_error("Index", index_name) - raise click.ClickException(str(e)) from e - @click.group(name="context-grounding") def context_grounding() -> None: - """Manage UiPath Context Grounding indexes and perform semantic search. + """Manage UiPath Context Grounding indexes. - Context Grounding indexes documents from storage buckets or cloud drives - and makes them searchable via natural language queries (RAG). + Context Grounding indexes store and search contextual information + used to enhance AI-enabled automation processes. \b - Commands: - list - List all indexes in a folder - retrieve - Get details of a specific index - search - Query an index with natural language - ingest - Trigger re-ingestion of an index - delete - Delete an index + Two index types: + Regular - Persistent, backed by bucket or connection, created via 'create' + Ephemeral - Temporary, no Orchestrator folder, created from local files via 'create-ephemeral' \b Examples: uipath context-grounding list --folder-path "Shared" - uipath context-grounding retrieve --index "my-index" --folder-path "Shared" - uipath context-grounding search "how to process invoices" --index "my-index" --folder-path "Shared" - uipath context-grounding search "payment terms" --index "my-index" --folder-path "Shared" --limit 5 - uipath context-grounding ingest --index "my-index" --folder-path "Shared" - uipath context-grounding delete --index "my-index" --folder-path "Shared" --confirm - - \b - Folder context: - Set UIPATH_FOLDER_PATH to avoid passing --folder-path on every command: - export UIPATH_FOLDER_PATH="Shared" """ pass -@context_grounding.command("list") +@context_grounding.command(name="list") @common_service_options @service_command def list_indexes( @@ -92,228 +47,316 @@ def list_indexes( format: Optional[str], output: Optional[str], ) -> Any: - """List all context grounding indexes in a folder. + """List all context grounding indexes. \b Examples: uipath context-grounding list --folder-path "Shared" - uipath context-grounding list --folder-path "Shared" --format json """ client = ServiceCommandBase.get_client(ctx) - results = client.context_grounding.list( + return client.context_grounding.list_indexes( folder_path=folder_path, folder_key=folder_key, ) - # Table format: project to key fields for readability. - # JSON/CSV: return full objects so nothing is lost for scripting. - fmt = format or "table" - if fmt == "table": - return [ - { - "name": ix.name, - "last_ingestion_status": ix.last_ingestion_status, - "last_ingested": ix.last_ingested, - "description": ix.description, - } - for ix in results - ] - - return results - -@context_grounding.command("retrieve") -@click.option( - "--index", "index_name", required=True, help="Name of the index to retrieve." -) +@context_grounding.command(name="retrieve") +@click.option("--index-name", help="Name of the index to retrieve") +@click.option("--index-id", help="ID of the index to retrieve (ephemeral indexes only)") @common_service_options @service_command -def retrieve( +def retrieve_index( ctx: click.Context, - index_name: str, + index_name: Optional[str], + index_id: Optional[str], folder_path: Optional[str], folder_key: Optional[str], format: Optional[str], output: Optional[str], ) -> Any: - """Retrieve details of a context grounding index by name. + """Retrieve a context grounding index. + + \b + Two ways to specify the index: + Regular index: --index-name + --folder-path + Ephemeral index: --index-id \b Examples: - uipath context-grounding retrieve --index "my-index" --folder-path "Shared" - uipath context-grounding retrieve --index "my-index" --folder-path "Shared" --format json + uipath context-grounding retrieve --index-name my-index --folder-path "Shared" + uipath context-grounding retrieve --index-id abc-123-def-456 --format json """ + from httpx import HTTPStatusError + + if not index_name and not index_id: + raise click.UsageError("Either --index-name or --index-id must be provided.") + if index_name and index_id: + raise click.UsageError("Provide either --index-name or --index-id, not both.") + client = ServiceCommandBase.get_client(ctx) + if index_id: + from uipath.platform.context_grounding import ContextGroundingIndex + + raw = client.context_grounding.retrieve_by_id( + index_id, + folder_path=folder_path, + folder_key=folder_key, + ) + return ContextGroundingIndex(**raw) + + assert index_name is not None # validated above + name = index_name try: return client.context_grounding.retrieve( - name=index_name, + name=name, folder_path=folder_path, folder_key=folder_key, ) - except Exception as e: - _handle_retrieve_error(index_name, e) + except LookupError: + handle_not_found_error("Index", name) + except HTTPStatusError as e: + if e.response.status_code == 404: + handle_not_found_error("Index", name, e) + raise -@context_grounding.command("search") -@click.argument("query") +@context_grounding.command(name="create") +@click.option("--index-name", required=True, help="Name of the index to create") @click.option( - "--index", "index_name", required=True, help="Name of the index to search." + "--source-file", + type=click.Path(exists=True), + help="JSON file with connection source configuration (Google Drive, OneDrive, Dropbox, Confluence)", ) @click.option( - "--limit", - "-n", - "number_of_results", - type=click.IntRange(min=1), - default=10, - show_default=True, - help="Maximum number of results to return.", + "--bucket-source", + help="Bucket name for bucket-backed indexes", ) +@click.option("--description", default="", help="Description of the index") @click.option( - "--threshold", - type=float, - default=None, - help="Minimum relevance score threshold (0.0–1.0). Results below this score are excluded.", + "--extraction-strategy", + type=click.Choice(["LLMV4", "NativeV1"]), + default="LLMV4", + help="Extraction strategy (default: LLMV4)", ) +@click.option("--file-type", help="File type filter (e.g., 'pdf', 'txt')") @common_service_options @service_command -def search( +def create_index( ctx: click.Context, - query: str, index_name: str, - number_of_results: int, - threshold: Optional[float], + source_file: Optional[str], + bucket_source: Optional[str], + description: str, + extraction_strategy: str, + file_type: Optional[str], folder_path: Optional[str], folder_key: Optional[str], format: Optional[str], output: Optional[str], ) -> Any: - """Search an index with a natural language query. + """Create a new context grounding index (persistent). - QUERY is the natural language search string. + The created index lives in an Orchestrator folder. Ingestion must be + triggered separately after creation. + + \b + Two ways to specify the data source: + --bucket-source Bucket name for bucket-backed indexes + --source-file JSON file for connections (use 'source-schema' to see formats) \b Examples: - uipath context-grounding search "how to process invoices" --index "my-index" --folder-path "Shared" - uipath context-grounding search "payment terms" --index "my-index" --folder-path "Shared" --limit 5 - uipath context-grounding search "invoice" --index "my-index" --folder-path "Shared" --threshold 0.7 - uipath context-grounding search "approval workflow" --index "my-index" --folder-path "Shared" --format json - uipath context-grounding search "invoice policy" --index "my-index" --folder-path "Shared" -o results.json + uipath context-grounding create --index-name my-index --bucket-source my-bucket + uipath context-grounding create --index-name my-index --source-file config.json """ - from uipath.platform.errors import IngestionInProgressException + if source_file and bucket_source: + raise click.UsageError( + "Cannot use both --source-file and --bucket-source. Choose one." + ) + if not source_file and not bucket_source: + raise click.UsageError( + "Either --source-file or --bucket-source must be provided." + ) - client = ServiceCommandBase.get_client(ctx) + source: Any + if source_file: + from pydantic import TypeAdapter - try: - results = client.context_grounding.search( - name=index_name, - query=query, - number_of_results=number_of_results, - threshold=threshold, - folder_path=folder_path, - folder_key=folder_key, + from uipath.platform.context_grounding import BucketSourceConfig, SourceConfig + + with open(source_file) as f: + source_data = json_module.load(f) + + source = TypeAdapter(SourceConfig).validate_python(source_data) + else: + from uipath.platform.context_grounding import BucketSourceConfig + + assert bucket_source is not None # validated above + source = BucketSourceConfig( + bucket_name=bucket_source, + folder_path=folder_path or "", + file_type=file_type, ) - except IngestionInProgressException as e: - raise click.ClickException( - f"Index '{index_name}' is currently being ingested. " - "Please wait for ingestion to complete and try again." - ) from e - except Exception as e: - _handle_retrieve_error(index_name, e) - - if not results: - click.echo("No results found.", err=True) - return None - - # Table format: show only human-readable columns, truncate content. - # JSON/CSV: return full objects so nothing is lost for scripting. - fmt = format or "table" - if fmt == "table": - rows = [] - for r in results: - content = r.content - if len(content) > 120: - content = content[:120] + "…" - rows.append( - { - "score": round(r.score, 3) if r.score is not None else None, - "source": r.source, - "page_number": r.page_number, - "content": content, - } - ) - return rows - return results + client = ServiceCommandBase.get_client(ctx) + result = client.context_grounding.create_index( + name=index_name, + source=source, + description=description, + extraction_strategy=extraction_strategy, + folder_path=folder_path, + folder_key=folder_key, + ) + return result + + +_SOURCE_TYPES = { + "google_drive": "GoogleDriveSourceConfig", + "onedrive": "OneDriveSourceConfig", + "dropbox": "DropboxSourceConfig", + "confluence": "ConfluenceSourceConfig", +} -@context_grounding.command("ingest") +@context_grounding.command(name="source-schema") @click.option( - "--index", "index_name", required=True, help="Name of the index to re-ingest." + "--type", + "source_type", + type=click.Choice(list(_SOURCE_TYPES.keys())), + help="Show schema for a specific source type (omit to show all)", +) +def source_schema(source_type: Optional[str]) -> None: + """Show JSON source file formats for connection-backed indexes. + + Use this to see the required fields for --source-file when creating + an index backed by Google Drive, OneDrive, Dropbox, or Confluence. + + \b + Examples: + uipath context-grounding source-schema --type google_drive + """ + import json + + from uipath.platform.context_grounding import context_grounding_payloads as payloads + + types_to_show = ( + {source_type: _SOURCE_TYPES[source_type]} if source_type else _SOURCE_TYPES + ) + + for type_key, class_name in types_to_show.items(): + cls = getattr(payloads, class_name) + schema = cls.model_json_schema() + required = schema.get("required", []) + props = schema.get("properties", {}) + example = {} + for field_name, info in props.items(): + if field_name in required: + default = info.get("default") + example[field_name] = ( + default + if default is not None + else f"<{info.get('description', field_name)}>" + ) + click.echo(f"\n{type_key}:", err=False) + click.echo(json.dumps(example, indent=2), err=False) + + click.echo("", err=False) + + +@context_grounding.command(name="create-ephemeral") +@click.option( + "--usage", + required=True, + type=click.Choice(["DeepRAG", "BatchRAG"]), + help="Task type for the ephemeral index", +) +@click.option( + "--files", + multiple=True, + required=True, + type=click.Path(exists=True), + help="Local file paths to upload as attachments (repeatable)", ) @common_service_options @service_command -def ingest( +def create_ephemeral( ctx: click.Context, - index_name: str, + usage: str, + files: tuple[str, ...], folder_path: Optional[str], folder_key: Optional[str], format: Optional[str], output: Optional[str], -) -> None: - """Trigger re-ingestion of a context grounding index. +) -> Any: + """Create an ephemeral index from local files (temporary). + + Uploads files as attachments and creates a temporary index. Reference it + in other commands with --index-id (no folder, no name). Ingestion starts + automatically. Poll with 'retrieve --index-id ' until + lastIngestionStatus is Successful before starting a task. + + \b + Supported file types: + DeepRAG: PDF, TXT + BatchRAG: CSV \b Examples: - uipath context-grounding ingest --index "my-index" --folder-path "Shared" + uipath context-grounding create-ephemeral --usage DeepRAG --files doc1.pdf --files doc2.pdf + uipath context-grounding create-ephemeral --usage BatchRAG --files data.csv """ - from uipath.platform.errors import IngestionInProgressException + from pathlib import Path + + from uipath.platform.context_grounding import EphemeralIndexUsage + + allowed_extensions = { + "DeepRAG": {".pdf", ".txt"}, + "BatchRAG": {".csv"}, + } + allowed = allowed_extensions[usage] + for file_path in files: + ext = Path(file_path).suffix.lower() + if ext not in allowed: + raise click.UsageError( + f"File '{Path(file_path).name}' has unsupported extension '{ext}' for {usage}. " + f"Supported: {', '.join(sorted(allowed))}" + ) client = ServiceCommandBase.get_client(ctx) - try: - index = client.context_grounding.retrieve( - name=index_name, + attachment_ids = [] + for file_path in files: + file_name = Path(file_path).name + click.echo(f"Uploading '{file_name}'...", err=True) + attachment_id = client.attachments.upload( + name=file_name, + source_path=file_path, folder_path=folder_path, folder_key=folder_key, ) - if not index.id: - raise click.ClickException( - f"Index '{index_name}' has no ID and cannot be ingested." - ) - if index.in_progress_ingestion(): - raise click.ClickException( - f"Index '{index_name}' is already being ingested." - ) - client.context_grounding.ingest_data( - index=index, - folder_path=folder_path, - folder_key=folder_key, - ) - except click.ClickException: - raise - except IngestionInProgressException as e: - # Catches the 409 race condition from ingest_data() itself. - raise click.ClickException( - f"Index '{index_name}' is already being ingested." - ) from e - except Exception as e: - _handle_retrieve_error(index_name, e) + attachment_ids.append(str(attachment_id)) + click.echo(f"Uploaded '{file_name}' (ID: {attachment_id})", err=True) - click.echo(f"Triggered ingestion for index '{index_name}'.", err=True) + click.echo("Creating ephemeral index...", err=True) + index = client.context_grounding.create_ephemeral_index( + usage=EphemeralIndexUsage(usage), + attachments=attachment_ids, + ) + click.echo(f"Ephemeral index created: {index.id}", err=True) + return index -@context_grounding.command("delete") -@click.option( - "--index", "index_name", required=True, help="Name of the index to delete." -) -@click.option("--confirm", is_flag=True, help="Skip confirmation prompt.") + +@context_grounding.command(name="delete") +@click.option("--index-name", required=True, help="Name of the index to delete") +@click.option("--confirm", is_flag=True, help="Skip confirmation prompt") @click.option( - "--dry-run", is_flag=True, help="Show what would be deleted without deleting." + "--dry-run", is_flag=True, help="Show what would be deleted without deleting" ) @common_service_options @service_command -def delete( +def delete_index( ctx: click.Context, index_name: str, confirm: bool, @@ -322,48 +365,421 @@ def delete( folder_key: Optional[str], format: Optional[str], output: Optional[str], -) -> None: - """Delete a context grounding index by name. +) -> Any: + """Delete a context grounding index. \b Examples: - uipath context-grounding delete --index "my-index" --folder-path "Shared" --confirm - uipath context-grounding delete --index "my-index" --folder-path "Shared" --dry-run + uipath context-grounding delete --index-name my-index --confirm + uipath context-grounding delete --index-name my-index --dry-run """ + from httpx import HTTPStatusError + client = ServiceCommandBase.get_client(ctx) - # Resolve the index object first — surfaces not-found before prompting the user. + # First retrieve to verify index exists try: - index = client.context_grounding.retrieve( + client.context_grounding.retrieve( name=index_name, folder_path=folder_path, folder_key=folder_key, ) - except Exception as e: - _handle_retrieve_error(index_name, e) - - if not index.id: - raise click.ClickException( - f"Index '{index_name}' has no ID and cannot be deleted." - ) + except LookupError: + handle_not_found_error("Index", index_name) + except HTTPStatusError as e: + if e.response.status_code == 404: + handle_not_found_error("Index", index_name, e) + raise - # dry-run and confirmation after index is confirmed to exist and have an ID. + # Handle dry-run if dry_run: - click.echo(f"Would delete index '{index_name}'.", err=True) + click.echo(f"Would delete index '{index_name}'", err=True) return + # Handle confirmation if not confirm: if not click.confirm(f"Delete index '{index_name}'?"): click.echo("Deletion cancelled.") return - try: - client.context_grounding.delete_index( - index=index, - folder_path=folder_path, - folder_key=folder_key, - ) - except Exception as e: - raise click.ClickException(f"Failed to delete index '{index_name}': {e}") from e + # Perform delete + client.context_grounding.delete_by_name( + name=index_name, + folder_path=folder_path, + folder_key=folder_key, + ) + + click.echo(f"Deleted index '{index_name}'", err=True) + + +@context_grounding.command(name="ingest") +@click.option("--index-name", required=True, help="Name of the index to ingest") +@common_service_options +@service_command +def ingest_index( + ctx: click.Context, + index_name: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> None: + """Trigger ingestion on a context grounding index. + + Ingestion runs asynchronously. Use 'retrieve' to poll lastIngestionStatus + until it reaches Successful or Failed. + + \b + Examples: + uipath context-grounding ingest --index-name my-index --folder-path "Shared" + """ + client = ServiceCommandBase.get_client(ctx) + client.context_grounding.ingest_by_name( + name=index_name, + folder_path=folder_path, + folder_key=folder_key, + ) + + click.echo(f"Ingestion triggered for index '{index_name}'", err=True) + + +@context_grounding.command(name="search") +@click.option("--index-name", required=True, help="Name of the index to search") +@click.option("--query", required=True, help="Search query in natural language") +@click.option( + "--limit", + type=click.IntRange(min=1), + default=10, + help="Maximum number of results (default: 10)", +) +@click.option( + "--threshold", + type=click.FloatRange(min=0.0, max=1.0), + default=0.0, + help="Minimum similarity threshold (default: 0.0)", +) +@click.option( + "--search-mode", + type=click.Choice(["Auto", "Semantic"]), + default="Auto", + help="Search mode (default: Auto)", +) +@common_service_options +@service_command +def search_index( + ctx: click.Context, + index_name: str, + query: str, + limit: int, + threshold: float, + search_mode: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> Any: + """Search a context grounding index (regular indexes only). + + \b + Examples: + uipath context-grounding search --index-name my-index --query "What is the revenue?" + uipath context-grounding search --index-name my-index --query "results" --limit 5 + """ + from uipath.platform.context_grounding import SearchMode + + client = ServiceCommandBase.get_client(ctx) + return client.context_grounding.unified_search( + name=index_name, + query=query, + number_of_results=limit, + threshold=threshold, + search_mode=SearchMode(search_mode), + folder_path=folder_path, + folder_key=folder_key, + ) + + +# --- Deep RAG nested group --- + + +@click.group(name="deep-rag") +def deep_rag() -> None: + """Manage Deep RAG tasks. + + Deep RAG performs multi-document research and synthesis on context + grounding indexes. + + \b + Examples: + uipath context-grounding deep-rag start --help + """ + pass + + +@deep_rag.command(name="start") +@click.option("--index-name", help="Name of the context grounding index") +@click.option( + "--index-id", help="ID of the context grounding index (ephemeral indexes only)" +) +@click.option("--task-name", required=True, help="Name for the Deep RAG task") +@click.option("--prompt", required=True, help="Task prompt describing what to research") +@click.option( + "--glob-pattern", + default="**", + help="Glob pattern to filter files in the index (default: **)", +) +@click.option( + "--citation-mode", + type=click.Choice(["Skip", "Inline"]), + default="Skip", + help="Citation mode (default: Skip)", +) +@common_service_options +@service_command +def start_deep_rag( + ctx: click.Context, + index_name: Optional[str], + index_id: Optional[str], + task_name: str, + prompt: str, + glob_pattern: str, + citation_mode: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> Any: + """Start a Deep RAG task on an index. + + \b + Two ways to specify the index: + Regular index: --index-name + --folder-path + Ephemeral index: --index-id + + \b + Examples: + uipath context-grounding deep-rag start --index-name my-index --folder-path Shared --task-name my-task --prompt "Summarize" + uipath context-grounding deep-rag start --index-id abc-123 --task-name my-task --prompt "Summarize" + """ + from uipath.platform.context_grounding import CitationMode + + if not index_name and not index_id: + raise click.UsageError("Either --index-name or --index-id must be provided.") + if index_name and index_id: + raise click.UsageError("Provide either --index-name or --index-id, not both.") + + client = ServiceCommandBase.get_client(ctx) + result = client.context_grounding.start_deep_rag( + name=task_name, + prompt=prompt, + index_name=index_name, + index_id=index_id, + glob_pattern=glob_pattern, + citation_mode=CitationMode(citation_mode), + folder_path=folder_path, + folder_key=folder_key, + ) + + click.echo(f"Deep RAG task started: {result.id}", err=True) + + return result + + +@deep_rag.command(name="retrieve") +@click.option("--task-id", required=True, help="ID of the Deep RAG task") +@common_service_options +@service_command +def retrieve_deep_rag( + ctx: click.Context, + task_id: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> Any: + """Retrieve a Deep RAG task result (status, summary, citations). + + \b + Examples: + uipath context-grounding deep-rag retrieve --task-id abc-123-def-456 + """ + client = ServiceCommandBase.get_client(ctx) + return client.context_grounding.retrieve_deep_rag(id=task_id) + + +context_grounding.add_command(deep_rag) + + +# --- Batch Transform nested group --- + + +@click.group(name="batch-transform") +def batch_transform() -> None: + """Manage Batch Transform tasks. + + Batch Transform processes and transforms CSV files from context + grounding indexes. + + \b + Examples: + uipath context-grounding batch-transform start --help + """ + pass + + +@batch_transform.command(name="start") +@click.option("--index-name", help="Name of the context grounding index") +@click.option( + "--index-id", help="ID of the context grounding index (ephemeral indexes only)" +) +@click.option("--task-name", required=True, help="Name for the Batch Transform task") +@click.option("--prompt", required=True, help="Task prompt describing what to process") +@click.option( + "--columns-file", + required=True, + type=click.Path(exists=True), + help="JSON file defining output columns (see format above)", +) +@click.option("--target-file", help="Specific file name to target in the index") +@click.option( + "--prefix", + help="Storage bucket folder path prefix for filtering files", +) +@click.option( + "--web-search", + is_flag=True, + help="Enable web search grounding", +) +@common_service_options +@service_command +def start_batch_transform( + ctx: click.Context, + index_name: Optional[str], + index_id: Optional[str], + task_name: str, + prompt: str, + columns_file: str, + target_file: Optional[str], + prefix: Optional[str], + web_search: bool, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> Any: + """Start a Batch Transform task on an index. + + The index must contain CSV files. Only one file is processed per task. + + \b + Two ways to specify the index: + Regular index: --index-name + --folder-path + Ephemeral index: --index-id + + \b + --columns-file is a JSON array defining output columns: + [ + {"name": "entity", "description": "Extracted entity name"}, + {"name": "category", "description": "Entity category"} + ] + + \b + Examples: + uipath context-grounding batch-transform start --index-name my-index --task-name my-task --prompt "Extract" --columns-file cols.json + uipath context-grounding batch-transform start --index-id abc-123 --task-name my-task --prompt "Extract" --columns-file cols.json + """ + from uipath.platform.context_grounding import BatchTransformOutputColumn + + if not index_name and not index_id: + raise click.UsageError("Either --index-name or --index-id must be provided.") + if index_name and index_id: + raise click.UsageError("Provide either --index-name or --index-id, not both.") + + with open(columns_file) as f: + columns_data = json_module.load(f) + + output_columns = [BatchTransformOutputColumn(**col) for col in columns_data] + + client = ServiceCommandBase.get_client(ctx) + result = client.context_grounding.start_batch_transform( + name=task_name, + prompt=prompt, + index_name=index_name, + index_id=index_id, + output_columns=output_columns, + storage_bucket_folder_path_prefix=prefix, + target_file_name=target_file, + enable_web_search_grounding=web_search, + folder_path=folder_path, + folder_key=folder_key, + ) + + click.echo(f"Batch Transform task started: {result.id}", err=True) + + return result + + +@batch_transform.command(name="retrieve") +@click.option("--task-id", required=True, help="ID of the Batch Transform task") +@common_service_options +@service_command +def retrieve_batch_transform( + ctx: click.Context, + task_id: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> Any: + """Retrieve a Batch Transform task status. + + \b + Examples: + uipath context-grounding batch-transform retrieve --task-id abc-123-def-456 + """ + client = ServiceCommandBase.get_client(ctx) + return client.context_grounding.retrieve_batch_transform(id=task_id) + + +@batch_transform.command(name="download") +@click.option("--task-id", required=True, help="ID of the Batch Transform task") +@click.option( + "--output-file", + required=True, + type=click.Path(), + help="Local destination path for the result file", +) +@common_service_options +@service_command +def download_batch_transform( + ctx: click.Context, + task_id: str, + output_file: str, + folder_path: Optional[str], + folder_key: Optional[str], + format: Optional[str], + output: Optional[str], +) -> None: + """Download a Batch Transform result file. + + \b + Examples: + uipath context-grounding batch-transform download --task-id abc-123 --output-file result.csv + """ + client = ServiceCommandBase.get_client(ctx) + client.context_grounding.download_batch_transform_result( + id=task_id, + destination_path=output_file, + ) + + click.echo(f"Downloaded to '{output_file}'", err=True) + + +context_grounding.add_command(batch_transform) + - click.echo(f"Deleted index '{index_name}'.", err=True) +def __getattr__(name: str) -> click.Group: + """Allow lazy loading with hyphenated command name.""" + if name == "context-grounding": + return context_grounding + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/packages/uipath/tests/cli/contract/test_sdk_cli_alignment.py b/packages/uipath/tests/cli/contract/test_sdk_cli_alignment.py index 9390f0b06..52699d42e 100644 --- a/packages/uipath/tests/cli/contract/test_sdk_cli_alignment.py +++ b/packages/uipath/tests/cli/contract/test_sdk_cli_alignment.py @@ -118,18 +118,43 @@ def assert_cli_sdk_alignment( # Parameter mappings: CLI param name → SDK param name # Used when CLI uses more user-friendly names than SDK PARAM_MAPPINGS = { - # context-grounding: --index maps to SDK 'name', --folder maps to SDK 'folder_path' + # context-grounding: --index-name maps to SDK 'name' "context-grounding_list": {}, "context-grounding_retrieve": { "index_name": "name", }, + "context-grounding_create": { + "index_name": "name", + }, "context-grounding_search": { "index_name": "name", }, - # ingest/delete: SDK takes an index *object*, not a name. CLI's --index is a - # lookup key used to call retrieve() internally — it has no direct SDK counterpart. - "context-grounding_ingest": {}, - "context-grounding_delete": {}, + "context-grounding_ingest": { + "index_name": "name", + }, + "context-grounding_delete": { + "index_name": "name", + }, + # deep-rag/batch-transform: --task-name maps to SDK 'name' + "context-grounding_deep-rag_start": { + "task_name": "name", + }, + "context-grounding_deep-rag_retrieve": { + "task_id": "id", + }, + "context-grounding_batch-transform_start": { + "task_name": "name", + "prefix": "storage_bucket_folder_path_prefix", + "target_file": "target_file_name", + "web_search": "enable_web_search_grounding", + }, + "context-grounding_batch-transform_retrieve": { + "task_id": "id", + }, + "context-grounding_batch-transform_download": { + "task_id": "id", + "output_file": "destination_path", + }, "buckets_files_download": { "bucket_name": "name", "remote_path": "blob_file_path", @@ -166,13 +191,17 @@ def assert_cli_sdk_alignment( # SDK parameters to exclude for specific commands # Used when SDK has optional params that CLI doesn't expose SDK_EXCLUSIONS = { - # context-grounding: ingest/delete take an index *object*; CLI uses --index name - # then retrieves internally — so the SDK 'index' object param is not a CLI option. "context-grounding_list": set(), "context-grounding_retrieve": set(), - "context-grounding_search": set(), - "context-grounding_ingest": {"index"}, - "context-grounding_delete": {"index"}, + "context-grounding_create": {"source", "embeddings_enabled", "is_encrypted"}, + "context-grounding_search": {"scope", "number_of_results"}, + "context-grounding_ingest": set(), + "context-grounding_delete": set(), + "context-grounding_deep-rag_start": set(), + "context-grounding_deep-rag_retrieve": {"index_name"}, + "context-grounding_batch-transform_start": {"output_columns"}, + "context-grounding_batch-transform_retrieve": {"index_name"}, + "context-grounding_batch-transform_download": {"index_name", "validate_status"}, "buckets_list": { "name", "skip", @@ -208,11 +237,42 @@ def assert_cli_sdk_alignment( "service,command,sdk_class,sdk_method", [ # Context Grounding - ("context-grounding", "list", "ContextGroundingService", "list"), + ("context-grounding", "list", "ContextGroundingService", "list_indexes"), ("context-grounding", "retrieve", "ContextGroundingService", "retrieve"), - ("context-grounding", "search", "ContextGroundingService", "search"), - ("context-grounding", "ingest", "ContextGroundingService", "ingest_data"), - ("context-grounding", "delete", "ContextGroundingService", "delete_index"), + ("context-grounding", "create", "ContextGroundingService", "create_index"), + ("context-grounding", "search", "ContextGroundingService", "unified_search"), + ("context-grounding", "ingest", "ContextGroundingService", "ingest_by_name"), + ("context-grounding", "delete", "ContextGroundingService", "delete_by_name"), + ( + "context-grounding_deep-rag", + "start", + "ContextGroundingService", + "start_deep_rag", + ), + ( + "context-grounding_deep-rag", + "retrieve", + "ContextGroundingService", + "retrieve_deep_rag", + ), + ( + "context-grounding_batch-transform", + "start", + "ContextGroundingService", + "start_batch_transform", + ), + ( + "context-grounding_batch-transform", + "retrieve", + "ContextGroundingService", + "retrieve_batch_transform", + ), + ( + "context-grounding_batch-transform", + "download", + "ContextGroundingService", + "download_batch_transform_result", + ), # Buckets - bucket operations ("buckets", "list", "BucketsService", "list"), ("buckets", "retrieve", "BucketsService", "retrieve"), @@ -251,7 +311,11 @@ def test_service_command_params_match_sdk(service, command, sdk_class, sdk_metho # cli_services.buckets.commands["files"].commands[command]). # All other service keys resolve directly to a cli_services attribute # (hyphens are converted to underscores: "context-grounding" → context_grounding). - NESTED_SERVICES = {"buckets_files"} + NESTED_SERVICES = { + "buckets_files", + "context-grounding_deep-rag", + "context-grounding_batch-transform", + } def _get_service_group(name: str) -> click.Group: """Resolve a cli_services attribute, converting hyphens to underscores.""" @@ -268,10 +332,12 @@ def _get_service_group(name: str) -> click.Group: # CLI-only options per command (beyond the global CLI_ONLY_OPTIONS set). CLI_EXCLUSIONS: dict[str, set[str]] = { - # ingest/delete use --index to look up the index object via retrieve(); - # index_name is a CLI-only lookup key with no direct SDK parameter counterpart. - "context-grounding_ingest": {"index_name"}, - "context-grounding_delete": {"index_name"}, + # retrieve: --index-id routes to retrieve_by_id, not retrieve — exclude from alignment + "context-grounding_retrieve": {"index_id"}, + # create: CLI-only options that build the source object internally + "context-grounding_create": {"source_file", "bucket_source", "file_type"}, + # batch-transform start: --columns-file is CLI-only (reads JSON, builds output_columns) + "context-grounding_batch-transform_start": {"columns_file"}, } # Get mappings and exclusions for this command diff --git a/packages/uipath/tests/cli/integration/test_context_grounding_commands.py b/packages/uipath/tests/cli/integration/test_context_grounding_commands.py index 75cf88670..7296d7080 100644 --- a/packages/uipath/tests/cli/integration/test_context_grounding_commands.py +++ b/packages/uipath/tests/cli/integration/test_context_grounding_commands.py @@ -1,20 +1,17 @@ """Integration tests for context-grounding CLI commands. -These tests verify end-to-end functionality of the context-grounding service -commands, including proper context handling, error messages, and output formatting. +Tests verify end-to-end functionality of all context-grounding commands, +including proper option handling, error messages, and output formatting. """ +import json from unittest.mock import MagicMock, patch import pytest from click.testing import CliRunner from uipath._cli import cli -from uipath.platform.context_grounding import ( - ContextGroundingIndex, - ContextGroundingQueryResponse, -) -from uipath.platform.errors import IngestionInProgressException +from uipath.platform.context_grounding import ContextGroundingIndex @pytest.fixture @@ -30,24 +27,22 @@ def mock_client(): client_instance = MagicMock() mock.return_value = client_instance client_instance.context_grounding = MagicMock() + client_instance.attachments = MagicMock() yield client_instance def _make_index(name="my-index", status="Completed", description="Test index"): - """Helper: build a mock ContextGroundingIndex. - - Uses spec=ContextGroundingIndex so MagicMock is not mistaken for an - Iterator by format_output (MagicMock implements __iter__ by default). - in_progress_ingestion() returns False by default (ingestion complete). - """ + """Build a mock ContextGroundingIndex.""" index = MagicMock(spec=ContextGroundingIndex) index.id = "test-index-id" index.name = name index.last_ingestion_status = status index.last_ingested = None index.description = description + index.folder_key = "test-folder-key" index.in_progress_ingestion.return_value = False index.model_dump.return_value = { + "id": "test-index-id", "name": name, "last_ingestion_status": status, "description": description, @@ -55,27 +50,6 @@ def _make_index(name="my-index", status="Completed", description="Test index"): return index -def _make_result(source="doc.pdf", page="1", content="Some content", score=0.95): - """Helper: build a mock ContextGroundingQueryResponse. - - Uses spec=ContextGroundingQueryResponse for the same reason as _make_index. - """ - result = MagicMock(spec=ContextGroundingQueryResponse) - result.source = source - result.page_number = page - result.content = content - result.score = score - result.model_dump.return_value = { - "source": source, - "page_number": page, - "content": content, - "score": score, - } - return result - - -# --------------------------------------------------------------------------- -# retrieve # --------------------------------------------------------------------------- # list # --------------------------------------------------------------------------- @@ -83,816 +57,778 @@ def _make_result(source="doc.pdf", page="1", content="Some content", score=0.95) class TestListCommand: def test_list_basic(self, runner, mock_client, mock_env_vars): - """list returns all indexes in a folder.""" - mock_client.context_grounding.list.return_value = [ - _make_index(name="index-one", status="Completed"), - _make_index(name="index-two", status="Queued"), + mock_client.context_grounding.list_indexes.return_value = [ + _make_index(name="index-one"), + _make_index(name="index-two"), ] - result = runner.invoke( - cli, - ["context-grounding", "list", "--folder-path", "Shared"], + cli, ["context-grounding", "list", "--folder-path", "Shared"] ) - assert result.exit_code == 0 assert "index-one" in result.output assert "index-two" in result.output - # table columns projected - assert "last_ingestion_status" in result.output - assert "last_ingested" in result.output - # raw fields not in table - assert "data_source" not in result.output - mock_client.context_grounding.list.assert_called_once_with( - folder_path="Shared", - folder_key=None, - ) - - def test_list_json_format(self, runner, mock_client, mock_env_vars): - """list with --format json emits JSON.""" - mock_client.context_grounding.list.return_value = [ - _make_index(name="index-one"), - ] - - result = runner.invoke( - cli, - [ - "context-grounding", - "list", - "--folder-path", - "Shared", - "--format", - "json", - ], + mock_client.context_grounding.list_indexes.assert_called_once_with( + folder_path="Shared", folder_key=None ) - assert result.exit_code == 0 - assert "index-one" in result.output - def test_list_empty(self, runner, mock_client, mock_env_vars): - """list with no indexes returns empty output gracefully.""" - mock_client.context_grounding.list.return_value = [] - + mock_client.context_grounding.list_indexes.return_value = [] result = runner.invoke( - cli, - ["context-grounding", "list", "--folder-path", "Shared"], + cli, ["context-grounding", "list", "--folder-path", "Shared"] ) - assert result.exit_code == 0 def test_list_with_folder_key(self, runner, mock_client, mock_env_vars): - """list passes folder_key when --folder-key is provided.""" - mock_client.context_grounding.list.return_value = [] - folder_key = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" - - result = runner.invoke( - cli, - ["context-grounding", "list", "--folder-key", folder_key], - ) - + mock_client.context_grounding.list_indexes.return_value = [] + fk = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" + result = runner.invoke(cli, ["context-grounding", "list", "--folder-key", fk]) assert result.exit_code == 0 - mock_client.context_grounding.list.assert_called_once_with( - folder_path=None, - folder_key=folder_key, + mock_client.context_grounding.list_indexes.assert_called_once_with( + folder_path=None, folder_key=fk ) # --------------------------------------------------------------------------- +# retrieve +# --------------------------------------------------------------------------- class TestRetrieveCommand: - def test_retrieve_basic(self, runner, mock_client, mock_env_vars): - """retrieve returns index details.""" + def test_retrieve_by_name(self, runner, mock_client, mock_env_vars): mock_client.context_grounding.retrieve.return_value = _make_index() - result = runner.invoke( cli, [ "context-grounding", "retrieve", - "--index", + "--index-name", "my-index", "--folder-path", "Shared", ], ) - assert result.exit_code == 0 assert "my-index" in result.output mock_client.context_grounding.retrieve.assert_called_once_with( - name="my-index", - folder_path="Shared", - folder_key=None, + name="my-index", folder_path="Shared", folder_key=None ) - def test_retrieve_json_format(self, runner, mock_client, mock_env_vars): - """retrieve with --format json emits JSON.""" - mock_client.context_grounding.retrieve.return_value = _make_index() - + def test_retrieve_by_id(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve_by_id.return_value = { + "id": "abc-123", + "name": "ephemeral-index", + "lastIngestionStatus": "Successful", + } result = runner.invoke( cli, [ "context-grounding", "retrieve", - "--index", - "my-index", - "--folder-path", - "Shared", + "--index-id", + "abc-123", "--format", "json", ], ) - assert result.exit_code == 0 - assert "my-index" in result.output + mock_client.context_grounding.retrieve_by_id.assert_called_once() - def test_retrieve_missing_index_flag_fails(self, runner, mock_env_vars): - """retrieve without --index shows usage error.""" - result = runner.invoke( - cli, ["context-grounding", "retrieve", "--folder-path", "Shared"] + def test_retrieve_no_identifier_fails(self, runner, mock_env_vars): + result = runner.invoke(cli, ["context-grounding", "retrieve"]) + assert result.exit_code != 0 + assert ( + "index-name" in result.output.lower() or "index-id" in result.output.lower() ) + def test_retrieve_both_identifiers_fails(self, runner, mock_env_vars): + result = runner.invoke( + cli, + ["context-grounding", "retrieve", "--index-name", "X", "--index-id", "Y"], + ) assert result.exit_code != 0 - assert "index" in result.output.lower() or "missing" in result.output.lower() def test_retrieve_not_found(self, runner, mock_client, mock_env_vars): - """retrieve surfaces a clean not-found error (SDK bare Exception).""" - mock_client.context_grounding.retrieve.side_effect = Exception( - "ContextGroundingIndex not found" - ) - + mock_client.context_grounding.retrieve.side_effect = LookupError("not found") result = runner.invoke( cli, [ "context-grounding", "retrieve", - "--index", - "no-such-index", + "--index-name", + "missing", "--folder-path", "Shared", ], ) - assert result.exit_code != 0 assert "not found" in result.output.lower() def test_retrieve_not_found_http_404(self, runner, mock_client, mock_env_vars): - """retrieve surfaces a clean not-found error (HTTPStatusError 404).""" - from unittest.mock import MagicMock - from httpx import HTTPStatusError, Request, Response - mock_response = MagicMock(spec=Response) - mock_response.status_code = 404 + resp = Response(404, request=Request("GET", "http://test")) mock_client.context_grounding.retrieve.side_effect = HTTPStatusError( - "404", request=MagicMock(spec=Request), response=mock_response + "Not found", request=resp.request, response=resp ) - result = runner.invoke( cli, [ "context-grounding", "retrieve", - "--index", - "no-such-index", + "--index-name", + "missing", "--folder-path", "Shared", ], ) - assert result.exit_code != 0 assert "not found" in result.output.lower() - def test_retrieve_with_folder_key(self, runner, mock_client, mock_env_vars): - """retrieve passes folder_key when --folder-key is provided.""" - mock_client.context_grounding.retrieve.return_value = _make_index() - folder_key = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" +# --------------------------------------------------------------------------- +# create +# --------------------------------------------------------------------------- + + +class TestCreateCommand: + def test_create_with_bucket_source(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.create_index.return_value = _make_index( + name="new-index" + ) result = runner.invoke( cli, [ "context-grounding", - "retrieve", - "--index", - "my-index", - "--folder-key", - folder_key, + "create", + "--index-name", + "new-index", + "--bucket-source", + "my-bucket", + "--folder-path", + "Shared", ], ) - assert result.exit_code == 0 - mock_client.context_grounding.retrieve.assert_called_once_with( - name="my-index", - folder_path=None, - folder_key=folder_key, + mock_client.context_grounding.create_index.assert_called_once() + call_kwargs = mock_client.context_grounding.create_index.call_args[1] + assert call_kwargs["name"] == "new-index" + + def test_create_both_sources_fails(self, runner, mock_env_vars, tmp_path): + f = tmp_path / "config.json" + f.write_text("{}") + result = runner.invoke( + cli, + [ + "context-grounding", + "create", + "--index-name", + "X", + "--bucket-source", + "B", + "--source-file", + str(f), + ], + ) + assert result.exit_code != 0 + assert ( + "cannot use both" in result.output.lower() + or "choose one" in result.output.lower() ) + def test_create_no_source_fails(self, runner, mock_env_vars): + result = runner.invoke( + cli, + [ + "context-grounding", + "create", + "--index-name", + "X", + "--folder-path", + "Shared", + ], + ) + assert result.exit_code != 0 + + def test_create_missing_index_name_fails(self, runner, mock_env_vars): + result = runner.invoke( + cli, + ["context-grounding", "create", "--bucket-source", "B"], + ) + assert result.exit_code != 0 + # --------------------------------------------------------------------------- -# search +# create-ephemeral # --------------------------------------------------------------------------- -class TestSearchCommand: - def test_search_basic(self, runner, mock_client, mock_env_vars): - """search returns results table.""" - mock_client.context_grounding.search.return_value = [ - _make_result( - source="invoice.pdf", content="Pay within 30 days", score=0.92 - ), - _make_result(source="policy.pdf", content="Approval required", score=0.85), - ] +class TestCreateEphemeralCommand: + def test_create_ephemeral_basic(self, runner, mock_client, mock_env_vars, tmp_path): + test_file = tmp_path / "doc.pdf" + test_file.write_text("test content") + + import uuid + + mock_client.attachments.upload.return_value = uuid.uuid4() + mock_client.context_grounding.create_ephemeral_index.return_value = _make_index( + name="ephemeral" + ) result = runner.invoke( cli, [ "context-grounding", - "search", - "process an invoice", - "--index", - "my-index", - "--folder-path", - "Shared", + "create-ephemeral", + "--usage", + "DeepRAG", + "--files", + str(test_file), ], ) - assert result.exit_code == 0 - assert "invoice.pdf" in result.output - assert "policy.pdf" in result.output - mock_client.context_grounding.search.assert_called_once_with( - name="my-index", - query="process an invoice", - number_of_results=10, - threshold=None, - folder_path="Shared", - folder_key=None, - ) - - def test_search_with_limit_option(self, runner, mock_client, mock_env_vars): - """--limit N is forwarded to the SDK as number_of_results.""" - mock_client.context_grounding.search.return_value = [ - _make_result(content="Short answer"), - ] + mock_client.attachments.upload.assert_called_once() + mock_client.context_grounding.create_ephemeral_index.assert_called_once() + + def test_create_ephemeral_multiple_files( + self, runner, mock_client, mock_env_vars, tmp_path + ): + f1 = tmp_path / "a.csv" + f2 = tmp_path / "b.csv" + f1.write_text("a") + f2.write_text("b") + + import uuid + + mock_client.attachments.upload.return_value = uuid.uuid4() + mock_client.context_grounding.create_ephemeral_index.return_value = ( + _make_index() + ) result = runner.invoke( cli, [ "context-grounding", - "search", - "payment terms", - "--index", - "my-index", - "--folder-path", - "Shared", - "--limit", - "3", + "create-ephemeral", + "--usage", + "BatchRAG", + "--files", + str(f1), + "--files", + str(f2), ], ) + assert result.exit_code == 0 + assert mock_client.attachments.upload.call_count == 2 + + def test_create_ephemeral_missing_usage_fails( + self, runner, mock_env_vars, tmp_path + ): + f = tmp_path / "doc.pdf" + f.write_text("x") + result = runner.invoke( + cli, + ["context-grounding", "create-ephemeral", "--files", str(f)], + ) + assert result.exit_code != 0 + + +# --------------------------------------------------------------------------- +# source-schema +# --------------------------------------------------------------------------- + +class TestSourceSchemaCommand: + def test_source_schema_all(self, runner): + result = runner.invoke(cli, ["context-grounding", "source-schema"]) assert result.exit_code == 0 - mock_client.context_grounding.search.assert_called_once_with( - name="my-index", - query="payment terms", - number_of_results=3, - threshold=None, - folder_path="Shared", - folder_key=None, - ) - - def test_search_json_format_no_truncation(self, runner, mock_client, mock_env_vars): - """JSON output contains full content, not truncated.""" - long_content = "x" * 300 - mock_client.context_grounding.search.return_value = [ - _make_result(content=long_content), - ] + assert "google_drive" in result.output + assert "onedrive" in result.output + assert "dropbox" in result.output + assert "confluence" in result.output + assert "connection_id" in result.output + def test_source_schema_specific_type(self, runner): + result = runner.invoke( + cli, ["context-grounding", "source-schema", "--type", "confluence"] + ) + assert result.exit_code == 0 + assert "confluence" in result.output + assert "space_id" in result.output + assert "google_drive" not in result.output + + +# --------------------------------------------------------------------------- +# delete +# --------------------------------------------------------------------------- + + +class TestDeleteCommand: + def test_delete_with_confirm(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve.return_value = _make_index() result = runner.invoke( cli, [ "context-grounding", - "search", - "query", - "--index", + "delete", + "--index-name", "my-index", "--folder-path", "Shared", - "--format", - "json", + "--confirm", ], ) - assert result.exit_code == 0 - assert long_content in result.output - - def test_search_table_format_truncates_content( - self, runner, mock_client, mock_env_vars - ): - """Table output shows score/source/page/content only; content truncated at 120 chars.""" - long_content = "A" * 200 - mock_client.context_grounding.search.return_value = [ - _make_result(source="doc.pdf", page="3", content=long_content, score=0.92), - ] + mock_client.context_grounding.delete_by_name.assert_called_once_with( + name="my-index", folder_path="Shared", folder_key=None + ) + def test_delete_dry_run(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve.return_value = _make_index() result = runner.invoke( cli, [ "context-grounding", - "search", - "query", - "--index", + "delete", + "--index-name", "my-index", "--folder-path", "Shared", - "--format", - "table", + "--dry-run", ], ) - assert result.exit_code == 0 - # content is truncated - assert long_content not in result.output - assert "…" in result.output - # only human-readable columns rendered - assert "score" in result.output - assert "source" in result.output - assert "page_number" in result.output - assert "doc.pdf" in result.output - # raw fields not present in table - assert "metadata" not in result.output - assert "reference" not in result.output - - def test_search_empty_results(self, runner, mock_client, mock_env_vars): - """search with no results prints a helpful message and exits 0. - - Note: the message is emitted via click.echo(..., err=True) so it goes to - stderr. CliRunner mixes stderr into result.output by default, which is why - the assertion below works. If this runner were created with mix_stderr=False - the assertion would need to check result.stderr instead. - """ - mock_client.context_grounding.search.return_value = [] + assert "would delete" in result.output.lower() + mock_client.context_grounding.delete_by_name.assert_not_called() + def test_delete_prompts_no_cancels(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve.return_value = _make_index() result = runner.invoke( cli, [ "context-grounding", - "search", - "unknown query", - "--index", + "delete", + "--index-name", "my-index", "--folder-path", "Shared", ], + input="n\n", ) - assert result.exit_code == 0 - # message goes to stderr; CliRunner mixes stderr into output by default - assert "no results" in result.output.lower() - - def test_search_ingestion_in_progress_error( - self, runner, mock_client, mock_env_vars - ): - """search surfaces a clean error when index is being ingested.""" - mock_client.context_grounding.search.side_effect = IngestionInProgressException( - index_name="my-index" - ) + mock_client.context_grounding.delete_by_name.assert_not_called() + def test_delete_prompts_yes_deletes(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve.return_value = _make_index() result = runner.invoke( cli, [ "context-grounding", - "search", - "query", - "--index", + "delete", + "--index-name", "my-index", "--folder-path", "Shared", ], + input="y\n", ) + assert result.exit_code == 0 + mock_client.context_grounding.delete_by_name.assert_called_once() - assert result.exit_code != 0 - assert ( - "ingested" in result.output.lower() or "ingestion" in result.output.lower() - ) - - def test_search_not_found(self, runner, mock_client, mock_env_vars): - """search surfaces a clean not-found error for missing index.""" - mock_client.context_grounding.search.side_effect = Exception( - "ContextGroundingIndex not found" - ) - + def test_delete_not_found(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.retrieve.side_effect = LookupError("not found") result = runner.invoke( cli, [ "context-grounding", - "search", - "query", - "--index", - "no-such-index", + "delete", + "--index-name", + "missing", "--folder-path", "Shared", + "--confirm", ], ) - assert result.exit_code != 0 assert "not found" in result.output.lower() - def test_search_missing_index_flag_fails(self, runner, mock_env_vars): - """search without --index shows usage error.""" - result = runner.invoke( - cli, ["context-grounding", "search", "query", "--folder-path", "Shared"] - ) - + def test_delete_missing_index_name_fails(self, runner, mock_env_vars): + result = runner.invoke(cli, ["context-grounding", "delete", "--confirm"]) assert result.exit_code != 0 - assert "index" in result.output.lower() or "missing" in result.output.lower() - def test_search_missing_query_fails(self, runner, mock_env_vars): - """search without QUERY positional arg shows usage error.""" + +# --------------------------------------------------------------------------- +# ingest +# --------------------------------------------------------------------------- + + +class TestIngestCommand: + def test_ingest_basic(self, runner, mock_client, mock_env_vars): result = runner.invoke( cli, [ "context-grounding", - "search", - "--index", + "ingest", + "--index-name", "my-index", "--folder-path", "Shared", ], ) + assert result.exit_code == 0 + assert "ingestion triggered" in result.output.lower() + mock_client.context_grounding.ingest_by_name.assert_called_once_with( + name="my-index", folder_path="Shared", folder_key=None + ) + def test_ingest_missing_index_name_fails(self, runner, mock_env_vars): + result = runner.invoke(cli, ["context-grounding", "ingest"]) assert result.exit_code != 0 - def test_search_with_folder_key(self, runner, mock_client, mock_env_vars): - """search passes folder_key when --folder-key is provided.""" - mock_client.context_grounding.search.return_value = [_make_result()] - folder_key = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" - + def test_ingest_with_folder_key(self, runner, mock_client, mock_env_vars): + fk = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" result = runner.invoke( cli, [ "context-grounding", - "search", - "query", - "--index", + "ingest", + "--index-name", "my-index", "--folder-key", - folder_key, + fk, ], ) - assert result.exit_code == 0 - mock_client.context_grounding.search.assert_called_once_with( - name="my-index", - query="query", - number_of_results=10, - threshold=None, - folder_path=None, - folder_key=folder_key, + mock_client.context_grounding.ingest_by_name.assert_called_once_with( + name="my-index", folder_path=None, folder_key=fk ) # --------------------------------------------------------------------------- -# ingest +# search # --------------------------------------------------------------------------- -class TestIngestCommand: - def test_ingest_basic(self, runner, mock_client, mock_env_vars): - """ingest triggers ingestion and prints confirmation.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.ingest_data.return_value = None - +class TestSearchCommand: + def test_search_by_name(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.unified_search.return_value = [] result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", + "search", + "--index-name", "my-index", "--folder-path", "Shared", + "--query", + "revenue", ], ) - assert result.exit_code == 0 - assert "my-index" in result.output - mock_client.context_grounding.ingest_data.assert_called_once_with( - index=index, - folder_path="Shared", - folder_key=None, - ) - - def test_ingest_already_in_progress_fast_fail( - self, runner, mock_client, mock_env_vars - ): - """ingest fails fast (no HTTP call) when retrieve shows ingestion in progress.""" - index = _make_index(status="In Progress") - index.in_progress_ingestion.return_value = True - mock_client.context_grounding.retrieve.return_value = index + mock_client.context_grounding.unified_search.assert_called_once() + call_kwargs = mock_client.context_grounding.unified_search.call_args[1] + assert call_kwargs["name"] == "my-index" + assert call_kwargs["query"] == "revenue" + def test_search_with_limit(self, runner, mock_client, mock_env_vars): + mock_client.context_grounding.unified_search.return_value = [] result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", - "my-index", + "search", + "--index-name", + "X", "--folder-path", "Shared", + "--query", + "test", + "--limit", + "3", ], ) + assert result.exit_code == 0 + call_kwargs = mock_client.context_grounding.unified_search.call_args[1] + assert call_kwargs["number_of_results"] == 3 - assert result.exit_code != 0 - assert "already" in result.output.lower() or "ingested" in result.output.lower() - mock_client.context_grounding.ingest_data.assert_not_called() - - def test_ingest_already_in_progress(self, runner, mock_client, mock_env_vars): - """ingest surfaces a clean error when the API reports 409 (race condition).""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.ingest_data.side_effect = ( - IngestionInProgressException(index_name="my-index", search_operation=False) - ) - + def test_search_missing_query_fails(self, runner, mock_env_vars): result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", - "my-index", + "search", + "--index-name", + "X", "--folder-path", "Shared", ], ) - assert result.exit_code != 0 - assert "already" in result.output.lower() or "ingested" in result.output.lower() - def test_ingest_not_found(self, runner, mock_client, mock_env_vars): - """ingest surfaces a clean not-found error.""" - mock_client.context_grounding.retrieve.side_effect = Exception( - "ContextGroundingIndex not found" - ) + +# --------------------------------------------------------------------------- +# deep-rag start +# --------------------------------------------------------------------------- + + +class TestDeepRagStartCommand: + def test_start_by_name(self, runner, mock_client, mock_env_vars): + mock_result = MagicMock() + mock_result.id = "task-123" + mock_result.model_dump.return_value = {"id": "task-123", "status": "Queued"} + mock_client.context_grounding.start_deep_rag.return_value = mock_result result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", - "no-such-index", + "deep-rag", + "start", + "--index-name", + "my-index", "--folder-path", "Shared", + "--task-name", + "my-task", + "--prompt", + "Summarize", ], ) + assert result.exit_code == 0 + call_kwargs = mock_client.context_grounding.start_deep_rag.call_args[1] + assert call_kwargs["name"] == "my-task" + assert call_kwargs["index_name"] == "my-index" + assert call_kwargs["index_id"] is None - assert result.exit_code != 0 - assert "not found" in result.output.lower() - - def test_ingest_index_has_no_id(self, runner, mock_client, mock_env_vars): - """ingest raises a clean error if the retrieved index has no ID (avoids silent no-op).""" - index = _make_index() - index.id = None - mock_client.context_grounding.retrieve.return_value = index + def test_start_by_id(self, runner, mock_client, mock_env_vars): + mock_result = MagicMock() + mock_result.id = "task-456" + mock_result.model_dump.return_value = {"id": "task-456"} + mock_client.context_grounding.start_deep_rag.return_value = mock_result result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", - "my-index", - "--folder-path", - "Shared", + "deep-rag", + "start", + "--index-id", + "idx-789", + "--task-name", + "my-task", + "--prompt", + "Summarize", ], ) + assert result.exit_code == 0 + call_kwargs = mock_client.context_grounding.start_deep_rag.call_args[1] + assert call_kwargs["index_id"] == "idx-789" + assert call_kwargs["index_name"] is None - assert result.exit_code != 0 - assert "no id" in result.output.lower() or "cannot" in result.output.lower() - mock_client.context_grounding.ingest_data.assert_not_called() - - def test_ingest_with_folder_key(self, runner, mock_client, mock_env_vars): - """ingest passes folder_key when --folder-key is provided.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.ingest_data.return_value = None - folder_key = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" - + def test_start_no_index_fails(self, runner, mock_env_vars): result = runner.invoke( cli, [ "context-grounding", - "ingest", - "--index", - "my-index", - "--folder-key", - folder_key, + "deep-rag", + "start", + "--task-name", + "T", + "--prompt", + "P", ], ) + assert result.exit_code != 0 - assert result.exit_code == 0 - mock_client.context_grounding.retrieve.assert_called_once_with( - name="my-index", - folder_path=None, - folder_key=folder_key, - ) - mock_client.context_grounding.ingest_data.assert_called_once_with( - index=index, - folder_path=None, - folder_key=folder_key, - ) - - def test_ingest_missing_index_flag_fails(self, runner, mock_env_vars): - """ingest without --index shows usage error.""" + def test_start_both_index_fails(self, runner, mock_env_vars): result = runner.invoke( - cli, ["context-grounding", "ingest", "--folder-path", "Shared"] + cli, + [ + "context-grounding", + "deep-rag", + "start", + "--index-name", + "X", + "--index-id", + "Y", + "--task-name", + "T", + "--prompt", + "P", + ], ) - assert result.exit_code != 0 # --------------------------------------------------------------------------- -# delete +# deep-rag retrieve # --------------------------------------------------------------------------- -class TestDeleteCommand: - def test_delete_with_confirm(self, runner, mock_client, mock_env_vars): - """delete --confirm removes the index without prompting.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.delete_index.return_value = None +class TestDeepRagRetrieveCommand: + def test_retrieve(self, runner, mock_client, mock_env_vars): + mock_result = MagicMock() + mock_result.model_dump.return_value = {"id": "task-123", "status": "Successful"} + mock_client.context_grounding.retrieve_deep_rag.return_value = mock_result result = runner.invoke( cli, - [ - "context-grounding", - "delete", - "--index", - "my-index", - "--folder-path", - "Shared", - "--confirm", - ], + ["context-grounding", "deep-rag", "retrieve", "--task-id", "task-123"], ) - assert result.exit_code == 0 - assert "my-index" in result.output - mock_client.context_grounding.delete_index.assert_called_once_with( - index=index, - folder_path="Shared", - folder_key=None, + mock_client.context_grounding.retrieve_deep_rag.assert_called_once_with( + id="task-123" ) - def test_delete_dry_run(self, runner, mock_client, mock_env_vars): - """delete --dry-run prints what would be deleted without deleting.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index + def test_retrieve_missing_task_id_fails(self, runner, mock_env_vars): + result = runner.invoke(cli, ["context-grounding", "deep-rag", "retrieve"]) + assert result.exit_code != 0 - result = runner.invoke( - cli, - [ - "context-grounding", - "delete", - "--index", - "my-index", - "--folder-path", - "Shared", - "--dry-run", - ], - ) - assert result.exit_code == 0 - assert "would delete" in result.output.lower() - mock_client.context_grounding.delete_index.assert_not_called() +# --------------------------------------------------------------------------- +# batch-transform start +# --------------------------------------------------------------------------- + - def test_delete_prompts_without_confirm(self, runner, mock_client, mock_env_vars): - """delete without --confirm or --dry-run prompts the user; 'n' cancels.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index +class TestBatchTransformStartCommand: + def test_start_by_name(self, runner, mock_client, mock_env_vars, tmp_path): + cols_file = tmp_path / "cols.json" + cols_file.write_text( + json.dumps([{"name": "entity", "description": "Entity name"}]) + ) + + mock_result = MagicMock() + mock_result.id = "bt-123" + mock_result.model_dump.return_value = {"id": "bt-123"} + mock_client.context_grounding.start_batch_transform.return_value = mock_result result = runner.invoke( cli, [ "context-grounding", - "delete", - "--index", + "batch-transform", + "start", + "--index-name", "my-index", "--folder-path", "Shared", + "--task-name", + "my-task", + "--prompt", + "Extract entities", + "--columns-file", + str(cols_file), ], - input="n\n", ) - assert result.exit_code == 0 - assert "cancelled" in result.output.lower() - mock_client.context_grounding.delete_index.assert_not_called() + call_kwargs = mock_client.context_grounding.start_batch_transform.call_args[1] + assert call_kwargs["name"] == "my-task" + assert call_kwargs["index_name"] == "my-index" - def test_delete_prompts_yes_deletes(self, runner, mock_client, mock_env_vars): - """delete without --confirm but answering 'y' at the prompt deletes the index.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.delete_index.return_value = None + def test_start_by_id(self, runner, mock_client, mock_env_vars, tmp_path): + cols_file = tmp_path / "cols.json" + cols_file.write_text(json.dumps([{"name": "col1", "description": "Col"}])) + + mock_result = MagicMock() + mock_result.id = "bt-456" + mock_result.model_dump.return_value = {"id": "bt-456"} + mock_client.context_grounding.start_batch_transform.return_value = mock_result result = runner.invoke( cli, [ "context-grounding", - "delete", - "--index", - "my-index", - "--folder-path", - "Shared", + "batch-transform", + "start", + "--index-id", + "idx-789", + "--task-name", + "my-task", + "--prompt", + "Extract", + "--columns-file", + str(cols_file), ], - input="y\n", ) - assert result.exit_code == 0 - mock_client.context_grounding.delete_index.assert_called_once() - - def test_delete_with_folder_key(self, runner, mock_client, mock_env_vars): - """delete passes folder_key when --folder-key is provided.""" - index = _make_index() - mock_client.context_grounding.retrieve.return_value = index - mock_client.context_grounding.delete_index.return_value = None - folder_key = "a1b2c3d4-e5f6-7890-abcd-ef1234567890" + call_kwargs = mock_client.context_grounding.start_batch_transform.call_args[1] + assert call_kwargs["index_id"] == "idx-789" + assert call_kwargs["index_name"] is None + def test_start_no_index_fails(self, runner, mock_env_vars, tmp_path): + cols_file = tmp_path / "cols.json" + cols_file.write_text("[]") result = runner.invoke( cli, [ "context-grounding", - "delete", - "--index", - "my-index", - "--folder-key", - folder_key, - "--confirm", + "batch-transform", + "start", + "--task-name", + "T", + "--prompt", + "P", + "--columns-file", + str(cols_file), ], ) + assert result.exit_code != 0 - assert result.exit_code == 0 - mock_client.context_grounding.retrieve.assert_called_once_with( - name="my-index", - folder_path=None, - folder_key=folder_key, - ) - mock_client.context_grounding.delete_index.assert_called_once_with( - index=index, - folder_path=None, - folder_key=folder_key, - ) - def test_delete_index_has_no_id(self, runner, mock_client, mock_env_vars): - """delete raises a clean error if the retrieved index has no ID (avoids silent no-op).""" - index = _make_index() - index.id = None - mock_client.context_grounding.retrieve.return_value = index +# --------------------------------------------------------------------------- +# batch-transform retrieve +# --------------------------------------------------------------------------- + + +class TestBatchTransformRetrieveCommand: + def test_retrieve(self, runner, mock_client, mock_env_vars): + mock_result = MagicMock() + mock_result.model_dump.return_value = {"id": "bt-123", "status": "Successful"} + mock_client.context_grounding.retrieve_batch_transform.return_value = ( + mock_result + ) result = runner.invoke( cli, - [ - "context-grounding", - "delete", - "--index", - "my-index", - "--folder-path", - "Shared", - "--confirm", - ], + ["context-grounding", "batch-transform", "retrieve", "--task-id", "bt-123"], + ) + assert result.exit_code == 0 + mock_client.context_grounding.retrieve_batch_transform.assert_called_once_with( + id="bt-123" ) - assert result.exit_code != 0 - assert "no id" in result.output.lower() or "cannot" in result.output.lower() - mock_client.context_grounding.delete_index.assert_not_called() - def test_delete_not_found(self, runner, mock_client, mock_env_vars): - """delete surfaces a clean not-found error.""" - mock_client.context_grounding.retrieve.side_effect = Exception( - "ContextGroundingIndex not found" - ) +# --------------------------------------------------------------------------- +# batch-transform download +# --------------------------------------------------------------------------- + +class TestBatchTransformDownloadCommand: + def test_download(self, runner, mock_client, mock_env_vars, tmp_path): + out_file = str(tmp_path / "result.csv") result = runner.invoke( cli, [ "context-grounding", - "delete", - "--index", - "no-such-index", - "--folder-path", - "Shared", - "--confirm", + "batch-transform", + "download", + "--task-id", + "bt-123", + "--output-file", + out_file, ], ) + assert result.exit_code == 0 + mock_client.context_grounding.download_batch_transform_result.assert_called_once_with( + id="bt-123", destination_path=out_file + ) + assert "downloaded" in result.output.lower() - assert result.exit_code != 0 - assert "not found" in result.output.lower() - - def test_delete_missing_index_flag_fails(self, runner, mock_env_vars): - """delete without --index shows usage error.""" + def test_download_missing_output_file_fails(self, runner, mock_env_vars): result = runner.invoke( - cli, ["context-grounding", "delete", "--folder-path", "Shared", "--confirm"] + cli, + ["context-grounding", "batch-transform", "download", "--task-id", "bt-123"], ) - assert result.exit_code != 0 @@ -903,38 +839,45 @@ def test_delete_missing_index_flag_fails(self, runner, mock_env_vars): class TestHelpText: def test_group_help(self, runner): - """context-grounding group has correct help text.""" result = runner.invoke(cli, ["context-grounding", "--help"]) - - assert result.exit_code == 0 - assert "retrieve" in result.output - assert "search" in result.output - assert "ingest" in result.output - assert "delete" in result.output - - def test_search_help(self, runner): - """search command exposes all expected options.""" - result = runner.invoke(cli, ["context-grounding", "search", "--help"]) - assert result.exit_code == 0 - assert "--index" in result.output - assert "--limit" in result.output - assert "--folder-path" in result.output - assert "--folder-key" in result.output - assert "--format" in result.output + assert "Two index types" in result.output + assert "Regular" in result.output + assert "Ephemeral" in result.output + for cmd in [ + "list", + "retrieve", + "create", + "create-ephemeral", + "delete", + "ingest", + "search", + "source-schema", + "deep-rag", + "batch-transform", + ]: + assert cmd in result.output def test_retrieve_help(self, runner): - """retrieve command exposes --index and folder options.""" result = runner.invoke(cli, ["context-grounding", "retrieve", "--help"]) - assert result.exit_code == 0 - assert "--index" in result.output - assert "--folder-path" in result.output + assert "--index-name" in result.output + assert "--index-id" in result.output - def test_delete_help(self, runner): - """delete command exposes --confirm and --dry-run.""" - result = runner.invoke(cli, ["context-grounding", "delete", "--help"]) + def test_deep_rag_start_help(self, runner): + result = runner.invoke( + cli, ["context-grounding", "deep-rag", "start", "--help"] + ) + assert result.exit_code == 0 + assert "--index-name" in result.output + assert "--index-id" in result.output + assert "--task-name" in result.output + assert "--prompt" in result.output + def test_batch_transform_start_help(self, runner): + result = runner.invoke( + cli, ["context-grounding", "batch-transform", "start", "--help"] + ) assert result.exit_code == 0 - assert "--confirm" in result.output - assert "--dry-run" in result.output + assert "--columns-file" in result.output + assert "--task-name" in result.output diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 318e1bd25..3214375f9 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2543,7 +2543,7 @@ wheels = [ [[package]] name = "uipath" -version = "2.10.38" +version = "2.10.39" source = { editable = "." } dependencies = [ { name = "applicationinsights" }, @@ -2682,7 +2682,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.1.15" +version = "0.1.16" source = { editable = "../uipath-platform" } dependencies = [ { name = "httpx" }, From 6cdfa5704aa219f00f867fb7465145cece65acd2 Mon Sep 17 00:00:00 2001 From: Jay Patel <21jaypatel@gmail.com> Date: Tue, 31 Mar 2026 11:59:43 -0400 Subject: [PATCH 2/3] fix(context-grounding): address PR review feedback Remove duplicate decorators on ingest_by_name, use safe .get("value", []) in list_indexes/list_indexes_async, and drop redundant folder_key re-resolution in _list_spec. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../context_grounding/_context_grounding_service.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py b/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py index 627143429..75e525d3b 100644 --- a/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -1723,7 +1723,8 @@ def list_indexes( ).json() return [ - ContextGroundingIndex.model_validate(item) for item in response["value"] + ContextGroundingIndex.model_validate(item) + for item in response.get("value", []) ] @resource_override(resource_type="index") @@ -1761,7 +1762,8 @@ async def list_indexes_async( ).json() return [ - ContextGroundingIndex.model_validate(item) for item in response["value"] + ContextGroundingIndex.model_validate(item) + for item in response.get("value", []) ] @resource_override(resource_type="index") @@ -1814,8 +1816,6 @@ async def delete_by_name_async( index, folder_key=folder_key, folder_path=folder_path ) - @resource_override(resource_type="index") - @traced(name="contextgrounding_exists", run_type="uipath") @resource_override(resource_type="index") @traced(name="contextgrounding_ingest_by_name", run_type="uipath") def ingest_by_name( @@ -1903,10 +1903,7 @@ def _retrieve_across_folders_spec( def _list_spec( self, folder_key: Optional[str] = None, - folder_path: Optional[str] = None, ) -> RequestSpec: - folder_key = self._resolve_folder_key(folder_key, folder_path) - return RequestSpec( method="GET", endpoint=Endpoint("/ecs_/v2/indexes"), From 66ef66bacf9b2b0e544699452f502629c7aa36b6 Mon Sep 17 00:00:00 2001 From: Jay Patel <21jaypatel@gmail.com> Date: Tue, 31 Mar 2026 13:10:40 -0400 Subject: [PATCH 3/3] fix(context-grounding): remove unused import and use model_validate Remove unused BucketSourceConfig import and use model_validate() instead of direct constructor for robust nested model parsing. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../uipath/src/uipath/_cli/services/cli_context_grounding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py b/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py index 53e225c02..075ffe869 100644 --- a/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py +++ b/packages/uipath/src/uipath/_cli/services/cli_context_grounding.py @@ -103,7 +103,7 @@ def retrieve_index( folder_path=folder_path, folder_key=folder_key, ) - return ContextGroundingIndex(**raw) + return ContextGroundingIndex.model_validate(raw) assert index_name is not None # validated above name = index_name @@ -183,7 +183,7 @@ def create_index( if source_file: from pydantic import TypeAdapter - from uipath.platform.context_grounding import BucketSourceConfig, SourceConfig + from uipath.platform.context_grounding import SourceConfig with open(source_file) as f: source_data = json_module.load(f)