diff --git a/src/oss/python/migrate/langchain-v1.mdx b/src/oss/python/migrate/langchain-v1.mdx
index b013aa3e05..bdec3b0509 100644
--- a/src/oss/python/migrate/langchain-v1.mdx
+++ b/src/oss/python/migrate/langchain-v1.mdx
@@ -690,7 +690,7 @@ class Context:
 agent = create_agent(
     model=model,
     tools=tools,
-    context_schema=ContextSchema  # [!code highlight]
+    context_schema=Context  # [!code highlight]
 )
 
 result = agent.invoke(