dao-ai 0.0.6__py3-none-any.whl → 0.0.7__py3-none-any.whl

dao_ai/nodes.py

@@ -10,6 +10,7 @@ from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt import create_react_agent
+from langgraph.runtime import Runtime
 from langmem import create_manage_memory_tool, create_search_memory_tool
 from langmem.short_term import SummarizationNode
 from langmem.short_term.summarization import TokenCounter
@@ -26,7 +27,7 @@ from dao_ai.config import (
 from dao_ai.guardrails import reflection_guardrail, with_guardrails
 from dao_ai.hooks.core import create_hooks
 from dao_ai.prompts import make_prompt
-from dao_ai.state import IncomingState, SharedState
+from dao_ai.state import Context, IncomingState, SharedState
 from dao_ai.tools import create_tools
 
 
@@ -53,6 +54,7 @@ def summarization_node(app_model: AppModel) -> RunnableLike:
     )
 
     summarization_model: LanguageModelLike = chat_history.model.as_chat_model()
+
     node: RunnableLike = SummarizationNode(
         model=summarization_model,
         max_tokens=max_tokens,
@@ -67,7 +69,7 @@ def summarization_node(app_model: AppModel) -> RunnableLike:
 
 
 def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLike:
-    def call_agent(state: SharedState, config: RunnableConfig) -> SharedState:
+    def call_agent(state: SharedState, runtime: Runtime[Context]) -> SharedState:
         logger.debug(f"Calling agent {agent.name} with summarized messages")
 
         # Get the summarized messages from the summarization node
@@ -79,7 +81,7 @@ def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLi
             "messages": messages,
         }
 
-        response: dict[str, Any] = agent.invoke(input=input, config=config)
+        response: dict[str, Any] = agent.invoke(input=input, context=runtime.context)
         response_messages = response.get("messages", [])
         logger.debug(f"Agent returned {len(response_messages)} messages")
 
@@ -147,9 +149,9 @@ def create_agent_node(
         prompt=make_prompt(agent.prompt),
         tools=tools,
         store=True,
-        state_schema=SharedState,
-        config_schema=RunnableConfig,
         checkpointer=True,
+        state_schema=SharedState,
+        context_schema=Context,
         pre_model_hook=pre_agent_hook,
         post_model_hook=post_agent_hook,
     )
