refactor(datafabric): inject ontology into system prompt, drop fetch_ontology tool

Author: sankalp-uipath

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

@@ -133,15 +133,15 @@ def build_sql_context(
     )
 
 
-def format_sql_context(ctx: SQLContext, ontology_names: list[str] | None = None) -> str:
+def format_sql_context(ctx: SQLContext, ontology_text: str = "") -> str:
     """Format a SQLContext as text for system prompt injection.
 
     Args:
         ctx: The built SQL context (entities, prompts, constraints).
-        ontology_names: Names of the ontologies attached to this agent. When
-            non-empty, an "Available Ontology" section is added telling the LLM
-            to call ``fetch_ontology`` before writing SQL — mirroring how the
-            entity set is surfaced below.
+        ontology_text: The fetched ontology OWL content. When non-empty, an
+            "Available Ontology" section embeds it as the authoritative schema
+            the LLM should ground its SQL on — mirroring how the entity set is
+            surfaced below.
     """
     lines: list[str] = []
 
@@ -151,19 +151,16 @@ def format_sql_context(ctx: SQLContext, ontology_names: list[str] | None = None)
         lines.append(ctx.base_system_prompt)
         lines.append("")
 
-    if ontology_names:
-        names = ", ".join(f"`{n}`" for n in ontology_names)
+    if ontology_text:
         lines.append(
             "## Available Ontology (authoritative semantic schema)\n\n"
-            f"This agent has a semantic ontology attached for these entities: "
-            f"{names}. It 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 below, which omits value formats and semantics.\n\n"
-            "**Before writing any SQL, call the `fetch_ontology` tool once** to "
-            "load it, then base your column names, filter values, and joins on "
-            "what it says. The entity tables below are a quick reference only; "
-            "the ontology is the source of truth when they disagree."
+            "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("")
 
@@ -220,7 +217,7 @@ def build(
     resource_description: str = "",
     base_system_prompt: str = "",
     prompt_version: str | None = None,
-    ontology_names: list[str] | None = None,
+    ontology_text: str = "",
 ) -> str:
     """Build the full SQL prompt text for the inner sub-graph LLM.
 
@@ -234,10 +231,9 @@ def build(
         base_system_prompt: Optional system prompt from the outer agent.
         prompt_version: Optional version key (e.g. ``"v0"``, ``"v1"``).
             Defaults to the registry's default.
-        ontology_names: Names of ontologies attached to this agent. When
-            non-empty, an "Available Ontology" section instructs the LLM to call
-            ``fetch_ontology`` before writing SQL. Pass these only when the
-            fetch_ontology tool is actually bound.
+        ontology_text: The fetched ontology OWL content. When non-empty, an
+            "Available Ontology" section embeds it so the LLM grounds its SQL on
+            the ontology. Empty string → no ontology section.
 
     Returns:
         Formatted prompt string for the inner LLM system message.
@@ -251,4 +247,4 @@ def build(
         base_system_prompt,
         prompt_version=prompt_version,
     )
-    return format_sql_context(ctx, ontology_names=ontology_names)
+    return format_sql_context(ctx, ontology_text=ontology_text)

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

@@ -29,14 +29,11 @@
 from langgraph.graph.message import add_messages
 from langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel
-from uipath.core.feature_flags import FeatureFlags
 from uipath.platform.entities import EntitiesService, Entity
 
 from ..datafabric_query_tool import DataFabricQueryTool
 from . import datafabric_prompt_builder
-from .datafabric_tool import DATAFABRIC_ONTOLOGY_FF
 from .models import DataFabricExecuteSqlInput
-from .ontology_fetch_tool import create_ontology_fetch_tool
 
 logger = logging.getLogger(__name__)
 
@@ -91,40 +88,25 @@ def __init__(
         max_iterations: int = 25,
         resource_description: str = "",
         base_system_prompt: str = "",
-        ontologies: list[tuple[str, str | None]] | None = None,
+        ontology_text: str = "",
     ) -> None:
         self._max_iterations = max_iterations
         self._execute_sql_tool = self._create_execute_sql_tool(
             entities_service, entities
         )
-        # Inner toolset: always execute_sql; optionally an LLM-decided
-        # fetch_ontology tool, added only when ontologies are configured AND the
-        # DataFabricOntologyEnabled feature flag is on. The flag decides which
-        # graph gets built — off → the original entities-only graph.
-        inner_tools: list[BaseTool] = [self._execute_sql_tool]
-        ontology_names: list[str] = []
-        if ontologies and FeatureFlags.is_flag_enabled(
-            DATAFABRIC_ONTOLOGY_FF, default=False
-        ):
-            inner_tools.append(
-                create_ontology_fetch_tool(entities_service, ontologies)
-            )
-            ontology_names = [name for name, _ in ontologies]
-        self._tools_by_name: dict[str, BaseTool] = {
-            tool.name: tool for tool in inner_tools
-        }
-        # Surface the ontology in the system prompt only when its fetch tool is
-        # actually bound — otherwise the LLM is told to call a tool it lacks.
+        # The ontology (when configured and enabled) is fetched deterministically
+        # upstream and embedded directly in the system prompt — the inner agent
+        # still has a single tool, execute_sql.
         self._system_message = SystemMessage(
             content=datafabric_prompt_builder.build(
                 entities,
                 resource_description,
                 base_system_prompt,
-                ontology_names=ontology_names,
+                ontology_text=ontology_text,
             )
         )
         self._inner_llm = llm.model_copy(update={"disable_streaming": True}).bind_tools(
-            inner_tools
+            [self._execute_sql_tool]
         )
 
         # Build and compile the graph
@@ -155,69 +137,36 @@ async def tool_node(self, state: DataFabricSubgraphState) -> dict[str, Any]:
         results = await asyncio.gather(
             *[self._execute_tool_call(tc) for tc in last.tool_calls]
         )
-        # End as soon as ANY tool call is a terminal success (a row-returning
-        # execute_sql). `any` not `all`: a non-terminal tool (e.g. fetch_ontology)
-        # co-issued in the same turn must not prevent a successful SQL from ending
-        # the loop.
-        any_succeeded = any(success for _, success in results)
-        # When short-circuiting to END, return ONLY the terminal-success
-        # ToolMessages so the outer agent's result is the query rows — not a
-        # co-issued fetch_ontology's OWL. On a non-terminal turn keep all messages
-        # so the inner LLM can use them on its next pass.
-        if any_succeeded:
-            tool_messages = [msg for msg, success in results if success]
-        else:
-            tool_messages = [msg for msg, _ in results]
+        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": any_succeeded,
+            "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 is a terminal success.
-
-        Dispatches by tool name so the sub-graph can host more than one tool
-        (e.g. ``execute_sql`` and ``fetch_ontology``). Only a successful
-        ``execute_sql`` that returned rows is terminal; every other tool
-        (including ontology fetch) reports ``False`` so the router loops back to
-        the inner LLM, letting it use the result to write or refine SQL.
-        """
-        name = tool_call.get("name", "")
+        """Execute a single tool call and report whether it succeeded."""
         args = tool_call.get("args", {})
-        tool = self._tools_by_name.get(name)
-        if tool is None:
-            return (
-                ToolMessage(
-                    content=f"Unknown tool: {name}",
-                    tool_call_id=tool_call["id"],
-                    name=name,
-                ),
-                False,
-            )
         try:
-            result = await tool.ainvoke(args)
+            result = await self._execute_sql_tool.ainvoke(args)
         except ValueError as e:
-            if name == self._execute_sql_tool.name:
-                result = {
-                    "records": [],
-                    "total_count": 0,
-                    "error": str(e),
-                    "sql_query": args.get("sql_query", ""),
-                }
-            else:
-                result = f"Tool '{name}' failed: {e}"
+            result = {
+                "records": [],
+                "total_count": 0,
+                "error": str(e),
+                "sql_query": args.get("sql_query", ""),
+            }
         succeeded = (
-            name == self._execute_sql_tool.name
-            and isinstance(result, dict)
+            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=name,
+                name="execute_sql",
             ),
             succeeded,
         )
