diff --git a/src/uipath_langchain/agent/tools/__init__.py b/src/uipath_langchain/agent/tools/__init__.py index 6c653b201..f434ed140 100644 --- a/src/uipath_langchain/agent/tools/__init__.py +++ b/src/uipath_langchain/agent/tools/__init__.py @@ -1,9 +1,6 @@ """Tool creation and management for LowCode agents.""" from .context_tool import create_context_tool -from .datafabric_tool import ( - fetch_entity_schemas, -) from .escalation_tool import create_escalation_tool from .extraction_tool import create_ixp_extraction_tool from .integration_tool import create_integration_tool @@ -30,7 +27,6 @@ "create_escalation_tool", "create_ixp_extraction_tool", "create_ixp_escalation_tool", - "fetch_entity_schemas", "UiPathToolNode", "ToolWrapperMixin", "wrap_tools_with_error_handling", diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py b/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py index 4a72ede98..fccbda389 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py @@ -2,10 +2,8 @@ from .datafabric_tool import ( create_datafabric_query_tool, - fetch_entity_schemas, ) __all__ = [ "create_datafabric_query_tool", - "fetch_entity_schemas", ] diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py index 928dea1e7..089ab0ab5 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py @@ -23,7 +23,7 @@ from langgraph.graph.message import add_messages from langgraph.graph.state import CompiledStateGraph from pydantic import BaseModel -from uipath.platform.entities import Entity, QueryRoutingOverrideContext +from uipath.platform.entities import EntitiesService, Entity from ..datafabric_query_tool import DataFabricQueryTool from . import datafabric_prompt_builder @@ -42,18 +42,14 @@ class DataFabricSubgraphState(BaseModel): class QueryExecutor: """Executes SQL queries against Data Fabric.""" - def __init__(self, routing_context: QueryRoutingOverrideContext) -> None: - from uipath.platform import UiPath - - self._sdk = UiPath() - self._routing_context = routing_context + def __init__(self, entities_service: EntitiesService) -> None: + self._entities = entities_service async def __call__(self, sql_query: str) -> dict[str, Any]: logger.debug("execute_sql called with SQL: %s", sql_query) try: - records = await self._sdk.entities.query_entity_records_async( + records = await self._entities.query_entity_records_async( sql_query=sql_query, - routing_context=self._routing_context, ) return { "records": records, @@ -81,14 +77,14 @@ def __init__( self, llm: BaseChatModel, entities: list[Entity], - routing_context: QueryRoutingOverrideContext, + entities_service: EntitiesService, max_iterations: int = 25, resource_description: str = "", base_system_prompt: str = "", ) -> None: self._max_iterations = max_iterations self._execute_sql_tool = self._create_execute_sql_tool( - routing_context, entities + entities_service, entities ) self._system_message = SystemMessage( content=datafabric_prompt_builder.build( @@ -175,7 +171,7 @@ def router(self, state: DataFabricSubgraphState) -> str: def _create_execute_sql_tool( self, - routing_context: QueryRoutingOverrideContext, + entities_service: EntitiesService, entities: list[Entity], ) -> BaseTool: """Create the inner ``execute_sql`` tool.""" @@ -188,7 +184,7 @@ def _create_execute_sql_tool( "tables and columns. Retry with a corrected query on errors." ), args_schema=DataFabricExecuteSqlInput, - coroutine=QueryExecutor(routing_context), + coroutine=QueryExecutor(entities_service), metadata={"tool_type": "datafabric_sql"}, ) @@ -196,7 +192,7 @@ def _create_execute_sql_tool( def create( llm: BaseChatModel, entities: list[Entity], - routing_context: QueryRoutingOverrideContext, + entities_service: EntitiesService, max_iterations: int = 25, resource_description: str = "", base_system_prompt: str = "", @@ -205,7 +201,7 @@ def create( graph = DataFabricGraph( llm, entities, - routing_context, + entities_service, max_iterations, resource_description, base_system_prompt, diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py index f67b148af..bea8123d9 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py @@ -1,8 +1,7 @@ """Data Fabric tool creation and resource detection. -This module provides: -1. An agentic ``query_datafabric`` tool with inner LLM sub-graph -2. Entity schema fetching from the Data Fabric API +This module provides an agentic ``query_datafabric`` tool with an inner +LLM sub-graph. The tool accepts natural language queries, runs an inner LangGraph sub-graph for SQL generation + execution + self-correction, and @@ -21,7 +20,7 @@ from langchain_core.tools import BaseTool from langgraph.graph.state import CompiledStateGraph from uipath.agent.models.agent import AgentContextResourceConfig -from uipath.platform.entities import Entity, EntityRouting, QueryRoutingOverrideContext +from uipath.platform.entities import DataFabricEntityItem from ..base_uipath_structured_tool import BaseUiPathStructuredTool from .models import DataFabricQueryInput @@ -34,20 +33,19 @@ class DataFabricTextQueryHandler: """Manages lazy initialization and invocation of the Data Fabric sub-graph. - On first call, fetches entity schemas from the DF API and compiles - the inner LangGraph sub-graph. Subsequent calls reuse the cached graph. + On first call, resolves entity schemas and routing via the platform + layer and compiles the inner LangGraph sub-graph. Subsequent calls + reuse the cached graph. """ def __init__( self, - entity_identifiers: list[str], - routing_context: QueryRoutingOverrideContext, + entity_set: list[DataFabricEntityItem], llm: BaseChatModel, resource_description: str = "", base_system_prompt: str = "", ) -> None: - self._entity_identifiers = entity_identifiers - self._routing_context = routing_context + self._entity_set = entity_set self._llm = llm self._resource_description = resource_description self._base_system_prompt = base_system_prompt @@ -55,7 +53,7 @@ def __init__( self._init_lock = asyncio.Lock() async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]: - """Lazy-init: fetch schemas + build sub-graph on first call. + """Lazy-init: resolve entities + build sub-graph on first call. Uses asyncio.Lock because the outer agent supports parallel tool calls — two concurrent invocations could race on first call. @@ -67,18 +65,23 @@ async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]: if self._compiled is not None: return self._compiled + from uipath.platform import UiPath + from .datafabric_subgraph import DataFabricGraph - entities = await fetch_entity_schemas(self._entity_identifiers) - if not entities: + sdk = UiPath() + resolution = await sdk.entities.resolve_entity_set_async( + self._entity_set + ) + if not resolution.entities: raise ValueError( "No Data Fabric entity schemas could be fetched. " "Check entity identifiers and permissions." ) self._compiled = DataFabricGraph.create( llm=self._llm, - entities=entities, - routing_context=self._routing_context, + entities=resolution.entities, + entities_service=resolution.entities_service, resource_description=self._resource_description, base_system_prompt=self._base_system_prompt, ) @@ -98,44 +101,6 @@ async def __call__(self, user_query: str) -> str: return "Unable to generate an answer from the available data." -async def _fetch_single_entity(sdk: Any, identifier: str) -> Entity | None: - """Fetch a single entity by identifier, returning None on failure.""" - try: - entity = await sdk.entities.retrieve_async(identifier) - logger.info("Fetched schema for entity '%s'", entity.display_name) - return entity - except Exception: - logger.warning("Failed to fetch entity '%s'", identifier, exc_info=True) - return None - - -async def fetch_entity_schemas(entity_identifiers: list[str]) -> list[Entity]: - """Fetch entity metadata from Data Fabric concurrently.""" - from uipath.platform import UiPath - - sdk = UiPath() - results = await asyncio.gather( - *[_fetch_single_entity(sdk, eid) for eid in entity_identifiers] - ) - return [e for e in results if e is not None] - - -def _build_routing_context( - resource: AgentContextResourceConfig, -) -> QueryRoutingOverrideContext: - """Build query routing context from entity set items. - - Maps each entity to its folder so the backend resolves - entities at folder level instead of tenant level. - """ - return QueryRoutingOverrideContext( - entity_routings=[ - EntityRouting(entity_name=item.name, folder_id=item.folder_id) - for item in (resource.entity_set or []) - ] - ) - - def create_datafabric_query_tool( resource: AgentContextResourceConfig, llm: BaseChatModel, @@ -153,8 +118,7 @@ def create_datafabric_query_tool( """ config = agent_config or {} handler = DataFabricTextQueryHandler( - entity_identifiers=resource.datafabric_entity_identifiers, - routing_context=_build_routing_context(resource), + entity_set=resource.entity_set or [], llm=llm, resource_description=resource.description or "", base_system_prompt=config.get(BASE_SYSTEM_PROMPT, ""),