x

langchain-ai · Sep 14, 2024 · 2fdbc81 · 2fdbc81
1 parent 6ee8d4e
commit 2fdbc81
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/tests/unit_tests/test_serialization.py b/tests/unit_tests/test_serialization.py
@@ -7,12 +7,13 @@
 from langchain_core.documents.base import Document
 from langchain_core.messages import HumanMessage, HumanMessageChunk, SystemMessage
 from langchain_core.outputs import ChatGeneration
+from pydantic import BaseModel
+
 from langserve.serialization import (
     WellKnownLCObject,
     WellKnownLCSerializer,
     load_events,
 )
-from pydantic import BaseModel
 
 
 def test_document_serialization() -> None:

diff --git a/tests/unit_tests/test_server_client.py b/tests/unit_tests/test_server_client.py
@@ -122,17 +122,17 @@ def _replace_run_id_in_stream_resp(streamed_resp: str) -> str:
     return streamed_resp.replace(uuid, "<REPLACED>")
 
 
-def _null_run_id_recursively(decoded_response: Any) -> None:
+def _null_run_id_and_metadata_recursively(decoded_response: Any) -> None:
     """Recursively traverse the object and delete any keys called run_id"""
     if isinstance(decoded_response, dict):
         for key, value in decoded_response.items():
-            if key == "run_id":
+            if key in {"run_id", "__langserve_version"}:
                 decoded_response[key] = None
             else:
-                _null_run_id_recursively(value)
+                _null_run_id_and_metadata_recursively(value)
     elif isinstance(decoded_response, list):
         for item in decoded_response:
-            _null_run_id_recursively(item)
+            _null_run_id_and_metadata_recursively(item)
 
 
 @pytest.fixture(scope="module")
@@ -1141,15 +1141,15 @@ async def add_one(x: int) -> int:
         assert response.status_code == 200
         decoded_response = response.json()
         # Remove any run_id from the response recursively
-        _null_run_id_recursively(decoded_response)
+        _null_run_id_and_metadata_recursively(decoded_response)
         assert decoded_response == {
             "callback_events": [
                 {
                     "inputs": 1,
                     "kwargs": {"name": "add_one", "run_type": None},
                     "metadata": {
                         "__langserve_endpoint": "invoke",
-                        "__langserve_version": "0.3.0",
+                        "__langserve_version": None,
                         "__useragent": "python-httpx/0.27.2",
                     },
                     "parent_run_id": None,
@@ -1201,7 +1201,7 @@ async def add_one(x: int) -> int:
         assert response.status_code == 200
         decoded_response = response.json()
         # Remove any run_id from the response recursively
-        _null_run_id_recursively(decoded_response)
+        _null_run_id_and_metadata_recursively(decoded_response)
         del decoded_response["metadata"]["run_ids"]
         assert decoded_response == {
             "callback_events": [
@@ -1211,7 +1211,7 @@ async def add_one(x: int) -> int:
                         "kwargs": {"name": "add_one", "run_type": None},
                         "metadata": {
                             "__langserve_endpoint": "batch",
-                            "__langserve_version": "0.3.0",
+                            "__langserve_version": None,
                             "__useragent": "python-httpx/0.27.2",
                         },
                         "parent_run_id": None,
@@ -1236,7 +1236,7 @@ async def add_one(x: int) -> int:
                         "kwargs": {"name": "add_one", "run_type": None},
                         "metadata": {
                             "__langserve_endpoint": "batch",
-                            "__langserve_version": "0.3.0",
+                            "__langserve_version": None,
                             "__useragent": "python-httpx/0.27.2",
                         },
                         "parent_run_id": None,