diff --git a/docs/data_fabric_lineage.md b/docs/data_fabric_lineage.md new file mode 100644 index 000000000..7dd87e595 --- /dev/null +++ b/docs/data_fabric_lineage.md @@ -0,0 +1,41 @@ +# Data Fabric query evidence + +The Data Fabric natural-language query tool produces a versioned local evidence +artifact alongside its existing text result. The V1 artifact is inferred from +the entity metadata returned by `resolve_entity_set_async` and the records +returned by `query_entity_records_async`; it is not backend-authoritative +lineage. + +The text returned to the model and existing clients is unchanged. Supporting +runtimes can read `ToolMessage.artifact`, and the complete envelope is also +stored in `inner_state.tools_storage` under `(tool_name, execution_id)`. Trace +events contain only bounded identifier mappings, hashes, counts, and integrity +digests. They contain neither raw SQL, query results, credentials, nor error +details. + +Because source identifiers are authorization-scoped, cached source mappings and +joins are included only after the existing backend query call succeeds. A +failure-only execution keeps its record-free attempt ledger but does not release +the cached source graph; a successful zero-row result is still a valid success. + +Evidence IDs are valid only for the execution that created them. Consumers must +validate them against that execution's artifact and the source-reference IDs +authorized for the current response. Unknown major versions, invented IDs, +cross-execution IDs, and references not used by the execution are rejected. +Clients without Data Fabric citation support continue to return plain text. + +## Future backend boundary + +Backend-authoritative lineage is deliberately deferred. A future integration +must use the backend's actual response contract for: + +- backend execution IDs; +- normalized physical query plans; +- authoritative output-column lineage; +- source-system timestamps; and +- optional record-level evidence. + +No speculative Python SDK methods or response models are introduced here. When +those fields are available, the local envelope can consume them and change its +authority from inferred or partial to authoritative while preserving the V1 +answer-content boundary. diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py index 629f71f47..bf55668c6 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_subgraph.py @@ -13,7 +13,10 @@ import asyncio import logging -from typing import Annotated, Any +import time +from contextlib import contextmanager +from operator import add +from typing import Annotated, Any, Iterator from langchain_core.language_models import BaseChatModel from langchain_core.messages import ( @@ -30,12 +33,31 @@ from langgraph.graph.state import CompiledStateGraph from pydantic import BaseModel from uipath.platform.entities import EntitiesService, Entity +from uipath.platform.errors import DataFabricError, EnrichedException from ..datafabric_query_tool import DataFabricQueryTool from . import datafabric_prompt_builder -from .models import DataFabricExecuteSqlInput +from .lineage import ( + DataFabricSchemaSnapshot, + build_attempt, + build_schema_snapshot, + infer_column_lineage, + sha256_text, + utc_now, +) +from .models import ( + DataFabricColumnLineageV1, + DataFabricExecuteSqlInput, + DataFabricQueryAttemptV1, +) logger = logging.getLogger(__name__) +CATEGORY_MARKER = "(category: " + + +@contextmanager +def _noop_context() -> Iterator[None]: + yield None class DataFabricSubgraphState(BaseModel): @@ -44,36 +66,148 @@ class DataFabricSubgraphState(BaseModel): messages: Annotated[list[AnyMessage], add_messages] = [] iteration_count: int = 0 last_tool_success: bool = False + last_error_category: str = "" + last_error_detail: str = "" + execution_id: str = "" + attempts: Annotated[list[DataFabricQueryAttemptV1], add] = [] + column_lineage: Annotated[list[DataFabricColumnLineageV1], add] = [] + resolved_reference_count: Annotated[int, add] = 0 + total_reference_count: Annotated[int, add] = 0 + forced_partial: bool = False + + +class QueryExecutionResult(dict[str, Any]): + """Dict-compatible result with non-rendered error classification metadata.""" + + def __init__(self, *args: Any, error_category: str | None = None, **kwargs: Any): + super().__init__(*args, **kwargs) + self.error_category = error_category class QueryExecutor: """Executes SQL queries against Data Fabric.""" - def __init__(self, entities_service: EntitiesService) -> None: + def __init__( + self, entities_service: EntitiesService, entities: list[Entity] | None = None + ) -> None: self._entities = entities_service + resolved_entities = entities or [] + self._entity_attrs: dict[str, int] = { + "df.entity_count": len(resolved_entities), + "df.native_entity_count": sum( + 1 for entity in resolved_entities if not entity.external_fields + ), + "df.federated_entity_count": sum( + 1 for entity in resolved_entities if entity.external_fields + ), + } async def __call__(self, sql_query: str) -> dict[str, Any]: - logger.debug("execute_sql called with SQL: %s", sql_query) + query_hash = sha256_text(sql_query) + logger.debug("execute_sql called with query hash %s", query_hash) try: - # Relationship (FK) fields are typed as their scalar id so the SQL the - # agent writes can join on `relationshipField = Other.Id`. - records = await self._entities.query_entity_records_async( - sql_query=sql_query, - relationships_as_scalar=True, + from opentelemetry import trace as otel_trace + + tracer = otel_trace.get_tracer("uipath_langchain.datafabric") + except ImportError: + tracer = None + + span_context = ( + tracer.start_as_current_span( + "Data Fabric SQL query", + attributes={ + "openinference.span.kind": "TOOL", + "span_type": "datafabricQuery", + "uipath.custom_instrumentation": True, + "df.query_hash": query_hash, + **self._entity_attrs, + }, ) - 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, - } + if tracer + else _noop_context() + ) + + with span_context as span: + try: + records = await self._entities.query_entity_records_async( + sql_query=sql_query, + relationships_as_scalar=True, + ) + if span is not None: + span.set_attribute("df.row_count", len(records)) + span.set_attribute("df.success", True) + return QueryExecutionResult( + records=records, + total_count=len(records), + sql_query=sql_query, + ) + except Exception as error: + return self._handle_query_error(error, span, sql_query) + + def _handle_query_error( + self, error: Exception, span: Any, sql_query: str + ) -> QueryExecutionResult: + """Classify a failure without exposing raw errors in telemetry.""" + logger.error("SQL query failed (%s)", type(error).__name__) + data_fabric_error = ( + DataFabricError.from_enriched_exception(error) + if isinstance(error, EnrichedException) + else None + ) + category = ( + data_fabric_error.category.value + if data_fabric_error is not None + else type(error).__name__ + ) + detail = self._build_error_detail(error, data_fabric_error) + if span is not None: + self._record_error_span(span, error, data_fabric_error) + return QueryExecutionResult( + records=[], + total_count=0, + error=detail, + sql_query=sql_query, + error_category=category, + ) + + @staticmethod + def _record_error_span( + span: Any, error: Exception, data_fabric_error: DataFabricError | None + ) -> None: + """Record policy-safe failure attributes on an OTEL span.""" + category = ( + data_fabric_error.category.value + if data_fabric_error is not None + else type(error).__name__ + ) + span.set_attribute("df.success", False) + span.set_attribute("df.error.category", category) + span.set_attribute("df.error.type", type(error).__name__) + span.set_attribute("df.error.digest", sha256_text(str(error))) + if data_fabric_error is not None and data_fabric_error.code: + span.set_attribute("df.error.code", data_fabric_error.code) + + from opentelemetry.trace import Status, StatusCode + + span.set_status(Status(StatusCode.ERROR, "Data Fabric query failed")) + + @staticmethod + def _build_error_detail( + error: Exception, data_fabric_error: DataFabricError | None + ) -> str: + """Build the existing structured error detail for the inner LLM.""" + if data_fabric_error and data_fabric_error.code: + parts = [f"[{data_fabric_error.code}]"] + if data_fabric_error.category.value != "unknown": + parts.append(f"(category: {data_fabric_error.category.value})") + if data_fabric_error.message: + parts.append(data_fabric_error.message) + if data_fabric_error.is_retryable: + parts.append("— This error is transient, retry the same query.") + elif data_fabric_error.is_bad_sql: + parts.append("— Fix the SQL syntax and retry.") + return " ".join(parts) + return str(error) class DataFabricGraph: @@ -91,8 +225,10 @@ def __init__( max_iterations: int = 25, resource_description: str = "", base_system_prompt: str = "", + schema_snapshot: DataFabricSchemaSnapshot | None = None, ) -> None: self._max_iterations = max_iterations + self._schema_snapshot = schema_snapshot or build_schema_snapshot(entities) self._execute_sql_tool = self._create_execute_sql_tool( entities_service, entities ) @@ -131,33 +267,91 @@ async def tool_node(self, state: DataFabricSubgraphState) -> dict[str, Any]: return {"iteration_count": state.iteration_count} results = await asyncio.gather( - *[self._execute_tool_call(tc) for tc in last.tool_calls] + *[ + self._execute_tool_call( + tc, + execution_id=state.execution_id, + ordinal=state.iteration_count + index, + ) + for index, tc in enumerate(last.tool_calls, start=1) + ] ) - tool_messages = [msg for msg, _ in results] - all_succeeded = bool(results) and all(success for _, success in results) + tool_messages = [item[0] for item in results] + all_succeeded = bool(results) and all(item[1] for item in results) + last_category = "" + last_detail = "" + for item in reversed(results): + if not item[1] and item[8]: + last_category = item[7] + last_detail = item[8] + break return { "messages": tool_messages, "iteration_count": state.iteration_count + len(last.tool_calls), "last_tool_success": all_succeeded, + "last_error_category": last_category or state.last_error_category, + "last_error_detail": last_detail or state.last_error_detail, + "attempts": [item[2] for item in results], + "column_lineage": [column for item in results for column in item[3]], + "resolved_reference_count": sum(item[4] for item in results), + "total_reference_count": sum(item[5] for item in results), + "forced_partial": state.forced_partial or any(item[6] for item in results), } - async def _execute_tool_call(self, tool_call: ToolCall) -> tuple[ToolMessage, bool]: + async def _execute_tool_call( + self, tool_call: ToolCall, *, execution_id: str = "", ordinal: int = 1 + ) -> tuple[ + ToolMessage, + bool, + DataFabricQueryAttemptV1, + tuple[DataFabricColumnLineageV1, ...], + int, + int, + bool, + str, + str, + ]: """Execute a single tool call and report whether it succeeded.""" args = tool_call.get("args", {}) + sql_query = str(args.get("sql_query", "")) + started_at = utc_now() + started = time.perf_counter() 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 + result = QueryExecutionResult( + records=[], + total_count=0, + error=str(e), + sql_query=sql_query, + error_category=type(e).__name__, + ) + completed_at = utc_now() + duration_ms = (time.perf_counter() - started) * 1000.0 + succeeded = isinstance(result, dict) and not result.get("error") + attempt = build_attempt( + execution_id=execution_id, + tool_call_id=str(tool_call.get("id") or ""), + ordinal=ordinal, + sql=sql_query, + started_at=started_at, + completed_at=completed_at, + duration_ms=duration_ms, + result=result, ) + if succeeded: + columns, partial, resolved_count, reference_count = infer_column_lineage( + sql_query, self._schema_snapshot + ) + else: + columns, partial, resolved_count, reference_count = (), False, 0, 0 + error_detail = result.get("error", "") if isinstance(result, dict) else "" + error_category = getattr(result, "error_category", None) or "" + if error_detail and not error_category and CATEGORY_MARKER in error_detail: + start = error_detail.index(CATEGORY_MARKER) + len(CATEGORY_MARKER) + end = error_detail.find(")", start) + if end != -1: + error_category = error_detail[start:end] return ( ToolMessage( content=str(result), @@ -165,21 +359,27 @@ async def _execute_tool_call(self, tool_call: ToolCall) -> tuple[ToolMessage, bo name="execute_sql", ), succeeded, + attempt, + columns, + resolved_count, + reference_count, + partial, + error_category, + error_detail, ) 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." - ) - ) - ] - } + parts = [ + "I was unable to resolve the query after " + f"{state.iteration_count} SQL attempts." + ] + if state.last_error_category: + parts.append(f"Last error category: {state.last_error_category}.") + if state.last_error_detail: + parts.append(f"Last error: {state.last_error_detail[:300]}") + parts.append("Please try rephrasing the question or narrowing the scope.") + return {"messages": [AIMessage(content=" ".join(parts))]} def router(self, state: DataFabricSubgraphState) -> str: """Route from ``inner_llm`` to tool, termination, or END.""" @@ -217,7 +417,7 @@ def _create_execute_sql_tool( "tables and columns. Retry with a corrected query on errors." ), args_schema=DataFabricExecuteSqlInput, - coroutine=QueryExecutor(entities_service), + coroutine=QueryExecutor(entities_service, entities), metadata={"tool_type": "datafabric_sql"}, ) @@ -229,6 +429,7 @@ def create( max_iterations: int = 25, resource_description: str = "", base_system_prompt: str = "", + schema_snapshot: DataFabricSchemaSnapshot | None = None, ) -> CompiledStateGraph[Any]: """Create and return a compiled Data Fabric sub-graph.""" graph = DataFabricGraph( @@ -238,5 +439,6 @@ def create( max_iterations, resource_description, base_system_prompt, + schema_snapshot, ) return graph.compiled_graph diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py index aab4e4cfc..61dc87c74 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/datafabric_tool.py @@ -13,23 +13,48 @@ import asyncio import logging +import uuid +from dataclasses import dataclass from typing import Any from langchain_core.language_models import BaseChatModel -from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage from langchain_core.tools import BaseTool from langgraph.graph.state import CompiledStateGraph +from langgraph.types import Command from uipath.agent.models.agent import AgentContextResourceConfig from uipath.platform.entities import DataFabricEntityItem +from uipath_langchain.agent.react.types import AgentGraphState + from ..base_uipath_structured_tool import BaseUiPathStructuredTool -from .models import DataFabricQueryInput +from ..tool_node import ToolWrapperMixin, ToolWrapperReturnType +from . import datafabric_prompt_builder +from .lineage import ( + DataFabricSchemaSnapshot, + build_evidence_envelope, + build_schema_snapshot, + emit_trace_summary, + sha256_text, +) +from .models import DataFabricEvidenceEnvelopeV1, DataFabricQueryInput +from .prompts.registry import DEFAULT_PROMPT_VERSION logger = logging.getLogger(__name__) BASE_SYSTEM_PROMPT = "base_system_prompt" +class DataFabricStructuredTool(BaseUiPathStructuredTool, ToolWrapperMixin): + """Data Fabric tool with an outer wrapper for evidence artifacts/state.""" + + +@dataclass(frozen=True) +class _DataFabricExecution: + content: str + evidence: DataFabricEvidenceEnvelopeV1 | None + + class DataFabricTextQueryHandler: """Manages lazy initialization and invocation of the Data Fabric sub-graph. @@ -50,6 +75,8 @@ def __init__( self._resource_description = resource_description self._base_system_prompt = base_system_prompt self._compiled: CompiledStateGraph[Any] | None = None + self._schema_snapshot: DataFabricSchemaSnapshot | None = None + self._prompt_hash: str | None = None self._init_lock = asyncio.Lock() async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]: @@ -76,21 +103,47 @@ async def _ensure_datafabric_graph(self) -> CompiledStateGraph[Any]: "No Data Fabric entity schemas could be fetched. " "Check entity identifiers and permissions." ) + self._schema_snapshot = build_schema_snapshot(resolution.entities) + prompt = datafabric_prompt_builder.build( + resolution.entities, + self._resource_description, + self._base_system_prompt, + ) + self._prompt_hash = sha256_text(prompt) 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, + schema_snapshot=self._schema_snapshot, ) return self._compiled async def __call__(self, user_query: str) -> str: - logger.debug("query_datafabric called with: %s", user_query) + execution = await self.execute_with_evidence(user_query) + return execution.content + + async def execute_with_evidence( + self, + user_query: str, + *, + execution_id: str | None = None, + tool_name: str = "query_datafabric", + ) -> _DataFabricExecution: + """Run one query and construct evidence without changing returned text.""" + + logger.debug( + "query_datafabric called with query hash %s", sha256_text(user_query) + ) compiled_graph = await self._ensure_datafabric_graph() + resolved_execution_id = execution_id or str(uuid.uuid4()) result_state = await compiled_graph.ainvoke( - {"messages": [HumanMessage(content=user_query)]} + { + "messages": [HumanMessage(content=user_query)], + "execution_id": resolved_execution_id, + } ) messages = result_state["messages"] last_message = messages[-1] if messages else None @@ -105,17 +158,57 @@ async def __call__(self, user_query: str) -> str: if not isinstance(msg, ToolMessage): break trailing_tool_messages.append(msg) - return self._format_terminal_tool_messages( + content = self._format_terminal_tool_messages( list(reversed(trailing_tool_messages)) ) + return _DataFabricExecution( + content=content, + evidence=self._build_evidence( + result_state, resolved_execution_id, tool_name + ), + ) # On errors / max-iterations the terminal message is an AIMessage # carrying the natural-language explanation. for msg in reversed(messages): if isinstance(msg, AIMessage) and msg.content: - return str(msg.content) + return _DataFabricExecution( + content=str(msg.content), + evidence=self._build_evidence( + result_state, resolved_execution_id, tool_name + ), + ) - return "Unable to generate an answer from the available data." + return _DataFabricExecution( + content="Unable to generate an answer from the available data.", + evidence=self._build_evidence( + result_state, resolved_execution_id, tool_name + ), + ) + + def _build_evidence( + self, result_state: dict[str, Any], execution_id: str, tool_name: str + ) -> DataFabricEvidenceEnvelopeV1 | None: + if self._schema_snapshot is None or self._prompt_hash is None: + return None + attempts = result_state.get("attempts", []) + if not attempts: + return None + envelope = build_evidence_envelope( + execution_id=execution_id, + tool_name=tool_name, + snapshot=self._schema_snapshot, + attempts=attempts, + column_lineage=result_state.get("column_lineage", []), + prompt_version=DEFAULT_PROMPT_VERSION, + prompt_hash=self._prompt_hash, + llm=self._llm, + resolved_reference_count=result_state.get("resolved_reference_count", 0), + total_reference_count=result_state.get("total_reference_count", 0), + forced_partial=result_state.get("forced_partial", False), + ) + emit_trace_summary(envelope) + return envelope @staticmethod def _format_terminal_tool_messages(tool_messages: list[ToolMessage]) -> str: @@ -173,7 +266,39 @@ def create_datafabric_query_tool( entity_lines.append(line) entity_summary = "\n".join(entity_lines) - return BaseUiPathStructuredTool( + async def datafabric_evidence_wrapper( + _tool: BaseTool, + call: ToolCall, + _state: AgentGraphState, + ) -> ToolWrapperReturnType: + execution = await handler.execute_with_evidence( + str(call.get("args", {}).get("user_query", "")), + execution_id=call["id"], + tool_name=call["name"], + ) + message = ToolMessage( + content=execution.content, + artifact=execution.evidence, + name=call["name"], + tool_call_id=call["id"], + ) + if execution.evidence is None: + return {"messages": [message]} + return Command( + update={ + "messages": [message], + "inner_state": { + "tools_storage": { + ( + call["name"], + execution.evidence.execution_id, + ): execution.evidence + } + }, + } + ) + + tool = DataFabricStructuredTool( name=tool_name, description=( "Query the following Data Fabric entities using natural language:\n" @@ -185,3 +310,5 @@ def create_datafabric_query_tool( coroutine=handler, metadata={"tool_type": "datafabric_sql"}, ) + tool.set_tool_wrappers(awrapper=datafabric_evidence_wrapper) + return tool diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/lineage.py b/src/uipath_langchain/agent/tools/datafabric_tool/lineage.py new file mode 100644 index 000000000..e5d762925 --- /dev/null +++ b/src/uipath_langchain/agent/tools/datafabric_tool/lineage.py @@ -0,0 +1,1050 @@ +"""Local, non-authoritative Data Fabric lineage and evidence construction. + +Only metadata returned by entity-set resolution and records returned by the +existing query API are used. This module never performs authorization or +backend discovery and deliberately falls back to partial entity-level lineage +when an identifier cannot be resolved without guessing. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError, version +from typing import Any, Iterable, Sequence + +import sqlparse +from pydantic import BaseModel +from sqlparse.sql import Identifier, IdentifierList +from sqlparse.tokens import DML, Keyword + +from .models import ( + DataFabricAnswerCitationV1, + DataFabricColumnLineageV1, + DataFabricEvidenceEnvelopeV1, + DataFabricGenerationProvenanceV1, + DataFabricQueryAttemptV1, + DataFabricResultIntegrityV1, + DataFabricSourceJoinV1, + DataFabricSourceRefV1, +) + +_IDENTIFIER = r"[A-Za-z_][A-Za-z0-9_]*" +_DIRECT_PROJECTION_RE = re.compile( + rf"^\s*(?:(?P{_IDENTIFIER})\s*\.\s*)?" + rf"(?P{_IDENTIFIER})(?:\s+(?:AS\s+)?(?P{_IDENTIFIER}))?\s*$", + re.IGNORECASE, +) +_AGGREGATE_RE = re.compile( + rf"^\s*(?PCOUNT|SUM|AVG|MIN|MAX)\s*\(\s*" + rf"(?:(?:DISTINCT)\s+)?(?:(?P
{_IDENTIFIER})\s*\.\s*)?" + rf"(?P{_IDENTIFIER}|\*)\s*\)" + rf"(?:\s+(?:AS\s+)?(?P{_IDENTIFIER}))?\s*$", + re.IGNORECASE, +) +_EVIDENCE_TOKEN_RE = re.compile(r"\[\[df-evidence:([A-Za-z0-9_-]+)\]\]") +_TRACE_PAYLOAD_LIMIT = 8_000 + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def sha256_text(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + +def _json_default(value: Any) -> Any: + if isinstance(value, BaseModel): + return value.model_dump(mode="json") + if hasattr(value, "value"): + return value.value + return str(value) + + +def canonical_json(value: Any) -> str: + return json.dumps( + value, + default=_json_default, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + + +def digest_records(records: Any) -> str: + """Hash results without retaining or exposing record values.""" + + return sha256_text(canonical_json(records)) + + +def _get(value: Any, *names: str, default: Any = None) -> Any: + if value is None: + return default + if isinstance(value, dict): + lowered = {str(k).lower(): v for k, v in value.items()} + for name in names: + if name in value and value[name] is not None: + return value[name] + if name.lower() in lowered and lowered[name.lower()] is not None: + return lowered[name.lower()] + return default + for name in names: + if hasattr(value, name): + candidate = getattr(value, name) + if candidate is not None: + return candidate + return default + + +def _string(value: Any) -> str | None: + if value is None: + return None + if hasattr(value, "value"): + value = value.value + return str(value) + + +@dataclass(frozen=True) +class _SourceTemplate: + logical_entity_id: str | None + logical_entity_name: str + logical_field_id: str | None + logical_field_name: str | None + source_type: str + connector_id: str | None = None + connection_id: str | None = None + folder_id: str | None = None + external_object_id: str | None = None + external_object_name: str | None = None + primary_key: str | None = None + is_primary_source: bool | None = None + external_field_name: str | None = None + external_field_type: str | None = None + mapping_direction: str | None = None + mapping_complete: bool = True + + @property + def source_reference_id(self) -> str: + return "dfsrc_" + sha256_text(canonical_json(asdict(self)))[:24] + + def materialize(self, execution_id: str) -> DataFabricSourceRefV1: + source_id = self.source_reference_id + evidence_id = "dfev_" + sha256_text(f"1|{execution_id}|{source_id}")[:32] + return DataFabricSourceRefV1( + source_reference_id=source_id, + evidence_id=evidence_id, + **asdict(self), + ) + + +@dataclass(frozen=True) +class _JoinTemplate: + source_join_criteria_id: str | None + join_type: str | None + primary_object_id: str | None + primary_field_name: str | None + primary_source_reference_id: str | None + related_object_id: str | None + related_field_name: str | None + related_source_reference_id: str | None + + def materialize(self) -> DataFabricSourceJoinV1: + return DataFabricSourceJoinV1(**asdict(self)) + + +@dataclass(frozen=True) +class DataFabricSchemaSnapshot: + """Immutable authorization-scoped schema view for a cached graph.""" + + snapshot_id: str + resolved_at: str + schema_descriptors: tuple[str, ...] + sources: tuple[_SourceTemplate, ...] + joins: tuple[_JoinTemplate, ...] + + def materialize_sources( + self, execution_id: str + ) -> tuple[DataFabricSourceRefV1, ...]: + return tuple(source.materialize(execution_id) for source in self.sources) + + +def _external_field_entries(entity: Any) -> list[tuple[Any, Any, Any, Any]]: + """Return (field metadata, mapping, object, connection) tuples.""" + + entries: list[tuple[Any, Any, Any, Any]] = [] + for item in _get(entity, "external_fields", "externalFields", default=[]) or []: + grouped_fields = _get(item, "fields") + if grouped_fields is not None: + external_object = _get( + item, "external_object_detail", "externalObjectDetail" + ) + external_connection = _get( + item, "external_connection_detail", "externalConnectionDetail" + ) + for field in grouped_fields or []: + entries.append( + ( + _get(field, "field_metadata", "fieldMetadata"), + _get( + field, + "external_field_mapping_detail", + "externalFieldMappingDetail", + ), + external_object, + external_connection, + ) + ) + else: + entries.append( + ( + _get(item, "field_metadata", "fieldMetadata"), + _get( + item, + "external_field_mapping_detail", + "externalFieldMappingDetail", + ), + _get(item, "external_object_detail", "externalObjectDetail"), + _get( + item, + "external_connection_detail", + "externalConnectionDetail", + ), + ) + ) + return entries + + +def _build_source_templates(entity: Any) -> list[_SourceTemplate]: + entity_id = _string(_get(entity, "id")) + entity_name = str(_get(entity, "name")) + fields = _get(entity, "fields", default=[]) or [] + external_entries = _external_field_entries(entity) + by_internal_id: dict[str, list[tuple[Any, Any, Any, Any]]] = {} + by_field_name: dict[str, list[tuple[Any, Any, Any, Any]]] = {} + for entry in external_entries: + metadata, mapping, _, _ = entry + internal_id = _string( + _get(mapping, "internal_field_id", "internalFieldId") + ) or _string(_get(metadata, "id")) + field_name = _string(_get(metadata, "name")) + if internal_id: + by_internal_id.setdefault(internal_id, []).append(entry) + if field_name: + by_field_name.setdefault(field_name.lower(), []).append(entry) + + sources: list[_SourceTemplate] = [] + for field in fields: + field_id = _string(_get(field, "id")) + field_name = str(_get(field, "name")) + matches = ( + by_internal_id.get(field_id, []) if field_id else [] + ) or by_field_name.get(field_name.lower(), []) + is_external = bool(_get(field, "is_external_field", "isExternalField")) + if not matches and not is_external: + sources.append( + _SourceTemplate( + logical_entity_id=entity_id, + logical_entity_name=entity_name, + logical_field_id=field_id, + logical_field_name=field_name, + source_type="native", + ) + ) + continue + if not matches: + sources.append( + _SourceTemplate( + logical_entity_id=entity_id, + logical_entity_name=entity_name, + logical_field_id=field_id, + logical_field_name=field_name, + source_type="external", + mapping_complete=False, + ) + ) + continue + for _, mapping, external_object, external_connection in matches: + object_id = _string( + _get(external_object, "id") + or _get(mapping, "external_object_id", "externalObjectId") + ) + connector_id = _string( + _get(external_connection, "connector_id", "connectorId") + ) + connection_id = _string( + _get(external_connection, "connection_id", "connectionId") + ) + folder_id = _string(_get(external_connection, "folder_id", "folderId")) + object_name = _string( + _get( + external_object, + "external_object_name", + "externalObjectName", + ) + ) + external_field_name = _string( + _get(mapping, "external_field_name", "externalFieldName") + ) + complete = all( + ( + object_id, + connector_id, + connection_id, + folder_id, + object_name, + external_field_name, + ) + ) + sources.append( + _SourceTemplate( + logical_entity_id=entity_id, + logical_entity_name=entity_name, + logical_field_id=field_id, + logical_field_name=field_name, + source_type="external", + connector_id=connector_id, + connection_id=connection_id, + folder_id=folder_id, + external_object_id=object_id, + external_object_name=object_name, + primary_key=_string( + _get(external_object, "primary_key", "primaryKey") + ), + is_primary_source=_get( + external_object, "is_primary_source", "isPrimarySource" + ), + external_field_name=external_field_name, + external_field_type=_string( + _get(mapping, "external_field_type", "externalFieldType") + ), + mapping_direction=_string( + _get(mapping, "direction_type", "directionType") + ), + mapping_complete=complete, + ) + ) + return sources + + +def _match_source( + sources: Sequence[_SourceTemplate], object_id: str | None, field_name: str | None +) -> str | None: + candidates = [ + source + for source in sources + if source.external_object_id == object_id + and field_name is not None + and ( + (source.external_field_name or "").lower() == field_name.lower() + or (source.logical_field_name or "").lower() == field_name.lower() + ) + ] + return candidates[0].source_reference_id if len(candidates) == 1 else None + + +def _build_join_templates( + entities: Iterable[Any], sources: Sequence[_SourceTemplate] +) -> list[_JoinTemplate]: + joins: list[_JoinTemplate] = [] + primary_objects = { + source.external_object_id + for source in sources + if source.is_primary_source and source.external_object_id + } + for entity in entities: + entity_object_ids = { + source.external_object_id + for source in sources + if source.logical_entity_name == str(_get(entity, "name")) + and source.external_object_id + } + primary_object_id = next( + iter(sorted(entity_object_ids & primary_objects)), None + ) + for criterion in ( + _get(entity, "source_join_criteria", "sourceJoinCriteria", default=[]) or [] + ): + primary_field = _string(_get(criterion, "join_field_name", "joinFieldName")) + related_object = _string( + _get( + criterion, + "related_source_object_id", + "relatedSourceObjectId", + ) + ) + related_field = _string( + _get( + criterion, + "related_source_object_field_name", + "relatedSourceObjectFieldName", + "related_source_field_name", + "relatedSourceFieldName", + ) + ) + joins.append( + _JoinTemplate( + source_join_criteria_id=_string(_get(criterion, "id")), + join_type=_string(_get(criterion, "join_type", "joinType")), + primary_object_id=primary_object_id, + primary_field_name=primary_field, + primary_source_reference_id=_match_source( + sources, primary_object_id, primary_field + ), + related_object_id=related_object, + related_field_name=related_field, + related_source_reference_id=_match_source( + sources, related_object, related_field + ), + ) + ) + return joins + + +def build_schema_snapshot(entities: Iterable[Any]) -> DataFabricSchemaSnapshot: + """Build a canonical immutable snapshot from already-authorized SDK models.""" + + entity_list = list(entities) + sources = sorted( + ( + source + for entity in entity_list + for source in _build_source_templates(entity) + ), + key=lambda source: source.source_reference_id, + ) + joins = sorted( + _build_join_templates(entity_list, sources), + key=lambda join: canonical_json(asdict(join)), + ) + schema_descriptors = tuple( + sorted( + canonical_json( + { + "entity_id": _string(_get(entity, "id")), + "entity_name": _string(_get(entity, "name")), + "fields": sorted( + ( + { + "id": _string(_get(field, "id")), + "name": _string(_get(field, "name")), + "type": _string( + _get(_get(field, "sql_type", "sqlType"), "name") + ), + "is_external": bool( + _get( + field, + "is_external_field", + "isExternalField", + ) + ), + } + for field in (_get(entity, "fields", default=[]) or []) + ), + key=canonical_json, + ), + } + ) + for entity in entity_list + ) + ) + canonical = { + "version": "1.0", + "schema": schema_descriptors, + "sources": [asdict(source) for source in sources], + "joins": [asdict(join) for join in joins], + } + return DataFabricSchemaSnapshot( + snapshot_id=sha256_text(canonical_json(canonical)), + resolved_at=utc_now(), + schema_descriptors=schema_descriptors, + sources=tuple(sources), + joins=tuple(joins), + ) + + +def _split_expressions(value: str) -> list[str]: + expressions: list[str] = [] + start = 0 + depth = 0 + quote: str | None = None + for index, char in enumerate(value): + if quote: + if char == quote: + quote = None + continue + if char in ("'", '"'): + quote = char + elif char == "(": + depth += 1 + elif char == ")": + depth = max(0, depth - 1) + elif char == "," and depth == 0: + expressions.append(value[start:index].strip()) + start = index + 1 + tail = value[start:].strip() + if tail: + expressions.append(tail) + return expressions + + +def _parse_tables(statement: Any) -> tuple[list[str], dict[str, str]]: + tables: list[str] = [] + aliases: dict[str, str] = {} + expect_table = False + for token in statement.tokens: + if token.is_whitespace: + continue + if token.ttype in Keyword and ( + token.normalized == "FROM" or "JOIN" in token.normalized + ): + expect_table = True + continue + if not expect_table: + continue + identifiers = ( + list(token.get_identifiers()) + if isinstance(token, IdentifierList) + else [token] + ) + for identifier in identifiers: + if not isinstance(identifier, Identifier): + continue + table = identifier.get_real_name() + if not table: + continue + tables.append(table) + aliases[(identifier.get_alias() or table).lower()] = table + aliases[table.lower()] = table + expect_table = False + return list(dict.fromkeys(tables)), aliases + + +def _select_text(statement: Any) -> str | None: + collecting = False + parts: list[str] = [] + for token in statement.tokens: + if token.ttype is DML and token.normalized == "SELECT": + collecting = True + continue + if collecting and token.ttype in Keyword and token.normalized == "FROM": + break + if collecting: + parts.append(str(token)) + value = "".join(parts).strip() + return value or None + + +def _field_index( + sources: Sequence[_SourceTemplate], tables: Sequence[str] +) -> dict[tuple[str, str], list[_SourceTemplate]]: + table_set = {table.lower() for table in tables} + index: dict[tuple[str, str], list[_SourceTemplate]] = {} + for source in sources: + if ( + source.logical_entity_name.lower() in table_set + and source.logical_field_name + ): + index.setdefault( + ( + source.logical_entity_name.lower(), + source.logical_field_name.lower(), + ), + [], + ).append(source) + return index + + +def _resolve_field( + table: str | None, + field: str, + tables: Sequence[str], + aliases: dict[str, str], + index: dict[tuple[str, str], list[_SourceTemplate]], +) -> tuple[str, tuple[str, ...]] | None: + if table: + resolved_table = aliases.get(table.lower(), table) + candidates = index.get((resolved_table.lower(), field.lower()), []) + if not candidates: + return None + logical = ( + f"{candidates[0].logical_entity_name}.{candidates[0].logical_field_name}" + ) + return logical, tuple(source.source_reference_id for source in candidates) + matches = [ + (entity, values) + for (entity, name), values in index.items() + if name == field.lower() and entity in {item.lower() for item in tables} + ] + if len(matches) != 1: + return None + values = matches[0][1] + logical = f"{values[0].logical_entity_name}.{values[0].logical_field_name}" + return logical, tuple(source.source_reference_id for source in values) + + +def _wildcard_sources( + table: str | None, + tables: Sequence[str], + aliases: dict[str, str], + index: dict[tuple[str, str], list[_SourceTemplate]], +) -> tuple[tuple[str, ...], tuple[str, ...]]: + selected_tables = [aliases.get(table.lower(), table)] if table else list(tables) + logical_fields: list[str] = [] + refs: list[str] = [] + for (entity, _), sources in index.items(): + if entity not in {item.lower() for item in selected_tables}: + continue + logical_fields.append( + f"{sources[0].logical_entity_name}.{sources[0].logical_field_name}" + ) + refs.extend(source.source_reference_id for source in sources) + return tuple(sorted(set(logical_fields))), tuple(sorted(set(refs))) + + +def _masked_sql(statement: Any) -> str: + parts: list[str] = [] + for token in statement.flatten(): + if token.ttype and (token.ttype.parent and "Literal" in str(token.ttype)): + parts.append(" ") + else: + parts.append(str(token)) + return "".join(parts) + + +def _role_identifiers(masked_sql: str) -> list[tuple[str | None, str, str]]: + roles: list[tuple[str | None, str, str]] = [] + clauses = ( + ("filter", r"\bWHERE\b(.*?)(?=\bGROUP\s+BY\b|\bORDER\s+BY\b|\bLIMIT\b|$)"), + ("grouping", r"\bGROUP\s+BY\b(.*?)(?=\bORDER\s+BY\b|\bLIMIT\b|$)"), + ("ordering", r"\bORDER\s+BY\b(.*?)(?=\bLIMIT\b|$)"), + ( + "join", + r"\bON\b(.*?)(?=\b(?:LEFT|RIGHT|FULL|INNER|CROSS)?\s*JOIN\b|\bWHERE\b|\bGROUP\s+BY\b|\bORDER\s+BY\b|\bLIMIT\b|$)", + ), + ) + for role, pattern in clauses: + for match in re.finditer(pattern, masked_sql, flags=re.IGNORECASE | re.DOTALL): + text = match.group(1) + for qualified, table, field, bare in re.findall( + rf"(?:\b(({_IDENTIFIER})\.({_IDENTIFIER}))\b)|(?:\b({_IDENTIFIER})\b)", + text, + ): + if qualified: + roles.append((table, field, role)) + elif bare and bare.upper() not in { + "AND", + "BETWEEN", + "FALSE", + "ILIKE", + "LIKE", + "OR", + "NOT", + "IN", + "IS", + "NULL", + "TRUE", + "ASC", + "DESC", + "AS", + }: + roles.append((None, bare, role)) + return roles + + +def infer_column_lineage( + sql: str, snapshot: DataFabricSchemaSnapshot +) -> tuple[tuple[DataFabricColumnLineageV1, ...], bool, int, int]: + """Infer bounded SQL lineage, returning partial on every ambiguity.""" + + parsed = sqlparse.parse(sql) + if len(parsed) != 1: + return (), True, 0, 1 + statement = parsed[0] + tables, aliases = _parse_tables(statement) + select_text = _select_text(statement) + index = _field_index(snapshot.sources, tables) + if not tables or not select_text: + return (), True, 0, 1 + + columns: list[DataFabricColumnLineageV1] = [] + resolved_count = 0 + reference_count = 0 + partial = False + by_ref: dict[str, int] = {} + + for expression in _split_expressions(select_text): + reference_count += 1 + aggregate = _AGGREGATE_RE.match(expression) + direct = _DIRECT_PROJECTION_RE.match(expression) + if aggregate: + field = aggregate.group("field") + alias = aggregate.group("alias") + if field == "*": + wildcard_fields, refs = _wildcard_sources(None, tables, aliases, index) + columns.append( + DataFabricColumnLineageV1( + result_expression=f"{aggregate.group('function').upper()}(*)", + result_alias=alias, + logical_fields=wildcard_fields, + physical_source_reference_ids=refs, + roles=("projection", "aggregate"), + lineage_authority="partial", + ) + ) + partial = True + continue + resolved = _resolve_field( + aggregate.group("table"), field, tables, aliases, index + ) + if resolved: + logical, refs = resolved + resolved_count += 1 + item = DataFabricColumnLineageV1( + result_expression=f"{aggregate.group('function').upper()}({logical})", + result_alias=alias, + logical_fields=(logical,), + physical_source_reference_ids=refs, + roles=("projection", "aggregate"), + ) + columns.append(item) + for ref in refs: + by_ref[ref] = len(columns) - 1 + else: + partial = True + columns.append( + DataFabricColumnLineageV1( + result_expression="", + result_alias=alias, + roles=("projection", "aggregate"), + lineage_authority="partial", + ) + ) + continue + if expression.strip() == "*" or expression.strip().endswith(".*"): + table = expression.strip()[:-2] if expression.strip() != "*" else None + wildcard_fields, refs = _wildcard_sources(table, tables, aliases, index) + columns.append( + DataFabricColumnLineageV1( + result_expression=f"{table or '*'}.*", + logical_fields=wildcard_fields, + physical_source_reference_ids=refs, + roles=("projection",), + lineage_authority="partial", + ) + ) + partial = True + continue + if direct: + resolved = _resolve_field( + direct.group("table"), direct.group("field"), tables, aliases, index + ) + if resolved: + logical, refs = resolved + resolved_count += 1 + columns.append( + DataFabricColumnLineageV1( + result_expression=logical, + result_alias=direct.group("alias"), + logical_fields=(logical,), + physical_source_reference_ids=refs, + roles=("projection",), + ) + ) + for ref in refs: + by_ref[ref] = len(columns) - 1 + else: + partial = True + columns.append( + DataFabricColumnLineageV1( + result_expression="", + result_alias=direct.group("alias"), + roles=("projection",), + lineage_authority="partial", + ) + ) + continue + partial = True + columns.append( + DataFabricColumnLineageV1( + result_expression=f"", + roles=("projection",), + lineage_authority="partial", + ) + ) + + # Add non-projection roles without retaining predicates or literal values. + for table, field, role in _role_identifiers(_masked_sql(statement)): + reference_count += 1 + alias_index = next( + ( + index + for index, column in enumerate(columns) + if table is None + and column.result_alias + and column.result_alias.lower() == field.lower() + ), + None, + ) + if alias_index is not None: + resolved_count += 1 + item = columns[alias_index] + columns[alias_index] = item.model_copy( + update={"roles": tuple(dict.fromkeys((*item.roles, role)))} + ) + continue + resolved = _resolve_field(table, field, tables, aliases, index) + if not resolved: + partial = True + continue + logical, refs = resolved + resolved_count += 1 + matching_index = next((by_ref[ref] for ref in refs if ref in by_ref), None) + if matching_index is not None: + item = columns[matching_index] + columns[matching_index] = item.model_copy( + update={"roles": tuple(dict.fromkeys((*item.roles, role)))} + ) + else: + columns.append( + DataFabricColumnLineageV1( + result_expression=logical, + logical_fields=(logical,), + physical_source_reference_ids=refs, + roles=(role,), + ) + ) + for ref in refs: + by_ref[ref] = len(columns) - 1 + + if any( + not source.mapping_complete + for source in snapshot.sources + if source.source_reference_id + in {ref for column in columns for ref in column.physical_source_reference_ids} + ): + partial = True + return tuple(columns), partial, resolved_count, reference_count + + +def build_attempt( + *, + execution_id: str, + tool_call_id: str, + ordinal: int, + sql: str, + started_at: str, + completed_at: str, + duration_ms: float, + result: dict[str, Any], +) -> DataFabricQueryAttemptV1: + succeeded = not bool(result.get("error")) + records = result.get("records", []) + result_digest = digest_records(records) if succeeded else None + error_category = getattr(result, "error_category", None) + return DataFabricQueryAttemptV1( + attempt_id="dfattempt_" + + sha256_text(f"{execution_id}|{tool_call_id}|{ordinal}")[:24], + tool_call_id=tool_call_id, + ordinal=ordinal, + sql_hash=sha256_text(sql), + started_at=started_at, + completed_at=completed_at, + duration_ms=duration_ms, + outcome="success" if succeeded else "error", + row_count=int(result.get("total_count", len(records))), + error_category=_string(error_category), + result_digest=result_digest, + ) + + +def _package_version() -> str: + try: + return version("uipath-langchain") + except PackageNotFoundError: + return "unavailable" + + +def model_identity(llm: Any) -> str | None: + for name in ("model_name", "model", "model_id"): + value = getattr(llm, name, None) + if isinstance(value, str) and value: + return value + return None + + +def build_evidence_envelope( + *, + execution_id: str, + tool_name: str, + snapshot: DataFabricSchemaSnapshot, + attempts: Sequence[DataFabricQueryAttemptV1], + column_lineage: Sequence[DataFabricColumnLineageV1], + prompt_version: str, + prompt_hash: str, + llm: Any, + resolved_reference_count: int, + total_reference_count: int, + forced_partial: bool, +) -> DataFabricEvidenceEnvelopeV1: + ordered_attempts = tuple(sorted(attempts, key=lambda item: item.ordinal)) + successes = [item for item in ordered_attempts if item.outcome == "success"] + row_count = sum(item.row_count for item in successes) + combined_digest = sha256_text( + canonical_json([item.result_digest for item in successes]) + ) + # A successful backend read is the trusted-side authorization signal for + # releasing cached, authorization-scoped source identifiers. Failure-only + # executions retain the attempt ledger but not the cached source graph. + sources = snapshot.materialize_sources(execution_id) if successes else () + coverage = ( + min(1.0, resolved_reference_count / total_reference_count) + if total_reference_count + else 0.0 + ) + partial = ( + forced_partial + or coverage < 1.0 + or any(item.lineage_authority == "partial" for item in column_lineage) + ) + return DataFabricEvidenceEnvelopeV1( + execution_id=execution_id, + tool_name=tool_name, + created_at=utc_now(), + schema_resolved_at=snapshot.resolved_at, + schema_snapshot_id=snapshot.snapshot_id, + lineage_authority="partial" if partial else "inferred", + lineage_coverage=coverage, + attempts=ordered_attempts, + sources=sources, + joins=tuple(join.materialize() for join in snapshot.joins) if successes else (), + column_lineage=tuple(column_lineage), + generation_provenance=DataFabricGenerationProvenanceV1( + package_version=_package_version(), + model_identity=model_identity(llm), + prompt_version=prompt_version, + prompt_hash=prompt_hash, + schema_snapshot_id=snapshot.snapshot_id, + ), + result_integrity=DataFabricResultIntegrityV1( + successful_result_count=len(successes), + row_count=row_count, + result_digest=combined_digest, + ), + ) + + +def emit_trace_summary(envelope: DataFabricEvidenceEnvelopeV1) -> None: + """Emit bounded, record-free lineage telemetry on the current span.""" + + try: + from opentelemetry import trace + + span = trace.get_current_span() + if not span.is_recording(): + return + prefix = "datafabric.evidence." + span.set_attribute(prefix + "execution_id", envelope.execution_id) + span.set_attribute(prefix + "schema_snapshot_id", envelope.schema_snapshot_id) + span.set_attribute(prefix + "lineage_authority", envelope.lineage_authority) + span.set_attribute(prefix + "lineage_coverage", envelope.lineage_coverage) + span.set_attribute(prefix + "source_count", len(envelope.sources)) + span.set_attribute(prefix + "join_count", len(envelope.joins)) + span.set_attribute(prefix + "attempt_count", len(envelope.attempts)) + span.set_attribute(prefix + "row_count", envelope.result_integrity.row_count) + span.set_attribute( + prefix + "result_digest", envelope.result_integrity.result_digest + ) + if envelope.attempts: + span.set_attribute(prefix + "query_hash", envelope.attempts[-1].sql_hash) + + payload = { + "version": envelope.version, + "execution_id": envelope.execution_id, + "schema_snapshot_id": envelope.schema_snapshot_id, + "mappings": [ + { + "source_reference_id": source.source_reference_id, + "logical_entity_id": source.logical_entity_id, + "logical_field_id": source.logical_field_id, + "source_type": source.source_type, + "connector_id": source.connector_id, + "connection_id": source.connection_id, + "folder_id": source.folder_id, + "external_object_id": source.external_object_id, + } + for source in envelope.sources + ], + "joins": [ + { + "source_join_criteria_id": join.source_join_criteria_id, + "primary_object_id": join.primary_object_id, + "primary_source_reference_id": join.primary_source_reference_id, + "related_object_id": join.related_object_id, + "related_source_reference_id": join.related_source_reference_id, + } + for join in envelope.joins + ], + } + serialized = canonical_json(payload) + if len(serialized.encode("utf-8")) > _TRACE_PAYLOAD_LIMIT: + span.add_event( + "datafabric.lineage", + { + "payload_hash": sha256_text(serialized), + "source_count": len(envelope.sources), + "join_count": len(envelope.joins), + "trace_payload_truncated": True, + }, + ) + else: + span.add_event( + "datafabric.lineage", + {"payload": serialized, "trace_payload_truncated": False}, + ) + except (ImportError, AttributeError): + return + + +def validate_answer_citations_v1( + text: str, + envelope: DataFabricEvidenceEnvelopeV1, + *, + authorized_source_reference_ids: set[str] | None = None, +) -> tuple[str, tuple[DataFabricAnswerCitationV1, ...]]: + """Validate evidence tokens without changing plain text for unsupported clients. + + Callers with Data Fabric citation support may convert the returned citation + models. Callers without that support should ignore the second return value + and return ``text`` unchanged. + """ + + if envelope.version.split(".", 1)[0] != "1": + raise ValueError("Unsupported Data Fabric evidence major version") + evidence = {source.evidence_id: source for source in envelope.sources} + used_source_ids = { + source_id + for column in envelope.column_lineage + for source_id in column.physical_source_reference_ids + } + allowed = ( + used_source_ids + if authorized_source_reference_ids is None + else used_source_ids & authorized_source_reference_ids + ) + citations: list[DataFabricAnswerCitationV1] = [] + for evidence_id in _EVIDENCE_TOKEN_RE.findall(text): + source = evidence.get(evidence_id) + if source is None or source.source_reference_id not in allowed: + raise ValueError("Invalid or unauthorized Data Fabric evidence token") + citations.append( + DataFabricAnswerCitationV1( + evidence_id=evidence_id, + execution_id=envelope.execution_id, + source_reference_ids=(source.source_reference_id,), + ) + ) + return text, tuple(citations) diff --git a/src/uipath_langchain/agent/tools/datafabric_tool/models.py b/src/uipath_langchain/agent/tools/datafabric_tool/models.py index fbd0cbd59..f1a9a2589 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/models.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/models.py @@ -1,6 +1,8 @@ -"""Pydantic models for Data Fabric entity schemas.""" +"""Pydantic models for Data Fabric entity schemas and local evidence.""" -from pydantic import BaseModel, Field +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field NUMERIC_TYPES = frozenset({"int", "decimal", "float", "double", "bigint"}) TEXT_TYPES = frozenset({"varchar", "nvarchar", "text", "string", "ntext"}) @@ -110,3 +112,133 @@ class DataFabricExecuteSqlInput(BaseModel): "Use exact table and column names from the entity schemas." ), ) + + +class _EvidenceModelV1(BaseModel): + """Immutable base for the internal, versioned evidence contract.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + +class DataFabricSourceRefV1(_EvidenceModelV1): + """Logical Data Fabric field mapped to its locally known physical source.""" + + source_reference_id: str + evidence_id: str + logical_entity_id: str | None = None + logical_entity_name: str + logical_field_id: str | None = None + logical_field_name: str | None = None + source_type: Literal["native", "external"] + connector_id: str | None = None + connection_id: str | None = None + folder_id: str | None = None + external_object_id: str | None = None + external_object_name: str | None = None + primary_key: str | None = None + is_primary_source: bool | None = None + external_field_name: str | None = None + external_field_type: str | None = None + mapping_direction: str | None = None + mapping_complete: bool = True + + +class DataFabricSourceJoinV1(_EvidenceModelV1): + """A source-object join reported by entity-set resolution.""" + + source_join_criteria_id: str | None = None + join_type: str | None = None + primary_object_id: str | None = None + primary_field_name: str | None = None + primary_source_reference_id: str | None = None + related_object_id: str | None = None + related_field_name: str | None = None + related_source_reference_id: str | None = None + + +class DataFabricColumnLineageV1(_EvidenceModelV1): + """Conservative lineage inferred for one SQL result expression.""" + + result_expression: str + result_alias: str | None = None + logical_fields: tuple[str, ...] = () + physical_source_reference_ids: tuple[str, ...] = () + roles: tuple[ + Literal["projection", "filter", "join", "grouping", "ordering", "aggregate"], + ..., + ] = () + lineage_authority: Literal["inferred", "partial"] = "inferred" + + +class DataFabricQueryAttemptV1(_EvidenceModelV1): + """One ordered invocation of the inner execute-SQL tool.""" + + attempt_id: str + tool_call_id: str + ordinal: int + sql_hash: str + started_at: str + completed_at: str + duration_ms: float + outcome: Literal["success", "error"] + row_count: int + error_category: str | None = None + result_digest: str | None = None + + +class DataFabricGenerationProvenanceV1(_EvidenceModelV1): + """Locally observable inputs that produced an execution.""" + + package_name: str = "uipath-langchain" + package_version: str + model_identity: str | None = None + prompt_version: str + prompt_hash: str + schema_snapshot_id: str + backend_execution_id: None = None + backend_query_plan: None = None + + +class DataFabricResultIntegrityV1(_EvidenceModelV1): + """Record-free integrity summary for successful query results.""" + + successful_result_count: int + row_count: int + result_digest: str + record_contributors: None = None + source_system_timestamps: None = None + + +class DataFabricEvidenceEnvelopeV1(_EvidenceModelV1): + """Complete local evidence for one outer Data Fabric tool execution.""" + + version: Literal["1.0"] = "1.0" + execution_id: str + tool_name: str + created_at: str + schema_resolved_at: str + schema_snapshot_id: str + lineage_authority: Literal["inferred", "partial"] + lineage_coverage: float + attempts: tuple[DataFabricQueryAttemptV1, ...] = () + sources: tuple[DataFabricSourceRefV1, ...] = () + joins: tuple[DataFabricSourceJoinV1, ...] = () + column_lineage: tuple[DataFabricColumnLineageV1, ...] = () + generation_provenance: DataFabricGenerationProvenanceV1 + result_integrity: DataFabricResultIntegrityV1 + unavailable_backend_properties: tuple[str, ...] = ( + "backend_execution_id", + "normalized_physical_query_plan", + "authoritative_output_column_lineage", + "record_contributors", + "source_system_timestamps", + ) + + +class DataFabricAnswerCitationV1(_EvidenceModelV1): + """A citation validated against the current execution's evidence artifact.""" + + version: Literal["1.0"] = "1.0" + evidence_id: str + execution_id: str + source_reference_ids: tuple[str, ...] 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 index ed905f363..73ffff3ee 100644 --- a/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_subgraph.py +++ b/src/uipath_langchain/agent/tools/datafabric_tool/ontology/ontology_subgraph.py @@ -33,6 +33,7 @@ from uipath.platform.entities import EntitiesService, Entity from ...datafabric_query_tool import DataFabricQueryTool +from ..lineage import sha256_text from ..models import DataFabricExecuteSqlInput logger = logging.getLogger(__name__) @@ -53,7 +54,7 @@ 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) + logger.debug("execute_sql called with query hash %s", sha256_text(sql_query)) try: records = await self._entities.query_entity_records_async( sql_query=sql_query, @@ -64,7 +65,7 @@ async def __call__(self, sql_query: str) -> dict[str, Any]: "sql_query": sql_query, } except Exception as e: - logger.error("SQL query failed: %s", e) + logger.error("SQL query failed (%s)", type(e).__name__) return { "records": [], "total_count": 0, @@ -146,11 +147,7 @@ async def _execute_tool_call(self, tool_call: ToolCall) -> tuple[ToolMessage, bo "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 - ) + succeeded = isinstance(result, dict) and not result.get("error") return ( ToolMessage( content=str(result), diff --git a/tests/agent/tools/test_datafabric_lineage.py b/tests/agent/tools/test_datafabric_lineage.py new file mode 100644 index 000000000..5b8a7b267 --- /dev/null +++ b/tests/agent/tools/test_datafabric_lineage.py @@ -0,0 +1,460 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import ToolCall, ToolMessage +from langgraph.types import Command +from uipath.agent.models.agent import AgentContextResourceConfig + +from uipath_langchain.agent.react.types import AgentGraphState +from uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph import ( + DataFabricGraph, +) +from uipath_langchain.agent.tools.datafabric_tool.datafabric_tool import ( + DataFabricTextQueryHandler, + create_datafabric_query_tool, +) +from uipath_langchain.agent.tools.datafabric_tool.lineage import ( + build_evidence_envelope, + build_schema_snapshot, + canonical_json, + infer_column_lineage, + sha256_text, + validate_answer_citations_v1, +) +from uipath_langchain.agent.tools.datafabric_tool.models import ( + DataFabricQueryAttemptV1, +) + + +def _field(name: str, field_id: str, *, external: bool = False): + return SimpleNamespace( + id=field_id, + name=name, + is_external_field=external, + ) + + +def _entity( + name: str = "Invoice", + *, + fields=None, + external_fields=None, + joins=None, +): + return SimpleNamespace( + id=f"entity-{name.lower()}", + name=name, + fields=fields or [], + external_fields=external_fields or [], + source_join_criteria=joins or [], + ) + + +def _external_group( + *, + object_id: str, + object_name: str, + field_id: str, + logical_name: str, + external_name: str, + primary: bool, + connection_id: str = "connection-1", + folder_id: str = "folder-1", +): + return { + "externalObjectDetail": { + "id": object_id, + "externalObjectName": object_name, + "primaryKey": "ExternalId", + "isPrimarySource": primary, + }, + "externalConnectionDetail": { + "connectorId": "connector-1", + "connectionId": connection_id, + "folderId": folder_id, + }, + "fields": [ + { + "fieldMetadata": {"id": field_id, "name": logical_name}, + "externalFieldMappingDetail": { + "internalFieldId": field_id, + "externalObjectId": object_id, + "externalFieldName": external_name, + "externalFieldType": "string", + "directionType": "Read", + }, + } + ], + } + + +def _snapshot(): + entity = _entity( + fields=[ + _field("Id", "field-id"), + _field("Vendor", "field-vendor", external=True), + _field("Amount", "field-amount"), + ], + external_fields=[ + _external_group( + object_id="object-primary", + object_name="InvoiceObject", + field_id="field-vendor", + logical_name="Vendor", + external_name="supplier_name", + primary=True, + ), + _external_group( + object_id="object-related", + object_name="SupplierObject", + field_id="field-vendor", + logical_name="Vendor", + external_name="name", + primary=False, + connection_id="connection-2", + folder_id="folder-2", + ), + ], + joins=[ + { + "id": "join-1", + "joinType": "left", + "joinFieldName": "supplier_name", + "relatedSourceObjectId": "object-related", + "relatedSourceObjectFieldName": "name", + } + ], + ) + return build_schema_snapshot([entity]) + + +def _successful_attempt( + ordinal: int = 1, call_id: str = "inner-1" +) -> DataFabricQueryAttemptV1: + return DataFabricQueryAttemptV1( + attempt_id=f"attempt-{ordinal}", + tool_call_id=call_id, + ordinal=ordinal, + sql_hash=sha256_text(call_id), + started_at="2026-08-11T00:00:00+00:00", + completed_at="2026-08-11T00:00:01+00:00", + duration_ms=1.0, + outcome="success", + row_count=0, + result_digest=sha256_text("[]"), + ) + + +def test_snapshot_is_deterministic_and_preserves_external_graph(): + first = _snapshot() + second = _snapshot() + + assert first.snapshot_id == second.snapshot_id + assert first.resolved_at != "" + vendor_sources = [ + source for source in first.sources if source.logical_field_name == "Vendor" + ] + assert {source.external_object_id for source in vendor_sources} == { + "object-primary", + "object-related", + } + assert {source.connection_id for source in vendor_sources} == { + "connection-1", + "connection-2", + } + assert first.joins[0].source_join_criteria_id == "join-1" + assert first.joins[0].primary_source_reference_id is not None + assert first.joins[0].related_source_reference_id is not None + + +def test_snapshot_hash_changes_when_mapping_changes(): + original = _snapshot() + entity = _entity( + fields=[_field("Vendor", "field-vendor", external=True)], + external_fields=[ + _external_group( + object_id="object-primary", + object_name="InvoiceObject", + field_id="field-vendor", + logical_name="Vendor", + external_name="changed_supplier_name", + primary=True, + ) + ], + ) + + assert build_schema_snapshot([entity]).snapshot_id != original.snapshot_id + + +def test_missing_external_mapping_is_explicitly_incomplete(): + snapshot = build_schema_snapshot( + [_entity(fields=[_field("Vendor", "field-vendor", external=True)])] + ) + + assert snapshot.sources[0].source_type == "external" + assert snapshot.sources[0].mapping_complete is False + columns, partial, _, _ = infer_column_lineage( + "SELECT Vendor FROM Invoice", snapshot + ) + assert columns[0].physical_source_reference_ids + assert partial is True + + +def test_projection_alias_filter_group_order_and_aggregate_lineage(): + columns, partial, resolved, total = infer_column_lineage( + "SELECT Vendor AS supplier, SUM(Amount) AS total " + "FROM Invoice WHERE Vendor = 'secret literal' " + "GROUP BY Vendor ORDER BY total DESC", + _snapshot(), + ) + + assert partial is False + assert resolved == total + supplier = next(item for item in columns if item.result_alias == "supplier") + assert supplier.logical_fields == ("Invoice.Vendor",) + assert "filter" in supplier.roles + assert "grouping" in supplier.roles + aggregate = next(item for item in columns if item.result_alias == "total") + assert aggregate.result_expression == "SUM(Invoice.Amount)" + assert aggregate.roles == ("projection", "aggregate", "ordering") + assert "secret literal" not in canonical_json(columns) + + +def test_wildcard_and_unsupported_expression_are_partial_without_guessing(): + wildcard, wildcard_partial, _, _ = infer_column_lineage( + "SELECT * FROM Invoice", _snapshot() + ) + unsupported, unsupported_partial, _, _ = infer_column_lineage( + "SELECT Amount * 1.2 AS adjusted FROM Invoice", _snapshot() + ) + + assert wildcard_partial is True + assert wildcard[0].lineage_authority == "partial" + assert unsupported_partial is True + assert unsupported[0].logical_fields == () + assert unsupported[0].result_expression.startswith(" Entity: + """Create a minimal Entity for testing.""" + e = MagicMock(spec=Entity) + e.name = name + e.external_fields = external_fields + return e + + +def _make_entities_service( + records: list[dict[str, Any]] | None = None, + error: Exception | None = None, +) -> MagicMock: + svc = MagicMock() + if error: + svc.query_entity_records_async = AsyncMock(side_effect=error) + else: + svc.query_entity_records_async = AsyncMock(return_value=records or []) + return svc + + +# --------------------------------------------------------------------------- +# _noop_context +# --------------------------------------------------------------------------- + + +def test_noop_context_yields_none() -> None: + with _noop_context() as val: + assert val is None + + +# --------------------------------------------------------------------------- +# DataFabricSubgraphState +# --------------------------------------------------------------------------- + + +def test_state_defaults() -> None: + state = DataFabricSubgraphState() + assert state.messages == [] + assert state.iteration_count == 0 + assert state.last_tool_success is False + assert state.last_error_category == "" + assert state.last_error_detail == "" + + +# --------------------------------------------------------------------------- +# QueryExecutor.__init__ — entity attribute computation +# --------------------------------------------------------------------------- + + +def test_query_executor_entity_attrs_native_only() -> None: + entities = [_make_entity("Orders"), _make_entity("Products")] + svc = _make_entities_service() + qe = QueryExecutor(svc, entities) + assert qe._entity_attrs["df.entity_count"] == 2 + assert qe._entity_attrs["df.native_entity_count"] == 2 + assert qe._entity_attrs["df.federated_entity_count"] == 0 + assert all("entities" not in key for key in qe._entity_attrs) + +def test_query_executor_entity_attrs_mixed() -> None: + entities = [ + _make_entity("Orders"), + _make_entity("ExtTable", external_fields=["col1"]), + ] + svc = _make_entities_service() + qe = QueryExecutor(svc, entities) + assert qe._entity_attrs["df.native_entity_count"] == 1 + assert qe._entity_attrs["df.federated_entity_count"] == 1 + assert "ExtTable" not in str(qe._entity_attrs) + + +# --------------------------------------------------------------------------- +# QueryExecutor.__call__ — success path (no OTEL) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_query_executor_success_no_otel() -> None: + records = [{"id": 1}, {"id": 2}] + svc = _make_entities_service(records=records) + qe = QueryExecutor(svc, [_make_entity("T")]) + + with patch.dict( + "sys.modules", {"opentelemetry": None, "opentelemetry.trace": None} + ): + result = await qe("SELECT * FROM T") + + assert result["records"] == records + assert result["total_count"] == 2 + assert result["sql_query"] == "SELECT * FROM T" + assert "error" not in result + + +@pytest.mark.asyncio async def test_query_executor_requests_relationships_as_scalar() -> None: - """The Data Fabric tool always requests scalar relationship typing so the SQL - it writes can join on ``relationshipField = Other.Id``.""" - entities = MagicMock() - entities.query_entity_records_async = AsyncMock(return_value=[{"id": 1}]) + """Relationship fields are requested as scalar ids for SQL joins.""" + svc = _make_entities_service(records=[{"id": 1}]) - result = await QueryExecutor(entities)("SELECT id FROM TaskEntity LIMIT 10") + result = await QueryExecutor(svc, [])("SELECT id FROM TaskEntity LIMIT 10") - entities.query_entity_records_async.assert_awaited_once_with( + svc.query_entity_records_async.assert_awaited_once_with( sql_query="SELECT id FROM TaskEntity LIMIT 10", relationships_as_scalar=True, ) assert result["records"] == [{"id": 1}] + + +# --------------------------------------------------------------------------- +# QueryExecutor.__call__ — success path with OTEL span +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_query_executor_success_with_span() -> None: + records = [{"id": 1}] + svc = _make_entities_service(records=records) + qe = QueryExecutor(svc, [_make_entity("T")]) + + mock_span = MagicMock() + mock_tracer = MagicMock() + mock_tracer.start_as_current_span.return_value.__enter__ = MagicMock( + return_value=mock_span + ) + mock_tracer.start_as_current_span.return_value.__exit__ = MagicMock( + return_value=False + ) + + with patch("opentelemetry.trace.get_tracer", return_value=mock_tracer): + result = await qe("SELECT 1") + + assert result["total_count"] == 1 + set_attr_calls = { + call.args[0]: call.args[1] for call in mock_span.set_attribute.call_args_list + } + assert set_attr_calls["df.row_count"] == 1 + assert set_attr_calls["df.success"] is True + span_attributes = mock_tracer.start_as_current_span.call_args.kwargs["attributes"] + assert span_attributes["df.query_hash"] != "SELECT 1" + assert "df.sql_query" not in span_attributes + + +# --------------------------------------------------------------------------- +# QueryExecutor.__call__ — error path (no OTEL, plain exception) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_query_executor_error_no_otel() -> None: + svc = _make_entities_service(error=RuntimeError("connection timeout")) + qe = QueryExecutor(svc, [_make_entity("T")]) + + with patch.dict( + "sys.modules", {"opentelemetry": None, "opentelemetry.trace": None} + ): + result = await qe("SELECT * FROM T") + + assert result["records"] == [] + assert result["total_count"] == 0 + assert "connection timeout" in result["error"] + + +# --------------------------------------------------------------------------- +# QueryExecutor.__call__ — error path with EnrichedException + DataFabricError +# --------------------------------------------------------------------------- + + +def _make_enriched_exception(msg: str = "enriched error") -> EnrichedException: + """Create an EnrichedException with mocked httpx internals.""" + mock_response = MagicMock() + mock_response.status_code = 400 + mock_response.headers = {"content-type": "application/json"} + mock_response.content = msg.encode("utf-8") + mock_request = MagicMock() + mock_request.url = "https://datafabric_.example.com/query" + mock_request.method = "POST" + mock_error = MagicMock() + mock_error.response = mock_response + mock_error.request = mock_request + return EnrichedException(mock_error) + + +@pytest.mark.asyncio +async def test_query_executor_error_with_datafabric_error() -> None: + enriched = _make_enriched_exception() + + svc = _make_entities_service(error=enriched) + qe = QueryExecutor(svc, [_make_entity("T")]) + + fake_df_error = _FakeDataFabricError( + code="SQL_VALIDATION", + message="Invalid column reference", + trace_id="abc123", + category=_FakeCategory.BAD_SQL, + is_bad_sql=True, + ) + + with ( + patch.dict("sys.modules", {"opentelemetry": None, "opentelemetry.trace": None}), + patch( + "uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph.DataFabricError" + ) as mock_dfe_cls, + ): + mock_dfe_cls.from_enriched_exception.return_value = fake_df_error + result = await qe("SELECT bad_col FROM T") + + assert result["records"] == [] + assert "[SQL_VALIDATION]" in result["error"] + assert "Fix the SQL syntax" in result["error"] + + +# --------------------------------------------------------------------------- +# QueryExecutor._build_error_detail +# --------------------------------------------------------------------------- + + +def test_build_error_detail_no_df_error() -> None: + detail = QueryExecutor._build_error_detail(RuntimeError("boom"), None) + assert detail == "boom" + + +def test_build_error_detail_with_code_and_bad_sql() -> None: + df_err = _FakeDataFabricError( + code="SQL_VALIDATION", + message="Invalid column", + trace_id=None, + category=_FakeCategory.BAD_SQL, + is_bad_sql=True, + ) + detail = QueryExecutor._build_error_detail(RuntimeError("x"), df_err) # type: ignore[arg-type] + assert "[SQL_VALIDATION]" in detail + assert "(category: bad_sql)" in detail + assert "Invalid column" in detail + assert "Fix the SQL syntax" in detail + + +def test_build_error_detail_retryable() -> None: + df_err = _FakeDataFabricError( + code="TIMEOUT", + message="Request timed out", + trace_id="t1", + category=_FakeCategory.RETRYABLE, + is_retryable=True, + ) + detail = QueryExecutor._build_error_detail(RuntimeError("x"), df_err) # type: ignore[arg-type] + assert "[TIMEOUT]" in detail + assert "transient" in detail + + +def test_build_error_detail_unknown_category() -> None: + df_err = _FakeDataFabricError( + code="SOMETHING", + message="msg", + trace_id=None, + category=_FakeCategory.UNKNOWN, + ) + detail = QueryExecutor._build_error_detail(RuntimeError("x"), df_err) # type: ignore[arg-type] + assert "[SOMETHING]" in detail + # "unknown" category should NOT appear + assert "(category:" not in detail + + +def test_build_error_detail_no_code() -> None: + df_err = _FakeDataFabricError( + code=None, + message="msg", + trace_id=None, + category=_FakeCategory.UNKNOWN, + ) + # Falls through to str(exc) + detail = QueryExecutor._build_error_detail(RuntimeError("fallback"), df_err) # type: ignore[arg-type] + assert detail == "fallback" + + +def test_build_error_detail_no_message() -> None: + df_err = _FakeDataFabricError( + code="ERR", + message=None, + trace_id=None, + category=_FakeCategory.INFRASTRUCTURE, + ) + detail = QueryExecutor._build_error_detail(RuntimeError("x"), df_err) # type: ignore[arg-type] + assert "[ERR]" in detail + assert "(category: infrastructure)" in detail + + +# --------------------------------------------------------------------------- +# DataFabricGraph — routing +# --------------------------------------------------------------------------- + + +def _make_graph() -> DataFabricGraph: + """Create a DataFabricGraph with a mocked LLM.""" + llm = MagicMock(spec=["model_copy"]) + bound = MagicMock() + copy = MagicMock() + copy.bind_tools = MagicMock(return_value=bound) + llm.model_copy.return_value = copy + + entities = [_make_entity("Orders")] + svc = _make_entities_service() + + with patch( + "uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph.datafabric_prompt_builder" + ) as mock_pb: + mock_pb.build.return_value = "system prompt" + graph = DataFabricGraph(llm, entities, svc, max_iterations=3) + return graph + + +def test_router_to_tool() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", tool_calls=[{"id": "1", "name": "execute_sql", "args": {}}] + ) + ], + iteration_count=0, + ) + assert graph.router(state) == "inner_tool" + + +def test_router_to_termination() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", tool_calls=[{"id": "1", "name": "execute_sql", "args": {}}] + ) + ], + iteration_count=3, # at max + ) + assert graph.router(state) == "termination" + + +def test_router_to_end_no_tool_calls() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + messages=[AIMessage(content="final answer")], + ) + assert graph.router(state) == END + + +def test_router_to_end_empty_messages() -> None: + graph = _make_graph() + state = DataFabricSubgraphState(messages=[]) + assert graph.router(state) == END + + +def test_router_to_end_human_message() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + messages=[HumanMessage(content="hello")], + ) + assert graph.router(state) == END + + +# --------------------------------------------------------------------------- +# DataFabricGraph — tool_router +# --------------------------------------------------------------------------- + + +def test_tool_router_success_ends() -> None: + graph = _make_graph() + state = DataFabricSubgraphState(last_tool_success=True) + assert graph.tool_router(state) == END + + +def test_tool_router_failure_retries() -> None: + graph = _make_graph() + state = DataFabricSubgraphState(last_tool_success=False) + assert graph.tool_router(state) == "inner_llm" + + +# --------------------------------------------------------------------------- +# DataFabricGraph — termination_node +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_termination_node_basic() -> None: + graph = _make_graph() + state = DataFabricSubgraphState(iteration_count=5) + result = await graph.termination_node(state) + msg = result["messages"][0] + assert isinstance(msg, AIMessage) + assert "5 SQL attempts" in msg.content + assert "rephrasing" in msg.content + + +@pytest.mark.asyncio +async def test_termination_node_with_error_info() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + iteration_count=3, + last_error_category="bad_sql", + last_error_detail="Invalid column 'foo'", + ) + result = await graph.termination_node(state) + msg = result["messages"][0] + assert "bad_sql" in msg.content + assert "Invalid column 'foo'" in msg.content + + +# --------------------------------------------------------------------------- +# DataFabricGraph — tool_node +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_tool_node_no_tool_calls() -> None: + graph = _make_graph() + state = DataFabricSubgraphState( + messages=[HumanMessage(content="hi")], + iteration_count=2, + ) + result = await graph.tool_node(state) + assert result["iteration_count"] == 2 + assert "messages" not in result + + +@pytest.mark.asyncio +async def test_tool_node_success() -> None: + graph = _make_graph() + # Mock the execute_sql_tool to return a success result + graph._execute_sql_tool = MagicMock() + graph._execute_sql_tool.ainvoke = AsyncMock( + return_value={"records": [{"id": 1}], "total_count": 1, "sql_query": "SELECT 1"} + ) + + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tc1", + "name": "execute_sql", + "args": {"sql_query": "SELECT 1"}, + } + ], + ) + ], + iteration_count=0, + ) + result = await graph.tool_node(state) + assert result["last_tool_success"] is True + assert result["iteration_count"] == 1 + assert len(result["messages"]) == 1 + assert isinstance(result["messages"][0], ToolMessage) + + +@pytest.mark.asyncio +async def test_tool_node_failure_extracts_category() -> None: + graph = _make_graph() + graph._execute_sql_tool = MagicMock() + graph._execute_sql_tool.ainvoke = AsyncMock( + return_value={ + "records": [], + "total_count": 0, + "error": "[SQL_VALIDATION] (category: bad_sql) Invalid column", + "sql_query": "SELECT bad", + } + ) + + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tc1", + "name": "execute_sql", + "args": {"sql_query": "SELECT bad"}, + } + ], + ) + ], + ) + result = await graph.tool_node(state) + assert result["last_tool_success"] is False + assert result["last_error_category"] == "bad_sql" + assert "SQL_VALIDATION" in result["last_error_detail"] + + +@pytest.mark.asyncio +async def test_tool_node_value_error() -> None: + graph = _make_graph() + graph._execute_sql_tool = MagicMock() + graph._execute_sql_tool.ainvoke = AsyncMock(side_effect=ValueError("bad input")) + + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"id": "tc1", "name": "execute_sql", "args": {"sql_query": "X"}} + ], + ) + ], + ) + result = await graph.tool_node(state) + assert result["last_tool_success"] is False + assert "bad input" in result["last_error_detail"] + + +@pytest.mark.asyncio +async def test_tool_node_preserves_prior_error_on_no_new_error() -> None: + graph = _make_graph() + graph._execute_sql_tool = MagicMock() + # Return empty records with no error — not a success (total_count=0) but no error string + graph._execute_sql_tool.ainvoke = AsyncMock( + return_value={"records": [], "total_count": 0, "sql_query": "SELECT 1"} + ) + + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", + tool_calls=[{"id": "tc1", "name": "execute_sql", "args": {}}], + ) + ], + last_error_category="prior_cat", + last_error_detail="prior detail", + ) + result = await graph.tool_node(state) + # No new error, so prior values should be preserved + assert result["last_error_category"] == "prior_cat" + assert result["last_error_detail"] == "prior detail" + + +# --------------------------------------------------------------------------- +# DataFabricGraph.create +# --------------------------------------------------------------------------- + + +def test_create_returns_compiled_graph() -> None: + llm = MagicMock(spec=["model_copy"]) + bound = MagicMock() + copy = MagicMock() + copy.bind_tools = MagicMock(return_value=bound) + llm.model_copy.return_value = copy + + with patch( + "uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph.datafabric_prompt_builder" + ) as mock_pb: + mock_pb.build.return_value = "prompt" + compiled = DataFabricGraph.create( + llm, [_make_entity("T")], _make_entities_service() + ) + + assert compiled is not None + + +# --------------------------------------------------------------------------- +# CATEGORY_MARKER extraction +# --------------------------------------------------------------------------- + + +def test_category_marker_extraction() -> None: + error_str = "[SQL_VALIDATION] (category: bad_sql) Invalid column" + start = error_str.index(CATEGORY_MARKER) + len(CATEGORY_MARKER) + end = error_str.index(")", start) + assert error_str[start:end] == "bad_sql" + + +def test_category_marker_not_present() -> None: + error_str = "some generic error" + assert CATEGORY_MARKER not in error_str + + +def test_category_marker_malformed_no_closing_paren() -> None: + """Malformed error with marker but no closing ')' should not crash.""" + error_str = "[ERR] (category: bad_sql oops" + assert CATEGORY_MARKER in error_str + start = error_str.index(CATEGORY_MARKER) + len(CATEGORY_MARKER) + end = error_str.find(")", start) + assert end == -1 # no closing paren found + + +@pytest.mark.asyncio +async def test_tool_node_malformed_category_does_not_crash() -> None: + """Tool node should not crash on malformed category marker.""" + graph = _make_graph() + graph._execute_sql_tool = MagicMock() + graph._execute_sql_tool.ainvoke = AsyncMock( + return_value={ + "records": [], + "total_count": 0, + "error": "[ERR] (category: bad_sql oops no closing paren", + "sql_query": "SELECT bad", + } + ) + + state = DataFabricSubgraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tc1", + "name": "execute_sql", + "args": {"sql_query": "SELECT bad"}, + } + ], + ) + ], + ) + result = await graph.tool_node(state) + assert result["last_tool_success"] is False + # Category should be empty since parsing couldn't find closing ')' + assert result["last_error_category"] == "" + + +# --------------------------------------------------------------------------- +# QueryExecutor — error path with OTEL span active +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_query_executor_error_with_span_sets_attributes() -> None: + """Error path when an OTEL span is active — verifies span attributes are set.""" + enriched = _make_enriched_exception("sql error") + + svc = _make_entities_service(error=enriched) + qe = QueryExecutor(svc, [_make_entity("T")]) + + fake_df_error = _FakeDataFabricError( + code="SQL_VALIDATION", + message="bad column", + trace_id="trace-1", + category=_FakeCategory.BAD_SQL, + is_bad_sql=True, + ) + + mock_span = MagicMock() + mock_tracer = MagicMock() + mock_tracer.start_as_current_span.return_value.__enter__ = MagicMock( + return_value=mock_span + ) + mock_tracer.start_as_current_span.return_value.__exit__ = MagicMock( + return_value=False + ) + + with ( + patch("opentelemetry.trace.get_tracer", return_value=mock_tracer), + patch( + "uipath_langchain.agent.tools.datafabric_tool.datafabric_subgraph.DataFabricError" + ) as mock_dfe_cls, + ): + mock_dfe_cls.from_enriched_exception.return_value = fake_df_error + result = await qe("SELECT bad") + + assert result["records"] == [] + # Verify span attributes were set + set_attr_calls = { + call.args[0]: call.args[1] for call in mock_span.set_attribute.call_args_list + } + assert set_attr_calls["df.success"] is False + assert set_attr_calls["df.error.code"] == "SQL_VALIDATION" + assert set_attr_calls["df.error.category"] == "bad_sql" + assert set_attr_calls["df.error.type"] == "EnrichedException" + assert set_attr_calls["df.error.digest"] != "bad column" + assert "df.error.message" not in set_attr_calls + assert "df.error.trace_id" not in set_attr_calls + mock_span.record_exception.assert_not_called() + mock_span.set_status.assert_called_once() + + +@pytest.mark.asyncio +async def test_query_executor_error_with_span_no_df_error() -> None: + """Error path with OTEL span but a plain (non-EnrichedException) error.""" + svc = _make_entities_service(error=RuntimeError("timeout")) + qe = QueryExecutor(svc, [_make_entity("T")]) + + mock_span = MagicMock() + mock_tracer = MagicMock() + mock_tracer.start_as_current_span.return_value.__enter__ = MagicMock( + return_value=mock_span + ) + mock_tracer.start_as_current_span.return_value.__exit__ = MagicMock( + return_value=False + ) + + with patch("opentelemetry.trace.get_tracer", return_value=mock_tracer): + result = await qe("SELECT 1") + + assert result["error"] == "timeout" + set_attr_calls = { + call.args[0]: call.args[1] for call in mock_span.set_attribute.call_args_list + } + assert set_attr_calls["df.success"] is False + assert set_attr_calls["df.error.type"] == "RuntimeError" + assert set_attr_calls["df.error.digest"] != "timeout" + assert "df.error.raw" not in set_attr_calls + # No df_error attributes should be set + assert "df.error.code" not in set_attr_calls + + +# --------------------------------------------------------------------------- +# DataFabricGraph — llm_node +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_llm_node_invokes_inner_llm() -> None: + graph = _make_graph() + mock_response = AIMessage(content="I'll query the database") + graph._inner_llm = MagicMock() + graph._inner_llm.ainvoke = AsyncMock(return_value=mock_response) + + state = DataFabricSubgraphState(messages=[HumanMessage(content="How many orders?")]) + result = await graph.llm_node(state) + assert result["messages"] == [mock_response] + # Verify system message was prepended + call_args = graph._inner_llm.ainvoke.call_args[0][0] + assert call_args[0] == graph._system_message + assert call_args[1].content == "How many orders?"