Skip to content

Commit 889500d

Browse files
authored
Merge pull request #986 from UiPath/feat/deep-rag
feat: add deep rag support
2 parents d3732f1 + 4dcb44f commit 889500d

File tree

6 files changed

+589
-7
lines changed

6 files changed

+589
-7
lines changed

pyproject.toml

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath"
3-
version = "2.2.25"
3+
version = "2.2.26"
44
description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"
@@ -127,9 +127,9 @@ init_typed = true
127127
warn_required_dynamic_aliases = true
128128

129129
[tool.pytest.ini_options]
130-
#testpaths = ["tests"]
131-
#python_files = "test_*.py"
132-
#addopts = "-ra -q --cov"
130+
testpaths = ["tests"]
131+
python_files = "test_*.py"
132+
addopts = "-ra -q --cov"
133133
asyncio_default_fixture_loop_scope = "function"
134134
asyncio_mode = "auto"
135135

src/uipath/platform/context_grounding/__init__.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,13 @@
11
"""Init file for context grounding module."""
22

33
from ._context_grounding_service import ContextGroundingService
4-
from .context_grounding import ContextGroundingQueryResponse
4+
from .context_grounding import (
5+
CitationMode,
6+
ContextGroundingQueryResponse,
7+
DeepRagCreationResponse,
8+
DeepRagResponse,
9+
DeepRagStatus,
10+
)
511
from .context_grounding_index import ContextGroundingIndex
612
from .context_grounding_payloads import (
713
BaseSourceConfig,
@@ -42,4 +48,8 @@
4248
"OneDriveSourceConfig",
4349
"PreProcessing",
4450
"SourceConfig",
51+
"CitationMode",
52+
"DeepRagCreationResponse",
53+
"DeepRagStatus",
54+
"DeepRagResponse",
4555
]

src/uipath/platform/context_grounding/_context_grounding_service.py

