Начало работы с агентами ИИ Databricks

В этом руководстве описано, как построить агент ИИ, использующий извлечение и инструменты вместе. Набор данных состоит из подмножества документации Databricks, которая уже разделена на блоки. В этом руководстве вы создадите агент, который извлекает документы на основе ключевых слов.

Для простоты в этом руководстве используется простой подход в памяти на основе TF-IDF для извлечения ключевых слов и поиска ключевых слов по документам. В примере записной книжки содержится весь код, используемый в руководстве. Более реалистичный пример, использующий Mosaic AI Vector Search для масштабируемого индексирования и поиска документов, см. в примерах ChatAgent.

В этом руководстве рассматриваются некоторые основные проблемы создания создаваемых приложений ИИ:

• Упрощение процесса разработки для распространенных задач, таких как создание инструментов и выполнение агента отладки.
• Такие операционные проблемы, как:
  • Настройка агента отслеживания
  • Определение входных и выходных данных предсказуемым способом
  • Управление версиями зависимостей
  • Управление версиями и развертывание
• Измерение и улучшение качества и надежности агента.

См. руководство по навигации в приложении генеративного ИИ для полного руководства по процессу разработки агента.

пример записной книжки

Эта автономная записная книжка предназначена для быстрого работы с агентами ИИ Мозаики с помощью примера корпуса документов. Она готова к запуску без необходимости установки или предоставления данных.

Демонстрация агента ИИ Mosaic

Возьмите записную книжку

Создание агента и инструментов

Платформа агента ИИ Мозаики поддерживает множество различных платформ разработки. В этом примере используется LangGraph для иллюстрации концепций, но это не руководство по LangGraph.

Примеры других поддерживаемых платформ см. в примерах ChatAgent.

Первым шагом является создание агента. Необходимо указать клиент LLM и список инструментов. Пакет Python databricks-langchain включает клиентов, совместимых как с LangChain, так и с LangGraph, для использования с LLM от Databricks и инструментами, зарегистрированными в Unity Catalog.

Конечная точка должна быть API-интерфейсом модели Foundation или внешней моделью с помощью шлюза ИИ. См. раздел "Поддерживаемые модели".

from databricks_langchain import ChatDatabricks
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

Следующий код определяет функцию, которая создает агент, используя модель и некоторые инструменты; обсуждение внутренних процессов этого кода агента выходит за рамки данного руководства. Дополнительные сведения о том, как создать агента LangGraph, см. в документации по LangGraph .

from typing import Optional, Sequence, Union

from langchain_core.language_models import LanguageModelLike
from langchain_core.runnables import RunnableConfig, RunnableLambda
from langchain_core.tools import BaseTool
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
from mlflow.langchain.chat_agent_langgraph import ChatAgentState, ChatAgentToolNode


def create_tool_calling_agent(
  model: LanguageModelLike,
  tools: Union[ToolExecutor, Sequence[BaseTool]],
  agent_prompt: Optional[str] = None,
) -> CompiledGraph:
  model = model.bind_tools(tools)

  def routing_logic(state: ChatAgentState):
    last_message = state["messages"][-1]
    if last_message.get("tool_calls"):
      return "continue"
    else:
      return "end"

  if agent_prompt:
    system_message = {"role": "system", "content": agent_prompt}
    preprocessor = RunnableLambda(
      lambda state: [system_message] + state["messages"]
    )
  else:
    preprocessor = RunnableLambda(lambda state: state["messages"])
  model_runnable = preprocessor | model

  def call_model(
    state: ChatAgentState,
    config: RunnableConfig,
  ):
    response = model_runnable.invoke(state, config)

    return {"messages": [response]}

  workflow = StateGraph(ChatAgentState)

  workflow.add_node("agent", RunnableLambda(call_model))
  workflow.add_node("tools", ChatAgentToolNode(tools))

  workflow.set_entry_point("agent")
  workflow.add_conditional_edges(
    "agent",
    routing_logic,
    {
      "continue": "tools",
      "end": END,
    },
  )
  workflow.add_edge("tools", "agent")

  return workflow.compile()

Определение средств агента