@@ -284,7 +233,7 @@ def create(
         max_iterations: int = 25,
         resource_description: str = "",
         base_system_prompt: str = "",
-        ontologies: list[tuple[str, str | None]] | None = None,
+        ontology_text: str = "",
     ) -> CompiledStateGraph[Any]:
         """Create and return a compiled Data Fabric sub-graph."""
         graph = DataFabricGraph(
@@ -294,6 +243,6 @@ def create(
             max_iterations,
             resource_description,
             base_system_prompt,
-            ontologies,
+            ontology_text,
         )
         return graph.compiled_graph

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

@@ -95,9 +95,11 @@ async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]:
             if self._compiled is not None:
                 return self._compiled
 
+            from uipath.core.feature_flags import FeatureFlags
             from uipath.platform import UiPath
 
             from .datafabric_subgraph import DataFabricGraph
+            from .ontology_fetcher import fetch_ontology_text
 
             sdk = UiPath()
             resolution = await sdk.entities.resolve_entity_set_async(self._entity_set)
@@ -106,13 +108,23 @@ async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]:
                     "No Data Fabric entity schemas could be fetched. "
                     "Check entity identifiers and permissions."
                 )
+            # Deterministically fetch the ontology (when configured AND the flag
+            # is on) and embed it in the inner system prompt — the LLM never has
+            # to decide to fetch it.
+            ontology_text = ""
+            if self._ontologies and FeatureFlags.is_flag_enabled(
+                DATAFABRIC_ONTOLOGY_FF, default=False
+            ):
+                ontology_text = await fetch_ontology_text(
+                    resolution.entities_service, self._ontologies
+                )
             self._compiled = DataFabricGraph.create(
                 llm=self._llm,
                 entities=resolution.entities,
                 entities_service=resolution.entities_service,
                 resource_description=self._resource_description,
                 base_system_prompt=self._base_system_prompt,
-                ontologies=self._ontologies,
+                ontology_text=ontology_text,
             )
             return self._compiled
 

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

@@ -94,12 +94,3 @@ class DataFabricExecuteSqlInput(BaseModel):
             "Use exact table and column names from the entity schemas."
         ),
     )
-
-
-class OntologyFetchInput(BaseModel):
-    """Input schema for the ontology fetch tool — intentionally empty.
-
-    The ontology name is pinned from configuration, never supplied by the
-    LLM, so the model cannot redirect the fetch to an arbitrary resource. The
-    tool simply triggers a fetch of the configured ontology.
-    """