dao-ai 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

dao_ai/cli.py

@@ -460,6 +460,49 @@ def setup_logging(verbosity: int) -> None:
     logger.add(sys.stderr, level=level)
 
 
+def generate_bundle_from_template(config_path: Path, app_name: str) -> Path:
+    """
+    Generate an app-specific databricks.yaml from databricks.yaml.template.
+
+    This function:
+    1. Reads databricks.yaml.template (permanent template file)
+    2. Replaces __APP_NAME__ with the actual app name
+    3. Writes to databricks.yaml (overwrites if exists)
+    4. Returns the path to the generated file
+
+    The generated databricks.yaml is overwritten on each deployment and is not tracked in git.
+    Schema reference remains pointing to ./schemas/bundle_config_schema.json.
+
+    Args:
+        config_path: Path to the app config file
+        app_name: Normalized app name
+
+    Returns:
+        Path to the generated databricks.yaml file
+    """
+    cwd = Path.cwd()
+    template_path = cwd / "databricks.yaml.template"
+    output_path = cwd / "databricks.yaml"
+
+    if not template_path.exists():
+        logger.error(f"Template file {template_path} does not exist.")
+        sys.exit(1)
+
+    # Read template
+    with open(template_path, "r") as f:
+        template_content = f.read()
+
+    # Replace template variables
+    bundle_content = template_content.replace("__APP_NAME__", app_name)
+
+    # Write generated databricks.yaml (overwrite if exists)
+    with open(output_path, "w") as f:
+        f.write(bundle_content)
+
+    logger.info(f"Generated bundle configuration at {output_path} from template")
+    return output_path
+
+
 def run_databricks_command(
     command: list[str],
     profile: Optional[str] = None,
@@ -467,44 +510,55 @@ def run_databricks_command(
     target: Optional[str] = None,
     dry_run: bool = False,
 ) -> None:
-    """Execute a databricks CLI command with optional profile."""
+    """Execute a databricks CLI command with optional profile and target."""
+    config_path = Path(config) if config else None
+
+    if config_path and not config_path.exists():
+        logger.error(f"Configuration file {config_path} does not exist.")
+        sys.exit(1)
+
+    # Load app config and generate bundle from template
+    app_config: AppConfig = AppConfig.from_file(config_path) if config_path else None
+    normalized_name: str = normalize_name(app_config.app.name) if app_config else None
+
+    # Generate app-specific bundle from template (overwrites databricks.yaml temporarily)
+    if config_path and app_config:
+        generate_bundle_from_template(config_path, normalized_name)
+
+    # Use app name as target if not explicitly provided
+    # This ensures each app gets its own Terraform state in .databricks/bundle/<app-name>/
+    if not target and normalized_name:
+        target = normalized_name
+        logger.debug(f"Using app-specific target: {target}")
+
+    # Build databricks command (no -c flag needed, uses databricks.yaml in current dir)
     cmd = ["databricks"]
     if profile:
         cmd.extend(["--profile", profile])
+
     if target:
         cmd.extend(["--target", target])
-    cmd.extend(command)
-    if config:
-        config_path = Path(config)
 
-        if not config_path.exists():
-            logger.error(f"Configuration file {config_path} does not exist.")
-            sys.exit(1)
-
-        app_config: AppConfig = AppConfig.from_file(config_path)
+    cmd.extend(command)
 
-        # Always convert to path relative to notebooks directory
-        # Get absolute path of config file and current working directory
+    # Add config_path variable for notebooks
+    if config_path and app_config:
+        # Calculate relative path from notebooks directory to config file
         config_abs = config_path.resolve()
         cwd = Path.cwd()
         notebooks_dir = cwd / "notebooks"
 
-        # Calculate relative path from notebooks directory to config file
         try:
             relative_config = config_abs.relative_to(notebooks_dir)
         except ValueError:
-            # Config file is outside notebooks directory, calculate relative path
-            # Use os.path.relpath to get the relative path from notebooks_dir to config file
             relative_config = Path(os.path.relpath(config_abs, notebooks_dir))
 
         cmd.append(f'--var="config_path={relative_config}"')
 
-        normalized_name: str = normalize_name(app_config.app.name)
-        cmd.append(f'--var="app_name={normalized_name}"')
-
     logger.debug(f"Executing command: {' '.join(cmd)}")
 
     if dry_run:
+        logger.info(f"[DRY RUN] Would execute: {' '.join(cmd)}")
         return
 
     try:
@@ -531,6 +585,9 @@ def run_databricks_command(
     except FileNotFoundError:
         logger.error("databricks CLI not found. Please install the Databricks CLI.")
         sys.exit(1)
+    except Exception as e:
+        logger.error(f"Command execution failed: {e}")
+        sys.exit(1)
 
 
 def handle_bundle_command(options: Namespace) -> None:
@@ -539,6 +596,7 @@ def handle_bundle_command(options: Namespace) -> None:
     config: Optional[str] = options.config
     target: Optional[str] = options.target
     dry_run: bool = options.dry_run
+
     if options.deploy:
         logger.info("Deploying DAO AI asset bundle...")
         run_databricks_command(
@@ -546,8 +604,9 @@ def handle_bundle_command(options: Namespace) -> None:
         )
     if options.run:
         logger.info("Running DAO AI system with current configuration...")
+        # Use static job resource key that matches databricks.yaml (resources.jobs.deploy_job)
         run_databricks_command(
-            ["bundle", "run", "deploy-end-to-end"],
+            ["bundle", "run", "deploy_job"],
             profile,
             config,
             target,

dao_ai/config.py

@@ -30,12 +30,15 @@ from databricks_langchain import (
     DatabricksFunctionClient,
 )
 from langchain_core.language_models import LanguageModelLike
+from langchain_core.messages import BaseMessage, messages_from_dict
 from langchain_core.runnables.base import RunnableLike
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.store.base import BaseStore
 from loguru import logger
+from mlflow.genai.datasets import EvaluationDataset, create_dataset, get_dataset
+from mlflow.genai.prompts import PromptVersion, load_prompt
 from mlflow.models import ModelConfig
 from mlflow.models.resources import (
     DatabricksFunction,
@@ -49,6 +52,9 @@ from mlflow.models.resources import (
     DatabricksVectorSearchIndex,
 )
 from mlflow.pyfunc import ChatModel, ResponsesAgent
+from mlflow.types.responses import (
+    ResponsesAgentRequest,
+)
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -324,6 +330,10 @@ class LLMModel(BaseModel, IsDatabricksResource):
             "serving.serving-endpoints",
         ]
 
+    @property
+    def uri(self) -> str:
+        return f"databricks:/{self.name}"
+
     def as_resources(self) -> Sequence[DatabricksResource]:
         return [
             DatabricksServingEndpoint(
@@ -1181,17 +1191,32 @@ class PromptModel(BaseModel, HasFullName):
         from dao_ai.providers.databricks import DatabricksProvider
 
         provider: DatabricksProvider = DatabricksProvider()
-        prompt: str = provider.get_prompt(self)
-        return prompt
+        prompt_version = provider.get_prompt(self)
+        return prompt_version.to_single_brace_format()
 
     @property
     def full_name(self) -> str:
+        prompt_name: str = self.name
         if self.schema_model:
-            name: str = ""
-            if self.name:
-                name = f".{self.name}"
-            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}{name}"
-        return self.name
+            prompt_name = f"{self.schema_model.full_name}.{prompt_name}"
+        return prompt_name
+
+    @property
+    def uri(self) -> str:
+        prompt_uri: str = f"prompts:/{self.full_name}"
+
+        if self.alias:
+            prompt_uri = f"prompts:/{self.full_name}@{self.alias}"
+        elif self.version:
+            prompt_uri = f"prompts:/{self.full_name}/{self.version}"
+        else:
+            prompt_uri = f"prompts:/{self.full_name}@latest"
+
+        return prompt_uri
+
+    def as_prompt(self) -> PromptVersion:
+        prompt_version: PromptVersion = load_prompt(self.uri)
+        return prompt_version
 
     @model_validator(mode="after")
     def validate_mutually_exclusive(self):
@@ -1213,6 +1238,17 @@ class AgentModel(BaseModel):
     pre_agent_hook: Optional[FunctionHook] = None
     post_agent_hook: Optional[FunctionHook] = None
 
+    def as_runnable(self) -> RunnableLike:
+        from dao_ai.nodes import create_agent_node
+
+        return create_agent_node(self)
+
+    def as_responses_agent(self) -> ResponsesAgent:
+        from dao_ai.models import create_responses_agent
+
+        graph: CompiledStateGraph = self.as_runnable()
+        return create_responses_agent(graph)
+
 
 class SupervisorModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -1301,7 +1337,7 @@ class ChatPayload(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     input: Optional[list[Message]] = None
     messages: Optional[list[Message]] = None
-    custom_inputs: dict
+    custom_inputs: Optional[dict] = Field(default_factory=dict)
 
     @model_validator(mode="after")
     def validate_mutual_exclusion_and_alias(self) -> "ChatPayload":
@@ -1330,6 +1366,19 @@ class ChatPayload(BaseModel):
 
         return self
 
+    def as_messages(self) -> Sequence[BaseMessage]:
+        return messages_from_dict(
+            [{"type": m.role, "content": m.content} for m in self.messages]
+        )
+
+    def as_agent_request(self) -> ResponsesAgentRequest:
+        from mlflow.types.responses_helpers import Message as _Message
+
+        return ResponsesAgentRequest(
+            input=[_Message(role=m.role, content=m.content) for m in self.messages],
+            custom_inputs=self.custom_inputs,
+        )
+
 
 class ChatHistoryModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -1459,6 +1508,174 @@ class EvaluationModel(BaseModel):
     guidelines: list[GuidelineModel] = Field(default_factory=list)
 
 
+class EvaluationDatasetExpectationsModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    expected_response: Optional[str] = None
+    expected_facts: Optional[list[str]] = None
+
+    @model_validator(mode="after")
+    def validate_mutually_exclusive(self):
+        if self.expected_response is not None and self.expected_facts is not None:
+            raise ValueError("Cannot specify both expected_response and expected_facts")
+        return self
+
+
+class EvaluationDatasetEntryModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    inputs: ChatPayload
+    expectations: EvaluationDatasetExpectationsModel
+
+    def to_mlflow_format(self) -> dict[str, Any]:
+        """
+        Convert to MLflow evaluation dataset format.
+
+        Flattens the expectations fields to the top level alongside inputs,
+        which is the format expected by MLflow's Correctness scorer.
+
+        Returns:
+            dict: Flattened dictionary with inputs and expectation fields at top level
+        """
+        result: dict[str, Any] = {"inputs": self.inputs.model_dump()}
+
+        # Flatten expectations to top level for MLflow compatibility
+        if self.expectations.expected_response is not None:
+            result["expected_response"] = self.expectations.expected_response
+        if self.expectations.expected_facts is not None:
+            result["expected_facts"] = self.expectations.expected_facts
+
+        return result
+
+
+class EvaluationDatasetModel(BaseModel, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    name: str
+    data: Optional[list[EvaluationDatasetEntryModel]] = Field(default_factory=list)
+    overwrite: Optional[bool] = False
+
+    def as_dataset(self, w: WorkspaceClient | None = None) -> EvaluationDataset:
+        evaluation_dataset: EvaluationDataset
+        needs_creation: bool = False
+
+        try:
+            evaluation_dataset = get_dataset(name=self.full_name)
+            if self.overwrite:
+                logger.warning(f"Overwriting dataset {self.full_name}")
+                workspace_client: WorkspaceClient = w if w else WorkspaceClient()
+                logger.debug(f"Dropping table: {self.full_name}")
+                workspace_client.tables.delete(full_name=self.full_name)
+                needs_creation = True
+        except Exception:
+            logger.warning(
+                f"Dataset {self.full_name} not found, will create new dataset"
+            )
+            needs_creation = True
+
+        # Create dataset if needed (either new or after overwrite)
+        if needs_creation:
+            evaluation_dataset = create_dataset(name=self.full_name)
+            if self.data:
+                logger.debug(
+                    f"Merging {len(self.data)} entries into dataset {self.full_name}"
+                )
+                # Use to_mlflow_format() to flatten expectations for MLflow compatibility
+                evaluation_dataset.merge_records(
+                    [e.to_mlflow_format() for e in self.data]
+                )
+
+        return evaluation_dataset
+
+    @property
+    def full_name(self) -> str:
+        if self.schema_model:
+            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}.{self.name}"
+        return self.name
+
+
+class PromptOptimizationModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str
+    prompt: Optional[PromptModel] = None
+    agent: AgentModel
+    dataset: (
+        EvaluationDatasetModel | str
+    )  # Reference to dataset name (looked up in OptimizationsModel.training_datasets or MLflow)
+    reflection_model: Optional[LLMModel | str] = None
+    num_candidates: Optional[int] = 50
+    scorer_model: Optional[LLMModel | str] = None
+
+    def optimize(self, w: WorkspaceClient | None = None) -> PromptModel:
+        """
+        Optimize the prompt using MLflow's prompt optimization.
+
+        Args:
+            w: Optional WorkspaceClient for Databricks operations
+
+        Returns:
+            PromptModel: The optimized prompt model with new URI
+        """
+        from dao_ai.providers.base import ServiceProvider
+        from dao_ai.providers.databricks import DatabricksProvider
+
+        provider: ServiceProvider = DatabricksProvider(w=w)
+        optimized_prompt: PromptModel = provider.optimize_prompt(self)
+        return optimized_prompt
+
+    @model_validator(mode="after")
+    def set_defaults(self):
+        # If no prompt is specified, try to use the agent's prompt
+        if self.prompt is None:
+            if isinstance(self.agent.prompt, PromptModel):
+                self.prompt = self.agent.prompt
+            else:
+                raise ValueError(
+                    f"Prompt optimization '{self.name}' requires either an explicit prompt "
+                    f"or an agent with a prompt configured"
+                )
+
+        if self.reflection_model is None:
+            self.reflection_model = self.agent.model
+
+        if self.scorer_model is None:
+            self.scorer_model = self.agent.model
+
+        return self
+
+
+class OptimizationsModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    training_datasets: dict[str, EvaluationDatasetModel] = Field(default_factory=dict)
+    prompt_optimizations: dict[str, PromptOptimizationModel] = Field(
+        default_factory=dict
+    )
+
+    def optimize(self, w: WorkspaceClient | None = None) -> dict[str, PromptModel]:
+        """
+        Optimize all prompts in this configuration.
+
+        This method:
+        1. Ensures all training datasets are created/registered in MLflow
+        2. Runs each prompt optimization
+
+        Args:
+            w: Optional WorkspaceClient for Databricks operations
+
+        Returns:
+            dict[str, PromptModel]: Dictionary mapping optimization names to optimized prompts
+        """
+        # First, ensure all training datasets are created/registered in MLflow
+        logger.info(f"Ensuring {len(self.training_datasets)} training datasets exist")
+        for dataset_name, dataset_model in self.training_datasets.items():
+            logger.debug(f"Creating/updating dataset: {dataset_name}")
+            dataset_model.as_dataset()
+
+        # Run optimizations
+        results: dict[str, PromptModel] = {}
+        for name, optimization in self.prompt_optimizations.items():
+            results[name] = optimization.optimize(w)
+        return results
+
+
 class DatasetFormat(str, Enum):
     CSV = "csv"
     DELTA = "delta"
@@ -1537,6 +1754,7 @@ class AppConfig(BaseModel):
     agents: dict[str, AgentModel] = Field(default_factory=dict)
     app: Optional[AppModel] = None
     evaluation: Optional[EvaluationModel] = None
+    optimizations: Optional[OptimizationsModel] = None
     datasets: Optional[list[DatasetModel]] = Field(default_factory=list)
     unity_catalog_functions: Optional[list[UnityCatalogFunctionSqlModel]] = Field(
         default_factory=list

dao_ai/graph.py

@@ -79,7 +79,12 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
     for registered_agent in config.app.agents:
         agents.append(
             create_agent_node(
-                app=config.app, agent=registered_agent, additional_tools=[]
+                agent=registered_agent,
+                memory=config.app.orchestration.memory
+                if config.app.orchestration
+                else None,
+                chat_history=config.app.chat_history,
+                additional_tools=[],
             )
         )
         tools.append(
@@ -169,7 +174,12 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
         )
         agents.append(
             create_agent_node(
-                app=config.app, agent=registered_agent, additional_tools=handoff_tools
+                agent=registered_agent,
+                memory=config.app.orchestration.memory
+                if config.app.orchestration
+                else None,
+                chat_history=config.app.chat_history,
+                additional_tools=handoff_tools,
             )
         )
 

dao_ai/nodes.py

@@ -19,9 +19,9 @@ from loguru import logger
 from dao_ai.config import (
     AgentModel,
     AppConfig,
-    AppModel,
     ChatHistoryModel,
     FunctionHook,
+    MemoryModel,
     ToolModel,
 )
 from dao_ai.guardrails import reflection_guardrail, with_guardrails
@@ -31,12 +31,18 @@ from dao_ai.state import Context, IncomingState, SharedState
 from dao_ai.tools import create_tools
 
 
-def summarization_node(app_model: AppModel) -> RunnableLike:
-    chat_history: ChatHistoryModel | None = app_model.chat_history
+def summarization_node(chat_history: ChatHistoryModel) -> RunnableLike:
+    """
+    Create a summarization node for managing chat history.
+
+    Args:
+        chat_history: ChatHistoryModel configuration for summarization
+
+    Returns:
+        RunnableLike: A summarization node that processes messages
+    """
     if chat_history is None:
-        raise ValueError(
-            "AppModel must have chat_history configured to use summarization"
-        )
+        raise ValueError("chat_history must be provided to use summarization")
 
     max_tokens: int = chat_history.max_tokens
     max_tokens_before_summary: int | None = chat_history.max_tokens_before_summary
@@ -93,23 +99,26 @@ def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLi
 
 
 def create_agent_node(
-    app: AppModel,
     agent: AgentModel,
+    memory: Optional[MemoryModel] = None,
+    chat_history: Optional[ChatHistoryModel] = None,
     additional_tools: Optional[Sequence[BaseTool]] = None,
 ) -> RunnableLike:
     """
     Factory function that creates a LangGraph node for a specialized agent.
 
-    This creates a node function that handles user requests using a specialized agent
-    based on the provided agent_type. The function configures the agent with the
-    appropriate model, prompt, tools, and guardrails from the model_config.
+    This creates a node function that handles user requests using a specialized agent.
+    The function configures the agent with the appropriate model, prompt, tools, and guardrails.
+    If chat_history is provided, it creates a workflow with summarization node.
 
     Args:
-        model_config: Configuration containing models, prompts, tools, and guardrails
-        agent_type: Type of agent to create (e.g., "general", "product", "inventory")
+        agent: AgentModel configuration for the agent
+        memory: Optional MemoryModel for memory store configuration
+        chat_history: Optional ChatHistoryModel for chat history summarization
+        additional_tools: Optional sequence of additional tools to add to the agent
 
     Returns:
-        An agent callable function that processes state and returns responses
+        RunnableLike: An agent node that processes state and returns responses
     """
     logger.debug(f"Creating agent node for {agent.name}")
 
@@ -124,10 +133,10 @@ def create_agent_node(
         additional_tools = []
     tools: Sequence[BaseTool] = create_tools(tool_models) + additional_tools
 
-    if app.orchestration.memory and app.orchestration.memory.store:
+    if memory and memory.store:
         namespace: tuple[str, ...] = ("memory",)
-        if app.orchestration.memory.store.namespace:
-            namespace = namespace + (app.orchestration.memory.store.namespace,)
+        if memory.store.namespace:
+            namespace = namespace + (memory.store.namespace,)
         logger.debug(f"Memory store namespace: {namespace}")
 
         tools += [
@@ -145,13 +154,15 @@ def create_agent_node(
     )
     logger.debug(f"post_agent_hook: {post_agent_hook}")
 
+    checkpointer: bool = memory and memory.checkpointer is not None
+
     compiled_agent: CompiledStateGraph = create_react_agent(
         name=agent.name,
         model=llm,
         prompt=make_prompt(agent.prompt),
         tools=tools,
         store=True,
-        checkpointer=True,
+        checkpointer=checkpointer,
         state_schema=SharedState,
         context_schema=Context,
         pre_model_hook=pre_agent_hook,
@@ -166,8 +177,6 @@ def create_agent_node(
 
     agent_node: CompiledStateGraph
 
-    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
@@ -179,7 +188,7 @@ def create_agent_node(
             input=SharedState,
             output=SharedState,
         )
-        workflow.add_node("summarization", summarization_node(app))
+        workflow.add_node("summarization", summarization_node(chat_history))
         workflow.add_node(
             "agent",
             call_agent_with_summarized_messages(agent=compiled_agent),