diff --git a/docs/docs/integrations/tools/databricks.ipynb b/docs/docs/integrations/tools/databricks.ipynb index bb44a716f58..fafb14cd973 100644 --- a/docs/docs/integrations/tools/databricks.ipynb +++ b/docs/docs/integrations/tools/databricks.ipynb @@ -6,7 +6,7 @@ "source": [ "# Databricks Unity Catalog (UC)\n", "\n", - "This notebook shows how to use UC functions as LangChain tools.\n", + "This notebook shows how to use UC functions as LangChain tools, with both LangChain and LangGraph agent APIs.\n", "\n", "See Databricks documentation ([AWS](https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html)|[Azure](https://learn.microsoft.com/en-us/azure/databricks/sql/language-manual/sql-ref-syntax-ddl-create-sql-function)|[GCP](https://docs.gcp.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html)) to learn how to create SQL or Python functions in UC. Do not skip function and parameter comments, which are critical for LLMs to call functions properly.\n", "\n", @@ -34,11 +34,19 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], "source": [ - "%pip install --upgrade --quiet databricks-sdk langchain-community mlflow" + "%pip install --upgrade --quiet databricks-sdk langchain-community langchain-databricks langgraph mlflow" ] }, { @@ -47,7 +55,7 @@ "metadata": {}, "outputs": [], "source": [ - "from langchain_community.chat_models.databricks import ChatDatabricks\n", + "from langchain_databricks import ChatDatabricks\n", "\n", "llm = ChatDatabricks(endpoint=\"databricks-meta-llama-3-70b-instruct\")" ] @@ -58,6 +66,7 @@ "metadata": {}, "outputs": [], "source": [ + "from databricks.sdk import WorkspaceClient\n", "from langchain_community.tools.databricks import UCFunctionToolkit\n", "\n", "tools = (\n", @@ -76,9 +85,16 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, "source": [ - "(Optional) To increase the retry time for getting a function execution response, set environment variable UC_TOOL_CLIENT_EXECUTION_TIMEOUT. Default retry time value is 120s." + "(Optional) To increase the retry time for getting a function execution response, set environment variable UC_TOOL_CLIENT_EXECUTION_TIMEOUT. Default retry time value is 120s.", + + "## LangGraph agent example" ] }, { @@ -92,9 +108,68 @@ "os.environ[\"UC_TOOL_CLIENT_EXECUTION_TIMEOUT\"] = \"200\"" ] }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "source": [ + "## LangGraph agent example" + ] + }, { "cell_type": "code", "execution_count": 4, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'messages': [HumanMessage(content='36939 * 8922.4', additional_kwargs={}, response_metadata={}, id='1a10b10b-8e37-48c7-97a1-cac5006228d5'),\n", + " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_a8f3986f-4b91-40a3-8d6d-39f431dab69b', 'type': 'function', 'function': {'name': 'main__tools__python_exec', 'arguments': '{\"code\": \"print(36939 * 8922.4)\"}'}}]}, response_metadata={'prompt_tokens': 771, 'completion_tokens': 29, 'total_tokens': 800}, id='run-865c3613-20ba-4e80-afc8-fde1cfb26e5a-0', tool_calls=[{'name': 'main__tools__python_exec', 'args': {'code': 'print(36939 * 8922.4)'}, 'id': 'call_a8f3986f-4b91-40a3-8d6d-39f431dab69b', 'type': 'tool_call'}]),\n", + " ToolMessage(content='{\"format\": \"SCALAR\", \"value\": \"329584533.59999996\\\\n\", \"truncated\": false}', name='main__tools__python_exec', id='8b63d4c8-1a3d-46a5-a719-393b2ef36770', tool_call_id='call_a8f3986f-4b91-40a3-8d6d-39f431dab69b'),\n", + " AIMessage(content='The result of the multiplication is:\\n\\n329584533.59999996', additional_kwargs={}, response_metadata={'prompt_tokens': 846, 'completion_tokens': 22, 'total_tokens': 868}, id='run-22772404-611b-46e4-9956-b85e4a385f0f-0')]}" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from langgraph.prebuilt import create_react_agent\n", + "\n", + "agent = create_react_agent(\n", + " llm,\n", + " tools,\n", + " state_modifier=\"You are a helpful assistant. Make sure to use tool for information.\",\n", + ")\n", + "agent.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"36939 * 8922.4\"}]})" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "source": [ + "## LangChain agent example" + ] + }, + { + "cell_type": "code", + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -118,7 +193,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 6, "metadata": {}, "outputs": [ { @@ -132,7 +207,9 @@ "Invoking: `main__tools__python_exec` with `{'code': 'print(36939 * 8922.4)'}`\n", "\n", "\n", - "\u001b[0m\u001b[36;1m\u001b[1;3m{\"format\": \"SCALAR\", \"value\": \"329584533.59999996\\n\", \"truncated\": false}\u001b[0m\u001b[32;1m\u001b[1;3mThe result of the multiplication 36939 * 8922.4 is 329,584,533.60.\u001b[0m\n", + "\u001b[0m\u001b[36;1m\u001b[1;3m{\"format\": \"SCALAR\", \"value\": \"329584533.59999996\\n\", \"truncated\": false}\u001b[0m\u001b[32;1m\u001b[1;3mThe result of the multiplication is:\n", + "\n", + "329584533.59999996\u001b[0m\n", "\n", "\u001b[1m> Finished chain.\u001b[0m\n" ] @@ -141,10 +218,10 @@ "data": { "text/plain": [ "{'input': '36939 * 8922.4',\n", - " 'output': 'The result of the multiplication 36939 * 8922.4 is 329,584,533.60.'}" + " 'output': 'The result of the multiplication is:\\n\\n329584533.59999996'}" ] }, - "execution_count": 5, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -153,18 +230,11 @@ "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)\n", "agent_executor.invoke({\"input\": \"36939 * 8922.4\"})" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { "kernelspec": { - "display_name": "llm", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -178,9 +248,9 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.9" + "version": "3.11.10" } }, "nbformat": 4, - "nbformat_minor": 2 + "nbformat_minor": 4 }