diff --git a/pyproject.toml b/pyproject.toml index b1bc33f15..9e8e791d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,9 +5,9 @@ description = "Python SDK that enables developers to build and deploy LangGraph readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" dependencies = [ - "uipath>=2.13.0, <2.14.0", + "uipath>=2.13.2, <2.14.0", "uipath-core>=0.5.28, <0.6.0", - "uipath-platform>=0.2.0, <0.3.0", + "uipath-platform>=0.2.1, <0.3.0", "uipath-runtime>=0.12.1, <0.13.0", "langgraph>=1.1.8, <2.0.0", "langchain-core>=1.2.27, <2.0.0", @@ -23,6 +23,7 @@ dependencies = [ "mcp==1.26.0", "langchain-mcp-adapters==0.2.1", "pillow>=12.1.1", + "rdflib>=7.0.0, <8.0.0", "a2a-sdk>=0.2.0,<1.0.0", "uipath-langchain-client[openai]>=1.16.1,<1.17.0", ] diff --git a/sonar-project.properties b/sonar-project.properties index 159119f9a..5b4eba17e 100644 --- a/sonar-project.properties +++ b/sonar-project.properties @@ -10,4 +10,10 @@ sonar.python.coverage.reportPaths=coverage.xml sonar.exclusions=**/samples/**,**/testcases/**,**/template/** +# The Data Fabric ontology tool is an intentionally self-contained duplicate of +# the entity tool's sub-graph/prompt rendering, kept separate so the entity path +# (datafabric_subgraph.py / datafabric_tool.py) stays byte-identical to upstream. +# Exclude it from copy-paste detection so that deliberate isolation isn't flagged. +sonar.cpd.exclusions=**/datafabric_tool/ontology/** + sonar.sourceEncoding=UTF-8 diff --git a/src/uipath_langchain/agent/tools/context_tool.py b/src/uipath_langchain/agent/tools/context_tool.py index c22906835..7c1c1a508 100644 --- a/src/uipath_langchain/agent/tools/context_tool.py +++ b/src/uipath_langchain/agent/tools/context_tool.py @@ -158,9 +158,33 @@ def create_context_tool( ) -> StructuredTool | BaseTool | None: tool_name = sanitize_tool_name(resource.name) + # A Data Fabric ontology context is a standalone tool (flag-gated): the agent + # selects ontologies and the tool derives its entities from the R2RML mapping. + # With the flag off it is inert (returns None) and the agent runs without it. + if resource.context_type == AgentContextType.DATA_FABRIC_ONTOLOGY: + from uipath.core.feature_flags import FeatureFlags + + from .datafabric_tool.datafabric_tool import BASE_SYSTEM_PROMPT + from .datafabric_tool.ontology import DATAFABRIC_ONTOLOGY_FF + + if not FeatureFlags.is_flag_enabled(DATAFABRIC_ONTOLOGY_FF, default=False): + return None + if llm is None: + raise ValueError("Data Fabric ontology tools require an LLM instance") + + from .datafabric_tool.ontology import create_datafabric_ontology_tool + + return create_datafabric_ontology_tool( + resource, + llm, + tool_name=tool_name, + agent_config={BASE_SYSTEM_PROMPT: _extract_system_prompt(agent)}, + ) + 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 diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/__init__.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/__init__.py new file mode 100644 index 000000000..949b9fe1d --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/__init__.py @@ -0,0 +1,15 @@ +"""Standalone Data Fabric **ontology** tool. + +Grouped in its own subpackage so it does not touch the entity tool's +``datafabric_tool``/``datafabric_subgraph``. The agent selects ontologies; this +tool derives its entity allow-list from each ontology's R2RML mapping, grounds an +inner SQL sub-graph on the OWL + R2RML, and reuses the entity tool's execute-SQL +path. Feature-flag gated by :data:`DATAFABRIC_ONTOLOGY_FF`. +""" + +from .ontology_tool import DATAFABRIC_ONTOLOGY_FF, create_datafabric_ontology_tool + +__all__ = [ + "DATAFABRIC_ONTOLOGY_FF", + "create_datafabric_ontology_tool", +] diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_fetcher.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_fetcher.py new file mode 100644 index 000000000..f7099c024 --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_fetcher.py @@ -0,0 +1,67 @@ +"""Fetch ontology component files (OWL, R2RML) from Data Fabric for the ontology tool. + +A Data Fabric ontology exposes typed files — OWL (the semantic schema) and R2RML +(the ontology→table/column mapping) — via +``EntitiesService.get_ontology_file_async``. This module provides: + +* :func:`fetch_ontology_file` — a thin fetch of one file's raw content, which + **raises** on failure or oversize. Callers decide whether a given file is + critical (R2RML — it is also the entity allow-list) or optional (OWL — the + prompt can degrade without it). +* :func:`fence_ontology_block` — wraps content in a labelled block for prompt + injection, tagging the file type and media type (e.g. Turtle vs OWL Functional + Notation) so the LLM knows how to read it. + +Ontology names/folders are pinned from configuration, never supplied by the LLM. +""" + +from uipath.platform.entities import EntitiesService + +# Defensive cap per file so a malformed/oversized artifact can't blow up the +# prompt/token budget. OWL + R2RML for a real ontology are comfortably under this. +_MAX_ONTOLOGY_BYTES = 2_000_000 + + +async def fetch_ontology_file( + entities_service: EntitiesService, + name: str, + file_type: str, + folder_key: str | None, +) -> tuple[str, str]: + """Fetch one ontology file's ``(content, media_type)``. + + Args: + entities_service: Authenticated SDK service used for the REST call. + name: Ontology name (pinned from configuration). + file_type: The ontology file to fetch — ``"owl"`` or ``"r2rml"``. + folder_key: Key of the folder the ontology lives in. + + Returns: + ``(content, media_type)`` — the raw file body and its media type. + + Raises: + Exception: Propagates the underlying fetch error. Also raises + ``ValueError`` if the body exceeds :data:`_MAX_ONTOLOGY_BYTES`. + Callers apply the critical/optional policy (raise vs degrade). + """ + data = await entities_service.get_ontology_file_async(name, file_type, folder_key) + content = data.get("content") or "" + media_type = data.get("mediaType") or "" + if len(content.encode("utf-8")) > _MAX_ONTOLOGY_BYTES: + raise ValueError( + f"Ontology '{name}' {file_type} exceeds the size limit " + f"({_MAX_ONTOLOGY_BYTES} bytes)." + ) + return content, media_type + + +def fence_ontology_block( + name: str, file_type: str, content: str, media_type: str = "" +) -> str: + """Wrap ontology file content in a labelled block for prompt injection.""" + label = file_type.upper() + if media_type: + label = f"{label}, {media_type}" + return ( + f"--- {label}: {name} ---\n{content}\n--- END {file_type.upper()}: {name} ---" + ) diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_prompt_builder.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_prompt_builder.py new file mode 100644 index 000000000..3fb2dc68e --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_prompt_builder.py @@ -0,0 +1,146 @@ +"""Inner system-prompt builder for the Data Fabric **ontology** tool. + +Separate from the entity tool's ``datafabric_prompt_builder`` so the two prompts +can evolve independently. This builder assembles an ontology-grounded prompt: +the OWL (authoritative semantic schema) and the R2RML (ontology→table/column +mapping) come first, then the SQL strategy/constraints and the entity schema +tables. + +It reuses only ``build_sql_context`` (entity-context + strategy prompt +construction) from ``datafabric_prompt_builder``; the top-level layout, the +ontology-specific framing, and the entity-schema rendering are all owned here so +the shared entity-tool builder is not modified by this feature. +""" + +from uipath.platform.entities import Entity + +from ..datafabric_prompt_builder import build_sql_context +from ..models import SQLContext + + +def _render_entity_schema_sections(ctx: SQLContext) -> list[str]: + """Render the entity-schema tables + query patterns as prompt lines.""" + lines: list[str] = ["## All available Data Fabric Entities", ""] + + 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 | Description |") + lines.append("|-------|------|-------------|") + for field in entity.fields: + desc = (field.description or "").replace("|", r"\|").replace("\n", " ") + lines.append(f"| {field.name} | {field.display_type} | {desc} |") + + 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 lines + + +def format_ontology_context( + ctx: SQLContext, + ontology_text: str = "", + r2rml_text: str = "", +) -> str: + """Format a SQLContext + ontology artifacts as the ontology-tool prompt.""" + lines: list[str] = [] + + if ctx.base_system_prompt: + lines.append("## Agent Instructions") + lines.append("") + lines.append(ctx.base_system_prompt) + lines.append("") + + if ontology_text: + lines.append( + "## Available Ontology (authoritative semantic schema)\n\n" + "The ontology below is the authoritative source for the exact column " + "names, value formats (date formats, codes, zero-padding), allowed " + "values, and the relationships between entities — richer and more " + "reliable than the field list further down, which omits value formats " + "and semantics. Base your column names, filter values, and joins on " + "it; when it and the entity tables disagree, the ontology wins.\n\n" + f"{ontology_text}" + ) + lines.append("") + + if r2rml_text: + lines.append( + "## Ontology→Table Mapping (R2RML)\n\n" + "The R2RML below maps the ontology to the physical tables and columns " + "you must use in SQL: `rr:tableName` is the SQL table, `rr:column` is " + "the SQL column for each ontology predicate, and `rr:joinCondition` " + "(child/parent) gives the exact columns to join on. Use it to turn " + "ontology terms into precise SQL identifiers and joins.\n\n" + f"{r2rml_text}" + ) + 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("## Ontology description") + lines.append("") + lines.append(ctx.resource_description) + lines.append("") + + lines.extend(_render_entity_schema_sections(ctx)) + + return "\n".join(lines) + + +def build( + entities: list[Entity], + resource_description: str = "", + base_system_prompt: str = "", + ontology_text: str = "", + r2rml_text: str = "", + prompt_version: str | None = None, +) -> str: + """Build the ontology-tool inner system prompt. + + Args: + entities: Resolved Data Fabric entities (from the R2RML allow-list). + resource_description: Optional description of the ontology context. + base_system_prompt: Optional system prompt from the outer agent. + ontology_text: The fetched ontology OWL content (authoritative schema). + r2rml_text: The fetched ontology R2RML mapping (ontology→table/column). + prompt_version: Optional SQL-strategy prompt version key. + + Returns: + Formatted prompt string for the inner LLM system message. + """ + if not entities: + return "" + + ctx = build_sql_context( + entities, + resource_description, + base_system_prompt, + prompt_version=prompt_version, + ) + return format_ontology_context( + ctx, ontology_text=ontology_text, r2rml_text=r2rml_text + ) diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_r2rml.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_r2rml.py new file mode 100644 index 000000000..8cdf00a8b --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_r2rml.py @@ -0,0 +1,87 @@ +"""Parse a Data Fabric R2RML mapping into ``(entity_name, folder_path)`` pairs. + +Uses ``rdflib`` to parse the mapping as an RDF graph rather than scanning the +text, so authoring variations (predicate order, whitespace, one-line vs +multi-line blocks, ``@base``/prefix differences) are handled by a real Turtle +parser instead of regexes. + +The mapping follows the UiPath authoring contract: each entity is one +``rr:TriplesMap`` that declares exactly one ``rr:tableName`` (via +``rr:logicalTable``) naming the Data Fabric entity, and exactly one +``uipath:folderPath`` naming the folder that entity lives in. Because each pair +is read from a single TriplesMap subject, a table is always paired with the +folder from its own map (never mis-zipped across maps). + +The extracted ``(entity_name, folder_path)`` set is the closed allow-list of +entities the ontology tool is permitted to resolve and query. +""" + +from rdflib import Graph, Namespace +from rdflib.namespace import RDF + +# R2RML vocabulary and the UiPath extension carrying the folder path. +# These are RDF namespace IRIs — opaque identifiers that must match the terms in +# the parsed Turtle verbatim, not network endpoints. The R2RML namespace is +# defined as http by the W3C spec; switching to https would stop the terms from +# matching. NOSONAR suppresses the clear-text-protocol false positive (S5332). +RR = Namespace("http://www.w3.org/ns/r2rml#") # NOSONAR +UIPATH = Namespace("http://uipath.com/ns/datafabric#") # NOSONAR + + +class R2RMLParseError(ValueError): + """Raised when the R2RML mapping does not satisfy the authoring contract.""" + + +def parse_r2rml_entities(r2rml_text: str) -> list[tuple[str, str]]: + """Extract the ``(entity_name, folder_path)`` allow-list from an R2RML mapping. + + Parses the mapping as Turtle and, for each ``rr:TriplesMap``, reads the + single ``rr:tableName`` (under ``rr:logicalTable``) and the single + ``uipath:folderPath`` literal. Pairs are de-duplicated and returned in a + deterministic (sorted) order. + + Args: + r2rml_text: The R2RML mapping document (Turtle). + + Returns: + Sorted, de-duplicated ``(entity_name, folder_path)`` pairs. + + Raises: + R2RMLParseError: The document is not valid Turtle, contains no + ``rr:TriplesMap``, or a TriplesMap does not declare exactly one + ``rr:tableName`` and one ``uipath:folderPath``. + """ + graph = Graph() + try: + graph.parse(data=r2rml_text, format="turtle") + except Exception as e: # rdflib raises assorted parser errors + raise R2RMLParseError(f"R2RML mapping is not valid Turtle: {e}") from e + + triples_maps = list(graph.subjects(RDF.type, RR.TriplesMap)) + if not triples_maps: + raise R2RMLParseError("No `rr:TriplesMap` found in the R2RML mapping.") + + pairs: list[tuple[str, str]] = [] + seen: set[tuple[str, str]] = set() + + for triples_map in triples_maps: + table_names = [ + str(name) + for logical_table in graph.objects(triples_map, RR.logicalTable) + for name in graph.objects(logical_table, RR.tableName) + ] + folder_paths = [ + str(folder) for folder in graph.objects(triples_map, UIPATH.folderPath) + ] + if len(table_names) != 1 or len(folder_paths) != 1: + raise R2RMLParseError( + "Each TriplesMap must declare exactly one rr:tableName (via " + "rr:logicalTable) and one uipath:folderPath; found " + f"tableName={table_names} folderPath={folder_paths}." + ) + pair = (table_names[0], folder_paths[0]) + if pair not in seen: + seen.add(pair) + pairs.append(pair) + + return sorted(pairs) diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_subgraph.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_subgraph.py new file mode 100644 index 000000000..ed905f363 --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_subgraph.py @@ -0,0 +1,233 @@ +"""Inner LangGraph sub-graph for the Data Fabric **ontology** tool. + +A duplicate of the entity tool's ``datafabric_subgraph`` that is *prompt-agnostic*: +it accepts a ready-made ``system_prompt`` string rather than building one from +entities/resource_description internally. This lets the ontology tool ground the +inner LLM on the OWL + R2RML (assembled by ``ontology_prompt_builder``) without +modifying the shared entity-tool sub-graph. + +Behaviourally identical to the entity sub-graph: a self-contained ReAct loop that +translates natural-language questions into SQL, executes them via ``execute_sql``, +retries on errors, and short-circuits straight to END on a successful execution so +the outer agent receives the raw tool result. +""" + +import asyncio +import logging +from typing import Annotated, Any + +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import ( + AIMessage, + AnyMessage, + SystemMessage, + ToolCall, + ToolMessage, +) +from langchain_core.tools import BaseTool +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from langgraph.graph.state import CompiledStateGraph +from pydantic import BaseModel +from uipath.platform.entities import EntitiesService, Entity + +from ...datafabric_query_tool import DataFabricQueryTool +from ..models import DataFabricExecuteSqlInput + +logger = logging.getLogger(__name__) + + +class DataFabricSubgraphState(BaseModel): + """State for the inner Data Fabric ReAct sub-graph.""" + + messages: Annotated[list[AnyMessage], add_messages] = [] + iteration_count: int = 0 + last_tool_success: bool = False + + +class QueryExecutor: + """Executes SQL queries against Data Fabric.""" + + 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._entities.query_entity_records_async( + sql_query=sql_query, + ) + return { + "records": records, + "total_count": len(records), + "sql_query": sql_query, + } + except Exception as e: + logger.error("SQL query failed: %s", e) + return { + "records": [], + "total_count": 0, + "error": str(e), + "sql_query": sql_query, + } + + +class DataFabricGraph: + """Inner ReAct sub-graph for Data Fabric SQL execution (ontology tool). + + Each graph node is a method. The graph is compiled during __init__ + and available via the ``compiled`` property. Unlike the entity tool's + sub-graph, the system prompt is supplied ready-made by the caller. + """ + + def __init__( + self, + llm: BaseChatModel, + entities: list[Entity], + entities_service: EntitiesService, + system_prompt: str, + max_iterations: int = 25, + ) -> None: + self._max_iterations = max_iterations + self._execute_sql_tool = self._create_execute_sql_tool( + entities_service, entities + ) + self._system_message = SystemMessage(content=system_prompt) + self._inner_llm = llm.model_copy(update={"disable_streaming": True}).bind_tools( + [self._execute_sql_tool] + ) + + # Build and compile the graph + graph = StateGraph(DataFabricSubgraphState) + graph.add_node("inner_llm", self.llm_node) + graph.add_node("inner_tool", self.tool_node) + graph.add_node("termination", self.termination_node) + graph.add_edge(START, "inner_llm") + graph.add_conditional_edges( + "inner_llm", self.router, ["inner_tool", "termination", END] + ) + graph.add_conditional_edges("inner_tool", self.tool_router, ["inner_llm", END]) + graph.add_edge("termination", END) + self.compiled_graph: CompiledStateGraph[Any] = graph.compile() + + async def llm_node(self, state: DataFabricSubgraphState) -> dict[str, Any]: + """Invoke the inner LLM with the current message history.""" + messages = [self._system_message] + list(state.messages) + response = await self._inner_llm.ainvoke(messages) + return {"messages": [response]} + + async def tool_node(self, state: DataFabricSubgraphState) -> dict[str, Any]: + """Execute all tool calls from the last AIMessage concurrently.""" + last = state.messages[-1] + if not isinstance(last, AIMessage) or not last.tool_calls: + return {"iteration_count": state.iteration_count} + + results = await asyncio.gather( + *[self._execute_tool_call(tc) for tc in last.tool_calls] + ) + tool_messages = [msg for msg, _ in results] + all_succeeded = bool(results) and all(success for _, success in results) + return { + "messages": tool_messages, + "iteration_count": state.iteration_count + len(last.tool_calls), + "last_tool_success": all_succeeded, + } + + async def _execute_tool_call(self, tool_call: ToolCall) -> tuple[ToolMessage, bool]: + """Execute a single tool call and report whether it succeeded.""" + args = tool_call.get("args", {}) + try: + result = await self._execute_sql_tool.ainvoke(args) + except ValueError as e: + result = { + "records": [], + "total_count": 0, + "error": str(e), + "sql_query": args.get("sql_query", ""), + } + succeeded = ( + isinstance(result, dict) + and not result.get("error") + and result.get("total_count", 0) > 0 + ) + return ( + ToolMessage( + content=str(result), + tool_call_id=tool_call["id"], + name="execute_sql", + ), + succeeded, + ) + + async def termination_node(self, state: DataFabricSubgraphState) -> dict[str, Any]: + """Produce a clear message when max iterations is reached.""" + return { + "messages": [ + AIMessage( + content=( + "I was unable to resolve the query after " + f"{state.iteration_count} SQL attempts. " + "Please try rephrasing the question or narrowing the scope." + ) + ) + ] + } + + def router(self, state: DataFabricSubgraphState) -> str: + """Route from ``inner_llm`` to tool, termination, or END.""" + last = state.messages[-1] if state.messages else None + if isinstance(last, AIMessage) and last.tool_calls: + if state.iteration_count < self._max_iterations: + return "inner_tool" + return "termination" + return END + + def tool_router(self, state: DataFabricSubgraphState) -> str: + """Route from ``inner_tool``: short-circuit on success, retry on error. + + Skips the redundant LLM call that would otherwise reformat a + successful SQL result into prose — the outer agent receives the + raw tool output and produces the final natural-language answer. + Errors loop back to ``inner_llm`` so the retry path is preserved. + """ + if state.last_tool_success: + return END + return "inner_llm" + + def _create_execute_sql_tool( + self, + entities_service: EntitiesService, + entities: list[Entity], + ) -> BaseTool: + """Create the inner ``execute_sql`` tool.""" + entity_names = ", ".join(e.name for e in entities) + return DataFabricQueryTool( + name="execute_sql", + description=( + f"Execute a SQL SELECT query against Data Fabric entities: {entity_names}. " + "Refer to the entity schemas in the system message for available " + "tables and columns. Retry with a corrected query on errors." + ), + args_schema=DataFabricExecuteSqlInput, + coroutine=QueryExecutor(entities_service), + metadata={"tool_type": "datafabric_sql"}, + ) + + @staticmethod + def create( + llm: BaseChatModel, + entities: list[Entity], + entities_service: EntitiesService, + system_prompt: str, + max_iterations: int = 25, + ) -> CompiledStateGraph[Any]: + """Create and return a compiled Data Fabric ontology sub-graph.""" + graph = DataFabricGraph( + llm, + entities, + entities_service, + system_prompt, + max_iterations, + ) + return graph.compiled_graph diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_tool.py b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_tool.py new file mode 100644 index 000000000..ffcee205a --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_tool.py @@ -0,0 +1,323 @@ +"""Standalone Data Fabric Ontology tool. + +The agent selects *ontologies* (not entities). On first invocation this tool: + +1. fetches each ontology's **R2RML** (critical — it is the entity allow-list) and + **OWL** (optional — grounds the prompt), +2. parses R2RML into the closed ``(entity_name, folder_path)`` allow-list, +3. resolves each entity to its schema and builds a folder-scoped + ``EntitiesService`` (see :func:`resolve_ontology_entities`), +4. compiles the ontology sub-graph (:class:`ontology_subgraph.DataFabricGraph`), + grounded on both the OWL and the R2RML. + +Everything from the sub-graph down (``execute_sql`` → ``query_entity_records_async``) +mirrors the entity tool. Ontology names/folders come only from the agent +definition; the LLM never chooses which entity or folder to reach. + +This package duplicates the sub-graph and prompt builder so the entity tool's +``datafabric_tool``/``datafabric_subgraph`` are left untouched by this feature. +The only symbol borrowed from the entity tool is the shared ``BASE_SYSTEM_PROMPT`` +agent-config key, so both tools read the outer prompt from the same place. +""" + +import asyncio +import logging +from typing import TYPE_CHECKING, Any + +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.tools import BaseTool +from langgraph.graph.state import CompiledStateGraph +from uipath.agent.models.agent import AgentContextResourceConfig +from uipath.platform.entities import EntitiesService, Entity + +from ...base_uipath_structured_tool import BaseUiPathStructuredTool +from ..datafabric_tool import BASE_SYSTEM_PROMPT +from ..models import DataFabricQueryInput +from .ontology_fetcher import fence_ontology_block, fetch_ontology_file +from .ontology_r2rml import parse_r2rml_entities + +if TYPE_CHECKING: + from uipath.platform import UiPath + +logger = logging.getLogger(__name__) + +# Feature flag gating the whole ontology tool. Owned here (not in the entity +# tool's module) so this feature is fully self-contained; ``context_tool`` and +# the handler's defense-in-depth re-check both import it from this package. +DATAFABRIC_ONTOLOGY_FF = "DataFabricOntologyEnabled" + + +async def resolve_ontology_entities( + sdk: "UiPath", + pairs: list[tuple[str, str]], +) -> tuple[list[Entity], EntitiesService]: + """Resolve an R2RML allow-list into schemas + a folder-scoped service (b2). + + Resolves the folder paths to keys and fetches the entity schemas + *concurrently* (via ``asyncio.gather``) rather than serially, so the inner + graph is ready after roughly one round-trip instead of one per entity — + the main latency win on first invocation. Distinct folder paths are resolved + once each; entity schemas are then fetched all at once, preserving the input + order. + + Folder keys come only from the trusted folder resolution of the R2RML + ``uipath:folderPath`` — never from the LLM. + """ + # Phase 1: resolve each *distinct* folder path -> key, concurrently. + distinct_paths = list(dict.fromkeys(folder_path for _, folder_path in pairs)) + resolved_keys = await asyncio.gather( + *(sdk.folders.retrieve_key_async(folder_path=p) for p in distinct_paths) + ) + folder_key_by_path: dict[str, str] = {} + for folder_path, folder_key in zip(distinct_paths, resolved_keys, strict=True): + if not folder_key: + raise ValueError( + f"Folder path '{folder_path}' could not be resolved to a " + "folder key. Check the R2RML uipath:folderPath and permissions." + ) + folder_key_by_path[folder_path] = folder_key + + # Phase 2: fetch every entity schema concurrently (order preserved by gather). + folder_keys = [folder_key_by_path[folder_path] for _, folder_path in pairs] + entities: list[Entity] = list( + await asyncio.gather( + *( + sdk.entities.retrieve_by_name_async(name, folder_key=folder_key) + for (name, _), folder_key in zip(pairs, folder_keys, strict=True) + ) + ) + ) + # The SDK routes each query by entity *name* -> folder (EntityRouting), and an + # entity name is unique within its folder — which is also the SQL table + # identifier used inside a single query. The same name living in two different + # folders is a valid catalog-wide situation; each entity is still resolved and + # shown to the model with its own folder here. + folders_map: dict[str, str] = { + entity.name: folder_key + for entity, folder_key in zip(entities, folder_keys, strict=True) + } + + # Build the folder-scoped service ourselves via the public folders_map param + # (yields a FoldersMapRoutingStrategy). Reuse the SDK's already-resolved + # auth/config so this service talks to the same tenant as sdk.entities. + # NOTE: config/execution_context are read from the SDK's private attributes + # because there is currently no public factory for a folder-scoped + # EntitiesService with a folders_map. This is a deliberate, contained + # trade-off; if the SDK later exposes such a factory, switch to it here. + scoped_service = EntitiesService( + config=sdk._config, + execution_context=sdk._execution_context, + folders_service=sdk.folders, + folders_map=folders_map, + ) + return entities, scoped_service + + +class DataFabricOntologyQueryHandler: + """Lazy-init + invoke of the ontology-grounded Data Fabric sub-graph. + + On first call, fetches OWL + R2RML, derives the entity allow-list from R2RML, + resolves schemas + a folder-scoped service, and compiles the inner sub-graph. + Subsequent calls reuse the cached graph. + """ + + def __init__( + self, + ontologies: list[tuple[str, str | None]], + llm: BaseChatModel, + resource_description: str = "", + base_system_prompt: str = "", + ) -> None: + self._ontologies = ontologies + self._llm = llm + self._resource_description = resource_description + self._base_system_prompt = base_system_prompt + self._compiled: CompiledStateGraph[Any] | None = None + self._init_lock = asyncio.Lock() + + async def _ensure_graph(self) -> CompiledStateGraph[Any]: + """Lazy-init: fetch + parse + resolve + compile on first call. + + Uses ``asyncio.Lock`` because the outer agent supports parallel tool + calls — two concurrent invocations could otherwise race on first call. + """ + if self._compiled is not None: + return self._compiled + + async with self._init_lock: + if self._compiled is not None: + return self._compiled + + from uipath.core.feature_flags import FeatureFlags + from uipath.platform import UiPath + + from . import ontology_prompt_builder + from .ontology_subgraph import DataFabricGraph + + # Defense in depth: the tool is only created when the flag is on + # (the gate in context_tool). Re-check here — the last point before + # any ontology fetch/parse/resolution — so the feature can never do + # work with the flag disabled, however the tool was constructed. + if not FeatureFlags.is_flag_enabled(DATAFABRIC_ONTOLOGY_FF, default=False): + raise ValueError( + "Data Fabric ontology tool invoked while feature flag " + f"'{DATAFABRIC_ONTOLOGY_FF}' is disabled." + ) + + sdk = UiPath() + + owl_blocks: list[str] = [] + r2rml_blocks: list[str] = [] + pairs: list[tuple[str, str]] = [] + seen: set[tuple[str, str]] = set() + + for name, folder_key in self._ontologies: + # R2RML is critical (it is the entity allow-list) — fetch failure + # is fatal, not degraded. + r2rml_content, r2rml_media = await fetch_ontology_file( + sdk.entities, name, "r2rml", folder_key + ) + r2rml_blocks.append( + fence_ontology_block(name, "r2rml", r2rml_content, r2rml_media) + ) + for pair in parse_r2rml_entities(r2rml_content): + if pair not in seen: + seen.add(pair) + pairs.append(pair) + + # OWL only grounds the prompt — degrade to a note on failure. + try: + owl_content, owl_media = await fetch_ontology_file( + sdk.entities, name, "owl", folder_key + ) + owl_blocks.append( + fence_ontology_block(name, "owl", owl_content, owl_media) + ) + except Exception as e: + logger.warning("Ontology OWL fetch failed for %r: %s", name, e) + owl_blocks.append( + f"Ontology '{name}' OWL is unavailable ({type(e).__name__}); " + "rely on the R2RML mapping and entity schemas below." + ) + + if not pairs: + raise ValueError("Ontology R2RML declared no entities to query.") + + entities, scoped_service = await resolve_ontology_entities(sdk, pairs) + if not entities: + raise ValueError( + "No Data Fabric entity schemas could be resolved from the " + "ontology. Check the R2RML table names, folder paths, and " + "permissions." + ) + + system_prompt = ontology_prompt_builder.build( + entities, + resource_description=self._resource_description, + base_system_prompt=self._base_system_prompt, + ontology_text="\n\n".join(owl_blocks), + r2rml_text="\n\n".join(r2rml_blocks), + ) + self._compiled = DataFabricGraph.create( + llm=self._llm, + entities=entities, + entities_service=scoped_service, + system_prompt=system_prompt, + ) + return self._compiled + + async def __call__(self, user_query: str) -> str: + logger.debug("query_datafabric_ontology called with: %s", user_query) + + compiled_graph = await self._ensure_graph() + result_state = await compiled_graph.ainvoke( + {"messages": [HumanMessage(content=user_query)]} + ) + messages = result_state["messages"] + last_message = messages[-1] if messages else None + + # Happy path: the sub-graph short-circuits at END after a successful + # execute_sql, leaving a trailing batch of ToolMessages. Collapse the + # trailing batch into one synthetic message so the outer agent sees the + # full result set. + if isinstance(last_message, ToolMessage): + trailing_tool_messages: list[ToolMessage] = [] + for msg in reversed(messages): + if not isinstance(msg, ToolMessage): + break + trailing_tool_messages.append(msg) + return self._format_terminal_tool_messages( + list(reversed(trailing_tool_messages)) + ) + + for msg in reversed(messages): + if isinstance(msg, AIMessage) and msg.content: + return str(msg.content) + + return "Unable to generate an answer from the available data." + + @staticmethod + def _format_terminal_tool_messages(tool_messages: list[ToolMessage]) -> str: + """Build one returned message from the terminal ToolMessage batch.""" + non_empty_contents = [ + str(msg.content) for msg in tool_messages if getattr(msg, "content", None) + ] + if not non_empty_contents: + return "Unable to generate an answer from the available data." + if len(non_empty_contents) == 1: + return non_empty_contents[0] + + rendered_results = [ + f"Result {index}:\n{content}" + for index, content in enumerate(non_empty_contents, start=1) + ] + return ( + "Multiple SQL queries executed successfully. " + "Use all of the following results to answer the user's question.\n\n" + + "\n\n".join(rendered_results) + ) + + +def create_datafabric_ontology_tool( + resource: AgentContextResourceConfig, + llm: BaseChatModel, + tool_name: str = "query_datafabric_ontology", + agent_config: dict[str, str] | None = None, +) -> BaseTool: + """Create the standalone ontology tool from a ``datafabricontology`` context. + + Args: + resource: The Data Fabric ontology context resource configuration; its + ``ontology_set`` supplies the ``(name, folder_key)`` pairs. + llm: The language model for the inner SQL generation loop. + tool_name: Tool name surfaced to the outer agent. + agent_config: Optional dict; key ``base_system_prompt`` carries the outer + agent's system prompt. + """ + config = agent_config or {} + ontologies: list[tuple[str, str | None]] = [ + (item.name, item.folder_key) for item in (resource.ontology_set or []) + ] + handler = DataFabricOntologyQueryHandler( + ontologies=ontologies, + llm=llm, + resource_description=resource.description or "", + base_system_prompt=config.get(BASE_SYSTEM_PROMPT, ""), + ) + ontology_summary = "\n".join(f"- {name}" for name, _ in ontologies) + + return BaseUiPathStructuredTool( + name=tool_name, + description=( + "Query UiPath Data Fabric using the following ontologies with natural " + "language:\n" + f"{ontology_summary}\n" + "Describe the data you need; the tool grounds on each ontology and its " + "mapping, translates your request to SQL, executes it, and returns a " + "natural language answer." + ), + args_schema=DataFabricQueryInput, + coroutine=handler, + metadata={"tool_type": "datafabric_ontology_sql"}, + ) diff --git a/tests/agent/tools/test_datafabric_ontology_prompt_builder.py b/tests/agent/tools/test_datafabric_ontology_prompt_builder.py new file mode 100644 index 000000000..498fd79a1 --- /dev/null +++ b/tests/agent/tools/test_datafabric_ontology_prompt_builder.py @@ -0,0 +1,101 @@ +"""Tests for the ontology-tool inner prompt builder. + +Verifies the ontology prompt embeds OWL + R2RML (ontology-tool-only) and still +renders the shared entity-schema tables — separate from the entity tool's prompt. +""" + +from types import SimpleNamespace + +from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_prompt_builder import ( # noqa: E501 + build, +) + + +def _fake_field(**overrides): + return SimpleNamespace( + name="status", + display_name="Status", + sql_type=SimpleNamespace(name="varchar"), + description="The canonical workflow status", + allowed_values=["Open", "Closed"], + examples=["Open"], + good_for_aggregation=False, + good_for_grouping=True, + good_for_filtering=True, + is_foreign_key=False, + is_required=False, + is_unique=False, + is_hidden_field=False, + is_system_field=False, + **overrides, + ) + + +def _fake_entity(*fields, **overrides): + return SimpleNamespace( + id="entity-1", + name="Ticket", + display_name="Ticket", + description="Support tickets", + record_count=10, + fields=list(fields), + **overrides, + ) + + +def test_embeds_ontology_and_r2rml_sections_in_order(): + prompt = build( + [_fake_entity(_fake_field())], + ontology_text="OWL_BODY_XYZ", + r2rml_text="R2RML_BODY_XYZ", + ) + + assert "## Available Ontology" in prompt + assert "OWL_BODY_XYZ" in prompt + assert "## Ontology→Table Mapping (R2RML)" in prompt + assert "R2RML_BODY_XYZ" in prompt + # OWL section precedes the R2RML section. + assert prompt.index("Available Ontology") < prompt.index("Ontology→Table Mapping") + + +def test_still_renders_entity_schema_tables(): + prompt = build( + [_fake_entity(_fake_field())], + ontology_text="OWL", + r2rml_text="MAP", + ) + assert "## All available Data Fabric Entities" in prompt + assert "| status |" in prompt + assert "Ticket" in prompt + + +def test_omits_sections_when_empty(): + prompt = build([_fake_entity(_fake_field())]) + assert "## Available Ontology" not in prompt + assert "## Ontology→Table Mapping (R2RML)" not in prompt + + +def test_no_entities_returns_empty(): + assert build([], ontology_text="OWL", r2rml_text="MAP") == "" + + +def test_format_context_renders_agent_instructions_and_ontology_description(): + # `build()` folds resource_description into the strategy prompt (ctx.resource_description + # stays None), so exercise the Agent-Instructions + Ontology-description sections directly. + from uipath_langchain.agent.tools.datafabric_tool.models import SQLContext + from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_prompt_builder import ( + format_ontology_context, + ) + + ctx = SQLContext( + base_system_prompt="You are an agent.", + resource_description="Domain notes for the ontology.", + sql_expert_system_prompt="strategy", + constraints="constraints", + entity_contexts=[], + ) + prompt = format_ontology_context(ctx) + assert "## Agent Instructions" in prompt + assert "You are an agent." in prompt + assert "## Ontology description" in prompt + assert "Domain notes for the ontology." in prompt diff --git a/tests/agent/tools/test_datafabric_ontology_subgraph.py b/tests/agent/tools/test_datafabric_ontology_subgraph.py new file mode 100644 index 000000000..0f894b6e0 --- /dev/null +++ b/tests/agent/tools/test_datafabric_ontology_subgraph.py @@ -0,0 +1,182 @@ +"""Tests for the Data Fabric inner sub-graph (``DataFabricGraph``). + +The sub-graph is prompt-agnostic: the caller passes a pre-built ``system_prompt`` +(entity tool and ontology tool use separate builders) and the sub-graph only ever +binds ``execute_sql``. Covers the tool node's dispatch/terminal logic and compile. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.constants import END + +from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_subgraph import ( + DataFabricGraph, + DataFabricSubgraphState, + QueryExecutor, +) + + +@pytest.fixture +def entities_service(): + es = MagicMock() + es.query_entity_records_async = AsyncMock(return_value=[{"x": 1}]) + return es + + +@pytest.fixture +def make_graph(entities_service): + def _make(system_prompt="SYS"): + return DataFabricGraph( + llm=MagicMock(), + entities=[], + entities_service=entities_service, + system_prompt=system_prompt, + ) + + return _make + + +def _tc(name, args=None, cid="c1"): + return {"name": name, "args": args or {}, "id": cid, "type": "tool_call"} + + +def test_system_prompt_used_verbatim(entities_service): + graph = DataFabricGraph( + llm=MagicMock(), + entities=[], + entities_service=entities_service, + system_prompt="MY_PROMPT", + ) + assert graph._system_message.content == "MY_PROMPT" + + +async def test_execute_tool_call_sql_with_rows_is_terminal(make_graph): + graph = make_graph() + msg, ok = await graph._execute_tool_call( + _tc("execute_sql", {"sql_query": "SELECT 1"}) + ) + assert ok is True + + +async def test_execute_tool_call_sql_no_rows_not_terminal(make_graph, entities_service): + entities_service.query_entity_records_async = AsyncMock(return_value=[]) + graph = make_graph() + msg, ok = await graph._execute_tool_call( + _tc("execute_sql", {"sql_query": "SELECT 1"}) + ) + assert ok is False + + +async def test_execute_tool_call_sql_value_error_becomes_error_dict(make_graph): + # execute_sql raises ValueError on multiple statements; it must be caught and + # turned into an error result (non-terminal), not propagated. + graph = make_graph() + msg, ok = await graph._execute_tool_call( + _tc("execute_sql", {"sql_query": "SELECT 1; SELECT 2"}) + ) + assert ok is False + assert "error" in str(msg.content) + + +async def test_tool_node_terminal_on_sql_rows(make_graph): + graph = make_graph() + ai = AIMessage( + content="", tool_calls=[_tc("execute_sql", {"sql_query": "SELECT 1"}, "a")] + ) + out = await graph.tool_node(DataFabricSubgraphState(messages=[ai])) + assert out["last_tool_success"] is True + assert len(out["messages"]) == 1 + assert out["messages"][0].name == "execute_sql" + + +def test_create_returns_compiled_graph(entities_service): + compiled = DataFabricGraph.create( + llm=MagicMock(), + entities=[], + entities_service=entities_service, + system_prompt="SYS", + ) + assert hasattr(compiled, "ainvoke") + + +# --- QueryExecutor ---------------------------------------------------------- + + +async def test_query_executor_returns_records(entities_service): + out = await QueryExecutor(entities_service)("SELECT 1") + assert out["records"] == [{"x": 1}] + assert out["total_count"] == 1 + assert out["sql_query"] == "SELECT 1" + + +async def test_query_executor_wraps_exception(entities_service): + entities_service.query_entity_records_async = AsyncMock( + side_effect=RuntimeError("db down") + ) + out = await QueryExecutor(entities_service)("SELECT 1") + assert out["total_count"] == 0 + assert "db down" in out["error"] + + +# --- graph nodes ------------------------------------------------------------ + + +async def test_llm_node_invokes_inner_llm(make_graph): + graph = make_graph() + graph._inner_llm.ainvoke = AsyncMock(return_value=AIMessage(content="hi")) + out = await graph.llm_node( + DataFabricSubgraphState(messages=[HumanMessage(content="q")]) + ) + assert out["messages"][0].content == "hi" + + +async def test_tool_node_without_tool_calls_is_noop(make_graph): + graph = make_graph() + state = DataFabricSubgraphState( + messages=[AIMessage(content="final")], iteration_count=3 + ) + assert await graph.tool_node(state) == {"iteration_count": 3} + + +async def test_termination_node_reports_attempts(make_graph): + graph = make_graph() + out = await graph.termination_node(DataFabricSubgraphState(iteration_count=7)) + assert "7 SQL attempts" in out["messages"][0].content + + +def test_router_to_tool_when_calls_under_limit(make_graph): + graph = make_graph() + ai = AIMessage(content="", tool_calls=[_tc("execute_sql", cid="a")]) + state = DataFabricSubgraphState(messages=[ai], iteration_count=0) + assert graph.router(state) == "inner_tool" + + +def test_router_to_termination_at_limit(make_graph): + graph = make_graph() # default max_iterations=25 + ai = AIMessage(content="", tool_calls=[_tc("execute_sql", cid="a")]) + state = DataFabricSubgraphState(messages=[ai], iteration_count=25) + assert graph.router(state) == "termination" + + +def test_router_ends_without_tool_calls(make_graph): + graph = make_graph() + state = DataFabricSubgraphState(messages=[AIMessage(content="final")]) + assert graph.router(state) == END + + +def test_router_ends_without_messages(make_graph): + graph = make_graph() + assert graph.router(DataFabricSubgraphState(messages=[])) == END + + +def test_tool_router_ends_on_success(make_graph): + graph = make_graph() + assert graph.tool_router(DataFabricSubgraphState(last_tool_success=True)) == END + + +def test_tool_router_loops_on_failure(make_graph): + graph = make_graph() + state = DataFabricSubgraphState(last_tool_success=False) + assert graph.tool_router(state) == "inner_llm" diff --git a/tests/agent/tools/test_datafabric_ontology_tool.py b/tests/agent/tools/test_datafabric_ontology_tool.py new file mode 100644 index 000000000..e808324cc --- /dev/null +++ b/tests/agent/tools/test_datafabric_ontology_tool.py @@ -0,0 +1,266 @@ +"""Tests for the standalone ontology tool (datafabric_tool/ontology/ontology_tool.py). + +Covers the factory (reads ontology_set → handler) and the b2 resolver +(folderPath→key caching, name→Entity, folders_map construction). +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +import uipath.platform as _uipath_platform +from langchain_core.messages import AIMessage, ToolMessage +from uipath.agent.models.agent import AgentContextResourceConfig +from uipath.core.feature_flags import FeatureFlags + +from uipath_langchain.agent.tools.datafabric_tool.ontology import ( + ontology_prompt_builder as _opb, +) +from uipath_langchain.agent.tools.datafabric_tool.ontology import ( + ontology_subgraph as _sg, +) +from uipath_langchain.agent.tools.datafabric_tool.ontology import ( + ontology_tool as datafabric_ontology_tool, +) +from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_tool import ( + DataFabricOntologyQueryHandler, + create_datafabric_ontology_tool, + resolve_ontology_entities, +) + +# --- factory ---------------------------------------------------------------- + + +def _ontology_resource(ontology_set): + return AgentContextResourceConfig.model_validate( + { + "$resourceType": "context", + "name": "California Schools Ontology", + "description": "schools domain", + "contextType": "datafabricontology", + "ontologySet": ontology_set, + } + ) + + +def test_factory_builds_handler_from_ontology_set(): + resource = _ontology_resource( + [ + {"name": "california_schools", "folderId": "f1"}, + {"name": "finance", "folderId": "f2"}, + ] + ) + tool = create_datafabric_ontology_tool(resource, MagicMock()) + + assert tool.name == "query_datafabric_ontology" + assert tool.coroutine._ontologies == [ # type: ignore[attr-defined] + ("california_schools", "f1"), + ("finance", "f2"), + ] + assert "california_schools" in tool.description + assert "finance" in tool.description + + +def test_factory_empty_ontology_set(): + tool = create_datafabric_ontology_tool(_ontology_resource([]), MagicMock()) + assert tool.coroutine._ontologies == [] # type: ignore[attr-defined] + + +# --- resolver (b2) ---------------------------------------------------------- + + +async def test_resolve_builds_folders_map_and_caches_folder_key(monkeypatch): + folder_keys = {"F/a": "key-a", "F/b": "key-b"} + + sdk = MagicMock() + sdk.folders.retrieve_key_async = AsyncMock( + side_effect=lambda folder_path: folder_keys[folder_path] + ) + sdk.entities.retrieve_by_name_async = AsyncMock( + side_effect=lambda name, folder_key: SimpleNamespace(name=name) + ) + + captured: dict[str, object] = {} + fake_service = object() + + def fake_entities_service(**kwargs): + captured.update(kwargs) + return fake_service + + monkeypatch.setattr( + datafabric_ontology_tool, "EntitiesService", fake_entities_service + ) + + entities, service = await resolve_ontology_entities( + sdk, [("alpha", "F/a"), ("beta", "F/b"), ("gamma", "F/a")] + ) + + assert [e.name for e in entities] == ["alpha", "beta", "gamma"] + assert captured["folders_map"] == { + "alpha": "key-a", + "beta": "key-b", + "gamma": "key-a", + } + assert service is fake_service + # F/a resolved once (cached), F/b once → 2 folder lookups for 3 entities. + assert sdk.folders.retrieve_key_async.await_count == 2 + assert sdk.entities.retrieve_by_name_async.await_count == 3 + + +async def test_ensure_graph_guarded_when_flag_off(monkeypatch): + # Defense in depth: even if the tool is somehow constructed, the handler + # refuses to do any ontology work when the flag is disabled. + from uipath.core.feature_flags import FeatureFlags + + monkeypatch.setattr(FeatureFlags, "is_flag_enabled", lambda *a, **k: False) + handler = DataFabricOntologyQueryHandler( + ontologies=[("california-schools", "f1")], llm=MagicMock() + ) + with pytest.raises(ValueError, match="disabled"): + await handler._ensure_graph() + + +async def test_resolve_raises_on_unresolved_folder(monkeypatch): + sdk = MagicMock() + sdk.folders.retrieve_key_async = AsyncMock(return_value=None) + monkeypatch.setattr(datafabric_ontology_tool, "EntitiesService", MagicMock()) + + try: + await resolve_ontology_entities(sdk, [("alpha", "F/missing")]) + raise AssertionError("expected ValueError") + except ValueError as e: + assert "could not be resolved" in str(e) + + +# --- _ensure_graph (fetch -> parse -> resolve -> compile) ------------------- + + +def _patch_ensure_graph_deps(monkeypatch, *, fetch=None, parse=None, resolve=None): + """Patch the flag, SDK, and the fetch/parse/resolve/build/compile seam so + _ensure_graph can run without any real IO. Returns the sentinel graph.""" + sentinel_graph = object() + monkeypatch.setattr(FeatureFlags, "is_flag_enabled", lambda *a, **k: True) + monkeypatch.setattr(_uipath_platform, "UiPath", lambda: MagicMock()) + monkeypatch.setattr( + datafabric_ontology_tool, + "fetch_ontology_file", + fetch or AsyncMock(return_value=("BODY", "text/turtle")), + ) + monkeypatch.setattr( + datafabric_ontology_tool, + "parse_r2rml_entities", + parse or (lambda _txt: [("frpm", "F/a")]), + ) + monkeypatch.setattr( + datafabric_ontology_tool, + "resolve_ontology_entities", + resolve or AsyncMock(return_value=([SimpleNamespace(name="frpm")], object())), + ) + monkeypatch.setattr(_opb, "build", lambda *a, **k: "SYSTEM_PROMPT") + monkeypatch.setattr(_sg.DataFabricGraph, "create", lambda **k: sentinel_graph) + return sentinel_graph + + +async def test_ensure_graph_happy_path(monkeypatch): + sentinel = _patch_ensure_graph_deps(monkeypatch) + handler = DataFabricOntologyQueryHandler( + ontologies=[("california-schools", "f1")], llm=MagicMock() + ) + compiled = await handler._ensure_graph() + assert compiled is sentinel + # cached on second call (no re-work) + assert await handler._ensure_graph() is sentinel + + +async def test_ensure_graph_owl_failure_degrades(monkeypatch): + async def fetch(_es, _name, file_type, _fk): + if file_type == "owl": + raise RuntimeError("owl boom") + return ("R2RML", "text/turtle") + + sentinel = _patch_ensure_graph_deps(monkeypatch, fetch=fetch) + handler = DataFabricOntologyQueryHandler( + ontologies=[("california-schools", "f1")], llm=MagicMock() + ) + # R2RML is critical (present) so it still compiles; OWL failure degrades. + assert await handler._ensure_graph() is sentinel + + +async def test_ensure_graph_no_entities_raises(monkeypatch): + _patch_ensure_graph_deps(monkeypatch, parse=lambda _txt: []) + handler = DataFabricOntologyQueryHandler( + ontologies=[("california-schools", "f1")], llm=MagicMock() + ) + with pytest.raises(ValueError, match="declared no entities"): + await handler._ensure_graph() + + +async def test_ensure_graph_returns_cached_without_work(monkeypatch): + handler = DataFabricOntologyQueryHandler(ontologies=[], llm=MagicMock()) + handler._compiled = "CACHED" # type: ignore[assignment] + # Flag intentionally OFF: if it did any work it would raise; cache short-circuits. + monkeypatch.setattr(FeatureFlags, "is_flag_enabled", lambda *a, **k: False) + assert await handler._ensure_graph() == "CACHED" + + +# --- __call__ (invoke + terminal message handling) -------------------------- + + +def _handler_with_state(monkeypatch, messages): + handler = DataFabricOntologyQueryHandler(ontologies=[], llm=MagicMock()) + graph = MagicMock() + graph.ainvoke = AsyncMock(return_value={"messages": messages}) + monkeypatch.setattr(handler, "_ensure_graph", AsyncMock(return_value=graph)) + return handler + + +async def test_call_collapses_terminal_tool_messages(monkeypatch): + handler = _handler_with_state( + monkeypatch, + [AIMessage(content=""), ToolMessage(content="rows!", tool_call_id="c1")], + ) + assert await handler("q") == "rows!" + + +async def test_call_returns_ai_message_content(monkeypatch): + handler = _handler_with_state(monkeypatch, [AIMessage(content="the answer")]) + assert await handler("q") == "the answer" + + +async def test_call_fallback_when_no_answer(monkeypatch): + handler = _handler_with_state(monkeypatch, [AIMessage(content="")]) + assert await handler("q") == "Unable to generate an answer from the available data." + + +async def test_ensure_graph_empty_resolution_raises(monkeypatch): + # pairs is non-empty (parse) but resolution yields no entities -> post-resolve guard. + _patch_ensure_graph_deps( + monkeypatch, resolve=AsyncMock(return_value=([], object())) + ) + handler = DataFabricOntologyQueryHandler( + ontologies=[("california-schools", "f1")], llm=MagicMock() + ) + with pytest.raises(ValueError, match="could be resolved"): + await handler._ensure_graph() + + +# --- terminal-message formatting -------------------------------------------- + + +def test_format_terminal_collapses_multiple_results(): + out = DataFabricOntologyQueryHandler._format_terminal_tool_messages( + [ + ToolMessage(content="rows-a", tool_call_id="1"), + ToolMessage(content="rows-b", tool_call_id="2"), + ] + ) + assert "Multiple SQL queries" in out + assert "Result 1:\nrows-a" in out + assert "Result 2:\nrows-b" in out + + +def test_format_terminal_all_empty_returns_fallback(): + out = DataFabricOntologyQueryHandler._format_terminal_tool_messages( + [ToolMessage(content="", tool_call_id="1")] + ) + assert out == "Unable to generate an answer from the available data." diff --git a/tests/agent/tools/test_datafabric_prompt_builder.py b/tests/agent/tools/test_datafabric_prompt_builder.py index 560e049d1..22eba34d0 100644 --- a/tests/agent/tools/test_datafabric_prompt_builder.py +++ b/tests/agent/tools/test_datafabric_prompt_builder.py @@ -60,6 +60,14 @@ def test_build_includes_domain_guidance_in_rendered_prompt(): assert "Use business-friendly ticket language." in prompt +def test_entity_prompt_never_contains_ontology_sections(): + # The entity tool's prompt is ontology-free; ontology content lives only in + # datafabric_ontology_prompt_builder. + prompt = build([_fake_entity(_fake_field())]) + assert "## Available Ontology" not in prompt + assert "## Ontology→Table Mapping (R2RML)" not in prompt + + def _system_field(name, type_name="datetimeoffset", **overrides): """A fake auto-added system/audit field (Id, CreateTime, ...).""" return _fake_field( diff --git a/tests/agent/tools/test_ontology_fetcher.py b/tests/agent/tools/test_ontology_fetcher.py new file mode 100644 index 000000000..1d65ebae0 --- /dev/null +++ b/tests/agent/tools/test_ontology_fetcher.py @@ -0,0 +1,85 @@ +"""Tests for ontology file fetching (datafabric_tool/ontology_fetcher.py).""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from uipath_langchain.agent.tools.datafabric_tool.ontology import ontology_fetcher +from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_fetcher import ( + fence_ontology_block, + fetch_ontology_file, +) + + +def _entities_service(content: str = "BODY", media_type: str = "text/turtle"): + es = MagicMock() + es.get_ontology_file_async = AsyncMock( + return_value={"content": content, "mediaType": media_type} + ) + return es + + +# --- fetch_ontology_file ---------------------------------------------------- + + +async def test_fetch_returns_content_and_media_type(): + es = _entities_service(content="OWLBODY", media_type="application/owl-functional") + + content, media_type = await fetch_ontology_file(es, "library", "owl", "folder-1") + + assert content == "OWLBODY" + assert media_type == "application/owl-functional" + es.get_ontology_file_async.assert_awaited_once_with("library", "owl", "folder-1") + + +async def test_fetch_r2rml_file_type_is_passed_through(): + es = _entities_service(content="@prefix rr: <> .") + + await fetch_ontology_file(es, "library", "r2rml", None) + + es.get_ontology_file_async.assert_awaited_once_with("library", "r2rml", None) + + +async def test_fetch_raises_on_underlying_error(): + es = MagicMock() + es.get_ontology_file_async = AsyncMock(side_effect=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await fetch_ontology_file(es, "library", "r2rml", None) + + +async def test_fetch_raises_when_oversized(monkeypatch): + monkeypatch.setattr(ontology_fetcher, "_MAX_ONTOLOGY_BYTES", 5) + es = _entities_service(content="0123456789") # 10 bytes > cap + + with pytest.raises(ValueError, match="size limit"): + await fetch_ontology_file(es, "library", "owl", None) + + +async def test_fetch_missing_content_defaults_empty(): + es = MagicMock() + es.get_ontology_file_async = AsyncMock(return_value={}) + + content, media_type = await fetch_ontology_file(es, "library", "owl", None) + + assert content == "" + assert media_type == "" + + +# --- fence_ontology_block --------------------------------------------------- + + +def test_fence_includes_type_name_and_media_type(): + block = fence_ontology_block("library", "owl", "CONTENT", "text/turtle") + + assert block.startswith("--- OWL, text/turtle: library ---") + assert "CONTENT" in block + assert block.endswith("--- END OWL: library ---") + + +def test_fence_without_media_type(): + block = fence_ontology_block("library", "r2rml", "MAP") + + assert block.startswith("--- R2RML: library ---") + assert "MAP" in block + assert block.endswith("--- END R2RML: library ---") diff --git a/tests/agent/tools/test_ontology_r2rml.py b/tests/agent/tools/test_ontology_r2rml.py new file mode 100644 index 000000000..1a8d1d590 --- /dev/null +++ b/tests/agent/tools/test_ontology_r2rml.py @@ -0,0 +1,113 @@ +"""Tests for the rdflib R2RML parser (datafabric_tool/ontology/ontology_r2rml.py).""" + +import pytest + +from uipath_langchain.agent.tools.datafabric_tool.ontology.ontology_r2rml import ( + R2RMLParseError, + parse_r2rml_entities, +) + +# Every document must declare the prefixes it uses — rdflib is a real Turtle +# parser (unlike the previous regex scanner), so undeclared prefixes are errors. +_PREFIXES = ( + "@prefix rr: .\n" + "@prefix uipath: .\n" + "@prefix map: .\n" + "@prefix onto: .\n" +) + + +def _block(subject: str, table: str, folder: str, folder_first: bool = False) -> str: + table_line = f' rr:logicalTable [ rr:tableName "{table}" ] ;' + folder_line = f' uipath:folderPath "{folder}" ;' + body = ( + f"{folder_line}\n{table_line}" + if folder_first + else f"{table_line}\n{folder_line}" + ) + return ( + f"{subject} a rr:TriplesMap ;\n{body}\n rr:subjectMap [ rr:class onto:X ] .\n" + ) + + +def _doc(*blocks: str) -> str: + return _PREFIXES + "\n".join(blocks) + + +def test_single_triplesmap(): + doc = _doc(_block("map:TM_a", "alpha", "Shared/Fin")) + assert parse_r2rml_entities(doc) == [("alpha", "Shared/Fin")] + + +def test_multiple_triplesmaps_sorted_deterministically(): + # rdflib is graph-based (no document order); the parser returns a sorted, + # de-duplicated allow-list so output is deterministic across runs. + doc = _doc(_block("map:TM_b", "beta", "F/b"), _block("map:TM_a", "alpha", "F/a")) + assert parse_r2rml_entities(doc) == [("alpha", "F/a"), ("beta", "F/b")] + + +def test_pairing_is_per_triplesmap_not_global(): + # folderPath before tableName in one map, after in the other. Pairing is per + # TriplesMap subject, so table always pairs with the folder from its own map. + doc = _doc( + _block("map:TM_a", "alpha", "F/a", folder_first=True), + _block("map:TM_b", "beta", "F/b", folder_first=False), + ) + assert parse_r2rml_entities(doc) == [("alpha", "F/a"), ("beta", "F/b")] + + +def test_duplicate_pairs_are_deduped(): + doc = _doc(_block("map:TM_a", "alpha", "F/a"), _block("map:TM_a2", "alpha", "F/a")) + assert parse_r2rml_entities(doc) == [("alpha", "F/a")] + + +def test_same_name_different_folder_kept_distinct(): + doc = _doc(_block("map:TM_a", "shared", "F/a"), _block("map:TM_b", "shared", "F/b")) + assert parse_r2rml_entities(doc) == [("shared", "F/a"), ("shared", "F/b")] + + +def test_missing_folder_path_raises(): + doc = _doc( + 'map:TM_a a rr:TriplesMap ;\n rr:logicalTable [ rr:tableName "alpha" ] .\n' + ) + with pytest.raises(R2RMLParseError): + parse_r2rml_entities(doc) + + +def test_missing_table_name_raises(): + doc = _doc('map:TM_a a rr:TriplesMap ;\n uipath:folderPath "F/a" .\n') + with pytest.raises(R2RMLParseError): + parse_r2rml_entities(doc) + + +def test_no_triplesmap_raises(): + with pytest.raises(R2RMLParseError, match="No `rr:TriplesMap`"): + parse_r2rml_entities("@prefix rr: .\n") + + +def test_invalid_turtle_raises(): + with pytest.raises(R2RMLParseError, match="not valid Turtle"): + parse_r2rml_entities("this is not turtle @@@ <") + + +def test_compact_one_line_triplesmap_parses(): + # rdflib handles compact single-line syntax the regex parser would have + # struggled with — a table paired with its folder on one physical line. + doc = _doc( + 'map:TM_a a rr:TriplesMap ; rr:logicalTable [ rr:tableName "alpha" ] ; ' + 'uipath:folderPath "F/a" .' + ) + assert parse_r2rml_entities(doc) == [("alpha", "F/a")] + + +def test_realistic_three_entity_mapping(): + doc = _doc( + _block("map:TM_frpm", "frpm", "Shared/CaliforniaSchools"), + _block("map:TM_satscores", "satscores", "Shared/CaliforniaSchools"), + _block("map:TM_schools", "schools", "Shared/CaliforniaSchools"), + ) + assert parse_r2rml_entities(doc) == [ + ("frpm", "Shared/CaliforniaSchools"), + ("satscores", "Shared/CaliforniaSchools"), + ("schools", "Shared/CaliforniaSchools"), + ] diff --git a/uv.lock b/uv.lock index 52bf943dc..11c3874f6 100644 --- a/uv.lock +++ b/uv.lock @@ -3541,6 +3541,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/b8/a0e2790ae249d6f38c9f66de7a211621a7ab2650217bcd04e1262f578a56/pyopenssl-26.2.0-py3-none-any.whl", hash = "sha256:4f9d971bc5298b8bc1fab282803da04bf000c755d4ad9d99b52de2569ca19a70", size = 55823, upload-time = "2026-05-04T23:06:08.395Z" }, ] +[[package]] +name = "pyparsing" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, +] + [[package]] name = "pysignalr" version = "1.3.0" @@ -3769,6 +3778,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] +[[package]] +name = "rdflib" +version = "7.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyparsing" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/98/f5/18bb77b7af9526add0c727a3b2048959847dc5fb030913e2918bf384fec3/rdflib-7.6.0.tar.gz", hash = "sha256:6c831288d5e4a5a7ece85d0ccde9877d512a3d0f02d7c06455d00d6d0ea379df", size = 4943826, upload-time = "2026-02-13T07:15:55.938Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/c2/6604a71269e0c1bd75656d5a001432d16f2cc5b8c057140ec797155c295e/rdflib-7.6.0-py3-none-any.whl", hash = "sha256:30c0a3ebf4c0e09215f066be7246794b6492e054e782d7ac2a34c9f70a15e0dd", size = 615416, upload-time = "2026-02-13T07:15:46.487Z" }, +] + [[package]] name = "referencing" version = "0.37.0" @@ -4432,7 +4453,7 @@ wheels = [ [[package]] name = "uipath" -version = "2.13.0" +version = "2.13.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "applicationinsights" }, @@ -4456,9 +4477,9 @@ dependencies = [ { name = "uipath-platform" }, { name = "uipath-runtime" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/15/c01594e4458b0963379f3f13fe5472b7740480d2497deba3cbda2f5ef04f/uipath-2.13.0.tar.gz", hash = "sha256:7e8ae855d25e1ab1716aa1e2a1e8ba98118c3b3d671772ebb5869d03cae32139", size = 4493516, upload-time = "2026-07-06T13:44:42.007Z" } +sdist = { url = "https://files.pythonhosted.org/packages/43/7c/733468447ac90661799ec4d5ef3fdeb744e8ada0a642648f174f59e9c03c/uipath-2.13.2.tar.gz", hash = "sha256:7ad3b65b2fd2bf9b0d44b9174cc1e65767db795356de2d72e7dc11c34794ddcf", size = 4495808, upload-time = "2026-07-07T06:53:15.751Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b6/97/54bfcddb08aa1203be74fed1d02d8dd4aace433716979ecc31a152cecd4f/uipath-2.13.0-py3-none-any.whl", hash = "sha256:313e1a3049b22326ab5dfa04a5892c79651a275b9ddeec1f896c28607fe771ce", size = 417567, upload-time = "2026-07-06T13:44:40.12Z" }, + { url = "https://files.pythonhosted.org/packages/4c/b6/34e34ea457cd69b96a941be9c225da59b829167b7122b8ba49c6263a6409/uipath-2.13.2-py3-none-any.whl", hash = "sha256:ec23c7110ea465bae434ef8a42ba5609df0017ef3dc3f6198ef38da395d22d18", size = 417700, upload-time = "2026-07-07T06:53:13.938Z" }, ] [[package]] @@ -4495,6 +4516,7 @@ dependencies = [ { name = "pillow" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, + { name = "rdflib" }, { name = "uipath" }, { name = "uipath-core" }, { name = "uipath-langchain-client", extra = ["openai"] }, @@ -4554,7 +4576,8 @@ requires-dist = [ { name = "pillow", specifier = ">=12.1.1" }, { name = "pydantic-settings", specifier = ">=2.6.0" }, { name = "python-dotenv", specifier = ">=1.0.1" }, - { name = "uipath", specifier = ">=2.13.0,<2.14.0" }, + { name = "rdflib", specifier = ">=7.0.0,<8.0.0" }, + { name = "uipath", specifier = ">=2.13.2,<2.14.0" }, { name = "uipath-core", specifier = ">=0.5.28,<0.6.0" }, { name = "uipath-langchain-client", extras = ["all"], marker = "extra == 'all'", specifier = ">=1.16.1,<1.17.0" }, { name = "uipath-langchain-client", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=1.16.1,<1.17.0" }, @@ -4563,7 +4586,7 @@ requires-dist = [ { name = "uipath-langchain-client", extras = ["google"], marker = "extra == 'vertex'", specifier = ">=1.16.1,<1.17.0" }, { name = "uipath-langchain-client", extras = ["openai"], specifier = ">=1.16.1,<1.17.0" }, { name = "uipath-langchain-client", extras = ["vertexai"], marker = "extra == 'vertex'", specifier = ">=1.16.1,<1.17.0" }, - { name = "uipath-platform", specifier = ">=0.2.0,<0.3.0" }, + { name = "uipath-platform", specifier = ">=0.2.1,<0.3.0" }, { name = "uipath-runtime", specifier = ">=0.12.1,<0.13.0" }, ] provides-extras = ["anthropic", "vertex", "bedrock", "fireworks", "all"] @@ -4647,7 +4670,7 @@ wheels = [ [[package]] name = "uipath-platform" -version = "0.2.0" +version = "0.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, @@ -4657,9 +4680,9 @@ dependencies = [ { name = "truststore" }, { name = "uipath-core" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/22/2e/86bee6d9aa5d00ef95e729f3d37c78b3dbe8ebf83519f88357be5240ae76/uipath_platform-0.2.0.tar.gz", hash = "sha256:86bb1dbb4267afe43719276809bc7ebe7e9f9885ca8238da13a9776848636576", size = 412145, upload-time = "2026-07-06T13:42:39.658Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/a6/b83b09393f092bfeb05ec1a34b344395a9a998a7f8bc71983a3b6337f222/uipath_platform-0.2.1.tar.gz", hash = "sha256:c277dc41ec1e74fbbfacb671c81c02f3e95ef14b239acc7f5f05136c12c14bad", size = 413613, upload-time = "2026-07-07T06:52:14.461Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/17/94/70e858975f022c44fa3d539632eb22a1a038ae39b4b187376e678371941a/uipath_platform-0.2.0-py3-none-any.whl", hash = "sha256:378986cb161521301560e32c2ea93950e1f14e8905bc7e3f6e034f548f2c6ac0", size = 271644, upload-time = "2026-07-06T13:42:38.246Z" }, + { url = "https://files.pythonhosted.org/packages/ff/79/92d99edd5500f48d0bbc9e823235e03269753fde22bce1012daa9f4c9b8c/uipath_platform-0.2.1-py3-none-any.whl", hash = "sha256:def420e2b8edcd193432983392b21c8de6340613a7a0831c8abdf1e10bf476b8", size = 273278, upload-time = "2026-07-07T06:52:12.562Z" }, ] [[package]]