From 9be08a1956fb2cd3e9b146fca246f160780ba8ae Mon Sep 17 00:00:00 2001
From: Paresh Chiramel <72655959+pareshchiramel@users.noreply.github.com>
Date: Tue, 2 Jan 2024 11:25:42 +1100
Subject: [PATCH] Update _retrieve_ref inside json_schema.py to include an
 isdigit() check (#14745)

- **Description:** Update _retrieve_ref inside json_schema.py to include
an isdigit() check
- **Issue:** This library is used inside dereference_refs inside
langchain_community.agent_toolkits.openapi.spec. When I read in a yaml
file which has references for "400", "401" etc; the line "out =
out[component]" causes a KeyError. The isdigit() check ensures that if
it is an integer like "400" or "401"; it converts it into integer before
using it as a key to prevent the error.
  - **Dependencies:** No dependencies
  - **Tag maintainer:** @baskaryan

---------

Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
---
 libs/core/langchain_core/utils/json_schema.py |  5 ++-
 .../unit_tests/utils/test_json_schema.py      | 32 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/libs/core/langchain_core/utils/json_schema.py b/libs/core/langchain_core/utils/json_schema.py
index 9628f9e521b..95313b3f954 100644
--- a/libs/core/langchain_core/utils/json_schema.py
+++ b/libs/core/langchain_core/utils/json_schema.py
@@ -13,7 +13,10 @@ def _retrieve_ref(path: str, schema: dict) -> dict:
         )
     out = schema
     for component in components[1:]:
-        out = out[component]
+        if component.isdigit():
+            out = out[int(component)]
+        else:
+            out = out[component]
     return deepcopy(out)
 
 
diff --git a/libs/core/tests/unit_tests/utils/test_json_schema.py b/libs/core/tests/unit_tests/utils/test_json_schema.py
index 3b0c7b4fe51..6f0a00fd99d 100644
--- a/libs/core/tests/unit_tests/utils/test_json_schema.py
+++ b/libs/core/tests/unit_tests/utils/test_json_schema.py
@@ -149,3 +149,35 @@ def test_dereference_refs_remote_ref() -> None:
     }
     with pytest.raises(ValueError):
         dereference_refs(schema)
+
+
+def test_dereference_refs_integer_ref() -> None:
+    schema = {
+        "type": "object",
+        "properties": {
+            "error_400": {"$ref": "#/$defs/400"},
+        },
+        "$defs": {
+            400: {
+                "type": "object",
+                "properties": {"description": "Bad Request"},
+            },
+        },
+    }
+    expected = {
+        "type": "object",
+        "properties": {
+            "error_400": {
+                "type": "object",
+                "properties": {"description": "Bad Request"},
+            },
+        },
+        "$defs": {
+            400: {
+                "type": "object",
+                "properties": {"description": "Bad Request"},
+            },
+        },
+    }
+    actual = dereference_refs(schema)
+    assert actual == expected