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

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")
@@ -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),