feat: add Data Fabric tool support (#726)

Author: milind-jain-uipath

pyproject.toml

@@ -7,7 +7,7 @@ requires-python = ">=3.11"
 dependencies = [
     "uipath>=2.10.29, <2.11.0",
     "uipath-core>=0.5.2, <0.6.0",
-    "uipath-platform>=0.1.10, <0.2.0",
+    "uipath-platform>=0.1.12, <0.2.0",
     "uipath-runtime>=0.10.0, <0.11.0",
     "langgraph>=1.0.0, <2.0.0",
     "langchain-core>=1.2.11, <2.0.0",

src/uipath_langchain/agent/tools/__init__.py

@@ -1,6 +1,9 @@
 """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
@@ -22,6 +25,7 @@
     "create_escalation_tool",
     "create_ixp_extraction_tool",
     "create_ixp_escalation_tool",
+    "fetch_entity_schemas",
     "UiPathToolNode",
     "ToolWrapperMixin",
 ]

src/uipath_langchain/agent/tools/context_tool.py

@@ -6,14 +6,18 @@
 
 from jsonpath_ng import parse  # type: ignore[import-untyped]
 from langchain_core.documents import Document
+from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import ToolCall
 from langchain_core.tools import BaseTool, StructuredTool
 from pydantic import BaseModel, Field, create_model
 from uipath.agent.models.agent import (
     AgentContextResourceConfig,
     AgentContextRetrievalMode,
+    AgentContextType,
+    AgentMessageRole,
     AgentToolArgumentArgumentProperties,
     AgentToolArgumentProperties,
+    LowCodeAgentDefinition,
 )
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
@@ -130,16 +134,48 @@ def is_static_query(resource: AgentContextResourceConfig) -> bool:
     return resource.settings.query.variant.lower() == "static"
 
 
-def create_context_tool(resource: AgentContextResourceConfig) -> StructuredTool:
+def _extract_system_prompt(agent: LowCodeAgentDefinition | None) -> str:
+    """Extract system prompt from agent definition messages."""
+    if agent is None:
+        return ""
+    return "\n\n".join(
+        msg.content
+        for msg in agent.messages
+        if msg.role == AgentMessageRole.SYSTEM and msg.content
+    )
+
+
+def create_context_tool(
+    resource: AgentContextResourceConfig,
+    llm: BaseChatModel | None = None,
+    agent: LowCodeAgentDefinition | None = None,
+) -> StructuredTool | BaseTool | None:
+    assert resource.context_type is not None
     tool_name = sanitize_tool_name(resource.name)
+
+    if resource.context_type == AgentContextType.DATA_FABRIC_ENTITY_SET:
+        if llm is None:
+            raise ValueError("Data Fabric entity set tools require an LLM instance")
+        from .datafabric_tool import create_datafabric_query_tool
+        from .datafabric_tool.datafabric_tool import BASE_SYSTEM_PROMPT
+
+        return create_datafabric_query_tool(
+            resource,
+            llm,
+            tool_name=tool_name,
+            agent_config={BASE_SYSTEM_PROMPT: _extract_system_prompt(agent)},
+        )
+
     assert resource.settings is not None
     retrieval_mode = resource.settings.retrieval_mode.lower()
+
     if retrieval_mode == AgentContextRetrievalMode.DEEP_RAG.value.lower():
         return handle_deep_rag(tool_name, resource)
-    elif retrieval_mode == AgentContextRetrievalMode.BATCH_TRANSFORM.value.lower():
+
+    if retrieval_mode == AgentContextRetrievalMode.BATCH_TRANSFORM.value.lower():
         return handle_batch_transform(tool_name, resource)
-    else:
-        return handle_semantic_search(tool_name, resource)
+
+    return handle_semantic_search(tool_name, resource)
 
 
 def handle_semantic_search(

src/uipath_langchain/agent/tools/datafabric_query_tool.py

@@ -0,0 +1,45 @@
+"""Data Fabric query tool with SQL syntax validation via sqlparse."""
+
+from typing import Any
+
+import sqlparse
+from langchain_core.callbacks import AsyncCallbackManagerForToolRun
+from langchain_core.runnables import RunnableConfig
+
+from .base_uipath_structured_tool import BaseUiPathStructuredTool
+
+
+def _validate_sql(sql: str) -> str | None:
+    """Validate SQL syntax using sqlparse.
+
+    Returns:
+        Error string if invalid, None if valid.
+    """
+    parsed = sqlparse.parse(sql)
+    if not parsed or not parsed[0].tokens:
+        return "Empty or unparseable SQL query"
+    return None
+
+
+class DataFabricQueryTool(BaseUiPathStructuredTool):
+    """Data Fabric query tool with SQL syntax validation.
+
+    Validates that the input SQL is parseable before delegating
+    to the underlying coroutine. On validation failure, raises
+    a ValueError so the caller can handle it as needed.
+    """
+
+    async def _arun(
+        __obj_internal_self__,
+        *args: Any,
+        config: RunnableConfig,
+        run_manager: AsyncCallbackManagerForToolRun | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        sql_query = kwargs.get("sql_query") or (args[0] if args else "")
+        error = _validate_sql(sql_query)
+        if error:
+            raise ValueError(error)
+        return await super()._arun(
+            *args, config=config, run_manager=run_manager, **kwargs
+        )

src/uipath_langchain/agent/tools/datafabric_tool/__init__.py

@@ -0,0 +1,11 @@
+"""Data Fabric tool module for entity-based SQL queries."""
+
+from .datafabric_tool import (
+    create_datafabric_query_tool,
+    fetch_entity_schemas,
+)
+
+__all__ = [
+    "create_datafabric_query_tool",
+    "fetch_entity_schemas",
+]

src/uipath_langchain/agent/tools/datafabric_tool/datafabric_prompt_builder.py

@@ -0,0 +1,195 @@
+"""Schema context building and formatting for Data Fabric entities.
+
+Converts raw Entity SDK objects into structured Pydantic models (SQLContext),
+then formats them as text for system prompt injection.
+
+Note: This module will go through refinements as we better understand
+the tool's performance characteristics and scoring in production.
+"""
+
+import logging
+
+from uipath.platform.entities import Entity
+
+from .datafabric_prompts import SQL_CONSTRAINTS, SQL_EXPERT_SYSTEM_PROMPT
+from .models import (
+    EntitySchema,
+    EntitySQLContext,
+    FieldSchema,
+    QueryPattern,
+    SQLContext,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def build_entity_context(entity: Entity) -> EntitySQLContext:
+    """Convert an Entity SDK object to schema + derived query patterns."""
+    field_schemas: list[FieldSchema] = []
+    numeric_field: str | None = None
+    text_field: str | None = None
+
+    for field in entity.fields or []:
+        if field.is_hidden_field or field.is_system_field:
+            continue
+        type_name = field.sql_type.name if field.sql_type else "unknown"
+        fs = FieldSchema(
+            name=field.name,
+            display_name=field.display_name,
+            type=type_name,
+            description=field.description,
+            is_foreign_key=field.is_foreign_key,
+            is_required=field.is_required,
+            is_unique=field.is_unique,
+            nullable=not field.is_required,
+        )
+        field_schemas.append(fs)
+
+        if not numeric_field and fs.is_numeric:
+            numeric_field = fs.name
+        if not text_field and fs.is_text:
+            text_field = fs.name
+
+    field_names = [f.name for f in field_schemas]
+    table = entity.name
+
+    group_field = text_field or (field_names[0] if field_names else "Category")
+    agg_field = numeric_field or (field_names[1] if len(field_names) > 1 else "Amount")
+    filter_field = text_field or (field_names[0] if field_names else "Name")
+    fields_sample = ", ".join(field_names[:5]) if field_names else "*"
+    count_col = field_names[0] if field_names else "id"
+
+    query_patterns = [
+        QueryPattern(
+            intent="Show all",
+            sql=f"SELECT {fields_sample} FROM {table} LIMIT 100",
+        ),
+        QueryPattern(
+            intent="Find by X",
+            sql=f"SELECT {fields_sample} FROM {table} WHERE {filter_field} = 'value' LIMIT 100",
+        ),
+        QueryPattern(
+            intent="Top N by Y",
+            sql=f"SELECT {fields_sample} FROM {table} ORDER BY {agg_field} DESC LIMIT N",
+        ),
+        QueryPattern(
+            intent="Count by X",
+            sql=f"SELECT {group_field}, COUNT({count_col}) as count FROM {table} GROUP BY {group_field}",
+        ),
+        QueryPattern(
+            intent="Top N segments",
+            sql=f"SELECT {group_field}, COUNT({count_col}) as count FROM {table} GROUP BY {group_field} ORDER BY count DESC LIMIT N",
+        ),
+        QueryPattern(
+            intent="Sum/Avg of Y",
+            sql=f"SELECT SUM({agg_field}) as total FROM {table}",
+        ),
+    ]
+
+    schema = EntitySchema(
+        id=entity.id,
+        entity_name=entity.name,
+        display_name=entity.display_name or entity.name,
+        description=entity.description,
+        record_count=entity.record_count,
+        fields=field_schemas,
+    )
+    return EntitySQLContext(entity_schema=schema, query_patterns=query_patterns)
+
+
+def build_sql_context(
+    entities: list[Entity],
+    resource_description: str = "",
+    base_system_prompt: str = "",
+) -> SQLContext:
+    """Build the full SQL context from entities, prompts, and constraints."""
+    return SQLContext(
+        base_system_prompt=base_system_prompt or None,
+        resource_description=resource_description or None,
+        sql_expert_system_prompt=SQL_EXPERT_SYSTEM_PROMPT,
+        constraints=SQL_CONSTRAINTS,
+        entity_contexts=[build_entity_context(e) for e in entities],
+    )
+
+
+def format_sql_context(ctx: SQLContext) -> str:
+    """Format a SQLContext as text for system prompt injection."""
+    lines: list[str] = []
+
+    if ctx.base_system_prompt:
+        lines.append("## Agent Instructions")
+        lines.append("")
+        lines.append(ctx.base_system_prompt)
+        lines.append("")
+
+    if ctx.sql_expert_system_prompt:
+        lines.append("## SQL Query Generation Guidelines")
+        lines.append("")
+        lines.append(ctx.sql_expert_system_prompt)
+        lines.append("")
+
+    if ctx.constraints:
+        lines.append("## SQL Constraints")
+        lines.append("")
+        lines.append(ctx.constraints)
+        lines.append("")
+
+    if ctx.resource_description:
+        lines.append("## Entity set description")
+        lines.append("")
+        lines.append(ctx.resource_description)
+        lines.append("")
+
+    lines.append("## All available Data Fabric Entities")
+    lines.append("")
+
+    for entity_ctx in ctx.entity_contexts:
+        entity = entity_ctx.entity_schema
+        lines.append(
+            f"### Entity: {entity.display_name} (SQL table: `{entity.entity_name}`)"
+        )
+        if entity.description:
+            lines.append(f"_{entity.description}_")
+        lines.append("")
+        lines.append("| Field | Type |")
+        lines.append("|-------|------|")
+
+        for field in entity.fields:
+            lines.append(f"| {field.name} | {field.display_type} |")
+
+        lines.append("")
+
+        lines.append(f"**Query Patterns for {entity.entity_name}:**")
+        lines.append("")
+        lines.append("| User Intent | SQL Pattern |")
+        lines.append("|-------------|-------------|")
+        for p in entity_ctx.query_patterns:
+            lines.append(f"| '{p.intent}' | `{p.sql}` |")
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+def build(
+    entities: list[Entity],
+    resource_description: str = "",
+    base_system_prompt: str = "",
+) -> str:
+    """Build the full SQL prompt text for the inner sub-graph LLM.
+
+    Combines agent system prompt, resource description, SQL guidelines,
+    constraints, entity schemas, and query patterns into a single prompt string.
+
+    Args:
+        entities: List of Entity objects with schema information.
+        resource_description: Optional description of the resource/entity set.
+        base_system_prompt: Optional system prompt from the outer agent.
+
+    Returns:
+        Formatted prompt string for the inner LLM system message.
+    """
+    if not entities:
+        return ""
+
+    ctx = build_sql_context(entities, resource_description, base_system_prompt)
+    return format_sql_context(ctx)