dao-ai 0.0.35__py3-none-any.whl → 0.1.0__py3-none-any.whl

dao_ai/genie/core.py

@@ -0,0 +1,35 @@
+"""
+Core Genie service implementation.
+
+This module provides the concrete implementation of GenieServiceBase
+that wraps the Databricks Genie SDK.
+"""
+
+import mlflow
+from databricks_ai_bridge.genie import Genie, GenieResponse
+
+from dao_ai.genie.cache import CacheResult, GenieServiceBase
+
+
+class GenieService(GenieServiceBase):
+    """Concrete implementation of GenieServiceBase using the Genie SDK."""
+
+    genie: Genie
+
+    def __init__(self, genie: Genie) -> None:
+        self.genie = genie
+
+    @mlflow.trace(name="genie_ask_question")
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> CacheResult:
+        """Ask question to Genie and return CacheResult (no caching at this level)."""
+        response: GenieResponse = self.genie.ask_question(
+            question, conversation_id=conversation_id
+        )
+        # No caching at this level - return cache miss
+        return CacheResult(response=response, cache_hit=False, served_by=None)
+
+    @property
+    def space_id(self) -> str:
+        return self.genie.space_id

dao_ai/graph.py

@@ -1,263 +1,37 @@
-from typing import Sequence
+"""
+Graph creation utilities for DAO AI multi-agent orchestration.
 
-from langchain_core.language_models import LanguageModelLike
-from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import BaseTool
-from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.graph import END, StateGraph
-from langgraph.graph.state import CompiledStateGraph
-from langgraph.store.base import BaseStore
-from langgraph_supervisor import create_handoff_tool as supervisor_handoff_tool
-from langgraph_supervisor import create_supervisor
-from langgraph_swarm import create_handoff_tool as swarm_handoff_tool
-from langgraph_swarm import create_swarm
-from langmem import create_manage_memory_tool, create_search_memory_tool
-from loguru import logger
-
-from dao_ai.config import (
-    AgentModel,
-    AppConfig,
-    OrchestrationModel,
-    SupervisorModel,
-    SwarmModel,
-)
-from dao_ai.nodes import (
-    create_agent_node,
-    message_hook_node,
-)
-from dao_ai.prompts import make_prompt
-from dao_ai.state import Context, IncomingState, OutgoingState, SharedState
-from dao_ai.tools import create_tools
-
-
-def route_message(state: SharedState) -> str:
-    if not state["is_valid"]:
-        return END
-    return "orchestration"
-
-
-def _handoffs_for_agent(agent: AgentModel, config: AppConfig) -> Sequence[BaseTool]:
-    handoff_tools: list[BaseTool] = []
-
-    handoffs: dict[str, Sequence[AgentModel | str]] = (
-        config.app.orchestration.swarm.handoffs or {}
-    )
-    agent_handoffs: Sequence[AgentModel | str] = handoffs.get(agent.name)
-    if agent_handoffs is None:
-        agent_handoffs = config.app.agents
-
-    for handoff_to_agent in agent_handoffs:
-        if isinstance(handoff_to_agent, str):
-            handoff_to_agent = next(
-                iter(config.find_agents(lambda a: a.name == handoff_to_agent)), None
-            )
-
-        if handoff_to_agent is None:
-            logger.warning(
-                f"Handoff agent {handoff_to_agent} not found in configuration for agent {agent.name}"
-            )
-            continue
-        if agent.name == handoff_to_agent.name:
-            continue
-        logger.debug(
-            f"Creating handoff tool from agent {agent.name} to {handoff_to_agent.name}"
-        )
-
-        # Use handoff_prompt if provided, otherwise create default description
-        handoff_description = handoff_to_agent.handoff_prompt or (
-            handoff_to_agent.description
-            if handoff_to_agent.description
-            else "general assistance and questions"
-        )
-
-        handoff_tools.append(
-            swarm_handoff_tool(
-                agent_name=handoff_to_agent.name,
-                description=f"Ask {handoff_to_agent.name} for help with: "
-                + handoff_description,
-            )
-        )
-    return handoff_tools
-
-
-def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
-    logger.debug("Creating supervisor graph")
-    agents: list[CompiledStateGraph] = []
-    tools: Sequence[BaseTool] = []
-    for registered_agent in config.app.agents:
-        agents.append(
-            create_agent_node(
-                agent=registered_agent,
-                memory=config.app.orchestration.memory
-                if config.app.orchestration
-                else None,
-                chat_history=config.app.chat_history,
-                additional_tools=[],
-            )
-        )
-        # Use handoff_prompt if provided, otherwise create default description
-        handoff_description = registered_agent.handoff_prompt or (
-            registered_agent.description
-            if registered_agent.description
-            else f"General assistance with {registered_agent.name} related tasks"
-        )
-
-        tools.append(
-            supervisor_handoff_tool(
-                agent_name=registered_agent.name,
-                description=handoff_description,
-            )
-        )
-
-    orchestration: OrchestrationModel = config.app.orchestration
-    supervisor: SupervisorModel = orchestration.supervisor
-
-    tools += create_tools(orchestration.supervisor.tools)
-
-    store: BaseStore = None
-    if orchestration.memory and orchestration.memory.store:
-        store = orchestration.memory.store.as_store()
-        logger.debug(f"Using memory store: {store}")
-        namespace: tuple[str, ...] = ("memory",)
-
-        if orchestration.memory.store.namespace:
-            namespace = namespace + (orchestration.memory.store.namespace,)
-            logger.debug(f"Memory store namespace: {namespace}")
-            tools += [
-                create_manage_memory_tool(namespace=namespace),
-                create_search_memory_tool(namespace=namespace),
-            ]
-
-    checkpointer: BaseCheckpointSaver = None
-    if orchestration.memory and orchestration.memory.checkpointer:
-        checkpointer = orchestration.memory.checkpointer.as_checkpointer()
-        logger.debug(f"Using checkpointer: {checkpointer}")
+This module provides backwards-compatible imports from the refactored
+orchestration package. New code should import directly from:
+    - dao_ai.orchestration
+    - dao_ai.orchestration.supervisor
+    - dao_ai.orchestration.swarm
 
-    prompt: str = supervisor.prompt
+See: https://docs.langchain.com/oss/python/langchain/multi-agent
+"""
 