Инструменты — это фундаментальная концепция для создания агентов. Они предоставляют возможность интеграции моделей машинного обучения (LLM) с кодом, написанным человеком. При предоставлении запроса и списка инструментов, LLM создает аргументы для вызова инструмента. Для получения дополнительной информации об инструментах и их использовании с агентами Mosaic AI см. инструменты агентов ИИ.

Первым шагом является создание средства извлечения ключевых слов на основе TF-IDF. В этом примере используется scikit-learn и средство каталога Unity.

Пакет databricks-langchain предоставляет удобный способ работы с инструментами каталога Unity. В следующем коде показано, как реализовать и зарегистрировать средство извлечения ключевых слов.

Примечание.

Рабочая область Databricks имеет встроенное средство system.ai.python_exec, которое можно использовать для расширения агентов с возможностью выполнения скриптов Python в изолированной среде выполнения. Другие полезные встроенные средства включают внешние подключения и функции ИИ .

from databricks_langchain.uc_ai import (
  DatabricksFunctionClient,
  UCFunctionToolkit,
  set_uc_function_client,
)

uc_client = DatabricksFunctionClient()
set_uc_function_client(client)

# Change this to your catalog and schema
CATALOG = "main"
SCHEMA = "my_schema"


def tfidf_keywords(text: str) -> list[str]:
  """
  Extracts keywords from the provided text using TF-IDF.

  Args:
    text (string): Input text.
  Returns:
    list[str]: List of extracted keywords in ascending order of importance.
  """
  from sklearn.feature_extraction.text import TfidfVectorizer

  def keywords(text, top_n=5):
    vec = TfidfVectorizer(stop_words="english")
    tfidf = vec.fit_transform([text])  # Convert text to TF-IDF matrix
    indices = tfidf.toarray().argsort()[0, -top_n:]  # Get indices of top N words
    return [vec.get_feature_names_out()[i] for i in indices]

  return keywords(text)


# Create the function in the Unity Catalog catalog and schema specified
# When you use `.create_python_function`, the provided function’s metadata
# (docstring, parameters, return type) are used to create a tool in the specified catalog and schema.
function_info = uc_client.create_python_function(
  func=tfidf_keywords,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True,  # Set to True to overwrite if the function already exists
)

print(function_info)

Ниже приведено описание приведенного выше кода:

1. Создает клиент, использующий каталог Unity в рабочей области Databricks в качестве реестра для создания и обнаружения инструментов.
2. Определяет функцию Python, которая выполняет извлечение ключевых слов TF-IDF.
3. Регистрирует функцию Python в качестве функции каталога Unity.

Этот рабочий процесс решает несколько распространенных проблем. Теперь у вас есть центральный реестр для инструментов, которые, как и другие объекты в каталоге Unity, можно управлять. Например, если у компании есть стандартный способ вычисления внутренней скорости возврата, его можно определить как функцию в каталоге Unity и предоставить доступ всем пользователям или агентам с помощью роли FinancialAnalyst.

Чтобы сделать этот инструмент пригодным для использования агентом LangChain, используйте [UCFunctionToolkit](/generative-ai/agent-framework/create-custom-tool.md), который создает коллекцию инструментов, чтобы LLM мог выбрать.

# Use ".*" here to specify all the tools in the schema, or
# explicitly list functions by name
# uc_tool_names = [f"{CATALOG}.{SCHEMA}.*"]
uc_tool_names = [f"{CATALOG}.{SCHEMA}.tfidf_keywords"]
uc_toolkit = UCFunctionToolkit(function_names=uc_tool_names)

В следующем коде показано, как протестировать средство:

uc_toolkit.tools[0].invoke({ "text": "The quick brown fox jumped over the lazy brown dog." })

Следующий код создает агент, использующий средство извлечения ключевых слов.

import mlflow
mlflow.langchain.autolog()

agent = create_tool_calling_agent(llm, tools=[*uc_toolkit.tools])

agent.invoke({"messages": [{"role": "user", "content":"What are the keywords for the sentence: 'the quick brown fox jumped over the lazy brown dog'?"}]})

В результирующей трассировке можно видеть, что LLM выбрал инструмент.

Использование трассировок для отладки агентов

MLflow Tracing — это мощный инструмент для отладки и наблюдения за генеративными ИИ-приложениями, включая агентов. Он записывает подробные сведения о выполнении с помощью диапазонов, которые инкапсулируют определенные сегменты кода и записывают входные данные, выходные данные и данные о времени.

