notebook/agentchat_databricks_dbrx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use AG2 in Databricks with DBRX\n",
    "\n",
    "![DBRX launch](https://www.databricks.com/sites/default/files/styles/max_2600x2600/public/2024-03/dbrx-blog-header-optimized.png)\n",
    "\n",
    "In March 2024, Databricks released [DBRX](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm), a general-purpose LLM that sets a new standard for open LLMs. While available as an open-source model on Hugging Face ([databricks/dbrx-instruct](https://huggingface.co/databricks/dbrx-instruct/tree/main) and [databricks/dbrx-base](https://huggingface.co/databricks/dbrx-base) ), customers of Databricks can also tap into the [Foundation Model APIs](https://docs.databricks.com/en/machine-learning/model-serving/score-foundation-models.html#query-a-chat-completion-model), which make DBRX available through an OpenAI-compatible, autoscaling REST API.\n",
    "\n",
    "[Autogen](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/conversable-agent/) is becoming a popular standard for agent creation. Built to support any \"LLM as a service\" that implements the OpenAI SDK, it can easily be extended to integrate with powerful open source models. \n",
    "\n",
    "This notebook will demonstrate a few basic examples of Autogen with DBRX, including the use of  `AssistantAgent`, `UserProxyAgent`, and `ConversableAgent`. These demos are not intended to be exhaustive - feel free to use them as a base to build upon!\n",
    "\n",
    "## Requirements\n",
    "AG2 must be installed on your Databricks cluster, and requires `Python>=3.9`. This example includes the `%pip` magic command to install: `%pip install pyautogen`, as well as other necessary libraries. \n",
    "\n",
    "This code has been tested on: \n",
    "* [Serverless Notebooks](https://docs.databricks.com/en/compute/serverless/index.html) (in public preview as of Apr 18, 2024)\n",
    "* Databricks Runtime 14.3 LTS ML [docs](https://docs.databricks.com/aws/en/release-notes/runtime/14.3lts-ml)\n",
    "\n",
    "This code can run in any Databricks workspace in a region where DBRX is available via pay-per-token APIs (or provisioned throughput). To check if your region is supported, see [Foundation Model Region Availability](https://docs.databricks.com/en/machine-learning/model-serving/model-serving-limits.html#foundation-model-apis-limits). If the above is true, the workspace must also be enabled by an admin for Foundation Model APIs [docs](https://docs.databricks.com/en/machine-learning/foundation-model-apis/index.html#requirements).\n",
    "\n",
    "## Tips\n",
    "* This notebook can be imported from github to a Databricks workspace and run directly. Use [sparse checkout mode with git](https://www.databricks.com/blog/2023/01/26/work-large-monorepos-sparse-checkout-support-databricks-repos.html) to import only this notebook or the examples directory. \n",
    "\n",
    "* Databricks recommends using [Secrets](https://docs.databricks.com/en/security/secrets/secrets.html) instead of storing tokens in plain text. \n",
    "\n",
    "## Contributor\n",
    "\n",
    "tj@databricks.com (Github: tj-cycyota)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "%pip install pyautogen==0.2.25 openai==1.21.2 typing_extensions==4.11.0 --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is recommended to restart the Python kernel after installs - uncomment and run the below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dbutils.library.restartPython()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup DBRX config list\n",
    "\n",
    "See Autogen docs for more information on the use of `config_list`: [LLM Configuration](https://docs.ag2.ai/docs/user-guide/advanced-concepts/llm-configuration-deep-dive#llm-configuration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from autogen import LLMConfig\n",
    "\n",
    "# Set environment variables with your current workspace host and a personal access token\n",
    "# To a secret you have already set up: dbutils.secrets.get('your_scope_name','databricks_host')\n",
    "\n",
    "# DATABRICKS_HOST format: \"https://{your workspace url}\" (no trailing slash)\n",
    "# AWS Workspace example: \"https://my-databricks-workspace.cloud.databricks.com\"\n",
    "# Azure Workspace example: \"https://adb-123456790123.12.azuredatabricks.net\"\n",
    "os.environ[\"DATABRICKS_HOST\"] = \"<FILL IN WITH YOUR WORKSPACE URL IN SUPPORTED DBRX REGION>\"\n",
    "\n",
    "# DATABRICKS_TOKEN format: \"dapi...\"\n",
    "# Temp token: dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()\n",
    "os.environ[\"DATABRICKS_TOKEN\"] = \"dapi....\"\n",
    "\n",
    "llm_config = LLMConfig(\n",
    "    config_list=[\n",
    "        {\n",
    "            \"model\": \"databricks-dbrx-instruct\",\n",
    "            \"api_key\": str(os.environ[\"DATABRICKS_TOKEN\"]),\n",
    "            \"base_url\": str(os.getenv(\"DATABRICKS_HOST\")) + \"/serving-endpoints\",\n",
    "        }\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hello World Example\n",
    "\n",
    "Our first example will be with a simple `UserProxyAgent` asking a question to an `AssistantAgent`. This is based on the tutorial demo [here](https://docs.ag2.ai/docs/home/quick-start).\n",
    "\n",
    "After sending the question and seeing a response, you can type `exit` to end the chat or continue to converse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import autogen\n",
    "\n",
    "# Create Assistant and User\n",
    "assistant = autogen.AssistantAgent(name=\"assistant\", llm_config=llm_config)\n",
    "\n",
    "user_proxy = autogen.UserProxyAgent(name=\"user\", code_execution_config=False)\n",
    "\n",
    "# Initiate chat from user_proxy side\n",
    "chat_result = user_proxy.initiate_chat(assistant, message=\"What is MLflow?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Coding Agent\n",
    "\n",
    "In this example, we will implement a \"coding agent\" that can execute code. You will see how this code is run alongside your notebook in your current workspace, taking advantage of the performance benefits of Databricks clusters. This is based off the demo [here](https://docs.ag2.ai/latest/docs/user-guide/models/mistralai).\n",
    "\n",
    "First, set up a directory: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "workdir = Path(\"coding\")\n",
    "print(workdir)\n",
    "workdir.mkdir(exist_ok=True)\n",
    "\n",
    "from autogen.coding import LocalCommandLineCodeExecutor\n",
    "\n",
    "code_executor = LocalCommandLineCodeExecutor(work_dir=workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, setup our agents and initiate a coding problem. Notice how the `UserProxyAgent` will take advantage of our `code_executor`; after the code is shown on screen, type Return/Enter in the chatbox to have it execute locally on your cluster via the bot's auto-reply. \n",
    "\n",
    "**Note**: with generative AI coding assistants, you should **always** manually read and review the code before executing it yourself, as LLM results are non-deterministic and may lead to unintended consequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import AssistantAgent, UserProxyAgent\n",
    "\n",
    "# Setting up the agents.\n",
    "user_proxy_agent = UserProxyAgent(\n",
    "    name=\"User\",\n",
    "    code_execution_config={\"executor\": code_executor},\n",
    "    is_termination_msg=lambda msg: \"TERMINATE\" in msg.get(\"content\"),\n",
    ")\n",
    "\n",
    "assistant_agent = AssistantAgent(\n",
    "    name=\"DBRX Assistant\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "chat_result = user_proxy_agent.initiate_chat(\n",
    "    assistant_agent,\n",
    "    message=\"Count how many prime numbers from 1 to 10000.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the python file that was created in our working directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%sh ls coding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%sh head coding/count_primes.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conversable Bots\n",
    "\n",
    "We can also implement the [two-agent chat pattern](https://docs.ag2.ai/docs/user-guide/basic-concepts/conversable-agent) using DBRX to \"talk to itself\" in a teacher/student exchange:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent\n",
    "\n",
    "# Setting up the agents.\n",
    "student_agent = ConversableAgent(\n",
    "    name=\"Student_Agent\",\n",
    "    system_message=\"You are a student willing to learn.\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "teacher_agent = ConversableAgent(\n",
    "    name=\"Teacher_Agent\",\n",
    "    system_message=\"You are a computer science teacher.\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# Initiate chat\n",
    "chat_result = student_agent.initiate_chat(\n",
    "    teacher_agent,\n",
    "    message=\"How does deep learning relate to artificial intelligence?\",\n",
    "    summary_method=\"last_msg\",\n",
    "    max_turns=1,  # Set to higher number to control back and forth\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implement Logging Display\n",
    "\n",
    "It can be useful to display chat logs to the notebook for debugging, and then persist those logs to a Delta table. The following section demonstrates how to extend the default AG2 logging libraries.\n",
    "\n",
    "First, we will implement a Python `class` that extends the capabilities of `autogen.runtime_logging` [docs](https://docs.ag2.ai/latest/docs/use-cases/notebooks/notebooks/agentchat_logging):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatabricksAutoGenLogger:\n",
    "    def __init__(self):\n",
    "        from pyspark.sql import SparkSession\n",
    "\n",
    "        self.spark = SparkSession.builder.getOrCreate()\n",
    "        self.logger_config = {\"dbname\": \"logs.db\"}\n",
    "\n",
    "    def start(self):\n",
    "        import autogen.runtime_logging\n",
    "\n",
    "        self.logging_session_id = autogen.runtime_logging.start(config=self.logger_config)\n",
    "        print(\"Logging session ID: \" + str(self.logging_session_id))\n",
    "\n",
    "    def stop(self):\n",
    "        import autogen.runtime_logging\n",
    "\n",
    "        autogen.runtime_logging.stop()\n",
    "\n",
    "    def _get_log(self, dbname=\"logs.db\", table=\"chat_completions\"):\n",
    "        import sqlite3\n",
    "\n",
    "        con = sqlite3.connect(dbname)\n",
    "        query = f\"SELECT * from {table} WHERE session_id == '{self.logging_session_id}' ORDER BY end_time DESC\"\n",
    "        cursor = con.execute(query)\n",
    "        rows = cursor.fetchall()\n",
    "        column_names = [description[0] for description in cursor.description]\n",
    "        data = [dict(zip(column_names, row)) for row in rows]\n",
    "        con.close()\n",
    "        return data\n",
    "\n",
    "    def display_session(self):\n",
    "        import pandas as pd\n",
    "\n",
    "        return pd.DataFrame(self._get_log())\n",
    "\n",
    "    def persist_results(self, target_delta_table: str, mode=\"append\"):\n",
    "        import pandas as pd\n",
    "\n",
    "        # Convert to Spark DF\n",
    "        sdf = self.spark.createDataFrame(pd.DataFrame(self._get_log()))\n",
    "\n",
    "        try:\n",
    "            sdf.write.format(\"delta\").mode(mode).saveAsTable(target_delta_table)\n",
    "            print(f\"Logs successfully written to table {target_delta_table} in {mode} mode\")\n",
    "        except Exception as e:\n",
    "            print(f\"An error occurred: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use the class above on our simplest example. Note the addition of logging `.start()` and `.stop()`, as well as try/except for error handling. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assistant = autogen.AssistantAgent(name=\"assistant\", llm_config=llm_config)\n",
    "user_proxy = autogen.UserProxyAgent(name=\"user\", code_execution_config=False)\n",
    "\n",
    "# Before initiating chat, start logging:\n",
    "logs = DatabricksAutoGenLogger()\n",
    "logs.start()\n",
    "try:\n",
    "    user_proxy.initiate_chat(assistant, message=\"What is MLflow?\", max_turns=1)\n",
    "except Exception as e:\n",
    "    print(f\"An error occurred: {e}\")\n",
    "logs.stop()\n",
    "# Display logs\n",
    "display(logs.display_session())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this, we have a simple framework to review and persist logs from our chats! Notice that in the `request` field above, we can also see the system prompt for the LLM - this can be useful for prompt engineering as well as debugging.\n",
    "\n",
    "Note that when you deploy this to Databricks Model Serving, model responses are auto-logged using [Lakehouse Monitoring](https://docs.databricks.com/en/lakehouse-monitoring/index.html); but the above approach provides a simple mechanism to log chats from the **client side**.\n",
    "\n",
    "Let's now persist these results to a Delta table in [Unity Catalog](https://docs.databricks.com/en/data-governance/unity-catalog/index.html):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "spark = SparkSession.builder.getOrCreate()  # Not needed in Databricks; session pre-provisioned in notebooks\n",
    "\n",
    "# Use 3-layer namespace: catalog.schema.table. The table will be created if it does not exist.\n",
    "target_delta_table = \"your_catalog.your_schema.autogen_logs\"\n",
    "logs.persist_results(target_delta_table=target_delta_table, mode=\"append\")\n",
    "\n",
    "# Display current rows in table\n",
    "display(spark.table(target_delta_table))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Closing Thoughts\n",
    "This notebook provides a few basic examples of using Autogen with DBRX, and we're excited to see how you can use this framework alongside leading open-source LLMs!\n",
    "\n",
    "### Limitations\n",
    "* Databricks Foundation Model API supports other open-source LLMs (Mixtral, Llama2, etc.), but the above code has not been tested on those.\n",
    "\n",
    "* As of April 2024, DBRX does not yet support tool/function calling abilities. To discuss this capability further, please reach out to your Databricks representative."
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Use Databricks DBRX and Foundation Model APIs to build AG2 applications backed by open-source LLMs.",
   "tags": [
    "integration",
    "code generation",
    "dbrx",
    "databricks",
    "open source",
    "lakehouse",
    "custom model",
    "data intelligence"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.11"
  },
  "notebook_environment": {},
  "nteract": {
   "version": "nteract-front-end@1.0.0"
  },
  "save_output": true,
  "skip_test": "Invalid environment: will only run in Databricks workspace after replacing variables",
  "spark_compute": {
   "compute_id": "/default",
   "session_options": {
    "conf": {},
    "enableDebugMode": false
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}