Add Memorize tool (#11722)

- **Description:** Add `Memorize` tool
  - **Tag maintainer:** @hwchase17

This PR added a new tool `Memorize` so that an agent can use it to
fine-tune itself. This tool requires `TrainableLLM` introduced in #11721

DEMO:
6a9003d5db

![image](https://github.com/langchain-ai/langchain/assets/601530/d6f0cb45-54df-4dcf-b143-f8aefb1e76e3)

docs/docs/integrations/tools/memorize.ipynb

@@ -0,0 +1,202 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Memorize\n",
+    "\n",
+    "Fine-tuning LLM itself to memorize information using unsupervised learning.\n",
+    "\n",
+    "This tool requires LLMs that support fine-tuning. Currently, only `langchain.llms import GradientLLM` is supported."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from langchain.llms import GradientLLM\n",
+    "from langchain.chains import LLMChain\n",
+    "from langchain.agents import AgentExecutor, AgentType, initialize_agent, load_tools\n",
+    "from langchain.memory import ConversationBufferMemory"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set the Environment API Key\n",
+    "Make sure to get your API key from Gradient AI. You are given $10 in free credits to test and fine-tune different models."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from getpass import getpass\n",
+    "\n",
+    "\n",
+    "if not os.environ.get(\"GRADIENT_ACCESS_TOKEN\",None):\n",
+    "    # Access token under https://auth.gradient.ai/select-workspace\n",
+    "    os.environ[\"GRADIENT_ACCESS_TOKEN\"] = getpass(\"gradient.ai access token:\")\n",
+    "if not os.environ.get(\"GRADIENT_WORKSPACE_ID\",None):\n",
+    "    # `ID` listed in `$ gradient workspace list`\n",
+    "    # also displayed after login at at https://auth.gradient.ai/select-workspace\n",
+    "    os.environ[\"GRADIENT_WORKSPACE_ID\"] = getpass(\"gradient.ai workspace id:\")\n",
+    "if not os.environ.get(\"GRADIENT_MODEL_ADAPTER_ID\",None):\n",
+    "    # `ID` listed in `$ gradient model list --workspace-id \"$GRADIENT_WORKSPACE_ID\"`\n",
+    "    os.environ[\"GRADIENT_MODEL_ID\"] = getpass(\"gradient.ai model id:\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Optional: Validate your Environment variables ```GRADIENT_ACCESS_TOKEN``` and ```GRADIENT_WORKSPACE_ID``` to get currently deployed models."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create the `GradientLLM` instance\n",
+    "You can specify different parameters such as the model name, max tokens generated, temperature, etc."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = GradientLLM(\n",
+    "    model_id=os.environ['GRADIENT_MODEL_ID'],\n",
+    "    # # optional: set new credentials, they default to environment variables\n",
+    "    # gradient_workspace_id=os.environ[\"GRADIENT_WORKSPACE_ID\"],\n",
+    "    # gradient_access_token=os.environ[\"GRADIENT_ACCESS_TOKEN\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load tools"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tools = load_tools([\"memorize\"], llm=llm)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initiate the Agent"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent = initialize_agent(\n",
+    "    tools,\n",
+    "    llm,\n",
+    "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
+    "    verbose=True,\n",
+    "    # memory=ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run the agent\n",
+    "Ask the agent to memorize a piece of text."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
+      "\u001b[32;1m\u001b[1;3mI should memorize this fact.\n",
+      "Action: Memorize\n",
+      "Action Input: Zara T\u001b[0m\n",
+      "Observation: \u001b[36;1m\u001b[1;3mTrain complete. Loss: 1.6853971333333335\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3mI now know the final answer.\n",
+      "Final Answer: Zara Tubikova set a world\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'Zara Tubikova set a world'"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "agent.run(\"Please remember the fact in detail:\\nWith astonishing dexterity, Zara Tubikova set a world record by solving a 4x4 Rubik's Cube variation blindfolded in under 20 seconds, employing only their feet.\")"
+   ]
+  }
+ ],
+ "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.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "a0a0263b650d907a3bfe41c0f8d6a63a071b884df3cfdc1579f00cdc1aed6b03"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

libs/langchain/langchain/agents/load_tools.py

@@ -58,6 +58,7 @@ from langchain.tools.wolfram_alpha.tool import WolframAlphaQueryRun
 from langchain.tools.openweathermap.tool import OpenWeatherMapQueryRun
 from langchain.tools.dataforseo_api_search import DataForSeoAPISearchRun
 from langchain.tools.dataforseo_api_search import DataForSeoAPISearchResults
+from langchain.tools.memorize.tool import Memorize
 from langchain.utilities.arxiv import ArxivAPIWrapper
 from langchain.utilities.golden_query import GoldenQueryAPIWrapper
 from langchain.utilities.pubmed import PubMedAPIWrapper
@@ -327,6 +328,10 @@ def _get_eleven_labs_text2speech(**kwargs: Any) -> BaseTool:
     return ElevenLabsText2SpeechTool(**kwargs)
 
 
+def _get_memorize(llm: BaseLanguageModel, **kwargs: Any) -> BaseTool:
+    return Memorize(llm=llm)
+
+
 def _get_google_cloud_texttospeech(**kwargs: Any) -> BaseTool:
     return GoogleCloudTextToSpeechTool(**kwargs)
 
@@ -338,6 +343,7 @@ _EXTRA_LLM_TOOLS: Dict[
     "news-api": (_get_news_api, ["news_api_key"]),
     "tmdb-api": (_get_tmdb_api, ["tmdb_bearer_token"]),
     "podcast-api": (_get_podcast_api, ["listen_api_key"]),
+    "memorize": (_get_memorize, []),
 }
 _EXTRA_OPTIONAL_TOOLS: Dict[str, Tuple[Callable[[KwArg(Any)], BaseTool], List[str]]] = {
     "wolfram-alpha": (_get_wolfram_alpha, ["wolfram_alpha_appid"]),

libs/langchain/langchain/tools/memorize/__init__.py

@@ -0,0 +1,5 @@
+"""Unsupervised learning based memorization."""
+
+from langchain.tools.memorize.tool import Memorize
+
+__all__ = ["Memorize"]

libs/langchain/langchain/tools/memorize/tool.py

@@ -0,0 +1,57 @@
+from abc import abstractmethod
+from typing import Any, Optional, Protocol, Sequence, runtime_checkable
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
+)
+from langchain.llms.gradient_ai import TrainResult
+from langchain.pydantic_v1 import Field
+from langchain.tools.base import BaseTool
+
+
+@runtime_checkable
+class TrainableLLM(Protocol):
+    @abstractmethod
+    def train_unsupervised(
+        self,
+        inputs: Sequence[str],
+        **kwargs: Any,
+    ) -> TrainResult:
+        ...
+
+    @abstractmethod
+    async def atrain_unsupervised(
+        self,
+        inputs: Sequence[str],
+        **kwargs: Any,
+    ) -> TrainResult:
+        ...
+
+
+class Memorize(BaseTool):
+    name: str = "Memorize"
+    description: str = (
+        "Useful whenever you observed novel information "
+        "from previous conversation history, "
+        "i.e., another tool's action outputs or human comments. "
+        "The action input should include observed information in detail, "
+        "then the tool will fine-tune yourself to remember it."
+    )
+    llm: TrainableLLM = Field()
+
+    def _run(
+        self,
+        information_to_learn: str,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+    ) -> str:
+        train_result = self.llm.train_unsupervised((information_to_learn,))
+        return f"Train complete. Loss: {train_result['loss']}"
+
+    async def _arun(
+        self,
+        information_to_learn: str,
+        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
+    ) -> str:
+        train_result = await self.llm.atrain_unsupervised((information_to_learn,))
+        return f"Train complete. Loss: {train_result['loss']}"