Для популярных библиотек, таких как LangChain, включите автоматическую трассировку с помощью mlflow.langchain.autolog(). Вы также можете использовать mlflow.start_span() для настройки трассировки. Например, можно добавить поля настраиваемых значений данных или метки для наблюдения. Код, который выполняется в контексте этого диапазона, получает связь с заданными полями. В этом примере TF-IDF, находящемся в памяти, дайте ему имя и укажите тип диапазона.

Чтобы узнать больше о трассировке, см. раздел о наблюдаемости агентов с использованием трассировки MLflow.

В следующем примере создается средство извлечения с помощью простого индекса TF-IDF, хранящегося в памяти. В нем демонстрируется автологирование выполнения инструментов и пользовательская трассировка спанов для дополнительной наблюдаемости.

from sklearn.feature_extraction.text import TfidfVectorizer
import mlflow
from langchain_core.tools import tool


documents = parsed_docs_df
doc_vectorizer = TfidfVectorizer(stop_words="english")
tfidf_matrix = doc_vectorizer.fit_transform(documents["content"])


@tool
def find_relevant_documents(query, top_n=5):
  """gets relevant documents for the query"""
  with mlflow.start_span(name="LittleIndex", span_type="RETRIEVER") as retriever_span:
    retriever_span.set_inputs({"query": query})
    retriever_span.set_attributes({"top_n": top_n})

    query_tfidf = doc_vectorizer.transform([query])
    similarities = (tfidf_matrix @ query_tfidf.T).toarray().flatten()
    ranked_docs = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)

    result = []
    for idx, score in ranked_docs[:top_n]:
      row = documents.iloc[idx]
      content = row["content"]
      doc_entry = {
        "page_content": content,
        "metadata": {
          "doc_uri": row["doc_uri"],
          "score": score,
        },
      }
      result.append(doc_entry)

    retriever_span.set_outputs(result)
    return result

Этот код использует специальный тип диапазона, RETRIEVER, зарезервированный для средств извлечения. Другие функции агента ИИ Мозаики (такие как ИИ-площадка, интерфейс для рецензирования и оценка) используют тип диапазона RETRIEVER для отображения результатов извлечения.

Ретриверы требуют указания их схемы для обеспечения совместимости с последующими функциями Databricks. Дополнительные сведения о mlflow.models.set_retriever_schemaсм. в разделе Укажите пользовательские схемы извлекателя.

import mlflow
from mlflow.models import set_retriever_schema

uc_toolkit = UCFunctionToolkit(function_names=[f"{CATALOG}.{SCHEMA}.*"])

graph = create_tool_calling_agent(llm, tools=[*uc_toolkit.tools, find_relevant_documents])

mlflow.langchain.autolog()
set_retriever_schema(
  primary_key="chunk_id",
  text_column="chunk_text",
  doc_uri="doc_uri",
  other_columns=["title"],
)

graph.invoke(input = {"messages": [("user", "How do the docs say I use llm judges on databricks?")]})

Определение агента

Следующим шагом является оценка агента и его подготовка к развертыванию. На высоком уровне это включает в себя следующее:

1. Определите предсказуемый API для агента с помощью сигнатуры.
2. Добавьте конфигурацию модели, что упрощает настройку параметров.
3. Запишите модель с зависимостями, которые обеспечивают ей воспроизводимую среду и позволяют настроить аутентификацию для взаимодействия с другими службами.

Интерфейс MLflow ChatAgent упрощает определение входных и выходных данных агента. Чтобы использовать его, определите ваш агент как подкласс ChatAgent, реализуя вывод для непотоковой передачи с помощью функции predict и вывод для потоковой передачи с помощью функции predict_stream.

ChatAgent не зависит от выбранной платформы разработки агента, что позволяет легко тестировать и использовать различные платформы и реализации агента. Единственное требование заключается в реализации predict и predict_stream интерфейсов.

Создание агента с помощью ChatAgent предоставляет ряд преимуществ, в том числе:

