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..f6cf31f80 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 Entity from ..datafabric_query_tool import DataFabricQueryTool from . import datafabric_prompt_builder @@ -42,18 +42,23 @@ class DataFabricSubgraphState(BaseModel): class QueryExecutor: """Executes SQL queries against Data Fabric.""" - def __init__(self, routing_context: QueryRoutingOverrideContext) -> None: + def __init__(self, folders_map: dict[str, str]) -> None: from uipath.platform import UiPath - - self._sdk = UiPath() - self._routing_context = routing_context + from uipath.platform.entities import EntitiesService + + sdk = UiPath() + self._entities = EntitiesService( + config=sdk._config, + execution_context=sdk._execution_context, + folders_service=sdk.folders, + folders_map=folders_map, + ) 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,15 +86,13 @@ def __init__( self, llm: BaseChatModel, entities: list[Entity], - routing_context: QueryRoutingOverrideContext, + folders_map: dict[str, str], 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 - ) + self._execute_sql_tool = self._create_execute_sql_tool(folders_map, entities) self._system_message = SystemMessage( content=datafabric_prompt_builder.build( entities, resource_description, base_system_prompt @@ -175,7 +178,7 @@ def router(self, state: DataFabricSubgraphState) -> str: def _create_execute_sql_tool( self, - routing_context: QueryRoutingOverrideContext, + folders_map: dict[str, str], entities: list[Entity], ) -> BaseTool: """Create the inner ``execute_sql`` tool.""" @@ -188,7 +191,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(folders_map), metadata={"tool_type": "datafabric_sql"}, ) @@ -196,7 +199,7 @@ def _create_execute_sql_tool( def create( llm: BaseChatModel, entities: list[Entity], - routing_context: QueryRoutingOverrideContext, + folders_map: dict[str, str], max_iterations: int = 25, resource_description: str = "", base_system_prompt: str = "", @@ -205,7 +208,7 @@ def create( graph = DataFabricGraph( llm, entities, - routing_context, + folders_map, 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..e163f1f09 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py @@ -21,7 +21,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 Entity from ..base_uipath_structured_tool import BaseUiPathStructuredTool from .models import DataFabricQueryInput @@ -41,13 +41,13 @@ class DataFabricTextQueryHandler: def __init__( self, entity_identifiers: list[str], - routing_context: QueryRoutingOverrideContext, + folders_map: dict[str, str], llm: BaseChatModel, resource_description: str = "", base_system_prompt: str = "", ) -> None: self._entity_identifiers = entity_identifiers - self._routing_context = routing_context + self._folders_map = folders_map self._llm = llm self._resource_description = resource_description self._base_system_prompt = base_system_prompt @@ -78,7 +78,7 @@ async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]: self._compiled = DataFabricGraph.create( llm=self._llm, entities=entities, - routing_context=self._routing_context, + folders_map=self._folders_map, resource_description=self._resource_description, base_system_prompt=self._base_system_prompt, ) @@ -120,20 +120,29 @@ async def fetch_entity_schemas(entity_identifiers: list[str]) -> list[Entity]: return [e for e in results if e is not None] -def _build_routing_context( +def _build_folders_map( resource: AgentContextResourceConfig, -) -> QueryRoutingOverrideContext: - """Build query routing context from entity set items. +) -> dict[str, str]: + """Build an entity-name-to-folder-id map from entity set items. - Maps each entity to its folder so the backend resolves - entities at folder level instead of tenant level. + Keys are always ``item.name`` (the entity name used in SQL generation) + so that ``with_folders_map()`` can match generated SQL table names. + When entity resource overwrites are active, the overwritten folder is + used instead of the original ``item.folder_id``. """ - return QueryRoutingOverrideContext( - entity_routings=[ - EntityRouting(entity_name=item.name, folder_id=item.folder_id) - for item in (resource.entity_set or []) - ] - ) + from uipath.platform.common._bindings import _resource_overwrites + + context_overwrites = _resource_overwrites.get() or {} + + folders_map: dict[str, str] = {} + for item in resource.entity_set or []: + overwrite = context_overwrites.get(f"entity.{item.id}") + if overwrite is not None: + folders_map[item.name] = overwrite.folder_identifier + else: + folders_map[item.name] = item.folder_id + + return folders_map def create_datafabric_query_tool( @@ -154,7 +163,7 @@ def create_datafabric_query_tool( config = agent_config or {} handler = DataFabricTextQueryHandler( entity_identifiers=resource.datafabric_entity_identifiers, - routing_context=_build_routing_context(resource), + folders_map=_build_folders_map(resource), llm=llm, resource_description=resource.description or "", base_system_prompt=config.get(BASE_SYSTEM_PROMPT, ""),