diff --git a/pyproject.toml b/pyproject.toml index 886c9b2a6..3ad3ce874 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ dev = [ "inflection>=0.5.1", "types-toml>=0.10.8", "pytest-timeout>=2.4.0", + "uipath", ] [tool.hatch.build.targets.wheel] @@ -148,3 +149,6 @@ name = "testpypi" url = "https://test.pypi.org/simple/" publish-url = "https://test.pypi.org/legacy/" explicit = true + +[tool.uv.sources] +uipath = { workspace = true } diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 2941577d3..b5336569f 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -85,6 +85,7 @@ class AgentContextRetrievalMode(str, Enum): STRUCTURED = "Structured" DEEP_RAG = "DeepRAG" BATCH_TRANSFORM = "BatchTransform" + DATA_FABRIC = "DataFabric" UNKNOWN = "Unknown" # fallback branch discriminator @@ -317,6 +318,7 @@ class AgentContextSettings(BaseCfg): AgentContextRetrievalMode.STRUCTURED, AgentContextRetrievalMode.DEEP_RAG, AgentContextRetrievalMode.BATCH_TRANSFORM, + AgentContextRetrievalMode.DATA_FABRIC, AgentContextRetrievalMode.UNKNOWN, ] = Field(alias="retrievalMode") threshold: float = Field(default=0) @@ -336,6 +338,10 @@ class AgentContextSettings(BaseCfg): output_columns: Optional[List[AgentContextOutputColumn]] = Field( None, alias="outputColumns" ) + # Data Fabric specific settings + entity_identifiers: Optional[List[str]] = Field( + None, alias="entityIdentifiers" + ) class AgentContextResourceConfig(BaseAgentResourceConfig): @@ -1162,6 +1168,7 @@ def _normalize_resources(v: Dict[str, Any]) -> None: "structured": "Structured", "deeprag": "DeepRAG", "batchtransform": "BatchTransform", + "datafabric": "DataFabric", "unknown": "Unknown", } diff --git a/src/uipath/platform/entities/_entities_service.py b/src/uipath/platform/entities/_entities_service.py index 2b31b830c..967542a6f 100644 --- a/src/uipath/platform/entities/_entities_service.py +++ b/src/uipath/platform/entities/_entities_service.py @@ -1,4 +1,4 @@ -from typing import Any, List, Optional, Type +from typing import Any, Dict, List, Optional, Type from httpx import Response @@ -389,6 +389,91 @@ class CustomerRecord: EntityRecord.from_data(data=record, model=schema) for record in records_data ] + @traced(name="entity_query_records", run_type="uipath") + def query_entity_records( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = entities_service.query_entity_records( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = entities_service.query_entity_records( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + response = self.request(spec.method, spec.endpoint, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + return records_data + + + @traced(name="query_entities_async", run_type="uipath") + async def query_entity_records_async( + self, + sql_query: str, + schema: Optional[Type[Any]] = None, + ) -> List[Dict[str, Any]]: + """Asynchronously query entity records using a SQL query. + + This method allows executing SQL queries directly against entity data + via the Data Fabric query endpoint. + + Args: + sql_query (str): The full SQL query to execute. Should be a valid + SELECT statement targeting the entity. + schema (Optional[Type[Any]]): Optional schema class for validation. + + Returns: + List[Dict[str, Any]]: A list of record dictionaries matching the query. + + Examples: + Basic query:: + + records = await entities_service.query_entity_records_async( + "SELECT * FROM Customers WHERE Status = 'Active' LIMIT 100" + ) + + Query with specific fields:: + + records = await entities_service.query_entity_records_async( + "SELECT OrderId, CustomerName, Amount FROM Orders WHERE Amount > 1000" + ) + """ + spec = self._query_entity_records_spec(sql_query) + headers = { + "X-UiPath-Internal-TenantName": self._url.tenant_name, + "X-UiPath-Internal-AccountName": self._url.org_name, + } + response = await self.request_async(spec.method, spec.endpoint, json=spec.json, headers=headers) + + records_data = response.json().get("results", []) + return records_data + @traced(name="entity_record_insert_batch", run_type="uipath") def insert_records( self, @@ -872,6 +957,16 @@ def _list_records_spec( params=({"start": start, "limit": limit}), ) + def _query_entity_records_spec( + self, + sql_query: str, + ) -> RequestSpec: + return RequestSpec( + method="POST", + endpoint=Endpoint("datafabric_/api/v1/query/execute"), + json={"query": sql_query}, + ) + def _insert_batch_spec(self, entity_key: str, records: List[Any]) -> RequestSpec: return RequestSpec( method="POST", diff --git a/tests/agent/models/test_agent.py b/tests/agent/models/test_agent.py index 63a8a917f..6af2b9663 100644 --- a/tests/agent/models/test_agent.py +++ b/tests/agent/models/test_agent.py @@ -2755,3 +2755,53 @@ def test_is_conversational_false_by_default(self): ) assert config.is_conversational is False + + +class TestDataFabricContextConfig: + """Tests for Data Fabric context resource configuration.""" + + def test_datafabric_retrieval_mode_exists(self): + """Test that DATA_FABRIC retrieval mode is defined.""" + assert AgentContextRetrievalMode.DATA_FABRIC == "DataFabric" + + def test_datafabric_context_config_parses(self): + """Test that Data Fabric context config parses correctly.""" + config = { + "$resourceType": "context", + "name": "Customer Data", + "description": "Query customer and order data", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", + "settings": { + "retrievalMode": "DataFabric", + "resultCount": 100, + "entityIdentifiers": ["customers-key", "orders-key"], + }, + } + + parsed = AgentContextResourceConfig.model_validate(config) + + assert parsed.name == "Customer Data" + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers == ["customers-key", "orders-key"] + + def test_datafabric_context_config_without_entity_identifiers(self): + """Test that entity_identifiers is optional.""" + config = { + "$resourceType": "context", + "name": "Test", + "description": "Test", + "isEnabled": True, + "folderPath": "Shared", + "indexName": "", + "settings": { + "retrievalMode": "DataFabric", + "resultCount": 10, + }, + } + + parsed = AgentContextResourceConfig.model_validate(config) + + assert parsed.settings.retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC + assert parsed.settings.entity_identifiers is None diff --git a/uv.lock b/uv.lock index 80cd0b9f6..9426a50ef 100644 --- a/uv.lock +++ b/uv.lock @@ -2580,6 +2580,7 @@ dev = [ { name = "termynal" }, { name = "tomli-w" }, { name = "types-toml" }, + { name = "uipath" }, { name = "virtualenv" }, ] @@ -2631,6 +2632,7 @@ dev = [ { name = "termynal", specifier = ">=0.13.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, { name = "types-toml", specifier = ">=0.10.8" }, + { name = "uipath", editable = "." }, { name = "virtualenv", specifier = ">=20.36.1" }, ]