Update patch_langmem.py

patch_langmem.py

@@ -1,48 +1,43 @@
 import os
-import re
-import sys
-from pathlib import Path
 
 def patch_langmem_for_python310():
-    if sys.version_info >= (3, 11):
-        print("✅ Python >=3.11 detected, no patching needed.")
-        return
-
-    try:
-        import langmem
-    except ImportError:
-        print("❌ langmem not installed.")
-        return
-
-    # Get langmem install path
-    langmem_path = Path(langmem.__file__).parent
-    extraction_path = langmem_path / "knowledge" / "extraction.py"
-
-    if not extraction_path.exists():
-        print("❌ Could not locate langmem extraction.py to patch.")
-        return
-
-    with open(extraction_path, "r", encoding="utf-8") as f:
-        content = f.read()
-
-    # Check if patch already applied
-    if "from typing_extensions import NotRequired" in content:
-        print("✅ langmem already patched.")
-        return
-
-    # Apply patch
-    content = re.sub(
-        r"import typing(\n|.*)",
-        r"import typing\nfrom typing_extensions import NotRequired\1",
-        content,
-    )
-    content = content.replace(
-        "existing: typing.NotRequired",
-        "existing: NotRequired"
-    )
-
-    with open(extraction_path, "w", encoding="utf-8") as f:
-        f.write(content)
-
-    print("✅ Patched langmem to support Python 3.10.")
-
+    # Path where langmem is installed inside the Docker container on HF Spaces
+    base_path = "/usr/local/lib/python3.10/site-packages/langmem/knowledge"
+
+    extraction_path = os.path.join(base_path, "extraction.py")
+    tools_path = os.path.join(base_path, "tools.py")
+
+    # === PATCH extraction.py ===
+    if os.path.exists(extraction_path):
+        with open(extraction_path, "r", encoding="utf-8") as file:
+            content = file.read()
+
+        if "typing.NotRequired" in content:
+            content = content.replace("typing.NotRequired", "NotRequired")
+            if "from typing_extensions import NotRequired" not in content:
+                content = content.replace("import uuid", "import uuid\nfrom typing_extensions import NotRequired")
+
+            with open(extraction_path, "w", encoding="utf-8") as file:
+                file.write(content)
+            print("✅ Patched extraction.py")
+        else:
+            print("✅ extraction.py already patched or not using NotRequired")
+
+    # === PATCH tools.py ===
+    if os.path.exists(tools_path):
+        with open(tools_path, "r", encoding="utf-8") as file:
+            content = file.read()
+
+        if "typing.Literal[*actions_permitted]" in content:
+            fixed = '''typing.Optional[
+        tuple[typing.Literal["create", "update", "delete"], ...]
+    ] = ("create", "update", "delete"),'''
+            content = content.replace("typing.Literal[*actions_permitted]", fixed)
+
+            with open(tools_path, "w", encoding="utf-8") as file:
+                file.write(content)
+            print("✅ Patched tools.py")
+        else:
+            print("✅ tools.py already patched or does not need fix")
+
+patch_langmem_for_python310()