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

dao_ai/agent_as_code.py

@@ -7,6 +7,9 @@ from mlflow.pyfunc import ResponsesAgent
 
 from dao_ai.config import AppConfig
 
+mlflow.set_registry_uri("databricks-uc")
+mlflow.set_tracking_uri("databricks")
+
 mlflow.langchain.autolog()
 
 model_config: ModelConfig = ModelConfig()

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(
@@ -387,6 +397,13 @@ class VectorSearchEndpoint(BaseModel):
     name: str
     type: VectorSearchEndpointType = VectorSearchEndpointType.STANDARD
 
+    @field_serializer("type")
+    def serialize_type(self, value: VectorSearchEndpointType) -> str:
+        """Ensure enum is serialized to string value."""
+        if isinstance(value, VectorSearchEndpointType):
+            return value.value
+        return str(value)
+
 
 class IndexModel(BaseModel, HasFullName, IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -878,6 +895,55 @@ class SearchParametersModel(BaseModel):
     query_type: Optional[str] = "ANN"
 
 
+class RerankParametersModel(BaseModel):
+    """
+    Configuration for reranking retrieved documents using FlashRank.
+
+    FlashRank provides fast, local reranking without API calls using lightweight
+    cross-encoder models. Reranking improves retrieval quality by reordering results
+    based on semantic relevance to the query.
+
+    Typical workflow:
+    1. Retrieve more documents than needed (e.g., 50 via num_results)
+    2. Rerank all retrieved documents
+    3. Return top_n best matches (e.g., 5)
+
+    Example:
+        ```yaml
+        retriever:
+          search_parameters:
+            num_results: 50  # Retrieve more candidates
+          rerank:
+            model: ms-marco-MiniLM-L-12-v2
+            top_n: 5  # Return top 5 after reranking
+        ```
+
+    Available models (from fastest to most accurate):
+    - "ms-marco-TinyBERT-L-2-v2" (fastest, smallest)
+    - "ms-marco-MiniLM-L-6-v2"
+    - "ms-marco-MiniLM-L-12-v2" (default, good balance)
+    - "rank-T5-flan" (most accurate, slower)
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    model: str = Field(
+        default="ms-marco-MiniLM-L-12-v2",
+        description="FlashRank model name. Default provides good balance of speed and accuracy.",
+    )
+    top_n: Optional[int] = Field(
+        default=None,
+        description="Number of documents to return after reranking. If None, uses search_parameters.num_results.",
+    )
+    cache_dir: Optional[str] = Field(
+        default="/tmp/flashrank_cache",
+        description="Directory to cache downloaded model weights.",
+    )
+    columns: Optional[list[str]] = Field(
+        default_factory=list, description="Columns to rerank using DatabricksReranker"
+    )
+
+
 class RetrieverModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     vector_store: VectorStoreModel
@@ -885,6 +951,10 @@ class RetrieverModel(BaseModel):
     search_parameters: SearchParametersModel = Field(
         default_factory=SearchParametersModel
     )
+    rerank: Optional[RerankParametersModel | bool] = Field(
+        default=None,
+        description="Optional reranking configuration. Set to true for defaults, or provide ReRankParametersModel for custom settings.",
+    )
 
     @model_validator(mode="after")
     def set_default_columns(self):
@@ -893,6 +963,13 @@ class RetrieverModel(BaseModel):
             self.columns = columns
         return self
 
+    @model_validator(mode="after")
+    def set_default_reranker(self):
+        """Convert bool to ReRankParametersModel with defaults."""
+        if isinstance(self.rerank, bool) and self.rerank:
+            self.rerank = RerankParametersModel()
+        return self
+
 
 class FunctionType(str, Enum):
     PYTHON = "python"
@@ -988,22 +1065,120 @@ class TransportType(str, Enum):
 class McpFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     type: Literal[FunctionType.MCP] = FunctionType.MCP
-
     transport: TransportType = TransportType.STREAMABLE_HTTP
     command: Optional[str] = "python"
     url: Optional[AnyVariable] = None
-    connection: Optional[ConnectionModel] = None
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
     pat: Optional[AnyVariable] = None
     client_id: Optional[AnyVariable] = None
     client_secret: Optional[AnyVariable] = None
     workspace_host: Optional[AnyVariable] = None
+    connection: Optional[ConnectionModel] = None
+    functions: Optional[SchemaModel] = None
+    genie_room: Optional[GenieRoomModel] = None
+    sql: Optional[bool] = None
+    vector_search: Optional[VectorStoreModel] = None
 
     @property
     def full_name(self) -> str:
         return self.name
 
+    def _get_workspace_host(self) -> str:
+        """
+        Get the workspace host, either from config or from workspace client.
+
+        If connection is provided, uses its workspace client.
+        Otherwise, falls back to creating a new workspace client.
+
+        Returns:
+            str: The workspace host URL without trailing slash
+        """
+        from databricks.sdk import WorkspaceClient
+
+        # Try to get workspace_host from config
+        workspace_host: str | None = (
+            value_of(self.workspace_host) if self.workspace_host else None
+        )
+
+        # If no workspace_host in config, get it from workspace client
+        if not workspace_host:
+            # Use connection's workspace client if available
+            if self.connection:
+                workspace_host = self.connection.workspace_client.config.host
+            else:
+                # Create a default workspace client
+                w: WorkspaceClient = WorkspaceClient()
+                workspace_host = w.config.host
+
+        # Remove trailing slash
+        return workspace_host.rstrip("/")
+
+    @property
+    def mcp_url(self) -> str:
+        """
+        Get the MCP URL for this function.
+
+        Returns the URL based on the configured source:
+        - If url is set, returns it directly
+        - If connection is set, constructs URL from connection
+        - If genie_room is set, constructs Genie MCP URL
+        - If sql is set, constructs DBSQL MCP URL (serverless)
+        - If vector_search is set, constructs Vector Search MCP URL
+        - If functions is set, constructs UC Functions MCP URL
+
+        URL patterns (per https://docs.databricks.com/aws/en/generative-ai/mcp/managed-mcp):
+        - Genie: https://{host}/api/2.0/mcp/genie/{space_id}
+        - DBSQL: https://{host}/api/2.0/mcp/sql (serverless, workspace-level)
+        - Vector Search: https://{host}/api/2.0/mcp/vector-search/{catalog}/{schema}
+        - UC Functions: https://{host}/api/2.0/mcp/functions/{catalog}/{schema}
+        - Connection: https://{host}/api/2.0/mcp/external/{connection_name}
+        """
+        # Direct URL provided
+        if self.url:
+            return self.url
+
+        # Get workspace host (from config, connection, or default workspace client)
+        workspace_host: str = self._get_workspace_host()
+
+        # UC Connection
+        if self.connection:
+            connection_name: str = self.connection.name
+            return f"{workspace_host}/api/2.0/mcp/external/{connection_name}"
+
+        # Genie Room
+        if self.genie_room:
+            space_id: str = value_of(self.genie_room.space_id)
+            return f"{workspace_host}/api/2.0/mcp/genie/{space_id}"
+
+        # DBSQL MCP server (serverless, workspace-level)
+        if self.sql:
+            return f"{workspace_host}/api/2.0/mcp/sql"
+
+        # Vector Search
+        if self.vector_search:
+            if (
+                not self.vector_search.index
+                or not self.vector_search.index.schema_model
+            ):
+                raise ValueError(
+                    "vector_search must have an index with a schema (catalog/schema) configured"
+                )
+            catalog: str = self.vector_search.index.schema_model.catalog_name
+            schema: str = self.vector_search.index.schema_model.schema_name
+            return f"{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}"
+
+        # UC Functions MCP server
+        if self.functions:
+            catalog: str = self.functions.catalog_name
+            schema: str = self.functions.schema_name
+            return f"{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}"
+
+        raise ValueError(
+            "No URL source configured. Provide one of: url, connection, genie_room, "
+            "sql, vector_search, or functions"
+        )
+
     @field_serializer("transport")
     def serialize_transport(self, value) -> str:
         if isinstance(value, TransportType):
@@ -1011,32 +1186,56 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
         return str(value)
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
-        if self.transport == TransportType.STREAMABLE_HTTP and not (
-            self.url or self.connection
-        ):
-            raise ValueError(
-                "url or connection must be provided for STREAMABLE_HTTP transport"
-            )
-        if self.transport == TransportType.STDIO and not self.command:
-            raise ValueError("command must not be provided for STDIO transport")
-        if self.transport == TransportType.STDIO and not self.args:
-            raise ValueError("args must not be provided for STDIO transport")
+    def validate_mutually_exclusive(self) -> "McpFunctionModel":
+        """Validate that exactly one URL source is provided."""
+        # Count how many URL sources are provided
+        url_sources: list[tuple[str, Any]] = [
+            ("url", self.url),
+            ("connection", self.connection),
+            ("genie_room", self.genie_room),
+            ("sql", self.sql),
+            ("vector_search", self.vector_search),
+            ("functions", self.functions),
+        ]
+
+        provided_sources: list[str] = [
+            name for name, value in url_sources if value is not None
+        ]
+
+        if self.transport == TransportType.STREAMABLE_HTTP:
+            if len(provided_sources) == 0:
+                raise ValueError(
+                    "For STREAMABLE_HTTP transport, exactly one of the following must be provided: "
+                    "url, connection, genie_room, sql, vector_search, or functions"
+                )
+            if len(provided_sources) > 1:
+                raise ValueError(
+                    f"For STREAMABLE_HTTP transport, only one URL source can be provided. "
+                    f"Found: {', '.join(provided_sources)}. "
+                    f"Please provide only one of: url, connection, genie_room, sql, vector_search, or functions"
+                )
+
+        if self.transport == TransportType.STDIO:
+            if not self.command:
+                raise ValueError("command must be provided for STDIO transport")
+            if not self.args:
+                raise ValueError("args must be provided for STDIO transport")
+
         return self
 
     @model_validator(mode="after")
-    def update_url(self):
+    def update_url(self) -> "McpFunctionModel":
         self.url = value_of(self.url)
         return self
 
     @model_validator(mode="after")
-    def update_headers(self):
+    def update_headers(self) -> "McpFunctionModel":
         for key, value in self.headers.items():
             self.headers[key] = value_of(value)
         return self
 
     @model_validator(mode="after")
-    def validate_auth_methods(self):
+    def validate_auth_methods(self) -> "McpFunctionModel":
         oauth_fields: Sequence[Any] = [
             self.client_id,
             self.client_secret,
@@ -1052,10 +1251,7 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
                 "Please provide either OAuth credentials or user credentials."
             )
 
-        if (has_oauth or has_user_auth) and not self.workspace_host:
-            raise ValueError(
-                "Workspace host must be provided when using OAuth or user credentials."
-            )
+        # Note: workspace_host is optional - it will be derived from workspace client if not provided
 
         return self
 
@@ -1181,17 +1377,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 +1424,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 +1552,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 +1694,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 +1940,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

@@ -62,11 +62,19 @@ def _handoffs_for_agent(agent: AgentModel, config: AppConfig) -> Sequence[BaseTo
         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_to_agent.handoff_prompt,
+                + handoff_description,
             )
         )
     return handoff_tools
@@ -79,13 +87,25 @@ 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=[],
             )
         )
+        # 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=registered_agent.handoff_prompt,
+                description=handoff_description,
             )
         )
 
@@ -169,7 +189,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,
             )
         )