diff --git a/2025-11-Multi-Agent-GenAI-System/03-create-multi-agent-with-genie.ipynb b/2025-11-Multi-Agent-GenAI-System/03-create-multi-agent-with-genie.ipynb deleted file mode 100644 index e73c89a..0000000 --- a/2025-11-Multi-Agent-GenAI-System/03-create-multi-agent-with-genie.ipynb +++ /dev/null @@ -1,876 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "bdb29bed-d9f7-4cf8-ae11-6fcc2b989eea", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "%pip install -U -qqq mlflow[databricks]==3.1.1 langgraph==0.5.4 databricks-langchain==0.6.0 databricks-agents==1.1.0 pydantic<2.12.0 uv\n", - "dbutils.library.restartPython()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "db2ef15b-42d8-4e4d-9b2f-5f507ba509bd", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "## Sales Support Multi-Agent Framework\n", - "\n", - " This notebook creates a multi-agent system for sales support with the following components:\n", - " - **Structured Data Agent**: Queries structured sales data (opportunities, accounts, activities)\n", - " - **Vector search Agent**: Retrieves information from unstructured documents (emails, meeting notes, feedback)\n", - " - **Supervisor**: Routes queries to appropriate agents and orchestrates responses\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "7e1b5b7d-4a97-4626-98d0-6b1cbc627637", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "%%writefile agent.py\n", - "\n", - "import warnings\n", - "import os\n", - "os.environ[\"DATABRICKS_DISABLE_NOTICE\"] = \"true\"\n", - "warnings.filterwarnings(\"ignore\", message=\".*notebook authentication token.*\")\n", - "\n", - "import functools\n", - "from typing import Any, Generator, Literal, Optional\n", - "\n", - "import mlflow\n", - "from databricks.sdk import WorkspaceClient\n", - "from databricks.vector_search.client import VectorSearchClient\n", - "from databricks_langchain import ChatDatabricks, VectorSearchRetrieverTool\n", - "from databricks_langchain.genie import GenieAgent\n", - "from langchain_core.runnables import RunnableLambda\n", - "from langgraph.graph import END, StateGraph\n", - "from langgraph.graph.state import CompiledStateGraph\n", - "from langgraph.prebuilt import create_react_agent\n", - "from mlflow.entities import SpanType\n", - "from mlflow.langchain.chat_agent_langgraph import ChatAgentState\n", - "from mlflow.pyfunc import ChatAgent\n", - "from mlflow.types.agent import (\n", - " ChatAgentChunk,\n", - " ChatAgentMessage,\n", - " ChatAgentResponse,\n", - " ChatContext,\n", - ")\n", - "from pydantic import BaseModel\n", - "import time\n", - "\n", - "w = WorkspaceClient()\n", - "token = w.tokens.create(comment=f\"sdk-{time.time_ns()}\").token_value\n", - "\n", - "# Get catalog and schema from environment or config\n", - "CATALOG_NAME = 'andrea_tardif'\n", - "SCHEMA_NAME = 'workday_demos'\n", - "GENIE_SPACE_ID = '01f0c0291e6f1feabc4a8a46085cebd1'\n", - "MLFLOW_EXPERIMENT_NAME = f'multiagent_genie_{CATALOG_NAME}'\n", - "host='https://dbc-d079f94e-4181.cloud.databricks.com/'\n", - "\n", - "###################################################\n", - "## Configure LLM\n", - "###################################################\n", - "llm = ChatDatabricks(endpoint=\"databricks-claude-3-7-sonnet\")\n", - "\n", - "###################################################\n", - "## Create RAG Agent with Vector Search Tools\n", - "###################################################\n", - "\n", - "# Create retriever tools for each document type with disable_notice\n", - "email_retriever = VectorSearchRetrieverTool(\n", - " index_name=f\"{CATALOG_NAME}.{SCHEMA_NAME}.email_communications_index\",\n", - " columns=[\"content\", \"doc_uri\"],\n", - " name=\"email_search\",\n", - " description=(\n", - " \"Searches through email communications between sales reps and customers. \"\n", - " \"Use this to find information about: pricing discussions, objections, \"\n", - " \"follow-ups, proposal details, and customer email correspondence.\"\n", - " ),\n", - " disable_notice=True,\n", - ")\n", - "\n", - "meeting_notes_retriever = VectorSearchRetrieverTool(\n", - " index_name=f\"{CATALOG_NAME}.{SCHEMA_NAME}.meeting_notes_index\",\n", - " columns=[\"content\", \"doc_uri\"],\n", - " name=\"meeting_notes_search\",\n", - " description=(\n", - " \"Searches through meeting notes and call summaries. \"\n", - " \"Use this to find information about: customer meetings, demos, \"\n", - " \"discovery calls, requirements discussions, and decision-maker feedback.\"\n", - " ),\n", - " disable_notice=True,\n", - ")\n", - "\n", - "feedback_retriever = VectorSearchRetrieverTool(\n", - " index_name=f\"{CATALOG_NAME}.{SCHEMA_NAME}.customer_feedback_index\",\n", - " columns=[\"content\", \"doc_uri\"],\n", - " name=\"feedback_search\",\n", - " description=(\n", - " \"Searches through customer feedback and reviews. \"\n", - " \"Use this to find information about: customer satisfaction, \"\n", - " \"product impressions, concerns raised, and post-demo feedback.\"\n", - " ),\n", - " disable_notice=True,\n", - ")\n", - "\n", - "# Combine all RAG tools\n", - "rag_tools = [email_retriever, meeting_notes_retriever, feedback_retriever]\n", - "\n", - "rag_agent_description = (\n", - " \"Specializes in retrieving information from unstructured sales documents including \"\n", - " \"emails, meeting notes, and customer feedback. Use this agent for questions about: \"\n", - " \"customer communications, meeting discussions, feedback and concerns, proposal details, \"\n", - " \"and qualitative sales information.\"\n", - ")\n", - "\n", - "rag_agent = create_react_agent(llm, tools=rag_tools)\n", - "\n", - "###################################################\n", - "## Create SQL Agent for Structured Data\n", - "###################################################\n", - "\n", - "sql_agent_system_prompt = f\"\"\"You are a sales data analyst with access to Workday sales CRM data.\n", - "\n", - "Available tables in {CATALOG_NAME}.{SCHEMA_NAME}:\n", - "- sales_reps: Information about sales representatives\n", - "- customer_accounts: Customer account details and company information\n", - "- sales_opportunities: Sales opportunities and deal pipeline\n", - "- sales_activities: Sales activities and interactions\n", - "\n", - "When asked about structured data like:\n", - "- Sales metrics and KPIs\n", - "- Rep performance and quotas\n", - "- Customer demographics and firmographics\n", - "- Opportunity stages and values\n", - "- Activity tracking and history\n", - "\n", - "Provide analytical insights based on the available structured data tables.\n", - "\"\"\"\n", - "\n", - "genie_description = (\n", - " \"Specializes in analyzing structured sales data from CRM tables. \"\n", - " \"Use this agent for questions about: sales metrics, rep performance, \"\n", - " \"customer demographics, opportunity pipelines, deal values, quota attainment, \"\n", - " \"and quantitative sales analytics.\"\n", - ")\n", - "\n", - "genie_agent = GenieAgent(\n", - " genie_space_id=GENIE_SPACE_ID,\n", - " genie_agent_name=\"Genie\",\n", - " description=genie_description,\n", - " client=WorkspaceClient(\n", - " host=host,\n", - " token=token,\n", - " ),\n", - ")\n", - "\n", - "#############################\n", - "# Define the supervisor agent\n", - "#############################\n", - "\n", - "MAX_ITERATIONS = 4\n", - "\n", - "worker_descriptions = {\n", - " \"GenieAgent\": genie_description,\n", - " \"RAGAgent\": rag_agent_description,\n", - "}\n", - "\n", - "formatted_descriptions = \"\\n\".join(\n", - " f\"- {name}: {desc}\" for name, desc in worker_descriptions.items()\n", - ")\n", - "\n", - "system_prompt = f\"\"\"You are a strategic supervisor coordinating between specialized sales support agents.\n", - "\n", - "Your role is to:\n", - "1. Analyze the user's question to determine which agent(s) can best answer it\n", - "2. Route to the appropriate agent based on the question type\n", - "3. Ensure complete answers without redundant work\n", - "4. Synthesize information from multiple agents if needed\n", - "\n", - "Available agents:\n", - "{formatted_descriptions}\n", - "\n", - "Routing Guidelines:\n", - "- Use GenieAgent for: metrics, numbers, quotas, pipeline values, rep performance, account counts, etc.\n", - "- Use RAGAgent for: customer communications, meeting context, feedback, concerns, proposals, etc.\n", - "- You can route to multiple agents if the question requires both types of information\n", - "\n", - "Only respond with FINISH when:\n", - "- The user's question has been fully answered\n", - "- All necessary information has been gathered and processed\n", - "\n", - "Avoid routing to the same agent multiple times for the same information.\n", - "\n", - "Important:\n", - "- Do not choose FINISH until at least one specialized agent has been invoked.\n", - "- Prefer GenieAgent for numeric/metric queries; RAGAgent for unstructured text queries.\n", - "\"\"\"\n", - "\n", - "options = [\"FINISH\"] + list(worker_descriptions.keys())\n", - "FINISH = {\"next_node\": \"FINISH\"}\n", - "\n", - "@mlflow.trace(span_type=SpanType.AGENT, name=\"supervisor_agent\")\n", - "def supervisor_agent(state):\n", - " count = state.get(\"iteration_count\", 0) + 1\n", - " \n", - " if count > MAX_ITERATIONS:\n", - " return FINISH\n", - " \n", - " class NextNode(BaseModel):\n", - " next_node: Literal[tuple(options)]\n", - "\n", - " preprocessor = RunnableLambda(\n", - " lambda state: [{\"role\": \"system\", \"content\": system_prompt}] + state[\"messages\"]\n", - " )\n", - " supervisor_chain = preprocessor | llm.with_structured_output(NextNode)\n", - " result = supervisor_chain.invoke(state)\n", - " next_node = result.next_node\n", - " \n", - " # Prevent routing to the same node consecutively\n", - " if state.get(\"next_node\") == next_node:\n", - " return FINISH\n", - " \n", - " return {\n", - " \"iteration_count\": count,\n", - " \"next_node\": next_node\n", - " }\n", - "\n", - "#######################################\n", - "# Define multiagent graph structure\n", - "#######################################\n", - "\n", - "def agent_node(state, agent, name):\n", - " \"\"\"Execute agent and return results\"\"\"\n", - " result = agent.invoke({\"messages\": state[\"messages\"]})\n", - " return {\n", - " \"messages\": [\n", - " {\n", - " \"role\": \"assistant\",\n", - " \"content\": result[\"messages\"][-1].content,\n", - " \"name\": name,\n", - " }\n", - " ]\n", - " }\n", - "\n", - "def final_answer(state):\n", - " \"\"\"Generate final synthesized answer\"\"\"\n", - " prompt = (\n", - " \"Based on the information gathered by the specialized agents, \"\n", - " \"provide a comprehensive answer to the user's question. \"\n", - " \"Synthesize insights from all agents and present a clear, helpful response.\"\n", - " )\n", - " preprocessor = RunnableLambda(\n", - " lambda state: state[\"messages\"] + [{\"role\": \"user\", \"content\": prompt}]\n", - " )\n", - " final_answer_chain = preprocessor | llm\n", - " return {\"messages\": [final_answer_chain.invoke(state)]}\n", - "\n", - "class AgentState(ChatAgentState):\n", - " next_node: str\n", - " iteration_count: int\n", - "\n", - "# Create agent nodes\n", - "rag_node = functools.partial(agent_node, agent=rag_agent, name=\"RAGAgent\")\n", - "genie_node = functools.partial(agent_node, agent=genie_agent, name=\"GenieAgent\")\n", - "\n", - "# Build the workflow graph\n", - "workflow = StateGraph(AgentState)\n", - "workflow.add_node(\"GenieAgent\", genie_node)\n", - "workflow.add_node(\"RAGAgent\", rag_node)\n", - "workflow.add_node(\"supervisor\", supervisor_agent)\n", - "workflow.add_node(\"final_answer\", final_answer)\n", - "\n", - "workflow.set_entry_point(\"supervisor\")\n", - "\n", - "# Workers report back to supervisor\n", - "for worker in worker_descriptions.keys():\n", - " workflow.add_edge(worker, \"supervisor\")\n", - "\n", - "# Supervisor decides next node\n", - "workflow.add_conditional_edges(\n", - " \"supervisor\",\n", - " lambda x: x[\"next_node\"],\n", - " {**{k: k for k in worker_descriptions.keys()}, \"FINISH\": \"final_answer\"},\n", - ")\n", - "\n", - "workflow.add_edge(\"final_answer\", END)\n", - "multi_agent = workflow.compile()\n", - "\n", - "###################################\n", - "# Wrap in Databricks ChatAgent\n", - "###################################\n", - "\n", - "class LangGraphChatAgent(ChatAgent):\n", - " def __init__(self, agent: CompiledStateGraph):\n", - " self.agent = agent\n", - "\n", - " def predict(\n", - " self,\n", - " messages: list[ChatAgentMessage],\n", - " context: Optional[ChatContext] = None,\n", - " custom_inputs: Optional[dict[str, Any]] = None,\n", - " ) -> ChatAgentResponse:\n", - " request = {\n", - " \"messages\": [m.model_dump_compat(exclude_none=True) for m in messages]\n", - " }\n", - "\n", - " messages = []\n", - " for event in self.agent.stream(request, stream_mode=\"updates\"):\n", - " for node_data in event.values():\n", - " messages.extend(\n", - " ChatAgentMessage(**msg) for msg in node_data.get(\"messages\", [])\n", - " )\n", - " return ChatAgentResponse(messages=messages)\n", - "\n", - " def predict_stream(\n", - " self,\n", - " messages: list[ChatAgentMessage],\n", - " context: Optional[ChatContext] = None,\n", - " custom_inputs: Optional[dict[str, Any]] = None,\n", - " ) -> Generator[ChatAgentChunk, None, None]:\n", - " request = {\n", - " \"messages\": [m.model_dump_compat(exclude_none=True) for m in messages]\n", - " }\n", - " for event in self.agent.stream(request, stream_mode=\"updates\"):\n", - " for node_data in event.values():\n", - " yield from (\n", - " ChatAgentChunk(**{\"delta\": msg})\n", - " for msg in node_data.get(\"messages\", [])\n", - " )\n", - "\n", - "# Create the agent\n", - "mlflow.langchain.autolog()\n", - "AGENT = LangGraphChatAgent(multi_agent)\n", - "mlflow.models.set_model(AGENT)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "1eb839e2-725b-4ab0-80a2-f4a7c8fb4080", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "import mlflow\n", - "from agent import *\n", - "from databricks.sdk import WorkspaceClient \n", - "\n", - "w = WorkspaceClient()\n", - "current_user = w.current_user.me()\n", - "\n", - "experiment_fqdn = (\n", - " f\"/Workspace/Users/{current_user.user_name}/{MLFLOW_EXPERIMENT_NAME}\"\n", - ")\n", - "\n", - "# Check if the experiment exists\n", - "experiment = mlflow.get_experiment_by_name(experiment_fqdn)\n", - "\n", - "if experiment:\n", - " experiment_id = experiment.experiment_id\n", - "else:\n", - " # Create the experiment if it does not exist\n", - " experiment_id = mlflow.create_experiment(experiment_fqdn)\n", - "\n", - "mlflow.set_experiment(experiment_fqdn)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "5643ba67-edb0-4d98-af82-e0190bc97b5b", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "from IPython.display import Image, display\n", - "\n", - "display(Image(AGENT.agent.get_graph().draw_mermaid_png()))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "e0ff325d-72b3-4b69-8b40-e6edf7ed7b98", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "%md\n", - "\n", - "### Test the LangGraph Agent" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "2e121c60-ca5c-4f8e-b898-3569ca1cb4c4", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "sample_questions = [\n", - " \"How many outbound emails were sent by each sales rep last quarter?\",\n", - " \"Summarize the main concerns customers raised about implementation delays.\",\n", - " \"Extract any customer feedback that praises response time or support quality.\"\n", - "]\n", - "\n", - "input_example = {\n", - " \"messages\": [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": sample_questions[0],\n", - " }\n", - " ]\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "e6f88001-3798-4965-9474-cdbe85292ade", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "response = AGENT.predict(input_example)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "1b476203-0f02-49a8-bd7b-20bc5bfabe2e", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "input_example = {\n", - " \"messages\": [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": sample_questions[1],\n", - " }\n", - " ]\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "991d46c9-db33-437b-b4ba-a9c06887eef6", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "response = AGENT.predict(input_example)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "5b86bd96-a575-4aa9-b52b-7a8539cb8187", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "input_example = {\n", - " \"messages\": [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": sample_questions[2],\n", - " }\n", - " ]\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "c4b474b2-96a7-430c-bd08-8b18367b8df1", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "response = AGENT.predict(input_example)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "3a4c5be2-8644-4f55-afcc-61a2c7bdabb5", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "%md\n", - "\n", - "### Log into MLFlow" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "9268b3c7-f1a1-4246-980e-66f8c39c75a3", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "# Determine Databricks resources to specify for automatic auth passthrough at deployment time\n", - "import mlflow\n", - "from pkg_resources import get_distribution\n", - "from databricks_langchain import UnityCatalogTool, VectorSearchRetrieverTool\n", - "from mlflow.models.resources import (\n", - " DatabricksFunction,\n", - " DatabricksGenieSpace,\n", - " DatabricksServingEndpoint,\n", - " DatabricksVectorSearchIndex,\n", - " DatabricksSQLWarehouse,\n", - " DatabricksTable,\n", - ")\n", - "\n", - "resources = [\n", - " DatabricksServingEndpoint(endpoint_name=\"databricks-claude-3-7-sonnet\"),\n", - " DatabricksGenieSpace(genie_space_id=GENIE_SPACE_ID),\n", - " DatabricksSQLWarehouse(warehouse_id='4b9b953939869799'),\n", - "]\n", - "\n", - "# for tool in tools:\n", - "# if isinstance(tool, VectorSearchRetrieverTool):\n", - "# resources.extend(tool.resources)\n", - "# elif isinstance(tool, UnityCatalogTool):\n", - "# resources.append(DatabricksFunction(function_name=tool.uc_function_name))\n", - "\n", - "with mlflow.start_run():\n", - " logged_agent_info = mlflow.pyfunc.log_model(\n", - " name=f\"multi_agent_at\",\n", - " python_model=\"agent.py\",\n", - " # model_config=os.path.join(os.getcwd(), \"00-init-requirements\"),\n", - " input_example=input_example,\n", - " resources=resources,\n", - " pip_requirements=[\n", - " f\"databricks-connect=={get_distribution('databricks-connect').version}\",\n", - " f\"mlflow=={get_distribution('mlflow').version}\",\n", - " f\"databricks-langchain=={get_distribution('databricks-langchain').version}\",\n", - " f\"langgraph=={get_distribution('langgraph').version}\",\n", - " ]\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "038afed1-cfe8-4cda-bde7-4842d39b06fe", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "\n", - "### Local Validation" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "a7dff315-894b-47e2-bf33-c13dd479b66c", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "input_example = {\n", - " 'messages': [\n", - " {'role': 'user',\n", - " 'content': 'provide a summary of what do our policies say about customers eligible for promotions but receiving poor service?'\n", - " }\n", - " ]\n", - " }\n", - " \n", - "response = mlflow.models.predict(\n", - " model_uri=f\"runs:/{logged_agent_info.run_id}/multi_agent_at\",\n", - " input_data=input_example,\n", - " env_manager=\"uv\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "b59c2cd8-e94e-4b76-9166-92b4ade3fdd8", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "\n", - "### Register Model" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "479f654b-4dbb-40c8-9a13-5e5ca87f8bd4", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "mlflow.set_registry_uri(\"databricks-uc\")\n", - "\n", - "# Define the catalog, schema, and model name for your UC model\n", - "catalog = catalog_name\n", - "schema = schema_name\n", - "model_name = f\"multi_agent_demo_{catalog_name}\"\n", - "UC_MODEL_NAME = f\"{catalog}.{schema}.{model_name}\"\n", - "\n", - "# register the model to UC\n", - "uc_registered_model_info = mlflow.register_model(\n", - " model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": {}, - "inputWidgets": {}, - "nuid": "20b3882e-1d8e-4b46-b0e6-cd0edd94863d", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "source": [ - "\n", - "### Create Endpoint" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "application/vnd.databricks.v1+cell": { - "cellMetadata": { - "byteLimit": 2048000, - "rowLimit": 10000 - }, - "inputWidgets": {}, - "nuid": "d88590e5-c50d-43f6-9cb3-405447f8cf9b", - "showTitle": false, - "tableResultSettingsMap": {}, - "title": "" - } - }, - "outputs": [], - "source": [ - "from databricks import agents\n", - "\n", - "agents.deploy(\n", - " UC_MODEL_NAME,\n", - " uc_registered_model_info.version,\n", - " tags={\"endpointSource\": \"playground\"},\n", - " endpoint_name = f\"multi_agent_{catalog_name}\",\n", - " environment_vars={\n", - " \"DATABRICKS_GENIE_PAT\": token\n", - " },\n", - ")" - ] - } - ], - "metadata": { - "application/vnd.databricks.v1+notebook": { - "computePreferences": { - "hardware": { - "accelerator": null, - "gpuPoolId": null, - "memory": null - } - }, - "dashboards": [], - "environmentMetadata": { - "base_environment": "", - "environment_version": "2" - }, - "inputWidgetPreferences": null, - "language": "python", - "notebookMetadata": { - "mostRecentlyExecutedCommandWithImplicitDF": { - "commandId": -1, - "dataframes": [ - "_sqldf" - ] - }, - "pythonIndentUnit": 4 - }, - "notebookName": "03-create-multi-agent-with-genie", - "widgets": {} - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 0 -}