feat: init context registry

Author: milind-jain-uipath

src/uipath_langchain/agent/react/init_context_registry.py

@@ -0,0 +1,69 @@
+"""Registry for resource types that contribute init-time context.
+
+Resource modules self-register by calling ``register_init_context_provider``
+at module level. The INIT node calls ``gather_init_context`` to collect
+additional context from all registered providers, without needing to know
+which resource types participate.
+"""
+
+import logging
+from typing import Protocol, Sequence
+
+from uipath.agent.models.agent import BaseAgentResourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class InitContextProvider(Protocol):
+    """Contract for a resource type's init-time context builder."""
+
+    async def __call__(
+        self,
+        resources: Sequence[BaseAgentResourceConfig],
+    ) -> str | None: ...
+
+
+_registry: dict[str, InitContextProvider] = {}
+
+
+def register_init_context_provider(
+    name: str,
+    provider: InitContextProvider,
+) -> None:
+    """Register a provider that contributes init-time context.
+
+    Args:
+        name: Identifier for logging and deduplication.
+        provider: Async callable matching ``InitContextProvider``.
+    """
+    if name in _registry:
+        raise ValueError(f"Init context provider '{name}' is already registered")
+    _registry[name] = provider
+    logger.debug("Registered init context provider: %s", name)
+
+
+async def gather_init_context(
+    resources: Sequence[BaseAgentResourceConfig],
+) -> str | None:
+    """Call all registered providers and merge their context contributions.
+
+    Args:
+        resources: The agent's resource configs.
+
+    Returns:
+        Merged context string, or None if no provider contributed.
+    """
+    parts: list[str] = []
+    for name, provider in _registry.items():
+        try:
+            result = await provider(resources)
+            if result:
+                parts.append(result)
+                logger.info(
+                    "Init context provider '%s' contributed %d chars",
+                    name,
+                    len(result),
+                )
+        except Exception:
+            logger.exception("Init context provider '%s' failed; skipping", name)
+    return "\n\n".join(parts) if parts else None

src/uipath_langchain/agent/react/init_node.py

@@ -24,32 +24,19 @@ def create_init_node(
     resources_for_init: AgentResources | None = None,
 ):
     async def graph_state_init(state: Any) -> Any:
-        # --- Data Fabric schema fetch (INIT-time) ---
-        schema_context: str | None = None
+        # --- Gather init-time context from registered providers ---
+        additional_context: str | None = None
         if resources_for_init:
-            from uipath_langchain.agent.tools.datafabric_tool import (
-                fetch_entity_schemas,
-                format_schemas_for_context,
-                get_datafabric_entity_identifiers_from_resources,
-            )
+            from .init_context_registry import gather_init_context
 
-            entity_identifiers = get_datafabric_entity_identifiers_from_resources(
-                resources_for_init
-            )
-            if entity_identifiers:
-                logger.info(
-                    "Fetching Data Fabric schemas for %d identifier(s)",
-                    len(entity_identifiers),
-                )
-                entities = await fetch_entity_schemas(entity_identifiers)
-                schema_context = format_schemas_for_context(entities)
+            additional_context = await gather_init_context(resources_for_init)
 
         # --- Resolve messages ---
         resolved_messages: Sequence[SystemMessage | HumanMessage] | Overwrite
         if callable(messages):
-            if schema_context:
+            if additional_context:
                 resolved_messages = list(
-                    messages(state, additional_context=schema_context)
+                    messages(state, additional_context=additional_context)
                 )
             else:
                 resolved_messages = list(messages(state))

src/uipath_langchain/agent/tools/__init__.py

@@ -2,7 +2,6 @@
 
 from .context_tool import create_context_tool
 from .datafabric_tool import (
-    create_datafabric_tools,
     fetch_entity_schemas,
     format_schemas_for_context,
     get_datafabric_contexts,
@@ -23,7 +22,6 @@
     "create_tools_from_resources",
     "create_tool_node",
     "create_context_tool",
-    "create_datafabric_tools",
     "open_mcp_tools",
     "create_process_tool",
     "create_integration_tool",

src/uipath_langchain/agent/tools/context_tool.py

@@ -12,6 +12,7 @@
 from uipath.agent.models.agent import (
     AgentContextResourceConfig,
     AgentContextRetrievalMode,
+    AgentContextType,
     AgentToolArgumentArgumentProperties,
     AgentToolArgumentProperties,
 )
@@ -134,21 +135,25 @@ def is_static_query(resource: AgentContextResourceConfig) -> bool:
     return resource.settings.query.variant.lower() == "static"
 
 
-def create_context_tool(resource: AgentContextResourceConfig) -> StructuredTool:
-    tool_name = sanitize_tool_name(resource.name)
+def create_context_tool(
+    resource: AgentContextResourceConfig,
+) -> StructuredTool | BaseTool:
+    if resource.context_type == AgentContextType.DATA_FABRIC_ENTITY_SET:
+        from .datafabric_tool import create_datafabric_query_tool
+
+        return create_datafabric_query_tool()
+
     assert resource.settings is not None
+    tool_name = sanitize_tool_name(resource.name)
     retrieval_mode = resource.settings.retrieval_mode.lower()
+
     if retrieval_mode == AgentContextRetrievalMode.DEEP_RAG.value.lower():
         return handle_deep_rag(tool_name, resource)
-    elif retrieval_mode == AgentContextRetrievalMode.BATCH_TRANSFORM.value.lower():
+
+    if retrieval_mode == AgentContextRetrievalMode.BATCH_TRANSFORM.value.lower():
         return handle_batch_transform(tool_name, resource)
-    elif retrieval_mode == AgentContextRetrievalMode.DATA_FABRIC.value.lower():
-        raise ValueError(
-            "Data Fabric context should be handled via create_datafabric_tools(), "
-            "not create_context_tool()"
-        )
-    else:
-        return handle_semantic_search(tool_name, resource)
+
+    return handle_semantic_search(tool_name, resource)
 
 
 def handle_semantic_search(

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

@@ -1,17 +1,50 @@
 """Data Fabric tool module for entity-based SQL queries."""
 
+import logging
+from typing import Sequence
+
+from uipath.agent.models.agent import BaseAgentResourceConfig
+
+from uipath_langchain.agent.react.init_context_registry import (
+    register_init_context_provider,
+)
+
 from .datafabric_tool import (
-    create_datafabric_tools,
+    create_datafabric_query_tool,
     fetch_entity_schemas,
-    format_schemas_for_context,
     get_datafabric_contexts,
     get_datafabric_entity_identifiers_from_resources,
 )
+from .schema_context import format_schemas_for_context
 
 __all__ = [
-    "create_datafabric_tools",
+    "create_datafabric_query_tool",
     "fetch_entity_schemas",
     "format_schemas_for_context",
     "get_datafabric_contexts",
     "get_datafabric_entity_identifiers_from_resources",
 ]
+
+_logger = logging.getLogger(__name__)
+
+
+# --- Init-time context self-registration ---
+
+
+async def _datafabric_init_context_provider(
+    resources: Sequence[BaseAgentResourceConfig],
+) -> str | None:
+    """Fetch and format DataFabric entity schemas for system prompt injection."""
+    entity_identifiers = get_datafabric_entity_identifiers_from_resources(resources)
+    if not entity_identifiers:
+        return None
+
+    _logger.info(
+        "Fetching Data Fabric schemas for %d identifier(s)",
+        len(entity_identifiers),
+    )
+    entities = await fetch_entity_schemas(entity_identifiers)
+    return format_schemas_for_context(entities)
+
+
+register_init_context_provider("datafabric", _datafabric_init_context_provider)