Lines changed: 213 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,12 @@
1717
)
1818
from ..orchestrator._buckets_service import BucketsService
1919
from ..orchestrator._folder_service import FolderService
20-
from .context_grounding import ContextGroundingQueryResponse
20+
from .context_grounding import (
21+
CitationMode,
22+
ContextGroundingQueryResponse,
23+
DeepRagCreationResponse,
24+
DeepRagResponse,
25+
)
2126
from .context_grounding_index import ContextGroundingIndex
2227
from .context_grounding_payloads import (
2328
BucketDataSource,
@@ -438,6 +443,171 @@ async def create_index_async(
438443

439444
return ContextGroundingIndex.model_validate(response.json())
440445

446+
@traced(name="contextgrounding_retrieve_deep_rag", run_type="uipath")
447+
@resource_override(resource_type="index", resource_identifier="index_name")
448+
def retrieve_deep_rag(
449+
self,
450+
id: str,
451+
*,
452+
index_name: str | None = None,
453+
) -> DeepRagResponse:
454+
"""Retrieves a Deep RAG task.
455+
456+
Args:
457+
id (str): The id of the Deep RAG task.
458+
index_name (Optional[str]): Index name hint for resource override.
459+
460+
Returns:
461+
DeepRagResponse: The Deep RAG task response.
462+
"""
463+
spec = self._deep_rag_retrieve_spec(
464+
id=id,
465+
)
466+
response = self.request(
467+
spec.method,
468+
spec.endpoint,
469+
params=spec.params,
470+
json=spec.json,
471+
headers=spec.headers,
472+
)
473+
return DeepRagResponse.model_validate(response.json())
474+
475+
@traced(name="contextgrounding_retrieve_deep_rag_async", run_type="uipath")
476+
@resource_override(resource_type="index", resource_identifier="index_name")
477+
async def retrieve_deep_rag_async(
478+
self,
479+
id: str,
480+
*,
481+
index_name: str | None = None,
482+
) -> DeepRagResponse:
483+
"""Asynchronously retrieves a Deep RAG task.
484+
485+
Args:
486+
id (str): The id of the Deep RAG task.
487+
index_name (Optional[str]): Index name hint for resource override.
488+
489+
Returns:
490+
DeepRagResponse: The Deep RAG task response.
491+
"""
492+
spec = self._deep_rag_retrieve_spec(
493+
id=id,
494+
)
495+
496+
response = await self.request_async(
497+
spec.method,
498+
spec.endpoint,
499+
params=spec.params,
500+
json=spec.json,
501+
headers=spec.headers,
502+
)
503+
504+
return DeepRagResponse.model_validate(response.json())
505+
506+
@traced(name="contextgrounding_start_deep_rag", run_type="uipath")
507+
@resource_override(resource_type="index", resource_identifier="index_name")
508+
def start_deep_rag(
509+
self,
510+
name: str,
511+
index_name: str,
512+
prompt: str,
513+
glob_pattern: str = "*",
514+
citation_mode: CitationMode = CitationMode.SKIP,
515+
folder_key: str | None = None,
516+
folder_path: str | None = None,
517+
) -> DeepRagCreationResponse:
518+
"""Starts a Deep RAG task on the targeted index.
519+
520+
Args:
521+
name (str): The name of the Deep RAG task.
522+
index_name (str): The name of the context index to search in.
523+
prompt (str): Describe the task: what to research across documents, what to synthesize and how to cite sources.
524+
glob_pattern (str): The glob pattern to search in the index. Defaults to "*".
525+
citation_mode (CitationMode): The citation mode to use. Defaults to SKIP.
526+
folder_key (str, optional): The folder key where the index resides. Defaults to None.
527+
folder_path (str, optional): The folder path where the index resides. Defaults to None.
528+
529+
Returns:
530+
DeepRagCreationResponse: The Deep RAG task creation response.
531+
"""
532+
index = self.retrieve(
533+
index_name, folder_key=folder_key, folder_path=folder_path
534+
)
535+
if index and index.in_progress_ingestion():
536+
raise IngestionInProgressException(index_name=index_name)
537+
538+
spec = self._deep_rag_creation_spec(
539+
index_id=index.id,
540+
name=name,
541+
glob_pattern=glob_pattern,
542+
prompt=prompt,
543+
citation_mode=citation_mode,
544+
folder_key=folder_key,
545+
folder_path=folder_path,
546+
)
547+
548+
response = self.request(
549+
spec.method,
550+
spec.endpoint,
551+
json=spec.json,
552+
params=spec.params,
553+
headers=spec.headers,
554+
)
555+
556+
return DeepRagCreationResponse.model_validate(response.json())
557+
558+
@traced(name="contextgrounding_start_deep_rag_async", run_type="uipath")
559+
@resource_override(resource_type="index", resource_identifier="index_name")
560+
async def start_deep_rag_async(
561+
self,
562+
name: str,
563+
index_name: str,
564+
prompt: str,
565+
glob_pattern: str = "*",
566+
citation_mode: CitationMode = CitationMode.SKIP,
567+
folder_key: str | None = None,
568+
folder_path: str | None = None,
569+
) -> DeepRagCreationResponse:
570+
"""Asynchronously starts a Deep RAG task on the targeted index.
571+
572+
Args:
573+
name (str): The name of the Deep RAG task.
574+
index_name (str): The name of the context index to search in.
575+
name (str): The name of the Deep RAG task.
576+
prompt (str): Describe the task: what to research across documents, what to synthesize and how to cite sources.
577+
glob_pattern (str): The glob pattern to search in the index. Defaults to "*".
578+
citation_mode (CitationMode): The citation mode to use. Defaults to SKIP.
579+
folder_key (str, optional): The folder key where the index resides. Defaults to None.
580+
folder_path (str, optional): The folder path where the index resides. Defaults to None.
581+
582+
Returns:
583+
DeepRagCreationResponse: The Deep RAG task creation response.
584+
"""
585+
index = await self.retrieve_async(
586+
index_name, folder_key=folder_key, folder_path=folder_path
587+
)
588+
if index and index.in_progress_ingestion():
589+
raise IngestionInProgressException(index_name=index_name)
590+
591+
spec = self._deep_rag_creation_spec(
592+
index_id=index.id,
593+
name=name,
594+
glob_pattern=glob_pattern,
595+
prompt=prompt,
596+
citation_mode=citation_mode,
597+
folder_key=folder_key,
598+
folder_path=folder_path,
599+
)
600+
601+
response = await self.request_async(
602+
spec.method,
603+
spec.endpoint,
604+
params=spec.params,
605+
json=spec.json,
606+
headers=spec.headers,
607+
)
608+
609+
return DeepRagCreationResponse.model_validate(response.json())
610+
441611
@traced(name="contextgrounding_search", run_type="uipath")
442612
@resource_override(resource_type="index")
443613
def search(
@@ -875,6 +1045,48 @@ def _search_spec(
8751045
},
8761046
)
8771047

1048+
def _deep_rag_creation_spec(
1049+
self,
1050+
index_id: str,
1051+
name: str,
1052+
glob_pattern: str,
1053+
prompt: str,
1054+
citation_mode: CitationMode,
1055+
folder_key: str | None = None,
1056+
folder_path: str | None = None,
1057+
) -> RequestSpec:
1058+
folder_key = self._resolve_folder_key(folder_key, folder_path)
1059+
1060+
return RequestSpec(
1061+
method="POST",
1062+
endpoint=Endpoint(f"/ecs_/v2/indexes/{index_id}/createDeepRag"),
1063+
json={
1064+
"name": name,
1065+
"prompt": prompt,
1066+
"globPattern": glob_pattern,
1067+
"citationMode": citation_mode.value,
1068+
},
1069+
params={
1070+
"$select": "id,lastDeepRagStatus,createdDate",
1071+
},
1072+
headers={
1073+
**header_folder(folder_key, None),
1074+
},
1075+
)
1076+
1077+
def _deep_rag_retrieve_spec(
1078+
self,
1079+
id: str,
1080+
) -> RequestSpec:
1081+
return RequestSpec(
1082+
method="GET",
1083+
endpoint=Endpoint(f"/ecs_/v2/deeprag/{id}"),
1084+
params={
1085+
"$expand": "content",
1086+
"$select": "content,name,createdDate,lastDeepRagStatus",
1087+
},
1088+
)
1089+
8781090
def _resolve_folder_key(self, folder_key, folder_path):
8791091
if folder_key is None and folder_path is not None:
8801092
folder_key = self._folders_service.retrieve_key(folder_path=folder_path)

src/uipath/platform/context_grounding/context_grounding.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,73 @@
11
"""Context Grounding response payload models."""
22

3+
from enum import Enum
34
from typing import Optional
45

56
from pydantic import BaseModel, ConfigDict, Field
67

78

9+
class CitationMode(str, Enum):
10+
"""Enum representing possible citation modes."""
11+
12+
SKIP = "Skip"
13+
INLINE = "Inline"
14+
15+
16+
class DeepRagStatus(str, Enum):
17+
"""Enum representing possible deep RAG tasks status."""
18+
19+
QUEUED = "Queued"
20+
IN_PROGRESS = "InProgress"
21+
SUCCESSFUL = "Successful"
22+
FAILED = "Failed"
23+
24+
25+
class DeepRagContent(BaseModel):
26+
"""Model representing a deep RAG task content."""
27+
28+
text: str
29+
citations: list[str]
30+
31+
model_config = ConfigDict(
32+
validate_by_name=True,
33+
validate_by_alias=True,
34+
use_enum_values=True,
35+
arbitrary_types_allowed=True,
36+
extra="allow",
37+
)
38+
39+
40+
class DeepRagResponse(BaseModel):
41+
"""Model representing a deep RAG task response."""
42+
43+
model_config = ConfigDict(
44+
validate_by_name=True,
45+
validate_by_alias=True,
46+
use_enum_values=True,
47+
arbitrary_types_allowed=True,
48+
extra="allow",
49+
)
50+
name: str
51+
created_date: str = Field(alias="createdDate")
52+
last_deep_rag_status: DeepRagStatus = Field(alias="lastDeepRagStatus")
53+
content: DeepRagContent | None = Field(alias="content")
54+
55+
56+
class DeepRagCreationResponse(BaseModel):
57+
"""Model representing a deep RAG task creation response."""
58+
59+
model_config = ConfigDict(
60+
validate_by_name=True,
61+
validate_by_alias=True,
62+
use_enum_values=True,
63+
arbitrary_types_allowed=True,
64+
extra="allow",
65+
)
66+
id: str
67+
last_deep_rag_status: DeepRagStatus = Field(alias="lastDeepRagStatus")
68+
created_date: str = Field(alias="createdDate")
69+
70+
871
class ContextGroundingMetadata(BaseModel):
972
"""Model representing metadata for a Context Grounding query response."""
1073

0 commit comments

Comments
 (0)