-    model: LanguageModelLike = supervisor.model.as_chat_model()
-    supervisor_workflow: StateGraph = create_supervisor(
-        supervisor_name="supervisor",
-        prompt=make_prompt(base_system_prompt=prompt),
-        agents=agents,
-        model=model,
-        tools=tools,
-        state_schema=SharedState,
-        config_schema=RunnableConfig,
-        output_mode="last_message",
-        add_handoff_messages=False,
-        add_handoff_back_messages=False,
-        context_schema=Context,
-        # output_mode="full",
-        # add_handoff_messages=True,
-        # add_handoff_back_messages=True,
-    )
-
-    supervisor_node: CompiledStateGraph = supervisor_workflow.compile(
-        checkpointer=checkpointer, store=store
-    )
-
-    workflow: StateGraph = StateGraph(
-        SharedState,
-        input_schema=IncomingState,
-        output_schema=OutgoingState,
-        context_schema=Context,
-    )
-
-    workflow.add_node("message_hook", message_hook_node(config=config))
-
-    workflow.add_node("orchestration", supervisor_node)
-    workflow.add_conditional_edges(
-        "message_hook",
-        route_message,
-        {
-            "orchestration": "orchestration",
-            END: END,
-        },
-    )
-    workflow.set_entry_point("message_hook")
-
-    return workflow.compile(checkpointer=checkpointer, store=store)
-
-
-def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
-    logger.debug("Creating swarm graph")
-    agents: list[CompiledStateGraph] = []
-    for registered_agent in config.app.agents:
-        handoff_tools: Sequence[BaseTool] = _handoffs_for_agent(
-            agent=registered_agent, config=config
-        )
-        agents.append(
-            create_agent_node(
-                agent=registered_agent,
-                memory=config.app.orchestration.memory
-                if config.app.orchestration
-                else None,
-                chat_history=config.app.chat_history,
-                additional_tools=handoff_tools,
-            )
-        )
-
-    orchestration: OrchestrationModel = config.app.orchestration
-    swarm: SwarmModel = orchestration.swarm
-
-    store: BaseStore = None
-    if orchestration.memory and orchestration.memory.store:
-        store = orchestration.memory.store.as_store()
-        logger.debug(f"Using memory store: {store}")
-
-    checkpointer: BaseCheckpointSaver = None
-    if orchestration.memory and orchestration.memory.checkpointer:
-        checkpointer = orchestration.memory.checkpointer.as_checkpointer()
-        logger.debug(f"Using checkpointer: {checkpointer}")
-
-    default_agent: AgentModel = swarm.default_agent
-    if isinstance(default_agent, AgentModel):
-        default_agent = default_agent.name
-
-    swarm_workflow: StateGraph = create_swarm(
-        agents=agents,
-        default_active_agent=default_agent,
-        state_schema=SharedState,
-        context_schema=Context,
-    )
-
-    swarm_node: CompiledStateGraph = swarm_workflow.compile(
-        checkpointer=checkpointer, store=store
-    )
-
-    workflow: StateGraph = StateGraph(
-        SharedState,
-        input_schema=IncomingState,
-        output_schema=OutgoingState,
-        context_schema=Context,
-    )
-
-    workflow.add_node("message_hook", message_hook_node(config=config))
-    workflow.add_node("orchestration", swarm_node)
-
-    workflow.add_conditional_edges(
-        "message_hook",
-        route_message,
-        {
-            "orchestration": "orchestration",
-            END: END,
-        },
-    )
+from langgraph.graph.state import CompiledStateGraph
 