@@ -165,17 +167,17 @@ def create_agent_node(
     chat_history: ChatHistoryModel = app.chat_history
 
     if chat_history is None:
+        logger.debug("No chat history configured, using compiled agent directly")
         agent_node = compiled_agent
     else:
+        logger.debug("Creating agent node with chat history summarization")
         workflow: StateGraph = StateGraph(
             SharedState,
             config_schema=RunnableConfig,
             input=SharedState,
             output=SharedState,
         )
-        workflow.add_node(
-            "summarization", summarization_node(chat_history=chat_history)
-        )
+        workflow.add_node("summarization", summarization_node(app))
         workflow.add_node(
             "agent",
             call_agent_with_summarized_messages(agent=compiled_agent),
@@ -191,7 +193,7 @@ def message_hook_node(config: AppConfig) -> RunnableLike:
     message_hooks: Sequence[Callable[..., Any]] = create_hooks(config.app.message_hooks)
 
     @mlflow.trace()
-    def message_hook(state: IncomingState, config: RunnableConfig) -> SharedState:
+    def message_hook(state: IncomingState, runtime: Runtime[Context]) -> SharedState:
         logger.debug("Running message validation")
         response: dict[str, Any] = {"is_valid": True, "message_error": None}
 
@@ -201,7 +203,7 @@ def message_hook_node(config: AppConfig) -> RunnableLike:
                 try:
                     hook_response: dict[str, Any] = message_hook(
                         state=state,
-                        config=config,
+                        runtime=runtime,
                     )
                     response.update(hook_response)
                     logger.debug(f"Hook response: {hook_response}")

dao_ai/providers/databricks.py

@@ -355,6 +355,20 @@ class DatabricksProvider(ServiceProvider):
 
         latest_version: int = get_latest_model_version(registered_model_name)
 
+        # Check if endpoint exists to determine deployment strategy
+        endpoint_exists: bool = False
+        try:
+            agents.get_deployments(endpoint_name)
+            endpoint_exists = True
+            logger.debug(
+                f"Endpoint {endpoint_name} already exists, updating without tags to avoid conflicts..."
+            )
+        except Exception:
+            logger.debug(
+                f"Endpoint {endpoint_name} doesn't exist, creating new with tags..."
+            )
+
+        # Deploy - skip tags for existing endpoints to avoid conflicts
         agents.deploy(
             endpoint_name=endpoint_name,
             model_name=registered_model_name,
@@ -362,7 +376,7 @@ class DatabricksProvider(ServiceProvider):
             scale_to_zero=scale_to_zero,
             environment_vars=environment_vars,
             workload_size=workload_size,
-            tags=tags,
+            tags=tags if not endpoint_exists else None,
         )
 
         registered_model_name: str = config.app.registered_model.full_name
@@ -526,9 +540,75 @@ class DatabricksProvider(ServiceProvider):
                 columns_to_sync=vector_store.columns,
             )
         else:
-            self.vsc.get_index(
+            logger.debug(
+                f"Index {vector_store.index.full_name} already exists, checking status and syncing..."
+            )
+            index = self.vsc.get_index(
                 vector_store.endpoint.name, vector_store.index.full_name
-            ).sync()
+            )
+
+            # Wait for index to be in a syncable state
+            import time
+
+            max_wait_time = 600  # 10 minutes
+            wait_interval = 10  # 10 seconds
+            elapsed = 0
+
+            while elapsed < max_wait_time:
+                try:
+                    index_status = index.describe()
+                    pipeline_status = index_status.get("status", {}).get(
+                        "detailed_state", "UNKNOWN"
+                    )
+                    logger.debug(f"Index pipeline status: {pipeline_status}")
+
+                    if pipeline_status in [
+                        "COMPLETED",
+                        "FAILED",
+                        "CANCELED",
+                        "ONLINE_PIPELINE_FAILED",
+                    ]:
+                        logger.debug(
+                            f"Index is ready to sync (status: {pipeline_status})"
+                        )
+                        break
+                    elif pipeline_status in [
+                        "WAITING_FOR_RESOURCES",
+                        "PROVISIONING",
+                        "INITIALIZING",
+                        "INDEXING",
+                        "ONLINE",
+                    ]:
+                        logger.debug(
+                            f"Index not ready yet (status: {pipeline_status}), waiting {wait_interval} seconds..."
+                        )
+                        time.sleep(wait_interval)
+                        elapsed += wait_interval
+                    else:
+                        logger.warning(
+                            f"Unknown pipeline status: {pipeline_status}, attempting sync anyway"
+                        )
+                        break
+                except Exception as status_error:
+                    logger.warning(
+                        f"Could not check index status: {status_error}, attempting sync anyway"
+                    )
+                    break
+
+            if elapsed >= max_wait_time:
+                logger.warning(
+                    f"Timed out waiting for index to be ready after {max_wait_time} seconds"
+                )
+
+            # Now attempt to sync
+            try:
+                index.sync()
+                logger.debug("Index sync completed successfully")
+            except Exception as sync_error:
+                if "not ready to sync yet" in str(sync_error).lower():
+                    logger.warning(f"Index still not ready to sync: {sync_error}")
+                else:
+                    raise sync_error
 
         logger.debug(
             f"index {vector_store.index.full_name} on table {vector_store.source_table.full_name} is ready"

dao_ai/state.py

@@ -1,6 +1,7 @@
 from langchain_core.messages import AnyMessage
 from langgraph.graph import MessagesState
 from langgraph.managed import RemainingSteps
+from pydantic import BaseModel
 
 
 class IncomingState(MessagesState): ...
@@ -29,3 +30,9 @@ class SharedState(MessagesState):
 
     is_valid: bool  # message validation node
     message_error: str
+
+
+class Context(BaseModel):
+    user_id: str | None = None
+    thread_id: str | None = None
+    store_num: int | None = None

dao_ai/tools/__init__.py

@@ -1,14 +1,12 @@
 from dao_ai.hooks.core import create_hooks
 from dao_ai.tools.agent import create_agent_endpoint_tool
 from dao_ai.tools.core import (
-    create_factory_tool,
-    create_mcp_tools,
-    create_python_tool,
     create_tools,
-    create_uc_tools,
     search_tool,
 )
 from dao_ai.tools.genie import create_genie_tool
+from dao_ai.tools.mcp import create_mcp_tools
+from dao_ai.tools.python import create_factory_tool, create_python_tool
 from dao_ai.tools.time import (
     add_time_tool,
     current_time_tool,
@@ -18,6 +16,7 @@ from dao_ai.tools.time import (
     time_in_timezone_tool,
     time_until_tool,
 )
+from dao_ai.tools.unity_catalog import create_uc_tools
 from dao_ai.tools.vector_search import create_vector_search_tool
 
 __all__ = [

dao_ai/tools/core.py

@@ -1,118 +1,15 @@
-import asyncio
 from collections import OrderedDict
-from typing import Any, Callable, Optional, Sequence
+from typing import Sequence
 
-from databricks_langchain import (
-    DatabricksFunctionClient,
-    UCFunctionToolkit,
-)
 from langchain_community.tools import DuckDuckGoSearchRun
-from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.base import RunnableLike
-from langchain_core.tools import BaseTool
-from langchain_core.tools import tool as create_tool
-from langchain_mcp_adapters.client import MultiServerMCPClient
-from langgraph.prebuilt.interrupt import HumanInterrupt, HumanInterruptConfig
-from langgraph.types import interrupt
 from loguru import logger
-from mcp.types import ListToolsResult, Tool
 
 from dao_ai.config import (
     AnyTool,
-    BaseFunctionModel,
-    FactoryFunctionModel,
-    HumanInTheLoopModel,
-    McpFunctionModel,
-    PythonFunctionModel,
     ToolModel,
-    TransportType,
-    UnityCatalogFunctionModel,
 )
 from dao_ai.hooks.core import create_hooks
-from dao_ai.utils import load_function
-
-
-def add_human_in_the_loop(
-    tool: RunnableLike,
-    *,
-    interrupt_config: HumanInterruptConfig | None = None,
-    review_prompt: Optional[str] = "Please review the tool call",
-) -> BaseTool:
-    """
-    Wrap a tool with human-in-the-loop functionality.
-    This function takes a tool (either a callable or a BaseTool instance) and wraps it
-    with a human-in-the-loop mechanism. When the tool is invoked, it will first
-    request human review before executing the tool's logic. The human can choose to
-    accept, edit the input, or provide a custom response.
-
-    Args:
-        tool (Callable[..., Any] | BaseTool): _description_
-        interrupt_config (HumanInterruptConfig | None, optional): _description_. Defaults to None.
-
-    Raises:
-        ValueError: _description_
-
-    Returns:
-        BaseTool: _description_
-    """
-    if not isinstance(tool, BaseTool):
-        tool = create_tool(tool)
-
-    if interrupt_config is None:
-        interrupt_config = {
-            "allow_accept": True,
-            "allow_edit": True,
-            "allow_respond": True,
-        }
-
-    logger.debug(f"Wrapping tool {tool} with human-in-the-loop functionality")
-
-    @create_tool(tool.name, description=tool.description, args_schema=tool.args_schema)
-    def call_tool_with_interrupt(config: RunnableConfig, **tool_input) -> Any:
-        logger.debug(f"call_tool_with_interrupt: {tool.name} with input: {tool_input}")
-        request: HumanInterrupt = {
-            "action_request": {
-                "action": tool.name,
-                "args": tool_input,
-            },
-            "config": interrupt_config,
-            "description": review_prompt,
-        }
-
-        logger.debug(f"Human interrupt request: {request}")
-        response: dict[str, Any] = interrupt([request])[0]
-        logger.debug(f"Human interrupt response: {response}")
-
-        if response["type"] == "accept":
-            tool_response = tool.invoke(tool_input, config=config)
-        elif response["type"] == "edit":
-            tool_input = response["args"]["args"]
-            tool_response = tool.invoke(tool_input, config=config)
-        elif response["type"] == "response":
-            user_feedback = response["args"]
-            tool_response = user_feedback
-        else:
-            raise ValueError(f"Unknown interrupt response type: {response['type']}")
-
-        return tool_response
-
-    return call_tool_with_interrupt
-
-
-def as_human_in_the_loop(
-    tool: RunnableLike, function: BaseFunctionModel | str
-) -> RunnableLike:
-    if isinstance(function, BaseFunctionModel):
-        human_in_the_loop: HumanInTheLoopModel | None = function.human_in_the_loop
-        if human_in_the_loop:
-            logger.debug(f"Adding human-in-the-loop to tool: {tool.name}")
-            tool = add_human_in_the_loop(
-                tool=tool,
-                interrupt_config=human_in_the_loop.interupt_config,
-                review_prompt=human_in_the_loop.review_prompt,
-            )
-    return tool
-
 
 tool_registry: dict[str, Sequence[RunnableLike]] = {}
 
@@ -157,196 +54,6 @@ def create_tools(tool_models: Sequence[ToolModel]) -> Sequence[RunnableLike]:
     return all_tools
 
 
-def create_mcp_tools(
-    function: McpFunctionModel,
-) -> Sequence[RunnableLike]:
-    """
-    Create tools for invoking Databricks MCP functions.
-
-    Uses session-based approach to handle authentication token expiration properly.
-    """
-    logger.debug(f"create_mcp_tools: {function}")
-
-    def _create_fresh_connection() -> dict[str, Any]:
-        logger.debug("Creating fresh connection...")
-        """Create connection config with fresh authentication headers."""
-        if function.transport == TransportType.STDIO:
-            return {
-                "command": function.command,
-                "args": function.args,
-                "transport": function.transport,
-            }
-
-        # For HTTP transport, generate fresh headers
-        headers = function.headers.copy() if function.headers else {}
-
-        if "Authorization" not in headers:
-            logger.debug("Generating fresh authentication token for MCP function")
-
-            from dao_ai.config import value_of
-            from dao_ai.providers.databricks import DatabricksProvider
-
-            try:
-                provider = DatabricksProvider(
-                    workspace_host=value_of(function.workspace_host),
-                    client_id=value_of(function.client_id),
-                    client_secret=value_of(function.client_secret),
-                    pat=value_of(function.pat),
-                )
-                headers["Authorization"] = f"Bearer {provider.create_token()}"
-                logger.debug("Generated fresh authentication token")
-            except Exception as e:
-                logger.error(f"Failed to create fresh token: {e}")
-        else:
-            logger.debug("Using existing authentication token")
-
-        response = {
-            "url": function.url,
-            "transport": function.transport,
-            "headers": headers,
-        }
-
-        return response
-
-    # Get available tools from MCP server
-    async def _list_mcp_tools():
-        connection = _create_fresh_connection()
-        client = MultiServerMCPClient({function.name: connection})
-
-        try:
-            async with client.session(function.name) as session:
-                return await session.list_tools()
-        except Exception as e:
-            logger.error(f"Failed to list MCP tools: {e}")
-            return []
-
-    try:
-        mcp_tools: list | ListToolsResult = asyncio.run(_list_mcp_tools())
-        if isinstance(mcp_tools, ListToolsResult):
-            mcp_tools = mcp_tools.tools
-
-        logger.debug(f"Retrieved {len(mcp_tools)} MCP tools")
-    except Exception as e:
-        logger.error(f"Failed to get tools from MCP server: {e}")
-        raise RuntimeError(
-            f"Failed to list MCP tools for function '{function.name}' with transport '{function.transport}' and URL '{function.url}': {e}"
-        )
-
-    # Create wrapper tools with fresh session per invocation
-    def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
-        @create_tool(
-            mcp_tool.name,
-            description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
-            args_schema=mcp_tool.inputSchema,
-        )
-        def tool_wrapper(**kwargs):
-            """Execute MCP tool with fresh session and authentication."""
-            logger.debug(f"Invoking MCP tool {mcp_tool.name} with fresh session")
-
-            async def _invoke():
-                connection = _create_fresh_connection()
-                client = MultiServerMCPClient({function.name: connection})
-
-                try:
-                    async with client.session(function.name) as session:
-                        return await session.call_tool(mcp_tool.name, kwargs)
-                except Exception as e:
-                    logger.error(f"MCP tool {mcp_tool.name} failed: {e}")
-                    raise
-
-            return asyncio.run(_invoke())
-
-        return as_human_in_the_loop(tool_wrapper, function)
-
-    return [_create_tool_wrapper(tool) for tool in mcp_tools]
-
-
-def create_factory_tool(
-    function: FactoryFunctionModel,
-) -> RunnableLike:
-    """
-    Create a factory tool from a FactoryFunctionModel.
-    This factory function dynamically loads a Python function and returns it as a callable tool.
-    Args:
-        function: FactoryFunctionModel instance containing the function details
-    Returns:
-        A callable tool function that wraps the specified factory function
-    """
-    logger.debug(f"create_factory_tool: {function}")
-
-    factory: Callable[..., Any] = load_function(function_name=function.full_name)
-    tool: Callable[..., Any] = factory(**function.args)
-    tool = as_human_in_the_loop(
-        tool=tool,
-        function=function,
-    )
-    return tool
-
-
-def create_python_tool(
-    function: PythonFunctionModel | str,
-) -> RunnableLike:
-    """
-    Create a Python tool from a Python function model.
-    This factory function wraps a Python function as a callable tool that can be
-    invoked by agents during reasoning.
-    Args:
-        function: PythonFunctionModel instance containing the function details
-    Returns:
-        A callable tool function that wraps the specified Python function
-    """
-    logger.debug(f"create_python_tool: {function}")
-
-    if isinstance(function, PythonFunctionModel):
-        function = function.full_name
-
-    # Load the Python function dynamically
-    tool: Callable[..., Any] = load_function(function_name=function)
-
-    tool = as_human_in_the_loop(
-        tool=tool,
-        function=function,
-    )
-    return tool
-
-
-def create_uc_tools(
-    function: UnityCatalogFunctionModel | str,
-) -> Sequence[RunnableLike]:
-    """
-    Create LangChain tools from Unity Catalog functions.
-
-    This factory function wraps Unity Catalog functions as LangChain tools,
-    making them available for use by agents. Each UC function becomes a callable
-    tool that can be invoked by the agent during reasoning.
-
-    Args:
-        function: UnityCatalogFunctionModel instance containing the function details
-
-    Returns:
-        A sequence of BaseTool objects that wrap the specified UC functions
-    """
-
-    logger.debug(f"create_uc_tools: {function}")
-
-    if isinstance(function, UnityCatalogFunctionModel):
-        function = function.full_name
-
-    client: DatabricksFunctionClient = DatabricksFunctionClient()
-
-    toolkit: UCFunctionToolkit = UCFunctionToolkit(
-        function_names=[function], client=client
-    )
-
-    tools = toolkit.tools or []
-
-    logger.debug(f"Retrieved tools: {tools}")
-
-    tools = [as_human_in_the_loop(tool=tool, function=function) for tool in tools]
-
-    return tools
-
-
 def search_tool() -> RunnableLike:
     logger.debug("search_tool")
     return DuckDuckGoSearchRun(output_format="list")

dao_ai/tools/human_in_the_loop.py

@@ -0,0 +1,96 @@
+from typing import Any, Optional
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.base import RunnableLike
+from langchain_core.tools import BaseTool
+from langchain_core.tools import tool as create_tool
+from langgraph.prebuilt.interrupt import HumanInterrupt, HumanInterruptConfig
+from langgraph.types import interrupt
+from loguru import logger
+
+from dao_ai.config import (
+    BaseFunctionModel,
+    HumanInTheLoopModel,
+)
+
+
+def add_human_in_the_loop(
+    tool: RunnableLike,
+    *,
+    interrupt_config: HumanInterruptConfig | None = None,
+    review_prompt: Optional[str] = "Please review the tool call",
+) -> BaseTool:
+    """
+    Wrap a tool with human-in-the-loop functionality.
+    This function takes a tool (either a callable or a BaseTool instance) and wraps it
+    with a human-in-the-loop mechanism. When the tool is invoked, it will first
+    request human review before executing the tool's logic. The human can choose to
+    accept, edit the input, or provide a custom response.
+
+    Args:
+        tool (Callable[..., Any] | BaseTool): _description_
+        interrupt_config (HumanInterruptConfig | None, optional): _description_. Defaults to None.
+
+    Raises:
+        ValueError: _description_
+
+    Returns:
+        BaseTool: _description_
+    """
+    if not isinstance(tool, BaseTool):
+        tool = create_tool(tool)
+
+    if interrupt_config is None:
+        interrupt_config = {
+            "allow_accept": True,
+            "allow_edit": True,
+            "allow_respond": True,
+        }
+
+    logger.debug(f"Wrapping tool {tool} with human-in-the-loop functionality")
+
+    @create_tool(tool.name, description=tool.description, args_schema=tool.args_schema)
+    def call_tool_with_interrupt(config: RunnableConfig, **tool_input) -> Any:
+        logger.debug(f"call_tool_with_interrupt: {tool.name} with input: {tool_input}")
+        request: HumanInterrupt = {
+            "action_request": {
+                "action": tool.name,
+                "args": tool_input,
+            },
+            "config": interrupt_config,
+            "description": review_prompt,
+        }
+
+        logger.debug(f"Human interrupt request: {request}")
+        response: dict[str, Any] = interrupt([request])[0]
+        logger.debug(f"Human interrupt response: {response}")
+
+        if response["type"] == "accept":
+            tool_response = tool.invoke(tool_input, config=config)
+        elif response["type"] == "edit":
+            tool_input = response["args"]["args"]
+            tool_response = tool.invoke(tool_input, config=config)
+        elif response["type"] == "response":
+            user_feedback = response["args"]
+            tool_response = user_feedback
+        else:
+            raise ValueError(f"Unknown interrupt response type: {response['type']}")
+
+        return tool_response
+
+    return call_tool_with_interrupt
+
+
+def as_human_in_the_loop(
+    tool: RunnableLike, function: BaseFunctionModel | str
+) -> RunnableLike:
+    if isinstance(function, BaseFunctionModel):
+        human_in_the_loop: HumanInTheLoopModel | None = function.human_in_the_loop
+        if human_in_the_loop:
+            logger.debug(f"Adding human-in-the-loop to tool: {tool.name}")
+            tool = add_human_in_the_loop(
+                tool=tool,
+                interrupt_config=human_in_the_loop.interupt_config,
+                review_prompt=human_in_the_loop.review_prompt,
+            )
+    return tool