Return run IDs from invoke/batch endpoints

langserve/schema.py

@@ -1,3 +1,5 @@
+from typing import List
+
 try:
     from pydantic.v1 import BaseModel
 except ImportError:
@@ -21,3 +23,39 @@ class CustomUserType(BaseModel):
         the server will keep the decoded type as a pydantic model instead
         of converting it into a dict.
     """
+
+
+class SharedResponseMetadata(BaseModel):
+    """
+    Any response metadata should inherit from this class. Response metadata
+    represents non-output data that may be useful to some clients, but
+    ignorable to most. For example, the run_ids associated with each run
+    kicked off by the associated request.
+
+    SharedResponseMetadata is an abstraction to represent any metadata
+    representing a LangServe response shared across all outputs in said
+    response.
+    """
+
+    pass
+
+
+class SingletonResponseMetadata(SharedResponseMetadata):
+    """
+    Represents response metadata used for just single input/output LangServe
+    responses.
+    """
+
+    # Represents the parent run id for a given request
+    run_id: str
+
+
+class BatchResponseMetadata(SharedResponseMetadata):
+    """
+    Represents response metadata used for batches of input/output LangServe
+    responses.
+    """
+
+    # Represents each parent run id for a given request, in
+    # the same order in which they were received
+    run_ids: List[str]

langserve/server.py

@@ -34,7 +34,11 @@
 
 from langserve.callbacks import AsyncEventAggregatorCallback, CallbackEventDict
 from langserve.lzstring import LZString
-from langserve.schema import CustomUserType
+from langserve.schema import (
+    BatchResponseMetadata,
+    CustomUserType,
+    SingletonResponseMetadata,
+)
 
 try:
     from pydantic.v1 import BaseModel, create_model
@@ -292,6 +296,23 @@ def _with_validation_error_translation() -> Generator[None, None, None]:
         raise RequestValidationError(e.errors(), body=e.model)
 
 
+def _get_base_run_id_as_str(
+    event_aggregator: AsyncEventAggregatorCallback,
+) -> Optional[str]:
+    """
+    Uses `event_aggregator` to determine the base run ID for a given run. Returns
+    the run_id as a string, or None if it does not exist.
+    """
+    # The first run in the callback_events list corresponds to the
+    # overall trace for request
+    if event_aggregator.callback_events and event_aggregator.callback_events[0].get(
+        "run_id"
+    ):
+        return str(event_aggregator.callback_events[0].get("run_id"))
+    else:
+        raise AssertionError("No run_id found for the given run")
+
+
 # PUBLIC API
 
 
@@ -494,6 +515,9 @@ async def invoke(
             # Callbacks are scrubbed and exceptions are converted to serializable format
             # before returned in the response.
             callback_events=callback_events,
+            metadata=SingletonResponseMetadata(
+                run_id=_get_base_run_id_as_str(event_aggregator)
+            ),
         )
 
     @app.post(
@@ -585,6 +609,9 @@ async def batch(
         return BatchResponse(
             output=well_known_lc_serializer.dumpd(output),
             callback_events=callback_events,
+            metadata=BatchResponseMetadata(
+                run_ids=[_get_base_run_id_as_str(agg) for agg in aggregators]
+            ),
         )
 
     @app.post(

langserve/validation.py

@@ -27,6 +27,8 @@
     RunInfo,
 )
 
+from langserve.schema import BatchResponseMetadata, SingletonResponseMetadata
+
 try:
     from pydantic.v1 import BaseModel, Field, create_model
 except ImportError:
@@ -207,6 +209,16 @@ def create_invoke_response_model(
             List[CallbackEvent],
             Field(..., description="Callback events generated by the server side."),
         ),
+        metadata=(
+            SingletonResponseMetadata,
+            Field(
+                ...,
+                description=(
+                    "Metadata about the response that may be useful to "
+                    "specific clients"
+                ),
+            ),
+        ),
     )
     invoke_response_type.update_forward_refs()
     return invoke_response_type
@@ -241,6 +253,15 @@ def create_batch_response_model(
                 ),
             ),
         ),
+        metadata=(
+            BatchResponseMetadata,
+            Field(
+                ...,
+                description=(
+                    "Metadata about the response that may be useful to specific clients"
+                ),
+            ),
+        ),
     )
     batch_response_type.update_forward_refs()
     return batch_response_type

tests/unit_tests/test_server_client.py

@@ -1,6 +1,7 @@
 """Test the server and client together."""
 import asyncio
 import json
+import uuid
 from asyncio import AbstractEventLoop
 from contextlib import asynccontextmanager, contextmanager
 from typing import Any, Dict, Iterator, List, Optional, Union
@@ -1438,3 +1439,31 @@ async def test_using_router() -> None:
     )
 
     app.include_router(router)
+
+
+def _is_valid_uuid(uuid_as_str: str) -> bool:
+    try:
+        uuid.UUID(str(uuid_as_str))
+        return True
+    except ValueError:
+        return False
+
+
+@pytest.mark.asyncio
+async def test_invoke_returns_run_id(app: FastAPI) -> None:
+    """Test the server directly via HTTP requests."""
+    async with get_async_test_client(app, raise_app_exceptions=True) as async_client:
+        response = await async_client.post("/invoke", json={"input": 1})
+        run_id = response.json()["metadata"]["run_id"]
+        assert _is_valid_uuid(run_id)
+
+
+@pytest.mark.asyncio
+async def test_batch_returns_run_id(app: FastAPI) -> None:
+    """Test the server directly via HTTP requests."""
+    async with get_async_test_client(app, raise_app_exceptions=True) as async_client:
+        response = await async_client.post("/batch", json={"inputs": [1, 2]})
+        run_ids = response.json()["metadata"]["run_ids"]
+        assert len(run_ids) == 2
+        for run_id in run_ids:
+            assert _is_valid_uuid(run_id)