-    workflow.set_entry_point("message_hook")
+from dao_ai.config import AppConfig
+from dao_ai.orchestration import create_orchestration_graph
 
-    return swarm_node
 
-    # return workflow.compile(checkpointer=checkpointer, store=store)
+def create_dao_ai_graph(config: AppConfig) -> CompiledStateGraph:
+    """
+    Create the main DAO AI graph based on the orchestration configuration.
 
+    This factory function creates either a supervisor or swarm graph
+    depending on the configuration.
 
-def create_dao_ai_graph(config: AppConfig) -> CompiledStateGraph:
-    orchestration: OrchestrationModel = config.app.orchestration
-    if orchestration.supervisor:
-        return _create_supervisor_graph(config)
+    Args:
+        config: The application configuration
 
-    if orchestration.swarm:
-        return _create_swarm_graph(config)
+    Returns:
+        A compiled LangGraph state machine
 
-    raise ValueError("No valid orchestration model found in the configuration.")
+    Note:
+        This function is provided for backwards compatibility.
+        New code should use `create_orchestration_graph` from
+        `dao_ai.orchestration` instead.
+    """
+    return create_orchestration_graph(config)

dao_ai/hooks/__init__.py

@@ -1,11 +1,17 @@
+"""
+Hook utilities for DAO AI.
+
+For validation hooks, use middleware instead:
+- dao_ai.middleware.UserIdValidationMiddleware
+- dao_ai.middleware.ThreadIdValidationMiddleware
+- dao_ai.middleware.FilterLastHumanMessageMiddleware
+"""
+
 from dao_ai.hooks.core import (
     create_hooks,
-    filter_last_human_message_hook,
     null_hook,
     null_initialization_hook,
     null_shutdown_hook,
-    require_thread_id_hook,
-    require_user_id_hook,
 )
 
 __all__ = [
@@ -13,7 +19,4 @@ __all__ = [
     "null_hook",
     "null_initialization_hook",
     "null_shutdown_hook",
-    "require_thread_id_hook",
-    "require_user_id_hook",
-    "filter_last_human_message_hook",
 ]

dao_ai/hooks/core.py

@@ -1,20 +1,32 @@
-import json
+"""
+Hook utilities for DAO AI.
+
+This module provides the create_hooks function for resolving FunctionHook
+references to callable functions. Individual validation hooks have been
+migrated to middleware - see dao_ai.middleware.message_validation.
+"""
+
 from typing import Any, Callable, Sequence
 
-from langchain_core.messages import BaseMessage, HumanMessage, RemoveMessage
-from langgraph.runtime import Runtime
 from loguru import logger
 
 from dao_ai.config import AppConfig, FunctionHook, PythonFunctionModel
-from dao_ai.messages import last_human_message
-from dao_ai.state import Context
 
 
 def create_hooks(
     function_hooks: FunctionHook | list[FunctionHook] | None,
 ) -> Sequence[Callable[..., Any]]:
+    """
+    Resolve FunctionHook references to callable functions.
+
+    Args:
+        function_hooks: A single FunctionHook or list of FunctionHooks to resolve
+
+    Returns:
+        Sequence of callable functions
+    """
     logger.debug(f"Creating hooks from: {function_hooks}")
-    hooks: Sequence[Callable[..., Any]] = []
+    hooks: list[Callable[..., Any]] = []
     if not function_hooks:
         return []
     if not isinstance(function_hooks, (list, tuple, set)):
@@ -27,197 +39,17 @@ def create_hooks(
     return hooks
 
 
-def null_hook(state: dict[str, Any], runtime: Runtime[Context]) -> dict[str, Any]:
+def null_hook(state: dict[str, Any], config: Any) -> dict[str, Any]:
+    """A no-op hook that returns an empty dict."""
     logger.debug("Executing null hook")
     return {}
 
 
 def null_initialization_hook(config: AppConfig) -> None:
+    """A no-op initialization hook."""
     logger.debug("Executing null initialization hook")
 
 
 def null_shutdown_hook(config: AppConfig) -> None:
+    """A no-op shutdown hook."""
     logger.debug("Executing null shutdown hook")
-
-
-def filter_last_human_message_hook(
-    state: dict[str, Any], runtime: Runtime[Context]
-) -> dict[str, Any]:
-    """
-    Filter messages to keep only the last human message.
-
-    This hook removes all messages except for the most recent human message,
-    which can be useful for scenarios where you want to process only the
-    latest user input without conversation history.
-
-    Args:
-        state: The current state containing messages
-        runtime: The runtime context (unused in this hook)
-
-    Returns:
-        Updated state with filtered messages
-    """
-    logger.debug("Executing filter_last_human_message hook")
-
-    messages: list[BaseMessage] = state.get("messages", [])
-
-    if not messages:
-        logger.debug("No messages found in state")
-        return state
-
-    # Use the helper function to find the last human message
-    last_message: HumanMessage = last_human_message(messages)
-
-    if last_message is None:
-        logger.debug("No human messages found in state")
-        # Return empty messages if no human message found
-        updated_state = state.copy()
-        updated_state["messages"] = []
-        return updated_state
-
-    logger.debug(f"Filtered {len(messages)} messages down to 1 (last human message)")
-
-    removed_messages: Sequence[BaseMessage] = [
-        RemoveMessage(id=message.id)
-        for message in messages
-        if message.id != last_message.id
-    ]
-
-    updated_state: dict[str, Sequence[BaseMessage]] = {"messages": removed_messages}
-
-    return updated_state
-
-
-def require_user_id_hook(
-    state: dict[str, Any], runtime: Runtime[Context]
-) -> dict[str, Any]:
-    logger.debug("Executing user_id validation hook")
-
-    context: Context = runtime.context or Context()
-
-    user_id: str | None = context.user_id
-
-    if not user_id:
-        logger.error("User ID is required but not provided in the configuration.")
-
-        # Create corrected configuration using any provided context parameters
-        corrected_config = {
-            "configurable": {
-                "thread_id": context.thread_id or "1",
-                "user_id": "my_user_id",
-                "store_num": context.store_num or 87887,
-            }
-        }
-
-        # Format as JSON for copy-paste
-        corrected_config_json = json.dumps(corrected_config, indent=2)
-
-        error_message = f"""
-## Authentication Required
-
-A **user_id** is required to process your request. Please provide your user ID in the configuration.
-
-### Required Configuration Format
-
-Please include the following JSON in your request configuration:
-
-```json
-{corrected_config_json}
-```
-
-### Field Descriptions
-- **user_id**: Your unique user identifier (required)
-- **thread_id**: Conversation thread identifier (required)
-- **store_num**: Your store number (required)
-
-Please update your configuration and try again.
-        """.strip()
-
-        raise ValueError(error_message)
-
-    if "." in user_id:
-        logger.error(f"User ID '{user_id}' contains invalid character '.'")
-
-        # Create a corrected version of the user_id
-        corrected_user_id = user_id.replace(".", "_")
-
-        # Create corrected configuration for the error message
-
-        # Corrected config with fixed user_id
-        corrected_config = {
-            "configurable": {
-                "thread_id": context.thread_id or "1",
-                "user_id": corrected_user_id,
-                "store_num": context.store_num or 87887,
-            }
-        }
-
-        # Format as JSON for copy-paste
-        corrected_config_json = json.dumps(corrected_config, indent=2)
-
-        error_message = f"""
-## Invalid User ID Format
-
-The **user_id** cannot contain a dot character ('.'). Please provide a valid user ID without dots.
-
-### Corrected Configuration (Copy & Paste This)
-```json
-{corrected_config_json}
-```
-
-Please update your user_id and try again.
-        """.strip()
-
-        raise ValueError(error_message)
-
-    return {}
-
-
-def require_thread_id_hook(
-    state: dict[str, Any], runtime: Runtime[Context]
-) -> dict[str, Any]:
-    logger.debug("Executing thread_id validation hook")
-
-    context: Context = runtime.context or Context()
-
-    thread_id: str | None = context.thread_id
-
-    if not thread_id:
-        logger.error("Thread ID is required but not provided in the configuration.")
-
-        # Create corrected configuration using any provided context parameters
-        corrected_config = {
-            "configurable": {
-                "thread_id": "1",
-                "user_id": context.user_id or "my_user_id",
-                "store_num": context.store_num or 87887,
-            }
-        }
-
-        # Format as JSON for copy-paste
-        corrected_config_json = json.dumps(corrected_config, indent=2)
-
-        error_message = f"""
-## Authentication Required
-
-A **thread_id** is required to process your request. Please provide your thread ID in the configuration.
-
-### Required Configuration Format
-
-Please include the following JSON in your request configuration:
-
-```json
-{corrected_config_json}
-```
-
-### Field Descriptions
-- **thread_id**: Conversation thread identifier (required)
-- **user_id**: Your unique user identifier (required)
-- **store_num**: Your store number (required)
-
-Please update your configuration and try again.
-        """.strip()
-
-        raise ValueError(error_message)
-
-    return {}

dao_ai/memory/__init__.py

@@ -3,10 +3,20 @@ from dao_ai.memory.base import (
     StoreManagerBase,
 )
 from dao_ai.memory.core import CheckpointManager, StoreManager
+from dao_ai.memory.databricks import (
+    AsyncDatabricksCheckpointSaver,
+    AsyncDatabricksStore,
+    DatabricksCheckpointerManager,
+    DatabricksStoreManager,
+)
 
 __all__ = [
     "CheckpointManagerBase",
     "StoreManagerBase",
     "CheckpointManager",
     "StoreManager",
+    "AsyncDatabricksCheckpointSaver",
+    "AsyncDatabricksStore",
+    "DatabricksCheckpointerManager",
+    "DatabricksStoreManager",
 ]

dao_ai/memory/core.py

@@ -1,8 +1,6 @@
 from typing import Any
 
-from databricks_langchain import (
-    DatabricksEmbeddings,
-)
+from databricks_langchain import DatabricksEmbeddings
 from langchain_core.embeddings.embeddings import Embeddings
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import InMemorySaver
@@ -60,7 +58,7 @@ class StoreManager:
 
     @classmethod
     def instance(cls, store_model: StoreModel) -> StoreManagerBase:
-        store_manager: StoreManagerBase = None
+        store_manager: StoreManagerBase | None = None
         match store_model.type:
             case StorageType.MEMORY:
                 store_manager = cls.store_managers.get(store_model.name)
@@ -78,6 +76,13 @@ class StoreManager:
                     cls.store_managers[store_model.database.instance_name] = (
                         store_manager
                     )
+            case StorageType.LAKEBASE:
+                from dao_ai.memory.databricks import DatabricksStoreManager
+
+                store_manager = cls.store_managers.get(store_model.name)
+                if store_manager is None:
+                    store_manager = DatabricksStoreManager(store_model)
+                    cls.store_managers[store_model.name] = store_manager
             case _:
                 raise ValueError(f"Unknown store type: {store_model.type}")
 
@@ -89,7 +94,7 @@ class CheckpointManager:
 
     @classmethod
     def instance(cls, checkpointer_model: CheckpointerModel) -> CheckpointManagerBase:
-        checkpointer_manager: CheckpointManagerBase = None
+        checkpointer_manager: CheckpointManagerBase | None = None
         match checkpointer_model.type:
             case StorageType.MEMORY:
                 checkpointer_manager = cls.checkpoint_managers.get(
@@ -115,6 +120,19 @@ class CheckpointManager:
                     cls.checkpoint_managers[
                         checkpointer_model.database.instance_name
                     ] = checkpointer_manager
+            case StorageType.LAKEBASE:
+                from dao_ai.memory.databricks import DatabricksCheckpointerManager
+
+                checkpointer_manager = cls.checkpoint_managers.get(
+                    checkpointer_model.name
+                )
+                if checkpointer_manager is None:
+                    checkpointer_manager = DatabricksCheckpointerManager(
+                        checkpointer_model
+                    )
+                    cls.checkpoint_managers[checkpointer_model.name] = (
+                        checkpointer_manager
+                    )
             case _:
                 raise ValueError(f"Unknown store type: {checkpointer_model.type}")