Skip to content

Instantly share code, notes, and snippets.

@Daethyra

Daethyra/nest_runs_within_tools.ipynb

Last active December 8, 2023 04:45
Show Gist options
  • Save Daethyra/91288d366522e43c572875fd41ef9c14 to your computer and use it in GitHub Desktop.
Save Daethyra/91288d366522e43c572875fd41ef9c14 to your computer and use it in GitHub Desktop.
Download ZIP
Group and trace nested calls via `run_manager`
Display the source blob
Display the rendered blob
Raw
nest_runs_within_tools.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/Daethyra/91288d366522e43c572875fd41ef9c14/copy-of-nest_runs_within_tools.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"id": "ea796b03-1259-4742-b9db-25530cc93c2e",
"metadata": {
"id": "ea796b03-1259-4742-b9db-25530cc93c2e"
},
"source": [
"# Tracing Nested Calls within Tools\n",
"[![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/langchain-ai/langsmith-cookbook/blob/main/tracing-examples/nesting-tools/nest_runs_within_tools.ipynb)\n",
"\n",
"Tools are interfaces an agent can use to interact with any function. Often those functions include additional calls to a retriever, LLM, or other resource you'd like to trace. To ensure the nested calls are all grouped within the same trace, just use the `run_manager`!\n",
"\n",
"The following example shows how to pass callbacks through a custom tool to a nested chain. This ensures that the chain's trace is properly included within the agent trace rather than appearing separately.\n",
"\n",
"The resulting trace looks like the following:\n",
"\n",
"![Trace](https://github.com/langchain-ai/langsmith-cookbook/blob/main/tracing-examples/nesting-tools/img/trace.png?raw=1)"
]
},
{
"cell_type": "markdown",
"id": "57c8bbee-34c7-48ff-8338-203182759415",
"metadata": {
"id": "57c8bbee-34c7-48ff-8338-203182759415"
},
"source": [
"## Prerequisites\n",
"\n",
"This example uses OpenAI and LangSmith. Please make sure the dependencies and API keys are configured appropriately before continuing."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fca82ce5-4892-46f1-bda3-21fd14a0d7ae",
"metadata": {
"id": "fca82ce5-4892-46f1-bda3-21fd14a0d7ae"
},
"outputs": [],
"source": [
"%pip install -U langchain openai"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "6f747c6d-cbbf-449f-9470-c40ea825b6e6",
"metadata": {
"id": "6f747c6d-cbbf-449f-9470-c40ea825b6e6"
},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"LANGCHAIN_API_KEY\"] = \"LANGCHAIN_API_KEY\"\n",
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
"os.environ[\"OPENAI_API_KEY\"] = \"OPENAI_API_KEY\""
]
},
{
"cell_type": "markdown",
"id": "3157e9f9-0aa2-4746-b5bd-e313eaa78e6e",
"metadata": {
"id": "3157e9f9-0aa2-4746-b5bd-e313eaa78e6e"
},
"source": [
"## 1. Define tool\n",
"\n",
"We will show how to properly structure your nested calls both when using the decorator and when subclassing the `BaseTool` class.\n",
"\n",
"### Option 1: Using the @tool decorator\n",
"\n",
"The `@tool` decorator is the most concise way to define a LangChain tool. To propagate callbacks\n",
"through the tool function, simply include the \"callbacks\" option in the wrapped function. Below is an example."
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "b7aebf19-2815-4f05-8c6b-eb838ea268b6",
"metadata": {
"id": "b7aebf19-2815-4f05-8c6b-eb838ea268b6"
},
"outputs": [],
"source": [
"import uuid\n",
"\n",
"import requests\n",
"from langchain.callbacks.manager import Callbacks\n",
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.prompts import ChatPromptTemplate\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.tools import tool\n",
"\n",
"\n",
"@tool\n",
"def summarize_tool(url: str, callbacks: Callbacks = None):\n",
" \"\"\"Summarize a website.\"\"\"\n",
" text = requests.get(url).text\n",
" summary_chain = (\n",
" ChatPromptTemplate.from_template(\n",
" \"Summarize the following text:\\n<TEXT {uid}>\\n\" \"{text}\" \"\\n</TEXT {uid}>\"\n",
" ).partial(uid=lambda: uuid.uuid4())\n",
" | ChatOpenAI(model=\"gpt-3.5-turbo-1106\")\n",
" | StrOutputParser()\n",
" ).with_config(run_name=\"Summarize Text\")\n",
" return summary_chain.invoke(\n",
" {\"text\": text},\n",
" {\"callbacks\": callbacks},\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "f83877b9-34a7-46a7-8741-e844f788d8ef",
"metadata": {
"id": "f83877b9-34a7-46a7-8741-e844f788d8ef"
},
"source": [
"### Option 2: Subclass BaseTool\n",
"\n",
"If you directly subclass `BaseTool`, you have more control over the class behavior and state management. You can choose to propagate callbacks by accepting a \"run_manager\" argument in your `_run` method.\n",
"\n",
"Below is an equivalent definition of our tool using inheritance."
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "f872f468-89fa-4d68-a1fc-11fe74263189",
"metadata": {
"id": "f872f468-89fa-4d68-a1fc-11fe74263189"
},
"outputs": [],
"source": [
"import uuid\n",
"from typing import Optional\n",
"\n",
"import aiohttp\n",
"import requests\n",
"from langchain.callbacks.manager import (\n",
" AsyncCallbackManagerForToolRun,\n",
" CallbackManagerForToolRun,\n",
")\n",
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.prompts import ChatPromptTemplate\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.schema.runnable import Runnable\n",
"from langchain.tools import BaseTool\n",
"from pydantic import Field\n",
"\n",
"\n",
"def _default_summary_chain():\n",
" \"\"\"An LLM chain that summarizes the input text\"\"\"\n",
" return (\n",
" ChatPromptTemplate.from_template(\n",
" \"Summarize the following text:\\n<TEXT {uid}>\\n\" \"{text}\" \"\\n</TEXT {uid}>\"\n",
" ).partial(uid=lambda: uuid.uuid4())\n",
" | ChatOpenAI(model=\"gpt-3.5-turbo-1106\")\n",
" | StrOutputParser()\n",
" ).with_config(run_name=\"Summarize Text\")\n",
"\n",
"\n",
"class CustomSummarizer(BaseTool):\n",
" name = \"summary_tool\"\n",
" description = \"summarize a website\"\n",
" summary_chain: Runnable = Field(default_factory=_default_summary_chain)\n",
"\n",
" def _run(\n",
" self,\n",
" url: str,\n",
" run_manager: Optional[CallbackManagerForToolRun] = None,\n",
" ) -> str:\n",
" \"\"\"Use the tool.\"\"\"\n",
" text = requests.get(url).text\n",
" callbacks = run_manager.get_child() if run_manager else None\n",
" return self.summary_chain.invoke(\n",
" {\"text\": text},\n",
" {\"callbacks\": callbacks},\n",
" )\n",
"\n",
" async def _arun(\n",
" self,\n",
" url: str,\n",
" run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n",
" ) -> str:\n",
" \"\"\"Use the tool asynchronously.\"\"\"\n",
" async with aiohttp.ClientSession() as session:\n",
" async with session.get(url) as response:\n",
" text = await response.text()\n",
" callbacks = run_manager.get_child() if run_manager else None\n",
" return await self.summary_chain.ainvoke(\n",
" {\"text\": text},\n",
" {\"callbacks\": callbacks},\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "d0d9e507-14c6-4bce-b8e1-2fdc7553d412",
"metadata": {
"id": "d0d9e507-14c6-4bce-b8e1-2fdc7553d412"
},
"source": [
"## 2. Define agent\n",
"\n",
"We will construct a simple agent using runnables and OpenAI functions, following the [Agents overview](https://python.langchain.com/docs/modules/agents/) in the LangChain documentation. The specifics aren't important to this example."
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "997babf2-41d4-4047-baa3-cac4f13fec84",
"metadata": {
"id": "997babf2-41d4-4047-baa3-cac4f13fec84"
},
"outputs": [],
"source": [
"from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
"\n",
"prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
" (\"system\", \"You are very powerful assistant, but bad at summarizing thingws.\"),\n",
" (\"user\", \"{input}\"),\n",
" MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
" ]\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "7e9398d7-d641-46e0-a81a-fa7fecd76a4b",
"metadata": {
"id": "7e9398d7-d641-46e0-a81a-fa7fecd76a4b"
},
"outputs": [],
"source": [
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.tools.render import format_tool_to_openai_function\n",
"\n",
"# Configure the LLM with access to the appropriate function definitions\n",
"llm = ChatOpenAI(temperature=0)\n",
"tools = [summarize_tool, CustomSummarizer()]\n",
"llm_with_tools = llm.bind(functions=[format_tool_to_openai_function(t) for t in tools])"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "05b6b794-e359-486e-ae79-f26d771095dc",
"metadata": {
"id": "05b6b794-e359-486e-ae79-f26d771095dc"
},
"outputs": [],
"source": [
"from langchain.agents import AgentExecutor\n",
"from langchain.agents.format_scratchpad import format_to_openai_functions\n",
"from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser\n",
"\n",
"agent = (\n",
" {\n",
" \"input\": lambda x: x[\"input\"],\n",
" \"agent_scratchpad\": lambda x: format_to_openai_functions(\n",
" x[\"intermediate_steps\"]\n",
" ),\n",
" \"chat_history\": lambda x: x.get(\"chat_history\") or [],\n",
" }\n",
" | prompt\n",
" | llm_with_tools\n",
" | OpenAIFunctionsAgentOutputParser()\n",
").with_config(run_name=\"Agent\")\n",
"\n",
"agent_executor = AgentExecutor(agent=agent, tools=tools)"
]
},
{
"cell_type": "markdown",
"id": "4eec4034-f866-4900-a008-7ca9a77584bc",
"metadata": {
"id": "4eec4034-f866-4900-a008-7ca9a77584bc"
},
"source": [
"## 3. Invoke\n",
"\n",
"Now its time to call the agent. All callbacks, including the LangSmith tracer, will be passed to the child runnables."
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "1f7fec7e-1773-4af4-aab4-ca83ff8fbf60",
"metadata": {
"id": "1f7fec7e-1773-4af4-aab4-ca83ff8fbf60",
"outputId": "12b9ff21-86e9-44ff-893e-10c000a280f8",
"colab": {
"base_uri": "https://localhost:8080/"
}
},
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"{'input': \"What's this about? https://blog.langchain.dev/langchain-expression-language/\",\n",
" 'output': 'The blog post titled \"LangChain Expression Language\" on the LangChain Blog introduces a new syntax called LangChain Expression Language (LCEL) for creating chains with composition. The post explains the benefits of chaining and provides examples of using LCEL and its features. It also mentions a webinar about LCEL and provides a link to access the \"LangChain Teacher\" application for learning LCEL.'}"
]
},
"metadata": {},
"execution_count": 20
}
],
"source": [
"agent_executor.invoke(\n",
" {\n",
" \"input\": f\"What's this about? https://blog.langchain.dev/langchain-expression-language/\"\n",
" }\n",
")"
]
},
{
"cell_type": "markdown",
"id": "8a5ade81-7320-4934-80b4-a171097741ae",
"metadata": {
"id": "8a5ade81-7320-4934-80b4-a171097741ae"
},
"source": [
"The resulting LangSmith trace should look something like the following:\n",
"\n",
"[![Full run trace](https://github.com/langchain-ai/langsmith-cookbook/blob/main/tracing-examples/nesting-tools/img/full_run.png?raw=1)](https://smith.langchain.com/public/d0212988-a6d5-4d21-9751-89e77bfa58fa/r)"
]
},
{
"cell_type": "markdown",
"id": "d23004db-c044-4b37-b9be-0ca126767ecc",
"metadata": {
"id": "d23004db-c044-4b37-b9be-0ca126767ecc"
},
"source": [
"## Conclusion\n",
"\n",
"In this example, you created a tool in two ways: using the decorator and by subclassing `BaseTool`. You configured the callbacks so the nested call to the \"Summarize Text\" chain is traced correctly.\n",
"\n",
"LangSmith uses LangChain's callbacks system to trace the execution of your application. To trace nested components, the callbacks have to be passed to that component. Any time you see a trace show up on the top level when it ought to be nested, it's likely that somewhere the callbacks weren't correctly passed between components.\n",
"\n",
"This is all made easy when composing functions and other calls as runnables (i.e., [LangChain expression language](https://python.langchain.com/docs/expression_language/))."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.2"
},
"colab": {
"provenance": [],
"include_colab_link": true
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment