From a68afb3bb2a03a7c6b96d9b4a2f12215e84a235d Mon Sep 17 00:00:00 2001 From: Harshit Rohatgi Date: Wed, 12 Aug 2026 01:16:20 +0530 Subject: [PATCH] feat: add coded agent Data Fabric query tool --- samples/README.md | 3 + samples/datafabric-coded-agent/README.md | 17 ++ samples/datafabric-coded-agent/graph.py | 24 +++ samples/datafabric-coded-agent/langgraph.json | 7 + samples/datafabric-coded-agent/pyproject.toml | 14 ++ samples/datafabric-coded-agent/uipath.json | 1 + src/uipath_langchain/agent/tools/__init__.py | 4 +- .../agent/tools/datafabric_tool/__init__.py | 8 +- .../datafabric_tool/datafabric_prompts.py | 2 +- .../tools/datafabric_tool/datafabric_tool.py | 151 ++++++++++++++---- .../tools/test_datafabric_prompt_builder.py | 8 + tests/agent/tools/test_datafabric_tool.py | 124 +++++++++++++- 12 files changed, 322 insertions(+), 41 deletions(-) create mode 100644 samples/datafabric-coded-agent/README.md create mode 100644 samples/datafabric-coded-agent/graph.py create mode 100644 samples/datafabric-coded-agent/langgraph.json create mode 100644 samples/datafabric-coded-agent/pyproject.toml create mode 100644 samples/datafabric-coded-agent/uipath.json diff --git a/samples/README.md b/samples/README.md index 210fb5a2b..227acb222 100644 --- a/samples/README.md +++ b/samples/README.md @@ -9,6 +9,9 @@ This sample shows how to build an AI assistant using LangGraph and Tavily search ## [Company research agent](company-research-agent) This sample demonstrates how to create an agent that researches companies and develops outreach strategies using web search capabilities. +## [Data Fabric coded agent](datafabric-coded-agent) +This sample demonstrates a coded LangGraph agent using the same Data Fabric query-tool core as low-code entity-set contexts, including runtime entity and folder overrides. + ## [Email organizer agent](email-organizer-agent) This sample shows how to automate Outlook inbox organization with AI-powered rule suggestions and human-in-the-loop approval. diff --git a/samples/datafabric-coded-agent/README.md b/samples/datafabric-coded-agent/README.md new file mode 100644 index 000000000..2caf76338 --- /dev/null +++ b/samples/datafabric-coded-agent/README.md @@ -0,0 +1,17 @@ +# Data Fabric agent + +A minimal LangGraph agent that queries a Data Fabric entity with +`create_datafabric_tool` with direct entity references. + +Replace the entity ID, name, and folder key in `graph.py` with your entity values. + +Pass the coded agent's outer system prompt to both `create_agent` and +`create_datafabric_tool`. The latter forwards it to the tool's inner +SQL-generation graph so the same agent instructions apply at both levels. + +## Usage + +```bash +uv sync +uip codedagent run agent '{"messages":[{"role":"user","content":"List the names."}]}' +``` diff --git a/samples/datafabric-coded-agent/graph.py b/samples/datafabric-coded-agent/graph.py new file mode 100644 index 000000000..b09a47882 --- /dev/null +++ b/samples/datafabric-coded-agent/graph.py @@ -0,0 +1,24 @@ +from langchain.agents import create_agent +from uipath.platform.entities import DataFabricEntityItem + +from uipath_langchain.agent.tools import create_datafabric_tool +from uipath_langchain.chat import UiPathChat + +llm = UiPathChat(model="gpt-4.1-mini-2025-04-14") +system_prompt = "Answer questions using only the configured Data Fabric entities." + +datafabric_tool = create_datafabric_tool( + llm=llm, + name="query_agent_test", + description="Query the agentTest Data Fabric entity.", + base_system_prompt=system_prompt, + entities=[ + DataFabricEntityItem( + id="1312e893-8295-f111-9b33-0022482a9eea", + name="agentTest", + folder_key="379fec63-62b1-41ec-b2fc-718f8f7dda3c", + ) + ], +) + +graph = create_agent(llm, tools=[datafabric_tool], system_prompt=system_prompt) diff --git a/samples/datafabric-coded-agent/langgraph.json b/samples/datafabric-coded-agent/langgraph.json new file mode 100644 index 000000000..c465a881b --- /dev/null +++ b/samples/datafabric-coded-agent/langgraph.json @@ -0,0 +1,7 @@ +{ + "dependencies": ["."], + "graphs": { + "agent": "./graph.py:graph" + }, + "env": ".env" +} diff --git a/samples/datafabric-coded-agent/pyproject.toml b/samples/datafabric-coded-agent/pyproject.toml new file mode 100644 index 000000000..2e1e16925 --- /dev/null +++ b/samples/datafabric-coded-agent/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "datafabric-coded-agent" +version = "0.0.1" +description = "Coded agent that queries a UiPath Data Fabric entity" +authors = [{ name = "UiPath" }] +dependencies = [ + "uipath-langchain", +] +requires-python = ">=3.11" + +[dependency-groups] +dev = [ + "uipath-dev", +] diff --git a/samples/datafabric-coded-agent/uipath.json b/samples/datafabric-coded-agent/uipath.json new file mode 100644 index 000000000..0967ef424 --- /dev/null +++ b/samples/datafabric-coded-agent/uipath.json @@ -0,0 +1 @@ +{} diff --git a/src/uipath_langchain/agent/tools/__init__.py b/src/uipath_langchain/agent/tools/__init__.py index 7c6a7e37e..cf1e6d687 100644 --- a/src/uipath_langchain/agent/tools/__init__.py +++ b/src/uipath_langchain/agent/tools/__init__.py @@ -1,7 +1,8 @@ -"""Tool creation and management for LowCode agents.""" +"""Tool creation and management for low-code and coded agents.""" from .a2a import A2aClient, create_a2a_tools_and_clients, open_a2a_tools from .context_tool import create_context_tool +from .datafabric_tool import create_datafabric_tool from .escalation_tool import create_escalation_tool from .extraction_tool import create_ixp_extraction_tool from .integration_tool import create_integration_tool @@ -26,6 +27,7 @@ "create_tools_from_resources", "create_tool_node", "create_context_tool", + "create_datafabric_tool", "open_mcp_tools", "create_process_tool", "create_integration_tool", diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py b/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py index fccbda389..839e89a30 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/__init__.py @@ -1,9 +1,5 @@ """Data Fabric tool module for entity-based SQL queries.""" -from .datafabric_tool import ( - create_datafabric_query_tool, -) +from .datafabric_tool import create_datafabric_query_tool, create_datafabric_tool -__all__ = [ - "create_datafabric_query_tool", -] +__all__ = ["create_datafabric_query_tool", "create_datafabric_tool"] diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_prompts.py b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_prompts.py index 4dc078881..cfddebc22 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_prompts.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_prompts.py @@ -349,4 +349,4 @@ 8. **Explicit GROUP BY** - All non-aggregated columns in SELECT must be in GROUP BY 9. **Simple aggregations only** - No DISTINCT in aggregates 10. **ORDER BY only selected columns** - Cannot ORDER BY columns not in SELECT list -11. **ALWAYS include LIMIT** - Queries without WHERE must include a LIMIT clause (e.g., LIMIT 100). This applies to aggregates too (e.g., SELECT COUNT(col) FROM table LIMIT 1)""" +11. **Limit unbounded row queries** - Queries without WHERE that could return many rows must include a LIMIT clause (e.g., LIMIT 100). Scalar aggregate queries do not require LIMIT""" 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 aab4e4cfc..65c299556 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py @@ -13,6 +13,8 @@ import asyncio import logging +from collections.abc import Sequence +from dataclasses import dataclass from typing import Any from langchain_core.language_models import BaseChatModel @@ -30,6 +32,21 @@ BASE_SYSTEM_PROMPT = "base_system_prompt" +@dataclass(frozen=True, slots=True) +class _DataFabricToolConfig: + """Framework-neutral configuration for a Data Fabric query tool. + + Low-code contexts and coded-agent calls are normalized into this model before + the LangChain tool and its lazy query handler are created. + """ + + name: str + description: str + entities: tuple[DataFabricEntityItem, ...] + resource_description: str = "" + base_system_prompt: str = "" + + class DataFabricTextQueryHandler: """Manages lazy initialization and invocation of the Data Fabric sub-graph. @@ -139,49 +156,119 @@ def _format_terminal_tool_messages(tool_messages: list[ToolMessage]) -> str: ) +def _normalize_entities( + entities: Sequence[DataFabricEntityItem], +) -> tuple[DataFabricEntityItem, ...]: + """Copy entity references so caller mutations cannot change a built tool.""" + return tuple( + DataFabricEntityItem.model_validate(entity.model_dump(by_alias=True)) + for entity in entities + ) + + +def _default_tool_description(entities: Sequence[DataFabricEntityItem]) -> str: + entity_lines = [] + for entity in entities: + line = f"- {entity.name}" + if entity.description: + line += f": {entity.description}" + entity_lines.append(line) + entity_summary = "\n".join(entity_lines) + return ( + "Query the following Data Fabric entities using natural language:\n" + f"{entity_summary}\n" + "Describe what data you need and the tool will translate it to SQL, " + "execute the query, and return a natural language answer." + ) + + +def _build_datafabric_tool( + config: _DataFabricToolConfig, + llm: BaseChatModel, +) -> BaseTool: + """Build the shared LangChain tool used by coded and low-code agents.""" + handler = DataFabricTextQueryHandler( + entity_set=list(config.entities), + llm=llm, + resource_description=config.resource_description, + base_system_prompt=config.base_system_prompt, + ) + return BaseUiPathStructuredTool( + name=config.name, + description=config.description, + args_schema=DataFabricQueryInput, + coroutine=handler.__call__, + metadata={"tool_type": "datafabric_sql"}, + ) + + def create_datafabric_query_tool( resource: AgentContextResourceConfig, llm: BaseChatModel, tool_name: str = "query_datafabric", agent_config: dict[str, str] | None = None, ) -> BaseTool: - """Create the ``query_datafabric`` agentic tool. + """Create the low-code Data Fabric query tool from a context resource. + + Entity schemas and runtime binding overrides are resolved lazily on the first + invocation. Keep the resulting tool scoped to one agent execution so its + cached schema and routing cannot cross execution contexts. Args: - resource: The Data Fabric context resource configuration. - llm: The language model for the inner SQL generation loop. - tool_name: Sanitized tool name from the resource. - agent_config: Optional dict with agent-level config. - Key ``base_system_prompt`` carries the outer agent's system prompt. + resource: Low-code Data Fabric context resource. + llm: Language model for the inner SQL generation loop. + tool_name: LangChain tool name exposed to the agent. + agent_config: Optional agent-level configuration. Key + ``base_system_prompt`` carries the outer agent's system prompt. """ config = agent_config or {} - entity_set = [ - DataFabricEntityItem.model_validate(item.model_dump(by_alias=True)) - for item in (resource.entity_set or []) - ] - handler = DataFabricTextQueryHandler( - entity_set=entity_set, - llm=llm, - resource_description=resource.description or "", - base_system_prompt=config.get(BASE_SYSTEM_PROMPT, ""), + entity_set = _normalize_entities(resource.entity_set or []) + + return _build_datafabric_tool( + _DataFabricToolConfig( + name=tool_name, + description=_default_tool_description(entity_set), + entities=entity_set, + resource_description=resource.description or "", + base_system_prompt=config.get(BASE_SYSTEM_PROMPT, ""), + ), + llm, ) - entity_lines = [] - for e in entity_set: - line = f"- {e.name}" - if e.description: - line += f": {e.description}" - entity_lines.append(line) - entity_summary = "\n".join(entity_lines) - return BaseUiPathStructuredTool( - name=tool_name, - description=( - "Query the following Data Fabric entities using natural language:\n" - f"{entity_summary}\n" - "Describe what data you need and the tool will translate it to SQL, " - "execute the query, and return a natural language answer." + +def create_datafabric_tool( + *, + llm: BaseChatModel, + name: str, + description: str, + entities: Sequence[DataFabricEntityItem], + base_system_prompt: str, +) -> BaseTool: + """Create a Data Fabric query tool for a coded agent. + + Entity schemas are resolved lazily on the first invocation. Keep the tool + scoped to one agent execution so its cached schema and routing cannot cross + execution contexts. + + Pass the same outer-agent system prompt used to construct the coded agent so + its instructions are also available to the inner SQL-generation graph. + + Args: + llm: Language model for the inner SQL generation loop. + name: LangChain tool name exposed to the coded agent. + description: Description used by the outer agent to select the tool. + entities: Data Fabric entity references available to the tool. + base_system_prompt: Outer coded-agent system prompt forwarded to the + inner SQL-generation graph. + """ + entity_set = _normalize_entities(entities) + return _build_datafabric_tool( + _DataFabricToolConfig( + name=name, + description=description, + entities=entity_set, + resource_description=description, + base_system_prompt=base_system_prompt, ), - args_schema=DataFabricQueryInput, - coroutine=handler, - metadata={"tool_type": "datafabric_sql"}, + llm, ) diff --git a/tests/agent/tools/test_datafabric_prompt_builder.py b/tests/agent/tools/test_datafabric_prompt_builder.py index 601ca4b26..85a099492 100644 --- a/tests/agent/tools/test_datafabric_prompt_builder.py +++ b/tests/agent/tools/test_datafabric_prompt_builder.py @@ -125,6 +125,14 @@ def test_v1_prompt_documents_left_vs_inner_join_intent(): assert "INNER JOIN" in prompt +def test_v1_prompt_does_not_require_limit_for_scalar_aggregates(): + prompt = build([_fake_entity(_fake_field())]) + + assert "Scalar aggregate queries do not require LIMIT" in prompt + assert "This applies to aggregates too" not in prompt + assert "SELECT COUNT(col) FROM table LIMIT 1" not in prompt + + def test_relationship_subsection_absent_when_no_foreign_keys(): prompt = build([_fake_entity(_fake_field())]) diff --git a/tests/agent/tools/test_datafabric_tool.py b/tests/agent/tools/test_datafabric_tool.py index 57d83e2d0..f58db4d23 100644 --- a/tests/agent/tools/test_datafabric_tool.py +++ b/tests/agent/tools/test_datafabric_tool.py @@ -1,10 +1,20 @@ -from unittest.mock import MagicMock +from types import SimpleNamespace +from typing import Any, cast, get_type_hints +from unittest.mock import AsyncMock, MagicMock, patch import pytest from langchain_core.messages import AIMessage, ToolMessage +from uipath.agent.models.agent import AgentContextResourceConfig +from uipath.platform.entities import DataFabricEntityItem +import uipath_langchain.agent.tools as agent_tools +from uipath_langchain.agent.tools.base_uipath_structured_tool import ( + BaseUiPathStructuredTool, +) +from uipath_langchain.agent.tools.context_tool import create_context_tool from uipath_langchain.agent.tools.datafabric_tool.datafabric_tool import ( DataFabricTextQueryHandler, + create_datafabric_tool, ) @@ -16,6 +26,118 @@ async def ainvoke(self, _state): return self._result_state +def _entity() -> DataFabricEntityItem: + return DataFabricEntityItem( + id="1312e893-8295-f111-9b33-0022482a9eea", + entity_key="agentTest", + name="agentTest", + folder_key="379fec63-62b1-41ec-b2fc-718f8f7dda3c", + description="Agent test records", + ) + + +def test_coded_datafabric_factory_is_public(): + assert agent_tools.create_datafabric_tool is create_datafabric_tool + assert not hasattr(agent_tools, "create_datafabric_query_tool") + + +def test_create_datafabric_tool_builds_directly_configured_tool(): + entity = _entity() + tool = create_datafabric_tool( + llm=MagicMock(), + name="query_agent_test", + description="Query the agentTest entity.", + entities=[entity], + base_system_prompt="Answer only from Data Fabric.", + ) + + assert tool.name == "query_agent_test" + assert tool.description == "Query the agentTest entity." + assert tool.metadata == {"tool_type": "datafabric_sql"} + assert isinstance(tool, BaseUiPathStructuredTool) + assert tool.coroutine is not None + handler = cast(Any, tool.coroutine).__self__ + assert isinstance(handler, DataFabricTextQueryHandler) + assert get_type_hints(tool.coroutine) == {"user_query": str, "return": str} + assert handler._resource_description == "Query the agentTest entity." + assert handler._base_system_prompt == "Answer only from Data Fabric." + assert handler._entity_set == [_entity()] + + entity.name = "mutated-after-tool-creation" + assert handler._entity_set[0].name == "agentTest" + + +def test_create_datafabric_tool_preserves_empty_entity_set_behavior(): + tool = create_datafabric_tool( + llm=MagicMock(), + name="query_datafabric", + description="Query Data Fabric.", + entities=[], + base_system_prompt="Answer only from Data Fabric.", + ) + + assert tool.name == "query_datafabric" + + +def test_low_code_context_uses_resource_based_datafabric_query_factory(): + resource = AgentContextResourceConfig( + name="Agent Test Data", + description="Low-code Data Fabric context.", + contextType="datafabricentityset", + entitySet=[_entity()], + ) + + with patch( + "uipath_langchain.agent.tools.context_tool._extract_system_prompt", + return_value="Low-code system prompt.", + ): + tool = create_context_tool(resource, MagicMock()) + + assert isinstance(tool, BaseUiPathStructuredTool) + assert tool.name == "Agent_Test_Data" + assert tool.description == ( + "Query the following Data Fabric entities using natural language:\n" + "- agentTest: Agent test records\n" + "Describe what data you need and the tool will translate it to SQL, " + "execute the query, and return a natural language answer." + ) + assert tool.coroutine is not None + handler = cast(Any, tool.coroutine).__self__ + assert isinstance(handler, DataFabricTextQueryHandler) + assert handler._resource_description == "Low-code Data Fabric context." + assert handler._base_system_prompt == "Low-code system prompt." + assert handler._entity_set == [_entity()] + + +@pytest.mark.asyncio +async def test_datafabric_handler_resolves_entities_lazily_once(): + entity = _entity() + sdk = MagicMock() + sdk.entities.resolve_entity_set_async = AsyncMock( + return_value=SimpleNamespace( + entities=[MagicMock()], + entities_service=MagicMock(), + ) + ) + compiled = _FakeCompiledGraph({"messages": []}) + handler = DataFabricTextQueryHandler(entity_set=[entity], llm=MagicMock()) + + with ( + patch("uipath.platform.UiPath", return_value=sdk), + patch( + "uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph.DataFabricGraph.create", + return_value=compiled, + ) as create_graph, + ): + first = await handler._ensure_datafabric_graph() + second = await handler._ensure_datafabric_graph() + + assert first is compiled + assert second is compiled + sdk.entities.resolve_entity_set_async.assert_awaited_once_with([entity]) + create_graph.assert_called_once() + + @pytest.mark.asyncio async def test_datafabric_handler_returns_single_terminal_tool_message(): handler = DataFabricTextQueryHandler(