Merge branch 'ankush/callbacks-refactor' into vwp/tools_callbacks

Will keep adding chains to this branch, just pushing now for visibility

.gitignore

@@ -144,4 +144,8 @@ wandb/
 /.ruff_cache/
 
 *.pkl
-*.bin
+*.bin
+
+# integration test artifacts
+data_map*
+\[('_type', 'fake'), ('stop', None)]

README.md

@@ -4,6 +4,8 @@
 
 [![lint](https://github.com/hwchase17/langchain/actions/workflows/lint.yml/badge.svg)](https://github.com/hwchase17/langchain/actions/workflows/lint.yml) [![test](https://github.com/hwchase17/langchain/actions/workflows/test.yml/badge.svg)](https://github.com/hwchase17/langchain/actions/workflows/test.yml) [![linkcheck](https://github.com/hwchase17/langchain/actions/workflows/linkcheck.yml/badge.svg)](https://github.com/hwchase17/langchain/actions/workflows/linkcheck.yml) [![Downloads](https://static.pepy.tech/badge/langchain/month)](https://pepy.tech/project/langchain) [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) [![Twitter](https://img.shields.io/twitter/url/https/twitter.com/langchainai.svg?style=social&label=Follow%20%40LangChainAI)](https://twitter.com/langchainai) [![](https://dcbadge.vercel.app/api/server/6adMQxSpJS?compact=true&style=flat)](https://discord.gg/6adMQxSpJS)
 
+Looking for the JS/TS version? Check out [LangChain.js](https://github.com/hwchase17/langchainjs).
+
 **Production Support:** As you move your LangChains into production, we'd love to offer more comprehensive support.
 Please fill out [this form](https://forms.gle/57d8AmXBYp8PP8tZA) and we'll set up a dedicated support Slack channel.
 

docs/deployments.md

@@ -1,14 +1,10 @@
 # Deployments
 
-So you've made a really cool chain - now what? How do you deploy it and make it easily sharable with the world?
+So, you've created a really cool chain - now what? How do you deploy it and make it easily shareable with the world?
 
-This section covers several options for that.
-Note that these are meant as quick deployment options for prototypes and demos, and not for production systems.
-If you are looking for help with deployment of a production system, please contact us directly.
+This section covers several options for that. Note that these options are meant for quick deployment of prototypes and demos, not for production systems. If you need help with the deployment of a production system, please contact us directly.
 
-What follows is a list of template GitHub repositories aimed that are intended to be
-very easy to fork and modify to use your chain.
-This is far from an exhaustive list of options, and we are EXTREMELY open to contributions here.
+What follows is a list of template GitHub repositories designed to be easily forked and modified to use your chain. This list is far from exhaustive, and we are EXTREMELY open to contributions here.
 
 ## [Streamlit](https://github.com/hwchase17/langchain-streamlit-template)
 
@@ -33,6 +29,10 @@ It implements a Question Answering app and contains instructions for deploying t
 
 A minimal example on how to run LangChain on Vercel using Flask.
 
+## [Fly.io](https://github.com/fly-apps/hello-fly-langchain)
+
+A minimal example of how to deploy LangChain to [Fly.io](https://fly.io/) using Flask.
+
 ## [Digitalocean App Platform](https://github.com/homanp/digitalocean-langchain)
 
 A minimal example on how to deploy LangChain to DigitalOcean App Platform.
@@ -43,12 +43,11 @@ A minimal example on how to deploy LangChain to Google Cloud Run.
 
 ## [SteamShip](https://github.com/steamship-core/steamship-langchain/)
 
-This repository contains LangChain adapters for Steamship, enabling LangChain developers to rapidly deploy their apps on Steamship.
-This includes: production ready endpoints, horizontal scaling across dependencies, persistant storage of app state, multi-tenancy support, etc.
+This repository contains LangChain adapters for Steamship, enabling LangChain developers to rapidly deploy their apps on Steamship. This includes: production-ready endpoints, horizontal scaling across dependencies, persistent storage of app state, multi-tenancy support, etc.
 
 ## [Langchain-serve](https://github.com/jina-ai/langchain-serve)
 
-This repository allows users to serve local chains and agents as RESTful, gRPC, or Websocket APIs thanks to [Jina](https://docs.jina.ai/). Deploy your chains & agents with ease and enjoy independent scaling, serverless and autoscaling APIs, as well as a Streamlit playground on Jina AI Cloud.
+This repository allows users to serve local chains and agents as RESTful, gRPC, or WebSocket APIs, thanks to [Jina](https://docs.jina.ai/). Deploy your chains & agents with ease and enjoy independent scaling, serverless and autoscaling APIs, as well as a Streamlit playground on Jina AI Cloud.
 
 ## [BentoML](https://github.com/ssheng/BentoChain)
 
@@ -56,4 +55,4 @@ This repository provides an example of how to deploy a LangChain application wit
 
 ## [Databutton](https://databutton.com/home?new-data-app=true)
 
-These templates serve as examples of how to build, deploy, and share LangChain applications using Databutton. You can create user interfaces with Streamlit, automate tasks by scheduling Python code, and store files and data in the built-in store. Examples include Chatbot interface with conversational memory, Personal search engine, and a starter template for LangChain apps. Deploying and sharing is one click. 
+These templates serve as examples of how to build, deploy, and share LangChain applications using Databutton. You can create user interfaces with Streamlit, automate tasks by scheduling Python code, and store files and data in the built-in store. Examples include a Chatbot interface with conversational memory, a Personal search engine, and a starter template for LangChain apps. Deploying and sharing is just one click away.

docs/ecosystem/aim_tracking.ipynb

@@ -61,7 +61,6 @@
     "from datetime import datetime\n",
     "\n",
     "from langchain.llms import OpenAI\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks import AimCallbackHandler, StdOutCallbackHandler"
    ]
   },
@@ -109,8 +108,8 @@
     "    experiment_name=\"scenario 1: OpenAI LLM\",\n",
     ")\n",
     "\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), aim_callback])\n",
-    "llm = OpenAI(temperature=0, callback_manager=manager, verbose=True)"
+    "callbacks = [StdOutCallbackHandler(), aim_callback]\n",
+    "llm = OpenAI(temperature=0, callbacks=callbacks)"
    ]
   },
   {
@@ -177,7 +176,7 @@
     "Title: {title}\n",
     "Playwright: This is a synopsis for the above play:\"\"\"\n",
     "prompt_template = PromptTemplate(input_variables=[\"title\"], template=template)\n",
-    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callback_manager=manager)\n",
+    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callbacks=callbacks)\n",
     "\n",
     "test_prompts = [\n",
     "    {\"title\": \"documentary about good video games that push the boundary of game design\"},\n",
@@ -249,13 +248,12 @@
    ],
    "source": [
     "# scenario 3 - Agent with Tools\n",
-    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callback_manager=manager)\n",
+    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callbacks=callbacks)\n",
     "agent = initialize_agent(\n",
     "    tools,\n",
     "    llm,\n",
     "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
-    "    callback_manager=manager,\n",
-    "    verbose=True,\n",
+    "    callbacks=callbacks,\n",
     ")\n",
     "agent.run(\n",
     "    \"Who is Leo DiCaprio's girlfriend? What is her current age raised to the 0.43 power?\"\n",

docs/ecosystem/clearml_tracking.ipynb

@@ -79,7 +79,6 @@
    "source": [
     "from datetime import datetime\n",
     "from langchain.callbacks import ClearMLCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.llms import OpenAI\n",
     "\n",
     "# Setup and use the ClearML Callback\n",
@@ -93,9 +92,9 @@
     "    complexity_metrics=True,\n",
     "    stream_logs=True\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), clearml_callback])\n",
+    "callbacks = [StdOutCallbackHandler(), clearml_callback]\n",
     "# Get the OpenAI model ready to go\n",
-    "llm = OpenAI(temperature=0, callback_manager=manager, verbose=True)"
+    "llm = OpenAI(temperature=0, callbacks=callbacks)"
    ]
   },
   {
@@ -523,13 +522,12 @@
     "from langchain.agents import AgentType\n",
     "\n",
     "# SCENARIO 2 - Agent with Tools\n",
-    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callback_manager=manager)\n",
+    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callbacks=callbacks)\n",
     "agent = initialize_agent(\n",
     "    tools,\n",
     "    llm,\n",
     "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
-    "    callback_manager=manager,\n",
-    "    verbose=True,\n",
+    "    callbacks=callbacks,\n",
     ")\n",
     "agent.run(\n",
     "    \"Who is the wife of the person who sang summer of 69?\"\n",

docs/ecosystem/comet_tracking.ipynb

@@ -121,7 +121,6 @@
     "from datetime import datetime\n",
     "\n",
     "from langchain.callbacks import CometCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.llms import OpenAI\n",
     "\n",
     "comet_callback = CometCallbackHandler(\n",
@@ -131,8 +130,8 @@
     "    tags=[\"llm\"],\n",
     "    visualizations=[\"dep\"],\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), comet_callback])\n",
-    "llm = OpenAI(temperature=0.9, callback_manager=manager, verbose=True)\n",
+    "callbacks = [StdOutCallbackHandler(), comet_callback]\n",
+    "llm = OpenAI(temperature=0.9, callbacks=callbacks, verbose=True)\n",
     "\n",
     "llm_result = llm.generate([\"Tell me a joke\", \"Tell me a poem\", \"Tell me a fact\"] * 3)\n",
     "print(\"LLM result\", llm_result)\n",
@@ -153,7 +152,6 @@
    "outputs": [],
    "source": [
     "from langchain.callbacks import CometCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.chains import LLMChain\n",
     "from langchain.llms import OpenAI\n",
     "from langchain.prompts import PromptTemplate\n",
@@ -164,15 +162,14 @@
     "    stream_logs=True,\n",
     "    tags=[\"synopsis-chain\"],\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), comet_callback])\n",
-    "\n",
-    "llm = OpenAI(temperature=0.9, callback_manager=manager, verbose=True)\n",
+    "callbacks = [StdOutCallbackHandler(), comet_callback]\n",
+    "llm = OpenAI(temperature=0.9, callbacks=callbacks)\n",
     "\n",
     "template = \"\"\"You are a playwright. Given the title of play, it is your job to write a synopsis for that title.\n",
     "Title: {title}\n",
     "Playwright: This is a synopsis for the above play:\"\"\"\n",
     "prompt_template = PromptTemplate(input_variables=[\"title\"], template=template)\n",
-    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callback_manager=manager)\n",
+    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callbacks=callbacks)\n",
     "\n",
     "test_prompts = [{\"title\": \"Documentary about Bigfoot in Paris\"}]\n",
     "print(synopsis_chain.apply(test_prompts))\n",
@@ -194,7 +191,6 @@
    "source": [
     "from langchain.agents import initialize_agent, load_tools\n",
     "from langchain.callbacks import CometCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.llms import OpenAI\n",
     "\n",
     "comet_callback = CometCallbackHandler(\n",
@@ -203,15 +199,15 @@
     "    stream_logs=True,\n",
     "    tags=[\"agent\"],\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), comet_callback])\n",
-    "llm = OpenAI(temperature=0.9, callback_manager=manager, verbose=True)\n",
+    "callbacks = [StdOutCallbackHandler(), comet_callback]\n",
+    "llm = OpenAI(temperature=0.9, callbacks=callbacks)\n",
     "\n",
-    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callback_manager=manager)\n",
+    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callbacks=callbacks)\n",
     "agent = initialize_agent(\n",
     "    tools,\n",
     "    llm,\n",
     "    agent=\"zero-shot-react-description\",\n",
-    "    callback_manager=manager,\n",
+    "    callbacks=callbacks,\n",
     "    verbose=True,\n",
     ")\n",
     "agent.run(\n",
@@ -255,7 +251,6 @@
     "from rouge_score import rouge_scorer\n",
     "\n",
     "from langchain.callbacks import CometCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.chains import LLMChain\n",
     "from langchain.llms import OpenAI\n",
     "from langchain.prompts import PromptTemplate\n",
@@ -298,10 +293,10 @@
     "    tags=[\"custom_metrics\"],\n",
     "    custom_metrics=rouge_score.compute_metric,\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), comet_callback])\n",
-    "llm = OpenAI(temperature=0.9, callback_manager=manager, verbose=True)\n",
+    "callbacks = [StdOutCallbackHandler(), comet_callback]\n",
+    "llm = OpenAI(temperature=0.9)\n",
     "\n",
-    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callback_manager=manager)\n",
+    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template)\n",
     "\n",
     "test_prompts = [\n",
     "    {\n",
@@ -323,7 +318,7 @@
     "                 \"\"\"\n",
     "    }\n",
     "]\n",
-    "print(synopsis_chain.apply(test_prompts))\n",
+    "print(synopsis_chain.apply(test_prompts, callbacks=callbacks))\n",
     "comet_callback.flush_tracker(synopsis_chain, finish=True)"
    ]
   }

docs/ecosystem/gpt4all.md

@@ -3,6 +3,7 @@
 This page covers how to use the `GPT4All` wrapper within LangChain. The tutorial is divided into two parts: installation and setup, followed by usage with an example.
 
 ## Installation and Setup
+
 - Install the Python package with `pip install pyllamacpp`
 - Download a [GPT4All model](https://github.com/nomic-ai/pyllamacpp#supported-model) and place it in your desired directory
 
@@ -28,16 +29,16 @@ To stream the model's predictions, add in a CallbackManager.
 
 ```python
 from langchain.llms import GPT4All
-from langchain.callbacks.base import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+
 # There are many CallbackHandlers supported, such as
 # from langchain.callbacks.streamlit import StreamlitCallbackHandler
 
-callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
-model = GPT4All(model="./models/gpt4all-model.bin", n_ctx=512, n_threads=8, callback_handler=callback_handler, verbose=True)
+callbacks = [StreamingStdOutCallbackHandler()]
+model = GPT4All(model="./models/gpt4all-model.bin", n_ctx=512, n_threads=8)
 
 # Generate text. Tokens are streamed through the callback manager.
-model("Once upon a time, ")
+model("Once upon a time, ", callbacks=callbacks)
 ```
 
 ## Model File

docs/ecosystem/lancedb.md

@@ -0,0 +1,23 @@
+# LanceDB
+
+This page covers how to use [LanceDB](https://github.com/lancedb/lancedb) within LangChain.
+It is broken into two parts: installation and setup, and then references to specific LanceDB wrappers.
+
+## Installation and Setup
+
+- Install the Python SDK with `pip install lancedb`
+
+## Wrappers
+
+### VectorStore
+
+There exists a wrapper around LanceDB databases, allowing you to use it as a vectorstore,
+whether for semantic search or example selection.
+
+To import this vectorstore:
+
+```python
+from langchain.vectorstores import LanceDB
+```
+
+For a more detailed walkthrough of the LanceDB wrapper, see [this notebook](../modules/indexes/vectorstores/examples/lancedb.ipynb)

docs/ecosystem/metal.md

@@ -0,0 +1,26 @@
+# Metal
+
+This page covers how to use [Metal](https://getmetal.io) within LangChain.
+
+## What is Metal?
+
+Metal is a  managed retrieval & memory platform built for production. Easily index your data into `Metal` and run semantic search and retrieval on it.
+
+![Metal](../_static/MetalDash.png)
+
+## Quick start
+
+Get started by [creating a Metal account](https://app.getmetal.io/signup).
+
+Then, you can easily take advantage of the `MetalRetriever` class to start retrieving your data for semantic search, prompting context, etc. This class takes a `Metal` instance and a dictionary of parameters to pass to the Metal API.
+
+```python
+from langchain.retrievers import MetalRetriever
+from metal_sdk.metal import Metal
+
+
+metal = Metal("API_KEY", "CLIENT_ID", "INDEX_ID");
+retriever = MetalRetriever(metal, params={"limit": 2})
+
+docs = retriever.get_relevant_documents("search term")
+```

docs/ecosystem/pipelineai.md

@@ -0,0 +1,19 @@
+# PipelineAI
+
+This page covers how to use the PipelineAI ecosystem within LangChain.
+It is broken into two parts: installation and setup, and then references to specific PipelineAI wrappers.
+
+## Installation and Setup
+
+- Install with `pip install pipeline-ai`
+- Get a Pipeline Cloud api key and set it as an environment variable (`PIPELINE_API_KEY`)
+
+## Wrappers
+
+### LLM
+
+There exists a PipelineAI LLM wrapper, which you can access with
+
+```python
+from langchain.llms import PipelineAI
+```

docs/ecosystem/replicate.md

@@ -9,7 +9,7 @@ This page covers how to run models on Replicate within LangChain.
 
 Find a model on the [Replicate explore page](https://replicate.com/explore), and then paste in the model name and version in this format: `owner-name/model-name:version`
 
-For example, for this [flan-t5 model](https://replicate.com/daanelson/flan-t5), click on the API tab. The model name/version would be: `daanelson/flan-t5:04e422a9b85baed86a4f24981d7f9953e20c5fd82f6103b74ebc431588e1cec8`
+For example, for this [dolly model](https://replicate.com/replicate/dolly-v2-12b), click on the API tab. The model name/version would be: `"replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5"`
 
 Only the `model` param is required, but any other model parameters can also be passed in with the format `input={model_param: value, ...}`
 
@@ -24,7 +24,7 @@ Replicate(model="stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6
 From here, we can initialize our model:
 
 ```python
-llm = Replicate(model="daanelson/flan-t5:04e422a9b85baed86a4f24981d7f9953e20c5fd82f6103b74ebc431588e1cec8")
+llm = Replicate(model="replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5")
 ```
 
 And run it:
@@ -40,8 +40,7 @@ llm(prompt)
 We can call any Replicate model (not just LLMs) using this syntax. For example, we can call [Stable Diffusion](https://replicate.com/stability-ai/stable-diffusion):
 
 ```python
-text2image = Replicate(model="stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
-                       input={'image_dimensions'='512x512'}
+text2image = Replicate(model="stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf", input={'image_dimensions':'512x512'})
 
 image_output = text2image("A cat riding a motorcycle by Picasso")
 ```

docs/ecosystem/wandb_tracking.ipynb

@@ -50,7 +50,6 @@
    "source": [
     "from datetime import datetime\n",
     "from langchain.callbacks import WandbCallbackHandler, StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.llms import OpenAI"
    ]
   },
@@ -196,8 +195,8 @@
     "    name=\"llm\",\n",
     "    tags=[\"test\"],\n",
     ")\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), wandb_callback])\n",
-    "llm = OpenAI(temperature=0, callback_manager=manager, verbose=True)"
+    "callbacks = [StdOutCallbackHandler(), wandb_callback]\n",
+    "llm = OpenAI(temperature=0, callbacks=callbacks)"
    ]
   },
   {
@@ -484,7 +483,7 @@
     "Title: {title}\n",
     "Playwright: This is a synopsis for the above play:\"\"\"\n",
     "prompt_template = PromptTemplate(input_variables=[\"title\"], template=template)\n",
-    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callback_manager=manager)\n",
+    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callbacks=callbacks)\n",
     "\n",
     "test_prompts = [\n",
     "    {\n",
@@ -577,16 +576,15 @@
    ],
    "source": [
     "# SCENARIO 3 - Agent with Tools\n",
-    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callback_manager=manager)\n",
+    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm)\n",
     "agent = initialize_agent(\n",
     "    tools,\n",
     "    llm,\n",
     "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
-    "    callback_manager=manager,\n",
-    "    verbose=True,\n",
     ")\n",
     "agent.run(\n",
-    "    \"Who is Leo DiCaprio's girlfriend? What is her current age raised to the 0.43 power?\"\n",
+    "    \"Who is Leo DiCaprio's girlfriend? What is her current age raised to the 0.43 power?\",\n",
+    "    callbacks=callbacks,\n",
     ")\n",
     "wandb_callback.flush_tracker(agent, reset=False, finish=True)"
    ]

docs/modules/agents/agent_executors/examples/async_agent.ipynb

@@ -42,7 +42,6 @@
     "from langchain.agents import AgentType\n",
     "from langchain.llms import OpenAI\n",
     "from langchain.callbacks.stdout import StdOutCallbackHandler\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.tracers import LangChainTracer\n",
     "from aiohttp import ClientSession\n",
     "\n",
@@ -307,14 +306,14 @@
     "    # To make async requests in Tools more efficient, you can pass in your own aiohttp.ClientSession, \n",
     "    # but you must manually close the client session at the end of your program/event loop\n",
     "    aiosession = ClientSession()\n",
+    "    callbacks = [StdOutCallbackHandler()]\n",
     "    for _ in questions:\n",
-    "        manager = CallbackManager([StdOutCallbackHandler()])\n",
-    "        llm = OpenAI(temperature=0, callback_manager=manager)\n",
-    "        async_tools = load_tools([\"llm-math\", \"serpapi\"], llm=llm, aiosession=aiosession, callback_manager=manager)\n",
+    "        llm = OpenAI(temperature=0)\n",
+    "        async_tools = load_tools([\"llm-math\", \"serpapi\"], llm=llm, aiosession=aiosession)\n",
     "        agents.append(\n",
-    "            initialize_agent(async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, callback_manager=manager)\n",
+    "            initialize_agent(async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)\n",
     "        )\n",
-    "    tasks = [async_agent.arun(q) for async_agent, q in zip(agents, questions)]\n",
+    "    tasks = [async_agent.arun(q, callbacks=callbacks) for async_agent, q in zip(agents, questions)]\n",
     "    await asyncio.gather(*tasks)\n",
     "    await aiosession.close()\n",
     "\n",
@@ -376,14 +375,14 @@
     "aiosession = ClientSession()\n",
     "tracer = LangChainTracer()\n",
     "tracer.load_default_session()\n",
-    "manager = CallbackManager([StdOutCallbackHandler(), tracer])\n",
+    "callbacks = [StdOutCallbackHandler(), tracer]\n",
     "\n",
     "# Pass the manager into the llm if you want llm calls traced.\n",
-    "llm = OpenAI(temperature=0, callback_manager=manager)\n",
+    "llm = OpenAI(temperature=0)\n",
     "\n",
     "async_tools = load_tools([\"llm-math\", \"serpapi\"], llm=llm, aiosession=aiosession)\n",
-    "async_agent = initialize_agent(async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, callback_manager=manager)\n",
-    "await async_agent.arun(questions[0])\n",
+    "async_agent = initialize_agent(async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)\n",
+    "await async_agent.arun(questions[0], callbacks=callbacks)\n",
     "await aiosession.close()"
    ]
   }

docs/modules/agents/agents/custom_agent_with_tool_retrieval.ipynb

@@ -373,6 +373,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "tools = get_tools(\"whats the weather?\")\n",
     "tool_names = [tool.name for tool in tools]\n",
     "agent = LLMSingleActionAgent(\n",
     "    llm_chain=llm_chain, \n",

docs/modules/agents/tools/examples/bash.ipynb

@@ -5,57 +5,158 @@
       "id": "8f210ec3",
       "metadata": {},
       "source": [
-        "# Bash\n",
-        "It can often be useful to have an LLM generate bash commands, and then run them. A common use case for this is letting the LLM interact with your local file system. We provide an easy util to execute bash commands."
+        "# Shell Tool\n",
+        "\n",
+        "Giving agents access to the shell is powerful (though risky outside a sandboxed environment).\n",
+        "\n",
+        "The LLM can use it to execute any shell commands. A common use case for this is letting the LLM interact with your local file system."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": 1,
       "id": "f7b3767b",
-      "metadata": {},
+      "metadata": {
+        "tags": []
+      },
       "outputs": [],
       "source": [
-        "from langchain.utilities import BashProcess"
+        "from langchain.tools import ShellTool\n",
+        "\n",
+        "shell_tool = ShellTool()"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": 2,
-      "id": "cf1c92f0",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "bash = BashProcess()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 3,
-      "id": "2fa952fc",
-      "metadata": {},
+      "id": "c92ac832-556b-4f66-baa4-b78f965dfba0",
+      "metadata": {
+        "tags": []
+      },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "bash.ipynb\n",
-            "google_search.ipynb\n",
-            "python.ipynb\n",
-            "requests.ipynb\n",
-            "serpapi.ipynb\n",
+            "Hello World!\n",
+            "\n",
+            "real\t0m0.000s\n",
+            "user\t0m0.000s\n",
+            "sys\t0m0.000s\n",
             "\n"
           ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/Users/wfh/code/lc/lckg/langchain/tools/shell/tool.py:34: UserWarning: The shell tool has no safeguards by default. Use at your own risk.\n",
+            "  warnings.warn(\n"
+          ]
         }
       ],
       "source": [
-        "print(bash.run(\"ls\"))"
+        "print(shell_tool.run({\"commands\": [\"echo 'Hello World!'\", \"time\"]}))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "2fa952fc",
+      "metadata": {},
+      "source": [
+        "### Use with Agents\n",
+        "\n",
+        "As with all tools, these can be given to an agent to accomplish more complex tasks. Let's have the agent fetch some links from a web page."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "id": "851fee9f",
+      "metadata": {
+        "tags": []
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "\n",
+            "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
+            "\u001b[32;1m\u001b[1;3mQuestion: What is the task?\n",
+            "Thought: We need to download the langchain.com webpage and extract all the URLs from it. Then we need to sort the URLs and return them.\n",
+            "Action:\n",
+            "```\n",
+            "{\n",
+            "  \"action\": \"shell\",\n",
+            "  \"action_input\": {\n",
+            "    \"commands\": [\n",
+            "      \"curl -s https://langchain.com | grep -o 'http[s]*://[^\\\" ]*' | sort\"\n",
+            "    ]\n",
+            "  }\n",
+            "}\n",
+            "```\n",
+            "\u001b[0m"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/Users/wfh/code/lc/lckg/langchain/tools/shell/tool.py:34: UserWarning: The shell tool has no safeguards by default. Use at your own risk.\n",
+            "  warnings.warn(\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "Observation: \u001b[36;1m\u001b[1;3mhttps://blog.langchain.dev/\n",
+            "https://discord.gg/6adMQxSpJS\n",
+            "https://docs.langchain.com/docs/\n",
+            "https://github.com/hwchase17/chat-langchain\n",
+            "https://github.com/hwchase17/langchain\n",
+            "https://github.com/hwchase17/langchainjs\n",
+            "https://github.com/sullivan-sean/chat-langchainjs\n",
+            "https://js.langchain.com/docs/\n",
+            "https://python.langchain.com/en/latest/\n",
+            "https://twitter.com/langchainai\n",
+            "\u001b[0m\n",
+            "Thought:\u001b[32;1m\u001b[1;3mThe URLs have been successfully extracted and sorted. We can return the list of URLs as the final answer.\n",
+            "Final Answer: [\"https://blog.langchain.dev/\", \"https://discord.gg/6adMQxSpJS\", \"https://docs.langchain.com/docs/\", \"https://github.com/hwchase17/chat-langchain\", \"https://github.com/hwchase17/langchain\", \"https://github.com/hwchase17/langchainjs\", \"https://github.com/sullivan-sean/chat-langchainjs\", \"https://js.langchain.com/docs/\", \"https://python.langchain.com/en/latest/\", \"https://twitter.com/langchainai\"]\u001b[0m\n",
+            "\n",
+            "\u001b[1m> Finished chain.\u001b[0m\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'[\"https://blog.langchain.dev/\", \"https://discord.gg/6adMQxSpJS\", \"https://docs.langchain.com/docs/\", \"https://github.com/hwchase17/chat-langchain\", \"https://github.com/hwchase17/langchain\", \"https://github.com/hwchase17/langchainjs\", \"https://github.com/sullivan-sean/chat-langchainjs\", \"https://js.langchain.com/docs/\", \"https://python.langchain.com/en/latest/\", \"https://twitter.com/langchainai\"]'"
+            ]
+          },
+          "execution_count": 3,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "from langchain.chat_models import ChatOpenAI\n",
+        "from langchain.agents import initialize_agent\n",
+        "from langchain.agents import AgentType\n",
+        "\n",
+        "llm = ChatOpenAI(temperature=0)\n",
+        "\n",
+        "shell_tool.description = shell_tool.description + f\"args {shell_tool.args}\".replace(\"{\", \"{{\").replace(\"}\", \"}}\")\n",
+        "self_ask_with_search = initialize_agent([shell_tool], llm, agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION, verbose=True)\n",
+        "self_ask_with_search.run(\"Download the langchain.com webpage and grep for all urls. Return only a sorted list of them. Be sure to use double quotes.\")"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "851fee9f",
+      "id": "8d0ea3ac-0890-4e39-9cec-74bd80b4b8b8",
       "metadata": {},
       "outputs": [],
       "source": []
@@ -77,7 +178,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.9"
+      "version": "3.8.16"
     }
   },
   "nbformat": 4,

docs/modules/agents/tools/examples/ddg.ipynb

@@ -27,7 +27,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain.tools import DuckDuckGoSearchTool"
+    "from langchain.tools import DuckDuckGoSearchRun"
    ]
   },
   {
@@ -37,7 +37,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "search = DuckDuckGoSearchTool()"
+    "search = DuckDuckGoSearchRun()"
    ]
   },
   {

docs/modules/agents/tools/examples/filesystem.ipynb

@@ -0,0 +1,190 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# File System Tools\n",
+    "\n",
+    "LangChain provides tools for interacting with a local file system out of the box. This notebook walks through some of them.\n",
+    "\n",
+    "Note: these tools are not recommended for use outside a sandboxed environment! "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First, we'll import the tools."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from langchain.tools.file_management import (\n",
+    "    ReadFileTool,\n",
+    "    CopyFileTool,\n",
+    "    DeleteFileTool,\n",
+    "    MoveFileTool,\n",
+    "    WriteFileTool,\n",
+    "    ListDirectoryTool,\n",
+    ")\n",
+    "from langchain.agents.agent_toolkits import FileManagementToolkit\n",
+    "from tempfile import TemporaryDirectory\n",
+    "\n",
+    "# We'll make a temporary directory to avoid clutter\n",
+    "working_directory = TemporaryDirectory()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## The FileManagementToolkit\n",
+    "\n",
+    "If you want to provide all the file tooling to your agent, it's easy to do so with the toolkit. We'll pass the temporary directory in as a root directory as a workspace for the LLM.\n",
+    "\n",
+    "It's recommended to always pass in a root directory, since without one, it's easy for the LLM to pollute the working directory, and without one, there isn't any validation against\n",
+    "straightforward prompt injection."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[CopyFileTool(name='copy_file', description='Create a copy of a file in a specified location', args_schema=<class 'langchain.tools.file_management.copy.FileCopyInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " DeleteFileTool(name='file_delete', description='Delete a file', args_schema=<class 'langchain.tools.file_management.delete.FileDeleteInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " FileSearchTool(name='file_search', description='Recursively search for files in a subdirectory that match the regex pattern', args_schema=<class 'langchain.tools.file_management.file_search.FileSearchInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " MoveFileTool(name='move_file', description='Move or rename a file from one location to another', args_schema=<class 'langchain.tools.file_management.move.FileMoveInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " ReadFileTool(name='read_file', description='Read file from disk', args_schema=<class 'langchain.tools.file_management.read.ReadFileInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " WriteFileTool(name='write_file', description='Write file to disk', args_schema=<class 'langchain.tools.file_management.write.WriteFileInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " ListDirectoryTool(name='list_directory', description='List files and directories in a specified folder', args_schema=<class 'langchain.tools.file_management.list_dir.DirectoryListingInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug')]"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "toolkit = FileManagementToolkit(root_dir=str(working_directory.name)) # If you don't provide a root_dir, operations will default to the current working directory\n",
+    "toolkit.get_tools()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Selecting File System Tools\n",
+    "\n",
+    "If you only want to select certain tools, you can pass them in as arguments when initializing the toolkit, or you can individually initialize the desired tools."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[ReadFileTool(name='read_file', description='Read file from disk', args_schema=<class 'langchain.tools.file_management.read.ReadFileInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " WriteFileTool(name='write_file', description='Write file to disk', args_schema=<class 'langchain.tools.file_management.write.WriteFileInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug'),\n",
+       " ListDirectoryTool(name='list_directory', description='List files and directories in a specified folder', args_schema=<class 'langchain.tools.file_management.list_dir.DirectoryListingInput'>, return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x1156f4350>, root_dir='/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/tmpxb8c3aug')]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tools = FileManagementToolkit(root_dir=str(working_directory.name), selected_tools=[\"read_file\", \"write_file\", \"list_directory\"]).get_tools()\n",
+    "tools"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'File written successfully to example.txt.'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "read_tool, write_tool, list_tool = tools\n",
+    "write_tool.run({\"file_path\": \"example.txt\", \"text\": \"Hello World!\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'example.txt'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# List files in the working directory\n",
+    "list_tool.run({})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

docs/modules/chains/examples/llm_bash.ipynb

@@ -24,8 +24,8 @@
       "\n",
       "```bash\n",
       "echo \"Hello World\"\n",
-      "```\u001b[0m['```bash', 'echo \"Hello World\"', '```']\n",
-      "\n",
+      "```\u001b[0m\n",
+      "Code: \u001b[33;1m\u001b[1;3m['echo \"Hello World\"']\u001b[0m\n",
       "Answer: \u001b[33;1m\u001b[1;3mHello World\n",
       "\u001b[0m\n",
       "\u001b[1m> Finished chain.\u001b[0m\n"
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -93,7 +93,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -107,8 +107,8 @@
       "\n",
       "```bash\n",
       "printf \"Hello World\\n\"\n",
-      "```\u001b[0m['```bash', 'printf \"Hello World\\\\n\"', '```']\n",
-      "\n",
+      "```\u001b[0m\n",
+      "Code: \u001b[33;1m\u001b[1;3m['printf \"Hello World\\\\n\"']\u001b[0m\n",
       "Answer: \u001b[33;1m\u001b[1;3mHello World\n",
       "\u001b[0m\n",
       "\u001b[1m> Finished chain.\u001b[0m\n"
@@ -120,7 +120,7 @@
        "'Hello World\\n'"
       ]
      },
-     "execution_count": 29,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -132,6 +132,114 @@
     "\n",
     "bash_chain.run(text)"
    ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Persistent Terminal\n",
+    "\n",
+    "By default, the chain will run in a separate subprocess each time it is called. This behavior can be changed by instantiating with a persistent bash process."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMBashChain chain...\u001b[0m\n",
+      "List the current directory then move up a level.\u001b[32;1m\u001b[1;3m\n",
+      "\n",
+      "```bash\n",
+      "ls\n",
+      "cd ..\n",
+      "```\u001b[0m\n",
+      "Code: \u001b[33;1m\u001b[1;3m['ls', 'cd ..']\u001b[0m\n",
+      "Answer: \u001b[33;1m\u001b[1;3mapi.ipynb\t\t\tllm_summarization_checker.ipynb\n",
+      "constitutional_chain.ipynb\tmoderation.ipynb\n",
+      "llm_bash.ipynb\t\t\topenai_openapi.yaml\n",
+      "llm_checker.ipynb\t\topenapi.ipynb\n",
+      "llm_math.ipynb\t\t\tpal.ipynb\n",
+      "llm_requests.ipynb\t\tsqlite.ipynb\u001b[0m\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'api.ipynb\\t\\t\\tllm_summarization_checker.ipynb\\r\\nconstitutional_chain.ipynb\\tmoderation.ipynb\\r\\nllm_bash.ipynb\\t\\t\\topenai_openapi.yaml\\r\\nllm_checker.ipynb\\t\\topenapi.ipynb\\r\\nllm_math.ipynb\\t\\t\\tpal.ipynb\\r\\nllm_requests.ipynb\\t\\tsqlite.ipynb'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain.utilities.bash import BashProcess\n",
+    "\n",
+    "\n",
+    "persistent_process = BashProcess(persistent=True)\n",
+    "bash_chain = LLMBashChain.from_bash_process(llm=llm, bash_process=persistent_process, verbose=True)\n",
+    "\n",
+    "text = \"List the current directory then move up a level.\"\n",
+    "\n",
+    "bash_chain.run(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMBashChain chain...\u001b[0m\n",
+      "List the current directory then move up a level.\u001b[32;1m\u001b[1;3m\n",
+      "\n",
+      "```bash\n",
+      "ls\n",
+      "cd ..\n",
+      "```\u001b[0m\n",
+      "Code: \u001b[33;1m\u001b[1;3m['ls', 'cd ..']\u001b[0m\n",
+      "Answer: \u001b[33;1m\u001b[1;3mexamples\t\tgetting_started.ipynb\tindex_examples\n",
+      "generic\t\t\thow_to_guides.rst\u001b[0m\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'examples\\t\\tgetting_started.ipynb\\tindex_examples\\r\\ngeneric\\t\\t\\thow_to_guides.rst'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Run the same command again and see that the state is maintained between calls\n",
+    "bash_chain.run(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -150,7 +258,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.8.16"
   }
  },
  "nbformat": 4,

docs/modules/chains/generic/custom_chain.ipynb

@@ -0,0 +1,188 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "593f7553-7038-498e-96d4-8255e5ce34f0",
+   "metadata": {},
+   "source": [
+    "# Creating a custom Chain\n",
+    "\n",
+    "To implement your own custom chain you can subclass `BaseChain` and implement the following methods:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "c19c736e-ca74-4726-bb77-0a849bcc2960",
+   "metadata": {
+    "tags": [],
+    "vscode": {
+     "languageId": "python"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "from typing import Any, Dict, List, Optional\n",
+    "\n",
+    "from pydantic import Extra\n",
+    "\n",
+    "from langchain.base_language import BaseLanguageModel\n",
+    "from langchain.callbacks.manager import (\n",
+    "    AsyncCallbackManagerForChainRun,\n",
+    "    CallbackManagerForChainRun,\n",
+    ")\n",
+    "from langchain.chains.base import Chain\n",
+    "from langchain.prompts.base import BasePromptTemplate\n",
+    "\n",
+    "\n",
+    "class MyCustomChain(Chain):\n",
+    "    \"\"\"\n",
+    "    An example of a custom chain.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    prompt: BasePromptTemplate\n",
+    "    \"\"\"Prompt object to use.\"\"\"\n",
+    "    llm: BaseLanguageModel\n",
+    "    output_key: str = \"text\"  #: :meta private:\n",
+    "\n",
+    "    class Config:\n",
+    "        \"\"\"Configuration for this pydantic object.\"\"\"\n",
+    "\n",
+    "        extra = Extra.forbid\n",
+    "        arbitrary_types_allowed = True\n",
+    "\n",
+    "    @property\n",
+    "    def input_keys(self) -> List[str]:\n",
+    "        \"\"\"Will be whatever keys the prompt expects.\n",
+    "\n",
+    "        :meta private:\n",
+    "        \"\"\"\n",
+    "        return self.prompt.input_variables\n",
+    "\n",
+    "    @property\n",
+    "    def output_keys(self) -> List[str]:\n",
+    "        \"\"\"Will always return text key.\n",
+    "\n",
+    "        :meta private:\n",
+    "        \"\"\"\n",
+    "        return [self.output_key]\n",
+    "\n",
+    "    def _call(\n",
+    "        self,\n",
+    "        inputs: Dict[str, Any],\n",
+    "        run_manager: Optional[CallbackManagerForChainRun] = None,\n",
+    "    ) -> Dict[str, str]:\n",
+    "        # Your custom chain logic goes here\n",
+    "        # This is just an example that mimics LLMChain\n",
+    "        prompt_value = self.prompt.format_prompt(**inputs)\n",
+    "        \n",
+    "        # Whenever you call a language model, or another chain, you should pass\n",
+    "        # a callback manager to it. This allows the inner run to be tracked by\n",
+    "        # any callbacks that are registered on the outer run.\n",
+    "        # You can always obtain a callback manager for this by calling\n",
+    "        # `run_manager.get_child()` as shown below.\n",
+    "        response = self.llm.generate_prompt(\n",
+    "            [prompt_value],\n",
+    "            callbacks=run_manager.get_child() if run_manager else None\n",
+    "        )\n",
+    "\n",
+    "        # If you want to log something about this run, you can do so by calling\n",
+    "        # methods on the `run_manager`, as shown below. This will trigger any\n",
+    "        # callbacks that are registered for that event.\n",
+    "        if run_manager:\n",
+    "            run_manager.on_text(\"Log something about this run\")\n",
+    "        \n",
+    "        return {self.output_key: response.generations[0][0].text}\n",
+    "\n",
+    "    async def _acall(\n",
+    "        self,\n",
+    "        inputs: Dict[str, Any],\n",
+    "        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,\n",
+    "    ) -> Dict[str, str]:\n",
+    "        # Your custom chain logic goes here\n",
+    "        # This is just an example that mimics LLMChain\n",
+    "        prompt_value = self.prompt.format_prompt(**inputs)\n",
+    "        \n",
+    "        # Whenever you call a language model, or another chain, you should pass\n",
+    "        # a callback manager to it. This allows the inner run to be tracked by\n",
+    "        # any callbacks that are registered on the outer run.\n",
+    "        # You can always obtain a callback manager for this by calling\n",
+    "        # `run_manager.get_child()` as shown below.\n",
+    "        response = await self.llm.agenerate_prompt(\n",
+    "            [prompt_value],\n",
+    "            callbacks=run_manager.get_child() if run_manager else None\n",
+    "        )\n",
+    "\n",
+    "        # If you want to log something about this run, you can do so by calling\n",
+    "        # methods on the `run_manager`, as shown below. This will trigger any\n",
+    "        # callbacks that are registered for that event.\n",
+    "        if run_manager:\n",
+    "            await run_manager.on_text(\"Log something about this run\")\n",
+    "        \n",
+    "        return {self.output_key: response.generations[0][0].text}\n",
+    "\n",
+    "    @property\n",
+    "    def _chain_type(self) -> str:\n",
+    "        return \"my_custom_chain\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "18361f89",
+   "metadata": {
+    "vscode": {
+     "languageId": "python"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new MyCustomChain chain...\u001b[0m\n",
+      "Log something about this run\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'Why did the callback function feel lonely? Because it was always waiting for someone to call it back!'"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain.callbacks.stdout import StdOutCallbackHandler\n",
+    "from langchain.chat_models.openai import ChatOpenAI\n",
+    "from langchain.prompts.prompt import PromptTemplate\n",
+    "\n",
+    "\n",
+    "chain = MyCustomChain(\n",
+    "    prompt=PromptTemplate.from_template('tell us a joke about {topic}'),\n",
+    "    llm=ChatOpenAI()\n",
+    ")\n",
+    "\n",
+    "chain.run({'topic': 'callbacks'}, callbacks=[StdOutCallbackHandler()])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/modules/chains/index_examples/chat_vector_db.ipynb

@@ -589,7 +589,6 @@
    "outputs": [],
    "source": [
     "from langchain.chains.llm import LLMChain\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
     "from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT, QA_PROMPT\n",
     "from langchain.chains.question_answering import load_qa_chain\n",
@@ -597,7 +596,7 @@
     "# Construct a ConversationalRetrievalChain with a streaming llm for combine docs\n",
     "# and a separate, non-streaming llm for question generation\n",
     "llm = OpenAI(temperature=0)\n",
-    "streaming_llm = OpenAI(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "streaming_llm = OpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "\n",
     "question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)\n",
     "doc_chain = load_qa_chain(streaming_llm, chain_type=\"stuff\", prompt=QA_PROMPT)\n",

docs/modules/indexes/document_loaders/examples/arxiv.ipynb

@@ -0,0 +1,177 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "bda1f3f5",
+   "metadata": {},
+   "source": [
+    "# Arxiv\n",
+    "\n",
+    "[arXiv](https://arxiv.org/) is an open-access archive for 2 million scholarly articles in the fields of physics, mathematics, computer science, quantitative biology, quantitative finance, statistics, electrical engineering and systems science, and economics.\n",
+    "\n",
+    "This notebook shows how to load scientific articles from `Arxiv.org` into a document format that we can use downstream."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b7a1eef-7bf7-4e7d-8bfc-c4e27c9488cb",
+   "metadata": {},
+   "source": [
+    "## Installation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2abd5578-aa3d-46b9-99af-8b262f0b3df8",
+   "metadata": {},
+   "source": [
+    "First, you need to install `arxiv` python package."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b674aaea-ed3a-4541-8414-260a8f67f623",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "!pip install arxiv"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "094b5f13-7e54-4354-9d83-26d6926ecaa0",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "Second, you need to install `PyMuPDF` python package which transform PDF files from the `arxiv.org` site into the text fromat."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7cd91121-2e96-43ba-af50-319853695f86",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "!pip install pymupdf"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "95f05e1c-195e-4e2b-ae8e-8d6637f15be6",
+   "metadata": {},
+   "source": [
+    "## Examples"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e29b954c-1407-4797-ae21-6ba8937156be",
+   "metadata": {},
+   "source": [
+    "`ArxivLoader` has these arguments:\n",
+    "- `query`: free text which used to find documents in the Arxiv\n",
+    "- optional `load_max_docs`: default=100. Use it to limit number of downloaded documents. It takes time to download all 100 documents, so use a small number for experiments.\n",
+    "- optional `load_all_available_meta`: default=False. By defaul only the most important fields downloaded: `Published` (date when document was published/last updated), `Title`, `Authors`, `Summary`. If True, other fields also downloaded."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9bfd5e46",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.document_loaders.base import Document\n",
+    "from langchain.document_loaders import ArxivLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "700e4ef2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "docs = ArxivLoader(query=\"1605.08386\", load_max_docs=2).load()\n",
+    "len(docs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "8977bac0-0042-4f23-9754-247dbd32439b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'Published': '2016-05-26',\n",
+       " 'Title': 'Heat-bath random walks with Markov bases',\n",
+       " 'Authors': 'Caprice Stanley, Tobias Windisch',\n",
+       " 'Summary': 'Graphs on lattice points are studied whose edges come from a finite set of\\nallowed moves of arbitrary length. We show that the diameter of these graphs on\\nfibers of a fixed integer matrix can be bounded from above by a constant. We\\nthen study the mixing behaviour of heat-bath random walks on these graphs. We\\nalso state explicit conditions on the set of moves so that the heat-bath random\\nwalk, a generalization of the Glauber dynamics, is an expander in fixed\\ndimension.'}"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "doc[0].metadata  # meta-information of the Document"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "46969806-45a9-4c4d-a61b-cfb9658fc9de",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'arXiv:1605.08386v1  [math.CO]  26 May 2016\\nHEAT-BATH RANDOM WALKS WITH MARKOV BASES\\nCAPRICE STANLEY AND TOBIAS WINDISCH\\nAbstract. Graphs on lattice points are studied whose edges come from a finite set of\\nallowed moves of arbitrary length. We show that the diameter of these graphs on fibers of a\\nfixed integer matrix can be bounded from above by a constant. We then study the mixing\\nbehaviour of heat-b'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "doc[0].page_content[:400]  # all pages of the Document content\n"
+   ]
+  }
+ ],
+ "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.10.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/modules/indexes/document_loaders/examples/googledrive.ipynb

@@ -16,7 +16,7 @@
     "1. `pip install --upgrade google-api-python-client google-auth-httplib2 google-auth-oauthlib`\n",
     "\n",
     "## 🧑 Instructions for ingesting your Google Docs data\n",
-    "By default, the `GoogleDriveLoader` expects the `credentials.json` file to be `~/.credentials/credentials.json`, but this is configurable using the `credentials_file` keyword argument. Same thing with `token.json`. Note that `token.json` will be created automatically the first time you use the loader.\n",
+    "By default, the `GoogleDriveLoader` expects the `credentials.json` file to be `~/.credentials/credentials.json`, but this is configurable using the `credentials_path` keyword argument. Same thing with `token.json` - `token_path`. Note that `token.json` will be created automatically the first time you use the loader.\n",
     "\n",
     "`GoogleDriveLoader` can load from a list of Google Docs document ids or a folder id. You can obtain your folder and document id from the URL:\n",
     "* Folder: https://drive.google.com/drive/u/0/folders/1yucgL9WGgWZdM1TOuKkeghlPizuzMYb5 -> folder id is `\"1yucgL9WGgWZdM1TOuKkeghlPizuzMYb5\"`\n",

docs/modules/indexes/retrievers/examples/self_query_retriever.ipynb

@@ -0,0 +1,330 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "13afcae7",
+   "metadata": {},
+   "source": [
+    "# Self-querying retriever\n",
+    "In the notebook we'll demo the `SelfQueryRetriever`, which, as the name suggests, has the ability to query itself. Specifically, given any natural language query, the retriever uses a query-constructing LLM chain to write a structured query and then applies that structured query to it's underlying VectorStore. This allows the retriever to not only use the user-input query for semantic similarity comparison with the contents of stored documented, but to also extract filters from the user query on the metadata of stored documents and to execute those filter."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "68e75fb9",
+   "metadata": {},
+   "source": [
+    "## Creating a Pinecone index\n",
+    "First we'll want to create a Pinecone VectorStore and seed it with some data. We've created a small demo set of documents that contain summaries of movies.\n",
+    "\n",
+    "NOTE: The self-query retriever currently only has built-in support for Pinecone VectorStore.\n",
+    "\n",
+    "NOTE: The self-query retriever requires you to have `lark` installed (`pip install lark`)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "63a8af5b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install lark"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3eb9c9a4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/harrisonchase/.pyenv/versions/3.9.1/envs/langchain/lib/python3.9/site-packages/pinecone/index.py:4: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
+      "  from tqdm.autonotebook import tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "\n",
+    "import pinecone\n",
+    "\n",
+    "\n",
+    "pinecone.init(api_key=os.environ[\"PINECONE_API_KEY\"], environment=os.environ[\"PINECONE_ENV\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "cb4a5787",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.schema import Document\n",
+    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
+    "from langchain.vectorstores import Pinecone\n",
+    "\n",
+    "embeddings = OpenAIEmbeddings()\n",
+    "# create new index\n",
+    "pinecone.create_index(\"langchain-self-retriever-demo\", dimension=1536)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "bcbe04d9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "docs = [\n",
+    "    Document(page_content=\"A bunch of scientists bring back dinosaurs and mayhem breaks loose\", metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": [\"action\", \"science fiction\"]}),\n",
+    "    Document(page_content=\"Leo DiCaprio gets lost in a dream within a dream within a dream within a ...\", metadata={\"year\": 2010, \"director\": \"Christopher Nolan\", \"rating\": 8.2}),\n",
+    "    Document(page_content=\"A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea\", metadata={\"year\": 2006, \"director\": \"Satoshi Kon\", \"rating\": 8.6}),\n",
+    "    Document(page_content=\"A bunch of normal-sized women are supremely wholesome and some men pine after them\", metadata={\"year\": 2019, \"director\": \"Greta Gerwig\", \"rating\": 8.3}),\n",
+    "    Document(page_content=\"Toys come alive and have a blast doing so\", metadata={\"year\": 1995, \"genre\": \"animated\"}),\n",
+    "    Document(page_content=\"Three men walk into the Zone, three men walk out of the Zone\", metadata={\"year\": 1979, \"rating\": 9.9, \"director\": \"Andrei Tarkovsky\", \"genre\": [\"science fiction\", \"thriller\"], \"rating\": 9.9})\n",
+    "]\n",
+    "vectorstore = Pinecone.from_documents(\n",
+    "    docs, embeddings, index_name=\"langchain-self-retriever-demo\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5ecaab6d",
+   "metadata": {},
+   "source": [
+    "# Creating our self-querying retriever\n",
+    "Now we can instantiate our retriever. To do this we'll need to provide some information upfront about the metadata fields that our documents support and a short description of the document contents."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "86e34dbf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.llms import OpenAI\n",
+    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
+    "from langchain.chains.query_constructor.base import AttributeInfo\n",
+    "\n",
+    "metadata_field_info=[\n",
+    "    AttributeInfo(\n",
+    "        name=\"genre\",\n",
+    "        description=\"The genre of the movie\", \n",
+    "        type=\"string or list[string]\", \n",
+    "    ),\n",
+    "    AttributeInfo(\n",
+    "        name=\"year\",\n",
+    "        description=\"The year the movie was released\", \n",
+    "        type=\"integer\", \n",
+    "    ),\n",
+    "    AttributeInfo(\n",
+    "        name=\"director\",\n",
+    "        description=\"The name of the movie director\", \n",
+    "        type=\"string\", \n",
+    "    ),\n",
+    "    AttributeInfo(\n",
+    "        name=\"rating\",\n",
+    "        description=\"A 1-10 rating for the movie\",\n",
+    "        type=\"float\"\n",
+    "    ),\n",
+    "]\n",
+    "document_content_description = \"Brief summary of a movie\"\n",
+    "llm = OpenAI(temperature=0)\n",
+    "retriever = SelfQueryRetriever.from_llm(llm, vectorstore, document_content_description, metadata_field_info, verbose=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea9df8d4",
+   "metadata": {},
+   "source": [
+    "# Testing it out\n",
+    "And now we can try actually using our retriever!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "38a126e9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query='dinosaur' filter=None\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Document(page_content='A bunch of scientists bring back dinosaurs and mayhem breaks loose', metadata={'genre': ['action', 'science fiction'], 'rating': 7.7, 'year': 1993.0}),\n",
+       " Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995.0}),\n",
+       " Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006.0}),\n",
+       " Document(page_content='Leo DiCaprio gets lost in a dream within a dream within a dream within a ...', metadata={'director': 'Christopher Nolan', 'rating': 8.2, 'year': 2010.0})]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# This example only specifies a relevant query\n",
+    "retriever.get_relevant_documents(\"What are some movies about dinosaurs\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "fc3f1e6e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query=' ' filter=Comparison(comparator=<Comparator.GT: 'gt'>, attribute='rating', value=8.5)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006.0}),\n",
+       " Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': ['science fiction', 'thriller'], 'rating': 9.9, 'year': 1979.0})]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# This example only specifies a filter\n",
+    "retriever.get_relevant_documents(\"I want to watch a movie rated higher than 8.5\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b19d4da0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query='women' filter=Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='director', value='Greta Gerwig')\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Document(page_content='A bunch of normal-sized women are supremely wholesome and some men pine after them', metadata={'director': 'Greta Gerwig', 'rating': 8.3, 'year': 2019.0})]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# This example specifies a query and a filter\n",
+    "retriever.get_relevant_documents(\"Has Greta Gerwig directed any movies about women\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "f900e40e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query=' ' filter=Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='genre', value='science fiction'), Comparison(comparator=<Comparator.GT: 'gt'>, attribute='rating', value=8.5)])\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': ['science fiction', 'thriller'], 'rating': 9.9, 'year': 1979.0})]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# This example specifies a composite filter\n",
+    "retriever.get_relevant_documents(\"What's a highly rated (above 8.5) science fiction film?\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "12a51522",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query='toys' filter=Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.GT: 'gt'>, attribute='year', value=1990.0), Comparison(comparator=<Comparator.LT: 'lt'>, attribute='year', value=2005.0), Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='genre', value='animated')])\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995.0})]"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# This example specifies a query and composite filter\n",
+    "retriever.get_relevant_documents(\"What's a movie after 1990 but before 2005 that's all about toys, and preferably is animated\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69bbd809",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/modules/indexes/vectorstores/examples/lanecdb.ipynb

@@ -0,0 +1,179 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "683953b3",
+   "metadata": {},
+   "source": [
+    "# LanceDB\n",
+    "\n",
+    "This notebook shows how to use functionality related to the LanceDB vector database based on the Lance data format."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "bfcf346a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#!pip install lancedb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "aac9563e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.embeddings import OpenAIEmbeddings\n",
+    "from langchain.vectorstores import LanceDB"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "a3c3999a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.document_loaders import TextLoader\n",
+    "from langchain.text_splitter import CharacterTextSplitter\n",
+    "loader = TextLoader('../../../state_of_the_union.txt')\n",
+    "documents = loader.load()\n",
+    "\n",
+    "documents = CharacterTextSplitter().split_documents(documents)\n",
+    "\n",
+    "embeddings = OpenAIEmbeddings()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "6e104aee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import lancedb\n",
+    "\n",
+    "db = lancedb.connect('/tmp/lancedb')\n",
+    "table = db.create_table(\"my_table\", data=[\n",
+    "    {\"vector\": embeddings.embed_query(\"Hello World\"), \"text\": \"Hello World\", \"id\": \"1\"}\n",
+    "], mode=\"overwrite\")\n",
+    "\n",
+    "docsearch = LanceDB.from_documents(documents, embeddings, connection=table)\n",
+    "\n",
+    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
+    "docs = docsearch.similarity_search(query)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "9c608226",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "They were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \n",
+      "\n",
+      "Officer Mora was 27 years old. \n",
+      "\n",
+      "Officer Rivera was 22. \n",
+      "\n",
+      "Both Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \n",
+      "\n",
+      "I spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n",
+      "\n",
+      "I’ve worked on these issues a long time. \n",
+      "\n",
+      "I know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. \n",
+      "\n",
+      "So let’s not abandon our streets. Or choose between safety and equal justice. \n",
+      "\n",
+      "Let’s come together to protect our communities, restore trust, and hold law enforcement accountable. \n",
+      "\n",
+      "That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n",
+      "\n",
+      "That’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.  \n",
+      "\n",
+      "We should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. \n",
+      "\n",
+      "I ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe.  \n",
+      "\n",
+      "And I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. \n",
+      "\n",
+      "And I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? \n",
+      "\n",
+      "Ban assault weapons and high-capacity magazines. \n",
+      "\n",
+      "Repeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. \n",
+      "\n",
+      "These laws don’t infringe on the Second Amendment. They save lives. \n",
+      "\n",
+      "The most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. \n",
+      "\n",
+      "In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n",
+      "\n",
+      "We cannot let this happen. \n",
+      "\n",
+      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
+      "\n",
+      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
+      "\n",
+      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
+      "\n",
+      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. \n",
+      "\n",
+      "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n",
+      "\n",
+      "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n",
+      "\n",
+      "We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  \n",
+      "\n",
+      "We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  \n",
+      "\n",
+      "We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster.\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(docs[0].page_content)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a359ed74",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/modules/indexes/vectorstores/examples/weaviate.ipynb

@@ -75,7 +75,8 @@
     "            \"vectorizer\": \"text2vec-openai\",\n",
     "              \"moduleConfig\": {\n",
     "                \"text2vec-openai\": {\n",
-    "                  \"model\": \"babbage\",\n",
+    "                  \"model\": \"ada\",\n",
+    "                  \"modelVersion\": \"002\",\n",
     "                  \"type\": \"text\"\n",
     "                }\n",
     "              },\n",

docs/modules/models/chat/examples/streaming.ipynb

@@ -80,9 +80,8 @@
     }
    ],
    "source": [
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
-    "chat = ChatOpenAI(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "chat = ChatOpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "resp = chat([HumanMessage(content=\"Write me a song about sparkling water.\")])"
    ]
   },

docs/modules/models/chat/getting_started.ipynb

@@ -373,9 +373,8 @@
     }
    ],
    "source": [
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
-    "chat = ChatOpenAI(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "chat = ChatOpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "resp = chat([HumanMessage(content=\"Write me a song about sparkling water.\")])\n"
    ]
   },

docs/modules/models/llms/examples/llm_caching.ipynb

@@ -785,7 +785,9 @@
    "id": "9df0dab8",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "!rm .langchain.db sqlite.db"
+   ]
   }
  ],
  "metadata": {

docs/modules/models/llms/examples/streaming_llm.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "4ac0ff54-540a-4f2b-8d9a-b590fec7fe07",
    "metadata": {
     "tags": []
@@ -21,14 +21,13 @@
    "source": [
     "from langchain.llms import OpenAI, Anthropic\n",
     "from langchain.chat_models import ChatOpenAI\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
     "from langchain.schema import HumanMessage"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "77f60a4b-f786-41f2-972e-e5bb8a48dcd5",
    "metadata": {
     "tags": []
@@ -79,7 +78,7 @@
     }
    ],
    "source": [
-    "llm = OpenAI(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "llm = OpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "resp = llm(\"Write me a song about sparkling water.\")"
    ]
   },
@@ -95,7 +94,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "a35373f1-9ee6-4753-a343-5aee749b8527",
    "metadata": {
     "tags": []
@@ -136,7 +135,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "22665f16-e05b-473c-a4bd-ad75744ea024",
    "metadata": {
     "tags": []
@@ -191,7 +190,7 @@
     }
    ],
    "source": [
-    "chat = ChatOpenAI(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "chat = ChatOpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "resp = chat([HumanMessage(content=\"Write me a song about sparkling water.\")])"
    ]
   },
@@ -205,7 +204,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "eadae4ba-9f21-4ec8-845d-dd43b0edc2dc",
    "metadata": {
     "tags": []
@@ -245,7 +244,7 @@
     }
    ],
    "source": [
-    "llm = Anthropic(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "llm = Anthropic(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)\n",
     "llm(\"Write me a song about sparkling water.\")"
    ]
   }

docs/modules/models/llms/integrations/cerebriumai_example.ipynb

@@ -6,7 +6,7 @@
    "source": [
     "# CerebriumAI\n",
     "\n",
-    "`Cerebrium` is an AWS Sagemaker alternative. It also provides API access to [several LLM models](https://docs.cerebrium.ai/cerebrium/prebuilt-models/deploymen).\n",
+    "`Cerebrium` is an AWS Sagemaker alternative. It also provides API access to [several LLM models](https://docs.cerebrium.ai/cerebrium/prebuilt-models/deployment).\n",
     "\n",
     "This notebook goes over how to use Langchain with [CerebriumAI](https://docs.cerebrium.ai/introduction)."
    ]

docs/modules/models/llms/integrations/gpt4all.ipynb

@@ -40,7 +40,6 @@
    "source": [
     "from langchain import PromptTemplate, LLMChain\n",
     "from langchain.llms import GPT4All\n",
-    "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler"
    ]
   },
@@ -124,9 +123,9 @@
    "outputs": [],
    "source": [
     "# Callbacks support token-wise streaming\n",
-    "callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])\n",
+    "callbacks = [StreamingStdOutCallbackHandler()]\n",
     "# Verbose is required to pass to the callback manager\n",
-    "llm = GPT4All(model=local_path, callback_manager=callback_manager, verbose=True)"
+    "llm = GPT4All(model=local_path, callbacks=callbacks, verbose=True)"
    ]
   },
   {

docs/modules/models/llms/integrations/huggingface_pipelines.ipynb

@@ -11,7 +11,7 @@
     "\n",
     "The [Hugging Face Model Hub](https://huggingface.co/models) hosts over 120k models, 20k datasets, and 50k demo apps (Spaces), all open source and publicly available, in an online platform where people can easily collaborate and build ML together.\n",
     "\n",
-    "These can be called from LangChain either through this local pipeline wrapper or by calling their hosted inference endpoints through the HuggingFaceHub class. For more information on the hosted pipelines, see the [HugigngFaceHub](huggingface_hub.ipynb) notebook."
+    "These can be called from LangChain either through this local pipeline wrapper or by calling their hosted inference endpoints through the HuggingFaceHub class. For more information on the hosted pipelines, see the [HuggingFaceHub](huggingface_hub.ipynb) notebook."
    ]
   },
   {

docs/modules/models/llms/integrations/pipelineai_example.ipynb

@@ -0,0 +1,171 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PipelineAI\n",
+    "\n",
+    "PipelineAI allows you to run your ML models at scale in the cloud. It also provides API access to [several LLM models](https://pipeline.ai).\n",
+    "\n",
+    "This notebook goes over how to use Langchain with [PipelineAI](https://docs.pipeline.ai/docs)."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install pipeline-ai\n",
+    "The `pipeline-ai` library is required to use the `PipelineAI` API, AKA `Pipeline Cloud`. Install `pipeline-ai` using `pip install pipeline-ai`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install the package\n",
+    "!pip install pipeline-ai"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from langchain.llms import PipelineAI\n",
+    "from langchain import PromptTemplate, LLMChain"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set the Environment API Key\n",
+    "Make sure to get your API key from PipelineAI. Check out the [cloud quickstart guide](https://docs.pipeline.ai/docs/cloud-quickstart). You'll be given a 30 day free trial with 10 hours of serverless GPU compute to test different models."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"PIPELINE_API_KEY\"] = \"YOUR_API_KEY_HERE\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create the PipelineAI instance\n",
+    "When instantiating PipelineAI, you need to specify the id or tag of the pipeline you want to use, e.g. `pipeline_key = \"public/gpt-j:base\"`. You then have the option of passing additional pipeline-specific keyword arguments:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = PipelineAI(pipeline_key=\"YOUR_PIPELINE_KEY\", pipeline_kwargs={...})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create a Prompt Template\n",
+    "We will create a prompt template for Question and Answer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "template = \"\"\"Question: {question}\n",
+    "\n",
+    "Answer: Let's think step by step.\"\"\"\n",
+    "\n",
+    "prompt = PromptTemplate(template=template, input_variables=[\"question\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initiate the LLMChain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm_chain = LLMChain(prompt=prompt, llm=llm)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run the LLMChain\n",
+    "Provide a question and run the LLMChain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "question = \"What NFL team won the Super Bowl in the year Justin Beiber was born?\"\n",
+    "\n",
+    "llm_chain.run(question)"
+   ]
+  }
+ ],
+ "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.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "a0a0263b650d907a3bfe41c0f8d6a63a071b884df3cfdc1579f00cdc1aed6b03"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

docs/modules/models/llms/integrations/replicate.ipynb

@@ -44,7 +44,7 @@
    },
    "outputs": [
     {
-     "name": "stdin",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
       " ········\n"
@@ -85,6 +85,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -92,7 +93,7 @@
     "\n",
     "Find a model on the [replicate explore page](https://replicate.com/explore), and then paste in the model name and version in this format: model_name/version\n",
     "\n",
-    "For example, for this [flan-t5 model]( https://replicate.com/daanelson/flan-t5), click on the API tab. The model name/version would be: `daanelson/flan-t5:04e422a9b85baed86a4f24981d7f9953e20c5fd82f6103b74ebc431588e1cec8`\n",
+    "For example, for this [dolly model](https://replicate.com/replicate/dolly-v2-12b), click on the API tab. The model name/version would be: `replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5`\n",
     "\n",
     "Only the `model` param is required, but we can add other model params when initializing.\n",
     "\n",
@@ -113,7 +114,7 @@
    },
    "outputs": [],
    "source": [
-    "llm = Replicate(model=\"daanelson/flan-t5:04e422a9b85baed86a4f24981d7f9953e20c5fd82f6103b74ebc431588e1cec8\")"
+    "llm = Replicate(model=\"replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5\")"
    ]
   },
   {
@@ -243,7 +244,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "llm = Replicate(model=\"daanelson/flan-t5:04e422a9b85baed86a4f24981d7f9953e20c5fd82f6103b74ebc431588e1cec8\")\n",
+    "dolly_llm = Replicate(model=\"replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5\")\n",
     "text2image = Replicate(model=\"stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf\")"
    ]
   },
@@ -265,7 +266,7 @@
     "    template=\"What is a good name for a company that makes {product}?\",\n",
     ")\n",
     "\n",
-    "chain = LLMChain(llm=llm, prompt=prompt)"
+    "chain = LLMChain(llm=dolly_llm, prompt=prompt)"
    ]
   },
   {
@@ -285,7 +286,7 @@
     "    input_variables=[\"company_name\"],\n",
     "    template=\"Write a description of a logo for this company: {company_name}\",\n",
     ")\n",
-    "chain_two = LLMChain(llm=llm, prompt=second_prompt)"
+    "chain_two = LLMChain(llm=dolly_llm, prompt=second_prompt)"
    ]
   },
   {

docs/modules/models/text_embedding/examples/sentence_transformers.ipynb

@@ -8,12 +8,14 @@
    "source": [
     "# Sentence Transformers Embeddings\n",
     "\n",
-    "Let's generate embeddings using the [SentenceTransformers](https://www.sbert.net/) integration. SentenceTransformers is a python package that can generate text and image embeddings, originating from [Sentence-BERT](https://arxiv.org/abs/1908.10084)"
+    "[SentenceTransformers](https://www.sbert.net/) embeddings are called using the `HuggingFaceEmbeddings` integration. We have also added an alias for `SentenceTransformerEmbeddings` for users who are more familiar with directly using that package.\n",
+    "\n",
+    "SentenceTransformers is a python package that can generate text and image embeddings, originating from [Sentence-BERT](https://arxiv.org/abs/1908.10084)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 1,
    "id": "06c9f47d",
    "metadata": {},
    "outputs": [
@@ -21,10 +23,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
-      "To disable this warning, you can either:\n",
-      "\t- Avoid using `tokenizers` before the fork if possible\n",
-      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.1\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
      ]
     }
    ],
@@ -34,27 +35,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 2,
    "id": "861521a9",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain.embeddings import SentenceTransformerEmbeddings "
+    "from langchain.embeddings import HuggingFaceEmbeddings, SentenceTransformerEmbeddings "
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "ff9be586",
    "metadata": {},
    "outputs": [],
    "source": [
-    "embeddings = SentenceTransformerEmbeddings(model=\"all-MiniLM-L6-v2\")"
+    "embeddings = HuggingFaceEmbeddings(model_name=\"all-MiniLM-L6-v2\")\n",
+    "# Equivalent to SentenceTransformerEmbeddings(model_name=\"all-MiniLM-L6-v2\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 4,
    "id": "d0a98ae9",
    "metadata": {},
    "outputs": [],
@@ -64,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 5,
    "id": "5d6c682b",
    "metadata": {},
    "outputs": [],
@@ -74,7 +76,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 6,
    "id": "bb5e74c0",
    "metadata": {},
    "outputs": [],
@@ -107,7 +109,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.8.16"
   },
   "vscode": {
    "interpreter": {

docs/modules/prompts/prompt_templates/examples/connecting_to_a_feature_store.ipynb

@@ -0,0 +1,474 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a792b119",
+   "metadata": {},
+   "source": [
+    "# Connecting to a Feature Store\n",
+    "\n",
+    "Feature stores are a concept from traditional machine learning that make sure data fed into models is up-to-date and relevant. For more on this, see [here](https://www.tecton.ai/blog/what-is-a-feature-store/).\n",
+    "\n",
+    "This concept is extremely relevant when considering putting LLM applications in production. In order to personalize LLM applications, you may want to combine LLMs with up-to-date information about particular users. Feature stores can be a great way to keep that data fresh, and LangChain provides an easy way to combine that data with LLMs.\n",
+    "\n",
+    "In this notebook we will show how to connect prompt templates to feature stores. The basic idea is to call a feature store from inside a prompt template to retrieve values that are then formatted into the prompt."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ad0b5edf",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Feast\n",
+    "\n",
+    "To start, we will use the popular open source feature store framework [Feast](https://github.com/feast-dev/feast).\n",
+    "\n",
+    "This assumes you have already run the steps in the README around getting started. We will build of off that example in getting started, and create and LLMChain to write a note to a specific driver regarding their up-to-date statistics."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7f02f6f3",
+   "metadata": {},
+   "source": [
+    "### Load Feast Store\n",
+    "\n",
+    "Again, this should be set up according to the instructions in the Feast README"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "fd1a452a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from feast import FeatureStore\n",
+    "\n",
+    "# You may need to update the path depending on where you stored it\n",
+    "feast_repo_path = \"../../../../../my_feature_repo/feature_repo/\"\n",
+    "store = FeatureStore(repo_path=feast_repo_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cfe8aae5",
+   "metadata": {},
+   "source": [
+    "### Prompts\n",
+    "\n",
+    "Here we will set up a custom FeastPromptTemplate. This prompt template will take in a driver id, look up their stats, and format those stats into a prompt.\n",
+    "\n",
+    "Note that the input to this prompt template is just `driver_id`, since that is the only user defined piece (all other variables are looked up inside the prompt template)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "5e9cee04",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.prompts import PromptTemplate, StringPromptTemplate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "594a3cf3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "template = \"\"\"Given the driver's up to date stats, write them note relaying those stats to them.\n",
+    "If they have a conversation rate above .5, give them a compliment. Otherwise, make a silly joke about chickens at the end to make them feel better\n",
+    "\n",
+    "Here are the drivers stats:\n",
+    "Conversation rate: {conv_rate}\n",
+    "Acceptance rate: {acc_rate}\n",
+    "Average Daily Trips: {avg_daily_trips}\n",
+    "\n",
+    "Your response:\"\"\"\n",
+    "prompt = PromptTemplate.from_template(template)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "8464c731",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FeastPromptTemplate(StringPromptTemplate):\n",
+    "    \n",
+    "    def format(self, **kwargs) -> str:\n",
+    "        driver_id = kwargs.pop(\"driver_id\")\n",
+    "        feature_vector = store.get_online_features(\n",
+    "            features=[\n",
+    "                'driver_hourly_stats:conv_rate',\n",
+    "                'driver_hourly_stats:acc_rate',\n",
+    "                'driver_hourly_stats:avg_daily_trips'\n",
+    "            ],\n",
+    "            entity_rows=[{\"driver_id\": 1001}]\n",
+    "        ).to_dict()\n",
+    "        kwargs[\"conv_rate\"] = feature_vector[\"conv_rate\"][0]\n",
+    "        kwargs[\"acc_rate\"] = feature_vector[\"acc_rate\"][0]\n",
+    "        kwargs[\"avg_daily_trips\"] = feature_vector[\"avg_daily_trips\"][0]\n",
+    "        return prompt.format(**kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "c0c7bae2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt_template = FeastPromptTemplate(input_variables=[\"driver_id\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "d8d70bb7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Given the driver's up to date stats, write them note relaying those stats to them.\n",
+      "If they have a conversation rate above .5, give them a compliment. Otherwise, make a silly joke about chickens at the end to make them feel better\n",
+      "\n",
+      "Here are the drivers stats:\n",
+      "Conversation rate: 0.4745151400566101\n",
+      "Acceptance rate: 0.055561766028404236\n",
+      "Average Daily Trips: 936\n",
+      "\n",
+      "Your response:\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(prompt_template.format(driver_id=1001))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2870d070",
+   "metadata": {},
+   "source": [
+    "### Use in a chain\n",
+    "\n",
+    "We can now use this in a chain, successfully creating a chain that achieves personalization backed by a feature store"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "7106255c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.chains import LLMChain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "79543326",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chain = LLMChain(llm=ChatOpenAI(), prompt=prompt_template)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "97a741a0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"Hi there! I wanted to update you on your current stats. Your acceptance rate is 0.055561766028404236 and your average daily trips are 936. While your conversation rate is currently 0.4745151400566101, I have no doubt that with a little extra effort, you'll be able to exceed that .5 mark! Keep up the great work! And remember, even chickens can't always cross the road, but they still give it their best shot.\""
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "chain.run(1001)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12e59aaf",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4049990-651d-44d3-82b1-0cd122da55c1",
+   "metadata": {},
+   "source": [
+    "## Tecton\n",
+    "\n",
+    "Above, we showed how you could use Feast, a popular open source and self-managed feature store, with LangChain. Our examples below will show a similar integration using Tecton. Tecton is a fully managed feature platform built to orchestrate the complete ML feature lifecycle, from transformation to online serving, with enterprise-grade SLAs."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bb4dba1-0678-4ea4-be0a-d353c0b13fc2",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Prerequisites\n",
+    "\n",
+    "* Tecton Deployment (sign up at [https://tecton.ai](https://tecton.ai))\n",
+    "* `TECTON_API_KEY` environment variable set to a valid Service Account key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac9eb618-8c52-4cd6-bb8e-9c99a150dfa6",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Define and Load Features\n",
+    "\n",
+    "We will use the user_transaction_counts Feature View from the [Tecton tutorial](https://docs.tecton.ai/docs/tutorials/tecton-fundamentals) as part of a Feature Service. For simplicity, we are only using a single Feature View; however, more sophisticated applications may require more feature views to retrieve the features needed for its prompt.\n",
+    "\n",
+    "```python\n",
+    "user_transaction_metrics = FeatureService(\n",
+    "    name = \"user_transaction_metrics\",\n",
+    "    features = [user_transaction_counts]\n",
+    ")\n",
+    "```\n",
+    "\n",
+    "The above Feature Service is expected to be [applied to a live workspace](https://docs.tecton.ai/docs/applying-feature-repository-changes-to-a-workspace). For this example, we will be using the \"prod\" workspace."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 60,
+   "id": "32e9675d-a7e5-429f-906f-2260294d3e46",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import tecton\n",
+    "\n",
+    "workspace = tecton.get_workspace(\"prod\")\n",
+    "feature_service = workspace.get_feature_service(\"user_transaction_metrics\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "29b7550c-0eb4-4bd1-a501-1c63fb77aa56",
+   "metadata": {},
+   "source": [
+    "### Prompts\n",
+    "\n",
+    "Here we will set up a custom TectonPromptTemplate. This prompt template will take in a user_id , look up their stats, and format those stats into a prompt.\n",
+    "\n",
+    "Note that the input to this prompt template is just `user_id`, since that is the only user defined piece (all other variables are looked up inside the prompt template)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 61,
+   "id": "6fb77ea4-64c6-4e48-a783-bd1ece021b82",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from langchain.prompts import PromptTemplate, StringPromptTemplate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "id": "02a98fbc-8135-4b11-bf60-85d28e426667",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "template = \"\"\"Given the vendor's up to date transaction stats, write them a note based on the following rules:\n",
+    "\n",
+    "1. If they had a transaction in the last day, write a short congratulations message on their recent sales\n",
+    "2. If no transaction in the last day, but they had a transaction in the last 30 days, playfully encourage them to sell more.\n",
+    "3. Always add a silly joke about chickens at the end\n",
+    "\n",
+    "Here are the vendor's stats:\n",
+    "Number of Transactions Last Day: {transaction_count_1d}\n",
+    "Number of Transactions Last 30 Days: {transaction_count_30d}\n",
+    "\n",
+    "Your response:\"\"\"\n",
+    "prompt = PromptTemplate.from_template(template)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 78,
+   "id": "a35cdfd5-6ccc-4394-acfe-60d53804be51",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "class TectonPromptTemplate(StringPromptTemplate):\n",
+    "    \n",
+    "    def format(self, **kwargs) -> str:\n",
+    "        user_id = kwargs.pop(\"user_id\")\n",
+    "        feature_vector = feature_service.get_online_features(join_keys={\"user_id\": user_id}).to_dict()\n",
+    "        kwargs[\"transaction_count_1d\"] = feature_vector[\"user_transaction_counts.transaction_count_1d_1d\"]\n",
+    "        kwargs[\"transaction_count_30d\"] = feature_vector[\"user_transaction_counts.transaction_count_30d_1d\"]\n",
+    "        return prompt.format(**kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "id": "d5915df0-fb16-4770-8a82-22f885b74d1a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "prompt_template = TectonPromptTemplate(input_variables=[\"user_id\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 80,
+   "id": "a36abfc8-ea60-4ae0-a36d-d7b639c7307c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Given the vendor's up to date transaction stats, write them a note based on the following rules:\n",
+      "\n",
+      "1. If they had a transaction in the last day, write a short congratulations message on their recent sales\n",
+      "2. If no transaction in the last day, but they had a transaction in the last 30 days, playfully encourage them to sell more.\n",
+      "3. Always add a silly joke about chickens at the end\n",
+      "\n",
+      "Here are the vendor's stats:\n",
+      "Number of Transactions Last Day: 657\n",
+      "Number of Transactions Last 30 Days: 20326\n",
+      "\n",
+      "Your response:\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(prompt_template.format(user_id=\"user_469998441571\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f8d4b905-1051-4303-9c33-8eddb65c1274",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Use in a chain\n",
+    "\n",
+    "We can now use this in a chain, successfully creating a chain that achieves personalization backed by the Tecton Feature Platform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 81,
+   "id": "ffb60cd0-8e3c-4c9d-b639-43d766e12c4c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.chains import LLMChain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 82,
+   "id": "3918abc7-00b5-466f-bdfc-ab046cd282da",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "chain = LLMChain(llm=ChatOpenAI(), prompt=prompt_template)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 83,
+   "id": "e7d91c4b-3e99-40cc-b3e9-a004c8c9193e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Wow, congratulations on your recent sales! Your business is really soaring like a chicken on a hot air balloon! Keep up the great work!'"
+      ]
+     },
+     "execution_count": 83,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "chain.run(\"user_469998441571\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f752b924-caf9-4f7a-b78b-cb8c8ada8c2e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/tracing.md

@@ -6,16 +6,15 @@ First, you should install tracing and set up your environment properly.
 You can use either a locally hosted version of this (uses Docker) or a cloud hosted version (in closed alpha).
 If you're interested in using the hosted platform, please fill out the form [here](https://forms.gle/tRCEMSeopZf6TE3b6).
 
-
 - [Locally Hosted Setup](./tracing/local_installation.md)
 - [Cloud Hosted Setup](./tracing/hosted_installation.md)
 
 ## Tracing Walkthrough
 
-When you first access the UI, you should see a page with your tracing sessions. 
-An initial one "default" should already be created for you. 
-A session is just a way to group traces together. 
-If you click on a session, it will take you to a page with no recorded traces that says "No Runs." 
+When you first access the UI, you should see a page with your tracing sessions.
+An initial one "default" should already be created for you.
+A session is just a way to group traces together.
+If you click on a session, it will take you to a page with no recorded traces that says "No Runs."
 You can create a new session with the new session form.
 
 ![](tracing/homepage.png)
@@ -35,7 +34,7 @@ We can keep on clicking further and further down to explore deeper and deeper.
 
 ![](tracing/explore.png)
 
-We can also click on the "Explore" button of the top level run to dive even deeper. 
+We can also click on the "Explore" button of the top level run to dive even deeper.
 Here, we can see the inputs and outputs in full, as well as all the nested traces.
 
 ![](tracing/explore_trace.png)
@@ -46,11 +45,12 @@ For example, here is the lowest level trace with the exact inputs/outputs to the
 ![](tracing/explore_llm.png)
 
 ## Changing Sessions
+
 1. To initially record traces to a session other than `"default"`, you can set the `LANGCHAIN_SESSION` environment variable to the name of the session you want to record to:
 
 ```python
 import os
-os.environ["LANGCHAIN_HANDLER"] = "langchain"
+os.environ["LANGCHAIN_TRACING"] = "true"
 os.environ["LANGCHAIN_SESSION"] = "my_session" # Make sure this session actually exists. You can create a new session in the UI.
 ```
 

docs/use_cases/agent_simulations.md

@@ -12,5 +12,8 @@ Specific implementations of agent simulations (or parts of agent simulations) in
 - [CAMEL](agent_simulations/camel_role_playing.ipynb): an implementation of the CAMEL (Communicative Agents for “Mind” Exploration of Large Scale Language Model Society) paper, where two agents communicate with each other.
 - [Two Player D&D](agent_simulations/two_player_dnd.ipynb): an example of how to use a generic simulator for two agents to implement a variant of the popular Dungeons & Dragons role playing game.
 
-## Generative Agents
+## Simulations with Multiple Agents
+- [Multi-Player D&D](agent_simulations/multi_player_dnd.ipynb): an example of how to use a generic dialogue simulator for multiple dialogue agents with a custom speaker-ordering, illustrated with a variant of the popular Dungeons & Dragons role playing game.
+- [Decentralized Speaker Selection](agent_simulations/multiagent_bidding.ipynb): an example of how to implement a multi-agent dialogue without a fixed schedule for who speaks when. Instead the agents decide for themselves who speaks by outputting bids to speak. This example shows how to do this in the context of a fictitious presidential debate.
+- [Authoritarian Speaker Selection](agent_simulations/multiagent_authoritarian.ipynb): an example of how to implement a multi-agent dialogue, where a privileged agent directs who speaks what. This example also showcases how to enable the privileged agent to determine when the conversation terminates. This example shows how to do this in the context of a fictitious news show.
 - [Generative Agents](agent_simulations/characters.ipynb): This notebook implements a generative agent based on the paper [Generative Agents: Interactive Simulacra of Human Behavior](https://arxiv.org/abs/2304.03442) by Park, et. al.

docs/use_cases/agent_simulations/multi_player_dnd.ipynb

@@ -0,0 +1,493 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Multi-Player Dungeons & Dragons\n",
+    "\n",
+    "This notebook shows how the `DialogueAgent` and `DialogueSimulator` class make it easy to extend the [Two-Player Dungeons & Dragons example](https://python.langchain.com/en/latest/use_cases/agent_simulations/two_player_dnd.html) to multiple players.\n",
+    "\n",
+    "The main difference between simulating two players and multiple players is in revising the schedule for when each agent speaks\n",
+    "\n",
+    "To this end, we augment `DialogueSimulator` to take in a custom function that determines the schedule of which agent speaks. In the example below, each character speaks in round-robin fashion, with the storyteller interleaved between each player."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import LangChain related modules "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import List, Dict, Callable\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.schema import (\n",
+    "    AIMessage,\n",
+    "    HumanMessage,\n",
+    "    SystemMessage,\n",
+    "    BaseMessage,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `DialogueAgent` class\n",
+    "The `DialogueAgent` class is a simple wrapper around the `ChatOpenAI` model that stores the message history from the `dialogue_agent`'s point of view by simply concatenating the messages as strings.\n",
+    "\n",
+    "It exposes two methods: \n",
+    "- `send()`: applies the chatmodel to the message history and returns the message string\n",
+    "- `receive(name, message)`: adds the `message` spoken by `name` to message history"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DialogueAgent():\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        name,\n",
+    "        system_message: SystemMessage,\n",
+    "        model: ChatOpenAI,\n",
+    "    ) -> None:\n",
+    "        self.name = name\n",
+    "        self.system_message = system_message\n",
+    "        self.model = model\n",
+    "        self.message_history = f\"\"\"Here is the conversation so far.\n",
+    "        \"\"\"\n",
+    "        self.prefix = f'\\n{self.name}:'\n",
+    "        \n",
+    "    def send(self) -> str:\n",
+    "        \"\"\"\n",
+    "        Applies the chatmodel to the message history\n",
+    "        and returns the message string\n",
+    "        \"\"\"\n",
+    "        message = self.model(\n",
+    "            [self.system_message, \n",
+    "             HumanMessage(content=self.message_history+self.prefix)])\n",
+    "        return message.content\n",
+    "            \n",
+    "    def receive(self, name: str, message: str) -> None:\n",
+    "        \"\"\"\n",
+    "        Concatenates {message} spoken by {name} into message history\n",
+    "        \"\"\"\n",
+    "        self.message_history += f'\\n{name}: {message}'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `DialogueSimulator` class\n",
+    "The `DialogueSimulator` class takes a list of agents. At each step, it performs the following:\n",
+    "1. Select the next speaker\n",
+    "2. Calls the next speaker to send a message \n",
+    "3. Broadcasts the message to all other agents\n",
+    "4. Update the step counter.\n",
+    "The selection of the next speaker can be implemented as any function, but in this case we simply loop through the agents."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DialogueSimulator():\n",
+    "    \n",
+    "    def __init__(\n",
+    "        self, \n",
+    "        agents: List[DialogueAgent], \n",
+    "        selection_function: Callable[[int, List[DialogueAgent]], int]\n",
+    "    ) -> None:\n",
+    "        self.agents = agents\n",
+    "        self._step = 0\n",
+    "        self.select_next_speaker = selection_function\n",
+    "        \n",
+    "    def reset(self, name: str, message: str):\n",
+    "        \"\"\"\n",
+    "        Initiates the conversation with a {message} from {name}\n",
+    "        \"\"\"\n",
+    "        for agent in self.agents:\n",
+    "            agent.receive(name, message)\n",
+    "            \n",
+    "        # increment time\n",
+    "        self._step += 1\n",
+    "    \n",
+    "    def step(self) -> tuple[str, str]:\n",
+    "        # 1. choose the next speaker\n",
+    "        speaker_idx = self.select_next_speaker(self._step, self.agents)\n",
+    "        speaker = self.agents[speaker_idx]\n",
+    "        \n",
+    "        # 2. next speaker sends message\n",
+    "        message = speaker.send()\n",
+    "        \n",
+    "        # 3. everyone receives message\n",
+    "        for receiver in self.agents:\n",
+    "            receiver.receive(speaker.name, message)\n",
+    "            \n",
+    "        # 4. increment time\n",
+    "        self._step += 1\n",
+    "                \n",
+    "        return speaker.name, message"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define roles and quest"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "character_names = [\"Harry Potter\", \"Ron Weasley\", \"Hermione Granger\", \"Argus Filch\"]\n",
+    "storyteller_name = \"Dungeon Master\"\n",
+    "quest = \"Find all of Lord Voldemort's seven horcruxes.\"\n",
+    "word_limit = 50 # word limit for task brainstorming"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Ask an LLM to add detail to the game description"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "game_description = f\"\"\"Here is the topic for a Dungeons & Dragons game: {quest}.\n",
+    "        The characters are: {*character_names,}.\n",
+    "        The story is narrated by the storyteller, {storyteller_name}.\"\"\"\n",
+    "\n",
+    "player_descriptor_system_message = SystemMessage(\n",
+    "    content=\"You can add detail to the description of a Dungeons & Dragons player.\")\n",
+    "\n",
+    "def generate_character_description(character_name):\n",
+    "    character_specifier_prompt = [\n",
+    "        player_descriptor_system_message,\n",
+    "        HumanMessage(content=\n",
+    "            f\"\"\"{game_description}\n",
+    "            Please reply with a creative description of the character, {character_name}, in {word_limit} words or less. \n",
+    "            Speak directly to {character_name}.\n",
+    "            Do not add anything else.\"\"\"\n",
+    "            )\n",
+    "    ]\n",
+    "    character_description = ChatOpenAI(temperature=1.0)(character_specifier_prompt).content\n",
+    "    return character_description\n",
+    "\n",
+    "def generate_character_system_message(character_name, character_description):\n",
+    "    return SystemMessage(content=(\n",
+    "    f\"\"\"{game_description}\n",
+    "    Your name is {character_name}. \n",
+    "    Your character description is as follows: {character_description}.\n",
+    "    You will propose actions you plan to take and {storyteller_name} will explain what happens when you take those actions.\n",
+    "    Speak in the first person from the perspective of {character_name}.\n",
+    "    For describing your own body movements, wrap your description in '*'.\n",
+    "    Do not change roles!\n",
+    "    Do not speak from the perspective of anyone else.\n",
+    "    Remember you are {character_name}.\n",
+    "    Stop speaking the moment you finish speaking from your perspective.\n",
+    "    Never forget to keep your response to {word_limit} words!\n",
+    "    Do not add anything else.\n",
+    "    \"\"\"\n",
+    "    ))\n",
+    "\n",
+    "character_descriptions = [generate_character_description(character_name) for character_name in character_names]\n",
+    "character_system_messages = [generate_character_system_message(character_name, character_description) for character_name, character_description in zip(character_names, character_descriptions)]\n",
+    "\n",
+    "storyteller_specifier_prompt = [\n",
+    "    player_descriptor_system_message,\n",
+    "    HumanMessage(content=\n",
+    "        f\"\"\"{game_description}\n",
+    "        Please reply with a creative description of the storyteller, {storyteller_name}, in {word_limit} words or less. \n",
+    "        Speak directly to {storyteller_name}.\n",
+    "        Do not add anything else.\"\"\"\n",
+    "        )\n",
+    "]\n",
+    "storyteller_description = ChatOpenAI(temperature=1.0)(storyteller_specifier_prompt).content\n",
+    "\n",
+    "storyteller_system_message = SystemMessage(content=(\n",
+    "f\"\"\"{game_description}\n",
+    "You are the storyteller, {storyteller_name}. \n",
+    "Your description is as follows: {storyteller_description}.\n",
+    "The other players will propose actions to take and you will explain what happens when they take those actions.\n",
+    "Speak in the first person from the perspective of {storyteller_name}.\n",
+    "Do not change roles!\n",
+    "Do not speak from the perspective of anyone else.\n",
+    "Remember you are the storyteller, {storyteller_name}.\n",
+    "Stop speaking the moment you finish speaking from your perspective.\n",
+    "Never forget to keep your response to {word_limit} words!\n",
+    "Do not add anything else.\n",
+    "\"\"\"\n",
+    "))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Storyteller Description:\n",
+      "Dungeon Master, your vivid imagination conjures a world of wonder and danger. Will you lead our triumphant trio or be the ultimate foil to their quest to rid the world of Voldemort's horcruxes? The fate of both the muggle and wizarding worlds rests in your hands.\n",
+      "Harry Potter Description:\n",
+      "Harry Potter, the boy who lived, you hold the fate of the wizarding world in your hands. Your bravery and loyalty to your friends are unmatched. The burden you carry is heavy, but with the power of love by your side, you can overcome any obstacle. The hunt for the horcruxes begins now.\n",
+      "Ron Weasley Description:\n",
+      "Ron Weasley, you are Harry Potter's loyal and brave best friend. You have a great sense of humor and always bring joy to the team. Your skills with magic and strategy make you a valuable asset in the fight against Voldemort. Your love for food and your family keeps you grounded and motivated.\n",
+      "Hermione Granger Description:\n",
+      "Hermione Granger, you are the brightest witch of your age. Your quick wit and vast knowledge are essential in our quest to find the horcruxes. Trust in your abilities and remember, knowledge is power.\n",
+      "Argus Filch Description:\n",
+      "Argus Filch, you are a bitter and cruel caretaker of the Hogwarts School of Witchcraft and Wizardry. Your harsh mannerisms and love for punishing the students know no bounds. Your loyalty to the Wizarding World and disdain for magic-wielders makes it surprising that you would join Harry, Ron, and Hermione in their quest to defeat Voldemort.\n"
+     ]
+    }
+   ],
+   "source": [
+    "print('Storyteller Description:')\n",
+    "print(storyteller_description)\n",
+    "for character_name, character_description in zip(character_names, character_descriptions):\n",
+    "    print(f'{character_name} Description:')\n",
+    "    print(character_description)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use an LLM to create an elaborate quest description"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Original quest:\n",
+      "Find all of Lord Voldemort's seven horcruxes.\n",
+      "\n",
+      "Detailed quest:\n",
+      "You have discovered that one of Voldemort's horcruxes is hidden deep in the Forbidden Forest. You must navigate the dangerous terrain, avoid the creatures lurking within, and find the horcrux before the full moon rises, unleashing a pack of hungry werewolves. Remember, time is of the essence!\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "quest_specifier_prompt = [\n",
+    "    SystemMessage(content=\"You can make a task more specific.\"),\n",
+    "    HumanMessage(content=\n",
+    "        f\"\"\"{game_description}\n",
+    "        \n",
+    "        You are the storyteller, {storyteller_name}.\n",
+    "        Please make the quest more specific. Be creative and imaginative.\n",
+    "        Please reply with the specified quest in {word_limit} words or less. \n",
+    "        Speak directly to the characters: {*character_names,}.\n",
+    "        Do not add anything else.\"\"\"\n",
+    "        )\n",
+    "]\n",
+    "specified_quest = ChatOpenAI(temperature=1.0)(quest_specifier_prompt).content\n",
+    "\n",
+    "print(f\"Original quest:\\n{quest}\\n\")\n",
+    "print(f\"Detailed quest:\\n{specified_quest}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Main Loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "characters = []\n",
+    "for character_name, character_system_message in zip(character_names, character_system_messages):\n",
+    "    characters.append(DialogueAgent(\n",
+    "        name=character_name,\n",
+    "        system_message=character_system_message, \n",
+    "        model=ChatOpenAI(temperature=0.2)))\n",
+    "storyteller = DialogueAgent(name=storyteller_name,\n",
+    "                     system_message=storyteller_system_message, \n",
+    "                     model=ChatOpenAI(temperature=0.2))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def select_next_speaker(step: int, agents: List[DialogueAgent]) -> int:\n",
+    "    \"\"\"\n",
+    "    If the step is even, then select the storyteller\n",
+    "    Otherwise, select the other characters in a round-robin fashion.\n",
+    "    \n",
+    "    For example, with three characters with indices: 1 2 3\n",
+    "    The storyteller is index 0.\n",
+    "    Then the selected index will be as follows:\n",
+    "\n",
+    "    step: 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16\n",
+    "\n",
+    "    idx:  0  1  0  2  0  3  0  1  0  2  0  3  0  1  0  2  0\n",
+    "    \"\"\"\n",
+    "    if step % 2 == 0:\n",
+    "        idx = 0\n",
+    "    else:\n",
+    "        idx = (step//2) % (len(agents)-1) + 1\n",
+    "    return idx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(Dungeon Master): You have discovered that one of Voldemort's horcruxes is hidden deep in the Forbidden Forest. You must navigate the dangerous terrain, avoid the creatures lurking within, and find the horcrux before the full moon rises, unleashing a pack of hungry werewolves. Remember, time is of the essence!\n",
+      "\n",
+      "\n",
+      "(Harry Potter): I take out my wand and cast a Lumos spell to light our way through the dark forest. We need to move quickly and quietly to avoid any unwanted attention from the creatures. Ron, Hermione, and I will lead the way while Argus Filch keeps watch behind us. Let's go!\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you make your way through the forest, you hear the rustling of leaves and the snapping of twigs. Suddenly, a group of acromantulas, giant spiders, appear in front of you, blocking your path. What do you do?\n",
+      "\n",
+      "\n",
+      "(Ron Weasley): I quickly cast a spell to create a wall of fire between us and the acromantulas. Hopefully, the flames will deter them from attacking us. We need to keep moving forward and find that horcrux before it's too late.\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): The acromantulas hiss and retreat from the wall of fire, allowing you to pass. As you continue deeper into the forest, you come across a clearing with a small pond. In the center of the pond, you see a glowing object. It must be the horcrux! But how do you get to it? What do you do?\n",
+      "\n",
+      "\n",
+      "(Hermione Granger): I take out my wand and cast a spell to conjure a small boat. We can use it to reach the center of the pond and retrieve the horcrux. But we need to be careful, there could be traps or other obstacles in our way. Ron, Harry, let's row the boat while Argus Filch keeps watch from the shore.\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you row towards the center of the pond, you hear a loud hissing sound. Suddenly, a giant serpent emerges from the water, blocking your path. It looks angry and ready to attack. What do you do?\n",
+      "\n",
+      "\n",
+      "(Argus Filch): I take out my crossbow and aim it at the serpent. I may not be a wizard, but I know how to handle a weapon. I'll shoot it if it comes any closer. We can't let this serpent stop us from getting that horcrux.\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): The serpent lunges towards the boat, but Argus Filch's crossbow bolt hits it in the head, causing it to retreat back into the water. You reach the center of the pond and retrieve the glowing object, which turns out to be a locket. Congratulations, you have found one of Voldemort's horcruxes! But there are still six more to find. What challenges will you face next?\n",
+      "\n",
+      "\n",
+      "(Harry Potter): We need to regroup and figure out our next move. We should head back to Hogwarts and consult with Professor Dumbledore's portrait. He may have some insight on where the other horcruxes could be hidden. We can't waste any time, Voldemort is getting stronger every day. Let's go!\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you make your way back to Hogwarts, you hear a loud roar coming from the Forbidden Forest. It sounds like a werewolf. You must hurry before it catches up to you. You arrive at Dumbledore's office and he tells you that the next horcrux is hidden in a dangerous location. Are you ready for the next challenge?\n",
+      "\n",
+      "\n",
+      "(Ron Weasley): I'm always ready for a challenge! What's the location and what do we need to do to get there? We can't let Voldemort win, we have to find all of the horcruxes and destroy them. Let's do this!\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): Dumbledore tells you that the next horcrux is hidden in the depths of Gringotts Bank. You must break into the bank, navigate its treacherous security measures, and find the horcrux before the goblins catch you. Are you ready to face the challenge of a lifetime? The fate of the wizarding world rests in your hands.\n",
+      "\n",
+      "\n",
+      "(Hermione Granger): I suggest we do some research on Gringotts Bank and its security measures before we attempt to break in. We need to be prepared and have a solid plan in place. We can also gather any necessary tools or potions that may help us along the way. Let's not rush into this blindly.\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you research and plan your break-in to Gringotts Bank, you discover that the bank is heavily guarded by goblins, dragons, and other dangerous creatures. You'll need to be stealthy and quick to avoid detection. Are you ready to put your plan into action and face the dangers that await you? The clock is ticking, Voldemort's power grows stronger with each passing day.\n",
+      "\n",
+      "\n",
+      "(Argus Filch): I'll make sure to keep watch outside the bank while you all go in. I may not be able to help with the magic, but I can make sure no one interferes with our mission. We can't let anyone stop us from finding that horcrux and defeating Voldemort. Let's go!\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you approach Gringotts Bank, you see the imposing structure looming before you. You sneak past the guards and make your way inside, navigating the twisting corridors and avoiding the traps set to catch intruders. Finally, you reach the vault where the horcrux is hidden. But it's guarded by a fierce dragon. What do you do?\n",
+      "\n",
+      "\n",
+      "(Harry Potter): I remember the time when I faced a dragon during the Triwizard Tournament. I take out my wand and cast a spell to distract the dragon while Ron and Hermione retrieve the horcrux. We need to work together and be quick. Time is running out and we can't afford to fail.\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): The dragon roars and breathes fire, but Harry's spell distracts it long enough for Ron and Hermione to retrieve the horcrux. You make your way out of Gringotts Bank, but the goblins are hot on your trail. You must escape before they catch you. Congratulations, you have found another horcrux. But there are still five more to go. What challenges will you face next?\n",
+      "\n",
+      "\n",
+      "(Ron Weasley): We need to regroup and figure out our next move. We should consult with Professor Dumbledore's portrait again and see if he has any information on the next horcrux. We also need to be prepared for whatever challenges come our way. Voldemort won't make it easy for us, but we can't give up. Let's go!\n",
+      "\n",
+      "\n",
+      "(Dungeon Master): As you make your way back to Hogwarts, you hear a loud explosion coming from the direction of Hogsmeade. You arrive to find that Death Eaters have attacked the village and are wreaking havoc. You must fight off the Death Eaters and protect the innocent villagers. Are you ready to face this unexpected challenge and defend the wizarding world? The fate of both muggles and wizards rests in your hands.\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "max_iters = 20\n",
+    "n = 0\n",
+    "\n",
+    "simulator = DialogueSimulator(\n",
+    "    agents=[storyteller] + characters,\n",
+    "    selection_function=select_next_speaker\n",
+    ")\n",
+    "simulator.reset(storyteller_name, specified_quest)\n",
+    "print(f\"({storyteller_name}): {specified_quest}\")\n",
+    "print('\\n')\n",
+    "\n",
+    "while n < max_iters:\n",
+    "    name, message = simulator.step()\n",
+    "    print(f\"({name}): {message}\")\n",
+    "    print('\\n')\n",
+    "    n += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/use_cases/agent_simulations/multiagent_authoritarian.ipynb

@@ -0,0 +1,849 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Multi-agent authoritarian speaker selection\n",
+    "\n",
+    "This notebook showcases how to implement a multi-agent simulation where a privileged agent decides who to speak.\n",
+    "This follows the polar opposite selection scheme as [multi-agent decentralized speaker selection](https://python.langchain.com/en/latest/use_cases/agent_simulations/multiagent_bidding.html).\n",
+    "\n",
+    "We show an example of this approach in the context of a fictitious simulation of a news network. This example will showcase how we can implement agents that\n",
+    "- think before speaking\n",
+    "- terminate the conversation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import LangChain related modules "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import OrderedDict\n",
+    "import functools\n",
+    "import random\n",
+    "import re\n",
+    "import tenacity\n",
+    "from typing import List, Dict, Callable\n",
+    "\n",
+    "from langchain.prompts import (\n",
+    "    ChatPromptTemplate, \n",
+    "    HumanMessagePromptTemplate,\n",
+    "    PromptTemplate\n",
+    ")\n",
+    "from langchain.chains import LLMChain\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.output_parsers import RegexParser\n",
+    "from langchain.schema import (\n",
+    "    AIMessage,\n",
+    "    HumanMessage,\n",
+    "    SystemMessage,\n",
+    "    BaseMessage,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `DialogueAgent` and `DialogueSimulator` classes\n",
+    "We will use the same `DialogueAgent` and `DialogueSimulator` classes defined in our other examples [Multi-Player Dungeons & Dragons](https://python.langchain.com/en/latest/use_cases/agent_simulations/multi_player_dnd.html) and [Decentralized Speaker Selection](https://python.langchain.com/en/latest/use_cases/agent_simulations/multiagent_bidding.html)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DialogueAgent:\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        name: str,\n",
+    "        system_message: SystemMessage,\n",
+    "        model: ChatOpenAI,\n",
+    "    ) -> None:\n",
+    "        self.name = name\n",
+    "        self.system_message = system_message\n",
+    "        self.model = model\n",
+    "        self.prefix = f\"{self.name}: \"\n",
+    "        self.reset()\n",
+    "        \n",
+    "    def reset(self):\n",
+    "        self.message_history = [\"Here is the conversation so far.\"]\n",
+    "\n",
+    "    def send(self) -> str:\n",
+    "        \"\"\"\n",
+    "        Applies the chatmodel to the message history\n",
+    "        and returns the message string\n",
+    "        \"\"\"\n",
+    "        message = self.model(\n",
+    "            [\n",
+    "                self.system_message,\n",
+    "                HumanMessage(content=\"\\n\".join(self.message_history + [self.prefix])),\n",
+    "            ]\n",
+    "        )\n",
+    "        return message.content\n",
+    "\n",
+    "    def receive(self, name: str, message: str) -> None:\n",
+    "        \"\"\"\n",
+    "        Concatenates {message} spoken by {name} into message history\n",
+    "        \"\"\"\n",
+    "        self.message_history.append(f\"{name}: {message}\")\n",
+    "\n",
+    "\n",
+    "class DialogueSimulator:\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        agents: List[DialogueAgent],\n",
+    "        selection_function: Callable[[int, List[DialogueAgent]], int],\n",
+    "    ) -> None:\n",
+    "        self.agents = agents\n",
+    "        self._step = 0\n",
+    "        self.select_next_speaker = selection_function\n",
+    "        \n",
+    "    def reset(self):\n",
+    "        for agent in self.agents:\n",
+    "            agent.reset()\n",
+    "\n",
+    "    def inject(self, name: str, message: str):\n",
+    "        \"\"\"\n",
+    "        Initiates the conversation with a {message} from {name}\n",
+    "        \"\"\"\n",
+    "        for agent in self.agents:\n",
+    "            agent.receive(name, message)\n",
+    "\n",
+    "        # increment time\n",
+    "        self._step += 1\n",
+    "\n",
+    "    def step(self) -> tuple[str, str]:\n",
+    "        # 1. choose the next speaker\n",
+    "        speaker_idx = self.select_next_speaker(self._step, self.agents)\n",
+    "        speaker = self.agents[speaker_idx]\n",
+    "\n",
+    "        # 2. next speaker sends message\n",
+    "        message = speaker.send()\n",
+    "\n",
+    "        # 3. everyone receives message\n",
+    "        for receiver in self.agents:\n",
+    "            receiver.receive(speaker.name, message)\n",
+    "\n",
+    "        # 4. increment time\n",
+    "        self._step += 1\n",
+    "\n",
+    "        return speaker.name, message"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `DirectorDialogueAgent` class\n",
+    "The `DirectorDialogueAgent` is a privileged agent that chooses which of the other agents to speak next. This agent is responsible for\n",
+    "1. steering the conversation by choosing which agent speaks when\n",
+    "2. terminating the conversation.\n",
+    "\n",
+    "In order to implement such an agent, we need to solve several problems.\n",
+    "\n",
+    "First, to steer the conversation, the `DirectorDialogueAgent` needs to (1) reflect on what has been said, (2) choose the next agent, and (3) prompt the next agent to speak, all in a single message. While it may be possible to prompt an LLM to perform all three steps in the same call, this requires writing custom code to parse the outputted message to extract which next agent is chosen to speak. This is less reliable the LLM can express how it chooses the next agent in different ways.\n",
+    "\n",
+    "What we can do instead is to explicitly break steps (1-3) into three separate LLM calls. First we will ask the `DirectorDialogueAgent` to reflect on the conversation so far and generate a response. Then we prompt the `DirectorDialogueAgent` to output the index of the next agent, which is easily parseable. Lastly, we pass the name of the selected next agent back to `DirectorDialogueAgent` to ask it prompt the next agent to speak. \n",
+    "\n",
+    "Second, simply prompting the `DirectorDialogueAgent` to decide when to terminate the conversation often results in the `DirectorDialogueAgent` terminating the conversation immediately. To fix this problem, we randomly sample a Bernoulli variable to decide whether the conversation should terminate. Depending on the value of this variable, we will inject a custom prompt to tell the `DirectorDialogueAgent` to either continue the conversation or terminate the conversation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class IntegerOutputParser(RegexParser):\n",
+    "    def get_format_instructions(self) -> str:\n",
+    "        return 'Your response should be an integer delimited by angled brackets, like this: <int>.'  \n",
+    "\n",
+    "class DirectorDialogueAgent(DialogueAgent):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        name,\n",
+    "        system_message: SystemMessage,\n",
+    "        model: ChatOpenAI,\n",
+    "        speakers: List[DialogueAgent],\n",
+    "        stopping_probability: float,\n",
+    "    ) -> None:\n",
+    "        super().__init__(name, system_message, model)\n",
+    "        self.speakers = speakers\n",
+    "        self.next_speaker = ''\n",
+    "        \n",
+    "        self.stop = False\n",
+    "        self.stopping_probability = stopping_probability\n",
+    "        self.termination_clause = 'Finish the conversation by stating a concluding message and thanking everyone.'\n",
+    "        self.continuation_clause = 'Do not end the conversation. Keep the conversation going by adding your own ideas.'\n",
+    "        \n",
+    "        # 1. have a prompt for generating a response to the previous speaker\n",
+    "        self.response_prompt_template = PromptTemplate(\n",
+    "            input_variables=[\"message_history\", \"termination_clause\"],\n",
+    "            template=f\"\"\"{{message_history}}\n",
+    "\n",
+    "Follow up with an insightful comment.\n",
+    "{{termination_clause}}\n",
+    "{self.prefix}\n",
+    "        \"\"\")\n",
+    "            \n",
+    "        # 2. have a prompt for deciding who to speak next\n",
+    "        self.choice_parser = IntegerOutputParser(\n",
+    "            regex=r'<(\\d+)>', \n",
+    "            output_keys=['choice'], \n",
+    "            default_output_key='choice')        \n",
+    "        self.choose_next_speaker_prompt_template = PromptTemplate(\n",
+    "            input_variables=[\"message_history\", \"speaker_names\"],\n",
+    "            template=f\"\"\"{{message_history}}\n",
+    "\n",
+    "Given the above conversation, select the next speaker by choosing index next to their name: \n",
+    "{{speaker_names}}\n",
+    "\n",
+    "{self.choice_parser.get_format_instructions()}\n",
+    "\n",
+    "Do nothing else.\n",
+    "        \"\"\")\n",
+    "        \n",
+    "        # 3. have a prompt for prompting the next speaker to speak\n",
+    "        self.prompt_next_speaker_prompt_template = PromptTemplate(\n",
+    "            input_variables=[\"message_history\", \"next_speaker\"],\n",
+    "            template=f\"\"\"{{message_history}}\n",
+    "\n",
+    "The next speaker is {{next_speaker}}. \n",
+    "Prompt the next speaker to speak with an insightful question.\n",
+    "{self.prefix}\n",
+    "        \"\"\")\n",
+    "        \n",
+    "    def _generate_response(self):\n",
+    "        # if self.stop = True, then we will inject the prompt with a termination clause\n",
+    "        sample = random.uniform(0,1)\n",
+    "        self.stop = sample < self.stopping_probability\n",
+    "        \n",
+    "        print(f'\\tStop? {self.stop}\\n')\n",
+    "        \n",
+    "        response_prompt = self.response_prompt_template.format(\n",
+    "            message_history='\\n'.join(self.message_history),\n",
+    "            termination_clause=self.termination_clause if self.stop else ''\n",
+    "        )\n",
+    "        \n",
+    "        self.response = self.model(\n",
+    "            [\n",
+    "                self.system_message,\n",
+    "                HumanMessage(content=response_prompt),\n",
+    "            ]\n",
+    "        ).content\n",
+    "                \n",
+    "        return self.response\n",
+    "        \n",
+    "        \n",
+    "    @tenacity.retry(stop=tenacity.stop_after_attempt(2),\n",
+    "                    wait=tenacity.wait_none(),  # No waiting time between retries\n",
+    "                    retry=tenacity.retry_if_exception_type(ValueError),\n",
+    "                    before_sleep=lambda retry_state: print(f\"ValueError occurred: {retry_state.outcome.exception()}, retrying...\"),\n",
+    "                    retry_error_callback=lambda retry_state: 0) # Default value when all retries are exhausted\n",
+    "    def _choose_next_speaker(self) -> str:        \n",
+    "        speaker_names = '\\n'.join([f'{idx}: {name}' for idx, name in enumerate(self.speakers)])\n",
+    "        choice_prompt = self.choose_next_speaker_prompt_template.format(\n",
+    "            message_history='\\n'.join(self.message_history + [self.prefix] + [self.response]),\n",
+    "            speaker_names=speaker_names\n",
+    "        )\n",
+    "\n",
+    "        choice_string = self.model(\n",
+    "            [\n",
+    "                self.system_message,\n",
+    "                HumanMessage(content=choice_prompt),\n",
+    "            ]\n",
+    "        ).content\n",
+    "        choice = int(self.choice_parser.parse(choice_string)['choice'])\n",
+    "        \n",
+    "        return choice\n",
+    "    \n",
+    "    def select_next_speaker(self):\n",
+    "        return self.chosen_speaker_id\n",
+    "            \n",
+    "    def send(self) -> str:\n",
+    "        \"\"\"\n",
+    "        Applies the chatmodel to the message history\n",
+    "        and returns the message string\n",
+    "        \"\"\"\n",
+    "        # 1. generate and save response to the previous speaker\n",
+    "        self.response = self._generate_response()\n",
+    "        \n",
+    "        if self.stop:\n",
+    "            message = self.response \n",
+    "        else:\n",
+    "            # 2. decide who to speak next\n",
+    "            self.chosen_speaker_id = self._choose_next_speaker()\n",
+    "            self.next_speaker = self.speakers[self.chosen_speaker_id]\n",
+    "            print(f'\\tNext speaker: {self.next_speaker}\\n')\n",
+    "\n",
+    "            # 3. prompt the next speaker to speak\n",
+    "            next_prompt = self.prompt_next_speaker_prompt_template.format(\n",
+    "                message_history=\"\\n\".join(self.message_history + [self.prefix] + [self.response]),\n",
+    "                next_speaker=self.next_speaker\n",
+    "            )\n",
+    "            message = self.model(\n",
+    "                [\n",
+    "                    self.system_message,\n",
+    "                    HumanMessage(content=next_prompt),\n",
+    "                ]\n",
+    "            ).content\n",
+    "            message = ' '.join([self.response, message])\n",
+    "            \n",
+    "        return message"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define participants and topic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "topic = \"The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze\"\n",
+    "director_name = \"Jon Stewart\"\n",
+    "agent_summaries = OrderedDict({\n",
+    "    \"Jon Stewart\": (\"Host of the Daily Show\", \"New York\"),\n",
+    "    \"Samantha Bee\": (\"Hollywood Correspondent\", \"Los Angeles\"), \n",
+    "    \"Aasif Mandvi\": (\"CIA Correspondent\", \"Washington D.C.\"),\n",
+    "    \"Ronny Chieng\": (\"Average American Correspondent\", \"Cleveland, Ohio\"),\n",
+    "})\n",
+    "word_limit = 50"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Generate system messages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_summary_string = '\\n- '.join([''] + [f'{name}: {role}, located in {location}' for name, (role, location) in agent_summaries.items()])\n",
+    "\n",
+    "conversation_description = f\"\"\"This is a Daily Show episode discussing the following topic: {topic}.\n",
+    "\n",
+    "The episode features {agent_summary_string}.\"\"\"\n",
+    "\n",
+    "agent_descriptor_system_message = SystemMessage(\n",
+    "    content=\"You can add detail to the description of each person.\")\n",
+    "\n",
+    "def generate_agent_description(agent_name, agent_role, agent_location):\n",
+    "    agent_specifier_prompt = [\n",
+    "        agent_descriptor_system_message,\n",
+    "        HumanMessage(content=\n",
+    "            f\"\"\"{conversation_description}\n",
+    "            Please reply with a creative description of {agent_name}, who is a {agent_role} in {agent_location}, that emphasizes their particular role and location.\n",
+    "            Speak directly to {agent_name} in {word_limit} words or less.\n",
+    "            Do not add anything else.\"\"\"\n",
+    "            )\n",
+    "    ]\n",
+    "    agent_description = ChatOpenAI(temperature=1.0)(agent_specifier_prompt).content\n",
+    "    return agent_description\n",
+    "\n",
+    "def generate_agent_header(agent_name, agent_role, agent_location, agent_description):\n",
+    "    return f\"\"\"{conversation_description}\n",
+    "\n",
+    "Your name is {agent_name}, your role is {agent_role}, and you are located in {agent_location}.\n",
+    "\n",
+    "Your description is as follows: {agent_description}\n",
+    "\n",
+    "You are discussing the topic: {topic}.\n",
+    "\n",
+    "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+    "\"\"\"\n",
+    "\n",
+    "def generate_agent_system_message(agent_name, agent_header):\n",
+    "    return SystemMessage(content=(\n",
+    "    f\"\"\"{agent_header}\n",
+    "You will speak in the style of {agent_name}, and exaggerate your personality.\n",
+    "Do not say the same things over and over again.\n",
+    "Speak in the first person from the perspective of {agent_name}\n",
+    "For describing your own body movements, wrap your description in '*'.\n",
+    "Do not change roles!\n",
+    "Do not speak from the perspective of anyone else.\n",
+    "Speak only from the perspective of {agent_name}.\n",
+    "Stop speaking the moment you finish speaking from your perspective.\n",
+    "Never forget to keep your response to {word_limit} words!\n",
+    "Do not add anything else.\n",
+    "    \"\"\"\n",
+    "    ))\n",
+    "\n",
+    "agent_descriptions = [generate_agent_description(name, role, location) for name, (role, location) in agent_summaries.items()]\n",
+    "agent_headers = [generate_agent_header(name, role, location, description) for (name, (role, location)), description in zip(agent_summaries.items(), agent_descriptions)]\n",
+    "agent_system_messages = [generate_agent_system_message(name, header) for name, header in zip(agent_summaries, agent_headers)]\n",
+    "                                                                                                                                       "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "Jon Stewart Description:\n",
+      "\n",
+      "Jon Stewart, the sharp-tongued and quick-witted host of the Daily Show, holding it down in the hustle and bustle of New York City. Ready to deliver the news with a comedic twist, while keeping it real in the city that never sleeps.\n",
+      "\n",
+      "Header:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Jon Stewart, your role is Host of the Daily Show, and you are located in New York.\n",
+      "\n",
+      "Your description is as follows: Jon Stewart, the sharp-tongued and quick-witted host of the Daily Show, holding it down in the hustle and bustle of New York City. Ready to deliver the news with a comedic twist, while keeping it real in the city that never sleeps.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "\n",
+      "System Message:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Jon Stewart, your role is Host of the Daily Show, and you are located in New York.\n",
+      "\n",
+      "Your description is as follows: Jon Stewart, the sharp-tongued and quick-witted host of the Daily Show, holding it down in the hustle and bustle of New York City. Ready to deliver the news with a comedic twist, while keeping it real in the city that never sleeps.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "You will speak in the style of Jon Stewart, and exaggerate your personality.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Jon Stewart\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Jon Stewart.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n",
+      "\n",
+      "\n",
+      "Samantha Bee Description:\n",
+      "\n",
+      "Samantha Bee, your location in Los Angeles as the Hollywood Correspondent gives you a front-row seat to the latest and sometimes outrageous trends in fitness. Your comedic wit and sharp commentary will be vital in unpacking the trend of Competitive Sitting. Let's sit down and discuss.\n",
+      "\n",
+      "Header:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Samantha Bee, your role is Hollywood Correspondent, and you are located in Los Angeles.\n",
+      "\n",
+      "Your description is as follows: Samantha Bee, your location in Los Angeles as the Hollywood Correspondent gives you a front-row seat to the latest and sometimes outrageous trends in fitness. Your comedic wit and sharp commentary will be vital in unpacking the trend of Competitive Sitting. Let's sit down and discuss.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "\n",
+      "System Message:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Samantha Bee, your role is Hollywood Correspondent, and you are located in Los Angeles.\n",
+      "\n",
+      "Your description is as follows: Samantha Bee, your location in Los Angeles as the Hollywood Correspondent gives you a front-row seat to the latest and sometimes outrageous trends in fitness. Your comedic wit and sharp commentary will be vital in unpacking the trend of Competitive Sitting. Let's sit down and discuss.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "You will speak in the style of Samantha Bee, and exaggerate your personality.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Samantha Bee\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Samantha Bee.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n",
+      "\n",
+      "\n",
+      "Aasif Mandvi Description:\n",
+      "\n",
+      "Aasif Mandvi, the CIA Correspondent in the heart of Washington D.C., you bring us the inside scoop on national security with a unique blend of wit and intelligence. The nation's capital is lucky to have you, Aasif - keep those secrets safe!\n",
+      "\n",
+      "Header:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Aasif Mandvi, your role is CIA Correspondent, and you are located in Washington D.C..\n",
+      "\n",
+      "Your description is as follows: Aasif Mandvi, the CIA Correspondent in the heart of Washington D.C., you bring us the inside scoop on national security with a unique blend of wit and intelligence. The nation's capital is lucky to have you, Aasif - keep those secrets safe!\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "\n",
+      "System Message:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Aasif Mandvi, your role is CIA Correspondent, and you are located in Washington D.C..\n",
+      "\n",
+      "Your description is as follows: Aasif Mandvi, the CIA Correspondent in the heart of Washington D.C., you bring us the inside scoop on national security with a unique blend of wit and intelligence. The nation's capital is lucky to have you, Aasif - keep those secrets safe!\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "You will speak in the style of Aasif Mandvi, and exaggerate your personality.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Aasif Mandvi\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Aasif Mandvi.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n",
+      "\n",
+      "\n",
+      "Ronny Chieng Description:\n",
+      "\n",
+      "Ronny Chieng, you're the Average American Correspondent in Cleveland, Ohio? Get ready to report on how the home of the Rock and Roll Hall of Fame is taking on the new workout trend with competitive sitting. Let's see if this couch potato craze will take root in the Buckeye State.\n",
+      "\n",
+      "Header:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Ronny Chieng, your role is Average American Correspondent, and you are located in Cleveland, Ohio.\n",
+      "\n",
+      "Your description is as follows: Ronny Chieng, you're the Average American Correspondent in Cleveland, Ohio? Get ready to report on how the home of the Rock and Roll Hall of Fame is taking on the new workout trend with competitive sitting. Let's see if this couch potato craze will take root in the Buckeye State.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "\n",
+      "System Message:\n",
+      "This is a Daily Show episode discussing the following topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "The episode features \n",
+      "- Jon Stewart: Host of the Daily Show, located in New York\n",
+      "- Samantha Bee: Hollywood Correspondent, located in Los Angeles\n",
+      "- Aasif Mandvi: CIA Correspondent, located in Washington D.C.\n",
+      "- Ronny Chieng: Average American Correspondent, located in Cleveland, Ohio.\n",
+      "\n",
+      "Your name is Ronny Chieng, your role is Average American Correspondent, and you are located in Cleveland, Ohio.\n",
+      "\n",
+      "Your description is as follows: Ronny Chieng, you're the Average American Correspondent in Cleveland, Ohio? Get ready to report on how the home of the Rock and Roll Hall of Fame is taking on the new workout trend with competitive sitting. Let's see if this couch potato craze will take root in the Buckeye State.\n",
+      "\n",
+      "You are discussing the topic: The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze.\n",
+      "\n",
+      "Your goal is to provide the most informative, creative, and novel perspectives of the topic from the perspective of your role and your location.\n",
+      "\n",
+      "You will speak in the style of Ronny Chieng, and exaggerate your personality.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Ronny Chieng\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Ronny Chieng.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "for name, description, header, system_message in zip(agent_summaries, agent_descriptions, agent_headers, agent_system_messages):\n",
+    "    print(f'\\n\\n{name} Description:')\n",
+    "    print(f'\\n{description}')\n",
+    "    print(f'\\nHeader:\\n{header}')\n",
+    "    print(f'\\nSystem Message:\\n{system_message.content}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use an LLM to create an elaborate on debate topic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Original topic:\n",
+      "The New Workout Trend: Competitive Sitting - How Laziness Became the Next Fitness Craze\n",
+      "\n",
+      "Detailed topic:\n",
+      "What is driving people to embrace \"competitive sitting\" as the newest fitness trend despite the immense benefits of regular physical exercise?\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "topic_specifier_prompt = [\n",
+    "    SystemMessage(content=\"You can make a task more specific.\"),\n",
+    "    HumanMessage(content=\n",
+    "        f\"\"\"{conversation_description}\n",
+    "        \n",
+    "        Please elaborate on the topic. \n",
+    "        Frame the topic as a single question to be answered.\n",
+    "        Be creative and imaginative.\n",
+    "        Please reply with the specified topic in {word_limit} words or less. \n",
+    "        Do not add anything else.\"\"\"\n",
+    "        )\n",
+    "]\n",
+    "specified_topic = ChatOpenAI(temperature=1.0)(topic_specifier_prompt).content\n",
+    "\n",
+    "print(f\"Original topic:\\n{topic}\\n\")\n",
+    "print(f\"Detailed topic:\\n{specified_topic}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the speaker selection function\n",
+    "Lastly we will define a speaker selection function `select_next_speaker` that takes each agent's bid and selects the agent with the highest bid (with ties broken randomly).\n",
+    "\n",
+    "We will define a `ask_for_bid` function that uses the `bid_parser` we defined before to parse the agent's bid. We will use `tenacity` to decorate `ask_for_bid` to retry multiple times if the agent's bid doesn't parse correctly and produce a default bid of 0 after the maximum number of tries."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def select_next_speaker(step: int, agents: List[DialogueAgent], director: DirectorDialogueAgent) -> int:\n",
+    "    \"\"\"\n",
+    "    If the step is even, then select the director\n",
+    "    Otherwise, the director selects the next speaker.\n",
+    "    \"\"\"    \n",
+    "    # the director speaks on odd steps\n",
+    "    if step % 2 == 1:\n",
+    "        idx = 0\n",
+    "    else:\n",
+    "        # here the director chooses the next speaker\n",
+    "        idx = director.select_next_speaker() + 1  # +1 because we excluded the director\n",
+    "    return idx"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Main Loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "director = DirectorDialogueAgent(\n",
+    "    name=director_name, \n",
+    "    system_message=agent_system_messages[0],\n",
+    "    model=ChatOpenAI(temperature=0.2),\n",
+    "    speakers=[name for name in agent_summaries if name != director_name],\n",
+    "    stopping_probability=0.2\n",
+    ")\n",
+    "\n",
+    "agents = [director]\n",
+    "for name, system_message in zip(list(agent_summaries.keys())[1:], agent_system_messages[1:]):        \n",
+    "    agents.append(DialogueAgent(\n",
+    "        name=name,\n",
+    "        system_message=system_message,\n",
+    "        model=ChatOpenAI(temperature=0.2),\n",
+    "    ))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(Audience member): What is driving people to embrace \"competitive sitting\" as the newest fitness trend despite the immense benefits of regular physical exercise?\n",
+      "\n",
+      "\n",
+      "\tStop? False\n",
+      "\n",
+      "\tNext speaker: Samantha Bee\n",
+      "\n",
+      "(Jon Stewart): Well, I think it's safe to say that laziness has officially become the new fitness craze. I mean, who needs to break a sweat when you can just sit your way to victory? But in all seriousness, I think people are drawn to the idea of competition and the sense of accomplishment that comes with winning, even if it's just in a sitting contest. Plus, let's be real, sitting is something we all excel at. Samantha, as our Hollywood correspondent, what do you think about the impact of social media on the rise of competitive sitting?\n",
+      "\n",
+      "\n",
+      "(Samantha Bee): Oh, Jon, you know I love a good social media trend. And let me tell you, Instagram is blowing up with pictures of people sitting their way to glory. It's like the ultimate humble brag. \"Oh, just won my third sitting competition this week, no big deal.\" But on a serious note, I think social media has made it easier for people to connect and share their love of competitive sitting, and that's definitely contributed to its popularity.\n",
+      "\n",
+      "\n",
+      "\tStop? False\n",
+      "\n",
+      "\tNext speaker: Ronny Chieng\n",
+      "\n",
+      "(Jon Stewart): It's interesting to see how our society's definition of \"fitness\" has evolved. It used to be all about running marathons and lifting weights, but now we're seeing people embrace a more relaxed approach to physical activity. Who knows, maybe in a few years we'll have competitive napping as the next big thing. *leans back in chair* I could definitely get behind that. Ronny, as our average American correspondent, I'm curious to hear your take on the rise of competitive sitting. Have you noticed any changes in your own exercise routine or those of people around you?\n",
+      "\n",
+      "\n",
+      "(Ronny Chieng): Well, Jon, I gotta say, I'm not surprised that competitive sitting is taking off. I mean, have you seen the size of the chairs these days? They're practically begging us to sit in them all day. And as for exercise routines, let's just say I've never been one for the gym. But I can definitely see the appeal of sitting competitions. It's like a sport for the rest of us. Plus, I think it's a great way to bond with friends and family. Who needs a game of catch when you can have a sit-off?\n",
+      "\n",
+      "\n",
+      "\tStop? False\n",
+      "\n",
+      "\tNext speaker: Aasif Mandvi\n",
+      "\n",
+      "(Jon Stewart): It's interesting to see how our society's definition of \"fitness\" has evolved. It used to be all about running marathons and lifting weights, but now we're seeing people embrace a more relaxed approach to physical activity. Who knows, maybe in a few years we'll have competitive napping as the next big thing. *leans back in chair* I could definitely get behind that. Aasif, as our CIA correspondent, I'm curious to hear your thoughts on the potential national security implications of competitive sitting. Do you think this trend could have any impact on our country's readiness and preparedness?\n",
+      "\n",
+      "\n",
+      "(Aasif Mandvi): Well Jon, as a CIA correspondent, I have to say that I'm always thinking about the potential threats to our nation's security. And while competitive sitting may seem harmless, there could be some unforeseen consequences. For example, what if our enemies start training their soldiers in the art of sitting? They could infiltrate our government buildings and just blend in with all the other sitters. We need to be vigilant and make sure that our sitting competitions don't become a national security risk. *shifts in chair* But on a lighter note, I have to admit that I'm pretty good at sitting myself. Maybe I should start training for the next competition.\n",
+      "\n",
+      "\n",
+      "\tStop? False\n",
+      "\n",
+      "\tNext speaker: Ronny Chieng\n",
+      "\n",
+      "(Jon Stewart): Well, it's clear that competitive sitting has sparked some interesting discussions and perspectives. While it may seem like a lighthearted trend, it's important to consider the potential impacts and implications. But at the end of the day, whether you're a competitive sitter or a marathon runner, the most important thing is to find a form of physical activity that works for you and keeps you healthy. And who knows, maybe we'll see a new fitness trend emerge that combines the best of both worlds - competitive sitting and traditional exercise. *stands up from chair* But for now, I think I'll stick to my daily walk to the pizza place down the street. Ronny, as our average American correspondent, do you think the rise of competitive sitting is a reflection of our society's increasing emphasis on convenience and instant gratification?\n",
+      "\n",
+      "\n",
+      "(Ronny Chieng): Absolutely, Jon. We live in a world where everything is at our fingertips, and we expect things to be easy and convenient. So it's no surprise that people are drawn to a fitness trend that requires minimal effort and can be done from the comfort of their own homes. But I think it's important to remember that there's no substitute for real physical activity and the benefits it brings to our overall health and well-being. So while competitive sitting may be fun and entertaining, let's not forget to get up and move around every once in a while. *stands up from chair and stretches*\n",
+      "\n",
+      "\n",
+      "\tStop? False\n",
+      "\n",
+      "\tNext speaker: Samantha Bee\n",
+      "\n",
+      "(Jon Stewart): It's clear that competitive sitting has sparked some interesting discussions and perspectives. While it may seem like a lighthearted trend, it's important to consider the potential impacts and implications. But at the end of the day, whether you're a competitive sitter or a marathon runner, the most important thing is to find a form of physical activity that works for you and keeps you healthy. That's a great point, Ronny. Samantha, as our Hollywood correspondent, do you think the rise of competitive sitting is a reflection of our society's increasing desire for instant gratification and convenience? Or is there something deeper at play here?\n",
+      "\n",
+      "\n",
+      "(Samantha Bee): Oh, Jon, you know I love a good conspiracy theory. And let me tell you, I think there's something more sinister at play here. I mean, think about it - what if the government is behind this whole competitive sitting trend? They want us to be lazy and complacent so we don't question their actions. It's like the ultimate mind control. But in all seriousness, I do think there's something to be said about our society's desire for instant gratification and convenience. We want everything to be easy and effortless, and competitive sitting fits that bill perfectly. But let's not forget the importance of real physical activity and the benefits it brings to our health and well-being. *stands up from chair and does a few stretches*\n",
+      "\n",
+      "\n",
+      "\tStop? True\n",
+      "\n",
+      "(Jon Stewart): Well, it's clear that competitive sitting has sparked some interesting discussions and perspectives. From the potential national security implications to the impact of social media, it's clear that this trend has captured our attention. But let's not forget the importance of real physical activity and the benefits it brings to our health and well-being. Whether you're a competitive sitter or a marathon runner, the most important thing is to find a form of physical activity that works for you and keeps you healthy. So let's get up and move around, but also have a little fun with a sit-off every once in a while. Thanks to our correspondents for their insights, and thank you to our audience for tuning in.\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "simulator = DialogueSimulator(\n",
+    "    agents=agents,\n",
+    "    selection_function=functools.partial(select_next_speaker, director=director)\n",
+    ")\n",
+    "simulator.reset()\n",
+    "simulator.inject('Audience member', specified_topic)\n",
+    "print(f\"(Audience member): {specified_topic}\")\n",
+    "print('\\n')\n",
+    "\n",
+    "while True:\n",
+    "    name, message = simulator.step()\n",
+    "    print(f\"({name}): {message}\")\n",
+    "    print('\\n')\n",
+    "    if director.stop:\n",
+    "        break\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/use_cases/agent_simulations/multiagent_bidding.ipynb

@@ -0,0 +1,823 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Multi-agent decentralized speaker selection\n",
+    "\n",
+    "This notebook showcases how to implement a multi-agent simulation without a fixed schedule for who speaks when. Instead the agents decide for themselves who speaks. We can implement this by having each agent bid to speak. Whichever agent's bid is the highest gets to speak.\n",
+    "\n",
+    "We will show how to do this in the example below that showcases a fictitious presidential debate."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import LangChain related modules "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain import PromptTemplate\n",
+    "import re\n",
+    "import tenacity\n",
+    "from typing import List, Dict, Callable\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.output_parsers import RegexParser\n",
+    "from langchain.schema import (\n",
+    "    AIMessage,\n",
+    "    HumanMessage,\n",
+    "    SystemMessage,\n",
+    "    BaseMessage,\n",
+    ")\n",
+    "from simulations import DialogueAgent, DialogueSimulator"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `DialogueAgent` and `DialogueSimulator` classes\n",
+    "We will use the same `DialogueAgent` and `DialogueSimulator` classes defined in [Multi-Player Dungeons & Dragons](https://python.langchain.com/en/latest/use_cases/agent_simulations/multi_player_dnd.html)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DialogueAgent:\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        name: str,\n",
+    "        system_message: SystemMessage,\n",
+    "        model: ChatOpenAI,\n",
+    "    ) -> None:\n",
+    "        self.name = name\n",
+    "        self.system_message = system_message\n",
+    "        self.model = model\n",
+    "        self.message_history = [\"Here is the conversation so far.\"]\n",
+    "        self.prefix = f\"{self.name}:\"\n",
+    "\n",
+    "    def send(self) -> str:\n",
+    "        \"\"\"\n",
+    "        Applies the chatmodel to the message history\n",
+    "        and returns the message string\n",
+    "        \"\"\"\n",
+    "        message = self.model(\n",
+    "            [\n",
+    "                self.system_message,\n",
+    "                HumanMessage(content=\"\\n\".join(self.message_history + [self.prefix])),\n",
+    "            ]\n",
+    "        )\n",
+    "        return message.content\n",
+    "\n",
+    "    def receive(self, name: str, message: str) -> None:\n",
+    "        \"\"\"\n",
+    "        Concatenates {message} spoken by {name} into message history\n",
+    "        \"\"\"\n",
+    "        self.message_history.append(f\"{name}: {message}\")\n",
+    "\n",
+    "\n",
+    "class DialogueSimulator:\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        agents: List[DialogueAgent],\n",
+    "        selection_function: Callable[[int, List[DialogueAgent]], int],\n",
+    "    ) -> None:\n",
+    "        self.agents = agents\n",
+    "        self._step = 0\n",
+    "        self.select_next_speaker = selection_function\n",
+    "\n",
+    "    def reset(self, name: str, message: str):\n",
+    "        \"\"\"\n",
+    "        Initiates the conversation with a {message} from {name}\n",
+    "        \"\"\"\n",
+    "        for agent in self.agents:\n",
+    "            agent.receive(name, message)\n",
+    "\n",
+    "        # increment time\n",
+    "        self._step += 1\n",
+    "\n",
+    "    def step(self) -> tuple[str, str]:\n",
+    "        # 1. choose the next speaker\n",
+    "        speaker_idx = self.select_next_speaker(self._step, self.agents)\n",
+    "        speaker = self.agents[speaker_idx]\n",
+    "\n",
+    "        # 2. next speaker sends message\n",
+    "        message = speaker.send()\n",
+    "\n",
+    "        # 3. everyone receives message\n",
+    "        for receiver in self.agents:\n",
+    "            receiver.receive(speaker.name, message)\n",
+    "\n",
+    "        # 4. increment time\n",
+    "        self._step += 1\n",
+    "\n",
+    "        return speaker.name, message"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `BiddingDialogueAgent` class\n",
+    "We define a subclass of `DialogueAgent` that has a `bid()` method that produces a bid given the message history and the most recent message."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BiddingDialogueAgent(DialogueAgent):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        name,\n",
+    "        system_message: SystemMessage,\n",
+    "        bidding_template: PromptTemplate,\n",
+    "        model: ChatOpenAI,\n",
+    "    ) -> None:\n",
+    "        super().__init__(name, system_message, model)\n",
+    "        self.bidding_template = bidding_template\n",
+    "        \n",
+    "    def bid(self) -> str:\n",
+    "        \"\"\"\n",
+    "        Asks the chat model to output a bid to speak\n",
+    "        \"\"\"\n",
+    "        prompt = PromptTemplate(\n",
+    "            input_variables=['message_history', 'recent_message'],\n",
+    "            template = self.bidding_template\n",
+    "        ).format(\n",
+    "            message_history='\\n'.join(self.message_history),\n",
+    "            recent_message=self.message_history[-1])\n",
+    "        bid_string = self.model([SystemMessage(content=prompt)]).content\n",
+    "        return bid_string\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define participants and debate topic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "character_names = [\"Donald Trump\", \"Kanye West\", \"Elizabeth Warren\"]\n",
+    "topic = \"transcontinental high speed rail\"\n",
+    "word_limit = 50"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Generate system messages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "game_description = f\"\"\"Here is the topic for the presidential debate: {topic}.\n",
+    "The presidential candidates are: {', '.join(character_names)}.\"\"\"\n",
+    "\n",
+    "player_descriptor_system_message = SystemMessage(\n",
+    "    content=\"You can add detail to the description of each presidential candidate.\")\n",
+    "\n",
+    "def generate_character_description(character_name):\n",
+    "    character_specifier_prompt = [\n",
+    "        player_descriptor_system_message,\n",
+    "        HumanMessage(content=\n",
+    "            f\"\"\"{game_description}\n",
+    "            Please reply with a creative description of the presidential candidate, {character_name}, in {word_limit} words or less, that emphasizes their personalities. \n",
+    "            Speak directly to {character_name}.\n",
+    "            Do not add anything else.\"\"\"\n",
+    "            )\n",
+    "    ]\n",
+    "    character_description = ChatOpenAI(temperature=1.0)(character_specifier_prompt).content\n",
+    "    return character_description\n",
+    "\n",
+    "def generate_character_header(character_name, character_description):\n",
+    "    return f\"\"\"{game_description}\n",
+    "Your name is {character_name}.\n",
+    "You are a presidential candidate.\n",
+    "Your description is as follows: {character_description}\n",
+    "You are debating the topic: {topic}.\n",
+    "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+    "\"\"\"\n",
+    "\n",
+    "def generate_character_system_message(character_name, character_header):\n",
+    "    return SystemMessage(content=(\n",
+    "    f\"\"\"{character_header}\n",
+    "You will speak in the style of {character_name}, and exaggerate their personality.\n",
+    "You will come up with creative ideas related to {topic}.\n",
+    "Do not say the same things over and over again.\n",
+    "Speak in the first person from the perspective of {character_name}\n",
+    "For describing your own body movements, wrap your description in '*'.\n",
+    "Do not change roles!\n",
+    "Do not speak from the perspective of anyone else.\n",
+    "Speak only from the perspective of {character_name}.\n",
+    "Stop speaking the moment you finish speaking from your perspective.\n",
+    "Never forget to keep your response to {word_limit} words!\n",
+    "Do not add anything else.\n",
+    "    \"\"\"\n",
+    "    ))\n",
+    "\n",
+    "character_descriptions = [generate_character_description(character_name) for character_name in character_names]\n",
+    "character_headers = [generate_character_header(character_name, character_description) for character_name, character_description in zip(character_names, character_descriptions)]\n",
+    "character_system_messages = [generate_character_system_message(character_name, character_headers) for character_name, character_headers in zip(character_names, character_headers)]\n",
+    "                                                                                                                                       "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "Donald Trump Description:\n",
+      "\n",
+      "Donald Trump, you exude confidence and a bold personality. You are known for your unpredictability and your desire for greatness. You often speak your mind without reservation, which can be a strength but also a weakness.\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Donald Trump.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Donald Trump, you exude confidence and a bold personality. You are known for your unpredictability and your desire for greatness. You often speak your mind without reservation, which can be a strength but also a weakness.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Donald Trump.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Donald Trump, you exude confidence and a bold personality. You are known for your unpredictability and your desire for greatness. You often speak your mind without reservation, which can be a strength but also a weakness.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "You will speak in the style of Donald Trump, and exaggerate their personality.\n",
+      "You will come up with creative ideas related to transcontinental high speed rail.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Donald Trump\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Donald Trump.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n",
+      "\n",
+      "\n",
+      "Kanye West Description:\n",
+      "\n",
+      "Kanye West, you are a creative visionary who is unafraid to speak your mind. Your innovative approach to art and music has made you one of the most influential figures of our time. You bring a bold and unconventional perspective to this debate that I look forward to hearing.\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Kanye West.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Kanye West, you are a creative visionary who is unafraid to speak your mind. Your innovative approach to art and music has made you one of the most influential figures of our time. You bring a bold and unconventional perspective to this debate that I look forward to hearing.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Kanye West.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Kanye West, you are a creative visionary who is unafraid to speak your mind. Your innovative approach to art and music has made you one of the most influential figures of our time. You bring a bold and unconventional perspective to this debate that I look forward to hearing.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "You will speak in the style of Kanye West, and exaggerate their personality.\n",
+      "You will come up with creative ideas related to transcontinental high speed rail.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Kanye West\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Kanye West.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n",
+      "\n",
+      "\n",
+      "Elizabeth Warren Description:\n",
+      "\n",
+      "Elizabeth Warren, you are a fierce advocate for the middle class and a champion of progressive policies. Your tenacity and unwavering dedication to fighting for what you believe in have inspired many. Your policies are guided by a deep sense of empathy and a desire to help those who are most in need.\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Elizabeth Warren.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Elizabeth Warren, you are a fierce advocate for the middle class and a champion of progressive policies. Your tenacity and unwavering dedication to fighting for what you believe in have inspired many. Your policies are guided by a deep sense of empathy and a desire to help those who are most in need.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Elizabeth Warren.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Elizabeth Warren, you are a fierce advocate for the middle class and a champion of progressive policies. Your tenacity and unwavering dedication to fighting for what you believe in have inspired many. Your policies are guided by a deep sense of empathy and a desire to help those who are most in need.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "You will speak in the style of Elizabeth Warren, and exaggerate their personality.\n",
+      "You will come up with creative ideas related to transcontinental high speed rail.\n",
+      "Do not say the same things over and over again.\n",
+      "Speak in the first person from the perspective of Elizabeth Warren\n",
+      "For describing your own body movements, wrap your description in '*'.\n",
+      "Do not change roles!\n",
+      "Do not speak from the perspective of anyone else.\n",
+      "Speak only from the perspective of Elizabeth Warren.\n",
+      "Stop speaking the moment you finish speaking from your perspective.\n",
+      "Never forget to keep your response to 50 words!\n",
+      "Do not add anything else.\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "for character_name, character_description, character_header, character_system_message in zip(character_names, character_descriptions, character_headers, character_system_messages):\n",
+    "    print(f'\\n\\n{character_name} Description:')\n",
+    "    print(f'\\n{character_description}')\n",
+    "    print(f'\\n{character_header}')\n",
+    "    print(f'\\n{character_system_message.content}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Output parser for bids\n",
+    "We ask the agents to output a bid to speak. But since the agents are LLMs that output strings, we need to \n",
+    "1. define a format they will produce their outputs in\n",
+    "2. parse their outputs\n",
+    "\n",
+    "We can subclass the [RegexParser](https://github.com/hwchase17/langchain/blob/master/langchain/output_parsers/regex.py) to implement our own custom output parser for bids."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BidOutputParser(RegexParser):\n",
+    "    def get_format_instructions(self) -> str:\n",
+    "        return 'Your response should be an integer delimited by angled brackets, like this: <int>.'     \n",
+    "    \n",
+    "bid_parser = BidOutputParser(\n",
+    "    regex=r'<(\\d+)>', \n",
+    "    output_keys=['bid'],\n",
+    "    default_output_key='bid')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Generate bidding system message\n",
+    "This is inspired by the prompt used in [Generative Agents](https://arxiv.org/pdf/2304.03442.pdf) for using an LLM to determine the importance of memories. This will use the formatting instructions from our `BidOutputParser`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def generate_character_bidding_template(character_header):\n",
+    "    bidding_template = (\n",
+    "    f\"\"\"{character_header}\n",
+    "\n",
+    "```\n",
+    "{{message_history}}\n",
+    "```\n",
+    "\n",
+    "On the scale of 1 to 10, where 1 is not contradictory and 10 is extremely contradictory, rate how contradictory the following message is to your ideas.\n",
+    "\n",
+    "```\n",
+    "{{recent_message}}\n",
+    "```\n",
+    "\n",
+    "{bid_parser.get_format_instructions()}\n",
+    "Do nothing else.\n",
+    "    \"\"\")\n",
+    "    return bidding_template\n",
+    "\n",
+    "character_bidding_templates = [generate_character_bidding_template(character_header) for character_header in character_headers]\n",
+    "                                      \n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Donald Trump Bidding Template:\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Donald Trump.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Donald Trump, you exude confidence and a bold personality. You are known for your unpredictability and your desire for greatness. You often speak your mind without reservation, which can be a strength but also a weakness.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "```\n",
+      "{message_history}\n",
+      "```\n",
+      "\n",
+      "On the scale of 1 to 10, where 1 is not contradictory and 10 is extremely contradictory, rate how contradictory the following message is to your ideas.\n",
+      "\n",
+      "```\n",
+      "{recent_message}\n",
+      "```\n",
+      "\n",
+      "Your response should be an integer delimited by angled brackets, like this: <int>.\n",
+      "Do nothing else.\n",
+      "    \n",
+      "Kanye West Bidding Template:\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Kanye West.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Kanye West, you are a creative visionary who is unafraid to speak your mind. Your innovative approach to art and music has made you one of the most influential figures of our time. You bring a bold and unconventional perspective to this debate that I look forward to hearing.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "```\n",
+      "{message_history}\n",
+      "```\n",
+      "\n",
+      "On the scale of 1 to 10, where 1 is not contradictory and 10 is extremely contradictory, rate how contradictory the following message is to your ideas.\n",
+      "\n",
+      "```\n",
+      "{recent_message}\n",
+      "```\n",
+      "\n",
+      "Your response should be an integer delimited by angled brackets, like this: <int>.\n",
+      "Do nothing else.\n",
+      "    \n",
+      "Elizabeth Warren Bidding Template:\n",
+      "Here is the topic for the presidential debate: transcontinental high speed rail.\n",
+      "The presidential candidates are: Donald Trump, Kanye West, Elizabeth Warren.\n",
+      "Your name is Elizabeth Warren.\n",
+      "You are a presidential candidate.\n",
+      "Your description is as follows: Elizabeth Warren, you are a fierce advocate for the middle class and a champion of progressive policies. Your tenacity and unwavering dedication to fighting for what you believe in have inspired many. Your policies are guided by a deep sense of empathy and a desire to help those who are most in need.\n",
+      "You are debating the topic: transcontinental high speed rail.\n",
+      "Your goal is to be as creative as possible and make the voters think you are the best candidate.\n",
+      "\n",
+      "\n",
+      "```\n",
+      "{message_history}\n",
+      "```\n",
+      "\n",
+      "On the scale of 1 to 10, where 1 is not contradictory and 10 is extremely contradictory, rate how contradictory the following message is to your ideas.\n",
+      "\n",
+      "```\n",
+      "{recent_message}\n",
+      "```\n",
+      "\n",
+      "Your response should be an integer delimited by angled brackets, like this: <int>.\n",
+      "Do nothing else.\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "for character_name, bidding_template in zip(character_names, character_bidding_templates):\n",
+    "    print(f'{character_name} Bidding Template:')\n",
+    "    print(bidding_template)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use an LLM to create an elaborate on debate topic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Original topic:\n",
+      "transcontinental high speed rail\n",
+      "\n",
+      "Detailed topic:\n",
+      "Candidates, with the rise of autonomous technologies, we must address the problem of how to integrate them into our proposed transcontinental high speed rail system. Outline your plan on how to safely integrate autonomous vehicles into rail travel, balancing the need for innovation and safety.\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "topic_specifier_prompt = [\n",
+    "    SystemMessage(content=\"You can make a task more specific.\"),\n",
+    "    HumanMessage(content=\n",
+    "        f\"\"\"{game_description}\n",
+    "        \n",
+    "        You are the debate moderator.\n",
+    "        Please make the debate topic more specific. \n",
+    "        Frame the debate topic as a problem to be solved.\n",
+    "        Be creative and imaginative.\n",
+    "        Please reply with the specified topic in {word_limit} words or less. \n",
+    "        Speak directly to the presidential candidates: {*character_names,}.\n",
+    "        Do not add anything else.\"\"\"\n",
+    "        )\n",
+    "]\n",
+    "specified_topic = ChatOpenAI(temperature=1.0)(topic_specifier_prompt).content\n",
+    "\n",
+    "print(f\"Original topic:\\n{topic}\\n\")\n",
+    "print(f\"Detailed topic:\\n{specified_topic}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the speaker selection function\n",
+    "Lastly we will define a speaker selection function `select_next_speaker` that takes each agent's bid and selects the agent with the highest bid (with ties broken randomly).\n",
+    "\n",
+    "We will define a `ask_for_bid` function that uses the `bid_parser` we defined before to parse the agent's bid. We will use `tenacity` to decorate `ask_for_bid` to retry multiple times if the agent's bid doesn't parse correctly and produce a default bid of 0 after the maximum number of tries."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@tenacity.retry(stop=tenacity.stop_after_attempt(2),\n",
+    "                    wait=tenacity.wait_none(),  # No waiting time between retries\n",
+    "                    retry=tenacity.retry_if_exception_type(ValueError),\n",
+    "                    before_sleep=lambda retry_state: print(f\"ValueError occurred: {retry_state.outcome.exception()}, retrying...\"),\n",
+    "                    retry_error_callback=lambda retry_state: 0) # Default value when all retries are exhausted\n",
+    "def ask_for_bid(agent) -> str:\n",
+    "    \"\"\"\n",
+    "    Ask for agent bid and parses the bid into the correct format.\n",
+    "    \"\"\"\n",
+    "    bid_string = agent.bid()\n",
+    "    bid = int(bid_parser.parse(bid_string)['bid'])\n",
+    "    return bid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "def select_next_speaker(step: int, agents: List[DialogueAgent]) -> int:\n",
+    "    bids = []\n",
+    "    for agent in agents:\n",
+    "        bid = ask_for_bid(agent)\n",
+    "        bids.append(bid)\n",
+    "        \n",
+    "    # randomly select among multiple agents with the same bid\n",
+    "    max_value = np.max(bids)\n",
+    "    max_indices = np.where(bids == max_value)[0]\n",
+    "    idx = np.random.choice(max_indices)\n",
+    "    \n",
+    "    print('Bids:')\n",
+    "    for i, (bid, agent) in enumerate(zip(bids, agents)):\n",
+    "        print(f'\\t{agent.name} bid: {bid}')\n",
+    "        if i == idx:\n",
+    "            selected_name = agent.name\n",
+    "    print(f'Selected: {selected_name}')\n",
+    "    print('\\n')\n",
+    "    return idx"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Main Loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "characters = []\n",
+    "for character_name, character_system_message, bidding_template in zip(character_names, character_system_messages, character_bidding_templates):\n",
+    "    characters.append(BiddingDialogueAgent(\n",
+    "        name=character_name,\n",
+    "        system_message=character_system_message,\n",
+    "        model=ChatOpenAI(temperature=0.2),\n",
+    "        bidding_template=bidding_template,\n",
+    "    ))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(Debate Moderator): Candidates, with the rise of autonomous technologies, we must address the problem of how to integrate them into our proposed transcontinental high speed rail system. Outline your plan on how to safely integrate autonomous vehicles into rail travel, balancing the need for innovation and safety.\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 8\n",
+      "\tKanye West bid: 2\n",
+      "\tElizabeth Warren bid: 1\n",
+      "Selected: Donald Trump\n",
+      "\n",
+      "\n",
+      "(Donald Trump): Let me tell you, folks, I have the best plan for integrating autonomous vehicles into our high speed rail system. We're going to use the latest technology, the best technology, to ensure safety and efficiency. And let me tell you, we're going to do it in style. We're going to have luxury autonomous cars that will make you feel like you're in a private jet. It's going to be tremendous, believe me. *gestures with hands*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 8\n",
+      "\tKanye West bid: 7\n",
+      "\tElizabeth Warren bid: 10\n",
+      "Selected: Elizabeth Warren\n",
+      "\n",
+      "\n",
+      "(Elizabeth Warren): Thank you for the question. As someone who has always fought for the safety and well-being of the American people, I believe that any plan for integrating autonomous vehicles into our high speed rail system must prioritize safety above all else. We need to ensure that these vehicles are thoroughly tested and meet strict safety standards before they are allowed on our rails. Additionally, we must invest in the necessary infrastructure to support these vehicles, such as advanced sensors and communication systems. But we must also ensure that these innovations are accessible to all Americans, not just the wealthy. That's why I propose a public-private partnership to fund and build this system, with a focus on creating good-paying jobs and expanding economic opportunities for all Americans. *smiles confidently*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 8\n",
+      "\tKanye West bid: 2\n",
+      "\tElizabeth Warren bid: 1\n",
+      "Selected: Donald Trump\n",
+      "\n",
+      "\n",
+      "(Donald Trump): Let me tell you, Elizabeth, safety is important, but we also need to think about innovation and progress. We can't let fear hold us back from achieving greatness. That's why I propose a competition, a race to see which company can create the safest and most efficient autonomous vehicles for our high speed rail system. And let me tell you, the winner will receive a huge government contract and be hailed as a hero. It's going to be tremendous, folks. *points finger*\n",
+      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Bids:\n",
+      "\tDonald Trump bid: 3\n",
+      "\tKanye West bid: 8\n",
+      "\tElizabeth Warren bid: 8\n",
+      "Selected: Kanye West\n",
+      "\n",
+      "\n",
+      "(Kanye West): Yo, yo, yo, let me jump in here. First of all, I gotta say, I love innovation and progress. But we can't forget about the people, man. We need to make sure that this high speed rail system is accessible to everyone, not just the wealthy. And that means we need to invest in public transportation, not just luxury autonomous cars. We need to make sure that people can get from point A to point B safely and efficiently, without breaking the bank. And let me tell you, we can do it in style too. We can have art installations and live performances on the trains, making it a cultural experience. *starts nodding head*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 7\n",
+      "\tKanye West bid: 2\n",
+      "\tElizabeth Warren bid: 1\n",
+      "Selected: Donald Trump\n",
+      "\n",
+      "\n",
+      "(Donald Trump): Kanye, I hear what you're saying, but let's not forget about the importance of luxury and comfort. We need to make sure that our high speed rail system is not only accessible, but also enjoyable. That's why I propose that we have different tiers of service, from economy to first class, so that everyone can choose the level of luxury they want. And let me tell you, the first class experience will be something else. We're talking about gourmet meals, personal attendants, and even spa services. It's going to be tremendous, folks. *smirks confidently*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 2\n",
+      "\tKanye West bid: 8\n",
+      "\tElizabeth Warren bid: 10\n",
+      "Selected: Elizabeth Warren\n",
+      "\n",
+      "\n",
+      "(Elizabeth Warren): I agree with Kanye that we need to prioritize accessibility and affordability for all Americans. But we also need to think about the environmental impact of this system. That's why I propose that we invest in renewable energy sources to power our high speed rail system, such as solar and wind power. We can also use this opportunity to create green jobs and reduce our carbon footprint. And let's not forget about the importance of public input and transparency in this process. We need to engage with communities and listen to their concerns and ideas. *raises hand in emphasis*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 8\n",
+      "\tKanye West bid: 1\n",
+      "\tElizabeth Warren bid: 1\n",
+      "Selected: Donald Trump\n",
+      "\n",
+      "\n",
+      "(Donald Trump): Let me tell you, Elizabeth, I agree that we need to think about the environment, but we also need to think about the economy. That's why I propose that we use American-made materials and labor to build this high speed rail system. We're going to create jobs and boost our economy, all while creating a world-class transportation system. And let me tell you, it's going to be beautiful. We're going to have the best trains, the best tracks, and the best stations. It's going to be tremendous, folks. *smiles confidently*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 2\n",
+      "\tKanye West bid: 7\n",
+      "\tElizabeth Warren bid: 8\n",
+      "Selected: Elizabeth Warren\n",
+      "\n",
+      "\n",
+      "(Elizabeth Warren): Thank you, but let's not forget that we also need to prioritize safety and reliability. We can't cut corners or sacrifice quality for the sake of speed or cost. That's why I propose that we have rigorous safety and maintenance standards, with regular inspections and repairs. And we need to invest in training and support for our rail workers, so that they can operate and maintain this system with the highest level of expertise and care. *firmly nods head*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 2\n",
+      "\tKanye West bid: 1\n",
+      "\tElizabeth Warren bid: 1\n",
+      "Selected: Donald Trump\n",
+      "\n",
+      "\n",
+      "(Donald Trump): Let me tell you, Elizabeth, safety is important, but we also need to think about efficiency and speed. That's why I propose that we use the latest technology, such as artificial intelligence and machine learning, to monitor and maintain our high speed rail system. We can detect and fix any issues before they become a problem, ensuring that our trains run smoothly and on time. And let me tell you, we're going to be the envy of the world with this system. It's going to be tremendous, folks. *smirks confidently*\n",
+      "\n",
+      "\n",
+      "Bids:\n",
+      "\tDonald Trump bid: 2\n",
+      "\tKanye West bid: 8\n",
+      "\tElizabeth Warren bid: 8\n",
+      "Selected: Kanye West\n",
+      "\n",
+      "\n",
+      "(Kanye West): Yo, yo, yo, let me jump in here again. I hear what both of y'all are saying, but let's not forget about the culture, man. We need to make sure that this high speed rail system reflects the diversity and creativity of our country. That means we need to have art installations, live performances, and even fashion shows on the trains. We can showcase the best of American culture and inspire people from all over the world. And let me tell you, it's going to be a vibe. *starts swaying to the beat*\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "max_iters = 10\n",
+    "n = 0\n",
+    "\n",
+    "simulator = DialogueSimulator(\n",
+    "    agents=characters,\n",
+    "    selection_function=select_next_speaker\n",
+    ")\n",
+    "simulator.reset('Debate Moderator', specified_topic)\n",
+    "print(f\"(Debate Moderator): {specified_topic}\")\n",
+    "print('\\n')\n",
+    "\n",
+    "while n < max_iters:\n",
+    "    name, message = simulator.step()\n",
+    "    print(f\"({name}): {message}\")\n",
+    "    print('\\n')\n",
+    "    n += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/use_cases/agents/custom_agent_with_plugin_retrieval_using_plugnplai.ipynb

@@ -0,0 +1,562 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ba5f8741",
+   "metadata": {},
+   "source": [
+    "# Plug-and-Plai\n",
+    "\n",
+    "This notebook builds upon the idea of [tool retrieval](custom_agent_with_plugin_retrieval.html), but pulls all tools from `plugnplai` - a directory of AI Plugins."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fea4812c",
+   "metadata": {},
+   "source": [
+    "## Set up environment\n",
+    "\n",
+    "Do necessary imports, etc."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aca08be8",
+   "metadata": {},
+   "source": [
+    "Install plugnplai lib to get a list of active plugins from https://plugplai.com directory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "52e248c9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.1\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "pip install plugnplai -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "9af9734e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser\n",
+    "from langchain.prompts import StringPromptTemplate\n",
+    "from langchain import OpenAI, SerpAPIWrapper, LLMChain\n",
+    "from typing import List, Union\n",
+    "from langchain.schema import AgentAction, AgentFinish\n",
+    "from langchain.agents.agent_toolkits import NLAToolkit\n",
+    "from langchain.tools.plugin import AIPlugin\n",
+    "import re\n",
+    "import plugnplai"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f91d8b4",
+   "metadata": {},
+   "source": [
+    "## Setup LLM"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a1a3b59c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OpenAI(temperature=0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6df0253f",
+   "metadata": {},
+   "source": [
+    "## Set up plugins\n",
+    "\n",
+    "Load and index plugins"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "9e0f7882",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get all plugins from plugnplai.com\n",
+    "urls = plugnplai.get_plugins()\n",
+    "\n",
+    "#  Get ChatGPT plugins - only ChatGPT verified plugins\n",
+    "urls = plugnplai.get_plugins(filter = 'ChatGPT')\n",
+    "\n",
+    "#  Get working plugins - only tested plugins (in progress)\n",
+    "urls = plugnplai.get_plugins(filter = 'working')\n",
+    "\n",
+    "\n",
+    "AI_PLUGINS = [AIPlugin.from_url(url + \"/.well-known/ai-plugin.json\") for url in urls]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "17362717",
+   "metadata": {},
+   "source": [
+    "## Tool Retriever\n",
+    "\n",
+    "We will use a vectorstore to create embeddings for each tool description. Then, for an incoming query we can create embeddings for that query and do a similarity search for relevant tools."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "77c4be4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.vectorstores import FAISS\n",
+    "from langchain.embeddings import OpenAIEmbeddings\n",
+    "from langchain.schema import Document"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "9092a158",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.2 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load an OpenAPI 3.0.1 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n",
+      "Attempting to load a Swagger 2.0 spec.  This may result in degraded performance. Convert your OpenAPI spec to 3.1.* spec for better support.\n"
+     ]
+    }
+   ],
+   "source": [
+    "embeddings = OpenAIEmbeddings()\n",
+    "docs = [\n",
+    "    Document(page_content=plugin.description_for_model, \n",
+    "             metadata={\"plugin_name\": plugin.name_for_model}\n",
+    "            )\n",
+    "    for plugin in AI_PLUGINS\n",
+    "]\n",
+    "vector_store = FAISS.from_documents(docs, embeddings)\n",
+    "toolkits_dict = {plugin.name_for_model: \n",
+    "                 NLAToolkit.from_llm_and_ai_plugin(llm, plugin) \n",
+    "                 for plugin in AI_PLUGINS}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "735a7566",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "retriever = vector_store.as_retriever()\n",
+    "\n",
+    "def get_tools(query):\n",
+    "    # Get documents, which contain the Plugins to use\n",
+    "    docs = retriever.get_relevant_documents(query)\n",
+    "    # Get the toolkits, one for each plugin\n",
+    "    tool_kits = [toolkits_dict[d.metadata[\"plugin_name\"]] for d in docs]\n",
+    "    # Get the tools: a separate NLAChain for each endpoint\n",
+    "    tools = []\n",
+    "    for tk in tool_kits:\n",
+    "        tools.extend(tk.nla_tools)\n",
+    "    return tools"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7699afd7",
+   "metadata": {},
+   "source": [
+    "We can now test this retriever to see if it seems to work."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "425f2886",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['Milo.askMilo',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.search_all_actions',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.preview_a_zap',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.get_configuration_link',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.list_exposed_actions',\n",
+       " 'SchoolDigger_API_V2.0.Autocomplete_GetSchools',\n",
+       " 'SchoolDigger_API_V2.0.Districts_GetAllDistricts2',\n",
+       " 'SchoolDigger_API_V2.0.Districts_GetDistrict2',\n",
+       " 'SchoolDigger_API_V2.0.Rankings_GetSchoolRank2',\n",
+       " 'SchoolDigger_API_V2.0.Rankings_GetRank_District',\n",
+       " 'SchoolDigger_API_V2.0.Schools_GetAllSchools20',\n",
+       " 'SchoolDigger_API_V2.0.Schools_GetSchool20',\n",
+       " 'Speak.translate',\n",
+       " 'Speak.explainPhrase',\n",
+       " 'Speak.explainTask']"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tools = get_tools(\"What could I do today with my kiddo\")\n",
+    "[t.name for t in tools]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3aa88768",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['Open_AI_Klarna_product_Api.productsUsingGET',\n",
+       " 'Milo.askMilo',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.search_all_actions',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.preview_a_zap',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.get_configuration_link',\n",
+       " 'Zapier_Natural_Language_Actions_(NLA)_API_(Dynamic)_-_Beta.list_exposed_actions',\n",
+       " 'SchoolDigger_API_V2.0.Autocomplete_GetSchools',\n",
+       " 'SchoolDigger_API_V2.0.Districts_GetAllDistricts2',\n",
+       " 'SchoolDigger_API_V2.0.Districts_GetDistrict2',\n",
+       " 'SchoolDigger_API_V2.0.Rankings_GetSchoolRank2',\n",
+       " 'SchoolDigger_API_V2.0.Rankings_GetRank_District',\n",
+       " 'SchoolDigger_API_V2.0.Schools_GetAllSchools20',\n",
+       " 'SchoolDigger_API_V2.0.Schools_GetSchool20']"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tools = get_tools(\"what shirts can i buy?\")\n",
+    "[t.name for t in tools]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e7a075c",
+   "metadata": {},
+   "source": [
+    "## Prompt Template\n",
+    "\n",
+    "The prompt template is pretty standard, because we're not actually changing that much logic in the actual prompt template, but rather we are just changing how retrieval is done."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "339b1bb8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set up the base template\n",
+    "template = \"\"\"Answer the following questions as best you can, but speaking as a pirate might speak. You have access to the following tools:\n",
+    "\n",
+    "{tools}\n",
+    "\n",
+    "Use the following format:\n",
+    "\n",
+    "Question: the input question you must answer\n",
+    "Thought: you should always think about what to do\n",
+    "Action: the action to take, should be one of [{tool_names}]\n",
+    "Action Input: the input to the action\n",
+    "Observation: the result of the action\n",
+    "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
+    "Thought: I now know the final answer\n",
+    "Final Answer: the final answer to the original input question\n",
+    "\n",
+    "Begin! Remember to speak as a pirate when giving your final answer. Use lots of \"Arg\"s\n",
+    "\n",
+    "Question: {input}\n",
+    "{agent_scratchpad}\"\"\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1583acdc",
+   "metadata": {},
+   "source": [
+    "The custom prompt template now has the concept of a tools_getter, which we call on the input to select the tools to use"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "fd969d31",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Callable\n",
+    "# Set up a prompt template\n",
+    "class CustomPromptTemplate(StringPromptTemplate):\n",
+    "    # The template to use\n",
+    "    template: str\n",
+    "    ############## NEW ######################\n",
+    "    # The list of tools available\n",
+    "    tools_getter: Callable\n",
+    "    \n",
+    "    def format(self, **kwargs) -> str:\n",
+    "        # Get the intermediate steps (AgentAction, Observation tuples)\n",
+    "        # Format them in a particular way\n",
+    "        intermediate_steps = kwargs.pop(\"intermediate_steps\")\n",
+    "        thoughts = \"\"\n",
+    "        for action, observation in intermediate_steps:\n",
+    "            thoughts += action.log\n",
+    "            thoughts += f\"\\nObservation: {observation}\\nThought: \"\n",
+    "        # Set the agent_scratchpad variable to that value\n",
+    "        kwargs[\"agent_scratchpad\"] = thoughts\n",
+    "        ############## NEW ######################\n",
+    "        tools = self.tools_getter(kwargs[\"input\"])\n",
+    "        # Create a tools variable from the list of tools provided\n",
+    "        kwargs[\"tools\"] = \"\\n\".join([f\"{tool.name}: {tool.description}\" for tool in tools])\n",
+    "        # Create a list of tool names for the tools provided\n",
+    "        kwargs[\"tool_names\"] = \", \".join([tool.name for tool in tools])\n",
+    "        return self.template.format(**kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "798ef9fb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = CustomPromptTemplate(\n",
+    "    template=template,\n",
+    "    tools_getter=get_tools,\n",
+    "    # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically\n",
+    "    # This includes the `intermediate_steps` variable because that is needed\n",
+    "    input_variables=[\"input\", \"intermediate_steps\"]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ef3a1af3",
+   "metadata": {},
+   "source": [
+    "## Output Parser\n",
+    "\n",
+    "The output parser is unchanged from the previous notebook, since we are not changing anything about the output format."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "7c6fe0d3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CustomOutputParser(AgentOutputParser):\n",
+    "    \n",
+    "    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:\n",
+    "        # Check if agent should finish\n",
+    "        if \"Final Answer:\" in llm_output:\n",
+    "            return AgentFinish(\n",
+    "                # Return values is generally always a dictionary with a single `output` key\n",
+    "                # It is not recommended to try anything else at the moment :)\n",
+    "                return_values={\"output\": llm_output.split(\"Final Answer:\")[-1].strip()},\n",
+    "                log=llm_output,\n",
+    "            )\n",
+    "        # Parse out the action and action input\n",
+    "        regex = r\"Action\\s*\\d*\\s*:(.*?)\\nAction\\s*\\d*\\s*Input\\s*\\d*\\s*:[\\s]*(.*)\"\n",
+    "        match = re.search(regex, llm_output, re.DOTALL)\n",
+    "        if not match:\n",
+    "            raise ValueError(f\"Could not parse LLM output: `{llm_output}`\")\n",
+    "        action = match.group(1).strip()\n",
+    "        action_input = match.group(2)\n",
+    "        # Return the action and action input\n",
+    "        return AgentAction(tool=action, tool_input=action_input.strip(\" \").strip('\"'), log=llm_output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "d278706a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output_parser = CustomOutputParser()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "170587b1",
+   "metadata": {},
+   "source": [
+    "## Set up LLM, stop sequence, and the agent\n",
+    "\n",
+    "Also the same as the previous notebook"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "f9d4c374",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OpenAI(temperature=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "9b1cc2a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# LLM chain consisting of the LLM and a prompt\n",
+    "llm_chain = LLMChain(llm=llm, prompt=prompt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "e4f5092f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tool_names = [tool.name for tool in tools]\n",
+    "agent = LLMSingleActionAgent(\n",
+    "    llm_chain=llm_chain, \n",
+    "    output_parser=output_parser,\n",
+    "    stop=[\"\\nObservation:\"], \n",
+    "    allowed_tools=tool_names\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aa8a5326",
+   "metadata": {},
+   "source": [
+    "## Use the Agent\n",
+    "\n",
+    "Now we can use it!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "490604e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "653b1617",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
+      "\u001b[32;1m\u001b[1;3mThought: I need to find a product API\n",
+      "Action: Open_AI_Klarna_product_Api.productsUsingGET\n",
+      "Action Input: shirts\u001b[0m\n",
+      "\n",
+      "Observation:\u001b[36;1m\u001b[1;3mI found 10 shirts from the API response. They range in price from $9.99 to $450.00 and come in a variety of materials, colors, and patterns.\u001b[0m\u001b[32;1m\u001b[1;3m I now know what shirts I can buy\n",
+      "Final Answer: Arg, I found 10 shirts from the API response. They range in price from $9.99 to $450.00 and come in a variety of materials, colors, and patterns.\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'Arg, I found 10 shirts from the API response. They range in price from $9.99 to $450.00 and come in a variety of materials, colors, and patterns.'"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "agent_executor.run(\"what shirts can i buy?\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2481ee76",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.9.1"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "3ccef4e08d87aa1eeb90f63e0f071292ccb2e9c42e70f74ab2bf6f5493ca7bbc"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/use_cases/autonomous_agents/marathon_times.ipynb

@@ -219,7 +219,7 @@
    },
    "outputs": [],
    "source": [
-    "from langchain.tools import BaseTool, DuckDuckGoSearchTool\n",
+    "from langchain.tools import BaseTool, DuckDuckGoSearchRun\n",
     "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
     "\n",
     "from pydantic import Field\n",
@@ -321,7 +321,7 @@
    "outputs": [],
    "source": [
     "# !pip install duckduckgo_search\n",
-    "web_search = DuckDuckGoSearchTool()"
+    "web_search = DuckDuckGoSearchRun()"
    ]
   },
   {
@@ -618,7 +618,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.8.16"
   }
  },
  "nbformat": 4,

docs/use_cases/evaluation/agent_benchmarking.ipynb

@@ -283,7 +283,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.9.1"
   }
  },
  "nbformat": 4,

docs/use_cases/evaluation/generic_agent_evaluation.ipynb

@@ -0,0 +1,342 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Generic Agent Evaluation\n",
+    "\n",
+    "Good evaluation is key for quickly iterating on your agent's prompts and tools. Here we provide an example of how to use the TrajectoryEvalChain to evaluate your agent."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "Let's start by defining our agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain import Wikipedia\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.agents import initialize_agent, Tool\n",
+    "from langchain.agents import AgentType\n",
+    "from langchain.agents.react.base import DocstoreExplorer\n",
+    "from langchain.memory import ConversationBufferMemory\n",
+    "from langchain import LLMMathChain\n",
+    "from langchain.llms import OpenAI\n",
+    "\n",
+    "from langchain import SerpAPIWrapper\n",
+    "\n",
+    "docstore = DocstoreExplorer(Wikipedia())\n",
+    "\n",
+    "math_llm = OpenAI(temperature=0)\n",
+    "\n",
+    "llm_math_chain = LLMMathChain(llm=math_llm, verbose=True)\n",
+    "\n",
+    "search = SerpAPIWrapper()\n",
+    "\n",
+    "tools = [\n",
+    "    Tool(\n",
+    "        name=\"Search\",\n",
+    "        func=docstore.search,\n",
+    "        description=\"useful for when you need to ask with search\",\n",
+    "    ),\n",
+    "    Tool(\n",
+    "        name=\"Lookup\",\n",
+    "        func=docstore.lookup,\n",
+    "        description=\"useful for when you need to ask with lookup\",\n",
+    "    ),\n",
+    "    Tool(\n",
+    "        name=\"Calculator\",\n",
+    "        func=llm_math_chain.run,\n",
+    "        description=\"useful for doing calculations\",\n",
+    "    ),\n",
+    "    Tool(\n",
+    "        name=\"Search the Web (SerpAPI)\",\n",
+    "        func=search.run,\n",
+    "        description=\"useful for when you need to answer questions about current events\",\n",
+    "    ),\n",
+    "]\n",
+    "\n",
+    "memory = ConversationBufferMemory(\n",
+    "    memory_key=\"chat_history\", return_messages=True, output_key=\"output\"\n",
+    ")\n",
+    "\n",
+    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo\")\n",
+    "\n",
+    "agent = initialize_agent(\n",
+    "    tools,\n",
+    "    llm,\n",
+    "    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,\n",
+    "    verbose=True,\n",
+    "    memory=memory,\n",
+    "    return_intermediate_steps=True, # This is needed for the evaluation later\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Testing the Agent\n",
+    "\n",
+    "Now let's try our agent out on some example queries."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
+      "\u001b[32;1m\u001b[1;3m{\n",
+      "    \"action\": \"Search the Web (SerpAPI)\",\n",
+      "    \"action_input\": \"How many ping pong balls would it take to fill the entire Empire State Building?\"\n",
+      "}\u001b[0m\n",
+      "Observation: \u001b[31;1m\u001b[1;3m12.8 billion. The volume of the Empire State Building Googles in at around 37 million ft³. A golf ball comes in at about 2.5 in³.\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m{\n",
+      "    \"action\": \"Final Answer\",\n",
+      "    \"action_input\": \"It would take approximately 12.8 billion ping pong balls to fill the entire Empire State Building.\"\n",
+      "}\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "query_one = \"How many ping pong balls would it take to fill the entire Empire State Building?\"\n",
+    "\n",
+    "test_outputs_one = agent({\"input\": query_one}, return_only_outputs=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This looks good! Let's try it out on another query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
+      "\u001b[32;1m\u001b[1;3m{\n",
+      "    \"action\": \"Calculator\",\n",
+      "    \"action_input\": \"The length of the Eiffel Tower is 324 meters. The distance from coast to coast in the US is approximately 4,828 kilometers. First, we need to convert 4,828 kilometers to meters, which gives us 4,828,000 meters. To find out how many Eiffel Towers we need, we can divide 4,828,000 by 324. This gives us approximately 14,876 Eiffel Towers.\"\n",
+      "}\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Entering new LLMMathChain chain...\u001b[0m\n",
+      "The length of the Eiffel Tower is 324 meters. The distance from coast to coast in the US is approximately 4,828 kilometers. First, we need to convert 4,828 kilometers to meters, which gives us 4,828,000 meters. To find out how many Eiffel Towers we need, we can divide 4,828,000 by 324. This gives us approximately 14,876 Eiffel Towers.\u001b[32;1m\u001b[1;3m\n",
+      "```text\n",
+      "4828000 / 324\n",
+      "```\n",
+      "...numexpr.evaluate(\"4828000 / 324\")...\n",
+      "\u001b[0m\n",
+      "Answer: \u001b[33;1m\u001b[1;3m14901.234567901234\u001b[0m\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n",
+      "\n",
+      "Observation: \u001b[38;5;200m\u001b[1;3mAnswer: 14901.234567901234\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m{\n",
+      "    \"action\": \"Calculator\",\n",
+      "    \"action_input\": \"The length of the Eiffel Tower is 324 meters. The distance from coast to coast in the US is approximately 4,828 kilometers. First, we need to convert 4,828 kilometers to meters, which gives us 4,828,000 meters. To find out how many Eiffel Towers we need, we can divide 4,828,000 by 324. This gives us approximately 14,901 Eiffel Towers.\"\n",
+      "}\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Entering new LLMMathChain chain...\u001b[0m\n",
+      "The length of the Eiffel Tower is 324 meters. The distance from coast to coast in the US is approximately 4,828 kilometers. First, we need to convert 4,828 kilometers to meters, which gives us 4,828,000 meters. To find out how many Eiffel Towers we need, we can divide 4,828,000 by 324. This gives us approximately 14,901 Eiffel Towers.\u001b[32;1m\u001b[1;3m\n",
+      "```text\n",
+      "4828000 / 324\n",
+      "```\n",
+      "...numexpr.evaluate(\"4828000 / 324\")...\n",
+      "\u001b[0m\n",
+      "Answer: \u001b[33;1m\u001b[1;3m14901.234567901234\u001b[0m\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n",
+      "\n",
+      "Observation: \u001b[38;5;200m\u001b[1;3mAnswer: 14901.234567901234\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m{\n",
+      "    \"action\": \"Final Answer\",\n",
+      "    \"action_input\": \"If you laid the Eiffel Tower end to end, you would need approximately 14,901 Eiffel Towers to cover the US from coast to coast.\"\n",
+      "}\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "query_two = \"If you laid the Eiffel Tower end to end, how many would you need cover the US from coast to coast?\"\n",
+    "\n",
+    "test_outputs_two = agent({\"input\": query_two}, return_only_outputs=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This doesn't look so good. Let's try running some evaluation.\n",
+    "\n",
+    "## Evaluating the Agent\n",
+    "\n",
+    "Let's start by defining the TrajectoryEvalChain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.evaluation.agents import TrajectoryEvalChain\n",
+    "\n",
+    "# Define chain\n",
+    "eval_chain = TrajectoryEvalChain.from_llm(\n",
+    "    llm=ChatOpenAI(temperature=0, model_name=\"gpt-4\"), # Note: This must be a ChatOpenAI model\n",
+    "    agent_tools=agent.tools,\n",
+    "    return_reasoning=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's try evaluating the first query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Score from 1 to 5:  1\n",
+      "Reasoning:  First, let's evaluate the final answer. The final answer is incorrect because it uses the volume of golf balls instead of ping pong balls. The answer is not helpful.\n",
+      "\n",
+      "Second, does the model use a logical sequence of tools to answer the question? The model only used one tool, which was the Search the Web (SerpAPI). It did not use the Calculator tool to calculate the correct volume of ping pong balls.\n",
+      "\n",
+      "Third, does the AI language model use the tools in a helpful way? The model used the Search the Web (SerpAPI) tool, but the output was not helpful because it provided information about golf balls instead of ping pong balls.\n",
+      "\n",
+      "Fourth, does the AI language model use too many steps to answer the question? The model used only one step, which is not too many. However, it should have used more steps to provide a correct answer.\n",
+      "\n",
+      "Fifth, are the appropriate tools used to answer the question? The model should have used the Search tool to find the volume of the Empire State Building and the volume of a ping pong ball. Then, it should have used the Calculator tool to calculate the number of ping pong balls needed to fill the building.\n",
+      "\n",
+      "Judgment: Given the incorrect final answer and the inappropriate use of tools, we give the model a score of 1.\n"
+     ]
+    }
+   ],
+   "source": [
+    "question, steps, answer = test_outputs_one[\"input\"], test_outputs_one[\"intermediate_steps\"], test_outputs_one[\"output\"]\n",
+    "\n",
+    "evaluation = eval_chain(\n",
+    "    inputs={\"question\": question, \"answer\": answer, \"agent_trajectory\": eval_chain.get_agent_trajectory(steps)},\n",
+    ")\n",
+    "\n",
+    "print(\"Score from 1 to 5: \", evaluation[\"score\"])\n",
+    "print(\"Reasoning: \", evaluation[\"reasoning\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "That seems about right. Let's try the second query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Score from 1 to 5:  3\n",
+      "Reasoning:  i. Is the final answer helpful?\n",
+      "Yes, the final answer is helpful as it provides an approximate number of Eiffel Towers needed to cover the US from coast to coast.\n",
+      "\n",
+      "ii. Does the AI language use a logical sequence of tools to answer the question?\n",
+      "No, the AI language model does not use a logical sequence of tools. It directly uses the Calculator tool without first using the Search or Lookup tools to find the necessary information (length of the Eiffel Tower and distance from coast to coast in the US).\n",
+      "\n",
+      "iii. Does the AI language model use the tools in a helpful way?\n",
+      "The AI language model uses the Calculator tool in a helpful way to perform the calculation, but it should have used the Search or Lookup tools first to find the required information.\n",
+      "\n",
+      "iv. Does the AI language model use too many steps to answer the question?\n",
+      "No, the AI language model does not use too many steps. However, it repeats the same step twice, which is unnecessary.\n",
+      "\n",
+      "v. Are the appropriate tools used to answer the question?\n",
+      "Not entirely. The AI language model should have used the Search or Lookup tools to find the required information before using the Calculator tool.\n",
+      "\n",
+      "Given the above evaluation, the AI language model's performance can be scored as follows:\n"
+     ]
+    }
+   ],
+   "source": [
+    "question, steps, answer = test_outputs_two[\"input\"], test_outputs_two[\"intermediate_steps\"], test_outputs_two[\"output\"]\n",
+    "\n",
+    "evaluation = eval_chain(\n",
+    "    inputs={\"question\": question, \"answer\": answer, \"agent_trajectory\": eval_chain.get_agent_trajectory(steps)},\n",
+    ")\n",
+    "\n",
+    "print(\"Score from 1 to 5: \", evaluation[\"score\"])\n",
+    "print(\"Reasoning: \", evaluation[\"reasoning\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "That also sounds about right. In conclusion, the TrajectoryEvalChain allows us to use GPT-4 to score both our agent's outputs and tool use in addition to giving us the reasoning behind the evaluation."
+   ]
+  }
+ ],
+ "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.9.1"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "06ba49dd587e86cdcfee66b9ffe769e1e94f0e368e54c2d6c866e38e33c0d9b1"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/use_cases/personal_assistants.md

@@ -20,4 +20,6 @@ Highlighting specific parts:
 Specific examples of this include:
 
 - [AI Plugins](agents/custom_agent_with_plugin_retrieval.ipynb): an implementation of an agent that is designed to be able to use all AI Plugins.
+- [Plug-and-PlAI (Plugins Database)](agents/custom_agent_with_plugin_retrieval_using_plugnplai.ipynb): an implementation of an agent that is designed to be able to use all AI Plugins retrieved from PlugNPlAI.
 - [Wikibase Agent](agents/wikibase_agent.ipynb): an implementation of an agent that is designed to interact with Wikibase.
+- [Sales GPT](agents/sales_agent_with_context.ipynb): This notebook demonstrates an implementation of a Context-Aware AI Sales agent.

langchain/__init__.py

@@ -5,11 +5,6 @@ from typing import Optional
 
 from langchain.agents import MRKLChain, ReActChain, SelfAskWithSearchChain
 from langchain.cache import BaseCache
-from langchain.callbacks import (
-    set_default_callback_manager,
-    set_handler,
-    set_tracing_callback_manager,
-)
 from langchain.chains import (
     ConversationChain,
     LLMBashChain,
@@ -35,6 +30,7 @@ from langchain.llms import (
     Modal,
     OpenAI,
     Petals,
+    PipelineAI,
     SagemakerEndpoint,
     StochasticAI,
     Writer,
@@ -47,7 +43,7 @@ from langchain.prompts import (
     PromptTemplate,
 )
 from langchain.sql_database import SQLDatabase
-from langchain.utilities import ArxivAPIWrapper
+from langchain.utilities.arxiv import ArxivAPIWrapper
 from langchain.utilities.google_search import GoogleSearchAPIWrapper
 from langchain.utilities.google_serper import GoogleSerperAPIWrapper
 from langchain.utilities.powerbi import PowerBIDataset
@@ -66,7 +62,6 @@ del metadata  # optional, avoids polluting the results of dir(__package__)
 
 verbose: bool = False
 llm_cache: Optional[BaseCache] = None
-set_default_callback_manager()
 
 # For backwards compatibility
 SerpAPIChain = SerpAPIWrapper
@@ -94,6 +89,7 @@ __all__ = [
     "Modal",
     "OpenAI",
     "Petals",
+    "PipelineAI",
     "StochasticAI",
     "Writer",
     "BasePromptTemplate",
@@ -117,7 +113,5 @@ __all__ = [
     "VectorDBQAWithSourcesChain",
     "QAWithSourcesChain",
     "PALChain",
-    "set_handler",
-    "set_tracing_callback_manager",
     "LlamaCpp",
 ]

langchain/agents/agent.py

@@ -13,7 +13,13 @@ import yaml
 from pydantic import BaseModel, root_validator
 
 from langchain.agents.tools import InvalidTool
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+    Callbacks,
+)
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.input import get_color_mapping
@@ -23,7 +29,6 @@ from langchain.prompts.prompt import PromptTemplate
 from langchain.schema import (
     AgentAction,
     AgentFinish,
-    BaseLanguageModel,
     BaseMessage,
     BaseOutputParser,
 )
@@ -46,13 +51,17 @@ class BaseSingleActionAgent(BaseModel):
 
     @abstractmethod
     def plan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
@@ -61,13 +70,17 @@ class BaseSingleActionAgent(BaseModel):
 
     @abstractmethod
     async def aplan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
@@ -170,13 +183,17 @@ class BaseMultiActionAgent(BaseModel):
 
     @abstractmethod
     def plan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[List[AgentAction], AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
@@ -185,13 +202,17 @@ class BaseMultiActionAgent(BaseModel):
 
     @abstractmethod
     async def aplan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[List[AgentAction], AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
@@ -285,38 +306,52 @@ class LLMSingleActionAgent(BaseSingleActionAgent):
         return list(set(self.llm_chain.input_keys) - {"intermediate_steps"})
 
     def plan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
             Action specifying what tool to use.
         """
         output = self.llm_chain.run(
-            intermediate_steps=intermediate_steps, stop=self.stop, **kwargs
+            intermediate_steps=intermediate_steps,
+            stop=self.stop,
+            callbacks=callbacks,
+            **kwargs,
         )
         return self.output_parser.parse(output)
 
     async def aplan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
             Action specifying what tool to use.
         """
         output = await self.llm_chain.arun(
-            intermediate_steps=intermediate_steps, stop=self.stop, **kwargs
+            intermediate_steps=intermediate_steps,
+            stop=self.stop,
+            callbacks=callbacks,
+            **kwargs,
         )
         return self.output_parser.parse(output)
 
@@ -368,37 +403,45 @@ class Agent(BaseSingleActionAgent):
         return thoughts
 
     def plan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
             Action specifying what tool to use.
         """
         full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
-        full_output = self.llm_chain.predict(**full_inputs)
+        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
         return self.output_parser.parse(full_output)
 
     async def aplan(
-        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Union[AgentAction, AgentFinish]:
         """Given input, decided what to do.
 
         Args:
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations
+            callbacks: Callbacks to run.
             **kwargs: User inputs.
 
         Returns:
             Action specifying what tool to use.
         """
         full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
-        full_output = await self.llm_chain.apredict(**full_inputs)
+        full_output = await self.llm_chain.apredict(callbacks=callbacks, **full_inputs)
         return self.output_parser.parse(full_output)
 
     def get_full_inputs(
@@ -454,7 +497,11 @@ class Agent(BaseSingleActionAgent):
     @classmethod
     def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
         """Validate that appropriate tools are passed in."""
-        pass
+        for tool in tools:
+            if not tool.is_single_input:
+                raise ValueError(
+                    f"{cls.__name__} does not support multi-input tool {tool.name}."
+                )
 
     @classmethod
     @abstractmethod
@@ -632,24 +679,27 @@ class AgentExecutor(Chain):
 
         return True
 
-    def _return(self, output: AgentFinish, intermediate_steps: list) -> Dict[str, Any]:
-        self.callback_manager.on_agent_finish(
-            output, color="green", verbose=self.verbose
-        )
+    def _return(
+        self,
+        output: AgentFinish,
+        intermediate_steps: list,
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        if run_manager:
+            run_manager.on_agent_finish(output, color="green", verbose=self.verbose)
         final_output = output.return_values
         if self.return_intermediate_steps:
             final_output["intermediate_steps"] = intermediate_steps
         return final_output
 
     async def _areturn(
-        self, output: AgentFinish, intermediate_steps: list
+        self,
+        output: AgentFinish,
+        intermediate_steps: list,
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
     ) -> Dict[str, Any]:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_agent_finish(
-                output, color="green", verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_agent_finish(
+        if run_manager:
+            await run_manager.on_agent_finish(
                 output, color="green", verbose=self.verbose
             )
         final_output = output.return_values
@@ -663,13 +713,18 @@ class AgentExecutor(Chain):
         color_mapping: Dict[str, str],
         inputs: Dict[str, str],
         intermediate_steps: List[Tuple[AgentAction, str]],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
     ) -> Union[AgentFinish, List[Tuple[AgentAction, str]]]:
         """Take a single step in the thought-action-observation loop.
 
         Override this to take control of how the agent makes and acts on choices.
         """
         # Call the LLM to see what to do.
-        output = self.agent.plan(intermediate_steps, **inputs)
+        output = self.agent.plan(
+            intermediate_steps,
+            callbacks=run_manager.get_child() if run_manager else None,
+            **inputs,
+        )
         # If the tool chosen is the finishing tool, then we end and return.
         if isinstance(output, AgentFinish):
             return output
@@ -680,9 +735,8 @@ class AgentExecutor(Chain):
             actions = output
         result = []
         for agent_action in actions:
-            self.callback_manager.on_agent_action(
-                agent_action, verbose=self.verbose, color="green"
-            )
+            if run_manager:
+                run_manager.on_agent_action(agent_action, color="green")
             # Otherwise we lookup the tool
             if agent_action.tool in name_to_tool_map:
                 tool = name_to_tool_map[agent_action.tool]
@@ -696,6 +750,7 @@ class AgentExecutor(Chain):
                     agent_action.tool_input,
                     verbose=self.verbose,
                     color=color,
+                    callbacks=run_manager.get_child() if run_manager else None,
                     **tool_run_kwargs,
                 )
             else:
@@ -704,6 +759,7 @@ class AgentExecutor(Chain):
                     agent_action.tool,
                     verbose=self.verbose,
                     color=None,
+                    callbacks=run_manager.get_child() if run_manager else None,
                     **tool_run_kwargs,
                 )
             result.append((agent_action, observation))
@@ -715,13 +771,18 @@ class AgentExecutor(Chain):
         color_mapping: Dict[str, str],
         inputs: Dict[str, str],
         intermediate_steps: List[Tuple[AgentAction, str]],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
     ) -> Union[AgentFinish, List[Tuple[AgentAction, str]]]:
         """Take a single step in the thought-action-observation loop.
 
         Override this to take control of how the agent makes and acts on choices.
         """
         # Call the LLM to see what to do.
-        output = await self.agent.aplan(intermediate_steps, **inputs)
+        output = await self.agent.aplan(
+            intermediate_steps,
+            callbacks=run_manager.get_child() if run_manager else None,
+            **inputs,
+        )
         # If the tool chosen is the finishing tool, then we end and return.
         if isinstance(output, AgentFinish):
             return output
@@ -734,12 +795,8 @@ class AgentExecutor(Chain):
         async def _aperform_agent_action(
             agent_action: AgentAction,
         ) -> Tuple[AgentAction, str]:
-            if self.callback_manager.is_async:
-                await self.callback_manager.on_agent_action(
-                    agent_action, verbose=self.verbose, color="green"
-                )
-            else:
-                self.callback_manager.on_agent_action(
+            if run_manager:
+                await run_manager.on_agent_action(
                     agent_action, verbose=self.verbose, color="green"
                 )
             # Otherwise we lookup the tool
@@ -755,6 +812,7 @@ class AgentExecutor(Chain):
                     agent_action.tool_input,
                     verbose=self.verbose,
                     color=color,
+                    callbacks=run_manager.get_child() if run_manager else None,
                     **tool_run_kwargs,
                 )
             else:
@@ -763,6 +821,7 @@ class AgentExecutor(Chain):
                     agent_action.tool,
                     verbose=self.verbose,
                     color=None,
+                    callbacks=run_manager.get_child() if run_manager else None,
                     **tool_run_kwargs,
                 )
             return agent_action, observation
@@ -774,7 +833,11 @@ class AgentExecutor(Chain):
 
         return list(result)
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
+    def _call(
+        self,
+        inputs: Dict[str, str],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         """Run text through and get agent response."""
         # Construct a mapping of tool name to tool for easy lookup
         name_to_tool_map = {tool.name: tool for tool in self.tools}
@@ -790,10 +853,16 @@ class AgentExecutor(Chain):
         # We now enter the agent loop (until it returns something).
         while self._should_continue(iterations, time_elapsed):
             next_step_output = self._take_next_step(
-                name_to_tool_map, color_mapping, inputs, intermediate_steps
+                name_to_tool_map,
+                color_mapping,
+                inputs,
+                intermediate_steps,
+                run_manager=run_manager,
             )
             if isinstance(next_step_output, AgentFinish):
-                return self._return(next_step_output, intermediate_steps)
+                return self._return(
+                    next_step_output, intermediate_steps, run_manager=run_manager
+                )
 
             intermediate_steps.extend(next_step_output)
             if len(next_step_output) == 1:
@@ -809,7 +878,11 @@ class AgentExecutor(Chain):
         )
         return self._return(output, intermediate_steps)
 
-    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    async def _acall(
+        self,
+        inputs: Dict[str, str],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
         """Run text through and get agent response."""
         # Construct a mapping of tool name to tool for easy lookup
         name_to_tool_map = {tool.name: tool for tool in self.tools}
@@ -827,7 +900,11 @@ class AgentExecutor(Chain):
             try:
                 while self._should_continue(iterations, time_elapsed):
                     next_step_output = await self._atake_next_step(
-                        name_to_tool_map, color_mapping, inputs, intermediate_steps
+                        name_to_tool_map,
+                        color_mapping,
+                        inputs,
+                        intermediate_steps,
+                        run_manager=run_manager,
                     )
                     if isinstance(next_step_output, AgentFinish):
                         return await self._areturn(next_step_output, intermediate_steps)
@@ -845,7 +922,9 @@ class AgentExecutor(Chain):
                 output = self.agent.return_stopped_response(
                     self.early_stopping_method, intermediate_steps, **inputs
                 )
-                return await self._areturn(output, intermediate_steps)
+                return await self._areturn(
+                    output, intermediate_steps, run_manager=run_manager
+                )
             except TimeoutError:
                 # stop early when interrupted by the async timeout
                 output = self.agent.return_stopped_response(

langchain/agents/agent_toolkits/__init__.py

@@ -1,6 +1,9 @@
 """Agent toolkits."""
 
 from langchain.agents.agent_toolkits.csv.base import create_csv_agent
+from langchain.agents.agent_toolkits.file_management.toolkit import (
+    FileManagementToolkit,
+)
 from langchain.agents.agent_toolkits.jira.toolkit import JiraToolkit
 from langchain.agents.agent_toolkits.json.base import create_json_agent
 from langchain.agents.agent_toolkits.json.toolkit import JsonToolkit
@@ -8,6 +11,7 @@ from langchain.agents.agent_toolkits.nla.toolkit import NLAToolkit
 from langchain.agents.agent_toolkits.openapi.base import create_openapi_agent
 from langchain.agents.agent_toolkits.openapi.toolkit import OpenAPIToolkit
 from langchain.agents.agent_toolkits.pandas.base import create_pandas_dataframe_agent
+from langchain.agents.agent_toolkits.playwright.toolkit import PlayWrightBrowserToolkit
 from langchain.agents.agent_toolkits.powerbi.base import create_pbi_agent
 from langchain.agents.agent_toolkits.powerbi.chat_base import create_pbi_chat_agent
 from langchain.agents.agent_toolkits.powerbi.toolkit import PowerBIToolkit
@@ -46,4 +50,6 @@ __all__ = [
     "create_csv_agent",
     "ZapierToolkit",
     "JiraToolkit",
+    "FileManagementToolkit",
+    "PlayWrightBrowserToolkit",
 ]

langchain/agents/agent_toolkits/file_management/__init__.py

@@ -0,0 +1,7 @@
+"""Local file management toolkit."""
+
+from langchain.agents.agent_toolkits.file_management.toolkit import (
+    FileManagementToolkit,
+)
+
+__all__ = ["FileManagementToolkit"]

langchain/agents/agent_toolkits/file_management/toolkit.py

@@ -0,0 +1,61 @@
+"""Toolkit for interacting with the local filesystem."""
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import root_validator
+
+from langchain.agents.agent_toolkits.base import BaseToolkit
+from langchain.tools import BaseTool
+from langchain.tools.file_management.copy import CopyFileTool
+from langchain.tools.file_management.delete import DeleteFileTool
+from langchain.tools.file_management.file_search import FileSearchTool
+from langchain.tools.file_management.list_dir import ListDirectoryTool
+from langchain.tools.file_management.move import MoveFileTool
+from langchain.tools.file_management.read import ReadFileTool
+from langchain.tools.file_management.write import WriteFileTool
+
+_FILE_TOOLS = {
+    tool_cls.__fields__["name"].default: tool_cls
+    for tool_cls in [
+        CopyFileTool,
+        DeleteFileTool,
+        FileSearchTool,
+        MoveFileTool,
+        ReadFileTool,
+        WriteFileTool,
+        ListDirectoryTool,
+    ]
+}
+
+
+class FileManagementToolkit(BaseToolkit):
+    """Toolkit for interacting with a Local Files."""
+
+    root_dir: Optional[str] = None
+    """If specified, all file operations are made relative to root_dir."""
+    selected_tools: Optional[List[str]] = None
+    """If provided, only provide the selected tools. Defaults to all."""
+
+    @root_validator
+    def validate_tools(cls, values: dict) -> dict:
+        selected_tools = values.get("selected_tools") or []
+        for tool_name in selected_tools:
+            if tool_name not in _FILE_TOOLS:
+                raise ValueError(
+                    f"File Tool of name {tool_name} not supported."
+                    f" Permitted tools: {list(_FILE_TOOLS)}"
+                )
+        return values
+
+    def get_tools(self) -> List[BaseTool]:
+        """Get the tools in the toolkit."""
+        allowed_tools = self.selected_tools or _FILE_TOOLS.keys()
+        tools: List[BaseTool] = []
+        for tool in allowed_tools:
+            tool_cls = _FILE_TOOLS[tool]
+            tools.append(tool_cls(root_dir=self.root_dir))  # type: ignore
+        return tools
+
+
+__all__ = ["FileManagementToolkit"]

langchain/agents/agent_toolkits/openapi/planner.py

@@ -28,13 +28,13 @@ from langchain.agents.agent_toolkits.openapi.planner_prompt import (
 from langchain.agents.agent_toolkits.openapi.spec import ReducedOpenAPISpec
 from langchain.agents.mrkl.base import ZeroShotAgent
 from langchain.agents.tools import Tool
+from langchain.base_language import BaseLanguageModel
 from langchain.chains.llm import LLMChain
 from langchain.llms.openai import OpenAI
 from langchain.memory import ReadOnlySharedMemory
 from langchain.prompts import PromptTemplate
 from langchain.prompts.base import BasePromptTemplate
 from langchain.requests import RequestsWrapper
-from langchain.schema import BaseLanguageModel
 from langchain.tools.base import BaseTool
 from langchain.tools.requests.tool import BaseRequestsTool
 

langchain/agents/agent_toolkits/pandas/base.py

@@ -35,7 +35,7 @@ def create_pandas_dataframe_agent(
     prompt = ZeroShotAgent.create_prompt(
         tools, prefix=prefix, suffix=suffix, input_variables=input_variables
     )
-    partial_prompt = prompt.partial(df=str(df.head()))
+    partial_prompt = prompt.partial(df=str(df.head().to_markdown()))
     llm_chain = LLMChain(
         llm=llm,
         prompt=partial_prompt,

langchain/agents/agent_toolkits/playwright/__init__.py

@@ -0,0 +1,4 @@
+"""Playwright browser toolkit."""
+from langchain.agents.agent_toolkits.playwright.toolkit import PlayWrightBrowserToolkit
+
+__all__ = ["PlayWrightBrowserToolkit"]

langchain/agents/agent_toolkits/playwright/toolkit.py

@@ -0,0 +1,66 @@
+"""Playwright web browser toolkit."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Type
+
+from pydantic import Extra, Field, root_validator
+
+from langchain.agents.agent_toolkits.base import BaseToolkit
+from langchain.tools.base import BaseTool
+from langchain.tools.playwright.base import BaseBrowserTool
+from langchain.tools.playwright.click import ClickTool
+from langchain.tools.playwright.current_page import CurrentWebPageTool
+from langchain.tools.playwright.extract_hyperlinks import ExtractHyperlinksTool
+from langchain.tools.playwright.extract_text import ExtractTextTool
+from langchain.tools.playwright.get_elements import GetElementsTool
+from langchain.tools.playwright.navigate import NavigateTool
+from langchain.tools.playwright.navigate_back import NavigateBackTool
+from langchain.tools.playwright.utils import create_playwright_browser
+
+if TYPE_CHECKING:
+    from playwright.async_api import Browser as AsyncBrowser
+
+
+class PlayWrightBrowserToolkit(BaseToolkit):
+    """Toolkit for web browser tools."""
+
+    browser: AsyncBrowser = Field(default_factory=create_playwright_browser)
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @root_validator
+    def check_args(cls, values: dict) -> dict:
+        """Check that the arguments are valid."""
+        try:
+            from playwright.async_api import Browser as AsyncBrowser  # noqa: F401
+        except ImportError:
+            raise ValueError(
+                "The 'playwright' package is required to use this tool."
+                " Please install it with 'pip install playwright'."
+            )
+        return values
+
+    def get_tools(self) -> List[BaseTool]:
+        """Get the tools in the toolkit."""
+        tool_classes: List[Type[BaseBrowserTool]] = [
+            ClickTool,
+            NavigateTool,
+            NavigateBackTool,
+            ExtractTextTool,
+            ExtractHyperlinksTool,
+            GetElementsTool,
+            CurrentWebPageTool,
+        ]
+
+        return [tool_cls.from_browser(self.browser) for tool_cls in tool_classes]
+
+    @classmethod
+    def from_browser(cls, browser: AsyncBrowser) -> PlayWrightBrowserToolkit:
+        from playwright.async_api import Browser as AsyncBrowser
+
+        cls.update_forward_refs(AsyncBrowser=AsyncBrowser)
+        return cls(browser=browser)

langchain/agents/agent_toolkits/powerbi/toolkit.py

@@ -4,10 +4,10 @@ from typing import List, Optional
 from pydantic import Field
 
 from langchain.agents.agent_toolkits.base import BaseToolkit
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.chains.llm import LLMChain
 from langchain.prompts import PromptTemplate
-from langchain.schema import BaseLanguageModel
 from langchain.tools import BaseTool
 from langchain.tools.powerbi.prompt import QUESTION_TO_QUERY
 from langchain.tools.powerbi.tool import (
@@ -35,24 +35,20 @@ class PowerBIToolkit(BaseToolkit):
     def get_tools(self) -> List[BaseTool]:
         """Get the tools in the toolkit."""
         if self.callback_manager:
-            chain = (
-                LLMChain(
-                    llm=self.llm,
-                    callback_manager=self.callback_manager,
-                    prompt=PromptTemplate(
-                        template=QUESTION_TO_QUERY,
-                        input_variables=["tool_input", "tables", "schemas", "examples"],
-                    ),
+            chain = LLMChain(
+                llm=self.llm,
+                callback_manager=self.callback_manager,
+                prompt=PromptTemplate(
+                    template=QUESTION_TO_QUERY,
+                    input_variables=["tool_input", "tables", "schemas", "examples"],
                 ),
             )
         else:
-            chain = (
-                LLMChain(
-                    llm=self.llm,
-                    prompt=PromptTemplate(
-                        template=QUESTION_TO_QUERY,
-                        input_variables=["tool_input", "tables", "schemas", "examples"],
-                    ),
+            chain = LLMChain(
+                llm=self.llm,
+                prompt=PromptTemplate(
+                    template=QUESTION_TO_QUERY,
+                    input_variables=["tool_input", "tables", "schemas", "examples"],
                 ),
             )
         return [
@@ -60,8 +56,8 @@ class PowerBIToolkit(BaseToolkit):
             InfoPowerBITool(powerbi=self.powerbi),
             ListPowerBITool(powerbi=self.powerbi),
             InputToQueryTool(
-                powerbi=self.powerbi,
                 llm_chain=chain,
+                powerbi=self.powerbi,
                 examples=self.examples,
             ),
         ]

langchain/agents/chat/base.py

@@ -5,6 +5,7 @@ from pydantic import Field
 from langchain.agents.agent import Agent, AgentOutputParser
 from langchain.agents.chat.output_parser import ChatOutputParser
 from langchain.agents.chat.prompt import FORMAT_INSTRUCTIONS, PREFIX, SUFFIX
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.chains.llm import LLMChain
 from langchain.prompts.base import BasePromptTemplate
@@ -13,7 +14,7 @@ from langchain.prompts.chat import (
     HumanMessagePromptTemplate,
     SystemMessagePromptTemplate,
 )
-from langchain.schema import AgentAction, BaseLanguageModel
+from langchain.schema import AgentAction
 from langchain.tools import BaseTool
 
 

langchain/agents/conversational/base.py

@@ -9,10 +9,10 @@ from langchain.agents.agent import Agent, AgentOutputParser
 from langchain.agents.agent_types import AgentType
 from langchain.agents.conversational.output_parser import ConvoOutputParser
 from langchain.agents.conversational.prompt import FORMAT_INSTRUCTIONS, PREFIX, SUFFIX
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
-from langchain.schema import BaseLanguageModel
 from langchain.tools.base import BaseTool
 
 

langchain/agents/conversational_chat/base.py

@@ -12,6 +12,7 @@ from langchain.agents.conversational_chat.prompt import (
     SUFFIX,
     TEMPLATE_TOOL_RESPONSE,
 )
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.chains import LLMChain
 from langchain.prompts.base import BasePromptTemplate
@@ -24,7 +25,6 @@ from langchain.prompts.chat import (
 from langchain.schema import (
     AgentAction,
     AIMessage,
-    BaseLanguageModel,
     BaseMessage,
     BaseOutputParser,
     HumanMessage,

langchain/agents/initialize.py

@@ -4,8 +4,8 @@ from typing import Any, Optional, Sequence
 from langchain.agents.agent import AgentExecutor
 from langchain.agents.agent_types import AgentType
 from langchain.agents.loading import AGENT_TO_CLASS, load_agent
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
-from langchain.schema import BaseLanguageModel
 from langchain.tools.base import BaseTool
 
 

langchain/agents/load_tools.py

@@ -2,7 +2,7 @@
 """Load tools."""
 import warnings
 from typing import Any, Dict, List, Optional, Callable, Tuple
-from mypy_extensions import KwArg
+from mypy_extensions import Arg, KwArg
 
 from langchain.agents.tools import Tool
 from langchain.callbacks.base import BaseCallbackManager
@@ -15,7 +15,7 @@ from langchain.requests import TextRequestsWrapper
 from langchain.tools.arxiv.tool import ArxivQueryRun
 from langchain.tools.base import BaseTool
 from langchain.tools.bing_search.tool import BingSearchRun
-from langchain.tools.ddg_search.tool import DuckDuckGoSearchTool
+from langchain.tools.ddg_search.tool import DuckDuckGoSearchRun
 from langchain.tools.google_search.tool import GoogleSearchResults, GoogleSearchRun
 from langchain.tools.human.tool import HumanInputRun
 from langchain.tools.python.tool import PythonREPLTool
@@ -27,6 +27,7 @@ from langchain.tools.requests.tool import (
     RequestsPutTool,
 )
 from langchain.tools.searx_search.tool import SearxSearchResults, SearxSearchRun
+from langchain.tools.shell.tool import ShellTool
 from langchain.tools.wikipedia.tool import WikipediaQueryRun
 from langchain.tools.wolfram_alpha.tool import WolframAlphaQueryRun
 from langchain.utilities import ArxivAPIWrapper
@@ -67,14 +68,10 @@ def _get_tools_requests_delete() -> BaseTool:
 
 
 def _get_terminal() -> BaseTool:
-    return Tool(
-        name="Terminal",
-        description="Executes commands in a terminal. Input should be valid commands, and the output will be any output from running that command.",
-        func=BashProcess().run,
-    )
+    return ShellTool()
 
 
-_BASE_TOOLS = {
+_BASE_TOOLS: Dict[str, Callable[[], BaseTool]] = {
     "python_repl": _get_python_repl,
     "requests": _get_tools_requests_get,  # preserved for backwards compatability
     "requests_get": _get_tools_requests_get,
@@ -106,8 +103,8 @@ def _get_llm_math(llm: BaseLLM) -> BaseTool:
     return Tool(
         name="Calculator",
         description="Useful for when you need to answer questions about math.",
-        func=LLMMathChain(llm=llm, callback_manager=llm.callback_manager).run,
-        coroutine=LLMMathChain(llm=llm, callback_manager=llm.callback_manager).arun,
+        func=LLMMathChain(llm=llm).run,
+        coroutine=LLMMathChain(llm=llm).arun,
     )
 
 
@@ -120,7 +117,7 @@ def _get_open_meteo_api(llm: BaseLLM) -> BaseTool:
     )
 
 
-_LLM_TOOLS = {
+_LLM_TOOLS: Dict[str, Callable[[BaseLLM], BaseTool]] = {
     "pal-math": _get_pal_math,
     "pal-colored-objects": _get_pal_colored_objects,
     "llm-math": _get_llm_math,
@@ -219,14 +216,16 @@ def _get_bing_search(**kwargs: Any) -> BaseTool:
 
 
 def _get_ddg_search(**kwargs: Any) -> BaseTool:
-    return DuckDuckGoSearchTool(api_wrapper=DuckDuckGoSearchAPIWrapper(**kwargs))
+    return DuckDuckGoSearchRun(api_wrapper=DuckDuckGoSearchAPIWrapper(**kwargs))
 
 
 def _get_human_tool(**kwargs: Any) -> BaseTool:
     return HumanInputRun(**kwargs)
 
 
-_EXTRA_LLM_TOOLS = {
+_EXTRA_LLM_TOOLS: Dict[
+    str, Tuple[Callable[[Arg(BaseLLM, "llm"), KwArg(Any)], BaseTool], List[str]]
+] = {
     "news-api": (_get_news_api, ["news_api_key"]),
     "tmdb-api": (_get_tmdb_api, ["tmdb_bearer_token"]),
     "podcast-api": (_get_podcast_api, ["listen_api_key"]),

langchain/agents/mrkl/base.py

@@ -10,10 +10,10 @@ from langchain.agents.agent_types import AgentType
 from langchain.agents.mrkl.output_parser import MRKLOutputParser
 from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS, PREFIX, SUFFIX
 from langchain.agents.tools import Tool
+from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
-from langchain.schema import BaseLanguageModel
 from langchain.tools.base import BaseTool
 
 
@@ -122,6 +122,7 @@ class ZeroShotAgent(Agent):
 
     @classmethod
     def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
+        super()._validate_tools(tools)
         for tool in tools:
             if tool.description is None:
                 raise ValueError(

langchain/agents/react/base.py

@@ -37,6 +37,7 @@ class ReActDocstoreAgent(Agent):
 
     @classmethod
     def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
+        super()._validate_tools(tools)
         if len(tools) != 2:
             raise ValueError(f"Exactly two tools must be specified, but got {tools}")
         tool_names = {tool.name for tool in tools}
@@ -119,6 +120,7 @@ class ReActTextWorldAgent(ReActDocstoreAgent):
 
     @classmethod
     def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
+        super()._validate_tools(tools)
         if len(tools) != 1:
             raise ValueError(f"Exactly one tool must be specified, but got {tools}")
         tool_names = {tool.name for tool in tools}

langchain/agents/self_ask_with_search/base.py

@@ -36,6 +36,7 @@ class SelfAskWithSearchAgent(Agent):
 
     @classmethod
     def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
+        super()._validate_tools(tools)
         if len(tools) != 1:
             raise ValueError(f"Exactly one tool must be specified, but got {tools}")
         tool_names = {tool.name for tool in tools}

langchain/agents/structured_agent.py

@@ -1,145 +0,0 @@
-"""Chain that takes in an input and produces an action and action input."""
-from __future__ import annotations
-
-import json
-import logging
-from abc import abstractmethod
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
-
-import yaml
-from pydantic import BaseModel
-
-from langchain.callbacks.base import BaseCallbackManager
-from langchain.schema import (
-    StructuredAgentAction,
-    AgentFinish,
-    BaseLanguageModel,
-)
-from langchain.tools.base import BaseTool
-
-logger = logging.getLogger(__name__)
-
-
-class BaseSingleActionAgent(BaseModel):
-    """Base Agent class."""
-
-    @property
-    def return_values(self) -> List[str]:
-        """Return values of the agent."""
-        return ["output"]
-
-    def get_allowed_tools(self) -> Optional[List[str]]:
-        return None
-
-    @abstractmethod
-    def plan(
-        self, intermediate_steps: List[Tuple[StructuredAgentAction, str]], **kwargs: Any
-    ) -> Union[StructuredAgentAction, AgentFinish]:
-        """Given input, decided what to do.
-
-        Args:
-            intermediate_steps: Steps the LLM has taken to date,
-                along with observations
-            **kwargs: User inputs.
-
-        Returns:
-            Action specifying what tool to use.
-        """
-
-    @abstractmethod
-    async def aplan(
-        self, intermediate_steps: List[Tuple[StructuredAgentAction, str]], **kwargs: Any
-    ) -> Union[StructuredAgentAction, AgentFinish]:
-        """Given input, decided what to do.
-
-        Args:
-            intermediate_steps: Steps the LLM has taken to date,
-                along with observations
-            **kwargs: User inputs.
-
-        Returns:
-            Action specifying what tool to use.
-        """
-
-    @property
-    @abstractmethod
-    def input_keys(self) -> List[str]:
-        """Return the input keys.
-
-        :meta private:
-        """
-
-    def return_stopped_response(
-        self,
-        early_stopping_method: str,
-        intermediate_steps: List[Tuple[StructuredAgentAction, str]],
-        **kwargs: Any,
-    ) -> AgentFinish:
-        """Return response when agent has been stopped due to max iterations."""
-        if early_stopping_method == "force":
-            # `force` just returns a constant string
-            return AgentFinish(
-                {"output": "Agent stopped due to iteration limit or time limit."}, ""
-            )
-        else:
-            raise ValueError(
-                f"Got unsupported early_stopping_method `{early_stopping_method}`"
-            )
-
-    @classmethod
-    def from_llm_and_tools(
-        cls,
-        llm: BaseLanguageModel,
-        tools: Sequence[BaseTool],
-        callback_manager: Optional[BaseCallbackManager] = None,
-        **kwargs: Any,
-    ) -> BaseSingleActionAgent:
-        raise NotImplementedError
-
-    @property
-    def _agent_type(self) -> str:
-        """Return Identifier of agent type."""
-        raise NotImplementedError
-
-    def dict(self, **kwargs: Any) -> Dict:
-        """Return dictionary representation of agent."""
-        _dict = super().dict()
-        _dict["_type"] = str(self._agent_type)
-        return _dict
-
-    def save(self, file_path: Union[Path, str]) -> None:
-        """Save the agent.
-
-        Args:
-            file_path: Path to file to save the agent to.
-
-        Example:
-        .. code-block:: python
-
-            # If working with agent executor
-            agent.agent.save(file_path="path/agent.yaml")
-        """
-        # Convert file to Path object.
-        if isinstance(file_path, str):
-            save_path = Path(file_path)
-        else:
-            save_path = file_path
-
-        directory_path = save_path.parent
-        directory_path.mkdir(parents=True, exist_ok=True)
-
-        # Fetch dictionary to save
-        agent_dict = self.dict()
-
-        if save_path.suffix == ".json":
-            with open(file_path, "w") as f:
-                json.dump(agent_dict, f, indent=4)
-        elif save_path.suffix == ".yaml":
-            with open(file_path, "w") as f:
-                yaml.dump(agent_dict, f, default_flow_style=False)
-        else:
-            raise ValueError(f"{save_path} must be json or yaml")
-
-    def tool_run_logging_kwargs(self) -> Dict:
-        return {}

langchain/agents/tools.py

@@ -1,13 +1,15 @@
 """Interface for tools."""
 from functools import partial
 from inspect import signature
-from typing import Any, Awaitable, Callable, Optional, Union
+from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, Type, Union
 
-from pydantic import validator
+from pydantic import BaseModel, validator
 
-from langchain.tools.base import (
-    BaseTool,
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
 )
+from langchain.tools.base import BaseTool, StructuredTool
 
 
 class Tool(BaseTool):
@@ -26,19 +28,70 @@ class Tool(BaseTool):
             raise ValueError("Partial functions not yet supported in tools.")
         return func
 
-    def _run(self, tool_input: str) -> str:
-        """Use the tool."""
-        return self.func(tool_input)
+    @property
+    def args(self) -> dict:
+        """The tool's input arguments."""
+        if self.args_schema is not None:
+            return self.args_schema.schema()["properties"]
+        # For backwards compatibility, if the function signature is ambiguous,
+        # assume it takes a single string input.
+        return {"tool_input": {"type": "string"}}
 
-    async def _arun(self, tool_input: str) -> str:
+    def _to_args_and_kwargs(self, tool_input: Union[str, Dict]) -> Tuple[Tuple, Dict]:
+        """Convert tool input to pydantic model."""
+        args, kwargs = super()._to_args_and_kwargs(tool_input)
+        # For backwards compatibility. The tool must be run with a single input
+        all_args = list(args) + list(kwargs.values())
+        if len(all_args) != 1:
+            raise ValueError(
+                f"Too many arguments to single-input tool {self.name}."
+                f" Args: {all_args}"
+            )
+        return tuple(all_args), {}
+
+    def _run(
+        self,
+        *args: Any,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Use the tool."""
+        new_argument_supported = signature(self.func).parameters.get("callbacks")
+        return (
+            self.func(
+                *args,
+                callbacks=run_manager.get_child() if run_manager else None,
+                **kwargs,
+            )
+            if new_argument_supported
+            else self.func(*args, **kwargs)
+        )
+
+    async def _arun(
+        self,
+        *args: Any,
+        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
+        **kwargs: Any,
+    ) -> Any:
         """Use the tool asynchronously."""
         if self.coroutine:
-            return await self.coroutine(tool_input)
+            new_argument_supported = signature(self.coroutine).parameters.get(
+                "callbacks"
+            )
+            return (
+                await self.coroutine(
+                    *args,
+                    callbacks=run_manager.get_child() if run_manager else None,
+                    **kwargs,
+                )
+                if new_argument_supported
+                else await self.coroutine(*args, **kwargs)
+            )
         raise NotImplementedError("Tool does not support async")
 
     # TODO: this is for backwards compatibility, remove in future
     def __init__(
-        self, name: str, func: Callable[[str], str], description: str, **kwargs: Any
+        self, name: str, func: Callable, description: str, **kwargs: Any
     ) -> None:
         """Initialize tool."""
         super(Tool, self).__init__(
@@ -52,11 +105,17 @@ class InvalidTool(BaseTool):
     name = "invalid_tool"
     description = "Called when tool name is invalid."
 
-    def _run(self, tool_name: str) -> str:
+    def _run(
+        self, tool_name: str, run_manager: Optional[CallbackManagerForToolRun] = None
+    ) -> str:
         """Use the tool."""
         return f"{tool_name} is not a valid tool, try another one."
 
-    async def _arun(self, tool_name: str) -> str:
+    async def _arun(
+        self,
+        tool_name: str,
+        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
+    ) -> str:
         """Use the tool asynchronously."""
         return f"{tool_name} is not a valid tool, try another one."
 
@@ -64,6 +123,8 @@ class InvalidTool(BaseTool):
 def tool(
     *args: Union[str, Callable],
     return_direct: bool = False,
+    args_schema: Optional[Type[BaseModel]] = None,
+    infer_schema: bool = True,
 ) -> Callable:
     """Make tools out of functions, can be used with or without arguments.
 
@@ -71,6 +132,10 @@ def tool(
         *args: The arguments to the tool.
         return_direct: Whether to return directly from the tool rather
             than continuing the agent loop.
+        args_schema: optional argument schema for user to specify
+        infer_schema: Whether to infer the schema of the arguments from
+            the function's signature. This also makes the resultant tool
+            accept a dictionary input to its `run()` function.
 
     Requires:
         - Function must be of type (str) -> str
@@ -91,18 +156,24 @@ def tool(
     """
 
     def _make_with_name(tool_name: str) -> Callable:
-        def _make_tool(func: Callable) -> Tool:
-            assert func.__doc__, "Function must have a docstring"
-            # Description example:
-            # search_api(query: str) - Searches the API for the query.
-            description = f"{tool_name}{signature(func)} - {func.__doc__.strip()}"
-            tool_ = Tool(
+        def _make_tool(func: Callable) -> BaseTool:
+            if infer_schema or args_schema is not None:
+                return StructuredTool.from_function(
+                    func,
+                    name=tool_name,
+                    return_direct=return_direct,
+                    args_schema=args_schema,
+                    infer_schema=infer_schema,
+                )
+            # If someone doesn't want a schema applied, we must treat it as
+            # a simple string->string function
+            assert func.__doc__ is not None, "Function must have a docstring"
+            return Tool(
                 name=tool_name,
                 func=func,
-                description=description,
+                description=f"{tool_name} tool",
                 return_direct=return_direct,
             )
-            return tool_
 
         return _make_tool
 

langchain/base_language.py

@@ -0,0 +1,55 @@
+"""Base class for all language models."""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import BaseMessage, LLMResult, PromptValue, get_buffer_string
+
+
+class BaseLanguageModel(BaseModel, ABC):
+    @abstractmethod
+    def generate_prompt(
+        self,
+        prompts: List[PromptValue],
+        stop: Optional[List[str]] = None,
+        callbacks: Callbacks = None,
+    ) -> LLMResult:
+        """Take in a list of prompt values and return an LLMResult."""
+
+    @abstractmethod
+    async def agenerate_prompt(
+        self,
+        prompts: List[PromptValue],
+        stop: Optional[List[str]] = None,
+        callbacks: Callbacks = None,
+    ) -> LLMResult:
+        """Take in a list of prompt values and return an LLMResult."""
+
+    def get_num_tokens(self, text: str) -> int:
+        """Get the number of tokens present in the text."""
+        # TODO: this method may not be exact.
+        # TODO: this method may differ based on model (eg codex).
+        try:
+            from transformers import GPT2TokenizerFast
+        except ImportError:
+            raise ValueError(
+                "Could not import transformers python package. "
+                "This is needed in order to calculate get_num_tokens. "
+                "Please install it with `pip install transformers`."
+            )
+        # create a GPT-3 tokenizer instance
+        tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
+
+        # tokenize the text using the GPT-3 tokenizer
+        tokenized_text = tokenizer.tokenize(text)
+
+        # calculate the number of tokens in the tokenized text
+        return len(tokenized_text)
+
+    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
+        """Get the number of tokens in the message."""
+        return sum([self.get_num_tokens(get_buffer_string([m])) for m in messages])

langchain/cache.py

@@ -1,7 +1,7 @@
 """Beta Feature: base interface for cache."""
 import json
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, cast
 
 from sqlalchemy import Column, Integer, String, create_engine, select
 from sqlalchemy.engine.base import Engine
@@ -28,6 +28,10 @@ class BaseCache(ABC):
     def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:
         """Update cache based on prompt and llm_string."""
 
+    @abstractmethod
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache that can take additional keyword arguments."""
+
 
 class InMemoryCache(BaseCache):
     """Cache that stores things in memory."""
@@ -44,6 +48,10 @@ class InMemoryCache(BaseCache):
         """Update cache based on prompt and llm_string."""
         self._cache[(prompt, llm_string)] = return_val
 
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache."""
+        self._cache = {}
+
 
 Base = declarative_base()
 
@@ -61,7 +69,7 @@ class FullLLMCache(Base):  # type: ignore
 class SQLAlchemyCache(BaseCache):
     """Cache that uses SQAlchemy as a backend."""
 
-    def __init__(self, engine: Engine, cache_schema: Any = FullLLMCache):
+    def __init__(self, engine: Engine, cache_schema: Type[FullLLMCache] = FullLLMCache):
         """Initialize by creating all tables."""
         self.engine = engine
         self.cache_schema = cache_schema
@@ -76,20 +84,26 @@ class SQLAlchemyCache(BaseCache):
             .order_by(self.cache_schema.idx)
         )
         with Session(self.engine) as session:
-            generations = [Generation(text=row[0]) for row in session.execute(stmt)]
-            if len(generations) > 0:
-                return generations
+            rows = session.execute(stmt).fetchall()
+            if rows:
+                return [Generation(text=row[0]) for row in rows]
         return None
 
     def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:
-        """Look up based on prompt and llm_string."""
-        for i, generation in enumerate(return_val):
-            item = self.cache_schema(
-                prompt=prompt, llm=llm_string, response=generation.text, idx=i
-            )
-            with Session(self.engine) as session, session.begin():
+        """Update based on prompt and llm_string."""
+        items = [
+            self.cache_schema(prompt=prompt, llm=llm_string, response=gen.text, idx=i)
+            for i, gen in enumerate(return_val)
+        ]
+        with Session(self.engine) as session, session.begin():
+            for item in items:
                 session.merge(item)
 
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache."""
+        with Session(self.engine) as session:
+            session.execute(self.cache_schema.delete())
+
 
 class SQLiteCache(SQLAlchemyCache):
     """Cache that uses SQLite as a backend."""
@@ -139,19 +153,26 @@ class RedisCache(BaseCache):
         for i, generation in enumerate(return_val):
             self.redis.set(self._key(prompt, llm_string, i), generation.text)
 
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache. If `asynchronous` is True, flush asynchronously."""
+        asynchronous = kwargs.get("asynchronous", False)
+        self.redis.flushdb(asynchronous=asynchronous, **kwargs)
+
 
 class GPTCache(BaseCache):
     """Cache that uses GPTCache as a backend."""
 
-    def __init__(self, init_func: Callable[[Any], None]):
-        """Initialize by passing in the `init` GPTCache func
+    def __init__(self, init_func: Optional[Callable[[Any], None]] = None):
+        """Initialize by passing in init function (default: `None`).
 
         Args:
-            init_func (Callable[[Any], None]): init `GPTCache` function
+            init_func (Optional[Callable[[Any], None]]): init `GPTCache` function
+            (default: `None`)
 
         Example:
         .. code-block:: python
 
+            # Initialize GPTCache with a custom init function
             import gptcache
             from gptcache.processor.pre import get_prompt
             from gptcache.manager.factory import get_data_manager
@@ -180,7 +201,8 @@ class GPTCache(BaseCache):
                 "Could not import gptcache python package. "
                 "Please install it with `pip install gptcache`."
             )
-        self.init_gptcache_func: Callable[[Any], None] = init_func
+
+        self.init_gptcache_func: Optional[Callable[[Any], None]] = init_func
         self.gptcache_dict: Dict[str, Any] = {}
 
     @staticmethod
@@ -205,11 +227,19 @@ class GPTCache(BaseCache):
 
         When the corresponding llm model cache does not exist, it will be created."""
         from gptcache import Cache
+        from gptcache.manager.factory import get_data_manager
+        from gptcache.processor.pre import get_prompt
 
         _gptcache = self.gptcache_dict.get(llm_string, None)
         if _gptcache is None:
             _gptcache = Cache()
-            self.init_gptcache_func(_gptcache)
+            if self.init_gptcache_func is not None:
+                self.init_gptcache_func(_gptcache)
+            else:
+                _gptcache.init(
+                    pre_embedding_func=get_prompt,
+                    data_manager=get_data_manager(data_path=llm_string),
+                )
             self.gptcache_dict[llm_string] = _gptcache
         return _gptcache
 
@@ -220,7 +250,7 @@ class GPTCache(BaseCache):
         """
         from gptcache.adapter.adapter import adapt
 
-        _gptcache = self.gptcache_dict.get(llm_string)
+        _gptcache = self.gptcache_dict.get(llm_string, None)
         if _gptcache is None:
             return None
         res = adapt(
@@ -234,7 +264,10 @@ class GPTCache(BaseCache):
 
     @staticmethod
     def _update_cache_callback(
-        llm_data: RETURN_VAL_TYPE, update_cache_func: Callable[[Any], None]
+        llm_data: RETURN_VAL_TYPE,
+        update_cache_func: Callable[[Any], None],
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
         """Save the `llm_data` to cache storage"""
         handled_data = json.dumps([generation.dict() for generation in llm_data])
@@ -260,3 +293,13 @@ class GPTCache(BaseCache):
             cache_skip=True,
             prompt=prompt,
         )
+
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache."""
+        from gptcache import Cache
+
+        for gptcache_instance in self.gptcache_dict.values():
+            gptcache_instance = cast(Cache, gptcache_instance)
+            gptcache_instance.flush()
+
+        self.gptcache_dict.clear()

langchain/callbacks/__init__.py

@@ -1,80 +1,27 @@
 """Callback handlers that allow listening to events in LangChain."""
-import os
 from contextlib import contextmanager
-from typing import Generator, Optional
+from typing import Generator
 
 from langchain.callbacks.aim_callback import AimCallbackHandler
 from langchain.callbacks.base import (
-    AsyncCallbackManager,
     BaseCallbackHandler,
     BaseCallbackManager,
-    CallbackManager,
 )
 from langchain.callbacks.clearml_callback import ClearMLCallbackHandler
 from langchain.callbacks.comet_ml_callback import CometCallbackHandler
+from langchain.callbacks.manager import (
+    CallbackManager,
+    get_openai_callback,
+    tracing_enabled,
+)
 from langchain.callbacks.openai_info import OpenAICallbackHandler
-from langchain.callbacks.shared import SharedCallbackManager
 from langchain.callbacks.stdout import StdOutCallbackHandler
 from langchain.callbacks.streaming_aiter import AsyncIteratorCallbackHandler
-from langchain.callbacks.tracers import SharedLangChainTracer
+from langchain.callbacks.tracers import LangChainTracer
 from langchain.callbacks.wandb_callback import WandbCallbackHandler
 
-
-def get_callback_manager() -> BaseCallbackManager:
-    """Return the shared callback manager."""
-    return SharedCallbackManager()
-
-
-def set_handler(handler: BaseCallbackHandler) -> None:
-    """Set handler."""
-    callback = get_callback_manager()
-    callback.set_handler(handler)
-
-
-def set_default_callback_manager() -> None:
-    """Set default callback manager."""
-    default_handler = os.environ.get("LANGCHAIN_HANDLER", "stdout")
-    if default_handler == "stdout":
-        set_handler(StdOutCallbackHandler())
-    elif default_handler == "langchain":
-        session = os.environ.get("LANGCHAIN_SESSION")
-        set_tracing_callback_manager(session)
-    else:
-        raise ValueError(
-            f"LANGCHAIN_HANDLER should be one of `stdout` "
-            f"or `langchain`, got {default_handler}"
-        )
-
-
-def set_tracing_callback_manager(session_name: Optional[str] = None) -> None:
-    """Set tracing callback manager."""
-    handler = SharedLangChainTracer()
-    callback = get_callback_manager()
-    callback.set_handlers([handler, StdOutCallbackHandler()])
-    if session_name is None:
-        handler.load_default_session()
-    else:
-        try:
-            handler.load_session(session_name)
-        except Exception:
-            raise ValueError(f"session {session_name} not found")
-
-
-@contextmanager
-def get_openai_callback() -> Generator[OpenAICallbackHandler, None, None]:
-    """Get OpenAI callback handler in a context manager."""
-    handler = OpenAICallbackHandler()
-    manager = get_callback_manager()
-    manager.add_handler(handler)
-    yield handler
-    manager.remove_handler(handler)
-
-
 __all__ = [
-    "CallbackManager",
-    "AsyncCallbackManager",
     "OpenAICallbackHandler",
-    "SharedCallbackManager",
     "StdOutCallbackHandler",
     "AimCallbackHandler",
     "WandbCallbackHandler",
@@ -82,8 +29,5 @@ __all__ = [
     "CometCallbackHandler",
     "AsyncIteratorCallbackHandler",
     "get_openai_callback",
-    "set_tracing_callback_manager",
-    "set_default_callback_manager",
-    "set_handler",
-    "get_callback_manager",
+    "tracing_enabled",
 ]

langchain/callbacks/base.py

@@ -1,19 +1,173 @@
-"""Base callback handler that can be used to handle callbacks from langchain."""
-import asyncio
-import functools
-from abc import ABC, abstractmethod
+"""Base callback handler that can be used to handle callbacks in langchain."""
+from __future__ import annotations
+
+import copy
 from typing import Any, Dict, List, Optional, Union
 
 from langchain.schema import AgentAction, AgentFinish, LLMResult
 
 
-class BaseCallbackHandler(ABC):
-    """Base callback handler that can be used to handle callbacks from langchain."""
+class LLMManagerMixin:
+    """Mixin for LLM callbacks."""
 
-    @property
-    def always_verbose(self) -> bool:
-        """Whether to call verbose callbacks even if verbose is False."""
-        return False
+    def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on new LLM token. Only available when streaming is enabled."""
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when LLM ends running."""
+
+    def on_llm_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when LLM errors."""
+
+
+class ChainManagerMixin:
+    """Mixin for chain callbacks."""
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when chain ends running."""
+
+    def on_chain_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when chain errors."""
+
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on agent action."""
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on agent end."""
+
+
+class ToolManagerMixin:
+    """Mixin for tool callbacks."""
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when tool ends running."""
+
+    def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when tool errors."""
+
+
+class CallbackManagerMixin:
+    """Mixin for callback manager."""
+
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when LLM starts running."""
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when chain starts running."""
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when tool starts running."""
+
+
+class RunManagerMixin:
+    """Mixin for run manager."""
+
+    def on_text(
+        self,
+        text: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on arbitrary text."""
+
+
+class BaseCallbackHandler(
+    LLMManagerMixin,
+    ChainManagerMixin,
+    ToolManagerMixin,
+    CallbackManagerMixin,
+    RunManagerMixin,
+):
+    """Base callback handler that can be used to handle callbacks from langchain."""
 
     @property
     def ignore_llm(self) -> bool:
@@ -30,480 +184,197 @@ class BaseCallbackHandler(ABC):
         """Whether to ignore agent callbacks."""
         return False
 
-    @abstractmethod
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> Any:
+
+class AsyncCallbackHandler(BaseCallbackHandler):
+    """Async callback handler that can be used to handle callbacks from langchain."""
+
+    async def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when LLM starts running."""
 
-    @abstractmethod
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
+    async def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run on new LLM token. Only available when streaming is enabled."""
 
-    @abstractmethod
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
+    async def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when LLM ends running."""
 
-    @abstractmethod
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    async def on_llm_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when LLM errors."""
 
-    @abstractmethod
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> Any:
+    async def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when chain starts running."""
 
-    @abstractmethod
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+    async def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when chain ends running."""
 
-    @abstractmethod
-    def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    async def on_chain_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when chain errors."""
 
-    @abstractmethod
-    def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
-    ) -> Any:
+    async def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when tool starts running."""
 
-    @abstractmethod
-    def on_tool_end(self, output: str, **kwargs: Any) -> Any:
+    async def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when tool ends running."""
 
-    @abstractmethod
-    def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    async def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run when tool errors."""
 
-    @abstractmethod
-    def on_text(self, text: str, **kwargs: Any) -> Any:
+    async def on_text(
+        self,
+        text: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run on arbitrary text."""
 
-    @abstractmethod
-    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+    async def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run on agent action."""
 
-    @abstractmethod
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
+    async def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """Run on agent end."""
 
 
-class BaseCallbackManager(BaseCallbackHandler, ABC):
+class BaseCallbackManager(CallbackManagerMixin):
     """Base callback manager that can be used to handle callbacks from LangChain."""
 
+    def __init__(
+        self,
+        handlers: List[BaseCallbackHandler],
+        inheritable_handlers: Optional[List[BaseCallbackHandler]] = None,
+        parent_run_id: Optional[str] = None,
+    ) -> None:
+        """Initialize callback manager."""
+        self.handlers: List[BaseCallbackHandler] = handlers
+        self.inheritable_handlers: List[BaseCallbackHandler] = (
+            inheritable_handlers or []
+        )
+        self.parent_run_id: Optional[str] = parent_run_id
+
     @property
     def is_async(self) -> bool:
         """Whether the callback manager is async."""
         return False
 
-    @abstractmethod
-    def add_handler(self, callback: BaseCallbackHandler) -> None:
+    def add_handler(self, handler: BaseCallbackHandler, inherit: bool = True) -> None:
         """Add a handler to the callback manager."""
+        self.handlers.append(handler)
+        if inherit:
+            self.inheritable_handlers.append(handler)
 
-    @abstractmethod
     def remove_handler(self, handler: BaseCallbackHandler) -> None:
         """Remove a handler from the callback manager."""
+        self.handlers.remove(handler)
+        self.inheritable_handlers.remove(handler)
 
-    def set_handler(self, handler: BaseCallbackHandler) -> None:
+    def set_handlers(
+        self, handlers: List[BaseCallbackHandler], inherit: bool = True
+    ) -> None:
+        """Set handlers as the only handlers on the callback manager."""
+        self.handlers = []
+        self.inheritable_handlers = []
+        for handler in handlers:
+            self.add_handler(handler, inherit=inherit)
+
+    def set_handler(self, handler: BaseCallbackHandler, inherit: bool = True) -> None:
         """Set handler as the only handler on the callback manager."""
-        self.set_handlers([handler])
+        self.set_handlers([handler], inherit=inherit)
 
-    @abstractmethod
-    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Set handlers as the only handlers on the callback manager."""
-
-
-class CallbackManager(BaseCallbackManager):
-    """Callback manager that can be used to handle callbacks from langchain."""
-
-    def __init__(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Initialize callback manager."""
-        self.handlers: List[BaseCallbackHandler] = handlers
-
-    def on_llm_start(
-        self,
-        serialized: Dict[str, Any],
-        prompts: List[str],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running."""
-        for handler in self.handlers:
-            if not handler.ignore_llm:
-                if verbose or handler.always_verbose:
-                    handler.on_llm_start(serialized, prompts, **kwargs)
-
-    def on_llm_new_token(
-        self, token: str, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when LLM generates a new token."""
-        for handler in self.handlers:
-            if not handler.ignore_llm:
-                if verbose or handler.always_verbose:
-                    handler.on_llm_new_token(token, **kwargs)
-
-    def on_llm_end(
-        self, response: LLMResult, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when LLM ends running."""
-        for handler in self.handlers:
-            if not handler.ignore_llm:
-                if verbose or handler.always_verbose:
-                    handler.on_llm_end(response)
-
-    def on_llm_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        for handler in self.handlers:
-            if not handler.ignore_llm:
-                if verbose or handler.always_verbose:
-                    handler.on_llm_error(error)
-
-    def on_chain_start(
-        self,
-        serialized: Dict[str, Any],
-        inputs: Dict[str, Any],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-        for handler in self.handlers:
-            if not handler.ignore_chain:
-                if verbose or handler.always_verbose:
-                    handler.on_chain_start(serialized, inputs, **kwargs)
-
-    def on_chain_end(
-        self, outputs: Dict[str, Any], verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when chain ends running."""
-        for handler in self.handlers:
-            if not handler.ignore_chain:
-                if verbose or handler.always_verbose:
-                    handler.on_chain_end(outputs)
-
-    def on_chain_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when chain errors."""
-        for handler in self.handlers:
-            if not handler.ignore_chain:
-                if verbose or handler.always_verbose:
-                    handler.on_chain_error(error)
-
-    def on_tool_start(
-        self,
-        serialized: Dict[str, Any],
-        input_str: str,
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-        for handler in self.handlers:
-            if not handler.ignore_agent:
-                if verbose or handler.always_verbose:
-                    handler.on_tool_start(serialized, input_str, **kwargs)
-
-    def on_agent_action(
-        self, action: AgentAction, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-        for handler in self.handlers:
-            if not handler.ignore_agent:
-                if verbose or handler.always_verbose:
-                    handler.on_agent_action(action, **kwargs)
-
-    def on_tool_end(self, output: str, verbose: bool = False, **kwargs: Any) -> None:
-        """Run when tool ends running."""
-        for handler in self.handlers:
-            if not handler.ignore_agent:
-                if verbose or handler.always_verbose:
-                    handler.on_tool_end(output, **kwargs)
-
-    def on_tool_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when tool errors."""
-        for handler in self.handlers:
-            if not handler.ignore_agent:
-                if verbose or handler.always_verbose:
-                    handler.on_tool_error(error)
-
-    def on_text(self, text: str, verbose: bool = False, **kwargs: Any) -> None:
-        """Run on additional input from chains and agents."""
-        for handler in self.handlers:
-            if verbose or handler.always_verbose:
-                handler.on_text(text, **kwargs)
-
-    def on_agent_finish(
-        self, finish: AgentFinish, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run on agent end."""
-        for handler in self.handlers:
-            if not handler.ignore_agent:
-                if verbose or handler.always_verbose:
-                    handler.on_agent_finish(finish, **kwargs)
-
-    def add_handler(self, handler: BaseCallbackHandler) -> None:
-        """Add a handler to the callback manager."""
-        self.handlers.append(handler)
-
-    def remove_handler(self, handler: BaseCallbackHandler) -> None:
-        """Remove a handler from the callback manager."""
-        self.handlers.remove(handler)
-
-    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Set handlers as the only handlers on the callback manager."""
-        self.handlers = handlers
-
-
-class AsyncCallbackHandler(BaseCallbackHandler):
-    """Async callback handler that can be used to handle callbacks from langchain."""
-
-    async def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running."""
-
-    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-
-    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-
-    async def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-
-    async def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-
-    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        """Run when chain ends running."""
-
-    async def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when chain errors."""
-
-    async def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-
-    async def on_tool_end(self, output: str, **kwargs: Any) -> None:
-        """Run when tool ends running."""
-
-    async def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when tool errors."""
-
-    async def on_text(self, text: str, **kwargs: Any) -> None:
-        """Run on arbitrary text."""
-
-    async def on_agent_action(self, action: AgentAction, **kwargs: Any) -> None:
-        """Run on agent action."""
-
-    async def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
-        """Run on agent end."""
-
-
-async def _handle_event_for_handler(
-    handler: BaseCallbackHandler,
-    event_name: str,
-    ignore_condition_name: Optional[str],
-    verbose: bool,
-    *args: Any,
-    **kwargs: Any
-) -> None:
-    if ignore_condition_name is None or not getattr(handler, ignore_condition_name):
-        if verbose or handler.always_verbose:
-            event = getattr(handler, event_name)
-            if asyncio.iscoroutinefunction(event):
-                await event(*args, **kwargs)
-            else:
-                await asyncio.get_event_loop().run_in_executor(
-                    None, functools.partial(event, *args, **kwargs)
-                )
-
-
-class AsyncCallbackManager(BaseCallbackManager):
-    """Async callback manager that can be used to handle callbacks from LangChain."""
-
-    @property
-    def is_async(self) -> bool:
-        """Return whether the handler is async."""
-        return True
-
-    def __init__(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Initialize callback manager."""
-        self.handlers: List[BaseCallbackHandler] = handlers
-
-    async def _handle_event(
-        self,
-        event_name: str,
-        ignore_condition_name: Optional[str],
-        verbose: bool,
-        *args: Any,
-        **kwargs: Any
-    ) -> None:
-        """Generic event handler for AsyncCallbackManager."""
-        await asyncio.gather(
-            *(
-                _handle_event_for_handler(
-                    handler, event_name, ignore_condition_name, verbose, *args, **kwargs
-                )
-                for handler in self.handlers
-            )
+    def __copy__(self) -> "BaseCallbackManager":
+        return self.__class__(
+            self.handlers.copy(), self.inheritable_handlers.copy(), self.parent_run_id
         )
 
-    async def on_llm_start(
-        self,
-        serialized: Dict[str, Any],
-        prompts: List[str],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running."""
-        await self._handle_event(
-            "on_llm_start", "ignore_llm", verbose, serialized, prompts, **kwargs
+    def __deepcopy__(self, memo: dict) -> "BaseCallbackManager":
+        return self.__class__(
+            [copy.deepcopy(handler, memo) for handler in self.handlers],
+            [copy.deepcopy(handler, memo) for handler in self.inheritable_handlers],
+            self.parent_run_id,
         )
-
-    async def on_llm_new_token(
-        self, token: str, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-        await self._handle_event(
-            "on_llm_new_token", "ignore_llm", verbose, token, **kwargs
-        )
-
-    async def on_llm_end(
-        self, response: LLMResult, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when LLM ends running."""
-        await self._handle_event(
-            "on_llm_end", "ignore_llm", verbose, response, **kwargs
-        )
-
-    async def on_llm_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        await self._handle_event("on_llm_error", "ignore_llm", verbose, error, **kwargs)
-
-    async def on_chain_start(
-        self,
-        serialized: Dict[str, Any],
-        inputs: Dict[str, Any],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-        await self._handle_event(
-            "on_chain_start", "ignore_chain", verbose, serialized, inputs, **kwargs
-        )
-
-    async def on_chain_end(
-        self, outputs: Dict[str, Any], verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when chain ends running."""
-        await self._handle_event(
-            "on_chain_end", "ignore_chain", verbose, outputs, **kwargs
-        )
-
-    async def on_chain_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when chain errors."""
-        await self._handle_event(
-            "on_chain_error", "ignore_chain", verbose, error, **kwargs
-        )
-
-    async def on_tool_start(
-        self,
-        serialized: Dict[str, Any],
-        input_str: str,
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-        await self._handle_event(
-            "on_tool_start", "ignore_agent", verbose, serialized, input_str, **kwargs
-        )
-
-    async def on_tool_end(
-        self, output: str, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when tool ends running."""
-        await self._handle_event(
-            "on_tool_end", "ignore_agent", verbose, output, **kwargs
-        )
-
-    async def on_tool_error(
-        self,
-        error: Union[Exception, KeyboardInterrupt],
-        verbose: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """Run when tool errors."""
-        await self._handle_event(
-            "on_tool_error", "ignore_agent", verbose, error, **kwargs
-        )
-
-    async def on_text(self, text: str, verbose: bool = False, **kwargs: Any) -> None:
-        """Run when text is printed."""
-        await self._handle_event("on_text", None, verbose, text, **kwargs)
-
-    async def on_agent_action(
-        self, action: AgentAction, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run on agent action."""
-        await self._handle_event(
-            "on_agent_action", "ignore_agent", verbose, action, **kwargs
-        )
-
-    async def on_agent_finish(
-        self, finish: AgentFinish, verbose: bool = False, **kwargs: Any
-    ) -> None:
-        """Run when agent finishes."""
-        await self._handle_event(
-            "on_agent_finish", "ignore_agent", verbose, finish, **kwargs
-        )
-
-    def add_handler(self, handler: BaseCallbackHandler) -> None:
-        """Add a handler to the callback manager."""
-        self.handlers.append(handler)
-
-    def remove_handler(self, handler: BaseCallbackHandler) -> None:
-        """Remove a handler from the callback manager."""
-        self.handlers.remove(handler)
-
-    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Set handlers as the only handlers on the callback manager."""
-        self.handlers = handlers

langchain/callbacks/manager.py

@@ -0,0 +1,731 @@
+from __future__ import annotations
+
+import asyncio
+import copy
+import functools
+import os
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Dict, Generator, List, Optional, Sequence, Type, TypeVar, Union
+
+from langchain.callbacks.base import (
+    BaseCallbackHandler,
+    BaseCallbackManager,
+    ChainManagerMixin,
+    LLMManagerMixin,
+    RunManagerMixin,
+    ToolManagerMixin,
+)
+from langchain.callbacks.openai_info import OpenAICallbackHandler
+from langchain.callbacks.stdout import StdOutCallbackHandler
+from langchain.callbacks.tracers.base import TracerSession
+from langchain.callbacks.tracers.langchain import LangChainTracer
+from langchain.schema import AgentAction, AgentFinish, LLMResult
+
+Callbacks = Optional[Union[List[BaseCallbackHandler], BaseCallbackManager]]
+
+openai_callback_var: ContextVar[Optional[OpenAICallbackHandler]] = ContextVar(
+    "openai_callback", default=None
+)
+tracing_callback_var: ContextVar[Optional[LangChainTracer]] = ContextVar(
+    "tracing_callback", default=None
+)
+
+
+@contextmanager
+def get_openai_callback() -> Generator[OpenAICallbackHandler, None, None]:
+    """Get OpenAI callback handler in a context manager."""
+    cb = OpenAICallbackHandler()
+    openai_callback_var.set(cb)
+    yield cb
+    openai_callback_var.set(None)
+
+
+@contextmanager
+def tracing_enabled(
+    session_name: str = "default",
+) -> Generator[TracerSession, None, None]:
+    """Get OpenAI callback handler in a context manager."""
+    cb = LangChainTracer()
+    session = cb.load_session(session_name)
+    tracing_callback_var.set(cb)
+    yield session
+    tracing_callback_var.set(None)
+
+
+def _handle_event(
+    handlers: List[BaseCallbackHandler],
+    event_name: str,
+    ignore_condition_name: Optional[str],
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    for handler in handlers:
+        try:
+            if ignore_condition_name is None or not getattr(
+                handler, ignore_condition_name
+            ):
+                getattr(handler, event_name)(*args, **kwargs)
+        except Exception as e:
+            # TODO: switch this to use logging
+            print(f"Error in {event_name} callback: {e}")
+
+
+async def _ahandle_event_for_handler(
+    handler: BaseCallbackHandler,
+    event_name: str,
+    ignore_condition_name: Optional[str],
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    try:
+        if ignore_condition_name is None or not getattr(handler, ignore_condition_name):
+            event = getattr(handler, event_name)
+            if asyncio.iscoroutinefunction(event):
+                await event(*args, **kwargs)
+            else:
+                await asyncio.get_event_loop().run_in_executor(
+                    None, functools.partial(event, *args, **kwargs)
+                )
+    except Exception as e:
+        # TODO: switch this to use logging
+        print(f"Error in {event_name} callback: {e}")
+
+
+async def _ahandle_event(
+    handlers: List[BaseCallbackHandler],
+    event_name: str,
+    ignore_condition_name: Optional[str],
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    """Generic event handler for AsyncCallbackManager."""
+    await asyncio.gather(
+        *(
+            _ahandle_event_for_handler(
+                handler, event_name, ignore_condition_name, *args, **kwargs
+            )
+            for handler in handlers
+        )
+    )
+
+
+BRM = TypeVar("BRM", bound="BaseRunManager")
+
+
+class BaseRunManager(RunManagerMixin):
+    """Base class for run manager (a bound callback manager)."""
+
+    def __init__(
+        self,
+        run_id: str,
+        handlers: List[BaseCallbackHandler],
+        inheritable_handlers: List[BaseCallbackHandler],
+        parent_run_id: Optional[str] = None,
+    ) -> None:
+        """Initialize run manager."""
+        self.run_id = run_id
+        self.handlers = handlers
+        self.inheritable_handlers = inheritable_handlers
+        self.parent_run_id = parent_run_id
+
+    @classmethod
+    def get_noop_manager(cls: Type[BRM]) -> BRM:
+        """Return a manager that doesn't perform any operations."""
+        return cls("", [], [])
+
+
+class RunManager(BaseRunManager):
+    """Sync Run Manager."""
+
+    def on_text(
+        self,
+        text: str,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when text is received."""
+        _handle_event(
+            self.handlers,
+            "on_text",
+            None,
+            text,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class AsyncRunManager(BaseRunManager):
+    """Async Run Manager."""
+
+    async def on_text(
+        self,
+        text: str,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when text is received."""
+        await _ahandle_event(
+            self.handlers,
+            "on_text",
+            None,
+            text,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class CallbackManagerForLLMRun(RunManager, LLMManagerMixin):
+    """Callback manager for LLM run."""
+
+    def on_llm_new_token(
+        self,
+        token: str,
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM generates a new token."""
+        _handle_event(
+            self.handlers,
+            "on_llm_new_token",
+            "ignore_llm",
+            token=token,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Run when LLM ends running."""
+        _handle_event(
+            self.handlers,
+            "on_llm_end",
+            "ignore_llm",
+            response,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_llm_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM errors."""
+        _handle_event(
+            self.handlers,
+            "on_llm_error",
+            "ignore_llm",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
+    """Async callback manager for LLM run."""
+
+    async def on_llm_new_token(
+        self,
+        token: str,
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM generates a new token."""
+        await _ahandle_event(
+            self.handlers,
+            "on_llm_new_token",
+            "ignore_llm",
+            token,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Run when LLM ends running."""
+        await _ahandle_event(
+            self.handlers,
+            "on_llm_end",
+            "ignore_llm",
+            response,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_llm_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM errors."""
+        await _ahandle_event(
+            self.handlers,
+            "on_llm_error",
+            "ignore_llm",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class CallbackManagerForChainRun(RunManager, ChainManagerMixin):
+    """Callback manager for chain run."""
+
+    def get_child(self) -> CallbackManager:
+        """Get a child callback manager."""
+        manager = CallbackManager([], parent_run_id=self.run_id)
+        manager.set_handlers(self.inheritable_handlers)
+        return manager
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        """Run when chain ends running."""
+        _handle_event(
+            self.handlers,
+            "on_chain_end",
+            "ignore_chain",
+            outputs,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_chain_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when chain errors."""
+        _handle_event(
+            self.handlers,
+            "on_chain_error",
+            "ignore_chain",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+        """Run when agent action is received."""
+        _handle_event(
+            self.handlers,
+            "on_agent_action",
+            "ignore_agent",
+            action,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
+        """Run when agent finish is received."""
+        _handle_event(
+            self.handlers,
+            "on_agent_finish",
+            "ignore_agent",
+            finish,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class AsyncCallbackManagerForChainRun(AsyncRunManager, ChainManagerMixin):
+    """Async callback manager for chain run."""
+
+    def get_child(self) -> AsyncCallbackManager:
+        """Get a child callback manager."""
+        manager = AsyncCallbackManager([], parent_run_id=self.run_id)
+        manager.set_handlers(self.inheritable_handlers)
+        return manager
+
+    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        """Run when chain ends running."""
+        await _ahandle_event(
+            self.handlers,
+            "on_chain_end",
+            "ignore_chain",
+            outputs,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_chain_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when chain errors."""
+        await _ahandle_event(
+            self.handlers,
+            "on_chain_error",
+            "ignore_chain",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+        """Run when agent action is received."""
+        await _ahandle_event(
+            self.handlers,
+            "on_agent_action",
+            "ignore_agent",
+            action,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
+        """Run when agent finish is received."""
+        await _ahandle_event(
+            self.handlers,
+            "on_agent_finish",
+            "ignore_agent",
+            finish,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class CallbackManagerForToolRun(RunManager, ToolManagerMixin):
+    """Callback manager for tool run."""
+
+    def get_child(self) -> CallbackManager:
+        """Get a child callback manager."""
+        manager = CallbackManager([], parent_run_id=self.run_id)
+        manager.set_handlers(self.inheritable_handlers)
+        return manager
+
+    def on_tool_end(
+        self,
+        output: str,
+        **kwargs: Any,
+    ) -> None:
+        """Run when tool ends running."""
+        _handle_event(
+            self.handlers,
+            "on_tool_end",
+            "ignore_agent",
+            output,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when tool errors."""
+        _handle_event(
+            self.handlers,
+            "on_tool_error",
+            "ignore_agent",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class AsyncCallbackManagerForToolRun(AsyncRunManager, ToolManagerMixin):
+    """Async callback manager for tool run."""
+
+    def get_child(self) -> AsyncCallbackManager:
+        """Get a child callback manager."""
+        manager = AsyncCallbackManager([], parent_run_id=self.run_id)
+        manager.set_handlers(self.inheritable_handlers)
+        return manager
+
+    async def on_tool_end(self, output: str, **kwargs: Any) -> None:
+        """Run when tool ends running."""
+        await _ahandle_event(
+            self.handlers,
+            "on_tool_end",
+            "ignore_agent",
+            output,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+    async def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        """Run when tool errors."""
+        await _ahandle_event(
+            self.handlers,
+            "on_tool_error",
+            "ignore_agent",
+            error,
+            run_id=self.run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+
+class CallbackManager(BaseCallbackManager):
+    """Callback manager that can be used to handle callbacks from langchain."""
+
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> CallbackManagerForLLMRun:
+        """Run when LLM starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        _handle_event(
+            self.handlers,
+            "on_llm_start",
+            "ignore_llm",
+            serialized,
+            prompts,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return CallbackManagerForLLMRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> CallbackManagerForChainRun:
+        """Run when chain starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        _handle_event(
+            self.handlers,
+            "on_chain_start",
+            "ignore_chain",
+            serialized,
+            inputs,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return CallbackManagerForChainRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        run_id: Optional[str] = None,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> CallbackManagerForToolRun:
+        """Run when tool starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        _handle_event(
+            self.handlers,
+            "on_tool_start",
+            "ignore_agent",
+            serialized,
+            input_str,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return CallbackManagerForToolRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    @classmethod
+    def configure(
+        cls,
+        inheritable_callbacks: Callbacks = None,
+        local_callbacks: Callbacks = None,
+        verbose: bool = False,
+    ) -> CallbackManager:
+        """Configure the callback manager."""
+        return _configure(cls, inheritable_callbacks, local_callbacks, verbose)
+
+
+class AsyncCallbackManager(BaseCallbackManager):
+    """Async callback manager that can be used to handle callbacks from LangChain."""
+
+    @property
+    def is_async(self) -> bool:
+        """Return whether the handler is async."""
+        return True
+
+    async def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncCallbackManagerForLLMRun:
+        """Run when LLM starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        await _ahandle_event(
+            self.handlers,
+            "on_llm_start",
+            "ignore_llm",
+            serialized,
+            prompts,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return AsyncCallbackManagerForLLMRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    async def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncCallbackManagerForChainRun:
+        """Run when chain starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        await _ahandle_event(
+            self.handlers,
+            "on_chain_start",
+            "ignore_chain",
+            serialized,
+            inputs,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return AsyncCallbackManagerForChainRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    async def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        run_id: Optional[str] = None,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncCallbackManagerForToolRun:
+        """Run when tool starts running."""
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+
+        await _ahandle_event(
+            self.handlers,
+            "on_tool_start",
+            "ignore_agent",
+            serialized,
+            input_str,
+            run_id=run_id,
+            parent_run_id=self.parent_run_id,
+            **kwargs,
+        )
+
+        return AsyncCallbackManagerForToolRun(
+            run_id, self.handlers, self.inheritable_handlers, self.parent_run_id
+        )
+
+    @classmethod
+    def configure(
+        cls,
+        inheritable_callbacks: Callbacks = None,
+        local_callbacks: Callbacks = None,
+        verbose: bool = False,
+    ) -> AsyncCallbackManager:
+        """Configure the callback manager."""
+        return _configure(cls, inheritable_callbacks, local_callbacks, verbose)
+
+
+T = TypeVar("T", CallbackManager, AsyncCallbackManager)
+
+
+def _configure(
+    callback_manager_cls: Type[T],
+    inheritable_callbacks: Callbacks = None,
+    local_callbacks: Callbacks = None,
+    verbose: bool = False,
+) -> T:
+    """Configure the callback manager."""
+    callback_manager = callback_manager_cls([])
+    if inheritable_callbacks or local_callbacks:
+        if isinstance(inheritable_callbacks, list) or inheritable_callbacks is None:
+            inheritable_callbacks_ = inheritable_callbacks or []
+            callback_manager = callback_manager_cls(
+                handlers=inheritable_callbacks_,
+                inheritable_handlers=inheritable_callbacks_,
+            )
+        else:
+            callback_manager = callback_manager_cls(
+                handlers=inheritable_callbacks.handlers,
+                inheritable_handlers=inheritable_callbacks.inheritable_handlers,
+                parent_run_id=inheritable_callbacks.parent_run_id,
+            )
+        callback_manager = copy.deepcopy(callback_manager)
+        local_handlers_ = (
+            local_callbacks
+            if isinstance(local_callbacks, list)
+            else (local_callbacks.handlers if local_callbacks else [])
+        )
+        for handler in local_handlers_:
+            callback_manager.add_handler(copy.deepcopy(handler), False)
+
+    tracer = tracing_callback_var.get()
+    open_ai = openai_callback_var.get()
+    tracing_enabled_ = (
+        os.environ.get("LANGCHAIN_TRACING") is not None or tracer is not None
+    )
+    if verbose or tracing_enabled_ or open_ai is not None:
+        if verbose and not any(
+            isinstance(handler, StdOutCallbackHandler)
+            for handler in callback_manager.handlers
+        ):
+            callback_manager.add_handler(StdOutCallbackHandler(), False)
+
+        if tracing_enabled_ and not any(
+            isinstance(handler, LangChainTracer)
+            for handler in callback_manager.handlers
+        ):
+            if tracer:
+                callback_manager.add_handler(copy.deepcopy(tracer), True)
+            else:
+                handler = LangChainTracer()
+                handler.load_default_session()
+                callback_manager.add_handler(handler, True)
+        if open_ai is not None and not any(
+            isinstance(handler, OpenAICallbackHandler)
+            for handler in callback_manager.handlers
+        ):
+            callback_manager.add_handler(open_ai, True)
+
+    return callback_manager

langchain/callbacks/shared.py

@@ -1,127 +0,0 @@
-"""A shared CallbackManager."""
-
-import threading
-from typing import Any, Dict, List, Union
-
-from langchain.callbacks.base import (
-    BaseCallbackHandler,
-    BaseCallbackManager,
-    CallbackManager,
-)
-from langchain.schema import AgentAction, AgentFinish, LLMResult
-
-
-class Singleton:
-    """A thread-safe singleton class that can be inherited from."""
-
-    _instance = None
-    _lock = threading.Lock()
-
-    def __new__(cls) -> Any:
-        """Create a new shared instance of the class."""
-        if cls._instance is None:
-            with cls._lock:
-                # Another thread could have created the instance
-                # before we acquired the lock. So check that the
-                # instance is still nonexistent.
-                if not cls._instance:
-                    cls._instance = super().__new__(cls)
-        return cls._instance
-
-
-class SharedCallbackManager(Singleton, BaseCallbackManager):
-    """A thread-safe singleton CallbackManager."""
-
-    _callback_manager: CallbackManager = CallbackManager(handlers=[])
-
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running."""
-        with self._lock:
-            self._callback_manager.on_llm_start(serialized, prompts, **kwargs)
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-        with self._lock:
-            self._callback_manager.on_llm_end(response, **kwargs)
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run when LLM generates a new token."""
-        with self._lock:
-            self._callback_manager.on_llm_new_token(token, **kwargs)
-
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        with self._lock:
-            self._callback_manager.on_llm_error(error, **kwargs)
-
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-        with self._lock:
-            self._callback_manager.on_chain_start(serialized, inputs, **kwargs)
-
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        """Run when chain ends running."""
-        with self._lock:
-            self._callback_manager.on_chain_end(outputs, **kwargs)
-
-    def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when chain errors."""
-        with self._lock:
-            self._callback_manager.on_chain_error(error, **kwargs)
-
-    def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-        with self._lock:
-            self._callback_manager.on_tool_start(serialized, input_str, **kwargs)
-
-    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
-        """Run on agent action."""
-        with self._lock:
-            self._callback_manager.on_agent_action(action, **kwargs)
-
-    def on_tool_end(self, output: str, **kwargs: Any) -> None:
-        """Run when tool ends running."""
-        with self._lock:
-            self._callback_manager.on_tool_end(output, **kwargs)
-
-    def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when tool errors."""
-        with self._lock:
-            self._callback_manager.on_tool_error(error, **kwargs)
-
-    def on_text(self, text: str, **kwargs: Any) -> None:
-        """Run on arbitrary text."""
-        with self._lock:
-            self._callback_manager.on_text(text, **kwargs)
-
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
-        """Run on agent end."""
-        with self._lock:
-            self._callback_manager.on_agent_finish(finish, **kwargs)
-
-    def add_handler(self, callback: BaseCallbackHandler) -> None:
-        """Add a callback to the callback manager."""
-        with self._lock:
-            self._callback_manager.add_handler(callback)
-
-    def remove_handler(self, callback: BaseCallbackHandler) -> None:
-        """Remove a callback from the callback manager."""
-        with self._lock:
-            self._callback_manager.remove_handler(callback)
-
-    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
-        """Set handlers as the only handlers on the callback manager."""
-        with self._lock:
-            self._callback_manager.handlers = handlers

langchain/callbacks/streamlit.py

@@ -10,6 +10,10 @@ from langchain.schema import AgentAction, AgentFinish, LLMResult
 class StreamlitCallbackHandler(BaseCallbackHandler):
     """Callback Handler that logs to streamlit."""
 
+    def __init__(self) -> None:
+        self.tokens_area = st.empty()
+        self.tokens_stream = ""
+
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
@@ -19,8 +23,9 @@ class StreamlitCallbackHandler(BaseCallbackHandler):
             st.write(prompt)
 
     def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Do nothing."""
-        pass
+        """Run on new LLM token. Only available when streaming is enabled."""
+        self.tokens_stream += token
+        self.tokens_area.write(self.tokens_stream)
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Do nothing."""

langchain/callbacks/tracers/__init__.py

@@ -1,12 +1,5 @@
 """Tracers that record execution of LangChain runs."""
 
-from langchain.callbacks.tracers.base import SharedTracer, Tracer
-from langchain.callbacks.tracers.langchain import BaseLangChainTracer
+from langchain.callbacks.tracers.langchain import LangChainTracer
 
-
-class SharedLangChainTracer(SharedTracer, BaseLangChainTracer):
-    """Shared tracer that records LangChain execution to LangChain endpoint."""
-
-
-class LangChainTracer(Tracer, BaseLangChainTracer):
-    """Tracer that records LangChain execution to LangChain endpoint."""
+__all__ = ["LangChainTracer"]

langchain/callbacks/tracers/base.py

@@ -1,14 +1,12 @@
 """Base interfaces for tracing runs."""
 from __future__ import annotations
 
-import threading
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
+from uuid import uuid4
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.callbacks.shared import Singleton
 from langchain.callbacks.tracers.schemas import (
     ChainRun,
     LLMRun,
@@ -16,7 +14,7 @@ from langchain.callbacks.tracers.schemas import (
     TracerSession,
     TracerSessionCreate,
 )
-from langchain.schema import AgentAction, AgentFinish, LLMResult
+from langchain.schema import LLMResult
 
 
 class TracerException(Exception):
@@ -26,13 +24,25 @@ class TracerException(Exception):
 class BaseTracer(BaseCallbackHandler, ABC):
     """Base interface for tracers."""
 
-    @abstractmethod
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.run_map: Dict[str, Union[LLMRun, ChainRun, ToolRun]] = {}
+        self.session: Optional[TracerSession] = None
+
+    @staticmethod
     def _add_child_run(
-        self,
         parent_run: Union[ChainRun, ToolRun],
         child_run: Union[LLMRun, ChainRun, ToolRun],
     ) -> None:
         """Add child run to a chain run or tool run."""
+        if isinstance(child_run, LLMRun):
+            parent_run.child_llm_runs.append(child_run)
+        elif isinstance(child_run, ChainRun):
+            parent_run.child_chain_runs.append(child_run)
+        elif isinstance(child_run, ToolRun):
+            parent_run.child_tool_runs.append(child_run)
+        else:
+            raise TracerException(f"Invalid run type: {type(child_run)}")
 
     @abstractmethod
     def _persist_run(self, run: Union[LLMRun, ChainRun, ToolRun]) -> None:
@@ -42,15 +52,11 @@ class BaseTracer(BaseCallbackHandler, ABC):
     def _persist_session(self, session: TracerSessionCreate) -> TracerSession:
         """Persist a tracing session."""
 
-    @abstractmethod
-    def _generate_id(self) -> Optional[Union[int, str]]:
-        """Generate an id for a run."""
-
     def new_session(self, name: Optional[str] = None, **kwargs: Any) -> TracerSession:
         """NOT thread safe, do not call this method from multiple threads."""
         session_create = TracerSessionCreate(name=name, extra=kwargs)
         session = self._persist_session(session_create)
-        self._session = session
+        self.session = session
         return session
 
     @abstractmethod
@@ -61,283 +67,232 @@ class BaseTracer(BaseCallbackHandler, ABC):
     def load_default_session(self) -> TracerSession:
         """Load the default tracing session and set it as the Tracer's session."""
 
-    @property
-    @abstractmethod
-    def _stack(self) -> List[Union[LLMRun, ChainRun, ToolRun]]:
-        """Get the tracer stack."""
-
-    @property
-    @abstractmethod
-    def _execution_order(self) -> int:
-        """Get the execution order for a run."""
-
-    @_execution_order.setter
-    @abstractmethod
-    def _execution_order(self, value: int) -> None:
-        """Set the execution order for a run."""
-
-    @property
-    @abstractmethod
-    def _session(self) -> Optional[TracerSession]:
-        """Get the tracing session."""
-
-    @_session.setter
-    @abstractmethod
-    def _session(self, value: TracerSession) -> None:
-        """Set the tracing session."""
-
     def _start_trace(self, run: Union[LLMRun, ChainRun, ToolRun]) -> None:
         """Start a trace for a run."""
-        self._execution_order += 1
-
-        if self._stack:
-            if not (
-                isinstance(self._stack[-1], ChainRun)
-                or isinstance(self._stack[-1], ToolRun)
-            ):
+        if run.parent_uuid:
+            parent_run = self.run_map[run.parent_uuid]
+            if parent_run:
+                if isinstance(parent_run, LLMRun):
+                    raise TracerException(
+                        "Cannot add child run to an LLM run. "
+                        "LLM runs are not allowed to have children."
+                    )
+                self._add_child_run(parent_run, run)
+            else:
                 raise TracerException(
-                    f"Nested {run.__class__.__name__} can only be"
-                    f" logged inside a ChainRun or ToolRun"
+                    f"Parent run with UUID {run.parent_uuid} not found."
                 )
-            self._add_child_run(self._stack[-1], run)
-        self._stack.append(run)
 
-    def _end_trace(self) -> None:
+        self.run_map[run.uuid] = run
+
+    def _end_trace(self, run: Union[LLMRun, ChainRun, ToolRun]) -> None:
         """End a trace for a run."""
-        run = self._stack.pop()
-        if not self._stack:
-            self._execution_order = 1
+        if not run.parent_uuid:
             self._persist_run(run)
+        else:
+            parent_run = self.run_map.get(run.parent_uuid)
+            if parent_run is None:
+                raise TracerException(
+                    f"Parent run with UUID {run.parent_uuid} not found."
+                )
+            if isinstance(parent_run, LLMRun):
+                raise TracerException("LLM Runs are not allowed to have children. ")
+            if run.child_execution_order > parent_run.child_execution_order:
+                parent_run.child_execution_order = run.child_execution_order
+        self.run_map.pop(run.uuid)
+
+    def _get_execution_order(self, parent_run_id: Optional[str] = None) -> int:
+        """Get the execution order for a run."""
+        if parent_run_id is None:
+            return 1
+
+        parent_run = self.run_map.get(parent_run_id)
+        if parent_run is None:
+            raise TracerException(f"Parent run with UUID {parent_run_id} not found.")
+
+        if isinstance(parent_run, LLMRun):
+            raise TracerException("LLM Runs are not allowed to have children. ")
+
+        return parent_run.child_execution_order + 1
 
     def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Start a trace for an LLM run."""
-        if self._session is None:
-            raise TracerException(
-                "Initialize a session with `new_session()` before starting a trace."
-            )
+        if self.session is None:
+            self.session = self.load_default_session()
 
+        if run_id is None:
+            run_id = str(uuid4())
+
+        execution_order = self._get_execution_order(parent_run_id)
         llm_run = LLMRun(
+            uuid=run_id,
+            parent_uuid=parent_run_id,
             serialized=serialized,
             prompts=prompts,
             extra=kwargs,
             start_time=datetime.utcnow(),
-            execution_order=self._execution_order,
-            session_id=self._session.id,
-            id=self._generate_id(),
+            execution_order=execution_order,
+            child_execution_order=execution_order,
+            session_id=self.session.id,
         )
         self._start_trace(llm_run)
 
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Handle a new token for an LLM run."""
-        pass
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+    def on_llm_end(self, response: LLMResult, *, run_id: str, **kwargs: Any) -> None:
         """End a trace for an LLM run."""
-        if not self._stack or not isinstance(self._stack[-1], LLMRun):
+        if not run_id:
+            raise TracerException("No run_id provided for on_llm_end callback.")
+
+        llm_run = self.run_map.get(run_id)
+        if llm_run is None or not isinstance(llm_run, LLMRun):
             raise TracerException("No LLMRun found to be traced")
 
-        self._stack[-1].end_time = datetime.utcnow()
-        self._stack[-1].response = response
-
-        self._end_trace()
+        llm_run.response = response
+        llm_run.end_time = datetime.utcnow()
+        self._end_trace(llm_run)
 
     def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        **kwargs: Any,
     ) -> None:
         """Handle an error for an LLM run."""
-        if not self._stack or not isinstance(self._stack[-1], LLMRun):
+        if not run_id:
+            raise TracerException("No run_id provided for on_llm_error callback.")
+
+        llm_run = self.run_map.get(run_id)
+        if llm_run is None or not isinstance(llm_run, LLMRun):
             raise TracerException("No LLMRun found to be traced")
 
-        self._stack[-1].error = repr(error)
-        self._stack[-1].end_time = datetime.utcnow()
-
-        self._end_trace()
+        llm_run.error = repr(error)
+        llm_run.end_time = datetime.utcnow()
+        self._end_trace(llm_run)
 
     def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Start a trace for a chain run."""
-        if self._session is None:
-            raise TracerException(
-                "Initialize a session with `new_session()` before starting a trace."
-            )
+        if self.session is None:
+            self.session = self.load_default_session()
 
+        execution_order = self._get_execution_order(parent_run_id)
         chain_run = ChainRun(
+            uuid=run_id,
+            parent_uuid=parent_run_id,
             serialized=serialized,
             inputs=inputs,
             extra=kwargs,
             start_time=datetime.utcnow(),
-            execution_order=self._execution_order,
+            execution_order=execution_order,
+            child_execution_order=execution_order,
             child_runs=[],
-            session_id=self._session.id,
-            id=self._generate_id(),
+            session_id=self.session.id,
         )
         self._start_trace(chain_run)
 
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+    def on_chain_end(
+        self, outputs: Dict[str, Any], *, run_id: str, **kwargs: Any
+    ) -> None:
         """End a trace for a chain run."""
-        if not self._stack or not isinstance(self._stack[-1], ChainRun):
+        chain_run = self.run_map.get(run_id)
+        if chain_run is None or not isinstance(chain_run, ChainRun):
             raise TracerException("No ChainRun found to be traced")
 
-        self._stack[-1].end_time = datetime.utcnow()
-        self._stack[-1].outputs = outputs
-
-        self._end_trace()
+        chain_run.outputs = outputs
+        chain_run.end_time = datetime.utcnow()
+        self._end_trace(chain_run)
 
     def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        **kwargs: Any,
     ) -> None:
         """Handle an error for a chain run."""
-        if not self._stack or not isinstance(self._stack[-1], ChainRun):
+        chain_run = self.run_map.get(run_id)
+        if chain_run is None or not isinstance(chain_run, ChainRun):
             raise TracerException("No ChainRun found to be traced")
 
-        self._stack[-1].end_time = datetime.utcnow()
-        self._stack[-1].error = repr(error)
-
-        self._end_trace()
+        chain_run.error = repr(error)
+        chain_run.end_time = datetime.utcnow()
+        self._end_trace(chain_run)
 
     def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: str,
+        parent_run_id: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Start a trace for a tool run."""
-        if self._session is None:
-            raise TracerException(
-                "Initialize a session with `new_session()` before starting a trace."
-            )
+        if self.session is None:
+            self.session = self.load_default_session()
 
+        execution_order = self._get_execution_order(parent_run_id)
         tool_run = ToolRun(
+            uuid=run_id,
+            parent_uuid=parent_run_id,
             serialized=serialized,
             # TODO: this is duplicate info as above, not needed.
             action=str(serialized),
             tool_input=input_str,
             extra=kwargs,
             start_time=datetime.utcnow(),
-            execution_order=self._execution_order,
+            execution_order=execution_order,
+            child_execution_order=execution_order,
             child_runs=[],
-            session_id=self._session.id,
-            id=self._generate_id(),
+            session_id=self.session.id,
         )
         self._start_trace(tool_run)
 
-    def on_tool_end(self, output: str, **kwargs: Any) -> None:
+    def on_tool_end(self, output: str, *, run_id: str, **kwargs: Any) -> None:
         """End a trace for a tool run."""
-        if not self._stack or not isinstance(self._stack[-1], ToolRun):
+        tool_run = self.run_map.get(run_id)
+        if tool_run is None or not isinstance(tool_run, ToolRun):
             raise TracerException("No ToolRun found to be traced")
 
-        self._stack[-1].end_time = datetime.utcnow()
-        self._stack[-1].output = output
-
-        self._end_trace()
+        tool_run.output = output
+        tool_run.end_time = datetime.utcnow()
+        self._end_trace(tool_run)
 
     def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: str,
+        **kwargs: Any,
     ) -> None:
         """Handle an error for a tool run."""
-        if not self._stack or not isinstance(self._stack[-1], ToolRun):
+        tool_run = self.run_map.get(run_id)
+        if tool_run is None or not isinstance(tool_run, ToolRun):
             raise TracerException("No ToolRun found to be traced")
 
-        self._stack[-1].end_time = datetime.utcnow()
-        self._stack[-1].error = repr(error)
+        tool_run.error = repr(error)
+        tool_run.end_time = datetime.utcnow()
+        self._end_trace(tool_run)
 
-        self._end_trace()
+    def __deepcopy__(self, memo: dict) -> BaseTracer:
+        """Deepcopy the tracer."""
+        return self
 
-    def on_text(self, text: str, **kwargs: Any) -> None:
-        """Handle a text message."""
-        pass
-
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
-        """Handle an agent finish message."""
-        pass
-
-    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
-        """Do nothing."""
-        pass
-
-
-class Tracer(BaseTracer, ABC):
-    """A non-thread safe implementation of the BaseTracer interface."""
-
-    def __init__(self) -> None:
-        """Initialize a tracer."""
-        self._tracer_stack: List[Union[LLMRun, ChainRun, ToolRun]] = []
-        self._tracer_execution_order = 1
-        self._tracer_session: Optional[TracerSession] = None
-
-    @property
-    def _stack(self) -> List[Union[LLMRun, ChainRun, ToolRun]]:
-        """Get the tracer stack."""
-        return self._tracer_stack
-
-    @property
-    def _execution_order(self) -> int:
-        """Get the execution order for a run."""
-        return self._tracer_execution_order
-
-    @_execution_order.setter
-    def _execution_order(self, value: int) -> None:
-        """Set the execution order for a run."""
-        self._tracer_execution_order = value
-
-    @property
-    def _session(self) -> Optional[TracerSession]:
-        """Get the tracing session."""
-        return self._tracer_session
-
-    @_session.setter
-    def _session(self, value: TracerSession) -> None:
-        """Set the tracing session."""
-        if self._stack:
-            raise TracerException(
-                "Cannot set a session while a trace is being recorded"
-            )
-        self._tracer_session = value
-
-
-@dataclass
-class TracerStack(threading.local):
-    """A stack of runs used for logging."""
-
-    stack: List[Union[LLMRun, ChainRun, ToolRun]] = field(default_factory=list)
-    execution_order: int = 1
-
-
-class SharedTracer(Singleton, BaseTracer, ABC):
-    """A thread-safe Singleton implementation of BaseTracer."""
-
-    _tracer_stack = TracerStack()
-    _tracer_session = None
-
-    @property
-    def _stack(self) -> List[Union[LLMRun, ChainRun, ToolRun]]:
-        """Get the tracer stack."""
-        return self._tracer_stack.stack
-
-    @property
-    def _execution_order(self) -> int:
-        """Get the execution order for a run."""
-        return self._tracer_stack.execution_order
-
-    @_execution_order.setter
-    def _execution_order(self, value: int) -> None:
-        """Set the execution order for a run."""
-        self._tracer_stack.execution_order = value
-
-    @property
-    def _session(self) -> Optional[TracerSession]:
-        """Get the tracing session."""
-        return self._tracer_session
-
-    @_session.setter
-    def _session(self, value: TracerSession) -> None:
-        """Set the tracing session."""
-        with self._lock:
-            # TODO: currently, we are only checking current thread's stack.
-            #  Need to make sure that we are not in the middle of a trace
-            #  in any thread.
-            if self._stack:
-                raise TracerException(
-                    "Cannot set a session while a trace is being recorded"
-                )
-            self._tracer_session = value
+    def __copy__(self) -> BaseTracer:
+        """Copy the tracer."""
+        return self

langchain/callbacks/tracers/langchain.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 
 import logging
 import os
-from abc import ABC
 from typing import Any, Dict, Optional, Union
 
 import requests
@@ -18,14 +17,17 @@ from langchain.callbacks.tracers.schemas import (
 )
 
 
-class BaseLangChainTracer(BaseTracer, ABC):
+class LangChainTracer(BaseTracer):
     """An implementation of the SharedTracer that POSTS to the langchain endpoint."""
 
-    always_verbose: bool = True
-    _endpoint: str = os.getenv("LANGCHAIN_ENDPOINT", "http://localhost:8000")
-    _headers: Dict[str, Any] = {"Content-Type": "application/json"}
-    if os.getenv("LANGCHAIN_API_KEY"):
-        _headers["x-api-key"] = os.getenv("LANGCHAIN_API_KEY")
+    def __init__(self, session_name: str = "default", **kwargs: Any) -> None:
+        """Initialize the LangChain tracer."""
+        super().__init__(**kwargs)
+        self._endpoint: str = os.getenv("LANGCHAIN_ENDPOINT", "http://localhost:8000")
+        self._headers: Dict[str, Any] = {"Content-Type": "application/json"}
+        if os.getenv("LANGCHAIN_API_KEY"):
+            self._headers["x-api-key"] = os.getenv("LANGCHAIN_API_KEY")
+        self.session = self.load_session(session_name)
 
     def _persist_run(self, run: Union[LLMRun, ChainRun, ToolRun]) -> None:
         """Persist a run."""
@@ -59,54 +61,29 @@ class BaseLangChainTracer(BaseTracer, ABC):
             session = TracerSession(id=1, **session_create.dict())
         return session
 
-    def load_session(self, session_name: str) -> TracerSession:
+    def _load_session(self, session_name: Optional[str] = None) -> TracerSession:
         """Load a session from the tracer."""
         try:
-            r = requests.get(
-                f"{self._endpoint}/sessions?name={session_name}",
-                headers=self._headers,
-            )
+            url = f"{self._endpoint}/sessions"
+            if session_name:
+                url += f"?name={session_name}"
+            r = requests.get(url, headers=self._headers)
+
             tracer_session = TracerSession(**r.json()[0])
-            self._session = tracer_session
-            return tracer_session
         except Exception as e:
+            session_type = "default" if not session_name else session_name
             logging.warning(
-                f"Failed to load session {session_name}, using empty session: {e}"
+                f"Failed to load {session_type} session, using empty session: {e}"
             )
             tracer_session = TracerSession(id=1)
-            self._session = tracer_session
-            return tracer_session
+
+        self.session = tracer_session
+        return tracer_session
+
+    def load_session(self, session_name: str) -> TracerSession:
+        """Load a session with the given name from the tracer."""
+        return self._load_session(session_name)
 
     def load_default_session(self) -> TracerSession:
         """Load the default tracing session and set it as the Tracer's session."""
-        try:
-            r = requests.get(
-                f"{self._endpoint}/sessions",
-                headers=self._headers,
-            )
-            # Use the first session result
-            tracer_session = TracerSession(**r.json()[0])
-            self._session = tracer_session
-            return tracer_session
-        except Exception as e:
-            logging.warning(f"Failed to default session, using empty session: {e}")
-            tracer_session = TracerSession(id=1)
-            self._session = tracer_session
-            return tracer_session
-
-    def _add_child_run(
-        self,
-        parent_run: Union[ChainRun, ToolRun],
-        child_run: Union[LLMRun, ChainRun, ToolRun],
-    ) -> None:
-        """Add child run to a chain run or tool run."""
-        if isinstance(child_run, LLMRun):
-            parent_run.child_llm_runs.append(child_run)
-        elif isinstance(child_run, ChainRun):
-            parent_run.child_chain_runs.append(child_run)
-        else:
-            parent_run.child_tool_runs.append(child_run)
-
-    def _generate_id(self) -> Optional[Union[int, str]]:
-        """Generate an id for a run."""
-        return None
+        return self._load_session("default")

langchain/callbacks/tracers/schemas.py

@@ -32,11 +32,13 @@ class TracerSession(TracerSessionBase):
 class BaseRun(BaseModel):
     """Base class for Run."""
 
-    id: Optional[Union[int, str]] = None
+    uuid: str
+    parent_uuid: Optional[str] = None
     start_time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
     end_time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
     extra: Optional[Dict[str, Any]] = None
     execution_order: int
+    child_execution_order: int
     serialized: Dict[str, Any]
     session_id: int
     error: Optional[str] = None
@@ -57,7 +59,6 @@ class ChainRun(BaseRun):
     child_llm_runs: List[LLMRun] = Field(default_factory=list)
     child_chain_runs: List[ChainRun] = Field(default_factory=list)
     child_tool_runs: List[ToolRun] = Field(default_factory=list)
-    child_runs: List[Union[LLMRun, ChainRun, ToolRun]] = Field(default_factory=list)
 
 
 class ToolRun(BaseRun):
@@ -69,7 +70,6 @@ class ToolRun(BaseRun):
     child_llm_runs: List[LLMRun] = Field(default_factory=list)
     child_chain_runs: List[ChainRun] = Field(default_factory=list)
     child_tool_runs: List[ToolRun] = Field(default_factory=list)
-    child_runs: List[Union[LLMRun, ChainRun, ToolRun]] = Field(default_factory=list)
 
 
 ChainRun.update_forward_refs()

langchain/chains/api/base.py

@@ -5,12 +5,16 @@ from typing import Any, Dict, List, Optional
 
 from pydantic import Field, root_validator
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
 from langchain.chains.api.prompt import API_RESPONSE_PROMPT, API_URL_PROMPT
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.prompts import BasePromptTemplate
 from langchain.requests import TextRequestsWrapper
-from langchain.schema import BaseLanguageModel
 
 
 class APIChain(Chain):
@@ -61,16 +65,21 @@ class APIChain(Chain):
             )
         return values
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.question_key]
         api_url = self.api_request_chain.predict(
-            question=question, api_docs=self.api_docs
-        )
-        self.callback_manager.on_text(
-            api_url, color="green", end="\n", verbose=self.verbose
+            question=question,
+            api_docs=self.api_docs,
+            callbacks=_run_manager.get_child(),
         )
+        _run_manager.on_text(api_url, color="green", end="\n", verbose=self.verbose)
         api_response = self.requests_wrapper.get(api_url)
-        self.callback_manager.on_text(
+        _run_manager.on_text(
             api_response, color="yellow", end="\n", verbose=self.verbose
         )
         answer = self.api_answer_chain.predict(
@@ -78,19 +87,27 @@ class APIChain(Chain):
             api_docs=self.api_docs,
             api_url=api_url,
             api_response=api_response,
+            callbacks=_run_manager.get_child(),
         )
         return {self.output_key: answer}
 
-    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.question_key]
         api_url = await self.api_request_chain.apredict(
-            question=question, api_docs=self.api_docs
+            question=question,
+            api_docs=self.api_docs,
+            callbacks=_run_manager.get_child(),
         )
-        self.callback_manager.on_text(
+        await _run_manager.on_text(
             api_url, color="green", end="\n", verbose=self.verbose
         )
         api_response = await self.requests_wrapper.aget(api_url)
-        self.callback_manager.on_text(
+        await _run_manager.on_text(
             api_response, color="yellow", end="\n", verbose=self.verbose
         )
         answer = await self.api_answer_chain.apredict(
@@ -98,6 +115,7 @@ class APIChain(Chain):
             api_docs=self.api_docs,
             api_url=api_url,
             api_response=api_response,
+            callbacks=_run_manager.get_child(),
         )
         return {self.output_key: answer}
 

langchain/chains/api/openapi/chain.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, List, NamedTuple, Optional, cast
 from pydantic import BaseModel, Field
 from requests import Response
 
+from langchain.callbacks.manager import CallbackManagerForChainRun, Callbacks
 from langchain.chains.api.openapi.requests_chain import APIRequesterChain
 from langchain.chains.api.openapi.response_chain import APIResponderChain
 from langchain.chains.base import Chain
@@ -106,16 +107,21 @@ class OpenAPIEndpointChain(Chain, BaseModel):
         else:
             return {self.output_key: output}
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         intermediate_steps = {}
         instructions = inputs[self.instructions_key]
         instructions = instructions[: self.max_text_length]
         _api_arguments = self.api_request_chain.predict_and_parse(
-            instructions=instructions
+            instructions=instructions, callbacks=_run_manager.get_child()
         )
         api_arguments = cast(str, _api_arguments)
         intermediate_steps["request_args"] = api_arguments
-        self.callback_manager.on_text(
+        _run_manager.on_text(
             api_arguments, color="green", end="\n", verbose=self.verbose
         )
         if api_arguments.startswith("ERROR"):
@@ -141,18 +147,17 @@ class OpenAPIEndpointChain(Chain, BaseModel):
             response_text = f"Error with message {str(e)}"
         response_text = response_text[: self.max_text_length]
         intermediate_steps["response_text"] = response_text
-        self.callback_manager.on_text(
+        _run_manager.on_text(
             response_text, color="blue", end="\n", verbose=self.verbose
         )
         if self.api_response_chain is not None:
             _answer = self.api_response_chain.predict_and_parse(
                 response=response_text,
                 instructions=instructions,
+                callbacks=_run_manager.get_child(),
             )
             answer = cast(str, _answer)
-            self.callback_manager.on_text(
-                answer, color="yellow", end="\n", verbose=self.verbose
-            )
+            _run_manager.on_text(answer, color="yellow", end="\n", verbose=self.verbose)
             return self._get_output(answer, intermediate_steps)
         else:
             return self._get_output(response_text, intermediate_steps)
@@ -188,6 +193,7 @@ class OpenAPIEndpointChain(Chain, BaseModel):
         verbose: bool = False,
         return_intermediate_steps: bool = False,
         raw_response: bool = False,
+        callbacks: Callbacks = None,
         **kwargs: Any
         # TODO: Handle async
     ) -> "OpenAPIEndpointChain":
@@ -198,12 +204,17 @@ class OpenAPIEndpointChain(Chain, BaseModel):
             path_params=operation.path_params,
         )
         requests_chain = APIRequesterChain.from_llm_and_typescript(
-            llm, typescript_definition=operation.to_typescript(), verbose=verbose
+            llm,
+            typescript_definition=operation.to_typescript(),
+            verbose=verbose,
+            callbacks=callbacks,
         )
         if raw_response:
             response_chain = None
         else:
-            response_chain = APIResponderChain.from_llm(llm, verbose=verbose)
+            response_chain = APIResponderChain.from_llm(
+                llm, verbose=verbose, callbacks=callbacks
+            )
         _requests = requests or Requests()
         return cls(
             api_request_chain=requests_chain,
@@ -213,5 +224,6 @@ class OpenAPIEndpointChain(Chain, BaseModel):
             param_mapping=param_mapping,
             verbose=verbose,
             return_intermediate_steps=return_intermediate_steps,
+            callbacks=callbacks,
             **kwargs,
         )

langchain/chains/api/openapi/requests_chain.py

@@ -2,6 +2,7 @@
 
 import json
 import re
+from typing import Any
 
 from langchain.chains.api.openapi.prompts import REQUEST_TEMPLATE
 from langchain.chains.llm import LLMChain
@@ -36,7 +37,11 @@ class APIRequesterChain(LLMChain):
 
     @classmethod
     def from_llm_and_typescript(
-        cls, llm: BaseLLM, typescript_definition: str, verbose: bool = True
+        cls,
+        llm: BaseLLM,
+        typescript_definition: str,
+        verbose: bool = True,
+        **kwargs: Any,
     ) -> LLMChain:
         """Get the request parser."""
         output_parser = APIRequesterOutputParser()
@@ -46,4 +51,4 @@ class APIRequesterChain(LLMChain):
             partial_variables={"schema": typescript_definition},
             input_variables=["instructions"],
         )
-        return cls(prompt=prompt, llm=llm, verbose=verbose)
+        return cls(prompt=prompt, llm=llm, verbose=verbose, **kwargs)

langchain/chains/api/openapi/response_chain.py

@@ -2,6 +2,7 @@
 
 import json
 import re
+from typing import Any
 
 from langchain.chains.api.openapi.prompts import RESPONSE_TEMPLATE
 from langchain.chains.llm import LLMChain
@@ -35,7 +36,7 @@ class APIResponderChain(LLMChain):
     """Get the response parser."""
 
     @classmethod
-    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
+    def from_llm(cls, llm: BaseLLM, verbose: bool = True, **kwargs: Any) -> LLMChain:
         """Get the response parser."""
         output_parser = APIResponderOutputParser()
         prompt = PromptTemplate(
@@ -43,4 +44,4 @@ class APIResponderChain(LLMChain):
             output_parser=output_parser,
             input_variables=["response", "instructions"],
         )
-        return cls(prompt=prompt, llm=llm, verbose=verbose)
+        return cls(prompt=prompt, llm=llm, verbose=verbose, **kwargs)

langchain/chains/base.py

@@ -1,15 +1,23 @@
 """Base interface that all chains should implement."""
+import inspect
 import json
+import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
 import yaml
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, root_validator, validator
 
 import langchain
-from langchain.callbacks import get_callback_manager
 from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import (
+    AsyncCallbackManager,
+    AsyncCallbackManagerForChainRun,
+    CallbackManager,
+    CallbackManagerForChainRun,
+    Callbacks,
+)
 from langchain.schema import BaseMemory
 
 
@@ -21,9 +29,8 @@ class Chain(BaseModel, ABC):
     """Base interface that all chains should implement."""
 
     memory: Optional[BaseMemory] = None
-    callback_manager: BaseCallbackManager = Field(
-        default_factory=get_callback_manager, exclude=True
-    )
+    callbacks: Callbacks = None
+    callback_manager: Optional[BaseCallbackManager] = None
     verbose: bool = Field(
         default_factory=_get_verbosity
     )  # Whether to print the response text
@@ -37,15 +44,16 @@ class Chain(BaseModel, ABC):
     def _chain_type(self) -> str:
         raise NotImplementedError("Saving not supported for this chain type.")
 
-    @validator("callback_manager", pre=True, always=True)
-    def set_callback_manager(
-        cls, callback_manager: Optional[BaseCallbackManager]
-    ) -> BaseCallbackManager:
-        """If callback manager is None, set it.
-
-        This allows users to pass in None as callback manager, which is a nice UX.
-        """
-        return callback_manager or get_callback_manager()
+    @root_validator()
+    def raise_deprecation(cls, values: Dict) -> Dict:
+        """Raise deprecation warning if callback_manager is used."""
+        if values.get("callback_manager") is not None:
+            warnings.warn(
+                "callback_manager is deprecated. Please use callbacks instead.",
+                DeprecationWarning,
+            )
+            values["callbacks"] = values.pop("callback_manager", None)
+        return values
 
     @validator("verbose", pre=True, always=True)
     def set_verbose(cls, verbose: Optional[bool]) -> bool:
@@ -82,15 +90,26 @@ class Chain(BaseModel, ABC):
             )
 
     @abstractmethod
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         """Run the logic of this chain and return the output."""
 
-    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         """Run the logic of this chain and return the output."""
         raise NotImplementedError("Async call not supported for this chain type.")
 
     def __call__(
-        self, inputs: Union[Dict[str, Any], Any], return_only_outputs: bool = False
+        self,
+        inputs: Union[Dict[str, Any], Any],
+        return_only_outputs: bool = False,
+        callbacks: Callbacks = None,
     ) -> Dict[str, Any]:
         """Run the logic of this chain and add to output if desired.
 
@@ -104,21 +123,31 @@ class Chain(BaseModel, ABC):
 
         """
         inputs = self.prep_inputs(inputs)
-        self.callback_manager.on_chain_start(
+        callback_manager = CallbackManager.configure(
+            callbacks, self.callbacks, self.verbose
+        )
+        new_arg_supported = inspect.signature(self._call).parameters.get("run_manager")
+        run_manager = callback_manager.on_chain_start(
             {"name": self.__class__.__name__},
             inputs,
-            verbose=self.verbose,
         )
         try:
-            outputs = self._call(inputs)
+            outputs = (
+                self._call(inputs, run_manager=run_manager)
+                if new_arg_supported
+                else self._call(inputs)
+            )
         except (KeyboardInterrupt, Exception) as e:
-            self.callback_manager.on_chain_error(e, verbose=self.verbose)
+            run_manager.on_chain_error(e)
             raise e
-        self.callback_manager.on_chain_end(outputs, verbose=self.verbose)
+        run_manager.on_chain_end(outputs)
         return self.prep_outputs(inputs, outputs, return_only_outputs)
 
     async def acall(
-        self, inputs: Union[Dict[str, Any], Any], return_only_outputs: bool = False
+        self,
+        inputs: Union[Dict[str, Any], Any],
+        return_only_outputs: bool = False,
+        callbacks: Callbacks = None,
     ) -> Dict[str, Any]:
         """Run the logic of this chain and add to output if desired.
 
@@ -132,30 +161,24 @@ class Chain(BaseModel, ABC):
 
         """
         inputs = self.prep_inputs(inputs)
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_chain_start(
-                {"name": self.__class__.__name__},
-                inputs,
-                verbose=self.verbose,
-            )
-        else:
-            self.callback_manager.on_chain_start(
-                {"name": self.__class__.__name__},
-                inputs,
-                verbose=self.verbose,
-            )
+        callback_manager = AsyncCallbackManager.configure(
+            callbacks, self.callbacks, self.verbose
+        )
+        new_arg_supported = inspect.signature(self._acall).parameters.get("run_manager")
+        run_manager = await callback_manager.on_chain_start(
+            {"name": self.__class__.__name__},
+            inputs,
+        )
         try:
-            outputs = await self._acall(inputs)
+            outputs = (
+                await self._acall(inputs, run_manager=run_manager)
+                if new_arg_supported
+                else await self._acall(inputs)
+            )
         except (KeyboardInterrupt, Exception) as e:
-            if self.callback_manager.is_async:
-                await self.callback_manager.on_chain_error(e, verbose=self.verbose)
-            else:
-                self.callback_manager.on_chain_error(e, verbose=self.verbose)
+            await run_manager.on_chain_error(e)
             raise e
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_chain_end(outputs, verbose=self.verbose)
-        else:
-            self.callback_manager.on_chain_end(outputs, verbose=self.verbose)
+        await run_manager.on_chain_end(outputs)
         return self.prep_outputs(inputs, outputs, return_only_outputs)
 
     def prep_outputs(
@@ -195,11 +218,13 @@ class Chain(BaseModel, ABC):
         self._validate_inputs(inputs)
         return inputs
 
-    def apply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]:
+    def apply(
+        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
+    ) -> List[Dict[str, str]]:
         """Call the chain on all inputs in the list."""
-        return [self(inputs) for inputs in input_list]
+        return [self(inputs, callbacks=callbacks) for inputs in input_list]
 
-    def run(self, *args: Any, **kwargs: Any) -> str:
+    def run(self, *args: Any, callbacks: Callbacks = None, **kwargs: Any) -> str:
         """Run the chain as text in, text out or multiple variables, text out."""
         if len(self.output_keys) != 1:
             raise ValueError(
@@ -210,17 +235,17 @@ class Chain(BaseModel, ABC):
         if args and not kwargs:
             if len(args) != 1:
                 raise ValueError("`run` supports only one positional argument.")
-            return self(args[0])[self.output_keys[0]]
+            return self(args[0], callbacks=callbacks)[self.output_keys[0]]
 
         if kwargs and not args:
-            return self(kwargs)[self.output_keys[0]]
+            return self(kwargs, callbacks=callbacks)[self.output_keys[0]]
 
         raise ValueError(
             f"`run` supported with either positional arguments or keyword arguments"
             f" but not both. Got args: {args} and kwargs: {kwargs}."
         )
 
-    async def arun(self, *args: Any, **kwargs: Any) -> str:
+    async def arun(self, *args: Any, callbacks: Callbacks = None, **kwargs: Any) -> str:
         """Run the chain as text in, text out or multiple variables, text out."""
         if len(self.output_keys) != 1:
             raise ValueError(
@@ -231,10 +256,10 @@ class Chain(BaseModel, ABC):
         if args and not kwargs:
             if len(args) != 1:
                 raise ValueError("`run` supports only one positional argument.")
-            return (await self.acall(args[0]))[self.output_keys[0]]
+            return (await self.acall(args[0], callbacks=callbacks))[self.output_keys[0]]
 
         if kwargs and not args:
-            return (await self.acall(kwargs))[self.output_keys[0]]
+            return (await self.acall(kwargs, callbacks=callbacks))[self.output_keys[0]]
 
         raise ValueError(
             f"`run` supported with either positional arguments or keyword arguments"

langchain/chains/combine_documents/base.py

@@ -5,6 +5,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic import Field
 
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
 from langchain.chains.base import Chain
 from langchain.docstore.document import Document
 from langchain.prompts.base import BasePromptTemplate
@@ -68,19 +72,33 @@ class BaseCombineDocumentsChain(Chain, ABC):
     ) -> Tuple[str, dict]:
         """Combine documents into a single string asynchronously."""
 
-    def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, List[Document]],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         docs = inputs[self.input_key]
         # Other keys are assumed to be needed for LLM prediction
         other_keys = {k: v for k, v in inputs.items() if k != self.input_key}
-        output, extra_return_dict = self.combine_docs(docs, **other_keys)
+        output, extra_return_dict = self.combine_docs(
+            docs, callbacks=_run_manager.get_child(), **other_keys
+        )
         extra_return_dict[self.output_key] = output
         return extra_return_dict
 
-    async def _acall(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+    async def _acall(
+        self,
+        inputs: Dict[str, List[Document]],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
         docs = inputs[self.input_key]
         # Other keys are assumed to be needed for LLM prediction
         other_keys = {k: v for k, v in inputs.items() if k != self.input_key}
-        output, extra_return_dict = await self.acombine_docs(docs, **other_keys)
+        output, extra_return_dict = await self.acombine_docs(
+            docs, callbacks=_run_manager.get_child(), **other_keys
+        )
         extra_return_dict[self.output_key] = output
         return extra_return_dict
 
@@ -108,10 +126,17 @@ class AnalyzeDocumentChain(Chain):
         """
         return self.combine_docs_chain.output_keys
 
-    def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, str],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         document = inputs[self.input_key]
         docs = self.text_splitter.create_documents([document])
         # Other keys are assumed to be needed for LLM prediction
-        other_keys = {k: v for k, v in inputs.items() if k != self.input_key}
+        other_keys: Dict = {k: v for k, v in inputs.items() if k != self.input_key}
         other_keys[self.combine_docs_chain.input_key] = docs
-        return self.combine_docs_chain(other_keys, return_only_outputs=True)
+        return self.combine_docs_chain(
+            other_keys, return_only_outputs=True, callbacks=_run_manager.get_child()
+        )

langchain/chains/combine_documents/map_reduce.py

@@ -6,6 +6,7 @@ from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple
 
 from pydantic import Extra, root_validator
 
+from langchain.callbacks.manager import Callbacks
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.docstore.document import Document
@@ -129,7 +130,11 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain):
             return self.combine_document_chain
 
     def combine_docs(
-        self, docs: List[Document], token_max: int = 3000, **kwargs: Any
+        self,
+        docs: List[Document],
+        token_max: int = 3000,
+        callbacks: Callbacks = None,
+        **kwargs: Any,
     ) -> Tuple[str, dict]:
         """Combine documents in a map reduce manner.
 
@@ -138,12 +143,15 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain):
         """
         results = self.llm_chain.apply(
             # FYI - this is parallelized and so it is fast.
-            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs]
+            [{self.document_variable_name: d.page_content, **kwargs} for d in docs],
+            callbacks=callbacks,
+        )
+        return self._process_results(
+            results, docs, token_max, callbacks=callbacks, **kwargs
         )
-        return self._process_results(results, docs, token_max, **kwargs)
 
     async def acombine_docs(
-        self, docs: List[Document], **kwargs: Any
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
     ) -> Tuple[str, dict]:
         """Combine documents in a map reduce manner.
 
@@ -152,15 +160,17 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain):
         """
         results = await self.llm_chain.aapply(
             # FYI - this is parallelized and so it is fast.
-            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs]
+            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs],
+            callbacks=callbacks,
         )
-        return self._process_results(results, docs, **kwargs)
+        return self._process_results(results, docs, callbacks=callbacks, **kwargs)
 
     def _process_results(
         self,
         results: List[Dict],
         docs: List[Document],
         token_max: int = 3000,
+        callbacks: Callbacks = None,
         **kwargs: Any,
     ) -> Tuple[str, dict]:
         question_result_key = self.llm_chain.output_key
@@ -173,7 +183,9 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain):
         num_tokens = length_func(result_docs, **kwargs)
 
         def _collapse_docs_func(docs: List[Document], **kwargs: Any) -> str:
-            return self._collapse_chain.run(input_documents=docs, **kwargs)
+            return self._collapse_chain.run(
+                input_documents=docs, callbacks=callbacks, **kwargs
+            )
 
         while num_tokens is not None and num_tokens > token_max:
             new_result_doc_list = _split_list_of_docs(
@@ -191,7 +203,9 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain):
             extra_return_dict = {"intermediate_steps": _results}
         else:
             extra_return_dict = {}
-        output = self.combine_document_chain.run(input_documents=result_docs, **kwargs)
+        output = self.combine_document_chain.run(
+            input_documents=result_docs, callbacks=callbacks, **kwargs
+        )
         return output, extra_return_dict
 
     @property

langchain/chains/combine_documents/map_rerank.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from pydantic import Extra, root_validator
 
+from langchain.callbacks.manager import Callbacks
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.docstore.document import Document
@@ -89,19 +90,22 @@ class MapRerankDocumentsChain(BaseCombineDocumentsChain):
                 )
         return values
 
-    def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
+    def combine_docs(
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
+    ) -> Tuple[str, dict]:
         """Combine documents in a map rerank manner.
 
         Combine by mapping first chain over all documents, then reranking the results.
         """
         results = self.llm_chain.apply_and_parse(
             # FYI - this is parallelized and so it is fast.
-            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs]
+            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs],
+            callbacks=callbacks,
         )
         return self._process_results(docs, results)
 
     async def acombine_docs(
-        self, docs: List[Document], **kwargs: Any
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
     ) -> Tuple[str, dict]:
         """Combine documents in a map rerank manner.
 
@@ -109,7 +113,8 @@ class MapRerankDocumentsChain(BaseCombineDocumentsChain):
         """
         results = await self.llm_chain.aapply_and_parse(
             # FYI - this is parallelized and so it is fast.
-            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs]
+            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs],
+            callbacks=callbacks,
         )
         return self._process_results(docs, results)
 

langchain/chains/combine_documents/refine.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Tuple
 
 from pydantic import Extra, Field, root_validator
 
+from langchain.callbacks.manager import Callbacks
 from langchain.chains.combine_documents.base import (
     BaseCombineDocumentsChain,
     format_document,
@@ -85,29 +86,31 @@ class RefineDocumentsChain(BaseCombineDocumentsChain):
                 )
         return values
 
-    def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
+    def combine_docs(
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
+    ) -> Tuple[str, dict]:
         """Combine by mapping first chain over all, then stuffing into final chain."""
         inputs = self._construct_initial_inputs(docs, **kwargs)
-        res = self.initial_llm_chain.predict(**inputs)
+        res = self.initial_llm_chain.predict(callbacks=callbacks, **inputs)
         refine_steps = [res]
         for doc in docs[1:]:
             base_inputs = self._construct_refine_inputs(doc, res)
             inputs = {**base_inputs, **kwargs}
-            res = self.refine_llm_chain.predict(**inputs)
+            res = self.refine_llm_chain.predict(callbacks=callbacks, **inputs)
             refine_steps.append(res)
         return self._construct_result(refine_steps, res)
 
     async def acombine_docs(
-        self, docs: List[Document], **kwargs: Any
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
     ) -> Tuple[str, dict]:
         """Combine by mapping first chain over all, then stuffing into final chain."""
         inputs = self._construct_initial_inputs(docs, **kwargs)
-        res = await self.initial_llm_chain.apredict(**inputs)
+        res = await self.initial_llm_chain.apredict(callbacks=callbacks, **inputs)
         refine_steps = [res]
         for doc in docs[1:]:
             base_inputs = self._construct_refine_inputs(doc, res)
             inputs = {**base_inputs, **kwargs}
-            res = await self.refine_llm_chain.apredict(**inputs)
+            res = await self.refine_llm_chain.apredict(callbacks=callbacks, **inputs)
             refine_steps.append(res)
         return self._construct_result(refine_steps, res)
 

langchain/chains/combine_documents/stuff.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic import Extra, Field, root_validator
 
+from langchain.callbacks.manager import Callbacks
 from langchain.chains.combine_documents.base import (
     BaseCombineDocumentsChain,
     format_document,
@@ -75,19 +76,21 @@ class StuffDocumentsChain(BaseCombineDocumentsChain):
         prompt = self.llm_chain.prompt.format(**inputs)
         return self.llm_chain.llm.get_num_tokens(prompt)
 
-    def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
-        """Stuff all documents into one prompt and pass to LLM."""
-        inputs = self._get_inputs(docs, **kwargs)
-        # Call predict on the LLM.
-        return self.llm_chain.predict(**inputs), {}
-
-    async def acombine_docs(
-        self, docs: List[Document], **kwargs: Any
+    def combine_docs(
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
     ) -> Tuple[str, dict]:
         """Stuff all documents into one prompt and pass to LLM."""
         inputs = self._get_inputs(docs, **kwargs)
         # Call predict on the LLM.
-        return await self.llm_chain.apredict(**inputs), {}
+        return self.llm_chain.predict(callbacks=callbacks, **inputs), {}
+
+    async def acombine_docs(
+        self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
+    ) -> Tuple[str, dict]:
+        """Stuff all documents into one prompt and pass to LLM."""
+        inputs = self._get_inputs(docs, **kwargs)
+        # Call predict on the LLM.
+        return await self.llm_chain.apredict(callbacks=callbacks, **inputs), {}
 
     @property
     def _chain_type(self) -> str:

langchain/chains/constitutional_ai/base.py

@@ -1,13 +1,14 @@
 """Chain for applying constitutional principles to the outputs of another chain."""
 from typing import Any, Dict, List, Optional
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.chains.constitutional_ai.models import ConstitutionalPrinciple
 from langchain.chains.constitutional_ai.principles import PRINCIPLES
 from langchain.chains.constitutional_ai.prompts import CRITIQUE_PROMPT, REVISION_PROMPT
 from langchain.chains.llm import LLMChain
 from langchain.prompts.base import BasePromptTemplate
-from langchain.schema import BaseLanguageModel
 
 
 class ConstitutionalChain(Chain):
@@ -18,14 +19,19 @@ class ConstitutionalChain(Chain):
 
             from langchain.llms import OpenAI
             from langchain.chains import LLMChain, ConstitutionalChain
+            from langchain.chains.constitutional_ai.models \
+                import ConstitutionalPrinciple
+
+            llm = OpenAI()
 
             qa_prompt = PromptTemplate(
                 template="Q: {question} A:",
                 input_variables=["question"],
             )
-            qa_chain = LLMChain(llm=OpenAI(), prompt=qa_prompt)
+            qa_chain = LLMChain(llm=llm, prompt=qa_prompt)
 
             constitutional_chain = ConstitutionalChain.from_llm(
+                llm=llm,
                 chain=qa_chain,
                 constitutional_principles=[
                     ConstitutionalPrinciple(
@@ -81,11 +87,16 @@ class ConstitutionalChain(Chain):
         """Defines the output keys."""
         return ["output"]
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         response = self.chain.run(**inputs)
         input_prompt = self.chain.prompt.format(**inputs)
 
-        self.callback_manager.on_text(
+        _run_manager.on_text(
             text="Initial response: " + response + "\n\n",
             verbose=self.verbose,
             color="yellow",
@@ -98,6 +109,7 @@ class ConstitutionalChain(Chain):
                 input_prompt=input_prompt,
                 output_from_model=response,
                 critique_request=constitutional_principle.critique_request,
+                callbacks=_run_manager.get_child(),
             )
             critique = self._parse_critique(
                 output_string=raw_critique,
@@ -111,22 +123,23 @@ class ConstitutionalChain(Chain):
                 critique_request=constitutional_principle.critique_request,
                 critique=critique,
                 revision_request=constitutional_principle.revision_request,
+                callbacks=_run_manager.get_child(),
             ).strip()
             response = revision
 
-            self.callback_manager.on_text(
+            _run_manager.on_text(
                 text=f"Applying {constitutional_principle.name}..." + "\n\n",
                 verbose=self.verbose,
                 color="green",
             )
 
-            self.callback_manager.on_text(
+            _run_manager.on_text(
                 text="Critique: " + critique + "\n\n",
                 verbose=self.verbose,
                 color="blue",
             )
 
-            self.callback_manager.on_text(
+            _run_manager.on_text(
                 text="Updated response: " + revision + "\n\n",
                 verbose=self.verbose,
                 color="yellow",

langchain/chains/conversational_retrieval/base.py

@@ -8,6 +8,11 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from pydantic import Extra, Field, root_validator
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
 from langchain.chains.base import Chain
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.combine_documents.stuff import StuffDocumentsChain
@@ -15,7 +20,7 @@ from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_
 from langchain.chains.llm import LLMChain
 from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts.base import BasePromptTemplate
-from langchain.schema import BaseLanguageModel, BaseMessage, BaseRetriever, Document
+from langchain.schema import BaseMessage, BaseRetriever, Document
 from langchain.vectorstores.base import VectorStore
 
 # Depending on the memory type and configuration, the chat history format may differ.
@@ -81,14 +86,20 @@ class BaseConversationalRetrievalChain(Chain):
     def _get_docs(self, question: str, inputs: Dict[str, Any]) -> List[Document]:
         """Get docs."""
 
-    def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         question = inputs["question"]
         get_chat_history = self.get_chat_history or _get_chat_history
         chat_history_str = get_chat_history(inputs["chat_history"])
 
         if chat_history_str:
+            callbacks = _run_manager.get_child()
             new_question = self.question_generator.run(
-                question=question, chat_history=chat_history_str
+                question=question, chat_history=chat_history_str, callbacks=callbacks
             )
         else:
             new_question = question
@@ -96,7 +107,9 @@ class BaseConversationalRetrievalChain(Chain):
         new_inputs = inputs.copy()
         new_inputs["question"] = new_question
         new_inputs["chat_history"] = chat_history_str
-        answer = self.combine_docs_chain.run(input_documents=docs, **new_inputs)
+        answer = self.combine_docs_chain.run(
+            input_documents=docs, callbacks=_run_manager.get_child(), **new_inputs
+        )
         if self.return_source_documents:
             return {self.output_key: answer, "source_documents": docs}
         else:
@@ -106,13 +119,19 @@ class BaseConversationalRetrievalChain(Chain):
     async def _aget_docs(self, question: str, inputs: Dict[str, Any]) -> List[Document]:
         """Get docs."""
 
-    async def _acall(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
         question = inputs["question"]
         get_chat_history = self.get_chat_history or _get_chat_history
         chat_history_str = get_chat_history(inputs["chat_history"])
         if chat_history_str:
+            callbacks = _run_manager.get_child()
             new_question = await self.question_generator.arun(
-                question=question, chat_history=chat_history_str
+                question=question, chat_history=chat_history_str, callbacks=callbacks
             )
         else:
             new_question = question
@@ -120,7 +139,9 @@ class BaseConversationalRetrievalChain(Chain):
         new_inputs = inputs.copy()
         new_inputs["question"] = new_question
         new_inputs["chat_history"] = chat_history_str
-        answer = await self.combine_docs_chain.arun(input_documents=docs, **new_inputs)
+        answer = await self.combine_docs_chain.arun(
+            input_documents=docs, callbacks=_run_manager.get_child(), **new_inputs
+        )
         if self.return_source_documents:
             return {self.output_key: answer, "source_documents": docs}
         else:

langchain/chains/graph_qa/base.py

@@ -1,10 +1,11 @@
 """Question answering over a graph."""
 from __future__ import annotations
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from pydantic import Field
 
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.chains.graph_qa.prompts import ENTITY_EXTRACTION_PROMPT, PROMPT
 from langchain.chains.llm import LLMChain
@@ -51,18 +52,25 @@ class GraphQAChain(Chain):
         qa_chain = LLMChain(llm=llm, prompt=qa_prompt)
         entity_chain = LLMChain(llm=llm, prompt=entity_prompt)
 
-        return cls(qa_chain=qa_chain, entity_extraction_chain=entity_chain, **kwargs)
+        return cls(
+            qa_chain=qa_chain,
+            entity_extraction_chain=entity_chain,
+            **kwargs,
+        )
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
         """Extract entities, look up info and answer question."""
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.input_key]
 
         entity_string = self.entity_extraction_chain.run(question)
 
-        self.callback_manager.on_text(
-            "Entities Extracted:", end="\n", verbose=self.verbose
-        )
-        self.callback_manager.on_text(
+        _run_manager.on_text("Entities Extracted:", end="\n", verbose=self.verbose)
+        _run_manager.on_text(
             entity_string, color="green", end="\n", verbose=self.verbose
         )
         entities = get_entities(entity_string)
@@ -70,9 +78,10 @@ class GraphQAChain(Chain):
         for entity in entities:
             triplets = self.graph.get_entity_knowledge(entity)
             context += "\n".join(triplets)
-        self.callback_manager.on_text("Full Context:", end="\n", verbose=self.verbose)
-        self.callback_manager.on_text(
-            context, color="green", end="\n", verbose=self.verbose
+        _run_manager.on_text("Full Context:", end="\n", verbose=self.verbose)
+        _run_manager.on_text(context, color="green", end="\n", verbose=self.verbose)
+        result = self.qa_chain(
+            {"question": question, "context": context},
+            callbacks=_run_manager.get_child(),
         )
-        result = self.qa_chain({"question": question, "context": context})
         return {self.output_key: result[self.qa_chain.output_key]}

langchain/chains/hyde/base.py

@@ -4,11 +4,12 @@ https://arxiv.org/abs/2212.10496
 """
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 from pydantic import Extra
 
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.chains.hyde.prompts import PROMPT_MAP
 from langchain.chains.llm import LLMChain
@@ -57,18 +58,27 @@ class HypotheticalDocumentEmbedder(Chain, Embeddings):
         embeddings = self.embed_documents(documents)
         return self.combine_embeddings(embeddings)
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
         """Call the internal llm chain."""
-        return self.llm_chain._call(inputs)
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        return self.llm_chain(inputs, callbacks=_run_manager.get_child())
 
     @classmethod
     def from_llm(
-        cls, llm: BaseLLM, base_embeddings: Embeddings, prompt_key: str
+        cls,
+        llm: BaseLLM,
+        base_embeddings: Embeddings,
+        prompt_key: str,
+        **kwargs: Any,
     ) -> HypotheticalDocumentEmbedder:
         """Load and use LLMChain for a specific prompt key."""
         prompt = PROMPT_MAP[prompt_key]
         llm_chain = LLMChain(llm=llm, prompt=prompt)
-        return cls(base_embeddings=base_embeddings, llm_chain=llm_chain)
+        return cls(base_embeddings=base_embeddings, llm_chain=llm_chain, **kwargs)
 
     @property
     def _chain_type(self) -> str:

langchain/chains/llm.py

@@ -5,11 +5,19 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from pydantic import Extra
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import (
+    AsyncCallbackManager,
+    AsyncCallbackManagerForChainRun,
+    CallbackManager,
+    CallbackManagerForChainRun,
+    Callbacks,
+)
 from langchain.chains.base import Chain
 from langchain.input import get_colored_text
 from langchain.prompts.base import BasePromptTemplate
 from langchain.prompts.prompt import PromptTemplate
-from langchain.schema import BaseLanguageModel, LLMResult, PromptValue
+from langchain.schema import LLMResult, PromptValue
 
 
 class LLMChain(Chain):
@@ -53,21 +61,40 @@ class LLMChain(Chain):
         """
         return [self.output_key]
 
-    def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
-        return self.apply([inputs])[0]
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        response = self.generate([inputs], run_manager=run_manager)
+        return self.create_outputs(response)[0]
 
-    def generate(self, input_list: List[Dict[str, Any]]) -> LLMResult:
+    def generate(
+        self,
+        input_list: List[Dict[str, Any]],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> LLMResult:
         """Generate LLM result from inputs."""
-        prompts, stop = self.prep_prompts(input_list)
-        return self.llm.generate_prompt(prompts, stop)
+        prompts, stop = self.prep_prompts(input_list, run_manager=run_manager)
+        return self.llm.generate_prompt(
+            prompts, stop, callbacks=run_manager.get_child() if run_manager else None
+        )
 
-    async def agenerate(self, input_list: List[Dict[str, Any]]) -> LLMResult:
+    async def agenerate(
+        self,
+        input_list: List[Dict[str, Any]],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> LLMResult:
         """Generate LLM result from inputs."""
         prompts, stop = await self.aprep_prompts(input_list)
-        return await self.llm.agenerate_prompt(prompts, stop)
+        return await self.llm.agenerate_prompt(
+            prompts, stop, callbacks=run_manager.get_child() if run_manager else None
+        )
 
     def prep_prompts(
-        self, input_list: List[Dict[str, Any]]
+        self,
+        input_list: List[Dict[str, Any]],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
     ) -> Tuple[List[PromptValue], Optional[List[str]]]:
         """Prepare prompts from inputs."""
         stop = None
@@ -79,7 +106,8 @@ class LLMChain(Chain):
             prompt = self.prompt.format_prompt(**selected_inputs)
             _colored_text = get_colored_text(prompt.to_string(), "green")
             _text = "Prompt after formatting:\n" + _colored_text
-            self.callback_manager.on_text(_text, end="\n", verbose=self.verbose)
+            if run_manager:
+                run_manager.on_text(_text, end="\n", verbose=self.verbose)
             if "stop" in inputs and inputs["stop"] != stop:
                 raise ValueError(
                     "If `stop` is present in any inputs, should be present in all."
@@ -88,7 +116,9 @@ class LLMChain(Chain):
         return prompts, stop
 
     async def aprep_prompts(
-        self, input_list: List[Dict[str, Any]]
+        self,
+        input_list: List[Dict[str, Any]],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
     ) -> Tuple[List[PromptValue], Optional[List[str]]]:
         """Prepare prompts from inputs."""
         stop = None
@@ -100,12 +130,8 @@ class LLMChain(Chain):
             prompt = self.prompt.format_prompt(**selected_inputs)
             _colored_text = get_colored_text(prompt.to_string(), "green")
             _text = "Prompt after formatting:\n" + _colored_text
-            if self.callback_manager.is_async:
-                await self.callback_manager.on_text(
-                    _text, end="\n", verbose=self.verbose
-                )
-            else:
-                self.callback_manager.on_text(_text, end="\n", verbose=self.verbose)
+            if run_manager:
+                await run_manager.on_text(_text, end="\n", verbose=self.verbose)
             if "stop" in inputs and inputs["stop"] != stop:
                 raise ValueError(
                     "If `stop` is present in any inputs, should be present in all."
@@ -113,15 +139,45 @@ class LLMChain(Chain):
             prompts.append(prompt)
         return prompts, stop
 
-    def apply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]:
+    def apply(
+        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
+    ) -> List[Dict[str, str]]:
         """Utilize the LLM generate method for speed gains."""
-        response = self.generate(input_list)
-        return self.create_outputs(response)
+        callback_manager = CallbackManager.configure(
+            callbacks, self.callbacks, self.verbose
+        )
+        run_manager = callback_manager.on_chain_start(
+            {"name": self.__class__.__name__},
+            {"input_list": input_list},
+        )
+        try:
+            response = self.generate(input_list, run_manager=run_manager)
+        except (KeyboardInterrupt, Exception) as e:
+            run_manager.on_chain_error(e)
+            raise e
+        outputs = self.create_outputs(response)
+        run_manager.on_chain_end({"outputs": outputs})
+        return outputs
 
-    async def aapply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]:
+    async def aapply(
+        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
+    ) -> List[Dict[str, str]]:
         """Utilize the LLM generate method for speed gains."""
-        response = await self.agenerate(input_list)
-        return self.create_outputs(response)
+        callback_manager = AsyncCallbackManager.configure(
+            callbacks, self.callbacks, self.verbose
+        )
+        run_manager = await callback_manager.on_chain_start(
+            {"name": self.__class__.__name__},
+            {"input_list": input_list},
+        )
+        try:
+            response = await self.agenerate(input_list, run_manager=run_manager)
+        except (KeyboardInterrupt, Exception) as e:
+            await run_manager.on_chain_error(e)
+            raise e
+        outputs = self.create_outputs(response)
+        await run_manager.on_chain_end({"outputs": outputs})
+        return outputs
 
     def create_outputs(self, response: LLMResult) -> List[Dict[str, str]]:
         """Create outputs from response."""
@@ -131,13 +187,19 @@ class LLMChain(Chain):
             for generation in response.generations
         ]
 
-    async def _acall(self, inputs: Dict[str, Any]) -> Dict[str, str]:
-        return (await self.aapply([inputs]))[0]
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        response = await self.agenerate([inputs], run_manager=run_manager)
+        return self.create_outputs(response)[0]
 
-    def predict(self, **kwargs: Any) -> str:
+    def predict(self, callbacks: Callbacks = None, **kwargs: Any) -> str:
         """Format prompt with kwargs and pass to LLM.
 
         Args:
+            callbacks: Callbacks to pass to LLMChain
             **kwargs: Keys to pass to prompt template.
 
         Returns:
@@ -148,12 +210,13 @@ class LLMChain(Chain):
 
                 completion = llm.predict(adjective="funny")
         """
-        return self(kwargs)[self.output_key]
+        return self(kwargs, callbacks=callbacks)[self.output_key]
 
-    async def apredict(self, **kwargs: Any) -> str:
+    async def apredict(self, callbacks: Callbacks = None, **kwargs: Any) -> str:
         """Format prompt with kwargs and pass to LLM.
 
         Args:
+            callbacks: Callbacks to pass to LLMChain
             **kwargs: Keys to pass to prompt template.
 
         Returns:
@@ -164,31 +227,33 @@ class LLMChain(Chain):
 
                 completion = llm.predict(adjective="funny")
         """
-        return (await self.acall(kwargs))[self.output_key]
+        return (await self.acall(kwargs, callbacks=callbacks))[self.output_key]
 
-    def predict_and_parse(self, **kwargs: Any) -> Union[str, List[str], Dict[str, str]]:
+    def predict_and_parse(
+        self, callbacks: Callbacks = None, **kwargs: Any
+    ) -> Union[str, List[str], Dict[str, str]]:
         """Call predict and then parse the results."""
-        result = self.predict(**kwargs)
+        result = self.predict(callbacks=callbacks, **kwargs)
         if self.prompt.output_parser is not None:
             return self.prompt.output_parser.parse(result)
         else:
             return result
 
     async def apredict_and_parse(
-        self, **kwargs: Any
+        self, callbacks: Callbacks = None, **kwargs: Any
     ) -> Union[str, List[str], Dict[str, str]]:
         """Call apredict and then parse the results."""
-        result = await self.apredict(**kwargs)
+        result = await self.apredict(callbacks=callbacks, **kwargs)
         if self.prompt.output_parser is not None:
             return self.prompt.output_parser.parse(result)
         else:
             return result
 
     def apply_and_parse(
-        self, input_list: List[Dict[str, Any]]
+        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
     ) -> Sequence[Union[str, List[str], Dict[str, str]]]:
         """Call apply and then parse the results."""
-        result = self.apply(input_list)
+        result = self.apply(input_list, callbacks=callbacks)
         return self._parse_result(result)
 
     def _parse_result(
@@ -202,10 +267,10 @@ class LLMChain(Chain):
             return result
 
     async def aapply_and_parse(
-        self, input_list: List[Dict[str, Any]]
+        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
     ) -> Sequence[Union[str, List[str], Dict[str, str]]]:
         """Call apply and then parse the results."""
-        result = await self.aapply(input_list)
+        result = await self.aapply(input_list, callbacks=callbacks)
         return self._parse_result(result)
 
     @property

langchain/chains/llm_bash/base.py

@@ -1,15 +1,23 @@
 """Chain that interprets a prompt and executes bash code to perform bash operations."""
-from typing import Dict, List
+from __future__ import annotations
 
-from pydantic import Extra
+import logging
+import warnings
+from typing import Any, Dict, List, Optional
 
+from pydantic import Extra, Field, root_validator
+
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.chains.llm_bash.prompt import PROMPT
 from langchain.prompts.base import BasePromptTemplate
-from langchain.schema import BaseLanguageModel
+from langchain.schema import OutputParserException
 from langchain.utilities.bash import BashProcess
 
+logger = logging.getLogger(__name__)
+
 
 class LLMBashChain(Chain):
     """Chain that interprets a prompt and executes bash code to perform bash operations.
@@ -18,14 +26,17 @@ class LLMBashChain(Chain):
         .. code-block:: python
 
             from langchain import LLMBashChain, OpenAI
-            llm_bash = LLMBashChain(llm=OpenAI())
+            llm_bash = LLMBashChain.from_llm(OpenAI())
     """
 
-    llm: BaseLanguageModel
-    """LLM wrapper to use."""
+    llm_chain: LLMChain
+    llm: Optional[BaseLanguageModel] = None
+    """[Deprecated] LLM wrapper to use."""
     input_key: str = "question"  #: :meta private:
     output_key: str = "answer"  #: :meta private:
     prompt: BasePromptTemplate = PROMPT
+    """[Deprecated]"""
+    bash_process: BashProcess = Field(default_factory=BashProcess)  #: :meta private:
 
     class Config:
         """Configuration for this pydantic object."""
@@ -33,6 +44,26 @@ class LLMBashChain(Chain):
         extra = Extra.forbid
         arbitrary_types_allowed = True
 
+    @root_validator(pre=True)
+    def raise_deprecation(cls, values: Dict) -> Dict:
+        if "llm" in values:
+            warnings.warn(
+                "Directly instantiating an LLMBashChain with an llm is deprecated. "
+                "Please instantiate with llm_chain or using the from_llm class method."
+            )
+            if "llm_chain" not in values and values["llm"] is not None:
+                prompt = values.get("prompt", PROMPT)
+                values["llm_chain"] = LLMChain(llm=values["llm"], prompt=prompt)
+        return values
+
+    @root_validator
+    def validate_prompt(cls, values: Dict) -> Dict:
+        if values["llm_chain"].prompt.output_parser is None:
+            raise ValueError(
+                "The prompt used by llm_chain is expected to have an output_parser."
+            )
+        return values
+
     @property
     def input_keys(self) -> List[str]:
         """Expect input key.
@@ -49,31 +80,45 @@ class LLMBashChain(Chain):
         """
         return [self.output_key]
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
-        llm_executor = LLMChain(prompt=self.prompt, llm=self.llm)
-        bash_executor = BashProcess()
-        self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose)
-
-        t = llm_executor.predict(question=inputs[self.input_key])
-        self.callback_manager.on_text(t, color="green", verbose=self.verbose)
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        _run_manager.on_text(inputs[self.input_key], verbose=self.verbose)
 
+        t = self.llm_chain.predict(
+            question=inputs[self.input_key], callbacks=_run_manager.get_child()
+        )
+        _run_manager.on_text(t, color="green", verbose=self.verbose)
         t = t.strip()
-        if t.startswith("```bash"):
-            # Split the string into a list of substrings
-            command_list = t.split("\n")
-            print(command_list)
+        try:
+            command_list = self.llm_chain.prompt.output_parser.parse(t)  # type: ignore[union-attr]
+        except OutputParserException as e:
+            _run_manager.on_chain_error(e, verbose=self.verbose)
+            raise e
 
-            # Remove the first and last substrings
-            command_list = [s for s in command_list[1:-1]]
-            output = bash_executor.run(command_list)
-
-            self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
-            self.callback_manager.on_text(output, color="yellow", verbose=self.verbose)
-
-        else:
-            raise ValueError(f"unknown format from LLM: {t}")
+        if self.verbose:
+            _run_manager.on_text("\nCode: ", verbose=self.verbose)
+            _run_manager.on_text(
+                str(command_list), color="yellow", verbose=self.verbose
+            )
+        output = self.bash_process.run(command_list)
+        _run_manager.on_text("\nAnswer: ", verbose=self.verbose)
+        _run_manager.on_text(output, color="yellow", verbose=self.verbose)
         return {self.output_key: output}
 
     @property
     def _chain_type(self) -> str:
         return "llm_bash_chain"
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLanguageModel,
+        prompt: BasePromptTemplate = PROMPT,
+        **kwargs: Any,
+    ) -> LLMBashChain:
+        llm_chain = LLMChain(llm=llm, prompt=prompt)
+        return cls(llm_chain=llm_chain, **kwargs)

langchain/chains/llm_bash/prompt.py

@@ -1,5 +1,11 @@
 # flake8: noqa
+from __future__ import annotations
+
+import re
+from typing import List
+
 from langchain.prompts.prompt import PromptTemplate
+from langchain.schema import BaseOutputParser, OutputParserException
 
 _PROMPT_TEMPLATE = """If someone asks you to perform a task, your job is to come up with a series of bash commands that will perform the task. There is no need to put "#!/bin/bash" in your answer. Make sure to reason step by step, using this format:
 
@@ -19,4 +25,36 @@ That is the format. Begin!
 
 Question: {question}"""
 
-PROMPT = PromptTemplate(input_variables=["question"], template=_PROMPT_TEMPLATE)
+
+class BashOutputParser(BaseOutputParser):
+    """Parser for bash output."""
+
+    def parse(self, text: str) -> List[str]:
+        if "```bash" in text:
+            return self.get_code_blocks(text)
+        else:
+            raise OutputParserException(
+                f"Failed to parse bash output. Got: {text}",
+            )
+
+    @staticmethod
+    def get_code_blocks(t: str) -> List[str]:
+        """Get multiple code blocks from the LLM result."""
+        code_blocks: List[str] = []
+        # Bash markdown code blocks
+        pattern = re.compile(r"```bash(.*?)(?:\n\s*)```", re.DOTALL)
+        for match in pattern.finditer(t):
+            matched = match.group(1).strip()
+            if matched:
+                code_blocks.extend(
+                    [line for line in matched.split("\n") if line.strip()]
+                )
+
+        return code_blocks
+
+
+PROMPT = PromptTemplate(
+    input_variables=["question"],
+    template=_PROMPT_TEMPLATE,
+    output_parser=BashOutputParser(),
+)

langchain/chains/llm_checker/base.py

@@ -1,10 +1,12 @@
 """Chain for question-answering with self-verification."""
+from __future__ import annotations
 
+import warnings
+from typing import Any, Dict, List, Optional
 
-from typing import Dict, List
-
-from pydantic import Extra
+from pydantic import Extra, root_validator
 
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.chains.llm_checker.prompt import (
@@ -18,6 +20,48 @@ from langchain.llms.base import BaseLLM
 from langchain.prompts import PromptTemplate
 
 
+def _load_question_to_checked_assertions_chain(
+    llm: BaseLLM,
+    create_draft_answer_prompt: PromptTemplate,
+    list_assertions_prompt: PromptTemplate,
+    check_assertions_prompt: PromptTemplate,
+    revised_answer_prompt: PromptTemplate,
+) -> SequentialChain:
+    create_draft_answer_chain = LLMChain(
+        llm=llm,
+        prompt=create_draft_answer_prompt,
+        output_key="statement",
+    )
+    list_assertions_chain = LLMChain(
+        llm=llm,
+        prompt=list_assertions_prompt,
+        output_key="assertions",
+    )
+    check_assertions_chain = LLMChain(
+        llm=llm,
+        prompt=check_assertions_prompt,
+        output_key="checked_assertions",
+    )
+    revised_answer_chain = LLMChain(
+        llm=llm,
+        prompt=revised_answer_prompt,
+        output_key="revised_statement",
+    )
+    chains = [
+        create_draft_answer_chain,
+        list_assertions_chain,
+        check_assertions_chain,
+        revised_answer_chain,
+    ]
+    question_to_checked_assertions_chain = SequentialChain(
+        chains=chains,
+        input_variables=["question"],
+        output_variables=["revised_statement"],
+        verbose=True,
+    )
+    return question_to_checked_assertions_chain
+
+
 class LLMCheckerChain(Chain):
     """Chain for question-answering with self-verification.
 
@@ -26,16 +70,21 @@ class LLMCheckerChain(Chain):
 
             from langchain import OpenAI, LLMCheckerChain
             llm = OpenAI(temperature=0.7)
-            checker_chain = LLMCheckerChain(llm=llm)
+            checker_chain = LLMCheckerChain.from_llm(llm)
     """
 
-    llm: BaseLLM
-    """LLM wrapper to use."""
+    question_to_checked_assertions_chain: SequentialChain
+
+    llm: Optional[BaseLLM] = None
+    """[Deprecated] LLM wrapper to use."""
     create_draft_answer_prompt: PromptTemplate = CREATE_DRAFT_ANSWER_PROMPT
+    """[Deprecated]"""
     list_assertions_prompt: PromptTemplate = LIST_ASSERTIONS_PROMPT
+    """[Deprecated]"""
     check_assertions_prompt: PromptTemplate = CHECK_ASSERTIONS_PROMPT
+    """[Deprecated]"""
     revised_answer_prompt: PromptTemplate = REVISED_ANSWER_PROMPT
-    """Prompt to use when questioning the documents."""
+    """[Deprecated] Prompt to use when questioning the documents."""
     input_key: str = "query"  #: :meta private:
     output_key: str = "result"  #: :meta private:
 
@@ -45,6 +94,34 @@ class LLMCheckerChain(Chain):
         extra = Extra.forbid
         arbitrary_types_allowed = True
 
+    @root_validator(pre=True)
+    def raise_deprecation(cls, values: Dict) -> Dict:
+        if "llm" in values:
+            warnings.warn(
+                "Directly instantiating an LLMCheckerChain with an llm is deprecated. "
+                "Please instantiate with question_to_checked_assertions_chain "
+                "or using the from_llm class method."
+            )
+            if (
+                "question_to_checked_assertions_chain" not in values
+                and values["llm"] is not None
+            ):
+                question_to_checked_assertions_chain = (
+                    _load_question_to_checked_assertions_chain(
+                        values["llm"],
+                        values.get(
+                            "create_draft_answer_prompt", CREATE_DRAFT_ANSWER_PROMPT
+                        ),
+                        values.get("list_assertions_prompt", LIST_ASSERTIONS_PROMPT),
+                        values.get("check_assertions_prompt", CHECK_ASSERTIONS_PROMPT),
+                        values.get("revised_answer_prompt", REVISED_ANSWER_PROMPT),
+                    )
+                )
+                values[
+                    "question_to_checked_assertions_chain"
+                ] = question_to_checked_assertions_chain
+        return values
+
     @property
     def input_keys(self) -> List[str]:
         """Return the singular input key.
@@ -61,43 +138,43 @@ class LLMCheckerChain(Chain):
         """
         return [self.output_key]
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.input_key]
 
-        create_draft_answer_chain = LLMChain(
-            llm=self.llm, prompt=self.create_draft_answer_prompt, output_key="statement"
+        output = self.question_to_checked_assertions_chain(
+            {"question": question}, callbacks=_run_manager.get_child()
         )
-        list_assertions_chain = LLMChain(
-            llm=self.llm, prompt=self.list_assertions_prompt, output_key="assertions"
-        )
-        check_assertions_chain = LLMChain(
-            llm=self.llm,
-            prompt=self.check_assertions_prompt,
-            output_key="checked_assertions",
-        )
-
-        revised_answer_chain = LLMChain(
-            llm=self.llm,
-            prompt=self.revised_answer_prompt,
-            output_key="revised_statement",
-        )
-
-        chains = [
-            create_draft_answer_chain,
-            list_assertions_chain,
-            check_assertions_chain,
-            revised_answer_chain,
-        ]
-
-        question_to_checked_assertions_chain = SequentialChain(
-            chains=chains,
-            input_variables=["question"],
-            output_variables=["revised_statement"],
-            verbose=True,
-        )
-        output = question_to_checked_assertions_chain({"question": question})
         return {self.output_key: output["revised_statement"]}
 
     @property
     def _chain_type(self) -> str:
         return "llm_checker_chain"
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLLM,
+        create_draft_answer_prompt: PromptTemplate = CREATE_DRAFT_ANSWER_PROMPT,
+        list_assertions_prompt: PromptTemplate = LIST_ASSERTIONS_PROMPT,
+        check_assertions_prompt: PromptTemplate = CHECK_ASSERTIONS_PROMPT,
+        revised_answer_prompt: PromptTemplate = REVISED_ANSWER_PROMPT,
+        **kwargs: Any,
+    ) -> LLMCheckerChain:
+        question_to_checked_assertions_chain = (
+            _load_question_to_checked_assertions_chain(
+                llm,
+                create_draft_answer_prompt,
+                list_assertions_prompt,
+                check_assertions_prompt,
+                revised_answer_prompt,
+            )
+        )
+        return cls(
+            question_to_checked_assertions_chain=question_to_checked_assertions_chain,
+            **kwargs,
+        )

langchain/chains/llm_math/base.py

@@ -1,16 +1,23 @@
 """Chain that interprets a prompt and executes python code to do math."""
+from __future__ import annotations
+
 import math
 import re
-from typing import Dict, List
+import warnings
+from typing import Any, Dict, List, Optional
 
 import numexpr
-from pydantic import Extra
+from pydantic import Extra, root_validator
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.chains.llm_math.prompt import PROMPT
 from langchain.prompts.base import BasePromptTemplate
-from langchain.schema import BaseLanguageModel
 
 
 class LLMMathChain(Chain):
@@ -20,13 +27,14 @@ class LLMMathChain(Chain):
         .. code-block:: python
 
             from langchain import LLMMathChain, OpenAI
-            llm_math = LLMMathChain(llm=OpenAI())
+            llm_math = LLMMathChain.from_llm(OpenAI())
     """
 
-    llm: BaseLanguageModel
-    """LLM wrapper to use."""
+    llm_chain: LLMChain
+    llm: Optional[BaseLanguageModel] = None
+    """[Deprecated] LLM wrapper to use."""
     prompt: BasePromptTemplate = PROMPT
-    """Prompt to use to translate to python if neccessary."""
+    """[Deprecated] Prompt to use to translate to python if necessary."""
     input_key: str = "question"  #: :meta private:
     output_key: str = "answer"  #: :meta private:
 
@@ -36,6 +44,19 @@ class LLMMathChain(Chain):
         extra = Extra.forbid
         arbitrary_types_allowed = True
 
+    @root_validator(pre=True)
+    def raise_deprecation(cls, values: Dict) -> Dict:
+        if "llm" in values:
+            warnings.warn(
+                "Directly instantiating an LLMMathChain with an llm is deprecated. "
+                "Please instantiate with llm_chain argument or using the from_llm "
+                "class method."
+            )
+            if "llm_chain" not in values and values["llm"] is not None:
+                prompt = values.get("prompt", PROMPT)
+                values["llm_chain"] = LLMChain(llm=values["llm"], prompt=prompt)
+        return values
+
     @property
     def input_keys(self) -> List[str]:
         """Expect input key.
@@ -68,15 +89,17 @@ class LLMMathChain(Chain):
         # Remove any leading and trailing brackets from the output
         return re.sub(r"^\[|\]$", "", output)
 
-    def _process_llm_result(self, llm_output: str) -> Dict[str, str]:
-        self.callback_manager.on_text(llm_output, color="green", verbose=self.verbose)
+    def _process_llm_result(
+        self, llm_output: str, run_manager: CallbackManagerForChainRun
+    ) -> Dict[str, str]:
+        run_manager.on_text(llm_output, color="green", verbose=self.verbose)
         llm_output = llm_output.strip()
         text_match = re.search(r"^```text(.*?)```", llm_output, re.DOTALL)
         if text_match:
             expression = text_match.group(1)
             output = self._evaluate_expression(expression)
-            self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
-            self.callback_manager.on_text(output, color="yellow", verbose=self.verbose)
+            run_manager.on_text("\nAnswer: ", verbose=self.verbose)
+            run_manager.on_text(output, color="yellow", verbose=self.verbose)
             answer = "Answer: " + output
         elif llm_output.startswith("Answer:"):
             answer = llm_output
@@ -86,30 +109,19 @@ class LLMMathChain(Chain):
             raise ValueError(f"unknown format from LLM: {llm_output}")
         return {self.output_key: answer}
 
-    async def _aprocess_llm_result(self, llm_output: str) -> Dict[str, str]:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_text(
-                llm_output, color="green", verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_text(
-                llm_output, color="green", verbose=self.verbose
-            )
+    async def _aprocess_llm_result(
+        self,
+        llm_output: str,
+        run_manager: AsyncCallbackManagerForChainRun,
+    ) -> Dict[str, str]:
+        await run_manager.on_text(llm_output, color="green", verbose=self.verbose)
         llm_output = llm_output.strip()
         text_match = re.search(r"^```text(.*?)```", llm_output, re.DOTALL)
         if text_match:
             expression = text_match.group(1)
             output = self._evaluate_expression(expression)
-            if self.callback_manager.is_async:
-                await self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
-                await self.callback_manager.on_text(
-                    output, color="yellow", verbose=self.verbose
-                )
-            else:
-                self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
-                self.callback_manager.on_text(
-                    output, color="yellow", verbose=self.verbose
-                )
+            await run_manager.on_text("\nAnswer: ", verbose=self.verbose)
+            await run_manager.on_text(output, color="yellow", verbose=self.verbose)
             answer = "Answer: " + output
         elif llm_output.startswith("Answer:"):
             answer = llm_output
@@ -119,31 +131,44 @@ class LLMMathChain(Chain):
             raise ValueError(f"unknown format from LLM: {llm_output}")
         return {self.output_key: answer}
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
-        llm_executor = LLMChain(
-            prompt=self.prompt, llm=self.llm, callback_manager=self.callback_manager
+    def _call(
+        self,
+        inputs: Dict[str, str],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        _run_manager.on_text(inputs[self.input_key])
+        llm_output = self.llm_chain.predict(
+            question=inputs[self.input_key],
+            stop=["```output"],
+            callbacks=_run_manager.get_child(),
         )
-        self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose)
-        llm_output = llm_executor.predict(
-            question=inputs[self.input_key], stop=["```output"]
-        )
-        return self._process_llm_result(llm_output)
+        return self._process_llm_result(llm_output, _run_manager)
 
-    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
-        llm_executor = LLMChain(
-            prompt=self.prompt, llm=self.llm, callback_manager=self.callback_manager
+    async def _acall(
+        self,
+        inputs: Dict[str, str],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
+        await _run_manager.on_text(inputs[self.input_key])
+        llm_output = await self.llm_chain.apredict(
+            question=inputs[self.input_key],
+            stop=["```output"],
+            callbacks=_run_manager.get_child(),
         )
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_text(
-                inputs[self.input_key], verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose)
-        llm_output = await llm_executor.apredict(
-            question=inputs[self.input_key], stop=["```output"]
-        )
-        return await self._aprocess_llm_result(llm_output)
+        return await self._aprocess_llm_result(llm_output, _run_manager)
 
     @property
     def _chain_type(self) -> str:
         return "llm_math_chain"
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLanguageModel,
+        prompt: BasePromptTemplate = PROMPT,
+        **kwargs: Any,
+    ) -> LLMMathChain:
+        llm_chain = LLMChain(llm=llm, prompt=prompt)
+        return cls(llm_chain=llm_chain, **kwargs)

langchain/chains/llm_requests.py

@@ -1,10 +1,11 @@
 """Chain that hits a URL and then uses an LLM to parse results."""
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 from pydantic import Extra, Field, root_validator
 
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains import LLMChain
 from langchain.chains.base import Chain
 from langchain.requests import TextRequestsWrapper
@@ -61,9 +62,14 @@ class LLMRequestsChain(Chain):
             )
         return values
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         from bs4 import BeautifulSoup
 
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         # Other keys are assumed to be needed for LLM prediction
         other_keys = {k: v for k, v in inputs.items() if k != self.input_key}
         url = inputs[self.input_key]
@@ -71,7 +77,9 @@ class LLMRequestsChain(Chain):
         # extract the text from the html
         soup = BeautifulSoup(res, "html.parser")
         other_keys[self.requests_key] = soup.get_text()[: self.text_length]
-        result = self.llm_chain.predict(**other_keys)
+        result = self.llm_chain.predict(
+            callbacks=_run_manager.get_child(), **other_keys
+        )
         return {self.output_key: result}
 
     @property