• Поддержка потокового вывода
• Подробная история сообщений вызова инструментов: Возвращает несколько сообщений, включая промежуточные сообщения вызова инструментов, для повышения качества и более эффективного управления диалогами.
• поддержка многоагентных систем
• интеграция функций Databricks: изначальная совместимость с платформой ИИ, оценкой агентов и мониторингом агентов.
• Типизированные интерфейсы создания: напишите код агента, используя типизированные классы Python, и воспользуйтесь преимуществами автозаполнения интегрированной среды разработки и записной книжки.

Дополнительные сведения о создании ChatAgent см. в . Используйте ChatAgent для создания агентов.

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
  ChatAgentChunk,
  ChatAgentMessage,
  ChatAgentResponse,
  ChatContext,
)
from typing import Any, Optional


class DocsAgent(ChatAgent):
  def __init__(self, agent):
    self.agent = agent
    set_retriever_schema(
      primary_key="chunk_id",
      text_column="chunk_text",
      doc_uri="doc_uri",
      other_columns=["title"],
    )

  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    # ChatAgent has a built-in helper method to help convert framework-specific messages, like langchain BaseMessage to a python dictionary
    request = {"messages": self._convert_messages_to_dict(messages)}

    output = agent.invoke(request)
    # Here 'output' is already a ChatAgentResponse, but to make the ChatAgent signature explicit for this demonstration, the code returns a new instance
    return ChatAgentResponse(**output)

В следующем коде показано, как использовать ChatAgent.

AGENT = DocsAgent(agent=agent)
AGENT.predict(
  {
    "messages": [
      {"role": "user", "content": "What is DLT in Databricks?"},
    ]
  }
)

Настройка агентов с параметрами

Платформа агента позволяет управлять выполнением агента с параметрами. Это означает, что вы можете быстро протестировать различные конфигурации агента, например переключить конечные точки LLM или попробовать различные средства, не изменяя базовый код.

Следующий код создает словарь конфигурации, который задает параметры агента при инициализации модели.

Дополнительные сведения о параметризации кода агенток для развертывания в различных средах см. в разделе .

)

from mlflow.models import ModelConfig

baseline_config = {
  "endpoint_name": "databricks-meta-llama-3-1-70b-instruct",
  "temperature": 0.01,
  "max_tokens": 1000,
  "system_prompt": """You are a helpful assistant that answers questions about Databricks. Questions unrelated to Databricks are irrelevant.


  You answer questions using a set of tools. If needed, you ask the user follow-up questions to clarify their request.
  """,
  "tool_list": ["catalog.schema.*"],
}


class DocsAgent(ChatAgent):
  def __init__(self):
    self.config = ModelConfig(development_config=baseline_config)
    self.agent = self._build_agent_from_config()


def _build_agent_from_config(self):
  temperature = config.get("temperature", 0.01)
  max_tokens = config.get("max_tokens", 1000)
  system_prompt = config.get("system_prompt", """You are a helpful assistant.
    You answer questions using a set of tools. If needed you ask the user follow-up questions to clarify their request.""")
  llm_endpoint_name = config.get("endpoint_name", "databricks-meta-llama-3-3-70b-instruct")
  tool_list = config.get("tool_list", [])

  llm = ChatDatabricks(endpoint=llm_endpoint_name, temperature=temperature, max_tokens=max_tokens)
  toolkit = UCFunctionToolkit(function_names=tool_list)
  agent = create_tool_calling_agent(llm, tools=[*toolkit.tools, find_relevant_documents], prompt=system_prompt)

  return agent

Захегистрируйте агента

После определения агента он теперь готов к регистрации. В MLflow ведение журнала агента означает сохранение конфигурации агента (включая зависимости), чтобы его можно было использовать для оценки и развертывания.

Примечание.

При разработке агентов в записной книжке MLflow выводит зависимости агента из среды записной книжки.

Чтобы записать агент из записной книжки, можно написать весь код, определяющий модель в одной ячейке, а затем использовать команду %%writefile для сохранения определения агента в файле:

%%writefile agent.py
...
<Code that defines the agent>

Если агенту требуется доступ к внешним ресурсам, таким как каталог Unity для выполнения средства извлечения ключевых слов, необходимо настроить проверку подлинности агента, чтобы он смог получить доступ к ресурсам при развертывании.

Чтобы упростить проверку подлинности для ресурсов Databricks, включите сквозноеавтоматической проверки подлинности:

from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint


resources = [
  DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
  DatabricksFunction(function_name=tool.uc_function_name),
]


