diff --git a/docs/docs/how-tos/memory/semantic-search.ipynb b/docs/docs/how-tos/memory/semantic-search.ipynb
index 8e7a8d057..e7f4fc11f 100644
--- a/docs/docs/how-tos/memory/semantic-search.ipynb
+++ b/docs/docs/how-tos/memory/semantic-search.ipynb
@@ -208,7 +208,7 @@
     "from typing import Optional\n",
     "\n",
     "from langchain.chat_models import init_chat_model\n",
-    "from langchain_core.tools import InjectedToolArg\n",
+    "from langgraph.prebuilt import InjectedStore\n",
     "from langgraph.store.base import BaseStore\n",
     "from typing_extensions import Annotated\n",
     "\n",
@@ -232,7 +232,7 @@
     "    content: str,\n",
     "    *,\n",
     "    memory_id: Optional[uuid.UUID] = None,\n",
-    "    store: Annotated[BaseStore, InjectedToolArg],\n",
+    "    store: Annotated[BaseStore, InjectedStore],\n",
     "):\n",
     "    \"\"\"Upsert a memory in the database.\"\"\"\n",
     "    # The LLM can use this tool to store a new memory\n",