with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    artifact_path="agent",
    python_model="agent.py",
    pip_requirements=[
      "mlflow",
      "langchain",
      "langgraph",
      "databricks-langchain",
      "unitycatalog-langchain[databricks]",
      "pydantic",
    ],
    resources=resources,
  )

Дополнительные сведения об агентах ведения журнала см. в логирование на основе кода.

Оценка агента

Следующим шагом является оценка агента, чтобы узнать, как он работает. Оценка агента является сложной задачей и вызывает множество вопросов, таких как:

• Каковы правильные метрики для оценки качества? Как доверять выходным данным этих метрик?
• Мне нужно оценить много идей - как я...
  • Запустить оценку быстро, так чтобы большую часть моего времени не тратить на ожидание?
  • Быстро сравните эти разные версии моего агента по качеству, цене и задержке.
• Как быстро определить первопричину каких-либо проблем качества?

Как специалист по обработке и анализу данных или разработчику, вы не можете быть фактическим экспертом по теме. В остальной части этого раздела описываются средства оценки агента, которые помогут вам определить хорошие выходные данные.

Создайте оценочный набор

Чтобы определить, какое качество означает для агента, используйте метрики для измерения производительности агента в оценочном наборе. См. определение «качество»: оценочные наборы.

С помощью Оценки агента можно создавать искусственные наборы оценки и измерять качество, проводя их. Идея заключается в том, чтобы начать с фактов, таких как набор документов, и "работать назад", используя эти факты для создания набора вопросов. Вы можете задавать условия для вопросов, которые создаются, путем предоставления некоторых направлений.

from databricks.agents.evals import generate_evals_df
import pandas as pd


databricks_docs_url = "https://raw.githubusercontent.com/databricks/genai-cookbook/refs/heads/main/quick_start_demo/chunked_databricks_docs_filtered.jsonl"
parsed_docs_df = pd.read_json(databricks_docs_url, lines=True)


agent_description = f"""
The agent is a RAG chatbot that answers questions about Databricks. Questions unrelated to Databricks are irrelevant.
"""
question_guidelines = f"""
# User personas
- A developer who is new to the Databricks platform
- An experienced, highly technical Data Scientist or Data Engineer


# Example questions
- what API lets me parallelize operations over rows of a delta table?
- Which cluster settings will give me the best performance when using Spark?


# Additional Guidelines
- Questions should be succinct, and human-like
"""


num_evals = 25
evals = generate_evals_df(
  docs=parsed_docs_df[
    :500
  ],  # Pass your docs. They should be in a Pandas or Spark DataFrame with columns `content STRING` and `doc_uri STRING`.
  num_evals=num_evals,  # How many synthetic evaluations to generate
  agent_description=agent_description,
  question_guidelines=question_guidelines,
)

Созданные оценки включают следующие:

• Поле запроса, аналогичное упомянутому ранее ChatAgentRequest:

  {"messages":[{"content":"What command must be run at the start of your workload to explicitly target the Workspace Model Registry if your workspace default catalog is in Unity Catalog and you use Databricks Runtime 13.3 LTS or above?","role":"user"}]}
  

• Список "ожидаемого полученного содержимого". Схема извлекателя была определена с полями content и doc_uri.

  [{"content":"If your workspace’s [default catalog](https://docs.databricks.com/data-governance/unity-catalog/create-catalogs.html#view-the-current-default-catalog) is in Unity Catalog (rather than `hive_metastore`) and you are running a cluster using Databricks Runtime 13.3 LTS or above, models are automatically created in and loaded from the workspace default catalog, with no configuration required. To use the Workspace Model Registry in this case, you must explicitly target it by running `import mlflow; mlflow.set_registry_uri(\"databricks\")` at the start of your workload.","doc_uri":"https://docs.databricks.com/machine-learning/manage-model-lifecycle/workspace-model-registry.html"}]
  

• Список ожидаемых фактов. При сравнении двух ответов трудно найти небольшие различия между ними. Ожидаемые факты отделяют правильный ответ от частично правильного и неправильного ответов, улучшая как качество судей ИИ, так и опыт людей, работающих с агентом.

  ["The command must import the MLflow module.","The command must set the registry URI to \"databricks\"."]
  

• Поле source_id, которое здесь обозначено как SYNTHETIC_FROM_DOC. При создании более полных наборов оценки примеры будут поступать из разных источников, поэтому это поле различает их.

Более подробную информацию о создании наборов для оценки см. в разделе Синтезируйте наборы для оценки.

Оцените агента, используя судей LLM.

Ручная оценка производительности агента с использованием такого большого количества созданных примеров не масштабируется хорошо. В большом масштабе использование LLM в качестве судей является гораздо более разумным решением. Чтобы использовать встроенные судьи, доступные при использовании оценки агента, используйте следующий код:

with mlflow.start_run(run_name="my_agent"):
  eval_results = mlflow.evaluate(
    data=evals,  # Your evaluation set
    model=model_info.model_uri,  # Logged agent from above
    model_type="databricks-agent",  # activate Mosaic AI Agent Evaluation
)

Простой агент набрал 68% в целом. Результаты могут отличаться здесь в зависимости от используемой конфигурации. Выполнение эксперимента, чтобы сравнить три разных языковые модели (LLM) по затратам и качеству, так же просто, как изменение конфигурации и переоценка.

Попробуйте изменить конфигурацию модели, чтобы использовать другой параметр LLM, системного запроса или температуры.

Эти судьи могут быть настроены для того, чтобы следовать тем же рекомендациям, которые используют эксперты для оценки ответа. Дополнительную информацию о встроенных судьях ИИ LLM см. в .

С помощью функции оценки агентов можно настроить способ измерения качества конкретного агента с помощью настраиваемых метрик . Вы можете рассматривать оценку как тест интеграции и отдельные метрики как модульные тесты. В следующем примере используется логическая метрика для проверки того, использовал ли агент извлечение ключевых слов и извлекатель для данного запроса:

from databricks.agents.evals import metric

@metric
def uses_keywords_and_retriever(request, trace):
  retriever_spans = trace.search_spans(span_type="RETRIEVER")
  keyword_tool_spans = trace.search_spans(name=f"{CATALOG}__{SCHEMA}__tfidf_keywords")
  return len(keyword_tool_spans) > 0 and len(retriever_spans) > 0


# same evaluate as above, with the addition of 'extra_metrics'
with mlflow.start_run(run_name="my_agent"):
  eval_results = mlflow.evaluate(
    data=evals,  # Your evaluation set
    model=model_info.model_uri,  # Logged agent from above
    model_type="databricks-agent",  # activate Mosaic AI Agent Evaluation,
    extra_metrics=[uses_keywords_and_retriever],
  )

Обратите внимание, что агент никогда не использует извлечение ключевых слов. Как устранить эту проблему?

Развертывание и мониторинг агента

Когда вы будете готовы к тестированию агента с реальными пользователями, Agent Framework предоставляет готовое к работе решение для обслуживания агента в Службе модели ИИ Мозаики.

Развертывание агентов в Службе моделей обеспечивает следующие преимущества:

• Служба моделей управляет автомасштабированием, ведением журнала, управлением версиями и контролем доступа, что позволяет сосредоточиться на разработке агентов качества.
• Эксперты по темам могут использовать приложение проверки для взаимодействия с агентом и предоставления отзывов, которые можно включить в мониторинг и оценки.
• Вы можете контролировать работу агента, проводя оценки на реальном трафике. Хотя трафик пользователей не будет включать истинные данные, судьи LLM (и пользовательская метрика, которую вы создали) проводят оценку без надзора.

Следующий код развертывает агенты в конечной точке обслуживания. Дополнительные сведения см. в статье Развертывание агента для создания приложений ИИ.

from databricks import agents
import mlflow

# Connect to the Unity Catalog model registry
mlflow.set_registry_uri("databricks-uc")

# Configure UC model location
UC_MODEL_NAME = f"{CATALOG}.{SCHEMA}.getting_started_agent"
# REPLACE WITH UC CATALOG/SCHEMA THAT YOU HAVE `CREATE MODEL` permissions in

# Register to Unity Catalog
uc_registered_model_info = mlflow.register_model(
  model_uri=model_info.model_uri, name=UC_MODEL_NAME
)
# Deploy to enable the review app and create an API endpoint
deployment_info = agents.deploy(
  model_name=UC_MODEL_NAME, model_version=uc_registered_model_info.version
)